Elamur Research⌕
☰
home / papers / 2024 / Learning How Hard to Think: Input-Adaptive Allocation of LM Computation

Learning How Hard to Think: Input-Adaptive Allocation of LM Computation

Paper · 2024 · Mehul Damani, Idan Shenfeld, Andi Peng et al.

Summary

Costly decoding strategies such as best-of-k sampling with reranking, or calling a stronger model, improve LM outputs, but they are usually applied uniformly to every query. Some queries are trivial and others are beyond reach, so much of that extra compute is wasted. Damani et al. address this by deciding, before any generation, how much computation each input should get. They train a lightweight predictor, either a two-layer MLP probe on the base LM's last hidden state or a LoRA fine-tune, that estimates the marginal reward gained from each additional unit of budget. A greedy allocation algorithm, provably optimal because the constraints form a matroid, then distributes a fixed total budget across a batch (online) or through a fixed query-to-budget mapping (offline). The approach is evaluated on math, code (TACO) and chat (LMSYS-Chat) for adaptive best-of-k, and on two routing setups. Adaptive allocation saves up to half the compute at equal quality, or raises quality by up to 10% at the same budget. The findings also suggest that LM hidden states already encode how difficult a query is.

Key points

  • Two decoding procedures:
    • Adaptive best-of-k varies the number of samples passed to a reward-model reranker for each query.
    • Routing sends each query to either a cheap weak decoder or an expensive strong one.
  • Best-of-k setups:
    • Code: a TACO subset (10K train, 1K validation, 1K test) with Starcoder-15B and a cap of 100 samples.
    • Chat: LMSYS-Chat (50K train, 5K test) with Gemma-2B-it, scored by NCSOFT/Llama-3-OffsetBias-RM-8B with a cap of 8 samples.
    • Math: a cap of 128 samples per query.
  • Best-of-k compute savings: On Math and Code, adaptive sampling reaches the performance of non-adaptive baselines with up to 50% less compute. On Chat, it matches reward with up to 10% less compute.
  • Routing setups and results:
    • Model-size routing: Gemma-2B-it vs Gemma-7B-it.
    • Decoding-procedure routing: plain sampling vs value-augmented sampling (VAS), both with Llama-7B, on the harmless split of Anthropic-HH, scored by OpenAssistant/reward-model-deberta-v3-large-v2.
    • In both cases, the strong decoder's quality is matched while using it for only 50–75% of queries.
  • Low overhead: The predictor runs without generating any outputs. Its training labels come from many sampled responses per query, such as 128 for Math and 100 for Code.

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

 Institutions

🎓Massachusetts Institute of TechnologyCambridge · United States274 works · 82,013 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksArts and Humanities · 17,959 worksLanguage and Linguistics · 17,789 works

Cite

Cite as Mehul Damani, Idan Shenfeld, Andi Peng, Andreea Bobu and Jacob Andreas. 2024. Learning How Hard to Think: Input-Adaptive Allocation of LM Computation. arXiv:2410.04707.

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

@misc{damani2024learning,
  title = {{Learning How Hard to Think: Input-Adaptive Allocation of LM Computation}},
  author = {Mehul Damani, Idan Shenfeld, Andi Peng, Andreea Bobu and Jacob Andreas},
  year = {2024},
  doi = {10.48550/arxiv.2410.04707},
  eprint = {2410.04707},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2410.04707}
}

References

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

  1. Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, et al. Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05862, 2022. find
  2. Andrea Banino, Jan Balaguer, and Charles Blundell. Pondernet: Learning to ponder. In 8th ICML Workshop on Automated Machine Learning (AutoML), 2021.
  3. Ahmad Beirami, Alekh Agarwal, Jonathan Berant, Alexander D’Amour, Jacob Eisenstein, Chirag Nagpal, and Ananda Theertha Suresh. Theoretical guarantees on the best-of-n alignment policy. arXiv preprint arXiv:2401.01879, 2024. — [arXiv:2401.01879](https://arxiv.org/abs/2401.01879) find
  4. Bradley Brown, Jordan Juravsky, Ryan Ehrlich, Ronald Clark, Quoc V Le, Christopher Ré, and Azalia Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling. arXiv preprint arXiv:2407.21787, 2024.
  5. Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations, 2019. find
  6. Yilun Du, Shuang Li, Antonio Torralba, Joshua B. Tenenbaum, and Igor Mordatch. Improving factuality and reasoning in language models through multiagent debate. In Forty-first International Conference on Machine Learning, 2024.
  7. Jack Edmonds. Matroids and the greedy algorithm. Mathematical programming, 1:127–136, 1971. find
  8. Leo Gao, John Schulman, and Jacob Hilton. Scaling laws for reward model overoptimization. In International Conference on Machine Learning, pp. 10835–10866. PMLR, 2023.
  9. Alex Graves. Adaptive computation time for recurrent neural networks. arXiv:1603.08983, 2016. arXiv preprint Seungwook Han, Idan Shenfeld, Akash Srivastava, Yoon Kim, and Pulkit Agrawal. Value augmented sampling for language model alignment and personalization. arXiv preprint arXiv:2405.06639, 2024. find
  10. Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the MATH dataset. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
  11. Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022.
  12. Athul Paul Jacob, Yikang Shen, Gabriele Farina, and Jacob Andreas. The consensus game: Language model generation via equilibrium search. In The Twelfth International Conference on Learning Representations, 2024. find
  13. Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.06825, 2023. find
  14. Jia Li, Edward Beeching, Lewis Tunstall, Ben Lipkin, Roman Soletskyi, Shengyi Costa Huang, Kashif Rasul, Longhui Yu, Albert Jiang, Ziju Shen, Zihan Qin, Bin Dong, Li Zhou, Yann Fleureau, Guillaume Lample, and Stanislas Polu. Numinamath. [https://huggingface.co/AI-MO/NuminaMath-CoT](https://github.c… find
  15. Rongao Li, Jie Fu, Bo-Wen Zhang, Tao Huang, Zhihong Sun, Chen Lyu, Guang Liu, Zhi Jin, and Ge Li. Taco: Topics in algorithmic code generation dataset. arXiv preprint arXiv:2312.14852, 2023.
  16. Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, et al. Competition-level code generation with alphacode. Science, 378(6624):1092–1097, 2022.
  17. Jiacheng Liu, Andrew Cohen, Ramakanth Pasunuru, Yejin Choi, Hannaneh Hajishirzi, and Asli Celikyilmaz. Don’t throw away your value model! generating more preferable text with value-guided monte-carlo tree search decoding. In First Conference on Language Modeling, 2024. — Don't throw away your value model! Generating more preferable text with Value-Guided Monte-Carlo Tree Search decoding
  18. Ziqin Luo, Haixia Han, Haokun Zhao, Guochao Jiang, Chengyu Du, Tingyun Li, Jiaqing Liang, Deqing Yang, and Yanghua Xiao. Sed: Self-evaluation decoding enhances large language models for better generation. arXiv preprint arXiv:2405.16552, 2024. — [arXiv:2405.16552](https://arxiv.org/abs/2405.16552) find
  19. Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2112.00114, 2021.
  20. David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484–489, 2016. find
  21. Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling llm test-time compute optimally can be more effective than scaling model parameters. arXiv preprint arXiv:2408.03314, 2024.
  22. Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based on gemini research and technology. arXiv preprint arXiv:2403.08295, 2024. find
  23. Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc V Le, Ed H. Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. In The Eleventh International Conference on Learning Representations, 2023.
  24. Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824–24837, 2022.
  25. Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. Advances in Neural Information Processing Systems, 36, 2024. find
  26. Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Tianle Li, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zhuohan Li, Zi Lin, Eric Xing, Joseph E. Gonzalez, Ion Stoica, and Hao Zhang. LMSYS-chat-1m: A large-scale real-world LLM conversation dataset. In The Twelfth International Conference on Learning Repres…
  27. Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuurmans, Claire Cui, Olivier Bousquet, Quoc V Le, and Ed H. Chi. Least-to-most prompting enables complex reasoning in large language models. In The Eleventh International Conference on Learning Representations, 20…

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/damani-2024-learning-hard-think-input-adaptive/