EAGLE-1 解读

Paper：EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty

Publish：ICML 2024，Arxiv 2024.1.26

在 LLM 推理加速领域，有一类常见的方法是 Speculative Decoding，也就是投机解码。大致思想就是先用小模型一次性预测一些 Token，然后大模型再去并行验证，进而提升推理速度。在 Token-Level 的预测有一个问题，就是准确率低下。

EAGEL 的想法是，既然在 Token Level 预测效果一般，为什么不换到 Feature Level 呢？

加速流程细节

假设大模型输入的是 “How can”, 这个时候需要生成后面的 Token。

在 LLM 的倒数第二层（Transformer Block），经过 KQV 矩阵运算，得到了 “can” 位置的特征向量 。然后分成两路：

• 特征向量 进入大模型的全连接层，采样得到下一个真实的词 。
• 和 一起被送入Draft Model。

然后 Draft Model 将 和 进行拼接，然后通过轻量级自回归头，预测下一个词的特征向量 。接着这个 进入原始大模型的 LM Head, 采样下一个草稿词 ，然后继续循环。

image.png

Draft Model 在每一步采样时，会采样多个，因此最后会形成一个 Token Tree. 然后主模型把这个Token Tree 输入，跑一次完整的前向传播，能够计算所有草稿词的概率。

如何训练?

使用现有高质量对话数据集 ShareGPT，输入模型，记录倒数第二层的特征向量 和对应的Token。然后给 Draft Model 输入第 步的特征 和下一个 Token , 让 Draft Model 去预测下一步模型的隐藏状态向量。使用了两个损失：

• 回归损失：，让小模型的输出向量逼近真实向量
• 分类损失：，确保预测的特征向量经过 LM Head 映射后，得到的 Token 也是正确的。

实验结果如何？

image.png

值得注意的是，7B 模型已经很小，需要一个更小的模型来充当 Draft Model，但是很难找到，所以标记为 N/A。