用隐藏状态提前判断推理路径好坏，动态剪枝提升大模型推理效率
Hidden States as Early Signals: Step-level Trace Evaluation and Pruning for Efficient Test-Time Scaling
- 基于隐藏状态评估每一步推理质量，动态剪枝低效路径
- 相比自一致方法平均降低45%-70%延迟，准确率还更高
- 适合需要快速推理的部署场景，尤其对长链条任务有效
大型语言模型可通过生成多条推理轨迹实现测试时扩展以增强推理能力。然而，长轨迹与多重采样结合导致计算开销大、端到端延迟高。现有加速方法依赖相似性或置信度剪枝，但这些信号无法可靠反映轨迹质量。为此，我们提出STEP：基于隐藏状态的步骤级轨迹评估与剪枝框架，在生成过程中实时评估每一步并动态剔除低质量轨迹。我们训练了一个轻量级步骤评分器来估计轨迹质量，并设计了基于GPU内存的剪枝策略，在键值缓存占用接近饱和时触发剪枝，显著降低端到端延迟。在多个高难度推理基准上，STEP相比自一致方法平均降低45%-70%延迟，同时提升推理准确率。代码已开源。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。