Hidden States as Early Signals: Step-level Trace Evaluation and Pruning for Efficient Test-Time Scaling

• Self-consistency generates multiple reasoning traces and selects answers by majority voting, but lengthy traces and parallel sampling impose substantial computational cost and latency.Source: §1 · p. 1
• Prior pruning methods rely on similarity or confidence signals that may not reliably reflect reasoning quality.Source: §1 · pp. 1–2
• GPU memory saturation from multiple long traces causes preemption, waiting queues, redundant computation, and end-to-end latency.Source: §1 · p. 2
• STEP trains a lightweight scorer on reasoning-step hidden states to identify unpromising traces early during generation.Source: §1 · p. 2
• STEP triggers pruning as KV-cache accumulation saturates GPU memory, releasing resources to prevent preemption and queuing delays.Source: §1 · p. 2
• 45%–70%: STEP reduces average end-to-end inference latency versus self-consistency while improving reasoning accuracy by +0.4 to +7.5 percentage points.The evaluation covers AIME-25, HMMT-24/25, GPQA-Diamond, EquiBench, and DivLogicEval across three models.Source: §1 · p. 2

• Confidence-based methods prune low-confidence traces, whereas diversity-based methods remove similar traces to preserve answer diversity.Source: §2.1 · p. 2
• Confidence can suffer from model overconfidence and miscalibration, while surface-level similarity may not indicate reasoning redundancy.Source: §2.1 · p. 2
• Prior studies show that hidden states can encode correctness-related information and support lightweight prediction of intermediate-answer correctness.Source: §2.2 · p. 2
• STEP trains a lightweight step scorer on hidden-state boundary representations for continuous monitoring and early pruning during parallel scaling.Source: §2.2 · p. 3

• Hidden states in reasoning models contain rich signals indicative of reasoning quality, and the main latency bottleneck arises from inference-system behavior.Source: §3 · p. 3
• A 2-layer MLP trained on hidden states finds that scores from early reasoning steps effectively distinguish correct from incorrect traces.Source: §3 · p. 3
• Incorrect traces tend to be longer than correct traces, so early termination can reduce generated tokens while improving efficiency and accuracy.Source: §3 · p. 3
• KV-cache accumulation saturates GPU memory and causes inference engines to preempt traces, creating waiting time beyond token-generation costs.Source: §3 · p. 3

Framework overview

• STEP evaluates every reasoning step with a hidden-state-based scorer and uses a KV-cache-aware monitor to decide when pruning should occur.The framework addresses trace selection and pruning timing as separate design questions.Source: §4 · p. 3, §4.1 · p. 4

Step representation

• The scorer uses the last-layer hidden state at each step-end token, while trace steps are formed by splitting reasoning content at double-newline delimiters.Each step representation incorporates contextual information from preceding reasoning steps.Source: §4.1 · p. 4

Step scorer training

• Trace-level supervision propagates each trace’s correctness label to all its steps, with weighted loss compensating for the imbalance caused by longer incorrect traces.The weighting uses α = K−/K+ as the ratio of negative to positive step instances.Source: §4.1 · p. 4
• The two-layer MLP maps each step hidden state to a correctness probability, and the scorer is trained with weighted binary cross-entropy.The sigmoid-based scorer has trainable parameters and is intended to capture quality signals from hidden states.Source: §4.1 · p. 4

Memory-aware pruning

• When GPU memory is full and the next decoding step cannot be scheduled, STEP removes the trace with the lowest average step score and releases its KV cache.This memory-triggered policy is designed to avoid waiting queues caused by preemption and repeated trace resumption.Source: §4.2 · pp. 4–5

Scoring and answer selection

• During parallel generation, STEP maintains a running average of step scores and uses final trace-level scores in weighted majority voting after traces finish or are pruned.The average score captures quality evolution across steps and is less sensitive to individual-step variance.Source: §4.3 · p. 5

Experiment Setup

• Experiments cover six challenging benchmarks and three reasoning models, comparing STEP with CoT, SC, Slim-SC, and DeepConf.Evaluation reports accuracy, average output token usage, and inference latency; HMMT-24/25 results are averaged across both benchmarks.Source: §5.1 · p. 5, Table 1 · p. 6

Main Results

• STEP achieves the highest accuracy on most benchmark-model combinations and improves HMMT-25 accuracy over SC by 5.0%, 3.3%, and 0.8% across the three models.The method also achieves competitive GPQA-Diamond accuracy across models.Source: §5.2 · p. 6
• 45%–70% average latency reduction versus SC accompanies comparable or superior accuracy against Slim-SC and DeepConf.On Phi-4-reasoning-plus with HMMT-24, STEP achieves 58.3% accuracy in 630 seconds versus 2405 seconds for SC, a 3.8× speedup.Source: §5.2 · p. 7

Latency Scaling

• STEP reaches higher accuracy at a given time budget in most settings, including 70% accuracy at roughly 40% of SC latency on Qwen3-4B-Thinking-2507 with HMMT-25.On DeepSeek-R1-0528-Qwen3-8B with AIME-25, STEP reaches 85% accuracy at approximately 40% of SC’s latency for comparable performance.Source: §5.3.1 · p. 7, Figure 4 · p. 7

Ranking Ability of Step Scorer

• The hidden-state step scorer outperforms token confidence for pairwise ranking, with ranking performance improving as more reasoning steps become available.Its early-stage ranking accuracy indicates that hidden states provide useful information about solution correctness before the final answer.Source: §5.2 · p. 7, §5.3.2 · p. 8

Profiling Acceleration by Pruning

• STEP’s acceleration combines fewer generated tokens with elimination of waiting time, while memory-aware pruning removes the waiting queue and yields the lowest end-to-end latency.Accuracy remains stable across GPU memory budgets at 70.1 ± 1.8%, including settings that trigger pruning earlier.Source: §5.3.4 · p. 8, §5.3.5 · p. 9

• STEP combines hidden-state-based trace evaluation with GPU memory-aware pruning for efficient parallel scaling.The method trains a lightweight step scorer and triggers pruning according to GPU memory utilization.Source: §6 · p. 9
• 45%-70% average end-to-end inference latency reduction compared to self-consistency accompanies +0.4 to +7.5 percentage-point reasoning-accuracy gains.Source: §6 · p. 9

• Pseudo-labeling propagates trace-level correctness to individual steps, assuming all steps within a trace share the same quality.This weak supervision is noisy and may not reflect the true quality of individual reasoning steps.Source: Limitations · p. 9
• Memory-triggered latency improvements are tightly coupled to serving infrastructure, so their magnitude may differ across inference engines and multi-GPU configurations.Source: Limitations · p. 9
• Evaluation focuses on tasks with fixed-form verifiable answers, while open-ended generation would require adapting answer aggregation beyond majority voting.Source: Limitations · p. 9

• The training process uses hyperparameters listed in Table 5.Source: §A.1 · p. 11, Table 5 · p. 11
• Input dimensions match the model hidden-state sizes: 2560, 4096, and 5120 for the three evaluated LLMs.Source: §A.1 · p. 11
• Training problems come from HMMT 2012–2023 February competitions in Algebra, Combinatorics, and Geometry.Source: §A.2 · p. 11
• For each problem, 64 sampled solutions are verified with numeric matching and SymPy-based symbolic equivalence.Source: §A.2 · p. 11
• The scorer training set contains 5,000 correct and 5,000 incorrect traces for each LLM.Source: §A.2 · p. 11

• Sampling parameters remain fixed across methods, with details listed in Table 6.Source: §B.1 · p. 11, Table 6 · p. 12
• All methods receive the prompt: “Please reason step by step, and put your final answer within \\boxed{}.”Source: §B.2 · p. 12
• Slim-SC uses Random Pruning, while DeepConf-low terminates traces below a threshold retaining the top 10% highest-confidence warmup traces.Source: §B.3 · p. 12
• STEP places the scorer with the model runner on the same GPU and sends scores to the engine core for scheduler pruning decisions.Source: C System Integration · p. 12
• The implementation follows vLLM-V1, with the engine core and model runner residing in separate processes.Source: C System Integration · p. 12

D Additional Computational Overhead

• The step scorer is an auxiliary MLP invoked at every reasoning step, so its computation is measured against one LLM generation step.The comparison quantifies the scorer’s per-step overhead rather than treating it as a one-time cost.Source: D Additional Computational Overhead · p. 12
• The LLM’s approximate forward-generation cost is 2N, while the step-level MLP costs 2m(d + 1).Here, N is the number of non-embedding parameters, m is the MLP hidden dimension, and d is the LLM hidden dimension.Source: D Additional Computational Overhead · p. 12
• With m = 512, d on the order of 10^3, N on the order of billions, and t around 10^2, the relative overhead is below 10^-6.The ratio accounts for the average tokens per reasoning step through t.Source: D Additional Computational Overhead · p. 12
• The resulting ratio indicates that the computational overhead introduced by the step scorer is negligible.Source: D Additional Computational Overhead · p. 12

E Trace-level Score Dynamics

• Trace-level score dynamics are visualized on AIME-25 for Qwen3-4B-Thinking-2507 and DeepSeek-R1-0528-Qwen3-8B.The corresponding visualizations are provided in Figures 6 and 7.Source: E Trace-level Score Dynamics · p. 12, Figure 6 · p. 13, Figure 7 · p. 14
• Each subplot plots prefix mean step scores against token position grouped into 1024-token bins.Solutions use N=64 samples per problem.Source: E Trace-level Score Dynamics · p. 12
• Green and red lines represent average scores across correct and incorrect traces, respectively.Source: E Trace-level Score Dynamics · p. 12
• The results demonstrate that the step scorer separates promising reasoning paths from unpromising ones throughout generation.Source: E Trace-level Score Dynamics · p. 12