When2Think: Learning Difficulty-Aware Length Control for Efficient Hybrid Reasoning Models
Jaejun Shim, HyunJin Kim, Young Jin Kim, JinYeong Bak
cs.AI, cs.LG
2026-09-17
When2Think shapes hybrid-reasoning RL with per-item difficulty stats, lifting AIME24 Pass@3 from 46.0% to 56.0% while cutting tokens 27.9%.
Large reasoning models waste compute in a regular pattern: they pad easy items with unused chain-of-thought, then stop too early on hard ones. Uniform length penalties and hard Think/NoThink routers cut tokens, but they often cut the traces that hard items still need. The paper calls that cost an efficiency tax.
On AIME24, LC-R1 drops mean tokens from 14195 to 7972 versus R1-Distill-Qwen, and Pass@3 falls from 46.0% to 36.0%. AdaptThink is cheaper still at 5806 tokens, with Pass@3 at 44.7%, below the backbone. The failure is misallocation, not missing capacity.
When2Think is a post-training reward-shaping layer on a hybrid policy. Each rollout starts with a mode token. Think runs an explicit trace until an end-of-thinking marker; NoThink answers immediately. Training keeps AdaptThink's importance sampling: an exploration policy draws the mode token uniformly, later tokens come from the current policy, and updates reweight the trajectories so the optimizer sees the target policy.
The new piece is Instance-level Difficulty-Aware Control (IDAC). Before each epoch, a reference policy rolls K traces per problem and caches two numbers: reference accuracy α as a difficulty proxy, and mean length τ as a token budget. Online reward pays an efficiency bonus only when the answer is correct. The bonus decays exponentially with trace length relative to τ, scaled by α. Easy items (high α) get a steep penalty for long Think traces and are pushed toward NoThink. Hard items (low α) keep room to think. Wrong answers get no efficiency bonus, so short-and-wrong is not a cheap win.
There is no learned critic and no online reference-model call inside the update. The reward subtracts α as an instance baseline, then Batch-Wise Standardization (BWS) z-scores rewards across the minibatch at each trajectory index. Every token in a trace shares that advantage, and the policy is updated with a PPO-style clipped surrogate.
The backbone is the 1.5B DeepSeek-R1 distill of Qwen2.5-Math, trained on DeepScaleR (40k verifiable math items). Metrics are Pass@3, averaged over five runs.
| Method | AIME24 Pass@3 | AIME24 tokens | AIME25 Pass@3 |
| R1-Distill-Qwen | 46.0% | 14195 | 32.0% |
| LC-R1 | 36.0% | 7972 | 32.0% |
| AdaptThink | 44.7% | 5806 | 30.7% |
| DeepScaleR-Preview | 58.0% | 8473 | 39.3% |
| When2Think | 56.0% | 10236 | 40.0% |
Versus the backbone, AIME24 Pass@3 rises 10.0 points while tokens fall 27.9%; AIME25 reaches 40.0%. On MATH-500 Level 1, tokens drop from 1199 to 619 at 95.8% accuracy. On Level 5, tokens drop from 6830 to 4554 with 70.7% versus the backbone's 71.0%. The Think rate climbs with difficulty from about 0.2 to above 0.7. On perturbed GSM-Plus items the policy spends 462 extra tokens and lifts Pass@3 by 6.3 points to 85.7%.
Dropping importance sampling but keeping IDAC and BWS still hits 57.3% on AIME24, while GSM-Plus tokens rebound from 1052 to 1652. IS mainly drives NoThink on easy items; accuracy gains come from IDAC and BWS.
For teams already serving small R1-distill math models, this is a post-training recipe rather than a new architecture: cache reference stats, then train with verifiable rewards on when to think and how long. Efficiency is treated as allocation, not as "make the chain shorter."
It is incremental control at 1.5B, not a new reasoning paradigm. DeepScaleR-Preview, trained on the same data, scores 58.0% on AIME24 with 8473 tokens, both more accurate and shorter. When2Think's distinction is holding hard-item accuracy while shrinking easy-item traces. The authors note that larger models already allocate better, so the tax is most visible at small scale. In their 7B table, AIME24 Pass@3 rises 8.7 points versus a same-size R1-Distill, but hard-item tokens go up, which is not the 1.5B story.
The authors flag two limits. Rewards need a verifier, so open-ended tasks need a reward model or humans. Difficulty comes from offline reference stats, which drift when difficulty is ambiguous.
The main analysis stays at 1.5B; 7B numbers are barely discussed in the text. Precomputing reference rollouts each epoch has a cost the paper does not report. Both DeepScaleR and the eval suite are contest math, and some AIME years may overlap the train mix. DeepScaleR-Preview beating When2Think on AIME24 also shows that ordinary RFT already learns some allocation, so IDAC is not the only path.