🤖 AI

Reasoning Can Be Restored by Correcting a Few Decision Tokens

This paper demonstrates that the reasoning capabilities of large language models can be efficiently restored by identifying and intervening on a sparse set of early, high-uncertainty planning tokens where base models disagree with stronger reasoning models, using a lightweight inference-time delegation strategy that significantly boosts performance.

Original authors: Changshuo Shen, Leheng Sheng, Yuxin Chen, An Zhang, Xiang Wang

Published 2026-05-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Changshuo Shen, Leheng Sheng, Yuxin Chen, An Zhang, Xiang Wang

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

The Big Idea: The "GPS vs. Driver" Problem

Imagine you have two drivers trying to get to the same destination (solving a hard math problem).

Driver A (The Base Model): This driver is fast, efficient, and knows the local streets well. However, when faced with a complex, unfamiliar route, they tend to make a few critical mistakes early on—like turning left instead of right at the very first intersection. Once they make that wrong turn, they keep driving confidently down the wrong path, getting further and further lost.
Driver B (The Reasoning Model): This driver is a slow, methodical expert. They double-check every turn, plan the whole route carefully, and almost never get lost. But they are slow and expensive to hire for every single trip.

The Problem: We want Driver A to be as good as Driver B, but we don't want to pay the high cost of using Driver B for the entire trip.

The Discovery: The researchers found that Driver A doesn't actually fail the whole time. They drive perfectly fine for 90% of the journey. They only fail at a tiny number of specific moments—usually right at the beginning—where the road gets tricky. These are the "decision tokens."

The Core Finding: It's All About the First Few Turns

The paper analyzed millions of steps where the two drivers disagreed. They found three surprising things:

The Mistakes are Rare: Only about 8% of the words (tokens) generated by the fast driver were actually where the disagreement happened. The other 92% were fine.
The Mistakes Happen Early: The critical errors almost always happen at the very start of the answer. It's like the driver makes the wrong turn at the first intersection, and the rest of the trip is just them driving confidently in the wrong direction.
The Mistakes are "Planning" Moments: The errors happen when the driver is deciding what to do next (planning), not when they are just doing the math (execution). It's the moment of "Should I go left or right?" not the moment of "1 + 1 = 2."

The Solution: The "Spot-Check" Intervention

Instead of hiring the slow expert for the whole trip, the researchers proposed a clever trick called "Disagreement-Guided Token Intervention."

Here is how it works:

The fast driver (Base Model) starts driving.
A smart "spot-check" system watches the road. It compares the fast driver's next move with what the slow expert (Reasoning Model) would have done.
The Trigger: If the two drivers strongly disagree on the next step (a "spike" in disagreement), the system hits the brakes for a split second.
The Takeover: The slow expert steps in, says the one correct word or phrase needed to fix the direction, and then immediately hands the wheel back to the fast driver.
The Result: The fast driver continues the rest of the trip on the correct path, doing all the routine work itself.

The Analogy: The Editor and the Writer

Think of the Base Model as a fast, energetic writer who can write a whole story in seconds. But sometimes, they get the plot wrong in the first paragraph.
Think of the Reasoning Model as a slow, meticulous editor who knows the perfect story.

Instead of asking the editor to rewrite the whole story (which takes forever), the researchers found that if you just let the editor fix the first few sentences where the plot gets confusing, the writer can then finish the rest of the story perfectly on their own.

By fixing just 8% of the words (the critical decision points), the fast writer ended up performing better than a version of themselves that had been trained to be a "reasoning expert."

Why This Matters

The paper proves that "reasoning" isn't a magical ability that requires a massive brain for every single step. Instead, reasoning is mostly about making the right few choices at the start. Once the path is set correctly, the model can handle the rest easily.

By only intervening at the moments of high uncertainty (where the model is confused), we can restore the model's reasoning power with a tiny amount of help, making it much faster and cheaper than using a giant reasoning model for everything.

In short: You don't need to fix the whole car to get it to run; you just need to fix the steering wheel at the very beginning of the drive.

Technical Summary: Reasoning Can Be Restored by Correcting a Few Decision Tokens

Problem Statement

Large Reasoning Models (LRMs) consistently outperform their base Large Language Model (LLM) counterparts on challenging reasoning benchmarks, a gap often attributed to test-time scaling and reinforcement learning-based post-training (e.g., RLVR). However, the specific mechanisms driving this performance gap remain poorly understood. It is unclear where base models fail during token-by-token generation, whether reasoning capabilities are latent within base models, and how to efficiently bridge the performance gap without the computational cost of running a full LRM for every token.

Methodology

The authors propose a framework to analyze and bridge the base–LRM gap through token-level distributional disagreement.

1. Quantifying Disagreement

The study utilizes likelihood-based divergences, primarily Cross-Entropy (CE) and Reverse Kullback–Leibler (KL) divergence, to measure the disagreement between a base model ( $M_b$ ) and a stronger reasoning model ( $M_r$ ) at each generation step $t$ . The disagreement score $s_t$ is defined as:
$s_t = D_{CE}(p_b(\cdot | x_t), p_r(\cdot | x_t)) = -\sum_{y \in V} p_b(y | x_t) \log p_r(y | x_t)$
where $p_b$ and $p_r$ are the next-token probability distributions.

2. Disagreement-Guided Token Intervention

Based on the analysis of disagreement patterns, the authors introduce a sparse, inference-time delegation scheme called disagreement-guided token intervention.

Mechanism: The system defaults to the efficient base model ( $M_b$ ). However, when the disagreement score $s_t$ exceeds a calibrated global threshold $\tau$ and a local relative spike condition ( $s_t > \lambda \cdot \bar{s}_{t,W}$ ), a "one-token takeover" occurs.
Takeover: At these high-disagreement positions, the reasoning model ( $M_r$ ) generates the next token. Immediately after, control switches back to the base model.
Calibration: Thresholds are calibrated on a held-out set to target a specific "spike ratio" $r$ (e.g., the top 3% or 5% of disagreement scores), ensuring the intervention budget remains small.

Key Findings and Contributions

1. Sparsity and Concentration of Reasoning Gaps

The paper demonstrates that the performance gap between base and reasoning models is highly sparse.

Heavy-Tailed Distribution: A tiny fraction of tokens (e.g., $\sim8\%$ on Qwen3-0.6B) accounts for the vast majority of the distributional disagreement. The Gini coefficient for disagreement mass is extremely high ( $\approx 0.936$ ), indicating that most tokens accumulate little disagreement mass.
Early Concentration: High-disagreement spikes are front-loaded, occurring disproportionately early in the generation trajectory (mode at $u \approx 0.05$ ). This contradicts the intuition that reasoning gains stem primarily from long traces later in the process.

2. Alignment with Uncertainty and Planning

Disagreement spikes are not random; they exhibit specific structural properties:

Uncertainty Alignment: Disagreement spikes strongly overlap with the base model's intrinsic uncertainty (Shannon entropy).
Planning Enrichment: These critical tokens are heavily enriched for planning-related decisions (e.g., strategy selection, variable binding, decomposition). In Qwen3-0.6B, planning tokens constitute only $\sim1.9\%$ of the global stream but appear in $\sim14\%$ of disagreement spikes (a $7.46\times$ enrichment).

3. Predictive Power

Aggregating disagreement scores over the top- $K$ tokens serves as a strong predictor of reasoning failure. The disagreement-based score outperforms base-model entropy in predicting incorrect final answers (e.g., AUROC of 0.851 vs. 0.817 on GSM8K).

4. Performance Recovery via Sparse Intervention

The proposed intervention scheme yields disproportionate performance recovery:

Efficiency: Intervening on a small budget of tokens (e.g., $\bar{\rho} \approx 0.13$ , or 13% of tokens) allows a 0.6B base model to recover 157% of the Pass@8 gap relative to a same-size (0.6B) "Thinking" variant, effectively surpassing it.
Comparison to Baselines: Random token replacement or replacing only the first $\rho$ tokens (early-only) yields significantly lower gains, confirming that the selection of specific high-disagreement tokens is the critical factor, not merely the volume of intervention.
Generalization: The phenomenon transfers across model families (Qwen3 and LLaMA) and domains (math and scientific QA), though the magnitude of recovery varies.

Significance and Claims

The paper argues for a sparse control view of reasoning. It posits that base models possess latent reasoning machinery that is often misaligned at specific, high-uncertainty planning points. The "reasoning mode" of LRMs is not necessarily a fundamentally different architecture but rather a stabilization of these critical early decisions.

The authors claim that:

Reasoning is sparse: A small number of early planning commitments steer the entire reasoning trajectory.
Correction is sufficient: By selectively correcting only these high-disagreement decision tokens, a base model can recover the reasoning capabilities of a much stronger model without full post-training or continuous LRM inference.
Mechanistic Insight: The intervention acts as a lightweight "stop-and-check" mechanism, resolving ambiguities in variable binding or strategy selection before handing execution back to the base model.

The work suggests that the "reasoning gap" is not a lack of capability in the base model, but a failure to commit to the correct trajectory at specific branching points, which can be remedied through targeted, disagreement-guided intervention.