Research map / Large language models / Reasoning and chain of thought
Learning How Hard to Think: Input-Adaptive Allocation of LM Computation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Predicts per-query marginal reward gains of extra decoding compute from LM hidden states (or LoRA) and allocates budgets adaptively across best-of-k and routing decoders.
Problem
Decoding procedures such as best-of-k and routing spend the same compute on every query, although queries differ in how much extra computation helps; prior adaptive allocation needed costly a posteriori difficulty estimates.
Approach
Train a lightweight difficulty model (MLP probe on LM hidden states or LoRA-finetuned LM) to predict the vector of marginal rewards for each extra unit of compute before decoding, then allocate budgets greedily under an average-budget constraint (matroid greedy), online or via binned offline policies.
Key results
- On Math and Code, adaptive best-of-k matches uniform best-of-k success while using 25-50% less compute at high budgets.
- On Chat tranches subset, adaptive allocation reduces budget by 25-40% while matching best-of-k reward.
- In routing, the learned predictor calls the strong decoder only 50-75% of the time while matching its performance.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use adaptive best-of-k allocation instead of uniform best-of-k | Online and Offline Ada-BoK · Code, low-budget regime (B < 8) | expected success rate | Helped |
| Use offline adaptive allocation instead of uniform best-of-k | Offline Ada-BoK · Code | expected success rate | Helped |
| Use online adaptive allocation instead of uniform best-of-k | Online Ada-BoK · Code, high-budget regime | expected success rate | Hurt |
| Use online or offline adaptive allocation instead of uniform best-of-k | Ada-BoK · Math | expected success rate: 25--50% less computation | Helped |
| Use adaptive allocation instead of uniform best-of-k | Full Ada-BoK · Chat (Full) | expected reward: 0-10% reduction in compute budget | Helped |
| Use adaptive allocation instead of uniform best-of-k | Tranches Ada-BoK · Chat (Tranches) | expected reward: 25-40% reduction in compute budget | Helped |
| Route queries adaptively rather than randomly to the stronger decoder | Online Routing · Model-size routing and value-augmented sampling | expected reward: 25-40% reduction in calls to the more expensive decoding scheme | Helped |
| Use the two-stage math verification pipeline, adding an LLM evaluator for stage-one failures (ablation) | Math adaptive best-of-k training and evaluation · Math | verification quality and predictor training stability | Helped |
Surprising findings
- Online adaptive allocation in code falls below the best-of-k curve at high budgets because small prediction errors on near-zero success queries get large budgets.
- Routing sometimes outperforms the strong decoder, since the weak decoder occasionally wins on some queries.
Limitations
- Math experiments assume access to a verifier; practical deployment would need reward models.
- Online allocation is sensitive to prediction errors for queries with near-zero success probability.
- Chat gains are modest on the full test set (0-10% budget reduction).
Experimental setup
Datasets: TACO Numina-COT LMSYS-Chat Anthropic HH (harmless subset)
Benchmarks: TACO Numina-COT LMSYS-Chat Anthropic HH
Metrics: expected success rate expected reward loss of marginal reward predictor prediction accuracy
Baselines: Best-of-k (uniform allocation) Oracle allocation Random routing
Builds on: best-of-k sampling LoRA value-augmented sampling
Closest papers by meaning
- Sample Smart, Not Hard: Correctness-First Decoding for Better Reasoning in LLMs (ICLR 2026 · similarity 0.89)
- The Consensus Game: Language Model Generation via Equilibrium Search (ICLR 2024 · similarity 0.88)
- Improving Open-Ended Text Generation via Adaptive Decoding (ICML 2024 · similarity 0.88)
- On Speeding Up Language Model Evaluation (ICLR 2025 · similarity 0.88)
- Entropy-informed Decoding: Adaptive Information-Driven Branching (ICML 2026 · similarity 0.87)
- Towards Theoretical Understanding of Transformer Test-Time Computing: Investigation on In-Context Linear Regression (ICML 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.