Technical Summary: ReProbe – Efficient Test-Time Scaling of Multi-Step Reasoning by Probing Internal States of Large Language Models

1. Problem Statement

Large Language Models (LLMs) increasingly solve complex tasks using Chain-of-Thought (CoT) reasoning. However, a single flawed reasoning step can derail the entire solution. Test-Time Scaling (TTS) strategies, such as Best-of-N sampling and Beam Search, attempt to mitigate this by generating multiple reasoning trajectories and selecting the best one.

Current TTS relies heavily on Process Reward Models (PRMs) to verify the correctness of intermediate reasoning steps. PRMs face three critical limitations:

1. Computational Cost: They are large models (typically 1.5B–8B parameters) requiring significant GPU memory and inference time, often adding substantial overhead.
2. Data Scarcity & Cost: Training PRMs requires expensive step-level annotations, often derived from Monte-Carlo rollouts, human labeling, or consensus filtering, which are difficult to scale.
3. Limited Generalization: PRMs are often fine-tuned on narrow domains (e.g., mathematics) and struggle to generalize to out-of-domain (OOD) tasks like planning or general knowledge QA.

The paper asks: Can we achieve the performance of PRMs with a lightweight, efficient, and generalizable alternative that does not rely on external textual evaluation?

2. Methodology: ReProbe

The authors propose ReProbe, a lightweight, training-based reasoning probe that verifies reasoning steps by probing the internal states of the frozen target LLM rather than analyzing the generated text.

Core Architecture

ReProbe is a plug-and-play module consisting of:

1. Feature Extraction: It extracts internal signals from the frozen LLM at each token position. The authors investigate two feature sets:
   • Attn+Logit: Attention weights over the preceding 5 tokens and logits of the top-K candidate generations.
   • Hidden States: The hidden representations from all layers of the LLM.
2. Probe Network: A lightweight Transformer-based classifier (fewer than 10M parameters) that processes these features.
   • A linear projection layer adjusts feature dimensions.
   • A stack of Transformer layers captures contextual dependencies.
   • An aggregation layer pools token-level features into a step-level vector.
   • A two-layer classification network outputs a probability score indicating the likelihood of the step being correct.

Training Strategy

• Data Annotation: ReProbe can be trained using two annotation strategies:
  1. External Supervision: A larger, more capable LLM (e.g., DeepSeek-R1) acts as a judge to label the correctness of steps generated by the target LLM.
  2. Self-Supervision: The target LLM annotates its own generated steps.
• Efficiency: The training process is highly efficient. The underlying LLM remains frozen; only the probe parameters are updated. The authors report a total annotation cost of ~$200 and training compute of just 4 GH200 GPU-hours.
• Step Extraction: For non-thinking models, steps are extracted via pattern matching on newlines. For native thinking models (using <thought> tags), each sentence is treated as a step.

3. Key Contributions

1. Lightweight Verification: ReProbe achieves step-level verification performance comparable to or exceeding PRMs that are 150x to 810x larger (e.g., 10M parameters vs. 7B parameters).
2. Superior Generalization: Unlike domain-specific PRMs, ReProbe demonstrates strong Out-of-Domain (OOD) performance in planning and general knowledge QA, where large PRMs often fail.
3. Training Data Efficiency: ReProbe can be trained on automatically annotated data (self-supervised or via a stronger LLM) without the need for costly human annotations or complex Monte-Carlo rollouts.
4. Synergy with PRMs: The paper demonstrates that ReProbe and PRMs capture complementary signals (internal confidence vs. external knowledge). Combining them in a cascaded verification strategy further improves performance while reducing the computational cost of running heavy PRMs.

4. Experimental Results

The authors evaluated ReProbe on three domains: Mathematics (In-Domain), Planning (OOD), and General Knowledge QA (OOD), using models like Qwen3 and Phi-4.

• Step-Level Verification (PR-AUC):
  • ReProbe variants (especially those using Hidden States) consistently outperformed unsupervised Uncertainty Quantification (UQ) methods.
  • On OOD tasks (Planning, QA), ReProbe significantly surpassed the best large PRMs (e.g., Qwen2.5-Math-7B-PRM).
  • On In-Domain Math tasks, ReProbe matched or slightly trailed the strongest specialized PRMs but remained highly competitive with a fraction of the parameters.
• Test-Time Scaling (TTS):
  • Best-of-N: ReProbe-guided selection allowed Qwen3-8B to outperform its larger sibling (Qwen3-14B) on multiple benchmarks.
  • Beam Search: ReProbe achieved the highest decoding accuracy across all datasets, outperforming state-of-the-art PRMs.
• Efficiency: ReProbe inference is 2.6x to 25x faster than comparable PRMs, offering a massive reduction in inference latency and memory usage.

5. Significance and Conclusion

ReProbe represents a paradigm shift in test-time scaling by demonstrating that LLM internal states encode rich confidence signals about their own reasoning processes.

• Scalability: It offers a path to scalable, generalizable TTS that is not bottlenecked by the cost of training massive reward models.
• Introspection: It enables "introspective" LLMs that can self-monitor their reasoning quality without external critics.
• Practicality: By reducing the parameter count from billions to millions and eliminating the need for expensive human labeling, ReProbe makes advanced reasoning verification accessible for a wider range of applications and resource-constrained environments.

The work suggests that future reasoning systems may move away from heavy external reward models toward efficient, internal probes, potentially combined with external knowledge for a hybrid, highly efficient verification architecture.