dorsal/arxiv
View SchemaHidden States as Early Signals: Step-level Trace Evaluation and Pruning for Efficient Test-Time Scaling
| Authors | Zhixiang Liang, Beichen Huang, Zheng Wang, Minjia Zhang |
|---|---|
| Categories | |
| ArXiv ID | 2601.09093vv1 |
| URL | https://arxiv.org/abs/2601.09093 |
| 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
{
"annotation_id": "7e3f42d7-bfd8-4619-a663-263750677612",
"date_created": "2026-02-17T05:53:20.237000Z",
"date_modified": "2026-02-17T05:53:20.237000Z",
"file_hash": "cfe36cb65cd2713b5dc8fe3c71ac0bd14f0b48d75a369d83e6e44f16692f5fc3",
"private": false,
"record": {
"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",
"arxiv_id": "2601.09093",
"authors": [
"Zhixiang Liang",
"Beichen Huang",
"Zheng Wang",
"Minjia Zhang"
],
"categories": [
"cs.LG"
],
"license": "http://arxiv.org/licenses/nonexclusive-distrib/1.0/",
"title": "Hidden States as Early Signals: Step-level Trace Evaluation and Pruning for Efficient Test-Time Scaling",
"url": "https://arxiv.org/abs/2601.09093",
"version": "v1"
},
"schema_id": "dorsal/arxiv",
"source": {
"execution_id": "b4c7ef57-1775-4280-8618-8a3bcb3f37b6",
"id": "arXiv Dataset",
"type": "Model",
"variant": "snapshot-2026-01-17",
"version": "0.1.0"
},
"user_id": 1000002
}