投机解码EAGLE精读

题目：EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty [ICML2024] 发表于24.02

链接：https://arxiv.org/abs/2401.15077

本文一作对该工作的讲解：智源社区活动

LLM推理是自回归的，这个自回归是从token层开始的自回归（也就是词元被转成嵌入的那一层）。而EAGLE将这一自回归延后到了特征层（也就是经过了多层transformer layer后）。通过在特征层的自回归，省去了transformer layer的时间。在特征层自回归生成几个token之后，把它们当做草稿序列进行校验。

动机

EAGLE基于一些现象：

• feature层面上的自回归比在token层面上的更简单。“特征”指的是原始LLM的倒数第二层特征，位于LM Head之前。与自然语言的简单变换形式的token序列相比，特征序列表现出更多的规律性。在特征层面上进行自回归处理，然后使用原始LLM的head派生令牌，比直接自回归预测令牌能获得更好的结果。

• 采样过程中固有的不确定性显著限制了预测下一个特征的性能。在文本生成中，目标LLM预测标记的分布并相应地进行采样，引入了随机性。由于特征是高维且连续的，不能以同样的方式处理。例如下一个token可能是“am”或“always”，不同的选择会导致不同的特征序列，这在特征级别的自回归中引入了歧义。为了解决这个问题，EAGLE将一个时间步前的标记序列输入到草稿模型中，这包括了采样结果。

原理

由于feature是已经经过transformer layers计算过的，包含了序列的语义信息，因此更具有规律性，用一个更小的模型就能够理解它。所以EAGLE就是训练一个小模型，让它进行特征自回归，以达到和token自回归近似的结果。小模型就是一个单层的decoder，用来预测下一个token的特征。原理图如下：

上图中，how can经过一次正常自回归得到I，以及中间结果how的嵌入和can的嵌入。其中，can的嵌入就包含了how can的整体语义信息，再加上将I嵌入得到I的单独信息，二者拼接就可以交给草稿模型进行自回归了。

can的特征与I的嵌入都是1*隐藏层维度的向量，拼接后长度是2倍的隐藏层维度，因此自回归头的第一层是FC层，将它降维成隐藏层维度。第二层就是一个解码层，也就是单层transformer，它的输出是对下一个token的特征的预测。

把对下一个token的特征的预测交给LM Head，就能得到下一个token词表上概率的分布。

这样来看，每次草稿模型自回归，只需要消耗大约1.5个transformer层的时间。

训练

训练草稿模型，也就是图中的auto-regression head，使用了smooth L1损失，也就是

y是真实值，x是目标值。在本文中，y是下一个token的实际特征向量，x是草稿模型对下一个token特征的预测值。

预测特征是草图模型的中间目标，最终目标是预测生成一系列标记的标记。因此，本文还使用分类损失来直接优化这个最终目标：

EAGLE的自回归头理想情况下是通过目标LLM自回归生成的文本进行训练的，然而这种方法成本高昂。幸运的是，EAGLE对训练数据的敏感性较低(后面的消融实验证明，目标LLM自回归生成的对加速比仅有0.1的提升)，因此没有使用目标LLM生成的文本，而是使用了一个固定的训练数据集，大大减少了开销。在草拟阶段，EAGLE自回归地处理特征。特征中的不准确可能会导致错误累积。为了解决这个问题，在训练期间通过向目标LLM的特征添加从均匀分布U(−0.1, 0.1)中采样的随机噪声来进行数据增强

代码层面上，与transformers库中的代码相比，EAGLE进行了两处改动：

• tree mask，也就是把不属于自己序列的token掩盖掉

• KV cache的分配，改成了预分配cache，这个改动不影响加速比

EAGLE比Medusa还要快1.6倍，在不同任务的测评中，只有总结任务比PLD方法慢一些，其他任务上都是当前最优的方法

EAGLE-2对tree attention进行了优化，也就是用一个beam-search的思路，和EAGLE相比有20%-40%的提升。EAGLE-2的加速比达到了3.5-4倍，在代码任务上能达到4倍。