🤖 AI

Learning When to Think: Adaptive Reasoning for Test-Time Compute Allocation

This paper introduces a method for training reasoning language models to adaptively allocate test-time compute by learning to select between "NoThink," "Short," and "Long" reasoning modes via Group Relative Policy Optimization, achieving significant token reduction while maintaining or improving accuracy across various benchmarks.

Original authors: Gijs Kassenaar, Zhao Yang, Vincent François-Lavet

Published 2026-08-21

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Gijs Kassenaar, Zhao Yang, Vincent François-Lavet

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Large language models have become remarkably skilled at solving complex problems, often by breaking a question down into a long series of intermediate steps, much like a human working through a difficult math problem on a piece of paper. This step-by-step approach, known as chain-of-thought reasoning, allows these systems to reach high levels of accuracy on tasks involving mathematics and coding. However, this capability comes with a significant cost: the models often generate far more text than is necessary, consuming vast amounts of computing power and time. Even when a problem is simple, the model might produce a lengthy, verbose explanation, wasting resources on tasks that could be solved instantly. Researchers have long suspected that these systems lack the ability to judge how much effort a specific problem actually requires, leading to a situation where easy questions get the same heavy treatment as difficult ones.

A team of researchers at Vrije Universiteit Amsterdam set out to teach a model to make this judgment for itself. They worked with a relatively small reasoning model, containing 1.5 billion parameters, which had been trained to solve mathematical problems. Instead of forcing the model to always think for a fixed amount of time, they gave it a simple choice at the very start of its response: it could decide to answer immediately without thinking, reason briefly, or engage in extended, deep reasoning. The model was trained using a method called reinforcement learning, where it receives feedback based on whether its answer is correct and how efficiently it arrived there. The goal was to see if the model could learn to match its effort to the difficulty of the problem, saving time on easy tasks while still dedicating enough resources to solve hard ones.

The researchers designed a system where the model emits a single word as its first token to signal its chosen mode of operation. If the model selects "NoThink," it attempts to answer directly. If it chooses "Short," it is allowed to generate a brief explanation, but with a strict limit on how many words it can use. If it picks "Long," it is permitted to reason at length without a cap. To ensure the model actually stuck to these choices, the researchers imposed hard limits on the length of the response for the shorter modes. If a model tried to use the "Short" mode but kept writing past the limit, the system would mark the attempt as incorrect, regardless of whether the final answer was right. This forced the model to learn that the label it chose at the beginning had real consequences for its behavior.

To prevent the model from simply ignoring the choice and defaulting to the longest, safest option every time, the researchers adjusted the rewards it received. They made the "NoThink" and "Short" modes more rewarding for very short answers, while the "Long" mode offered a steady reward regardless of length. This created a landscape where the best strategy depended on the problem: for a trivial question, the quickest answer yielded the highest score, while a difficult problem required the extended effort of the "Long" mode to succeed. They also added a balancing mechanism that gently nudged the model to use all three options, ensuring it didn't collapse into using just one strategy for every single problem.

The results showed that the model successfully learned to sort problems by difficulty. On a set of 500 held-out math problems, the model began to route easy questions to the "NoThink" or "Short" modes and difficult questions to the "Long" mode. Over the course of training, the model's accuracy remained nearly identical to the original, unmodified version, but the average length of its responses dropped significantly. The researchers found that the new policy reduced the average number of words generated by 41 percent, cutting the response length from nearly 4,800 tokens down to about 2,800. This meant the model was solving the same number of problems correctly while using less than half the computational effort.

Crucially, this efficiency did not come at the cost of performance on problems the model had never seen before. When tested on a different set of elementary math word problems, the model saved an even larger amount of text, reducing the token count by 76 percent while maintaining high accuracy. On a set of extremely difficult competition math problems, where deep reasoning is almost always required, the model correctly chose the "Long" mode for nearly every question, matching the performance of the baseline without wasting time on shortcuts. The study demonstrates that a reasoning model can be trained to self-regulate its effort, allocating more thinking time only when the problem demands it. This suggests a path forward for making artificial intelligence more efficient, allowing systems to be both smart and economical, rather than simply verbose.

Technical Summary: Learning When to Think: Adaptive Reasoning for Test-Time Compute Allocation

Problem Statement

Reasoning language models, particularly those fine-tuned with Reinforcement Learning with Verifiable Rewards (RLVR), typically operate under a fixed token budget applied uniformly across all prompts. This static approach leads to two primary inefficiencies:

Over-computation: Easy problems receive excessive reasoning tokens, increasing inference latency and cost without improving accuracy.
Under-computation: Difficult problems may be cut off by the fixed ceiling before a solution is found.
Inefficiency in Small Models: Smaller models often exhibit "verbosity compensation," generating unnecessarily long explanations to compensate for limited capacity, and incorrect answers tend to be significantly longer than correct ones.

Existing solutions, such as temporal discounting in Group Relative Policy Optimization (GRPO), impose a uniform length penalty regardless of problem difficulty. This fails to adaptively allocate compute, pressuring easy and hard questions to be equally brief despite their differing optimal reasoning budgets. Furthermore, previous attempts at adaptive mode selection often rely on supervised pre-training stages, distillation, or complex curriculum staging to prevent "routing collapse," where the policy converges to a single mode.

Methodology

The paper proposes a Three-Mode Self-Routing GRPO framework where a 1.5B distilled reasoning model learns to allocate its own reasoning effort end-to-end, without a separate router network or supervised warm-up.

Core Mechanism

The model emits a discrete routing token as the first token of its response, selecting one of three modes:

NoThink: Answer as quickly as possible (minimal reasoning).
Short: Brief reasoning within a hard token cap.
Long: Extended reasoning without a hard cap.

Key Technical Components

To prevent the policy from collapsing into a single mode (a common failure in discrete routing), the authors employ three specific mechanisms:

Hard Per-Mode Token Caps:
- NoThink: Capped at 1,024 tokens.
- Short: Capped at 3,000 tokens.
- Long: Uncapped.
- Function: These caps ensure behavioral distinctness. If a "Short" response exceeds its cap, it is scored as incorrect regardless of content. This prevents the brief modes from silently absorbing hard problems that require more tokens, maintaining a genuine accuracy gap between modes.
Shaped Reward Surface:
- Each mode has a distinct base reward ( $b$ ) and a per-token discount factor ( $\gamma$ ) applied to the number of tokens generated after the routing token.
- Formulation: $r_i = b \cdot \gamma^{L_i}$ for correct answers, where $L_i$ is the token count.
- Parameters: $b_{NT}=1.3, b_{S}=1.2, b_{L}=1.0$ . The discounts are tuned so that the reward-optimal choice shifts at specific lengths (crossovers at ~800 and ~3,000 tokens). This creates a reward landscape where each mode is optimal for a specific difficulty/length range.
Advantage Balancing Term:
- To counteract the self-reinforcing nature of discrete choices, a balance term is added directly to the advantage estimate: $\hat{A}_i \leftarrow A_i + \beta_{bal}(p^* - f_{mode}(i))$ .
- This term penalizes over-represented modes and boosts under-represented ones. Crucially, it is applied asymmetrically: negative corrections apply only to incorrect rollouts, and positive corrections only to correct ones, ensuring the balance term does not distort the accuracy signal.
Training Modifications:
- No KL Penalty: The KL penalty toward the reference policy is removed ( $\beta=0$ ) to allow the policy to explore new behaviors (like answering quickly) that the base distilled model rarely produces.
- Token-Mean Aggregation: Instead of sequence-mean aggregation, token-mean aggregation is used to ensure long sequences do not receive weaker gradients than short ones.
- Forced Warmup: During initial steps, the model is forced to generate one rollout per mode to stabilize early dynamics before switching to free rollouts.

Key Contributions

End-to-End Adaptive Routing: Demonstrated that a 1.5B reasoning model can learn to route problems by difficulty using a single discrete token as the first output, trained entirely within GRPO without separate router networks, supervised pre-training, or distillation.
Stability Mechanisms: Showed that the combination of hard per-mode caps, shaped rewards, and an advantage-based balance term effectively prevents routing collapse, a persistent issue in discrete mode selection.
Efficiency on the Pareto Frontier: The resulting policy achieves accuracy comparable to the base model while significantly reducing token usage, positioning the model strictly above the accuracy-length Pareto frontier traced by single-mode baselines.

Experimental Results

The method was evaluated on a 1.5B model (DeepSeek-R1-Distill-Qwen-1.5B) trained on MATH and validated on MATH-500, GSM8K, and AIME.

Mode Emergence: Across three seeds, the three modes persisted without collapsing. By the final checkpoint, the distribution settled near 20% (NoThink), 32% (Short), and 47% (Long).
Difficulty Sorting: The routing was semantically meaningful.
- In-Distribution (MATH-500): The model allocated "NoThink" and "Short" to easier problems and "Long" to harder ones. Per-mode accuracy inverted during training: brief modes became more accurate than the Long mode on the problems they were assigned, indicating successful difficulty sorting.
- Correlation: Routing distribution correlated monotonically with MATH-500 difficulty levels (1–5), with "NoThink" dominating level 1 and "Long" dominating level 5.
Efficiency Gains:
- MATH-500: The policy achieved 0.782 accuracy (vs. 0.796 for the base) while reducing mean response length from 4,796 to 2,811 tokens (41% reduction).
- Pareto Frontier: The adaptive policy lay strictly above the Pareto frontier of single-mode baselines (NoThink-only, Short-only, Long-only), delivering higher accuracy for the same token budget or the same accuracy for fewer tokens.
Out-of-Distribution Transfer:
- GSM8K (Easier): Achieved a 76% token reduction compared to the base model while exceeding the accuracy of fixed-brief baselines at similar lengths.
- AIME (Harder): The router correctly identified that most problems required extended reasoning, allocating nearly all queries to "Long" and matching the frontier's accuracy with a 12% token reduction.

Significance and Claims

The paper claims to build a reasoning model that adaptively chooses how much to reason for each problem. The significance lies in:

Token Efficiency: Substantially cutting inference costs (up to 76% on easier tasks) without sacrificing accuracy.
Simplicity: Achieving adaptive behavior without complex external routers, supervised pre-training, or multi-stage curricula.
Generalizability: The routing mechanism transfers zero-shot to datasets with different difficulty distributions (GSM8K, AIME), saving compute where problems are easy and preserving budget where they are hard.

The authors note limitations, including the reliance on a single model size and task distribution, and the need for task-specific tuning of token caps. They conclude that the primary contribution is a token-efficiency mechanism that pushes the accuracy-length trade-off outward, rather than a mechanism that fundamentally increases the model's latent reasoning capability.