Reproduce “Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization” on OpenResearch

About this paper

Zishun Yu, Tengyu Xu, Di Jin, Karthik Abinav Sankararaman, Yun He, Wenxuan Zhou, Zhouhao Zeng, Eryk Helenowski and 4 more

Researchers at MetaAI developed Inference Budget-Constrained Policy Optimization (IBPO), a constrained Reinforcement Learning framework that enables large language models to adapt their reasoning complexity based on query difficulty. The approach achieved up to 5.74% absolute improvement in pass@1 accuracy on the MATH500 dataset while efficiently optimizing inference costs by dynamically allocating more reasoning budget to harder problems.

Key insights

• Framing adaptive reasoning as a constrained Reinforcement Learning problem enables principled optimization of inference budgets by maximizing utility under cost constraints.
• IBPO reveals a fundamental connection between constrained optimization, RL, and iterative weighted Supervised Fine-Tuning (SFT), providing a new interpretative lens for existing methods like RAFT and RFT.
• Models can learn to dynamically allocate more reasoning budget to harder problems, demonstrating the ability to 'understand' query difficulty and validate the 'Think Smarter not Harder' premise.

Summary by alphaXiv.