LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding

📄 arXiv: 2404.16710v4 📥 PDF

作者: Mostafa Elhoushi, Akshat Shrivastava, Diana Liskovich, Basil Hosmer, Bram Wasti, Liangzhen Lai, Anas Mahmoud, Bilge Acun, Saurabh Agarwal, Ahmed Roman, Ahmed A Aly, Beidi Chen, Carole-Jean Wu

分类: cs.CL, cs.AI, cs.LG

发布日期: 2024-04-25 (更新: 2024-10-18)

备注: ACL 2024

DOI: 10.18653/v1/2024.acl-long.681

🔗 代码/项目: GITHUB

💡 一句话要点

提出LayerSkip以加速大语言模型推理

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 推理加速 层丢弃 自我推测解码 早期退出 自然语言处理 计算效率

📋 核心要点

  1. 现有的大语言模型推理速度较慢，尤其是在处理复杂任务时，导致计算资源浪费和响应延迟。
  2. 论文提出LayerSkip，通过训练阶段的层丢弃和早期退出损失函数，优化推理过程，提升早期层的准确性。
  3. 实验结果显示，在多个任务上，LayerSkip实现了最高2.16倍的速度提升，显著提高了推理效率。

📝 摘要（中文）

我们提出了LayerSkip，这是一个端到端的解决方案，用于加速大型语言模型（LLMs）的推理。首先，在训练过程中，我们应用了层丢弃，早期层的丢弃率较低，后期层的丢弃率较高，并引入了一个早期退出损失函数，使所有变换器层共享同一个退出。其次，在推理过程中，我们展示了这种训练方法提高了早期层的早期退出准确性，而无需向模型添加任何辅助层或模块。最后，我们提出了一种新颖的自我推测解码解决方案，在早期层退出后，通过模型的剩余层进行验证和纠正。我们的自我推测解码方法相比其他推测解码方法具有更小的内存占用，并受益于草稿和验证阶段的共享计算和激活。我们在不同的Llama模型大小上进行了实验，展示了在CNN/DM文档摘要、编码和TOPv2语义解析任务上的速度提升。

🔬 方法详解

问题定义：本论文旨在解决大型语言模型推理过程中的速度瓶颈问题。现有方法在处理复杂任务时，往往需要经过所有层的计算，导致效率低下。

核心思路：LayerSkip的核心思路是通过在训练阶段引入层丢弃和早期退出机制，允许模型在推理时根据需要提前退出，从而减少不必要的计算。

技术框架：整体架构包括两个主要阶段：训练阶段和推理阶段。在训练阶段，应用层丢弃和早期退出损失函数；在推理阶段，模型根据早期层的输出决定是否退出，并利用剩余层进行验证和纠正。

关键创新：最重要的技术创新在于自我推测解码方法，它允许模型在早期层退出后，利用后续层进行验证，从而减少内存占用并提高效率。这一设计与传统的推测解码方法有本质区别。

关键设计：在训练过程中，采用低丢弃率的早期层和高丢弃率的后期层设置，损失函数设计为共享退出机制，确保模型在推理时能够有效利用早期层的输出。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，LayerSkip在不同任务上实现了显著的速度提升：在CNN/DM文档摘要任务上提升了2.16倍，在编码任务上提升了1.82倍，在TOPv2语义解析任务上提升了2.0倍。这些结果展示了LayerSkip在提高推理效率方面的有效性。

🎯 应用场景

LayerSkip的研究成果具有广泛的应用潜力，尤其在需要快速响应的自然语言处理任务中，如实时翻译、对话系统和文本摘要等。其高效的推理机制能够显著提升用户体验，并降低计算资源的消耗，推动大语言模型在实际应用中的普及。

📄 摘要（原文）

We present LayerSkip, an end-to-end solution to speed-up inference of large language models (LLMs). First, during training we apply layer dropout, with low dropout rates for earlier layers and higher dropout rates for later layers, and an early exit loss where all transformer layers share the same exit. Second, during inference, we show that this training recipe increases the accuracy of early exit at earlier layers, without adding any auxiliary layers or modules to the model. Third, we present a novel self-speculative decoding solution where we exit at early layers and verify and correct with remaining layers of the model. Our proposed self-speculative decoding approach has less memory footprint than other speculative decoding approaches and benefits from shared compute and activations of the draft and verification stages. We run experiments on different Llama model sizes on different types of training: pretraining from scratch, continual pretraining, finetuning on specific data domain, and finetuning on specific task. We implement our inference solution and show speedups of up to 2.16x on summarization for CNN/DM documents, 1.82x on coding, and 2.0x on TOPv2 semantic parsing task. We open source our code and checkpoints at https://github.com/facebookresearch/LayerSkip.