LLM 推测性解码原理解读
大型语言模型一次生成一个token。要生成一个单词,模型会读取所有已生成的内容,运行一次完整的前向传播来选择下一个token,然后重复这个过程。听起来昂贵的部分是计算,但在实际硬件上通常并非如此。慢的部分是内存。每一步都必须从内存中取出模型的权重来使用,对于大型模型来说,这种移动而非算术决定了速度。
这里隐藏着一个特性。因为解码步骤受限于移动权重而非计算,在一次传播中评估多个候选token几乎与评估一个token花费相同的墙钟时间(Leviathan等人,arXiv 2211.17192,第1节;Chen等人,arXiv 2302.01318)。权重已经加载了。你不妨在它们可用时检查多个token。
推测解码将这个特性转化为加速,并且它以这个领域罕见的属性实现:它产生的文本在数学上与原始模型自己会生成的文本完全相同。你获得了速度,而没有牺牲任何一个token的质量。本文将介绍其工作原理,以及EAGLE-3——当前实现这一目标的最佳方式。
1、核心思想:草稿,然后检查
推测解码运行两个模型。一个快速的小型草稿模型和你实际想要输出的大型目标模型。一轮流程如下(Leviathan,第2节;Chen,算法2):
- 草稿。小型模型连续提出接下来的几个token。这很便宜,因为草稿模型很小。
- 一次验证。大型目标模型在提示词和所有草稿token上运行一次。在单次传播中,它在每个草稿位置同时生成它自己选择该token时会使用的分布。
- 接受或纠正。从左到右遍历草稿token,决定保留哪些。
第三步是魔法所在,所以值得慢慢来。在每个位置,草稿提出了一个token,目标模型现在告诉我们它认为该token的可能性有多大。将目标模型对该token的概率称为p,草稿模型的概率称为q。(对于阅读源论文的人有一个符号警告:两篇基础论文交换了这些字母。我使用p表示目标模型,q表示草稿模型,与Leviathan和EAGLE论文一致。)
规则是:如果目标模型对该token的喜爱程度至少与草稿模型一样,则直接接受草稿token;否则以概率p/q接受。当第一次拒绝token时,丢弃草稿的其余部分,并从剩余概率(即目标分布中草稿供给不足的部分)中重新采样单个替代token。如果所有草稿token都被接受,你甚至可以从已经付费的传播中免费获得一个额外token(Leviathan,第2.3节和算法1;Chen,"修改拒绝采样")。
因此,每次目标传播至少提交一个token,最多提交所有草稿token加一个,永远不会比普通解码提交的token少,通常会多几个。
2、为什么它不会改变输出
这是我真正觉得优雅的部分,也是标题说"免费"的原因。那个接受或纠正规则的设计使得最终发出任何给定token的概率恰好等于目标模型自己的概率。接受的质量加上重新采样的校正总和回到p,对于任何草稿模型都是如此。较弱的草稿只是更常被拒绝,这会损失速度,但永远不会损失正确性(Leviathan,附录A.1;Chen,定理1)。
这就是整个技巧。草稿模型只是一个猜测者。目标模型仍然是决定写什么的唯一权威,因为每个token要么是目标模型认可的,要么是从目标模型自己的分布中重新采样的。更好的草稿使其更快,更差的草稿使其更慢,但你从目标模型单独获得的输出不会改变。它适用于贪心解码和使用温度、top-k或nucleus的采样,通过在测试前调整两个分布。
一个公平的问题是"免费"实际上花费了什么,因为没有什么是真正免费的。你在计算和内存上付出代价:你运行第二个模型,并验证可能被丢弃的额外token。这之所以负担得起,正是因为解码一开始就是内存受限的,所以额外的并行检查在墙钟时间上几乎是免费的。诚实的警告是,当你批量处理更多请求时,收益会减少,因为繁忙的服务器不再空闲等待内存。EAGLE-3自己的服务数字清楚地显示了这一点:在批量大小为2时报告为1.81倍,但在批量大小为64时为1.38倍(arXiv 2503.01840,表3,在一个服务器设置上)。免费午餐是真实的,但当你受延迟限制时最大,这对于单个用户来说大多数时候都是如此。
3、困难的部分是草稿
只有当草稿既便宜又足够正确,目标模型才会继续接受其token时,加速才有意义。制作好的草稿是真正的研究问题,一系列称为多token预测的方法围绕它发展起来。
这种模式始于块级并行解码(Stern等人,arXiv 1811.03115),它引入了预测多个、一次验证、保留最长一致前缀的循环,此后所有方法都建立在此基础上。Medusa(arXiv 2401.10774)在模型上添加了几个额外的预测头,五个,并使用称为树注意力的技巧一起检查许多候选延续,在其简单形式中报告约2.2倍的加速,在其训练形式中报告2.3至3.6倍的加速。Meta的多token预测(Gloeckle等人,arXiv 2404.19737)在共享主干上训练并行头。DeepSeek-V3(arXiv 2412.19437)保留顺序预测模块,可以重新用作草稿器,实现约1.8倍更快的生成。EAGLE是推动草稿质量最远的工作系列,其第三代是本文的重点。
4、从EAGLE-1到EAGLE-3
EAGLE-1(arXiv 2401.15077)提出了一个很好的观察:不要在token级别进行草稿,而是在目标模型的内部特征级别进行草稿,即输出下方的隐藏状态。它将刚采样的token与该特征一起馈入,以便草稿知道实际选择了什么,并训练一个小的草稿头来预测下一个特征。
EAGLE-2(arXiv 2406.16858)保留了该草稿模型,但使其提出的候选集更智能。它不是固定形状,而是增长一个动态草稿树,使用草稿自己的置信度作为每个分支被接受可能性的替代,扩展有希望的分支并重新排序。
EAGLE-3(arXiv 2503.01840)是当前一代,它进行了三项改变,这是"如何实现"的核心:
- 它直接预测token,而不是匹配特征向量。EAGLE-1和2强制草稿重现目标模型的特征向量之一,该约束限制了额外训练数据的帮助程度。放弃它让草稿准确性随着更多数据的训练而持续提升,作者将其报告为主要的可扩展性结果(第3.1节)。
- 它融合目标模型的多个层,而不仅仅是顶层。草稿接收来自目标模型的低、中、高层隐藏状态的组合,而不是单一层,为小型草稿模型提供了目标模型思考的更丰富画面(第3.1节)。
- 它按照将要比赛的方式进行练习。在训练期间,草稿通过与推理时相同的多步骤循环运行,在树形注意力掩码下将其自己的输出反馈作为下一步的输入。作者称此为"训练时测试"。这意味着草稿学会从自己的错误中恢复,而不是只看到完美的输入(第3.2节)。
5、一次解码轮,端到端
这就是上面动画逐步显示的内容。
从当前文本开始,草稿模型增长一个小的可能延续树,而不是单行。EAGLE-3报告它可以运行比前代更深的树,深度为8,而EAGLE-2的深度为6,同时保持相同数量的节点,因为其更好的草稿赢得了额外深度(第4.1节)。树保留大约十个扩展节点和大约五十个候选token,根据目标模型调整大小(作者对7B或8B、13B和70B目标分别使用60、50和48)。
然后整个树被展平,并由目标模型在单次前向传播中评分。特殊的树注意力掩码让每个候选token只能看到自己的祖先和原始上下文,从不看其兄弟分支,因此一次传播就像它们是独立的一样评分每条从根到叶的路径(EAGLE-1,第3.1和3.3节;Medusa)。然后将之前的接受或纠正规则应用于分支。提交单个最长的一致路径,其余的折叠,并在第一个被拒绝的token处重新采样一个校正。然后循环从新的、更长的上下文重复。
每次目标传播提交的token数量是决定加速的数字。作者称之为平均接受长度。在相同设置上,跨EAGLE代际,它从约3.96上升到4.83再到6.62(表1),在8B类目标上,EAGLE-3报告Llama-3.1–8B的平均值为6.23,DeepSeek-R1-Distill-8B为5.84(表1)。每次传播更多token意味着更多权重加载摊销,这就是端到端加速的来源:在一个编码基准测试上比普通解码快6.47倍,报告比EAGLE-2高20%到40%(第4.1节和摘要)。这些是作者在其硬件上的数字,不是我的。
6、为什么这值得了解
推测解码是这个领域中少数几个不必在快和好之间做出选择的地方之一。输出是目标模型的,逐token如此。所有巧妙之处都进入草稿,错误只会损失一点速度。EAGLE-3的贡献是一个更常正确且随着更多训练数据不断改进的草稿,这使得平均接受长度上升,延迟下降。
如果你提供或运行本地模型,这是在转向更小的模型或更重的量化之前值得尝试的杠杆,因为与那些不同,它不会改变输出。诚实的边界:报告的数字是作者在其硬件上的,会随着你的草稿、模型和负载而变化;收益在单流、延迟受限使用时最大,在高批量处理下较小;适合你目标的草稿比任何标题乘数更重要。这些都不会损害核心思想,这是我不断回顾的部分。你可以让一个大型模型快几倍生成,并在纸上证明你没有改变它所说的任何内容。
原文链接: Predicting Possible Futures = Faster LLM decoding for free: how EAGLE-3 speeds up generation
汇智网翻译整理,转载请标明出处