Source-linked AI summary

HSRM: Hidden-State Reward Models for Test-Time Verification

Xianzhi Li, Xiaodan Zhu

arXiv:2608.30841v1cs.AIcs.CL

TL;DR

Selecting correct solutions from multiple mathematical reasoning candidates is difficult and text-based verification adds substantial inference cost. HSRM instead ranks candidates from hidden states computed by a frozen generator, achieving comparable or better results than a much larger text-only verifier across most evaluated settings. Its lightweight design reuses generation-time representations, while evaluation remains focused on mathematical reasoning and has limited coverage of broader architectures and explicit deliberation.

  • Problem

    Selecting reliable solutions among multiple generated mathematical reasoning candidates remains difficult, while text-based verification can be costly.

  • Method

    HSRM uses a small Transformer to rank candidates from frozen-generator hidden states extracted at reasoning-step boundaries and trains on self-generated trajectories with correctness labels.

  • Results

    HSRM matches or outperforms a 55M-parameter text-only verifier in 15 of 16 generator–dataset settings while using about 2M parameters.

  • Takeaways & Limitations

    HSRM provides a low-cost complement to verifier-based test-time reasoning by reusing representations already computed during generation.

  • Takeaways & Limitations

    The study focuses on mathematical reasoning, and its explicit-deliberation evaluation is limited to two model scales on MATH-500.

Abstract

from arXiv · show

Large language models can often generate plausible mathematical reasoning traces, but reliably identifying the correct solution among multiple candidates remains a key challenge. Existing test-time reasoning pipelines typically rely on text-based verifiers that re-read each generated solution, making verification an expensive component of inference. Prior work has shown, however, that LLMs often encode correctness-related signals in their internal representations, including awareness of when their own answers are likely to be wrong. Building on this observation, we introduce HSRM, a lightweight hidden-state reward model that verifies candidate solutions by directly reading the generator's internal representations rather than re-processing its text. HSRM extracts hidden states from a frozen generator at reasoning-step boundaries and uses a small Transformer encoder to rank candidates. It is trained from self-generated trajectories with outcome labels, requiring neither human-written process supervision nor a large pretrained verifier. Across four mathematical reasoning benchmarks, HSRM matches or outperforms a 55M-parameter text-only energy verifier in 15 of 16 generator--dataset settings while using only about 2M parameters, providing an efficient alternative to text-only verification by reusing representations already computed during generation.

1 Introduction

HSRM addresses the cost and reliability challenges of selecting correct mathematical solutions by verifying candidates from generator hidden states rather than re-processing their text. It uses a small self-trained Transformer and performs competitively with much larger text-only verifiers.

  • Test-time sampling improves reasoning accuracy by selecting among multiple candidate solutions, but its effectiveness depends on verifier quality and cost.
  • Existing reasoning verifiers usually read generated text and can make verification a major component of inference cost.Outcome and process reward models score complete solutions or intermediate steps, with recent process reward models often implemented as large language models.
  • Generator hidden states may contain correctness information that enables verification without relying only on emitted text.Prior work reports internal signals related to models’ own answer correctness, truthfulness, factuality, and related properties.
  • HSRM scores candidates by extracting frozen-generator hidden states at reasoning-step boundaries and processing them with a small Transformer.It trains on the generator’s sampled trajectories using final-answer correctness labels, without human-written process supervision or a pre-existing verifier corpus.
  • 15 of 16 generator–dataset settings match or outperform a 55M-parameter text-only verifier while HSRM uses about 2M parameters.
  • Ablations indicate that HSRM’s advantage comes primarily from its hidden-state input representation rather than increased verifier size.

2 Related Work

Prior work primarily verifies reasoning from generated text, while newer approaches investigate internal representations as lower-cost signals for correctness. HSRM belongs to this emerging hidden-state verification line and reuses representations computed during generation.

  • Best-of-N reasoning samples multiple solutions and relies on a verifier to rank or aggregate them.
  • Text-based outcome and process reward models score complete solutions or intermediate steps, but large language model verifiers make verification costly.
  • Prior studies show that internal representations can encode information about a model’s knowledge, answer correctness, truthfulness, and factuality.
  • Hidden-state verification methods use probes, activation features, or lightweight reward heads for error detection, early exiting, training-free verification, or candidate selection.
  • HSRM avoids re-tokenizing and re-embedding sampled solutions by reading hidden states already produced during decoding.

3 Our Method

HSRM verifies best-of-N candidates by processing generator hidden states at reasoning-step boundaries instead of re-encoding solution text. A compact Transformer ranks candidates using cached representations and a tie-safe objective trained from correctness-labeled sampled trajectories.

  • HSRM Overview: HSRM reads per-step hidden states from a frozen generator and assigns each candidate a scalar score for best-of-N selection.The generator samples N candidates, and selection returns the candidate with the highest HSRM score.
  • Step-Boundary Hidden-State Extraction: Step-boundary extraction reduces verifier input length and requires no additional generator forward passes because the representations are already computed during generation.Only hidden states immediately before reasoning-step delimiters, including the final generated token, are collected.
  • HSRM Architecture: HSRM maps a step-level hidden-state sequence to a candidate score through projection, Transformer encoding, mean pooling, layer normalization, and linear readout.The default encoder has 2 layers, width 256, 4 attention heads, and roughly 2M parameters.
  • Training Objective: The tie-safe ranking loss requires every correct candidate to outrank incorrect candidates while treating correct candidates as ties.Problems containing no correct or no incorrect candidates are omitted because they provide no within-problem ranking signal.
  • Training and Inference: HSRM trains on cached hidden-state tensors, so verifier optimization proceeds without further generator calls.At inference, candidates and step-boundary states are cached and scored in a batched verifier pass.

4 Experimental Setup

The experiments evaluate HSRM with frozen Qwen3 generators across four mathematical benchmarks and compare it with text-based, external process-reward, and internal heuristic baselines. Evaluation uses best-of-8 selection, with verifier training based on sampled candidate pools and correctness labels.

  • Models and Benchmarks: Experiments use frozen Qwen3 generators at 1.7B, 4B, 8B, and 14B parameters across four mathematical reasoning benchmarks.The benchmarks are GSM8K, MATH-500, AIME, and OlympiadBench; thinking-mode evaluation is additional.
  • Training and Evaluation Protocol: Verifier training samples 64 candidate solutions per training problem, while reported in-distribution results average over five seeds.Zero-shot transfer trains on a source dataset and evaluates directly on a target dataset without target-domain verifier training.
  • Labeling: Candidate correctness labels are assigned by extracting final answers, normalizing formatting and equivalent expressions, and using an LLM judge for ambiguous cases.Reported results use these post-processed correctness labels unless otherwise stated.
  • Baselines: The primary text-only baseline is EORM 55M, an energy-based Transformer verifier trained with a Bradley–Terry ranking objective.EORM reads the full chain-of-thought text and scores candidates.
  • Baselines: The evaluation also compares HSRM with Qwen2.5-Math-PRM-7B and non-learned confidence, uncertainty, length, and reasoning-step heuristics.The 7B process reward model is treated as a strong external baseline rather than a capacity-matched comparison.
  • Metrics: At evaluation, verifier-best accuracy is measured at N = 8 by selecting the highest-scoring candidate among eight sampled solutions.Within-problem AUROC and oracle pass@N are also reported.

5 Results

Across generator scales and benchmark difficulty, HSRM generally outperforms the 55M-parameter text-only verifier while using about 2M parameters. It nearly matches the larger domain-trained PRM on GSM8K, but that PRM remains stronger on the hardest benchmarks.

  • GSM8K: HSRM outperforms EORM at every GSM8K generator scale and nearly matches Qwen2.5-Math-PRM-7B despite using about 2M parameters.As generator scale increases, HSRM approaches the oracle curve in the best-of-8 setting.
  • Harder Benchmarks: On MATH-500, AIME, and OlympiadBench, HSRM generally remains above EORM and improves with generator scale.Qwen2.5-Math-PRM-7B retains an advantage on these harder benchmarks, especially AIME.
  • Harder Benchmarks: HSRM closes a substantial portion of the gap to the larger domain-trained PRM while avoiding candidate-text re-encoding.The efficiency comes from using generator hidden states rather than reprocessing solution text.
  • Overall Findings: Figure 2 supports compact hidden-state verification as an alternative that can outperform a substantially larger text-only verifier across generator scales and benchmark difficulty.The paper identifies verifier efficiency as HSRM’s main attraction, while large domain-matched PRMs remain stronger on the hardest benchmarks.

6 Ablations

Ablations show that HSRM’s verification signal is distributed across generator layers and reasoning steps, while hidden-state inputs provide strong accuracy and ranking at substantially lower verification cost. Transfer and scaling experiments further show robust performance across data sizes, domains, and generator families.

  • Generator Layer: Correctness information appears across many generator layers, with strongest HSRM performance generally in the upper layers rather than at one uniquely optimal layer.The final layer remains competitive, but nearby upper layers can match or exceed it in several settings.
  • Input Modality: 86.11% best-of-8 accuracy and 0.691 within-problem AUROC from the simplest hidden-state verifier exceed the text-only verifier’s 82.23% and 0.519.This comparison attributes the improvement to the hidden-state representation rather than merely to verifier architecture.
  • Input Modality: Top-4-layer HSRM reaches 86.86% best-of-8 accuracy and 0.724 within-problem AUROC, outperforming the final-layer hidden-state variant.The result supports combining information across multiple upper layers and across the reasoning trajectory.
  • Efficiency Analysis: HSRM uses roughly five orders of magnitude fewer verification FLOPs per candidate than Qwen2.5-Math-PRM-7B while matching it at larger generator scales.It reads cached generator hidden states instead of re-encoding candidate text, placing it on the upper-left efficiency frontier.
  • Training Data: On GSM8K, increasing the verifier training pool from 25 to 500 problems raises best-of-8 accuracy from 83.4 to 86.2 and AUROC from 0.612 to 0.716.MATH-500 shows the same direction as K grows, although its curve is noisier.
  • Zero-Shot Transfer: HSRM outperforms EORM 55M across OlympiadBench transfer settings and remains close to Qwen2.5-Math-PRM-7B without target-domain training.The remaining performance gap is primarily to the oracle, indicating substantial candidate-generation and selection headroom.
  • Thinking-Mode Generators: Restricting thinking-mode extraction to post-deliberation representations improves AUROC from 0.669 to 0.736 for Qwen3-1.7B and from 0.672 to 0.884 for Qwen3-4B.This indicates that verification-relevant information is not equally stable throughout the full deliberation trace.
  • Cross-Family Evaluation: HSRM outperforms matched text-based EORM in all six Llama generator–dataset settings, with gains of 3.1–8.1 points on GSM8K and 3.6–4.5 points on MATH-500.The cross-family results suggest the exploited verification signal is not tied to one generator architecture or pretraining recipe.

7 Conclusion

HSRM verifies mathematical reasoning candidates from hidden states already produced by a frozen generator, avoiding separate text re-encoding. Across four benchmarks and Qwen and Llama generators, it matches or outperforms a much larger text-only verifier with about 2M parameters.

  • Conclusion: HSRM ranks mathematical reasoning candidates using hidden states from a frozen generator during decoding instead of re-encoding their text with a separate reward model.The method extracts representations at reasoning-step boundaries.
  • Conclusion: Across four benchmarks and Qwen and Llama generators, HSRM matches or outperforms a substantially larger text-only verifier while using about 2M parameters.The paper presents it as a low-cost complement to verifier-based test-time reasoning pipelines.

Limitations

The study is limited to mathematical reasoning, and its explicit-deliberation evaluation covers only two model scales on MATH-500. Broader domains, architectures, and deliberation settings remain for future evaluation.

  • Scope: The study focuses on mathematical reasoning, while broader domains and model architectures remain unevaluated.The controlled setting isolates generator hidden-state verification signals but limits scope.
  • Scope: The thinking-mode study is limited to two model scales on MATH-500, leaving broader explicit-deliberation evaluation as future work.

A Dataset Details and Split Indices

The dataset pipeline uses disjoint problem-level splits, sampled candidate trajectories, and binary final-answer correctness labels. A second-stage judge relabels unresolved candidates, changing about 3.1% of labels before training and evaluation.

  • Dataset and Splits: Training and evaluation splits are disjoint at the problem level, with N=64 candidates per training problem and N=8 per evaluation problem.The post-processed candidate corpus contains approximately 255K sampled solutions.
  • Labeling: Each candidate receives a binary final-answer correctness label through deterministic extraction and equivalence checks, followed by judge review for unresolved cases.The judge handles extraction failures and difficult formats such as fractions, square roots, units, and symbolic expressions.
  • Labeling: 7,788 labels changed during judge relabeling, including 7,650 incorrect-to-correct and 138 correct-to-incorrect changes.This produced a net increase of 7,512 correct labels.
  • Labeling: All reported results use post-relabel correctness labels rather than labels determined solely by answer-extraction rules.

C Architecture and Hyperparameter Details

HSRM uses a shared lightweight verifier design across generator scales, with a small Transformer operating on projected hidden-state representations. The experiments compare it with text-based verifiers, external process reward models, and simple generator-side scoring heuristics.

  • Architecture: HSRM’s step-level Transformer encoder and final score head are shared across generator scales; only the input projection depends on generator size.The projection is Win ∈ R^(dmodel×dgen).
  • Architecture: The default HSRM verifier uses dmodel=256, two Transformer layers, four attention heads, and dropout 0.1.
  • Data and training: The training recipe uses the same hyperparameters across datasets and generator scales unless otherwise stated.
  • Baselines: The evaluation compares HSRM with text-based learned verifiers, non-learned generator-side scorers, and an off-the-shelf Qwen2.5-Math-PRM-7B process reward model.The primary EORM comparison uses a 53.4M-parameter variant, while the smaller EORM variant has 2.6M parameters.
  • Data and training: MATH-500 uses 150 problems for verifier training and 350 held-out problems for headline results.
  • Baselines: Non-learned baselines include token log-probability, entropy, varentropy, response length, and reasoning-step counts.These scorers test whether confidence, uncertainty, or length signals suffice for best-of-N selection.
  • Baselines: A single fixed cheap scorer would generally perform worse than the table’s per-cell best-scoring envelope.

D Extended Results

Extended results show that HSRM remains competitive as candidate pools and baseline choices vary. It generally outperforms cheap heuristics, although gains narrow in low-accuracy or saturated settings.

  • N-scaling: Increasing the candidate pool beyond N=16 yields only small additional gains in many cells.The sampled pools appear close to their oracle ceiling in those settings.
  • Cheap scorers: 12 of 16 settings favor HSRM over the best cheap scorer, with 2 ties and 2 HSRM underperformances.The comparison gives cheap heuristics the favorable advantage of selecting the best scorer separately for each cell.
  • Cheap scorers: HSRM improves over the best heuristic at every generator scale on GSM8K and OlympiadBench.
  • Cheap scorers: No single cheap heuristic dominates across datasets and generator scales; selected scorers vary among log-probability, entropy, and length signals.This instability indicates that cheap scorers exploit dataset- and generator-specific artifacts rather than a robust verification signal.
  • Cheap scorers: HSRM uses the same learned hidden-state scoring mechanism across all settings.
  • Limitations: The main exceptions occur on AIME and a few saturated MATH-500 cells, where simple heuristics can sometimes match or exceed HSRM.AIME has low accuracy and small absolute differences, while saturated MATH-500 pools often contain strong solutions.

E Extended Ablations

Ablations examine the training objective, encoder architecture, and verifier cost. The results favor tie-safe pairwise ranking, show limited dependence on encoder capacity, and highlight HSRM’s reuse of generator representations.

  • Training objective: Tie-safe Bradley–Terry ranking is better aligned with best-of-N selection than BCE or ListMLE when multiple candidates share the same correctness label.It ranks correct candidates above incorrect ones without imposing an arbitrary total ordering among candidates with identical labels.
  • Training objective: Bradley–Terry achieves the strongest best-of-8 selection accuracy and within-problem AUROC on both GSM8K and MATH-500.On GSM8K, it reaches 85.93% best-of-8 accuracy and 0.718 within-problem AUROC.
  • Encoder architecture: DeepSet, Transformer d=128, and the default Transformer d=256 achieve comparable accuracy, with no architecture dominating every generator–dataset pair.
  • Encoder architecture: The ablation suggests that much of the verification signal is already present in step-boundary hidden states rather than created by a high-capacity encoder.
  • Encoder architecture: The d=256 Transformer is used by default because it is consistently competitive and offers a strong accuracy–AUROC trade-off.The ablation indicates that HSRM’s gains are not simply due to a larger encoder.
  • Verifier cost and efficiency: HSRM requires no additional generator forward passes after sampling because it reads hidden states already produced during generation.Text-based verifiers must separately re-encode each generated solution as text.
  • Verifier cost and efficiency: HSRM processes at most S ≤ 100 step vectors, whereas text-based verifiers process the full response length T with a separate encoder.For Qwen2.5-Math-PRM-7B, verification is comparable to an additional large-model pass over the full solution.
Loading 2608.30841v1…