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

TL;DR

STEP tackles the latency and cost of test-time scaling by pruning reasoning traces during generation using a hidden-state based step scorer and a GPU memory aware trigger. The approach aggregates per-step scores at step boundaries to form trace-level scores and prunes the weakest traces when KV cache memory is saturated, thereby eliminating waiting time as a major latency component. Across multiple models and challenging benchmarks, STEP delivers end-to-end latency reductions of $45\%$ to $70\%$ and improves reasoning accuracy, outperforming existing pruning baselines. This combination of early, signal-based pruning and system-level memory management enables more practical and scalable deployment of parallel reasoning with LLMs. The work demonstrates that hidden-state signals can reliably guide trace quality assessment and that memory-aware pruning is essential to reducing real-world latency.

Abstract

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

Figures (7)

• Figure 1: Comparison of accuracy versus latency across different methods on DeepSeek-R1-0528-Qwen3-8B. STEP achieves superior accuracy (averaged across AIME-25, HMMT-24/25, GPQA-D) while significantly reducing latency compared to baseline methods.
• Figure 2: (a) Hidden state score distributions for correct vs. incorrect reasoning traces on HMMT-25. Scores are computed by the scorer model as averages over the first 25%, 50%, and 75% of reasoning steps. (b) Token count comparison of correct and incorrect traces for AIME-25 Q28 using Qwen3-4B-Thinking-2507; incorrect traces average 42.5k tokens compared to 35.3k for correct ones. (c) Time distribution for generating one trace on the same setup; waiting time (59%) dominates over actual decoding (40%), with KV cache reconstruction accounting for 1%.
• Figure 3: Overview of the STEP framework. The step-level scoring module extracts hidden states at step boundaries and uses a trained step scorer to compute step-level scores, which are averaged to obtain trace-level scores. The KV-cache monitor triggers pruning when GPU memory is saturated, removing the trace with the lowest score and releasing its KV cache to prevent queuing delays.
• Figure 4: Latency scaling comparison between STEP and baseline methods on AIME-25 and HMMT-25 using Qwen3-4B-Thinking-2507 and DeepSeek-R1-0528-Qwen3-8B.
• Figure 5: Pairwise RankAcc of the hidden-state-based step scorer versus token-level confidence.