Elamur Research
☰
home / papers / 2024 / Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters

Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters

Written for the Research library (claude-opus-5-5) from the paper text. Original — arxiv.org/pdf/2408.03314

Summary

The paper asks how much an LLM can improve on difficult prompts when it gets extra computation at inference time, and how that budget is best spent. The authors use the MATH benchmark and PaLM 2 models fine-tuned for two specific capabilities. They study two ways to spend the extra compute. The first is search guided by a process-based reward model (PRM), which scores each individual solution step. The second is having the model revise its own earlier answers one after another. Which approach works best depends on how hard a question is for the base model. Revisions help most on easier questions, while parallel sampling or tree search pays off on harder ones. Their "compute-optimal" policy picks the approach for each prompt based on estimated difficulty, and it matches best-of-N with roughly four times less compute. With total FLOPs held equal, a small model given extra inference compute beats a ~14× larger model on easy and medium questions. On the hardest questions, extra pretraining still wins. These results help decide how to split compute between training and inference.

Key points

  • Setup: PaLM 2-S* was fine-tuned separately as a revision model and as a PRM verifier, and evaluated on MATH. The authors fine-tuned because simply prompting off-the-shelf models does not produce useful revisions. Questions were grouped into five difficulty bins based on how often the base model solves them.
  • Efficiency gain: Choosing the strategy per prompt makes test-time scaling 2–4× more efficient than a best-of-N baseline, for both revisions and PRM search.
  • Test-time vs. pretraining: The FLOPs-matched comparison against a ~14× larger model used three ratios of inference tokens to pretraining tokens (R = 0.16, 0.79, 22). Extra test-time compute wins at low R and on difficulty bins 1–3. Extra pretraining wins at high R and on bins 4–5.
  • Limitations: Estimating question difficulty itself costs a lot of inference compute. Gains on the hardest problems stay small, and PRM tree search was not tried together with revisions.
  • Future direction: The authors suggest distilling the outputs of test-time compute back into the base model, creating a self-improvement loop.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields: Mathematics & statistics

Cited by

503 papers of our library cite it (of the 140673 whose bibliographies we read); 17 of them have a page here, 486 not yet.

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2408.03314 · arXiv 2408.03314

@misc{snell2024scaling,
  title = {{Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters}},
  author = {Snell, Charlie and Lee, Jaehoon and Xu, Kelvin and Kumar, Aviral},
  year = {2024},
  doi = {10.48550/arxiv.2408.03314},
  eprint = {2408.03314},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2408.03314}
}

References

51 references, 22 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Training revision models with synthetic data. Coming soon, 2024. find
  2. C. Andrieu, N. De Freitas, A. Doucet, and M. I. Jordan. An introduction to mcmc for machine learning. 2003. find
  3. R. Anil, A. M. Dai, O. Firat, M. Johnson, D. Lepikhin, A. Passos, S. Shakeri, E. Taropa, P. Bailey, Z. Chen, E. Chu, J. H. Clark, L. E. Shafey, Y. Huang, K. Meier-Hellstern, G. Mishra, E. Moreira, M. Omernick, K. Robinson, S. Ruder, Y. Tay, K. Xiao, Y. Xu, Y. Zhang, G. H. Abrego, J. Ahn, J. Austin,… find our card
  4. Y. Bai, S. Kadavath, S. Kundu, A. Askell, J. Kernion, A. Jones, A. Chen, A. Goldie, A. Mirhoseini, C. McKinnon, C. Chen, C. Olsson, C. Olah, D. Hernandez, D. Drain, D. Ganguli, D. Li, E. Tran-Johnson, E. Perez, J. Kerr, J. Mueller, J. Ladish, J. Landau, K. Ndousse, K. Lukosuite, L. Lovitt, M. Sellit… find our card
  5. C. Blakeney, M. Paul, B. W. Larsen, S. Owen, and J. Frankle. Does your data spark joy? performance gains from domain upsampling at the end of training, 2024. URL https://arxiv.org/abs/ 2406.03476 . ↗
  6. G. Chen, M. Liao, C. Li, and K. Fan. Alphamath almost zero: process supervision without process, 2024. find our card
  7. K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, C. Hesse, and J. Schulman. Training verifiers to solve math word problems, 2021. find our card
  8. Y. Du, S. Li, A. Torralba, J. B. Tenenbaum, and I. Mordatch. Improving factuality and reasoning in language models through multiagent debate, 2023. find our card
  9. J. S. B. T. Evans. Heuristic and analytic processes in reasoning. British Journal of Psychology, 75(4): 451–468, 1984. find
  10. X. Feng, Z. Wan, M. Wen, S. M. McAleer, Y. Wen, W. Zhang, and J. Wang. Alphazero-like tree-search can guide large language model decoding and training, 2024. find
  11. L. Gao, A. Madaan, S. Zhou, U. Alon, P. Liu, Y. Yang, J. Callan, and G. Neubig. Pal: Program-aided language models, 2023. URL https://arxiv.org/abs/2211.10435 . ↗
  12. S. Goyal, Z. Ji, A. S. Rawat, A. K. Menon, S. Kumar, and V. Nagarajan. Think before you speak: Training language models with pause tokens, 2024. URL https://arxiv.org/abs/2310.02226 . ↗
  13. D. Hendrycks, C. Burns, S. Kadavath, A. Arora, S. Basart, E. Tang, D. Song, and J. Steinhardt. Measuring mathematical problem solving with the math dataset, 2021. find our card
  14. J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de Las Casas, L. A. Hendricks, J. Welbl, A. Clark, T. Hennigan, E. Noland, K. Millican, G. van den Driessche, B. Damoc, A. Guy, S. Osindero, K. Simonyan, E. Elsen, J. W. Rae, O. Vinyals, and L. Sifre. Training compute-opt… find
  15. J. Huang, X. Chen, S. Mishra, H. S. Zheng, A. W. Yu, X. Song, and D. Zhou. Large language models cannot self-correct reasoning yet, 2023. find our card
  16. A. L. Jones. Scaling scaling laws with board games, 2021. URL https://arxiv.org/abs/2104. 03113 . ↗
  17. D. Kahneman. Maps of bounded rationality: Psychology for behavioral economics. The American Economic Review, 93(5):1449–1475, 2003. find
  18. D. Kahneman. Thinking, fast and slow. Farrar, Straus and Giroux, New York, first paperback edition edition, 2013. find
  19. L. Kocsis and C. Szepesv’ari. Bandit based monte-carlo planning. In European conference on machine learning, pages 282–293. Springer, 2006. find
  20. A. Lewkowycz, A. Andreassen, D. Dohan, E. Dyer, H. Michalewski, V. Ramasesh, A. Slone, C. Anil, I. Schlag, T. Gutman-Solo, Y. Wu, B. Neyshabur, G. Gur-Ari, and V. Misra. Solving quantitative reasoning problems with language models, 2022. find our card
  21. Y. Li, Z. Lin, S. Zhang, Q. Fu, B. Chen, J.-G. Lou, and W. Chen. Making large language models better reasoners with step-aware verifier, 2023. find
  22. H. Lightman, V. Kosaraju, Y. Burda, H. Edwards, B. Baker, T. Lee, J. Leike, J. Schulman, I. Sutskever, and K. Cobbe. Let’s verify step by step, 2023. find our card
  23. A. Madaan, N. Tandon, P. Gupta, S. Hallinan, L. Gao, S. Wiegreffe, U. Alon, N. Dziri, S. Prabhumoye, Y. Yang, S. Gupta, B. P. Majumder, K. Hermann, S. Welleck, A. Yazdanbakhsh, and P. Clark. Self-refine: Iterative refinement with self-feedback, 2023. find
  24. N. McAleese, R. Pokorny, J. F. Cerón Uribe, E. Nitishinskaya, M. Trębacz, and J. Leike. Llm critics help catch llm bugs. OpenAI, 2024. find
  25. OpenAI. Gpt-4 technical report, 2024. find our card
  26. Y. Qin, S. Liang, Y. Ye, K. Zhu, L. Yan, Y. Lu, Y. Lin, X. Cong, X. Tang, B. Qian, S. Zhao, L. Hong, R. Tian, R. Xie, J. Zhou, M. Gerstein, D. Li, Z. Liu, and M. Sun. Toolllm: Facilitating large language models to master 16000+ real-world apis, 2023. URL https://arxiv.org/abs/2307.16789 . ↗ our card
  27. C. Qu, S. Dai, X. Wei, H. Cai, S. Wang, D. Yin, J. Xu, and J.-R. Wen. Tool learning with large language models: A survey, 2024. URL https://arxiv.org/abs/2405.17935 . ↗
  28. Y. Qu, T. Zhang, N. Garg, and A. Kumar. Recursive introspection: Teaching foundation models how to self-improve. 2024. find
  29. N. Sardana and J. Frankle. Beyond chinchilla-optimal: Accounting for inference in language model scaling laws, 2023. find
  30. W. Saunders, C. Yeh, J. Wu, S. Bills, L. Ouyang, J. Ward, and J. Leike. Self-critiquing models for assisting human evaluators, 2022. find our card
  31. A. Setlur, S. Garg, X. Geng, N. Garg, V. Smith, and A. Kumar. Rl on incorrect synthetic data scales the efficiency of llm math reasoning by eight-fold. arXiv preprint arXiv:2406.14532, 2024. find
  32. Z. Shao, P. Wang, Q. Zhu, R. Xu, J. Song, X. Bi, H. Zhang, M. Zhang, Y. K. Li, Y. Wu, and D. Guo. Deepseekmath: Pushing the limits of mathematical reasoning in open language models, 2024. find our card
  33. A. Sharma, S. Keh, E. Mitchell, C. Finn, K. Arora, and T. Kollar. A critical evaluation of ai feedback for aligning large language models, 2024. URL https://arxiv.org/abs/2402.12366 . — [arXiv:2402.12366](https://arxiv.org/abs/2402.12366) ↗
  34. N. Shinn, F. Cassano, E. Berman, A. Gopinath, K. Narasimhan, and S. Yao. Reflexion: Language agents with verbal reinforcement learning, 2023. find our card
  35. A. Singh, J. D. Co-Reyes, R. Agarwal, A. Anand, P. Patil, X. Garcia, P. J. Liu, J. Harrison, J. Lee, K. Xu, A. Parisi, A. Kumar, A. Alemi, A. Rizkowsky, A. Nova, B. Adlam, B. Bohnet, G. Elsayed, H. Sedghi, I. Mordatch, I. Simpson, I. Gur, J. Snoek, J. Pennington, J. Hron, K. Kenealy, K. Swersky, K.… find our card
  36. C. Snell, E. Wallace, D. Klein, and S. Levine. Predicting emergent capabilities by finetuning. Conference on Language Modeling 2024, 2024. find
  37. K. Stechly, M. Marquez, and S. Kambhampati. Gpt-4 doesn’t know it’s wrong: An analysis of iterative prompting for reasoning problems, 2023. find
  38. R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. Second edition, 2018. find
  39. G. Team. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context, 2024. find our card
  40. Y. Tian, B. Peng, L. Song, L. Jin, D. Yu, H. Mi, and D. Yu. Toward self-improvement of llms via imagination, searching, and criticizing, 2024. find
  41. H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, D. Bikel, L. Blecher, C. C. Ferrer, M. Chen, G. Cucurull, D. Esiobu, J. Fernandes, J. Fu, W. Fu, B. Fuller, C. Gao, V. Goswami, N. Goyal, A. Hartshorn, S. Hosseini, R. Hou, H. Inan,… find
  42. J. Uesato, N. Kushman, R. Kumar, F. Song, N. Siegel, L. Wang, A. Creswell, G. Irving, and I. Higgins. Solving math word problems with process- and outcome-based feedback, 2022. find our card
  43. K. Valmeekam, M. Marquez, and S. Kambhampati. Can large language models really improve by self-critiquing their own plans?, 2023. find
  44. P. Villalobos and D. Atkinson. Trading off compute in training and inference, 2023. URL https: //epochai.org/blog/trading-off-compute-in-training-and-inference . Accessed: 2024-07-03. find
  45. P. Wang, L. Li, Z. Shao, R. X. Xu, D. Dai, Y. Li, D. Chen, Y. Wu, and Z. Sui. Math-shepherd: Verify and reinforce llms step-by-step without human annotations, 2023. find our card
  46. R. Wang, E. Zelikman, G. Poesia, Y. Pu, N. Haber, and N. D. Goodman. Hypothesis search: Inductive reasoning with language models, 2024. URL https://arxiv.org/abs/2309.05660 . ↗
  47. J. Wei, X. Wang, D. Schuurmans, M. Bosma, B. Ichter, F. Xia, E. Chi, Q. Le, and D. Zhou. Chain-of-thought prompting elicits reasoning in large language models, 2023. find our card
  48. S. Yao, D. Yu, J. Zhao, I. Shafran, T. L. Griffiths, Y. Cao, and K. Narasimhan. Tree of thoughts: Deliberate problem solving with large language models, 2023. find
  49. Z. Yuan, H. Yuan, C. Li, G. Dong, K. Lu, C. Tan, C. Zhou, and J. Zhou. Scaling relationship on learning mathematical reasoning with large language models, 2023. find our card
  50. E. Zelikman, Y. Wu, J. Mu, and N. D. Goodman. Star: Bootstrapping reasoning with reasoning, 2022. find our card
  51. E. Zelikman, G. Harik, Y. Shao, V. Jayasiri, N. Haber, and N. D. Goodman. Quiet-star: Language models can teach themselves to think before speaking, 2024. URL https://arxiv.org/abs/ 2403.09629 . ↗ our card