# Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization

arXiv: https://arxiv.org/abs/2501.17974 · HTML: https://mlsanity.com/paper/2501.17974/

## Key facts

- Cited by 29 papers in the map (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 57 of its 83 references that are in the map.
- First posted on arXiv 29 Jan 2025.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)

## Builds on (most cited first, 30 of 57)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Self-Refine: Iterative Refinement with Self-Feedback](https://mlsanity.com/paper/2303.17651/) (2023)
- [Exploring Simple Siamese Representation Learning](https://mlsanity.com/paper/2011.10566/) (2020)
- [Fine-Tuning Language Models from Human Preferences](https://mlsanity.com/paper/1909.08593/) (2019)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Curiosity-driven Exploration by Self-supervised Prediction](https://mlsanity.com/paper/1705.05363/) (2017)
- [Large Language Models Cannot Self-Correct Reasoning Yet](https://mlsanity.com/paper/2310.01798/) (2023)
- [Encouraging Divergent Thinking in Large Language Models through Multi-Agent Debate](https://mlsanity.com/paper/2305.19118/) (2023)
- [STaR: Bootstrapping Reasoning With Reasoning](https://mlsanity.com/paper/2203.14465/) (2022)
- [Measuring Coding Challenge Competence With APPS](https://mlsanity.com/paper/2105.09938/) (2021)
- [Introduction to Online Convex Optimization](https://mlsanity.com/paper/1909.05207/) (2019)
- [Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1910.00177/) (2019)
- [RAFT: Reward rAnked FineTuning for Generative Foundation Model Alignment](https://mlsanity.com/paper/2304.06767/) (2023)
- [Reinforcement Learning with Deep Energy-Based Policies](https://mlsanity.com/paper/1702.08165/) (2017)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Introduction to Multi-Armed Bandits](https://mlsanity.com/paper/1904.07272/) (2019)
- [Training Language Models to Self-Correct via Reinforcement Learning](https://mlsanity.com/paper/2409.12917/) (2024)

## Built on by (most cited first, 29 of 29)

- [Stop Overthinking: A Survey on Efficient Reasoning for Large Language Models](https://mlsanity.com/paper/2503.16419/) (2025)
- [Towards Reasoning Era: A Survey of Long Chain-of-Thought for Reasoning Large Language Models](https://mlsanity.com/paper/2503.09567/) (2025)
- [A Survey of Efficient Reasoning for Large Reasoning Models: Language, Multimodality, and Beyond](https://mlsanity.com/paper/2503.21614/) (2025)
- [A Survey of Context Engineering for Large Language Models](https://mlsanity.com/paper/2507.13334/) (2025)
- [Efficient Inference for Large Reasoning Models: A Survey](https://mlsanity.com/paper/2503.23077/) (2025)
- [AdaCtrl: Towards Adaptive and Controllable Reasoning via Difficulty-Aware Budgeting](https://mlsanity.com/paper/2505.18822/) (2025)
- [Large Language Models for Planning: A Comprehensive and Systematic Survey](https://mlsanity.com/paper/2505.19683/) (2025)
- [Reasoning on a Budget: A Survey of Adaptive and Controllable Test-Time Compute in LLMs](https://mlsanity.com/paper/2507.02076/) (2025)
- [BudgetThinker: Empowering Budget-aware LLM Reasoning with Control Tokens](https://mlsanity.com/paper/2508.17196/) (2025)
- [A Comprehensive Survey on Trustworthiness in Reasoning with Large Language Models](https://mlsanity.com/paper/2509.03871/) (2025)
- [Counterfactual VLA: Self-Reflective Vision-Language-Action Model with Adaptive Reasoning](https://mlsanity.com/paper/2512.24426/) (2025)
- [Know What You Don't Know: Uncertainty Calibration of Process Reward Models](https://mlsanity.com/paper/2506.09338/) (2025)
- [Not All Thoughts are Generated Equal: Efficient LLM Reasoning via Multi-Turn Reinforcement Learning](https://mlsanity.com/paper/2505.11827/) (2025)
- [Generative AI Act II: Test Time Scaling Drives Cognition Engineering](https://mlsanity.com/paper/2504.13828/) (2025)
- [From Efficiency to Adaptivity: A Deeper Look at Adaptive Reasoning in Large Language Models](https://mlsanity.com/paper/2511.10788/) (2025)
- [Rethinking the Role of Prompting Strategies in LLM Test-Time Scaling: A Perspective of Probability Theory](https://mlsanity.com/paper/2505.10981/) (2025)
- [C3PO: Optimized Large Language Model Cascades with Probabilistic Cost Constraints for Reasoning](https://mlsanity.com/paper/2511.07396/) (2025)
- [Neural Chain-of-Thought Search: Searching the Optimal Reasoning Path to Enhance Large Language Models](https://mlsanity.com/paper/2601.11340/) (2026)
- [Ares: Adaptive Reasoning Effort Selection for Efficient LLM Agents](https://mlsanity.com/paper/2603.07915/) (2026)
- [DynScaling: Efficient Verifier-free Inference Scaling via Dynamic and Integrated Sampling](https://mlsanity.com/paper/2506.16043/) (2025)
- [DeepCompress: A Dual Reward Strategy for Dynamically Exploring and Compressing Reasoning Chains](https://mlsanity.com/paper/2510.27419/) (2025)
- [RADAR: Reasoning-Ability and Difficulty-Aware Routing for Reasoning LLMs](https://mlsanity.com/paper/2509.25426/) (2025)
- [Adaptive Rectification Sampling for Test-Time Compute Scaling](https://mlsanity.com/paper/2504.01317/) (2025)
- [UniScale: Adaptive Unified Inference Scaling via Online Joint Optimization of Model Routing and Test-Time Scaling](https://mlsanity.com/paper/2605.30898/) (2026)
- [EasyMath: A 0-shot Math Benchmark for SLMs](https://mlsanity.com/paper/2505.14852/) (2025)
- [HyPER: Bridging Exploration and Exploitation for Scalable LLM Reasoning with Hypothesis Path Expansion and Reduction](https://mlsanity.com/paper/2602.06527/) (2026)
- [Think Longer to Explore Deeper: Learn to Explore In-Context via Length-Incentivized Reinforcement Learning](https://mlsanity.com/paper/2602.11748/) (2026)
- [Post-Reasoning: Improving the Performance of Non-Thinking Models at No Cost](https://mlsanity.com/paper/2605.06165/) (2026)
- [LIMBO: Lifelong Inference-Time Memory and Budget Optimization for LLM Agents](https://mlsanity.com/paper/2609.14138/) (2026)

## Abstract (arXiv)

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.16x and 4.32x inference budgets respectively, relative to LLaMA3.1 8B Instruct. These improvements are approximately 2x those of self-consistency under the same budgets.