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

AI summary

Overview

• Research area: Test-time scaling (TTS) and step-level verification for multi-step LLM reasoning; a training-based alternative to Process Reward Models (PRMs) built on probing internal model signals.
• Technical level: Intermediate — the paper is readable without deep theory, but assumes familiarity with chain-of-thought prompting, best-of-N sampling, beam search, and reward modeling.
• Scope: One paper proposing a lightweight probe (fewer than 10M parameters) trained on a frozen LLM's internal states to score the correctness of individual reasoning steps, evaluated across mathematics, planning, and general knowledge QA.

What This Paper Is About

Large language models solve hard problems by writing long chains of reasoning, and a single bad intermediate step can ruin the final answer. Existing fix — Process Reward Models, which score each step — are expensive to annotate, expensive to run at inference, and often overfit to narrow domains like math. The authors ask whether a much smaller, cheaper model that reads the LLM's own internal states (hidden states, attention weights, logits) can verify reasoning steps just as well.

Key Contributions

1. A probing-based step verifier (ReProbe). A transformer-based classifier that sits on top of a frozen LLM and predicts how likely each generated reasoning step is to be correct, using internal signals already produced during generation rather than the text alone.
2. Demonstrated parameter efficiency. The probes contain fewer than 10M parameters — the paper reports they outperform PRMs up to 150× larger and remain competitive with PRMs up to 810× larger, where typical PRMs contain 1.5B–8B parameters.
3. Cost-efficient annotation alternatives. Training labels can come either from a larger external LLM (DeepSeek-R1) or self-supervised from the target model itself, avoiding the Monte-Carlo rollouts, human annotation, and proprietary data pipelines that PRMs typically require.
4. Strong out-of-domain generalization and complementarity with PRMs. ReProbes lead on OOD planning and QA step verification, and combining PRM scores with ReProbe scores via logistic regression improves results further, suggesting the two capture different aspects of reasoning quality.

Main Findings

• ReProbes match or beat far larger PRMs on step-level verification. For Qwen3-8B, the best ReProbe (Hidden States, Self-anno) reaches an overall PR-AUC of .604, versus .565 for Qwen2.5-Math-PRM-7B (860K samples) and .559 for Qwen2.5-Math-7B-PRM800k (265K samples). ReProbe uses 32K training samples.
• The advantage is largest out-of-domain. ReProbes score .667 average PR-AUC on the combined OOD tasks (planning and QA) versus .595 for Qwen2.5-Math-PRM-7B and .594 for Qwen2.5-Math-7B-PRM800k. On in-domain math, the strongest PRMs remain slightly ahead (.529–.565 for large PRMs versus .498 for the best ReProbe on ID average).
• Unsupervised uncertainty quantification underperforms. Every UQ baseline tested (MaxProb, MaxEntropy, Perplexity, Self-Certainty, CCP, P(True), Semantic Entropy, Lexical Similarity, Degree Matrix, Random) sits far below both small PRMs and ReProbes on PR-AUC, so UQ was excluded from the test-time scaling experiments.
• Beam search: ReProbe gives the strongest accuracy. With Qwen3-8B and B=5 beams, N=5 continuations, ReProbe (Hidden States, DeepSeek-anno) reaches 76.6 overall accuracy versus 71.6 for Qwen2.5-Math-7B-PRM800k and 70.7 for Qwen2.5-Math-PRM-7B.
• Best-of-N: ReProbes enable a smaller model to beat a bigger one. Qwen3-8B has a pass@1 lower bound of 59.8 and a pass@N upper bound of 72.1; Qwen3-14B pass@1 is 67.3. ReProbe-based BoN pushes Qwen3-8B to 62.8 (Attn+Logit, DeepSeek-anno), outperforming the larger Qwen3-14B on MATH, GSM8K, ProofNet, Meeting Planning, and ScienceQA. ReProbes rank best on all datasets except MATH, where they rank second.
• Self-supervised annotation is competitive with external annotation. The self-annotated probe reaches average performance comparable to the DeepSeek-R1-annotated variant, which the authors highlight as an efficient labeling path, particularly in OOD settings.
• Works on native thinking-mode models too. When reasoning is emitted in the model's native <think> format — with each sentence treated as a step and annotated by GPT-OSS-120B — ReProbes stay competitive with or ahead of PRMs. For Qwen3-32B, ReProbe reaches .585 overall PR-AUC versus .613 for Qwen2.5-Math-PRM-7B and .569 for Universal-PRM-Qwen2.5-Math-7B; on OOD it leads with .558.
• More data and more diverse data both help. Training ReProbes on 2K questions with 6K samples, the diverse subset (chosen by farthest-first traversal over Qwen3-Embedding-8B vectors) yields .507 overall PR-AUC versus .461 for the homogeneous subset.
• PRMs and ReProbes are complementary. A logistic regression combining both signals improves on either alone: on MATH, PRM1 alone gives .586, ReProbe + PRM1 gives .613; on ProofNet, PRM2 alone gives .310, ReProbe + PRM2 gives .327.
• Step-level architecture beats token-level. Comparing step-level ReProbe, a token-level linear probe, and a token-level transformer probe under identical setups, the step-level design performs best.
• Runs faster at inference. The current ReProbe implementation achieves a 2.6×–25× speedup over state-of-the-art PRMs.

Methodology in Plain English

The authors take a target LLM, freeze all of its weights, and harvest signals that already exist while it generates text. Two feature sets are tried: (1) Attn+Logit — attention weights over the five preceding tokens from all layers, plus logits over the top-K candidate next tokens; and (2) Hidden states — the layer-wise internal representations of each token.

These per-token features pass through a small pipeline: a linear projection layer, a stack of transformer layers to model context, a mean-pooling step that collapses token features into one vector per reasoning step, and a two-layer classification head with dropout and GeLU that outputs a correctness logit. Training uses binary cross-entropy with class weighting (to handle label imbalance), and only probe parameters update.

Data comes from 10.8K problems in the PRM800K training subset. Three reasoning trajectories are sampled per problem using nucleus sampling (top-k = 50, top-p = 0.95, temperature = 1.0), giving roughly 32K samples. Generation length is restricted during training for efficiency, and that restriction is removed at evaluation so results are not length-biased. Labels come from an LLM judge given the question, the chain-of-thought steps, the final answer, and the ground-truth answer — either an external model (DeepSeek-R1) or the target model judging itself. Step boundaries are found by pattern matching when the model writes one step per line, and by sentence splitting when models use native thinking mode.

The judge's reliability was checked against human annotations from a random PRM800K subset and a manually annotated set of 1,000 steps spanning QA, planning, and ProofNet: DeepSeek-R1 hit 95% accuracy on PRM800k and about 90% on the other datasets.

Evaluation covers three settings: step-level correctness via PR-AUC; best-of-N with N=10 for math and QA and N=5 for planning at temperature 1.0; and beam search with B=5 beams and N=5 sampled continuations per beam. Models tested are Qwen3-8B (non-thinking CoT mode), Qwen3-1.7B and Qwen3-32B (native thinking mode), and Phi-4. Baselines include two 1.5B PRMs, six 7–8B PRMs, and unsupervised UQ methods from the LM-Polygraph framework.

Why This Matters

Verification is the bottleneck in making test-time scaling work: if scoring candidate reasoning steps costs as much as generating them, the efficiency gains evaporate. ReProbe shows a scorer can be small, cheap, and trained without heavy human labeling or Monte-Carlo rollouts, and that internal model states carry real signal about the model's own confidence. The finding that PRM scores and probe scores are complementary is also notable — it suggests introspective confidence and external process knowledge are different, stackable things.

Real-world applications:

• Cost-sensitive reasoning deployments: replacing or augmenting 1.5B–8B PRM critics with fewer than 10M-parameter probes to cut GPU memory and inference-time latency (the paper reports a 2.6×–25× speedup).
• Domain-specific assistants: math tutoring, trip and meeting planning, calendar scheduling, and general knowledge QA — all domains actually evaluated here.
• Self-improving pipelines without human labels: self-supervised probes could let a deployed model flag its own unreliable reasoning steps without an annotation budget.
• Boosting smaller models on-device or under budget constraints: the results show an 8B model with ReProbe-based TTS beating a 14B model on several benchmarks, which matters when a larger model will not fit.

Industry relevance: the approach lowers the barrier to building reasoning systems with verification, since it avoids proprietary training data, costly human annotation, and running a second large model alongside the primary one.

Future Directions

• Scaling training data further. The authors note their data-scaling curves show diminishing but not saturated returns, especially on StrategyQA, and they deliberately stayed inside the PRM800K training set to keep comparisons to PRMs fair. Extending to new, diverse, high-quality questions — including outside math — is left open. They also cite budget limits: annotating 32K reasoning trajectories with DeepSeek-R1 cost more than $2000 USD.
• Cross-model transfer. ReProbes depend on the target LLM's specific internal states, so a probe trained for one model cannot be applied directly to another. The authors suggest sharing probes online so practitioners can download them, and suggest fine-tuning probes to adapt to customized or fine-tuned versions of target LLMs.
• Hybrid verifiers. Since combining PRM and ReProbe scores improved step-level PR-AUC, the paper calls for verifiers that merge introspective confidence with external process rewards.
• Longer reasoning chains. The authors report that both PRMs and ReProbes degrade slightly as reasoning length grows, which invites work on maintaining probe reliability over very long chains.

Target Audience

Researchers and engineers working on LLM reasoning, inference-time compute, and reward modeling; practitioners building retrieval-free verification or self-checking into deployed models; and anyone interested in model interpretability, since the work is direct evidence that internal states encode usable confidence signals about a model's own reasoning. The paper is most useful to readers who already know what best-of-N sampling, beam search, and PRMs are, but the core idea and results are accessible to a motivated general reader. All trained probes, code, prompts, human annotations, and processed datasets are open-sourced, with random seeds fixed to 1 or 42.

Authors’ abstract

LLMs can solve complex tasks by generating long, multi-step reasoning chains. Test-time scaling (TTS) can further improve performance by sampling multiple variants of intermediate reasoning steps, verifying their correctness, and selecting the best steps for continuation. However, existing verification approaches, such as Process Reward Models (PRMs), are computationally expensive and require large-scale human or model-generated annotations. We propose a lightweight alternative for step-level reasoning verification based on probing the internal states of LLMs. We train a transformer-based probe that uses the internal states of a frozen LLM to estimate the credibility of its reasoning steps during generation. Annotation can be provided either by a larger LLM (e.g., DeepSeek-R1) or in a self-supervised manner by the original model itself. The probes are lightweight, containing fewer than 10M parameters. Across multiple domains, including mathematics, planning, and general knowledge question answering, our probes match or exceed the performance of PRMs that are up to 810x larger. These results suggest that LLM internal states encode confidence in their reasoning processes and can serve as reliable signals for step verification, offering a promising path toward scalable, generalizable TTS and more introspective LLMs.