TinyGSM: achieving >80% on GSM8k with small language models

Introduction

One fascinating phenomenon regarding large language models (LLMs) is the emergence of capbilities as both the model and dataset sizes scale up (Wei et al., 2022a; Chan et al., 2022). Among many capabilities, mathematical reasoning is one particular aspect that has received tremendous attention Lewkowycz et al. (2022); Lightman et al. (2023). However, it is unclear to what extend scale is a necessity for mathematical reasoning, and the potential of small language models (SLMs) remains largely under-explored.

In this work, we push the boundaries of SLMs’ math reasoning capabilities. As a first step towards general mathematics, our primary testing ground is grade school math problems, to solve which require both mathematical understanding and language comprehension. The gold-standard benchmark in this regime is GSM8K (Cobbe et al., 2021), a collection of 8.8K grade-school math word problems (with a 7k-1k train-test split) that involve $2$ to $11$ reasoning steps. GSM8K has been widely regarded to be challenging for LLMs. Even though the questions appear rather simple for humans, there have been few models that achieve $>80\%$, and they are commonly of prohibitive sizes, e.g. $34$B and above (see Table 1).

Our goal is to break the $80\%$ barrier on GSM8K while keeping the model size friendly. As previous work shows (Gunasekar et al., 2023; Li et al., 2023; Eldan & Li, 2023), training data quality is one of the most important factors for enhancing performance of small models. In particular, prompt-engineered synthetic data generation from gigantic models such as GPT-3.5/4 enjoys the clear advantage of desirable data hygiene and controllable diversity. This constituents a teacher-student scenario where the student learns from teacher’s generations. On the tasks that the model model already excels at, their guided generations remain one of the highest quality data one can collect for training significantly smaller student models. It is also understood that the student model’s performance likely ends up inferior than the teacher, and may fall far short especially when the student is considerably smaller than the teacher (Mirzadeh et al., 2019; Gudibande et al., 2023) —after all, the teacher places an information-theoretic bottleneck on the student.

To our surprise, in the case of GSM8K, we are able to bridge the performance gap between the student and teacher, by utilizing a tiny amount of labeled real data (the original GSM8K training set of 7k questions) to train an independent verifier model. At test time, the verifier score and select among multiple candidate answers generated from the student, and then we output the highest score generation as the final submission. Note the idea of using a verifier is proposed by the seminal GSM8K paper (Cobbe et al., 2021), and here we demonstrate its power of bridging the teacher-student gap, and we conduct a more thorough examination of factors affecting its efficacy.

The contributions of our work are the following:

We introduce TinyGSM, a synthetic dataset containing GSM8K-style math word problems paired with Python solutions, generated fully by GPT-3.5-turbo. TinyGSM consists of 12.3M questions which amount to 1.8B tokens. We demonstrate TinyGSM’s high-quality by finetuning the Phi-1.5 1.3B model (before the use of verifiers) which improves its accuracy from 44.6% to 68.2% on the GSM8K test set. Notably, our smallest 125M model can also achieve 63.1% after finetuning on TinyGSM.

We demonstrate the power of verifiers on small-scale models. When integrated with a verifier for scoring generations, our models, named Phi-GSM models, achieve performance on par with other open source models that are orders of magnitude larger. In particular, our 1.3B model achieves 81.5% accuracy on GSM8K, as shown in Figure 1. This marks a new state-of-the-arts on billion-parameter-scale models, significantly outperforming existing open-source models and even rivaling the 77.4% accuracy of GPT-3.5, from which TinyGSM is generated. For verifier training, we identify data diversity as a crucial element for a verifier’s success, and find that the scaling of the verifier may be more effective than scaling of the generator: while scaling up from a 125M generator to a 1.3B generator only gives a 5.1% increase in performance (Table 1), scaling up the verifier from 125M to 1.3B leads to a 7.2% performance boost (Figure 4).

Related works

Distilling from synthetic data: While scaling up has been a useful strategy, it is possible to outpace conventional scaling laws by better use of data (Sorscher et al., 2022). In the data-scarce case, quality synthetic data serves as an effective workaround (Eldan & Li, 2023; Gunasekar et al., 2023), and the scaling in dataset size can compensate for a small model size (Edelman et al., 2023). Additionally, our work uses samples in the true distribution (i.e. the GSM8K train set) differently: given the small dataset size, we believe that the most sample-efficient way to utilize the true train set is to train a verifier—while the 7.4k samples in the GSM8K training set is too small for language model finetuning, it is sufficient for training a good quality verifier that provides 10% performance boost. While there have been potential concerns of learning from synthetic data such as loosing diversity or having a drifted distribution mean (Alemohammad et al., 2023; Shumailov et al., 2023), Alemohammad et al. (2023) showed that such degradation can be avoided by including fresh samples from the true distribution during training.

Math word problem datasets GSM8K (Cobbe et al., 2021) has been the most common used math word problem dataset for its quality and size. In comparison, earlier datasets such as MAWPS (Koncel-Kedziorski et al., 2016), ASDiv (Miao et al., 2020) and SVAMP (Patel et al., 2021) are either much smaller in size or of less difficulty. However, GSM8K questions are too clean to test for robustness. Motivated by the observation that language models are not robust to the presence of irrelevant context, Shi et al. (2023a) proposed GSM-IC (irrelevant context). Another problem is the GSM8K dataset itself is still too small for training language models. (Ni et al., 2023b) addressed this with self-sampled data. In a work concurrent to ours, Yu et al. (2023b) bootstraps an original dataset using various augmentation techniques, such as generating multiple answers for the solution, question rephrasing, and backward reasoning. The proposed MetaMath dataset consists of 40000 questions from GSM8K and MATH (Hendrycks et al., 2021). In comparison, TinyGSM is significantly larger, encompassing 12.3M questions (or equivalently 1.8B tokens).

Leveraging multiple generations: An important component of our method is to leverage multiple generation. This idea has been proven successful in many prior works. A notable example is “self-consistency” (Wang et al., 2022), which selects the most frequent response among candidates and integrates well with other methods such as progressive-hint prompting (Zheng et al., 2023) and model selection (Zhao et al., 2023). However, self-consistency was not particularly helpful in our setup as mentioned in Section 4.2. More related to and a direct inspiration of our work is Cobbe et al. (2021), which uses a verifier to select the best response among 100 candidates, leading to an 20% accuracy boost. Our work conducts a more thorough study on the design choices of the verifier, including data diversity and the effect of verifier sizes. Another design choice orthogonal to ours is the supervision signals, such as outcome-based supervision versus process supervision (Lightman et al., 2023).

Learning from partial or process supervision: In our experiments, we evaluate on the final accuracy only but train on full programs. Prior work has studied the effect of process versus outcome based supervision. Process-based supervision is shown to be particularly helpful for complex math problems (Lightman et al., 2023), though for general problems one needs to consider a cost-efficacy tradeoff (Uesato et al., 2022). When process supervision is not available, Ni et al. (2023b) proposed to learn from “self-sampled” solutions, which allows the model to learn from partially correct self-generated solutions selected based on the execution trace.

Self-improvement: Several works have explored the idea of “self-improvement” where a model evaluates and corrects its own generations, mostly relying on the self-debugging ability of GPT4. Examples include “self-refine” (Madaan et al., 2023) and “self-verify” (Weng et al., 2022; Zhou et al., 2023), both of which ask the model to iteratively provide feedback or verifications on its own generations and refine if needed. However, such self-improvement abilities have not been discovered in small language models. This motivated our use of a separate verifier model, which is initialized from the generative model but needs to be fully finetuned for verification.

Prompt-based methods: Prompt-based methods, which find prompts to improve the later conditional generations, are particularly effective for large models. Examples include in-context learning (Brown et al., 2020), where the model learns to perform novel tasks from few-shot examples provided in the prompt, as well as Chain-of-Thought (Wei et al., 2022b), which shows that explicitly generating intermediate steps can significantly help with reasoning abilities. However, similar to self-improvements, prompting is targeted at large language models and do not apply for SLMs.

The TinyGSM dataset

Our objective is to assess the capability of a small language model (SLM) on mathematical reasoning. Ideally, enhancing this mathematical reasoning ability should not compromise the model’s competence in language comprehension. This makes math word problems, which necessitate both mathematical and language understanding, a suitable test ground. We focus on the GSM8K dataset (Cobbe et al., 2021), consisting of around 8k grade-school math word problems. The math concepts in the dataset are elementary and within standard grade-school curricula, but the challenges posed by the natural language problem statement introduce an additional layer of complexity to the task.

Despite the high quality, the GSM8K training set only contains 7473 problems, which is too small for training a reasonably sized language model (Ni et al., 2023a). To alleviate the size issue, we augment the GSM8K training set using GPT-3.5-turbo generated synthetic problems.

We prompt GPT-3.5-turbo to generate problem variants similar to a given question (but not the solution) randomly sampled from the GSM8K training set. Each problem variant contains both a question and the corresponding solution written in Python, as shown in Figure 2.Note that the generated problems may be mathematically valid yet violating common sense. For example, some quantities may not be integers. Using code allows us to leverage a Python interpreter, circumventing language models’ known limitation regarding numerical calculations and code execution.

To enhance robustness, we also generated synthetic problems whose questions contain irrelevant information. This is achieved by augmenting the GSM-IC dataset (Shi et al., 2023a), which is an augmentation of GSM8K specifically designed to introduce irrelevant context (IC) to the question statement. These GSM-IC variants constitute to approximately one third of TinyGSM.

The resulting synthetic dataset contains 12.3M problems (i.e. question-solution pairs) This corresponds to 1.8B tokens, which costs around $3600 to generate according to GPT commercial pricing. with, based on the original 7.4k training set questions and their IC variants. For each question in the GSM8K train set, the prompt based on this question is shared across API calls, and the source of randomness comes entirely from the generation process. To encourage diversity, we use temperature sampling and specify in the prompt to encourage the problem variants to be grammatically diverse and contain multiple steps; the exact prompts are provided in Figure 3 and in Section A.1.

Filtering

To ensure the quality of the synthetic data in TinyGSM, we filter out problems that are too short or do not contain numbers, as well as code solutions which are not executable. Note that we do not check for the correctness of the question or the generated solutions, since the “ground truth” solution is not available. Given the effectiveness of self-consistency (Wang et al., 2022), one might want to filter the problems by keeping the ones which have majority vote only. We did not adopt this strategy since we find that GPT-3.5-turbo’s generations are only consistent on easy problems “Easy” problems refer to the ones for which a 350M model, trained on a part of TinyGSM, already produces same final answer as GPT-3.5-turbo. For example, for an early version of our 350M model, the model only achieves around 50% on the GSM8K test set, but can achieve more than 87% on synthetic questions with consistent answers. In other words, adding more easy problems like these will not help our 350M model bridge the performance gap between itself and GPT-3.5-turbo., hence such consistency filtering will remove challenging problems, resulting in a dataset that is too easy to be useful.

Solving grade school math with small language models

The 1.3B version of our phi-GSM models is able to achieve 81.5% accuracy on GSM8K, a dataset that remains challenging for small-scale models. The performance comes from sufficient good quality synthetic data and the use of a verifier, which we describe in this section.

We finetune the Phi-1.5 125M, 350M and 1.3B models on our TinyGSM from Section 3, and in particular, the 1.3B model reaches 68.2% accuracy.The Phi-1.5-small 350M and Phi-1.5-125M variants are pretrained on the same pretraining data as the Phi-1.5 1.3B model. Performance of training on TinyGSM from scratch is reported in Table 2. We use the Adam optimizer with FP16 during training, with a linear warm-up and a maximum learning rate of 1e-4, a weight decay of 0.01, and an effective batch size of 1024. The finetuning phase takes up to 20k steps in total. As shown in Figure 1, even without verifiers, our models are already competitive to models of size from 7B and larger. As an anecdote, an earlier and worse performing version of our Phi-GSM 1.3B model gets 94% (or 82.5% from 350M at pass@32, whereas the 750M CodeT5+ model (Wang et al., 2023) gets 73.8% (or 70.5% from 220M) at pass@100.

2 Improving small models with a verifier

While sufficient synthetic data can significantly boost model performance, the performance is still below 70%. Does further improvement necessitate larger model and more data then? There may be two concerns: First, there may be a diminishing return in adding extra parameters and data; for instance, while there is a 10% increase in performance when increasing from around one third of the final size of TinyGSM to two thirds, the final one third of the data provided only marginal gain. Moreover, even if the small language model is able to fully match the quality of the synthetic data, GPT-3.5-turbo itself can only achieves 77.4% test accuracy on GSM8K, which seemingly poses a limit on the performance of any models distilling from its generations.

In this section, we show that the use of a verifier can be an effective strategy orthogonal to introducing more and better data, and can even help SLMs exceed the accuracy of GPT-3.5-turbo generations. The main observation that the best of multiple generations significantly outperforms a single generation. These generations could be low-temperature generations from different checkpoints of a single run, where taking the best out of generations from 5 checkpoints of (an early version of) our 350M model reaches 75% accuracy, similar to findings in temporal ensembling (Laine & Aila, 2016) and snapshot ensembles (Huang et al., 2017). For utilizing multiple checkpoints, an option is to use model soup (Wortsman et al., 2022); however, a uniform soup did not improve the accuracy. Another option is to perform EMA, which has been shown effective in Block et al. (2023). We found that EMA was not helpful when applied to the 1k-step-interval checkpoints; more frequent averaging is likely required. The generations could also be from high-temperature generations based on a single checkpoint; for instance, the pass@32 accuracy of our 1.3B model is 94%.

This suggests a promising direction of leveraging multiple generations: we can obtain a great performance boost if we are able to identify the best generation. This idea is effective yet natural: The probabilistic nature of the generative process naturally leads to the fact that multiple generations of a language model are more likely to contain a correct solution than a single one. Empirically, it has been widely observed that pass@$k$ accuracy, namely, the accuracy by taking the best of $k$ generations, is often much higher than pass@1. The main challenge is that without knowing the labels, the definition of “best” is usually unclear. A workaround is to apply some form of self-selection, such as by choosing the one with the highest logit or the most consistent solution (Wang et al., 2022; Li et al., 2022). There is, however, a notable limitation: generations can be consistent and confident yet inaccurate, making the self-consistency approach through majority voting less effective (Li et al., 2022).

Given these observations and inspired by findings in (Cobbe et al., 2021), we propose to use a separate verifier for selecting candidate generations. For each base generation SLM, we train a verifier to predict whether a generation is a correct solution to the given question. During inference, we generate multiple candidate generations using temperature sampling, and select the one with the highest verifier score.

The training data consists of the SLM’s generations on the labele GSM8K training set questions, paired with the binary labels indicating whether a generation leads to the correct numerical answer. We sample 48 generations for each training set question. The binary label for each generation is based on the final execution result and the ground truth label only, and we do not verify the correctness of intermediate steps. Note that this is the only time where the GSM8K training set is directly utilized in training.

Training setup

The verifier is trained with a sequence-to-sequence task, where we use the binary label on the entire sequence to supervise each token. We find this approach improves consistently over training with a sequence classification task (i.e. only predicting a binary label on the entire sequence). The verifier model is initialized to be the same as the SLM, with an additional prediction head shared across all positions. All network parameters are updated during verifier training, which significantly outperforms alternatives where only the prediction head is updated, or where the network is trained from scratch.

Checkpoint and data diversity

The diversity of the training data proves to be important, even if this may result in slightly sacrificing the data qualityWorse quality data refers to generations from earlier checkpoints, or sampled with a higher temperature: we found that the accuracy for temperature=1 can be more than 10% worse than temperature=0 (i.e. greedy decoding).: we find that including data generated using various temperatures and different checkpoints improves the performance of the verifier. In particular, the training data for the verifier is generated from checkpoints at 1k, 6k and 12k steps, and both training and testing data use a mixture of data generated with temperature 0.5 and 0.7. Compared to training with generations from a single checkpoint, the use of multiple checkpoints improves the performance from 78.9% to 81.5%.

Note also that in general, while we train one verifier for each generative model, verifiers transfer reasonably well across models. For example, our 1.3B model was able to reach 78.1% with a 350M verifier trained on generations from a 350M model.

Generation model size vs verifier size

In Figure 4, we present results from a cross-examination of various generation model sizes + verifier model sizes. Interestingly, while the best accuracy is achieved with configuration with largest sizes, the verifier size seems to play a bigger role than the generation model size. The effect of model size scaling is surprisingly mild: as shown in Table 1, increasing the base generation model from 125M (Phi-1.5-tiny) to 1.3B (Phi-1.5) only gives a 6% boost. On the other hand, the verifier seems to be much more parameter efficient. For example, 125M generation model + 1.3B verifier can achieve 76.1%, while 1.3B generation model + 125M verifier gets only 71.7% Figure 4.

Robustness and decontamination

While we never use the GSM8K test set during training, TinyGSM consists entirely of synthetic data generated by GPT models, which may be contaminated since GPT-3.5-turbo may have been exposed to the test set during its own training, which would have led to some generated synthetic samples being replicating part of the test set. To prevent contamination, we decontaminate TinyGSM by checking for n-gram matches. We use $n=13$ following standard practices (Brown et al., 2020; Wei et al., 2021; Du et al., 2022), n-gram matching is not sufficient for guarding against some other types of contamination (e.g. with respect to paraphrasing). However, we are not aware of better checks. One alternative is to check embedding similarity, though our clustering results on CodeGen 350M (Nijkamp et al., 2022) embeddings suggest that the embedding mostly reflects the semantic (topics) rather than structural/functional similarity, making it unfit for checking similarity in math questions. To our knowledge, state-of-the-art papers on training set contamination only test for exact matching (Shi et al., 2023b; Oren et al., 2023), and checking for contamination beyond exact match remains an open problem. and remove punctuation and numbers before computing the matching. Out of the 11.0M unique synthetic questions The number of unique questions is smaller than the number of question-solution pairs since some questions were sampled more than once in the second step of the 2-step generation (Section A.1) and hence have multiple solutions. , 22 questions have a nonzero 13-gram match with the test set, and 38k questions (i.e. around 0.35% of the full set) have non-zero 8-gram matches. Examples of 13-gram matches are provided in Section A.2.

2 Evaluation on SVAMP

For evaluating robustness of our models, we test on the SVAMP (Simple Variations on Arithmetic Math word Problems) dataset (Patel et al., 2021), consisting of 1000 math word problem questions with a focus on arithmetics. SVAMP constructed by applying certain types of variations to a set of base questions. Even though the base questions are generally considered easier than GSM8K See Table 1 in Xie et al. (2023)., the variations may often confuse LLMs, thus making it a challenging benchmark for robustness. Our 1.3B model achieves 75.6% on SVAMP without further finetuning, indicating the robustness of the model.

Discussions

In this work, we showed a simple approach that enabled a 1.3B generation model to achieve 81.5% on the GSM8K dataset, setting a new state-of-the-art for small language models and raising the performance curve for scaling. Our approach consists of two simple steps: 1) collecting TinyGSM, a GPT-3.5 generated synthetic dataset which we will fully release, and 2) using a verifier that scores how likely a generation is correct, whose quality is boosted by utilizing diverse generations. Our results provide positive evidence that small language models have more potentials to be unlock and can be used for efficient. For future directions,

Leveraging different formats: TinyGSM uses Python code as solutions, inspired by the observation that language models tend to struggle at calculations. However, we found that different solution formats, i.e. code versus natural language, can be complementary: while code helps circumvent errors related to execution or calculation, it tends to perform worse at questions that require equation solving, likely due to the fact that the Python syntax does not naturally support equations. Properly combining both formats has the potential to further boost performance.

The effect of verifier size: Our results show that given a budget on the model size, scaling the verifier may be a more efficient use of the parameters. This counters our intuition that verification is an easier task than generation (which involves search), though there might be connections to findings in GAN training where the size of discriminator (Arora et al., 2018). Exploring the parameter efficiency in a generation model versus a verifier is an interesting future direction.

References

Appendix A Additional details on TinyGSM

The majority of the TinyGSM was generated using the prompt in Figure 3, where GPT-3.5-turbo is asked to generate both the question and the corresponding solution. The remaining data, including all data based on GSM-IC, is generated using a two-step process, where the first step prompts the model to generate question variants, and the second step asks to generate Python solutions given a question variant generated in the first step. The exact prompts are provided in Figure 7–Figure 9.

A.2 Contamination check: 13-gram collisions

There are 22 questions (out of 11.0M) with 13-gram collisions to test set questions. Examples are shown in Figure 10.

Appendix B Pretrained vs Random Init

In this section, we present a comparison of training on TinyGSM from a random initialization versus from a pretrained model.