Smaller, Weaker, Yet Better: Training LLM Reasoners via Compute-Optimal Sampling

Language models (LMs) have demonstrated impressive capabilities in reasoning tasks, but their success heavily relies on being trained on vast amounts of (problem, solution) pairs. Collecting this data from humans is a costly and time-consuming process. Recent studies have demonstrated the feasibility of synthetically generating this data using LMs themselves, offering a potentially more scalable and efficient approach to training data acquisition. One such widely-adopted approach is to sample multiple candidate solutions for a problem from an LM, filters them for final answer correctness, and finetune models on the correct solutions (Zelikman et al., 2022; Singh et al., 2023). Several works show that LMs trained with such synthetic solutions outperform those trained with human-written solutions (Yuan et al., 2023; Yu et al., 2023; Yue et al., 2023; Singh et al., 2023; Pang et al., 2024). Practitioners often sample solutions from strong LMs to ensure high quality (Teknium, 2023; Roziere et al., 2023; Mukherjee et al., 2023; Xu et al., 2023). However, sampling from strong LMs is expensive and resource-intensive, and limits the number of solutions that can be generated for practical sampling budgets.

In this paper, we explore an alternative sampling approach. Given a fixed compute budget, we investigate sampling from a weaker but cheaper (WC) model as opposed to the commonly-used approach of sampling from a stronger but more expensive (SE) model. We start by comparing data from WC vs SE across three axes that play crucial roles in the utility of such synthetic data: 1- coverage, the number of unique problems that are solved, 2- diversity, the average number of unique solutions we obtain per problem, and 3- false positive rate (FPR), the percentage of problems that arrive at the correct final answer but with a wrong reasoning. We find that since we can generate more samples from the WC model compared to the SE model under a fixed budget, the data from WC may exhibit higher coverage and diversity. However, due to the lower quality of the WC model, it may also have a higher FPR. As a particular example for the Gemma2 family (Team et al., 2024a, b) on the MATH dataset (Hendrycks et al., 2021), Gemma2-9B achieves $11\%$ higher coverage and $86\%$ higher diversity, but also with $7\%$ higher FPR compared to Gemma2-27B.

We then fine-tune models on data from SE and WC (see Figure 2) across diverse setups corresponding to three paradigms: 1) knowledge distillation, where a student LM learns from a teacher LM (Hinton et al., 2015); 2) self-improvement, where an LM learns from self-generated data (Huang et al., 2022); and 3) a new paradigm we introduce called Weak-to-Strong Improvement, where a strong student LM improves using synthetic data from a weaker teacher LM. Using two (WC, SE) model pairs, one from the Gemma2 family and another from the Gemini 1.5 family (Reid et al., 2024), we show on multiple benchmarks that training on WC-generated data consistently outperforms training on SE-generated data under the three setups, with relative gains of up to $31.6\%$ percent (see Figure 1 for a summary of the results). Our results indicate that it is more compute-optimal to sample from a WC model as opposed to the common-practice of sampling from a SE model. With the performance gap between small and large LMs getting narrower over time (especially at larger scales), our results establish a solid foundation for training the next generation of LM reasoners.

Let $\mathcal{D}=\{q_{i},a_{i}\}_{i=1}^{i=n}$ be a training dataset of size $n$ with reasoning questions $q_{i}$ and final answers (aka labels) $a_{i}$. A successful approach to leverage such data to improve models for reasoning is as follows. We sample multiple solutions for each $q_{i}$ at a non-zero temperature and create the synthetic data $\mathcal{D}_{G}=\{q_{i},\{(\hat{r}_{ij},\hat{a}_{ij})_{j=1}^{j=k}\}\}$, where $k$ is the number of samples, $\hat{r}_{ij}$ is the $j$-th reasoning chain (i.e. solution) generated by the model for $q_{i}$, and $\hat{a}_{ij}$ is the model’s final answer for $q_{i}$ in the $j$-th sample. Then, we filter the incorrect solutions by comparing $\hat{a}_{ij}$ to $a_{i}$ and removing the solutions whose final answer do not match that of the gold answer¹¹1While it is possible to use more sophisticated approaches for filtering (e.g., process-based or outcome-based reward model (Uesato et al., 2022)), in this work we focus on final answer correctness for filtering as it has shown to be strong.. Finally, we supervise finetune a model on the remaining data $\tilde{D}_{G}$ to maximize $J(\theta)=\mathbb{E}_{(q,r,a)\sim\tilde{D}_{G}}\left[\log(p_{\theta}(r,a|q))\right]$, i.e. the probability of generating the reasoning $r$ and final answer $a$ given the question $q$. This approach was first proposed in (Zelikman et al., 2022) and was then extended in multiple works including (Zelikman et al., 2024; Singh et al., 2023).

For a dataset $\mathcal{D}_{G}$, we compute $coverage@k$ (aka $pass@k$) (Chen et al., 2021) as $\mathbb{E}_{\mathcal{D}_{G}}\left[1-\binom{M-c}{k}/\binom{M}{k}\right]$ where $c$ is the number of solutions, out of $M$, with correct answers and $\mathbb{E}_{\mathcal{D}_{G}}[.]$ denotes the expectation over the problems and solutions in the generated dataset. Conceptually, $coverage@k$ measures the fraction of unique questions that have at least one correct solution, assuming that we sample $k$ solutions per question from the model. We also define $diversity@k$ as the average number of unique correct solutions we obtain per question when we sample $k$ solutions per question. Finally, we define false positive rate (FPR) as the percentage of solutions in $\tilde{D}_{G}$ where the reasoning is incorrect, despite the final answer being correct.

Different choices of the LM to sample solutions from and the LM to finetune lead to different setups. Knowledge Distillation (Hinton et al., 2015) corresponds to training a student LM on the synthetic data sampled from a stronger and larger LM. Self-Improvement (Huang et al., 2022) corresponds to training an LM on samples generated from itself.

To generate a dataset $\mathcal{D}_{G}$ with synthetic solutions from $\mathcal{D}$, one can leverage different models for generating solutions. Specifically, at a fixed sampling budget (FLOPs), one can generate more samples from a weaker but cheaper (WC) model or fewer samples from a stronger but more expensive (SE) model. Given a WC model with $P_{WC}$ parameters and SE with $P_{SE}$ parameters, we compute the sampling ratio at a fix budget for the two models, focusing on decoder-only transformer models (Vaswani, 2017). Following (Kaplan et al., 2020), we note that the FLOPs per inference token is $2P$, for a model with $P$ parameters. As a result, the FLOPs for $T$ inference tokens is $2PT$. Further, we assume that generating each solution requires an average of $W$ inference tokens for both models²²2This is a reasonable assumption given that the solution to a question is expected to be model-agnostic. We note, however, that it is possible for some questions that one model solves a question using a more optimal way compared to the other model thus producing a smaller solution.. Let $S_{WC}$ and $S_{SE}$ represent the number of samples we generate per question for the two models. The total cost of generating samples for the dataset $\mathcal{D}$ will then be $Cost_{WC}=n\times S_{WC}\times W\times(2P_{WC})$ and $Cost_{SE}=n\times S_{SE}\times W\times(2P_{SE})$ for the cheap and expensive models, respectively. At a fixed sampling budget, we have:

Equation 1 indicates that at a fixed sampling budget, for each question we can generate $P_{SE}/P_{WC}$ more samples from WC; the ratio scales linearly with the model parameters ratio³³3Note that this may also depend on the available hardware, which we ignore in this work.. Sampling more solutions from WC may increase the likelihood of correctly solving a larger subset of the problems (high coverage) and obtaining more correct solutions per question (high diversity).

Given a fixed budget, we can either generate fewer samples from a SE model or more samples from a WC model, and then finetune models for a fixed number of steps on the data from each of these models to measure and compare the utility of the data from each model. Specifically, we generate $P_{SE}/P_{WC}$ more samples from the WC model compared to the SE model. We consider three finetuning setups that consists of diverse finetuning paradigms. The paradigms include the widely used knowledge distillation, the emerging framework of self-improvement, and a novel weak-to-strong improvement paradigm we introduce in this work. We define weak-to-strong improvement (W2S-I) as enhancing the reasoning capabilities of a strong model using samples generated from a weaker model. The three setups are as follows:

• •

  Student-LM finetuning: Conventionally, the supervised finetuning data for training student LM is acquired from SE models to ensure high-quality (Teknium, 2023). However, we aim to understand whether WC models can replace SE models for distillation at the fixed sampling budget. To do so, we finetune a student LM separate from the WC and SE models on the WC and SE data, which corresponds to distillation in both the cases.

• •

  WC-LM finetuning: Prior work (Singh et al., 2023) has shown that finetuning a WC model through self-generated data lags behind distillation from SE data. However, their setup spends a higher sampling budget (FLOPs) on collecting data from the SE model than collecting SI data from the WC model. In this work, we revisit this finetuning setup under the fixed sampling budget and finetune the WC model on the WC and SE data at a fixed budget for both. Note that training the WC model on its own data corresponds to self-improvement whereas training WC on the data from SE corresponds to distillation. Hence, this setup compares self-improvement on WC data with distillation from SE data.

• •

  SE-LM finetuning: It is commonly believed that to improve a SE model, we either need synthetic data from the SE model itself or from an even stronger (and perhaps more expensive) model than SE. Here, we test an alternative approach to understand whether the synthetic data from the WC model can improve the SE model. To this end, we finetune the SE model on the WC and SE data. Training SE on data from WC corresponds to W2S-I and training SE on data from SE corresponds to self-improvement. Overall, this setup compares W2S-I by WC data with self-improvement by SE data.

A summary of the three setups and the finetuning paradigms that each case corresponds to can be found in Table 1.

Datasets: We utilize MATH (Hendrycks et al., 2021) and GSM-8K (Cobbe et al., 2021) as the reasoning datasets due to their wide adoption for mathematical problem solving. Specifically, MATH consists of competition level problems with various levels of difficulty (Level 1-5), and GSM-8K comprises of grade school level math problems. Each dataset contains $7500$ math problems in their training split. We evaluate the models on $500$ problems from the MATH test split (Lightman et al., 2023) and $1319$ problems from the GSM-8K test split. Further, we use $500$ problems from the MATH test split and $500$ problems from GSM-8K as the validation dataset. We also use the Functional MATH dataset (Srivastava et al., 2024) for a transfer study. Further, we present the results for a coding task in Appendix A.

Data Generation: We use Gemma2 models for synthetic data generation, with pretrained Gemma2-9B and Gemma2-27B acting as the WC and SE models respectively. We generate the solutions for the problems in the MATH using a 4-shot prompt and for GSM-8K using an 8-shot prompt. Since the 9B model is roughly 3 times smaller than the 27B model, at a fixed sampling compute budget we can sample $3\times$ more sample solutions per problem for Gemma2-9B. For our experiments, we consider two sampling budgets: a low budget, where we generate 1 and 3 candidate solutions per problem from Gemma2-27B and Gemma2-9B, respectively, and a high budget, where we generate 10 and 30 candidate solutions per problem. Further, we study the transfer of the reasoning capabilities for the models trained on MATH at the high sampling budget on the Functional MATH dataset.

Model Finetuning: We summarize the details for our finetuning setups in the Table 1. In the Student-LM finetuning setup, we finetune the Gemma-7B model (Team et al., 2024a) on data from Gemma2-9B (WC) and Gemma2-27B (SE). In addition, we use Gemma2-9B and Gemma2-27B for the WC-LM and SE-LM finetuning setups, respectively. Further, we train the LMs across different setups with the human-written solutions as a ground-truth baseline. We provide the finetuning details in Appendix F.

Synthetic Data Evaluation: To assess the quality of the synthetic data from the SE and WC models, we measure the false positive rate, as well as coverage and diversity at a fixed cost. From Equation 1, we know that sampling one solution from SE takes the same FLOPs as sampling $P_{SE}/P_{WC}$ solutions from WC. Therefore, we compare $coverage@k$ for SE to $coverage@(\frac{P_{SE}}{P_{WC}}k)$ for WC to allow a similar budget to both models. Specifically, we compare $coverage@k$ and $coverage@3k$ for our SE and WC models. Similarly we compare $diversity@k$ and $diversity@3k$ for our SE and WC models. Since FPR cannot be computed automatically, we compute it using two proxies: 1- a human evaluation on a subset of the data, where $50$ solutions from each model were selected randomly and rated for reasoning correctness by the authors, and 2- automatic evaluation where we sampled $500$ solutions and prompted Gemini-Pro-1.5 (Reid et al., 2024) to rate the correctness of the reasoning paths. To sample solutions, for the MATH dataset we selected uniformly from each diversity level. In our experiments, we find that the FPR estimates are close to each other for the human and automatic evaluation. We provide a few qualitative examples for the false positive instances in Appendix B.

Evaluating Finetuned Models: We use pass@1 accuracy to evaluate the performance of the finetuned LMs. Specifically, we generate a single solution for the problem (zero-shot) from the test split, using a sampling temperature of $0.0$ (greedy decoding) for the fine-tuned LM and measure the percentage of problems that where the final answer matches the golden final answer. We also report maj@k ($k=1,4,8,16$) for part of our experiments, where we generate $k$ solutions per problem at a sampling temperature of 0.7 and select the final answer that appears most among the $k$ samples.

We compare data from WC and SE models along several axes. First, we analyze the data along various quality metrics (§5.1). Subsequently, we present the supervised finetuning results for the different setups (§5.2). Finally, we perform ablation studies to study the impact of dataset size, sampling strategy, and the role of quality dimensions in the model performance (§5.3).

In the prior experiments, we focused on the synthetic data acquisition from open LMs. Here, we aim to show that data from the weaker SoTA LM can train better reasoners than stronger SoTA LM at a fixed sampling budget. To this end, we scale our method to sampling data from Gemini-1.5-Pro and Gemini-1.5-Flash. As the model sizes are not publicly available, we utilize the ratio between their pricing per output token as a proxy to perform compute-matched sampling. As of August 2024, we note that the price per million output tokens is $\$10.5$ and $\$0.3$ for Gemini-1.5-Pro and Gemini-1.5-Flash, respectively. Hence, we sample $1$ and $35$ solutions per problem from 1.5-Pro and 1.5-Flash, respectively. We conduct our experiments on the MATH dataset.

We perform knowledge distillation on the Gemma-7B, Gemma2-9B, and Gemma2-27B LMs with the synthetic data from the Pro (SE) and Flash (WC) models. We present the results in Figure 10. Interestingly, we find that finetuning with the WC data outperforms the SE data, achieving relative gains of $31.6\%$, $14.4\%$, and $10.9\%$ for Gemma-7B, Gemma2-9B, and Gemma2-27B, respectively. This can be attributed to the difference in the coverage of the models at the fixed sampling budget, which is $61.1\%$ and $81\%$ for 1.5-Pro and 1.5-Flash, respectively.

Reducing the price of data sampling. Further, we investigate training the LMs with the WC data that is less expensive than collecting $1$ solution per problem from the SE model. Specifically, we create a dataset by sampling $5$ solutions per problem from the Flash (WC) model, which is $7\times$ more economical than generating $1$ solution from the Pro (SE) model, in terms of the price ($). Upon training the LMs on the $0.15\times$ cost data regime (Figure 10), we find that training on this data can also outperform training with SC data, achieving relative gains of $19.1\%$, $9.8\%$, and $5.7\%$ for finetuning Gemma-7B, Gemma2-9B, and Gemma2-27B, respectively. This can be attributed to higher coverage of the weaker model ($69\%$), even in the more economical scenario, in comparison to the stronger model ($61.1\%$).

Takeaway: We demonstrate that price-matched sampling from weaker SoTA LMs produces superior reasoners compared to finetuning with data from stronger SoTA models.

We showed that for the current WC and SE models, training reasoners through sampling from WC models may be more compute-optimal. Here, we aim to discuss the relevance of these results for the future set of WC and SE models. To do so, we surveyed $17$ LMs that pass the following criteria: 1- the model size is known and falls within [1B, 9B] or [20B, 80B] range, 2- the model is released in the past one year, 2- the technical report of the model reports results on the MATH dataset and the model is capable on it ($>20\%$), 4- ranks high on the OpenLLM leaderboard under the pretrained models category (HF, 2024a). This resulted in models from seven families including Gemma-2 (Team et al., 2024b), LLaMA-3 (Dubey et al., 2024), Mixtral (Jiang et al., 2024), Qwen (Team, 2024; Yang et al., 2024a), Grok-1 (xAI, 2024), DeepSeek-v2 (Shao et al., 2024), and Yi (Young et al., 2024). We grouped these models into small LM (1B to 9B) and large LMs (20B to 80B). We then plotted in Figure 11 the model performances on the MATH dataset against their date of the publication release on arxiv and fitted trendlines to the data points representing the small and large LMs using the least squares method⁴⁴4We consider the number of active model parameters for mixture-of-experts LMs..

Our analysis reveals that, despite the variance, the trendline for the smaller LMs is steeper than that of the larger LMs. This indicates that the reasoning performance of the small LMs may be improving more rapidly over time compared to the larger LMs. The rapid rise in the performance of the small LMs can be attributed to factors such as the enhanced quality and scale of the pretraining data (e.g., LLaMA-3 employs 15T tokens), pruning and knowledge distillation (Muralidharan et al., 2024). With the performance gap between small and large LMs narrowing over time, we anticipate that our results will become even more relevant in the future.

LMs for reasoning. The ability to solve reasoning tasks has been a long standing goal of artificial intelligence (Reid et al., 2024; Achiam et al., 2023; Dubey et al., 2024; Team, 2024; Anthropic, 2024; AI, 2024). In this regard, LMs trained on the internet-scale data have achieved great success for math, code, and other reasoning tasks (Lewkowycz et al., 2022; Azerbayev et al., 2023; Kazemi et al., 2024). There have been several works that aim to enhance the reasoning capabilities of the LMs either via prompting (Kojima et al., 2022; Wang et al., 2022; Zheng et al., 2023; Kazemi et al., 2022) or finetuning (Yue et al., 2023; Yu et al., 2023). In this work, we focus on finetuning the LMs with task-specific datasets to build strong reasoners. Specifically, our method closely aligns with the widely adopted STaR (Zelikman et al., 2022) where the synthetic data from the LMs are used to elicit strong reasoning capabilities.

Finetuning LMs. Within the finetuning paradigm, there have been several works that improve reasoning with synthetic data. Broadly, these works focus on knowledge distillation from a strong but expensive LM (Wu et al., 2024; Yue et al., 2023) or self-improvement (Gulcehre et al., 2023; Singh et al., 2023). While it is common to filter the synthetic data for the final answer correctness (akin to Zelikman et al. (2022)), there are several works that aim to build task-specific verifiers to train strong reasoners (Lightman et al., 2023; Wu et al., 2024; Hosseini et al., 2024; Yuan et al., 2024). In this work, we explore the utility of the synthetic data from the weak but cheap LMs for training strong reasoners via knowledge distillation as well as self-improvement. However, we do not explore using model-based verifiers with the synthetic data for enhanced reasoning, and leave it as a future work.

Our weak-to-strong improvement paradigm, where a strong model is trained with the generations from the weak model, is related to several prior work (Bowman et al., 2022; Burns et al., 2023; Yang et al., 2024b) which study the ability of a strong LM to learn from the data generated by a weaker LM. However, the aim of these works is to recover the full capabilities of the strong model from weaker data, whereas we aim to enhance the strong model capabilities further. Additionally, our work studies compute-optimal sampling from weak and strong models, which is absent in previous work.

Large and small LMs. While training large LMs has led to significant advancements across various tasks, there has recently been a growing interest in developing capable small LMs (HF, 2024b; Javaheripi et al., 2023). Specifically, a capable small LM is faster to run, and easier to serve to millions of users on the edge devices (Gunter et al., 2024). As a result, several recent works aim to understand the utility of the weak but cheaper LMs in comparison to the strong but expensive LMs for reasoning. Specifically, Brown et al. (2024); Song et al. (2024); Snell et al. (2024) show that the solve rate of the small LMs can increase significantly with repeated sampling. In addition, Hassid et al. (2024) demonstrate that repeated generations from smaller LMs can outperform the data generated by larger LMs at a fixed sampling computational budget during inference for coding tasks. In this work, we go beyond these works and show the utility of the synthetic data from the small LMs for training strong reasoners across a diverse set of supervised finetuning setups.

In this work, we provide a framework for compute-optimal sampling from weak but cheap LM for reasoning tasks. Specifically, we show that at a fixed sampling compute budget, repeated sampling from a smaller model can achieve higher coverage and diversity than from a strong but more expensive model. Furthermore, our empirical findings highlight that fine-tuning LMs with data from the small LM can consistently outperform data from the large LM under the same compute budget. Our results can serve as a foundation for training LM reasoners, especially as the performance gap between small and large LMs continues to narrow over time.

This work was done during HB and AH’s internship at Google. We thank Hugo Larochelle and Hamidreza Alvari for feedback on this paper. We thank Chirag Nagpal, Katrin Tomanek, and Benjamin Estermann for support in setting up infrastructure for experimentation that was crucial for running our experiments.