Adaptive Layer-skipping in Pre-trained LLMs

本文提出FlexiDepth方法，通过插件式路由器和适配器实现预训练LLM的自适应层跳过，提高计算效率同时保持生成性能，并通过实验揭示了token类型对计算需求的影响。

Large Language Model, Transformer, Efficiency, Adaptive Systems, Tokenization, Pre-training

Xuan Luo, Weizhi Wang, Xifeng Yan

University of California, Santa Barbara

Generated by grok-3-mini-latest

Background Problem

大型语言模型（LLMs）在生成每个token时通常进行完整的Transformer解码器层前向传播，这种统一分配方式虽然简单但效率低下，因为不同token的计算需求存在显著差异。例如，生成重复token或固定短语可能只需要较少计算，而处理高不确定性或计算密集型token则需要更多资源。现有的层跳过方法（如基于统计信息或早停机制）虽然能减少计算成本，但忽略了一个根本问题：如何根据token类型动态调整计算需求。本文的工作起点是解决这一问题，通过自适应层跳过来优化LLMs的计算效率，同时避免性能下降，并揭示token生成中的计算需求模式。

Method

核心思想: FlexiDepth是一种插件式方法，旨在在不修改预训练LLM参数的情况下，实现自适应层跳过，以动态调整每个token的计算深度。具体实现包括在每个Transformer解码器层引入两个轻量级模块：路由器和适配器。路由器使用一个瓶颈MLP（bottlenecked MLP）计算门控分数，公式为$G = \sigma(\text{Router}(\text{Norm}(X)))$，其中Router定义为$\text{Router}(z) = W_{\text{f}} \cdot (W_{\uparrow} \cdot \text{Norm}(\tanh(W_{\downarrow}z)))$，然后通过预设阈值τ决定是否跳过该层。如果门控分数大于τ，token进入完整处理路径（包括注意力机制和前馈网络FFN），否则进入跳过路径。注意跳过: 在跳过注意力模块时，保留KV缓存以维护上下文完整性（如图3所示），确保后续token能访问完整信息；跳过FFN模块时，使用一个缩小16倍中间维度的轻量级适配器对隐藏状态进行表示对齐。层跳过损失: 为平衡效率和性能，引入跳过损失$\mathcal{L}_{skip} = \frac{1}{T} \sum_{t=1}^{T} \left( \sum_{l=1}^{L} g_t^l \right)^2$，并与语言建模损失结合，总体损失为$\mathcal{L} = \alpha \cdot \mathcal{L}_{skip} + \mathcal{L}_{lm}$，其中α控制跳过强度。

Experiment

实验设置: 本文在Llama-3-8B-Instruct模型（32层）上实现了FlexiDepth，将后16层转换为FlexiDepth层，使用Tulu-v2数据集训练路由器和适配器，训练参数包括学习率1e-4、AdamW优化器等。评估基准包括单token生成任务（如MMLU、HellaSwag、Winogrande）和多token生成任务（如GSM8K、HumanEval、CoQA），采用5-shot或zero-shot设置。基线方法包括LayerSkip、ShortGPT、LaCo和MindSkip，所有方法均配置为跳过相同层数。实验结果: FlexiDepth在跳过8层时保留100.7%的基准性能，显著优于基线（如ShortGPT在GSM8K上准确率降至0.001），尤其在多token任务中表现突出。进一步分析显示，FlexiDepth的层分配模式符合直觉：摘要任务平均使用更多层（28.65层），而复制任务使用较少（21.95层）；数学任务中，乘法需要更多层（23.90层） than 加法（22.45层）。消融实验确认了路由器设计、KV缓存和适配器的必要性：无适配器时性能降至28.1%。实验设计合理全面，数据集选择多样，结果与预期一致，证明了自适应深度的有效性，并提供了层分配数据集促进未来研究。

Further Thoughts

FlexiDepth的层跳过机制启发我们思考LLM内部计算的模块化特性，可能类似于人类认知中不同脑区的功能分工，例如早期层处理上下文理解，中间层执行任务特定计算，晚期层负责输出生成；此外，这项工作可扩展到混合专家模型（Mixture-of-Experts）中，优化专家激活以减少跨设备通信开销，或与硬件加速技术结合（如token分组和负载均衡）提升实际推理速度；未来可探索将这种自适应策略应用于多模态模型或强化学习场景，以实现更智能的资源分配，潜在地提升AI系统的能效和泛化能力。