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

We thank the reviewer for their insightful comments. We are motivated to see that the reviewer finds our work (a) diverse in terms of the experimentation, (b) interesting its empirical findings, and (c) clearly written and easy to follow.

Q: Matched number of samples vs matched compute

• We believe there might be a misunderstanding with the overall objective of our work and we wish to clarify that.
• Previous work has been comparing small and large models in what the reviewer calls “matched number of samples” setting, and found that data from the larger model is better. We also corroborate these results in Figure 18 of the paper by comparing small and large models in a number-matched setup (this is in addition to Figure 20 referenced by the reviewer and seems to be the result that the reviewer is asking for). Under a number-match setting, we also agree that it is better to use the large model.
• However, we argue that the previous practice of comparing models in a number-matched setup has been putting the small model at a disadvantage because the small model spends substantially less compute compared to the large model during data generation. A fair comparison would give the same amount of compute to both the small and the large model.
• We then compare small and large models in a fair setting where we allow both models the same amount of compute (i.e. compute-matched setup) and show that under a fair setting, the small model wins.
• Important Note: From what we understand, the reviewer seems to consider the number-matched setting to be the fair setting, and the compute-matched setting to be a setting that gives the small model an advantage. However, our work argues that the compute-matched setting is the fair setting and the previous number-matched setting has been putting the small models at a disadvantage and not allowing for the full potential to be realized.

Q: “that would amount to not being able to generate data for 8k-12.5k questions.”

• We would like to also clarify that we are generating samples for all the 8k-12.5k questions. In fact, we are considering settings where we generate up to 30 samples per question with Gemma2 9B and 35 samples per question with Gemini 1.5 Flash. This amounts to O(0.5M) samples for finetuning.

Q: Clarification regarding the confounding factors of variation

• We clarify that the compute-matched sampling aims to sample N times more solutions from the small LM in comparison to the large LM where N is the ratio of the large LM and small LM capacity. For our experiments, we considered Gemma2-9B and 27B, hence, N = 3. Thus, we highlight that the multiplication factor is directly related to the model capacities rather than it being a confounder in our analysis.
• Further, we clarify that some of our key results (Figure 1b and 8) also cover the state of the art language models i.e., Gemini-Flash and Pro. Here, we sampled 35x more solutions from the Flash model in comparison to the Pro model in the price-matched scenario. Hence, our experiments are applicable to open as well as closed language models.

I would like to thank the authors for the reply and especially the data on the Llama models. I just wanted to give an acknowledgement that I will respond in detail on Monday, since I am occupied with a parallel deadline until then.

We thank the reviewer for their diligent feedback. We are motivated to see that the reviewer finds our work: (a) well-structured and clearly written, (b) well-executed with convincing evidence, and (c) valuable contribution to the synthetic data generation research.

Q: Robustness of the conclusions

• We agree that models from different companies may have different properties making each of them suitable for a different use-case or task. While we do not provide results for comparing cross-company WC and SE models, our results so far help decide which model to use within the same family. For example, they help decide that within the Qwen2.5 family, one may be better of using a smaller model than the 7B model and within the Gemma2 family one may be better of using the 2B model instead of the larger models. Extending our results to models from different companies and understanding when it works as well as the failure modes is a great future direction.
• Note that the aims of the current work is to lay the foundation for compute-matched sampling to train LM reasoners. We perform experiments with WC and SE models from the same model family which is not an unreasonable assumption. In our paper, we show that WC data outperforms SE data from Gemma (open) as well as Gemini (closed) language models. We further point out that the Gemma and Gemini series models are quite different from each other based on the information available publicly. For instance, Gemma [1] models are purely text language models while Gemini models are natively multimodal in nature which would lead to an entirely different learned data distribution [2].
• To provide further evidence, we performed experiments with the Llama models. Specifically, for the MATH dataset, we generated 1 solution per problem from Llama-8B (SE) and 3 solutions per problem from Llama-3B (WC) in accordance with the compute-matched sampling ratio. Subsequently, we supervise finetuned Llama-1B, 3B, and 8B on the generated data from SE and WC models under the fixed sampling setup. We present the results on MATH500 test set below:

• Consistent with our original results, we find that training with the WC data is more compute-optimal than SE data across diverse finetuning setups including knowledge distillation, self-improvement, and weak-to-strong improvement. We will add these results to the revised paper.

• Also consistent with our original results, we see that the WC model has a coverage of 67% and a diversity of 2.2, whereas the SE model has a coverage of 49% and a diversity of 1.

• We leave the expansion to diverse pairs of the (WC, SE) models for the future work. We will add this discussion in Appendix A.

[1] Gemma2: https://arxiv.org/pdf/2408.00118
[2] Gemini 1.5: https://storage.googleapis.com/deepmind-media/gemini/gemini_v1_5_report.pdf

We thank the reviewer for their insightful comments. We are happy to note that the reviewer finds our work (a) significant in comparing SE and WC models, (b) impressive in terms of the findings that challenge the traditional beliefs, and (c) effective and robust in terms of the evaluation.

Q: Robustness of conclusions

• Thanks for the suggestion. To address the reviewer’s comments in the limited time, we performed experiments with the Llama models. Specifically, for the MATH dataset, we generated 1 solution per problem from Llama-8B (SE) and 3 solutions per problem from Llama-3B (WC) in accordance with the compute-matched sampling ratio. Subsequently, we supervise finetuned Llama-1B, 3B, and 8B on the generated data from SE and WC models under the fixed sampling setup. We present the results on MATH500 test set below:

• Consistent with our original results, we find that training with the WC data is more compute-optimal than SE data across diverse finetuning setups including knowledge distillation, self-improvement, and weak-to-strong improvement. We will add these results to the revised paper.
• Also consistent with our original results, we see that the WC model has a coverage of 67% and a diversity of 2.2, whereas the SE model has a coverage of 49% and a diversity of 1.
• We would like to also note that the models from the Gemma2 family, used in the original report, are quite capable open language models in their model capacity range. Specifically, Gemma2-9B achieves 36.6% in comparison to Mistral-7B’s 12.7% and LLaMA-2-7B’s 2.5% [1]. In addition, we experiment with state-of-the-art Gemini models (Pro/Flash) too and show that WC data outperform SE in this scenario. We believe these two families serve as solid evidence that our method works for open and closed language models, and the addition of the Llama results strengthen our claims even further.

[1] Gemma2: https://arxiv.org/pdf/2408.00118

Q: Scale of data collection

• We respectfully disagree with the reviewer’s take on this. We would like to clarify two important points:
  1- While the number of problems in our datasets may seem on the lower end (7.5K), we generate multiple solutions per problem for synthetic data creation. In particular, we generate 30 solutions and 10 solutions per problem from Gemma2-9B and Gemma2-27B for the two datasets, respectively, which amounts to 600K problem-solution pairs. In addition, we generate 35 solutions per problem from Gemini-Flash for the two datasets which amounts to 525K problem-solution pairs. We believe that this is quite large and at par with the size of the datasets used for supervised finetuning of the language models [4].
  2- Perhaps more importantly, we would like to clarify that the aim of the paper is NOT to reduce the cost of sampling data and/or make the approach scale to larger datasets. Instead, we argue that given a fixed budget (i.e. the budget one is willing to spend for training a model, in terms of compute or cost), it is more efficient usage of the budget if one spends it on sampling data from WC instead of the SE model for training language model reasoners as well as instruction following models. As a particular example, imagine one wants to spend 1000USD on tuning a model on their training data, and they have the choice of spending it on sampling from a WC or a SE model. Our results suggest that no matter how small or large their training data is, if they spend the 1000USD sampling from WC, they may end up with a better model than if they spend it sampling from SE.
• We also note that GSM-8K and MATH datasets are standard datasets in the language model reasoning literature [1,2,3].

[1] V-STaR: Training Verifiers for Self-Taught Reasoners: https://arxiv.org/abs/2402.06457
[2] STaR: https://arxiv.org/abs/2203.14465
[3] RestEM: https://arxiv.org/abs/2312.06585
[4] Alpaca: https://huggingface.co/datasets/tatsu-lab/alpaca

Q: Quality-diversity with WC and SE model

• In our early experiments, we found that changing the sampling temperatures from 0.7 (the one used for the final experiments) to either 0.5 (i.e. lower) or 1.0 (i.e. higher) would lead to lower performance than 0.7 for both WC and SE data. This finding corroborates the reviewer’s point that the WC model may be providing a superior quality-diversity trade-off compared to merely increasing the sampling temperature.

Q: Insights into WC and SE sampling.

• In this work, we lay the foundation for this direction by showing that, perhaps surprisingly and contrary to the common belief, data from smaller LMs can be more compute-optimal than data from larger MLs. As the reviewer mentioned, our work opens up the avenue for a great body of theoretical analysis on when WC may be superior than SE.
• So far, our understanding based on the experiments reported in the current work is that whether a WC can outperform a SE depends on how they compare in terms of coverage, diversity, and false positive rate (FPR). In particular, we believe the FPR of the WC model should be low (as evidenced by our experiments with filterings other than ground truth), and the coverage of the WC model should be higher than SE (when the two models have similar coverage, they tend to lead to similar performance).
• To get a more nuanced understanding of when WC outperforms SE, one possible experimental design is to compute the coverage, diversity, FPR, and potentially other features for several models and several datasets, and then finetune models on data from these models and compute the delta in their performance. Then fit lines/curves that predict the delta in performance with respect to the three properties. The weights assigned to these features can then be indicative of how important each feature is with respect to other features and allow for predicting whether a WC may outperform an SE only based on these features, before running any finetuning experiments. This is, however, quite computationally demanding and beyond the scope of the current work. We hope future work will dive deeper into this.
• We point that we study the notion of (a) coverage, (b) diversity, and (c) false positive rates to get insights into WC and SE sampling from Gemma and Gemini models (L209-237).
• Our experiments reveal that the WC sampling achieves higher coverage and diversity than SE sampling while SE sampling achieves a lower false positive rate (FPR) at the fixed sampling budget. Subsequently, our supervised finetuning results indicate that WC sampling can achieve better scores than SE sampling.
• While it is pertinent, we mention in our discussion (Appendix A) that the aim of our work is not to come up with the conditions under which WC sampling will outperform SE sampling.
• To achieve this, we will need to curate a large number of datasets with diverse coverage, diversity and FPR values and run many finetuning runs subsequently. However, this paper lays foundations for this future exploration.

Q: Performance on more complex reasoning tasks not fully explored.

• In this work, we consider the MATH reasoning dataset as our main testbed, which is standard in most of the LLM reasoning papers [3,6]. This dataset contains more complex problems in comparison to the high-school level GSM-8K dataset.
• In Figure 22, we show that the compute-matched sampling from the WC model increases the coverage across all the levels i.e., (easy: level 1 to complex: level 5).
• In addition, we evaluate our MATH-finetuned models on the Functional MATH dataset (Figure 6) to assess their generalization performance to distributionally-shifted problems. Note that Functional MATH has been shown to be quite challenging, even for the frontier models.
• Although several complex evaluation reasoning tasks exist (e.g., GPQA), the absence of widely accepted training datasets restricts our ability to effectively fine-tune models for these tasks.

Q: Theoretical analysis

• While we do have non-theoretical insights for now: the higher the coverage and diversity of WC and the lower the FPR, the more chance WC will work.
• In practice, this theoretical analysis will include many models with different coverage, diversity, and FPR. Subsequently, we can try to fit curves that predict the model performance.
• Then, the weights of these variables could tell us how important they are with respect to each other. However, this will require a lot of compute, which is out of scope of this work.

We thank the reviewer for their diligent feedback. We are motivated to see that the reviewer finds our work (a) novel with a fresh perspective on synthetic data generation, (b) experimentally rigorous, (c) practically impactful, and (d) rigorous in analysis and experimental validation.

Q: Reliance on ground truth, filtering strategies, and more tasks.

• We agree these are important directions for extending our work and we have indeed included experimental results both for 1- different filtering strategies beyond ground truth labels, and for 2- tasks beyond math and coding. These results are referenced on lines 407-413 of the main text, with the results being presented in the appendix due to space limitations. If the reviewer finds these parts of our result more interesting than any other part in the main text, we are happy to move them to the main text. In what follows, we provide a high-level summary of the results.
• We extend our results to scenarios lacking ground-truth final answers (mentioned in L411, and Appendix B). Specifically, we consider (a) no filtering, and (b) filtering using LM as a judge setting. In the latter case, we propose an approach to still match the computes, despite using an LM for judging model generated answers. Overall, the trends suggest that whether WC data is superior to SE data depends on the quality of the overall models and the finetuning setup in both the settings. In particular, we notice that Gemini-Flash generated data still outperforms Gemini-Pro generated data without access to ground-truth data. Future work can extend our results by using trained verifiers instead of llm-as-a-judge, in a similar compute-matched setup.
• Beyond reasoning, we consider the instruction-following setup where there is no notion of filtering. In Figure 13, we find that WC sampling achieves a better IFeval score than SE sampling for instruction-following too, thus showing that our results can be extended beyond reasoning.
• Lastly, we clarify that it is quite common to assume access to the ground-truth final answer in the LLM reasoning literature [1,2,3]. Specifically, many real-world datasets such as MATH and GSM-8K come with final answers, and many coding datasets come with test-cases that can be used to judge the correctness of a generated code. Note that in our work, we do not use human-written chain-of-thoughts to solve the math problems, we just utilize the final answers to filter the chain of thoughts that lead to incorrect final answers.

I appreciate your comprehensive responses to my concerns and questions. You have addressed the key points thoroughly and convincingly.

Theoretical Understanding & Framework

• You've clarified that while formal theoretical bounds weren't the paper's focus, you have meaningful empirical insights about the relationship between coverage, diversity, and FPR
• The suggested experimental design for understanding WC vs SE performance through feature analysis is promising
• Your explanation of compute-matched sampling ratios is clear and well-justified

Methodology & Evaluation

• The appendix results on different filtering strategies (including no filtering and LM-as-judge) significantly address my concerns about reliance on ground truth
• Your justification for the human evaluation sample size makes more sense given the correlation with large-scale automatic evaluation
• The prompting details and temperature choices are well-explained and consistent with prior work

Generalisation & Scope

• The results on instruction-following tasks demonstrate broader applicability than initially apparent
• Your work with MATH dataset (including level-specific analysis) and Functional MATH shows good coverage of complex reasoning
• The exploration across multiple model sizes (7B to 27B) and both open/closed models provides good coverage

Implementation & Practicality

• The cost analysis showing 0.15x cheaper Flash data can outperform Pro is particularly compelling
• The straightforward integration with existing training pipelines is encouraging
• The extension to different domains (math, code, instruction following) with minimal modifications strengthens the paper's practical value

Score Update: Given the thoroughness of your responses and the additional context provided, particularly around:

1. The broader applicability beyond just mathematical reasoning
2. The existence of results without ground truth filtering
3. The compelling cost-effectiveness analysis
4. The clear empirical insights into when WC sampling works better

I am raising my score from 6 to 8. The paper makes a stronger contribution than I initially assessed, with practical implications for making model training more efficient and accessible.

One suggestion for camera-ready: Consider moving some key implementation details and the non-ground-truth filtering results from the appendix to the main paper, as these strengthen your core arguments.