Source-linked AI summary
Hidden States as Early Signals: Step-level Trace Evaluation and Pruning for Efficient Test-Time Scaling
Zhixiang Liang, Beichen Huang, Zheng Wang, Minjia Zhang
TL;DR
Parallel test-time scaling is costly because lengthy traces and multiple samples create high computation and latency, while existing pruning signals may misidentify trace quality. STEP uses hidden-state step scoring with GPU-memory-aware pruning, reducing average end-to-end latency by 45%–70% versus self-consistency while improving reasoning accuracy.
Problem
Lengthy reasoning traces and multiple samples create high computation and latency, while existing similarity- and confidence-based pruning signals do not reliably indicate trace quality.
Method
STEP trains a lightweight scorer on step-boundary hidden states and prunes unpromising traces when GPU memory utilization approaches KV-cache saturation.
Results
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.
Takeaways & Limitations
Combining hidden-state-based early pruning with GPU memory-aware system optimization supports more efficient parallel scaling for complex reasoning tasks.
Takeaways & Limitations
The step scorer uses noisy pseudo-labels that assume all steps in a correct trace are good and all steps in an incorrect trace are bad.
Abstract
from arXiv · showhide
Large Language Models (LLMs) can enhance reasoning capabilities through test-time scaling by generating multiple traces. However, the combination of lengthy reasoning traces with multiple sampling introduces substantial computation and high end-to-end latency. Prior work on accelerating this process has relied on similarity-based or confidence-based pruning, but these signals do not reliably indicate trace quality. To address these limitations, we propose STEP: Step-level Trace Evaluation and Pruning, a novel pruning framework that evaluates reasoning steps using hidden states and dynamically prunes unpromising traces during generation. We train a lightweight step scorer to estimate trace quality, and design a GPU memory-aware pruning strategy that triggers pruning as the GPU memory is saturated by KV cache to reduce end-to-end latency. Experiments across challenging reasoning benchmarks demonstrate that STEP reduces end-to-end inference latency by 45%-70% on average compared to self-consistency while also improving reasoning accuracy. Our code is released at: https://github.com/Supercomputing-System-AI-Lab/STEP
1 Introduction
Test-time scaling improves LLM reasoning by generating multiple traces, but lengthy traces and parallel sampling create prohibitive computation and latency. STEP addresses unreliable trace-quality signals and GPU-memory-induced delays through hidden-state evaluation and memory-aware pruning.
- Self-consistency generates multiple reasoning traces and selects answers by majority voting, but lengthy traces and parallel sampling impose substantial computational cost and latency.
- Prior pruning methods rely on similarity or confidence signals that may not reliably reflect reasoning quality.
- GPU memory saturation from multiple long traces causes preemption, waiting queues, redundant computation, and end-to-end latency.
- STEP trains a lightweight scorer on reasoning-step hidden states to identify unpromising traces early during generation.
- STEP triggers pruning as KV-cache accumulation saturates GPU memory, releasing resources to prevent preemption and queuing delays.
- 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.
2 Related Work
Prior parallel-scaling methods prune traces using confidence or diversity signals, while recent work explores hidden states as indicators of reasoning quality. STEP extends this direction with continuous step-level monitoring and early pruning.
- Confidence-based methods prune low-confidence traces, whereas diversity-based methods remove similar traces to preserve answer diversity.
- Confidence can suffer from model overconfidence and miscalibration, while surface-level similarity may not indicate reasoning redundancy.
- Prior studies show that hidden states can encode correctness-related information and support lightweight prediction of intermediate-answer correctness.
- STEP trains a lightweight step scorer on hidden-state boundary representations for continuous monitoring and early pruning during parallel scaling.
3 Motivation
The motivation is twofold: hidden states may distinguish correct from incorrect reasoning early, while parallel scaling suffers from token inefficiency and GPU-memory-driven waiting. These observations motivate step-level evaluation and system-aware pruning.
- Hidden states in reasoning models contain rich signals indicative of reasoning quality, and the main latency bottleneck arises from inference-system behavior.
- A 2-layer MLP trained on hidden states finds that scores from early reasoning steps effectively distinguish correct from incorrect traces.
- Incorrect traces tend to be longer than correct traces, so early termination can reduce generated tokens while improving efficiency and accuracy.
- KV-cache accumulation saturates GPU memory and causes inference engines to preempt traces, creating waiting time beyond token-generation costs.
4 STEP
STEP evaluates reasoning traces from hidden states at step boundaries and prunes low-scoring traces when GPU memory limits further decoding. It combines step-level scoring, memory-aware pruning, and weighted answer selection to reduce inference inefficiency.
- 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.
- 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.
- 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.
- Step scorer training: 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.
- 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.
- 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.
5 Experiment
STEP is evaluated across multiple open reasoning models, challenging benchmarks, and baselines using accuracy, token usage, and latency. Results show improved accuracy, favorable accuracy-latency trade-offs, stronger early trace ranking, and lower end-to-end latency through memory-aware pruning.
- 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.
- 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.
- Main Results: 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.
- 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.
- 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.
- 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.
6 Conclusion
STEP combines hidden-state-based trace evaluation with GPU memory-aware pruning for efficient test-time scaling. Across six benchmarks and three model scales, it reduces latency while improving reasoning accuracy.
- 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.
- 45%-70% average end-to-end inference latency reduction compared to self-consistency accompanies +0.4 to +7.5 percentage-point reasoning-accuracy gains.
Limitations
STEP has limitations involving noisy step-level supervision, infrastructure-dependent latency gains, and evaluation restricted to fixed-form verifiable answers.
- 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.
- Memory-triggered latency improvements are tightly coupled to serving infrastructure, so their magnitude may differ across inference engines and multi-GPU configurations.
- Evaluation focuses on tasks with fixed-form verifiable answers, while open-ended generation would require adapting answer aggregation beyond majority voting.
A.1 Training Parameters
The step scorer is trained using fixed hyperparameters and a balanced dataset of automatically verified correct and incorrect mathematical traces.
- The training process uses hyperparameters listed in Table 5.
- Input dimensions match the model hidden-state sizes: 2560, 4096, and 5120 for the three evaluated LLMs.
- Training problems come from HMMT 2012–2023 February competitions in Algebra, Combinatorics, and Geometry.
- For each problem, 64 sampled solutions are verified with numeric matching and SymPy-based symbolic equivalence.
- The scorer training set contains 5,000 correct and 5,000 incorrect traces for each LLM.
B Experimental Settings
Experiments use fixed sampling settings, a standard prompt template, specified pruning baselines, and a vLLM-V1-based implementation with process-separated scheduling.
- Sampling parameters remain fixed across methods, with details listed in Table 6.
- All methods receive the prompt: “Please reason step by step, and put your final answer within \\boxed{}.”
- Slim-SC uses Random Pruning, while DeepConf-low terminates traces below a threshold retaining the top 10% highest-confidence warmup traces.
- STEP places the scorer with the model runner on the same GPU and sends scores to the engine core for scheduler pruning decisions.
- The implementation follows vLLM-V1, with the engine core and model runner residing in separate processes.
D Additional Computational Overhead
STEP’s auxiliary step scorer adds negligible computation relative to the underlying LLM because its per-step FLOP cost is extremely small under practical settings.
- 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.
- D Additional Computational Overhead: 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.
- D Additional Computational Overhead: 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.
- D Additional Computational Overhead: The resulting ratio indicates that the computational overhead introduced by the step scorer is negligible.
E Trace-level Score Dynamics
STEP’s step scorer separates promising from unpromising reasoning paths throughout generation across two models on AIME-25. The analysis uses prefix mean scores over token-position bins and compares correct with incorrect traces.
- 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.
- E Trace-level Score Dynamics: Each subplot plots prefix mean step scores against token position grouped into 1024-token bins.Solutions use N=64 samples per problem.
- E Trace-level Score Dynamics: Green and red lines represent average scores across correct and incorrect traces, respectively.
- E Trace-level Score Dynamics: The results demonstrate that the step scorer separates promising reasoning paths from unpromising ones throughout generation.