针对多阶段推理的延迟优化,提出动态跳层框架提升效率。
LiteStage: Latency-aware Layer Skipping for Multi-stage Reasoning
- 按阶段分配最优跳层预算,精准控制计算开销
- 在线置信度判断提前终止生成,减少冗余解码
- 在多个基准上实现更高效率与精度平衡
多阶段推理通过将复杂问题分解为连续子阶段,有效提升小型语言模型的推理能力,但带来延迟增加的问题。现有自适应加速技术(如层跳过)在此场景下难以兼顾效率与准确率,主要受两大挑战制约:(1) 各阶段对跳过的敏感性差异,(2) 生成冗余输出标记。为此,我们提出 LiteStage——一种面向多阶段推理的延迟感知层跳过框架。该框架结合阶段级离线搜索以分配最优层预算,并引入在线置信度驱动的生成早停机制,抑制不必要的解码过程。在 OBQA、CSQA 与 StrategyQA 三个基准上的实验表明,LiteStage 显著优于先前无训练的层跳过方法。
原文摘要 · Abstract (English)
Multi-stage reasoning has emerged as an effective strategy for enhancing the reasoning capability of small language models by decomposing complex problems into sequential sub-stages. However, this comes at the cost of increased latency. We observe that existing adaptive acceleration techniques, such as layer skipping, struggle to balance efficiency and accuracy in this setting due to two key challenges: (1) stage-wise variation in skip sensitivity, and (2) the generation of redundant output tokens. To address these, we propose LiteStage, a latency-aware layer skipping framework for multi-stage reasoning. LiteStage combines a stage-wise offline search that allocates optimal layer budgets with an online confidence-based generation early exit to suppress unnecessary decoding. Experiments on three benchmarks, e.g., OBQA, CSQA, and StrategyQA, show that LiteStage outperforms prior training-free layer skipping methods.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。