Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866909470769020928 |
|---|---|
| author | Yu, Zishun Xu, Tengyu Jin, Di Sankararaman, Karthik Abinav He, Yun Zhou, Wenxuan Zeng, Zhouhao Helenowski, Eryk Zhu, Chen Wang, Sinong Ma, Hao Fang, Han |
| author_facet | Yu, Zishun Xu, Tengyu Jin, Di Sankararaman, Karthik Abinav He, Yun Zhou, Wenxuan Zeng, Zhouhao Helenowski, Eryk Zhu, Chen Wang, Sinong Ma, Hao Fang, Han |
| contents | Solving mathematics problems has been an intriguing capability of large language models, and many efforts have been made to improve reasoning by extending reasoning length, such as through self-correction and extensive long chain-of-thoughts. While promising in problem-solving, advanced long reasoning chain models exhibit an undesired single-modal behavior, where trivial questions require unnecessarily tedious long chains of thought. In this work, we propose a way to allow models to be aware of inference budgets by formulating it as utility maximization with respect to an inference budget constraint, hence naming our algorithm Inference Budget-Constrained Policy Optimization (IBPO). In a nutshell, models fine-tuned through IBPO learn to ``understand'' the difficulty of queries and allocate inference budgets to harder ones. With different inference budgets, our best models are able to have a $4.14$\% and $5.74$\% absolute improvement ($8.08$\% and $11.2$\% relative improvement) on MATH500 using $2.16$x and $4.32$x inference budgets respectively, relative to LLaMA3.1 8B Instruct. These improvements are approximately $2$x those of self-consistency under the same budgets. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2501_17974 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization Yu, Zishun Xu, Tengyu Jin, Di Sankararaman, Karthik Abinav He, Yun Zhou, Wenxuan Zeng, Zhouhao Helenowski, Eryk Zhu, Chen Wang, Sinong Ma, Hao Fang, Han Artificial Intelligence Solving mathematics problems has been an intriguing capability of large language models, and many efforts have been made to improve reasoning by extending reasoning length, such as through self-correction and extensive long chain-of-thoughts. While promising in problem-solving, advanced long reasoning chain models exhibit an undesired single-modal behavior, where trivial questions require unnecessarily tedious long chains of thought. In this work, we propose a way to allow models to be aware of inference budgets by formulating it as utility maximization with respect to an inference budget constraint, hence naming our algorithm Inference Budget-Constrained Policy Optimization (IBPO). In a nutshell, models fine-tuned through IBPO learn to ``understand'' the difficulty of queries and allocate inference budgets to harder ones. With different inference budgets, our best models are able to have a $4.14$\% and $5.74$\% absolute improvement ($8.08$\% and $11.2$\% relative improvement) on MATH500 using $2.16$x and $4.32$x inference budgets respectively, relative to LLaMA3.1 8B Instruct. These improvements are approximately $2$x those of self-consistency under the same budgets. |
| title | Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization |
| topic | Artificial Intelligence |
| url | https://arxiv.org/abs/2501.17974 |