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

• URL: http://arxiv.org/abs/2601.09093v1
• Date: Wed, 14 Jan 2026 02:54:55 GMT
• Title: Hidden States as Early Signals: Step-level Trace Evaluation and Pruning for Efficient Test-Time Scaling
• Abstract summary: Large Language Models (LLMs) can enhance reasoning capabilities through test-time scaling by generating multiple traces.<n>The combination of lengthy reasoning traces with multiple sampling introduces substantial computation and high end-to-end latency.<n>We propose STEP: Step-level Trace Evaluation and Pruning, a novel pruning framework that evaluates reasoning steps using hidden states.
• Score: 19.080366193748127
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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

Related papers

• A Measurement Study of LLM Inference Trade-offs Across Edge Continuum Hardware [0.14913052010438635]
  Large language models (LLMs) are increasingly used as backends for intelligent web services.<n>This paper presents a controlled measurement study of self-hosted LLM inference across edge and near-edge deployment nodes.
  arXiv  Detail & Related papers  (2026-09-08T06:28:32Z)
• Thought-Level Beam Search for Reasoning [39.76902449987798]
  Test-time compute scaling is a primary driver of performance in large reasoning models.<n>We formalize test-time reasoning as a constrained compute allocation problem over partial trajectories.<n>We introduce Gambit, an inference algorithm that executes emphthought-level beam search
  arXiv  Detail & Related papers  (2026-08-08T09:00:01Z)
• DIRECT: When and Where Should You Allocate Test-Time Compute in Embodied Planners? [57.585275546688116]
  Vision-Language Models (VLMs) are increasingly deployed as high-level planners for embodied agents.<n>We argue that choosing when and where to spend test-time compute is central to bringing frontier performance to the real world.<n>We introduce DIRECT, a routing framework that uses multimodal scene context to allocate compute per prompt.
  arXiv  Detail & Related papers  (2026-06-10T17:58:49Z)
• SWE-TRACE: Optimizing Long-Horizon SWE Agents Through Rubric Process Reward Models and Heuristic Test-Time Scaling [13.335296846555204]
  Current pipelines are bottlenecked by unoptimized demonstration data, sparse execution rewards, and computationally prohibitive inference scaling.<n>We present SWE-TRACE, a unified framework optimizing the SWE agent lifecycle across data curation, reinforcement learning (RL), and test-time inference.<n> experiments on standard SWE benchmarks demonstrate that SWE-TRACE significantly advances the state-of-the-art, maximizing resolution rates while drastically reducing both token consumption inference latency.
  arXiv  Detail & Related papers  (2026-04-16T09:41:47Z)
• Adaptive Parallel Monte Carlo Tree Search for Efficient Test-time Compute Scaling [5.526787431354678]
  Monte Carlo Tree Search (MCTS) is an effective test-time compute scaling (TTCS) method for improving the reasoning performance of large language models.<n>Existing optimizations such as positive early exit, reduce latency in favorable cases but are less effective when search continues without meaningful progress.<n>We introduce it negative early exit, which prunes unproductive MCTS trajectories, and it adaptive boosting mechanism that reallocates reclaimed to reduce resource contention.
  arXiv  Detail & Related papers  (2026-04-01T05:52:38Z)
• Agentic Test-Time Scaling for WebAgents [65.5178428849495]
  We present Confidence-Aware Test-Time Scaling (CATTS), which uses vote-derived uncertainty to allocate compute only when decisions are genuinely contentious.<n>CATTS improves performance on WebArena-Lite and GoBrowse by up to 9.1% over React while using up to 2.3x fewer tokens than uniform scaling.
  arXiv  Detail & Related papers  (2026-02-12T18:58:30Z)
• DASH: Deterministic Attention Scheduling for High-throughput Reproducible LLM Training [22.898073682504023]
  In widely used attention implementations such as FlashAttention-3, the deterministic backward pass can incur up to a 37.9% throughput reduction.<n>We formulate the backward pass of deterministic attention as a scheduling problem on a Directed Acyclic Graph (DAG)<n>We present DASH (Deterministic Attention Scheduling for High-Throughput), which encapsulates two complementary scheduling strategies.
  arXiv  Detail & Related papers  (2026-01-29T15:10:13Z)
• Seer Self-Consistency: Advance Budget Estimation for Adaptive Test-Time Scaling [55.026048429595384]
  Test-time scaling improves the inference performance of Large Language Models (LLMs) but also incurs substantial computational costs.<n>We propose SeerSC, a dynamic self-consistency framework that simultaneously improves token efficiency and latency.
  arXiv  Detail & Related papers  (2025-11-12T13:57:43Z)
• dParallel: Learnable Parallel Decoding for dLLMs [77.24184219948337]
  Diffusion large language models (dLLMs) offer parallel token prediction and lower inference latency.<n>Existing open-source models still require nearly token-length decoding steps to ensure performance.<n>We introduce dParallel, a simple and effective method that unlocks the inherent parallelism of dLLMs for fast sampling.
  arXiv  Detail & Related papers  (2025-09-30T16:32:52Z)
• On the Bias of Next-Token Predictors Toward Systematically Inefficient Reasoning: A Shortest-Path Case Study [4.798155648915794]
  We study two key factors for improving reasoning in large language models.<n>We train decoder-only transformers on question-trace-answer triples using a custom tokenizer.<n>With the same training-token budget, models trained on inefficient traces generalize better to unseen graphs.
  arXiv  Detail & Related papers  (2025-07-07T18:00:06Z)
• $\texttt{SPECS}$: Faster Test-Time Scaling through Speculative Drafts [55.231201692232894]
  $textttSPECS$ is a latency-aware test-time scaling method inspired by speculative decoding.<n>Our results show that $textttSPECS$matches or surpasses beam search accuracy while reducing latency by up to $sim$19.1%.
  arXiv  Detail & Related papers  (2025-06-15T05:50:05Z)
• Tree Search-Based Policy Optimization under Stochastic Execution Delay [46.849634120584646]
  Delayed execution MDPs are a new formalism addressing random delays without resorting to state augmentation. We show that given observed delay values, it is sufficient to perform a policy search in the class of Markov policies. We devise DEZ, a model-based algorithm that optimize over the class of Markov policies.
  arXiv  Detail & Related papers  (2024-04-08T12:19:04Z)
• DRIVE: Dual Gradient-Based Rapid Iterative Pruning [2.209921757303168]
  Modern deep neural networks (DNNs) consist of millions of parameters, necessitating high-performance computing during training and inference. Traditional pruning methods that are applied post-training focus on streamlining inference, but there are recent efforts to leverage sparsity early on by pruning before training. We present Dual Gradient-Based Rapid Iterative Pruning (DRIVE), which leverages dense training for initial epochs to counteract the randomness inherent at the inception.
  arXiv  Detail & Related papers  (2024-04-01T20:44:28Z)
• Exploring Dynamic Transformer for Efficient Object Tracking [58.120191254379854]
  We propose DyTrack, a dynamic transformer framework for efficient tracking.<n>DyTrack automatically learns to configure proper reasoning routes for various inputs, gaining better utilization of the available computational budget.<n>Experiments on multiple benchmarks demonstrate that DyTrack achieves promising speed-precision trade-offs with only a single model.
  arXiv  Detail & Related papers  (2024-03-26T12:31:58Z)
• Label Delay in Online Continual Learning [77.05325581370893]
  A critical aspect often overlooked is the label delay, where new data may not be labeled due to slow and costly annotation processes. We introduce a new continual learning framework with explicit modeling of the label delay between data and label streams over time steps. We show experimentally that our method is the least affected by the label delay factor and in some cases successfully recovers the accuracy of the non-delayed counterpart.
  arXiv  Detail & Related papers  (2023-12-01T20:52:10Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.