[ICLR 2025] One-Token Verification: Probing the KV Cache for Efficient System-2 Reasoning

TL;DR

One-Token Verification (OTV) transforms a frozen LLM into its own verifier by inserting a special [ToT] token that "interrogates" the model's Key-Value cache. Unlike traditional verifiers that read text, OTV looks at the model's internal "memory," allowing it to predict if a reasoning step is correct in just one forward pass. This method slashes token costs by up to 90% via early stopping and consistently beats external reward models on elite math competitions like AIME.

The "Thinking" Paradox: Why Parallel Sampling is Broken

Today's top-tier reasoning models (like o1 or R1) rely on parallel thinking—generating multiple chains-of-thought and picking the best one. However, this creates two massive bottlenecks:

1. The Latency Trap: You usually have to wait for the entire 5,000-token trace to finish before a verifier can tell you it was wrong at token 50.
2. The Calibration Crisis: Models struggle to quantify their own uncertainty through logits alone, leading to "overthinking" on simple tasks or misplaced confidence in complex errors.

Prior attempts to solve this involve external Process Reward Models (PRMs). But these are expensive, separate models that don't "know" the internal state of the generator.

Methodology: Probing the Internal "Sufficiency Statistic"

The core insight of OTV is that the KV Cache is a sufficient statistic of the entire prefix. While the final hidden state of a token is a lossy summary, the KV cache contains the structured, layer-wise history of the reasoning process.

The [ToT] Mechanism

OTV introduces the Token of Truth ([ToT]). Mechanistically, the process works as follows:

1. Gated LoRA: A LoRA module is added to the LLM but remains inactive ($m_t=0$) during normal generation to avoid degrading the base model's logic.
2. Probing Pass: When verification is needed, the [ToT] token is appended. It activates the LoRA pathway ($m_t=1$), allowing it to perform cross-attention over the existing KV cache.
3. Scalar Regression: Instead of predicting the next word, a 3-layer MLP maps the [ToT] hidden state directly to a correctness score $c_t\in [0,1]$.

Parallelized Training

One of the most elegant parts of OTV is its training efficiency. By using a triangular probe mask, the authors can insert [ToT] tokens at every position in a training trace and calculate all verification scores in a single parallel batch, mimicking the standard shift-right training of Transformers.

Experimental Victories: Accuracy Meets Efficiency

1. Crushing the Baselines

On the grueling AIME (American Invitational Mathematics Examination) benchmarks, OTV significantly outperformed established external verifiers like AceMath-RM-7B and Qwen2.5-PRM.

2. The 90% Efficiency Gain

The real "killer app" for OTV is Early Termination. By using strategies like Halve@300 (halving the candidate pool every 300 tokens based on OTV scores), the system can prune "dead-end" reasoning paths early. The results show OTV achieves higher accuracy with 90% fewer tokens compared to full-length Best-of-N decoding.

Above: OTV shows clear stratification between correct (red) and incorrect (green) traces far earlier than previous methods.

Deep Insights: The "Shortened Trace" Preference

An interesting byproduct of OTV’s Linear Ramp labeling (where confidence is forced to grow from 0.5 to 1.0 linearly over the trace length) is that it naturally learns an inverse-length preference. Mathematical analysis in the paper proves that if two traces are both correct, OTV will assign a higher score to the shorter, more concise one. This effectively punishes "overthinking" and encourages the model to find the most efficient path to the solution.

Critical Analysis & Future Outlook

While OTV is a major step forward, its reliance on outcome-level pseudo-labels (assigning scores based on whether the final answer was right) is a double-edged sword. It requires no human process-supervision, but it might misinterpret a "correct step in a ultimately wrong trace" as a bad step.

Future Directions:

• Hybrid Supervision: Combining OTV's efficiency with sparse human process feedback.
• Ternary Logic: Expanding the output to "Correct/Incorrect/Unknown" to allow the model to specifically flag where it is "confused."
• Beyond Math: Applying OTV to code generation and multi-agent planning where the "state space" is even more complex.

OTV proves that the "internal voice" of an LLM knows more than it lets on—we just needed the right token to listen.