LLM推理加速新范式！推测解码（Speculative Decoding）最新综述

1. background
近年来，随着llm (large language model) 规模的逐渐增大（200m->7b->175b），llm的推理加速技术正逐步引起nlp学界的广泛关注。尤其是像chatgpt[1]，bard[2]这种线上实时交互的应用，llm的inference latency（推理耗时）极大程度地影响了用户的使用体验。那么，llm的latency主要来自哪里呢？
相关研究表明，llm推理主要是受内存带宽限制的（memory-bandwidth bound）[3][4]-- llm每个解码步所用的推理时间大部分并不是用于模型的前向计算，而是消耗在了将llm巨量的参数从gpu显存（high-bandwidth memory，hbm）迁移到高速缓存（cache）上（以进行运算操作）。也就是说，llm推理下的gpu并不是一个合格的打工人：他把每天大多数的时间都耗费在了早晚高峰堵车上，在公司没干啥实事儿（可不就是我摸鱼仙人:p）。
这个问题随着llm规模的增大愈发严重。并且，如下左图所示，目前llm常用的自回归解码（autoregressive decoding）在每个解码步只能生成一个token。这导致gpu计算资源利用率低下（->每个token的生成都需要重复读写llm的巨量参数），并且序列的生成时间随着序列长度的增加而线性增加。
图1: 自回归解码（左），推测解码（右）
2. speculative decoding
那么，如何更好地利用gpu资源，让它成为一个合格的打工人呢？相信大家心里已经有答案了：把公司当作家，减少通勤次数，就可以少摸鱼多打工了（泪目）。
推测解码（speculative decoding），作为2023年新兴的一项llm推理加速技术，正是提出了一种类似的解决方案：通过增加每个解码步llm计算的并行性，减少总的解码步数（即减少了llm参数的反复读写），从而实现推理加速。
如上右图所示，在每个解码步，推测解码首先高效地“推测”target llm（待加速的llm）未来多个解码步可能生成的token，然后再用target llm同时验证这些token。通过验证的token作为当前解码步的解码结果。如果“推测”足够准确，推测解码就可以在单个解码步并行生成多个token，从而实现llm推理加速。并且，使用target llm的验证过程可以在理论上保证解码结果和target llm自回归解码结果的完全一致[5][6]。
也就是说，推测解码在实现对target llm推理加速的同时，不损失llm的解码质量。这种优异的性质导致推测解码受到了学界和工业界的广泛关注，从2023年初至今涌现了许多优秀的研究工作和工程项目（如assisted generation[7]，medusa[8]，lookahead decoding[9]等等）。
考虑到推测解码领域2023年以来飞速的研究进展，我们撰写了一篇系统性的survey，给出推测解码的统一定义和通用算法，详细介绍了推测解码研究思路的演化，并对目前已有的研究工作进行了分类梳理。在下文中，我们将文章内容凝练为太长不看版——分享一些关于推测解码关键要素的看法，以及目前常用的研究思路，欢迎感兴趣的小伙伴一起讨论～
图2: 推测解码研究思路的演化
3. key facets of speculative decoding
首先，我们总结推测解码的定义：
推测解码是一种“先推测后验证” (draft-then-verify) 的解码算法：在每个解码步，该算法首先高效地“推测”target llm未来多个解码步的结果，然后用target llm同时进行验证，以加速推理。
也就是说，所有符合在每个解码步“高效推测->并行验证“模式的推理算法，都可以称为是推测解码（或其变体）。推测解码实现加速的关键要素，主要在于如下三点：
相比于生成单一token，llm并行计算额外引入的latency很小，甚至可以忽略；
“推测”的高效性&准确性：如何又快又准地“推测”llm未来多个解码步的生成结果；
“验证“策略的选择：如何在确保质量的同时，让尽可能多的“推测”token通过验证，提高解码并行性。
如上文所述，llm推理的主要latency瓶颈在于推理过程中参数的反复读写。在只考虑一个解码步的情况下，decoder-only llm的forward latency主要和decoder层数有关——层数越深，推理时间越长。相比于这两者，llm运算并行性带来的额外latency很小，这一点在非自回归解码的多个相关工作中有所讨论[10][11]。
因此，推测解码的算法设计主要考虑如下两点：“推测”（drafting）的高效性和准确性，以及“验证“策略（verification）的选择：
图3: 推测解码相关研究的归纳分类
4. “推测”的高效性和准确性
“推测“阶段（drafting）的目的是精准地“预测”llm未来多个解码步的生成结果，且不引入过多的latency。
因此，“推测”阶段的设计聚焦在“推测精度（accuracy）”和“推测耗时（latency）“的权衡上。一般来说，用以推测的模型越大，推测精度越高（即通过验证的token越多），但是推测阶段的耗时越大。如何在这两者之间达到权衡，使得推测解码总的加速比较高，是推测阶段主要关注的问题。
4.1 independent drafting
最简单的drafting思路是，拿一个跟target llm同系列的smaller lm进行“推测”[12][13]。比如opt-70b的加速可以用opt-125m进行推测，t5-xxl可以用t5-small。这样的好处是可以直接利用现有的模型资源，无需进行额外的训练。而且，由于同系列的模型使用相近的模型结构、分词方法、训练语料和训练流程，小模型本身就存在一定的和target llm之间的“行为相似性“（behavior alignment），适合用来作为高效的“推测“模型。
图4: https://huggingface.co/blog/assisted-generation
这一思路由google和deepmind同时提出[12][13]。作为speculative decoding的早期探索，这种“推测”思路易于实践和部署。并且，这两篇工作同时在理论上证明了推测解码不仅支持greedy decoding，还支持nucleus sampling的无损加速（我们下文会讲到）。这两种解码策略涵盖了llm应用的大多数场景。因此，这两篇工作极大地促进推测解码在llm推理加速中的应用，吸引了工业界和学术界的大量关注。
然而，同系列小模型的“推测”精度还有提升空间吗？
显然是有的。最直接的思路，就是去增强小模型和大模型之间的“行为相似性”（behavior alignment），让小模型模仿得“更像”一些。目前在这方面的研究进展集中在知识蒸馏（knowledge distillation）上：将target llm作为教师模型，小模型作为学生模型，通过知识蒸馏让小模型更加趋向于target llm的预测行为[14][15]。并且，知识蒸馏还可以有效地增强小模型的生成质量，通过减少低级的预测错误，增加通过验证的token数量。
4.2 self-drafting
然而，采用一个独立的“推测”模型也有缺点：
首先，并不是所有的llm都能找到现成的小模型，比如llama-7b。重新训练一个小模型需要较多的额外投入。
另外，引入一个额外的小模型增加了推理过程的计算复杂度，尤其不利于分布式部署场景。
因此，相关研究工作提出利用target llm自己进行“高效推测”。比如blockwise decoding[5]和medusa[8]在target llm最后一层decoder layer之上引入了多个额外的ffn heads（如下所示），使得模型可以在每个解码步并行生成多个token，作为“推测”结果。
图5: https://sites.google.com/view/medusa-llm
然而，这些ffn heads依然需要进行额外的训练。除了这两个工作，还有一些研究提出利用early-existing或者layer-skipping来进行“高效推测“[16][17]，甚至仅仅是在模型输入的最后插入多个[pad] token，从而实现并行的“推测”[18]。然而，“部署的便捷性”和“推测精度”之间依然存在一定的权衡关系。如何选择合适的“推测”策略，达到令人满意的加速效果，就见仁见智了。
感兴趣的友友可以移步具体论文查看细节，我们后续也准备提供一个公平的加速评测，给大家提供一个参考～
5. 验证策略的选择
“验证“阶段（verification）的首要目的是保证解码结果的质量。
让我们重新回顾推测解码的验证过程：
如下图所示，在给定“草稿”（即推测结果）时，llm的并行验证其实和训练阶段teacher-forcing的形式是一致的——在生成每个token时，都假设llm的前缀输入是正确的。比如，在验证第三个“推测”token时，llm以绿色前缀和两个黄色的推测“token作为前缀输入。以贪婪解码（greedy decoding）为例，以该前缀作为输入时，llm会自己生成一个概率最大的token。如果这个token（绿色）和第三个“推测”token相同，就说明第三个“推测”token通过了“验证”——这个token本来就是llm自己会生成的结果。
因此，第一个没有通过验证的“推测”token （图中的红色token）后续的“推测”token都将被丢弃。因为这个红色token不是llm自己会生成的结果，那么前缀正确性假设就被打破，这些后续token的验证都无法保证前缀输入是“正确”的了。
图6：recap of speculative decoding
由此可见，推测解码是可以保证最终解码结果和target llm原先的贪婪解码结果完全一致的。因此，贪婪解码经常被用于推测解码的demo展示[8]，用以清晰直观地表示推测解码在保持和target llm解码结果等同的前提下，实现了数倍的推理加速。
然而，严格要求和target llm解码结果完全匹配（exact-match）是最好的策略吗？
显然，并不是所有概率最大的token都是最合适的解码结果（比如beam search）。当推测模型的性能较好时，严格要求和target llm结果匹配会导致大量高质量的“推测”token被丢弃，仅仅是因为它们和target llm top-1解码结果不一致。这导致通过验证的“推测”token数量较小，从而影响推测解码的加速比。
因此，有一些工作提出可以适当地放松“验证”要求，使得更多高质量的“推测”token被接受，增大每个解码步通过验证的“推测”token数量，进一步提升加速比[12][14][15]。
除了支持贪婪解码，推测解码还可以在理论上保障和target llm nucleus sampling的分布相同[12][13]，具体证明感兴趣的朋友可以查看相关paper～。另外，相比于只验证单一的“推测”序列，相关研究还提出可以让target llm并行验证多条“推测”序列，从而进一步增大通过验证的“推测”token数量[19]。
6. 总结
表1: 推测解码算法总结
在上表中，我们给出目前常用的推测解码算法的总结～。作为一种新兴的推理加速算法，推测解码在实现对target llm推理加速的同时保障了解码结果的质量，具有广阔的应用前景和极大的科研潜力，个人比较看好～。然而，推测解码研究本身也存在许多尚未解答的问题，比如如何更好地实现target llm和“推测”模型之间的行为对齐、如何结合具体任务的特点设计相应的推测解码策略（比如多模态模型加速），都是值得思考的问题。