2025-11-17T14:34:12.785982

LiteStage: Latency-aware Layer Skipping for Multi-stage Reasoning

Kang, Song, Kim
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 achieves up to 1.70x speedup with less than 4.0% accuracy loss, outperforming prior training-free layer skipping methods.
academic

LiteStage: Latency-aware Layer Skipping for Multi-stage Reasoning

基本信息

摘要

多阶段推理已成为通过将复杂问题分解为顺序子阶段来增强小型语言模型推理能力的有效策略。然而，这以增加延迟为代价。作者观察到现有的自适应加速技术（如层跳跃）在这种设置下难以平衡效率和准确性，主要面临两个关键挑战：(1) 阶段间跳跃敏感性的差异，(2) 冗余输出token的生成。为解决这些问题，本文提出了LiteStage，一个面向多阶段推理的延迟感知层跳跃框架。LiteStage结合了分配最优层预算的阶段性离线搜索和基于置信度的在线生成早退机制来抑制不必要的解码。在OBQA、CSQA和StrategyQA三个基准上的实验表明，LiteStage实现了高达1.70×的加速，准确率损失少于4.0%，优于先前的无训练层跳跃方法。

研究背景与动机

问题定义

多阶段推理通过将复杂问题分解为多个连续的子问题来增强小型语言模型的推理能力。例如，TinyThinker采用三阶段推理：回忆(Recall)、分析(Analysis)和总结(Summary)。虽然这种方法有效提升了推理质量，但也不可避免地增加了推理延迟。

核心挑战

作者通过深入分析发现了两个关键问题：

  1. 阶段间跳跃敏感性差异：不同推理阶段对层跳跃的敏感程度存在显著差异。实验表明，Stage 3（总结阶段）对层跳跃最为敏感，而Stage 1（回忆阶段）相对鲁棒。
  2. 冗余token生成：层跳跃虽然降低了每token的计算成本，但往往导致生成更多token，反而增加了端到端延迟。

现有方法局限性

现有的层跳跃方法（如SkipDecode、UnifiedSkip、AdaSkip）通常采用统一的跳跃策略，无法适应多阶段推理中不同阶段的特性，导致：

  • 在敏感阶段过度压缩造成准确率急剧下降
  • 忽视了层跳跃引起的生成长度增加问题
  • 缺乏延迟感知的优化机制

核心贡献

  1. 提出LiteStage框架：首个专门针对多阶段推理的延迟感知层跳跃框架，有效解决了阶段间敏感性差异和冗余token生成问题。
  2. 阶段性层预算分配策略：设计了从最慢阶段到最快阶段的贪婪搜索算法，为每个推理阶段分配最优的层跳跃预算。
  3. 置信度驱动的生成早退机制：引入在线置信度监控，动态终止低置信度的冗余生成，进一步提升推理效率。
  4. 显著的性能提升：在三个基准数据集上实现1.16-1.70×加速，准确率损失仅0.4-4.0%，大幅超越现有无训练方法。

方法详解

任务定义

给定测试数据集D，目标是寻找阶段性层预算L，使得在给定准确率阈值ε内最小化推理延迟：

argmin_L (1/|D|) ∑_{d∈D} T(M_L(d))
subject to: A(M_L(d)) ≤ A(M(d)) - ε

其中T和A分别表示推理延迟和准确率，M_L和M分别表示应用层跳跃和完整层的模型。

模型架构

LiteStage包含两个互补组件：

1. 离线配置 (Offline Configuration)

Step 1: 层重要性估计

  • 采用子层级别的余弦相似度作为重要性代理
  • 分别计算多头自注意力(MHSA)和前馈网络(FFN)的重要性：
I^(j)_MHSA = (1/N) ∑_{n=0}^{N-1} cos(MHSA^(j)(x) + x, x)
I^(j)_FFN = (1/N) ∑_{n=0}^{N-1} cos(FFN^(j)(x) + x, x)

Step 2: 层预算搜索

  • 从最慢的推理阶段开始进行贪婪搜索
  • 构建准确率-延迟曲线，在满足准确率约束下选择最优延迟配置
  • 逐阶段优化，确保阶段间交互得到准确反映

2. 在线调整 (Online Adjustment)

Step 3: 生成早退

  • 维护最近n个token的置信度缓存
  • 计算平均置信度μ_Conf，当低于阈值时提前终止生成
  • 置信度定义为每个token的最大logit值

技术创新点

  1. 非均匀层预算分配：根据各阶段的敏感性差异，自适应分配不同的层跳跃预算，避免在敏感阶段过度压缩。
  2. 延迟感知优化：不仅考虑准确率，还综合考虑实际推理延迟，自动排除那些虽然跳跃更多层但延迟更高的配置。
  3. 动态生成控制：通过置信度监控主动控制生成长度，缓解层跳跃导致的冗余token问题。

实验设置

数据集

采用TinyThinker的三阶段推理流程，在三个问答基准上进行评估：

  • OpenBookQA (OBQA)：开放式问答任务
  • CommonSenseQA (CSQA)：常识推理问答
  • StrategyQA：策略性推理问答

评价指标

  • 准确率：问答正确率
  • 加速比：相对于完整层模型的推理速度提升
  • 延迟：端到端推理时间

对比方法

  • SkipDecode：渐进式深层跳跃
  • UnifiedSkip：周期性层跳跃
  • AdaSkip：基于余弦相似度的子层重要性估计

实现细节

  • 主要使用TinyLlama-1.1B-Chat-v1.0模型
  • 训练10个epoch，批大小16(OBQA/CSQA)或24(StrategyQA)
  • 学习率5×10^-5
  • 采用自一致性协议进行10次迭代评估
  • 置信度阈值设为0.5，缓存大小n=5

实验结果

主要结果

在三个基准数据集上，LiteStage均显著优于基线方法：

数据集基线准确率LiteStage准确率加速比
OBQA64.0%60.0%1.32×
CSQA54.8%53.2%1.16×
StrategyQA62.4%62.0%1.70×

关键发现

  1. 阶段敏感性差异：通过单阶段跳跃实验证实，Stage 3对层跳跃最敏感，其准确率曲线几乎决定了整体性能上界。
  2. 延迟悖论：更多的层跳跃并不总是带来更快的推理速度，由于生成长度增加，某些配置反而导致延迟增加。
  3. 置信度模式：层跳跃模型的token置信度呈单调递减趋势，而完整层模型在后期可能恢复置信度。

消融实验

非均匀层预算的效果：

  • 在相同层跳跃数量下，LiteStage的准确率显著高于统一跳跃策略
  • 随着跳跃层数增加，性能差距进一步扩大

生成早退的贡献：

  • 在轻度层跳跃时，生成早退影响很小（-0.5%解码步数）
  • 在重度层跳跃时，可减少高达82.5%的解码步数
  • 准确率基本保持稳定，个别情况下甚至有所提升

案例分析

通过CSQA的具体案例展示，生成早退能够有效截断低置信度的冗余文本，同时保持核心推理逻辑完整，最终答案保持一致。

相关工作

多阶段生成

  • TinyThinker：提出回忆-分析-总结的三阶段推理循环
  • DeAR：采用分解-分析-重思考的过程
  • CasCoD：级联方式蒸馏分解的思维链
  • Self-Discover：动态组织推理结构

层跳跃技术

基于训练的方法：

  • LayerSkip, DeeBERT, EE-LLM：中间层早退
  • Mixture-of-Depth：需要训练模型和路由器

无训练方法：

  • SkipDecode：渐进式深层跳跃
  • Unified Skipping：周期性跳跃
  • ShortGPT：基于余弦相似度
  • AdaSkip：子层级重要性估计

生成早退

现有方法主要针对冗长推理模型，缺乏对模型压缩导致的生成延长问题的关注。

结论与讨论

主要结论

  1. 多阶段推理中的非均匀敏感性：不同推理阶段对层压缩的敏感程度存在显著差异，需要差异化的优化策略。
  2. 延迟感知优化的必要性：单纯的层跳跃可能因生成长度增加而导致延迟恶化，需要综合考虑准确率和延迟。
  3. 生成控制的有效性：基于置信度的生成早退能够有效缓解层跳跃导致的冗余生成问题。

局限性

  1. 离线搜索开销：相比其他无训练方法，LiteStage的离线配置需要更多计算资源（约1-7.6小时）。
  2. 模型架构依赖：主要在Llama系列模型上验证，在Qwen等其他架构上效果有限。
  3. 适用范围限制：专门针对多阶段推理场景，对单阶段推理的适用性未充分验证。

未来方向

  1. 扩展到更多模型架构：研究不同架构的跳跃敏感性特征
  2. 动态预算分配：开发运行时自适应调整层预算的机制
  3. 多模态推理优化：将框架扩展到视觉-语言等多模态推理任务

深度评价

优点

  1. 问题识别准确：准确识别了多阶段推理中的关键瓶颈，包括阶段敏感性差异和冗余生成问题。
  2. 方法设计合理：离线-在线结合的框架设计巧妙，既保证了优化效果又控制了运行时开销。
  3. 实验设计充分：通过详细的动机实验、消融研究和案例分析，充分验证了方法的有效性。
  4. 实用价值高：作为无训练方法，具有良好的实际应用前景。

不足

  1. 理论分析不足：缺乏对阶段敏感性差异的理论解释，主要依赖经验观察。
  2. 参数设置启发式：置信度阈值、缓存大小等关键参数主要基于启发式设定，缺乏系统性分析。
  3. 泛化性有限：在不同模型架构上的表现差异较大，泛化能力有待提升。

影响力

  1. 学术贡献：首次系统性地研究了多阶段推理中的层跳跃优化问题，为相关研究提供了新的视角。
  2. 实用价值：为小型语言模型的高效推理提供了实用的解决方案，有助于推动边缘部署。
  3. 可复现性：提供了完整的代码实现，便于后续研究和应用。

适用场景

LiteStage特别适合以下场景：

  • 资源受限的边缘设备部署
  • 需要多阶段推理的复杂任务
  • 对延迟敏感的实时应用
  • 小型语言模型的推理加速

参考文献

论文引用了多个重要的相关工作，包括：

  • TinyThinker (Piao and Park, 2024)：多阶段推理的代表性工作
  • AdaSkip (He et al., 2025)：子层级层跳跃的最新方法
  • Mixture-of-Depths (Raposo et al., 2024)：动态计算分配的开创性工作

总体评价：本文针对多阶段推理中的层跳跃优化问题提出了创新性的解决方案，在理论洞察和实际效果方面都有显著贡献。虽然存在一些局限性，但为小型语言模型的高效推理开辟了新的研究方向，具有重要的学术价值和实用意义。