LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding

TL;DR

2024 年 Meta FAIR 提出了 LayerSkip，这是一种端到端的解决方案，用于加速大语言模型（LLMs）的推理过程

Paper name

LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding

Paper Reading Note

Paper URL:

https://arxiv.org/pdf/2404.16710

Code URL:

https://github.com/facebookresearch/LayerSkip

Introduction

背景

将 LLM 进一步加速部署至移动或边缘设备仍是研究热点。当前多数加速方案聚焦于减少非零权重数量（即稀疏化）、降低权重位宽（即量化）或削减每层注意力头数量（即头剪枝），仅有少量研究关注层数缩减。
本文探索通过在推理过程中提前退出来减少每 token 所需层数。与量化或稀疏化不同，层数缩减无需专用硬件或软件内核支持。

本文方案

提出了 LayerSkip，这是一种端到端的解决方案，用于加速大语言模型（LLMs）的推理过程
- 在训练过程中，采用了分层的 dropout 策略：对前面的 Transformer 层使用较低的 dropout 率，而对后面的层使用较高的 dropout 率，并引入了一种“早退出损失”（early exit loss），使所有 Transformer 层共享同一个退出点（exit point）
- 在推理阶段，展示了这种训练策略能够提高在较前层中提前退出的准确性，而且无需向模型中添加任何辅助层或模块
- 提出了一种新颖的自推测解码（self-speculative decoding）方法：在早期层进行输出预测，并利用模型后续的层对其进行验证和修正。提出的自推测解码方法相比其他推测解码方法具有更小的内存占用，并能从草案生成（drafting）和验证（verification）阶段之间的计算与激活值共享中受益。
实验显示该方案在不同任务上可实现1.34×至2.16×加速

Methods

出发点

以 HumanEval 编程数据集的首个提示为例，分析 Llama1 7B 模型各层行为（图2a）。当生成每个 token 时，通过将各 transformer 层的输出嵌入投影到语言模型头（含最终层归一化与线性层），经 softmax 获取预测 token。图 2b 显示：早期层预测常与最终结果无关，后期层预测逐渐收敛；多数 token 在到达最终层前已能正确预测；中间层存在预测反复现象（如 Token 05在 7 层预测 “range”，22-26 层改变主意后又回归）。统计显示平均每个 token 仅需 23.45 层（总 32 层），表明即使有完美预测器，最多节省 26% 计算量。因此一般不会提前退出。如图 2b 所示，即使是简单 token（如开启 for 循环的 Token 02）也需要全部 32 层来预测 “for”。本文提出在训练中采用分层丢弃策略——对深层施加更高丢弃率，使模型更依赖浅层特征。
过早退出的校正机制
- 无论是基于启发式规则还是改进训练流程，早期退出都可能降低准确率。理想方案应能验证早期预测的正确性，必要时通过执行剩余层进行校正。受推测解码技术启发，本文提出自推测解码：用早期退出机制自回归生成 token，用剩余层并行验证校正 token 组。该方法的优势在于验证 token 组的速度快于自回归生成。

4.1.1 层丢弃（Layer Dropout）

我们对常见的训练策略所做的第一个修改是引入层丢弃。因而，第  $l$  层在训练迭代  $t$  时的Transformer层操作变为：

$x_{l+1,t} = x_{l,t} + M(p_{l,t})f_l(x_{l,t}) \quad (1)$

其中， $p_{l,t}$ 是第  $l$  层在第  $t$  次迭代的丢弃率， $M (p)$ 是伯努利函数，以概率  $p$ 返回 0，以概率  $1 - p$ 返回 1。我们对批次中的每个样本分别应用丢弃操作：从批次中移除被丢弃的样本，对剩余样本应用 Transformer 操作 $f_l$ ，然后将输出与被丢弃的样本拼接起来。

为了在训练中实现更高的加速效果，我们为每个 GPU 使用相同的随机种子，这样每次迭代中每层 Transformer 将丢弃相同数量的样本。

丢弃率可以在每层  $l$  和每次迭代  $t$  中变化，定义如下：

$p_{l,t} = S(t) D(l) p_{\text{max}} \quad (2)$

其中 $p_{\text{max}}$ 是训练期间模型中最大丢弃率的超参数， $D (l)$ 是每层的缩放函数， $S (t)$ 是每时间步的缩放函数。

我们发现最佳的每层缩放方式是让丢弃率在各层中指数增加，从第 0 层的 0.0 增加到第 $L - 1$ 层的 1.0：

$\frac{e^{l \ln 2 / (L - 1)} - 1}{e^{\ln 2} - 1} \quad (3)$

对于时间上的缩放 $S (t)$ ，我们发现当从预训练模型进行持续预训练或微调时，最好不对时间进行缩放，即设 $S (t) = 1$ 。但如果是从头开始预训练，我们发现使用指数式课程（Exponential Curriculum） $S_{\text{exp}}(t)$ 能带来最好的精度表现：

$S_{\text{exp}}(t) = \frac{e^{t \ln 2 / T} - 1}{e^{\ln 2} - 1} \quad (4)$

4.1.2 提前退出损失（Early Exit Loss）

为了提升模型低层的预测能力，我们需要确保语言建模头（LM head） $g$ 能够解嵌（unembed）来自不同层的输出。因此，在训练中，我们在层丢弃的基础上引入了每层的提前退出损失（Early Exit Loss）。

我们在训练时直接监督每层与 LM head 的连接，从而使较低层也能参与语言建模任务。模型在迭代 $t$ 时的总损失如下：

$\sum_{l=0}^{L-1} \tilde{e}(t, l) \cdot J_{\text{CE}}(g(x_{l+1}), Y) \quad (5)$

其中 $\tilde{e}(t, l)$ 是每层的归一化损失权重，其总和为 1：

$\tilde{e}(t, l) = \frac{C(t, l) e(l)}{\sum_{i=0}^{L-1} C(t, i) e(i)} \quad (6)$

$C (t, l)$ 是一个二值课程函数，控制第 $l$ 层在第 $t$ 次迭代时是否启用提前退出。我们参考 Elbayad 等人（2020）的工作，设定逐层增加的缩放因子，使得后层的惩罚更大（因为后层更容易预测）：

$\begin{cases} e_{\text{scale}} \sum_{i=0}^{l} i, & \text{如果 } 0 \leq l < L - 1 \\ L - 1 + e_{\text{scale}} \sum_{i=0}^{L-2} i, & \text{如果 } l = L - 1 \end{cases}$

其中 $\leq e_{\text{scale}} \leq 1$ 是控制提前退出损失比例的超参数。注意，我们并未像 Elbayad et al.（2020）或 Schuster et al.（2022）那样为每层添加额外的 LM head，而是使用相同的 LM head。

提前退出损失课程（Curriculum）：我们发现，在所有迭代中为所有层添加提前退出损失会减慢训练速度并降低最终层的准确率。为此，我们引入课程函数 $C (t, l)$ ，并尝试了两种策略：

旋转式提前退出课程（Rotational Curriculum） $C_{\text{rot}, R}$ ：每 $R$ 层启用一次提前退出，并在每次迭代中循环轮换，从而每层每 $R$ 次迭代被启用一次。这样每次只进行 $\lceil L / R \rceil$ 次解嵌操作。
逐步式提前退出课程（Gradual Curriculum） $C_{\text{grad}}$ ：从第 $L - 1$ 层向第 0 层逐步启用，每隔 $T /2 L$ 次迭代启用一层。

总结一下，我们的训练策略包含以下超参数：

Layer Dropout：
- $p_{\text{max}}$ ：最后一层的最大丢弃率；
- $S (t)$ ：层丢弃课程，finetuning/持续预训练用 $S (t) = 1$ ，从头预训练用 $S(t)=S_{\text{exp}}(t)$ ；
Early Exit Loss：
- $e_{\text{scale}}$ ：控制早期层损失的缩放；
- $C (t, l)$ ：提前退出课程，选择 $C_{\text{rot}, R}(t, l)$ 或 $C_{\text{grad}}(t, l)$ ；
- $R$ ：在旋转课程中控制每多少层启用一次提前退出。

4.2 使用提前退出进行推理

在自回归解码时生成每个 token，我们只运行前 $E$ 层 Transformer，并跳转到模型的 LM head，即模型的最终输出为 $g(x_E)$ 。我们尝试了不同的 $E$ 值，并在“结果”部分提供准确率分析。
在这里插入图片描述

4.3 使用自我猜测解码（Self-Speculative Decoding）进行推理

我们通过引入层丢弃与提前退出损失，在训练时使得模型能够在推理时提前退出，从而加快生成速度，但这会带来准确率下降的问题。传统猜测式解码（Speculative Decoding，Leviathan 等人 2023；Chen 等人 2023）可以用一个较快但不那么精确的模型进行初步生成，再由完整模型验证，从而在不牺牲精度的情况下提升速度。但这通常需要两个模型，增加了内存和训练开销。

我们提出了一种自我猜测解码算法，在提前退出的基础上构建，仅使用一个模型，同时通过复用隐藏状态，在“草稿”和“验证”步骤中减少延迟。

在这里插入图片描述

该方法如图5所示，包含两个关键步骤：

自我草稿（Self-Drafting）：使用提前退出从模型中生成草稿 token；
自我验证（Self-Verification）：使用剩余层验证草稿 token 的预测。

为实现两阶段中的状态重用，我们设计了新的缓存复用机制，将 KV 缓存和退出层的查询信息合并。算法的详细说明见 §4.3.1 和 §4.3.2，伪代码见附录 A.4。

4.3.1 自我草稿

自我猜测解码的第一步是定义草稿 token 集 $D_0, ..., D_{d-1}$ 。我们通过提前退出计算前 $d$ 个草稿 token， $d$ 被称为“猜测数”。我们利用模型的前 $E$ 层进行自回归推理，生成草稿。

由于训练中采用了提前退出策略，模型本身就相当于在各个深度层训练了多个候选草稿子模型。我们可以在不同层退出，从而在延迟和准确率之间进行权衡。

4.3.2 自我验证

接下来是验证阶段。我们用完整模型对每个草稿 token 执行一次前向传播，预测下一个 token。我们比较草稿和验证结果的匹配情况，保留直到首次不匹配的所有草稿 token，再加上第一个验证 token，然后继续生成。

在我们提出的自我猜测解码中，验证阶段只需要计算草稿未使用的后 $L - E$ 层。为重用前 $E$ 层，我们在 KV 缓存上做了一些调整，见下文。

4.3.3 缓存复用（Reusing the Cache）

在自回归 Transformer 中，KV 缓存是高效生成的关键组件，它避免了每层重复计算。

由于草稿阶段使用了模型的前 $E$ 层，验证阶段使用后 $L - E$ 层，我们在两个阶段之间可以大量重用计算：

单一 KV 缓存：草稿和验证阶段都使用同一个模型且顺序一致，因此草稿阶段计算出的前 $E$ 层的 KV 缓存可以直接用于验证阶段，从而减少内存和延迟。
退出查询缓存（Exit Query Cache）：我们还引入“退出查询缓存”，保存第 $E - 1$ 层的查询向量，使验证阶段能直接从第 $E$ 层开始。这只需保存一个查询向量。我们将 KV 缓存与退出查询的组合称为 KVQ 缓存。