回顾推测解码

lixia01

推测解码已被证明是一种有效的技术，可以在不影响质量的情况下从 LLM 进行更快、更便宜的推理。它还被证明是一系列优化技术的有效范例。
大型语言模型 (LLM) 是近年来人工智能 (AI) 快速发展的核心。虽然具有开创性，但面向用户的产品面临的一个挑战是，由于规模庞大，这些大型模型的推理速度（即输出生成）很慢，这可能会导致用户体验不佳。
2022 年，我们发表了《通过推测解码从 Transformers 快速推理》，其中介绍了一种称为推测解码的技术，可以显著减少 LLM 的推理时间。该算法通过并行计算多个 token 来加速自回归模型的生成，而不会影响输出质量；事实上，该方法保证了相同的输出分布。使用相同的硬件更快地产生结果也意味着需要更少的机器来处理相同数量的流量，这又意味着服务相同模型的能源成本降低。
今天，我们回顾一下该方法的工作原理，探索其从推测执行中得到的启发，并重点介绍一些新的富有洞察力的发展。在我们最初的论文中，我们通过将其应用于翻译和摘要任务证明了这种方法的有效性，我们看到了约 2 到 3 倍的改进。从那时起，我们已经在许多 Google 产品中应用了推测解码，我们看到推理速度显著加快，同时保持了相同的响应质量。我们还看到整个行业都采用了推测解码，并见证了许多使用推测解码范式的富有洞察力和有效的应用和想法。
背景
LLM 每次生成一个 token，其中 token 通常是一个单词或单词的一部分。例如，使用常见的 token 生成器，“一个人的一小步，人类的一大步”这句话由 12 个 token 组成。这意味着要生成这句话，LLM 必须运行 12 次。每次这样的运行称为解码步骤。
LLM 越大，它就越有能力。然而，这些较大的模型也较慢，因为例如，每个解码步骤都需要读取模型的全部权重。这可能意味着模型需要为其生成的每个单词读取大约 1TB 的数据！
通常，LLM 的目标是生成许多单词，例如对话响应或文档摘要，并且由于每个标记都依赖于之前生成的标记，因此必须逐个生成它们，一遍又一遍地读取模型的所有权重。两个关键观察结果激发了我们的推测解码方法。
观察 1：有些代币比其他代币更容易生成
并非所有的 token 都一样：有些更难生成，有些则更容易生成。请考虑以下文本：
7的平方根是多少？ 7的平方根 是2.646。
生成强调的标记“ 7 ”相对容易；例如，我们可以注意到前面的标记“平方根”之前发生过，只需复制后面的标记即可。生成标记“ 2.646 ”更难；模型需要计算或记住答案。
这一观察结果表明，大型模型更好一些，因为在困难情况下（例如“ 2.646 ”）表现更好，但在许多简单情况下（例如“ 7 ”），小型模型可能为大型模型提供合理的近似值。
观察 2：LLM 推理的瓶颈通常是内存
机器学习硬件种类，TPU和GPU ，是高度并行的机器，通常每秒能够执行数百万亿次操作，而它们的内存带宽通常仅为每秒数万亿字节 - 低几个数量级。这意味着当使用现代硬件时，我们通常可以对从内存中读取的每个字节执行数百次操作。
相比之下，现代 LLM 所依赖的Transformer架构通常仅对推理过程中读取的每个字节执行少量操作，这意味着在现代硬件上从 LLM 生成输出时有充足的备用计算资源可用。
推测解码-2-操作请求
推测执行
基于在串行计算 token 的同时还有额外的并行计算资源可用的预期，我们的方法旨在通过并行计算多个 token 来提高并发性。该方法的灵感来自推测执行，这是一种优化技术，通过该技术，任务在验证是否真正需要之前或与验证是否真正需要的过程并行执行，从而提高并发性。推测执行的一个著名示例是现代流水线 CPU 中的分支预测。
为了使推测执行有效，我们需要一种有效的机制，可以建议可能需要执行的任务。更一般地，考虑推测执行的这个抽象设置，假设f ( X ) 和g ( Y ) 是冗长的操作：
Y = f（X）
Z = g（Y）
慢速函数f ( X ) 计算Y ，它是慢速函数g ( Y )的输入。在上述设置中，f ( X ) 和g ( Y ) 是同一函数。如果没有推测执行，我们需要串行评估它们。推测执行表明，给定任何快速近似函数f *( X )，我们可以在评估g ( f *( X ))的同时评估第一个慢速操作f ( X )。一旦f ( X ) 完成并且我们获得了正确的Y值，我们就可以检查快速近似f *( X )的输出是否也是Y，在这种情况下我们设法增加了并行化。如果f *( X ) 输出不同的值，我们可以简单地放弃g ( f *( X ))的计算并恢复到像串行情况一样计算g ( Y )。f * ( X )越有效，即它输出与f ( X )相同的值的可能性越高，就越有可能增加并发性。无论哪种方式，我们都能保证相同的输出。
推测性抽样
我们提出了一种将推测执行推广到随机设置的方法，即，需要以某个概率执行任务。我们将这种推广称为推测采样。考虑以下设置，与上面的设置相同，不同之处在于f ( X ) 现在输出一个概率分布，我们从中对函数g , Y 的输入进行采样：
Y ~ f ( X )
Z = g（Y）
与上面类似，给定任何快速近似值f *（X），这次输出一个概率分布，推测抽样允许我们在执行g（sample（f *（X ））的同时执行f ( X ) 。我们可以使用标准推测执行，并在样本不匹配的情况下放弃计算，但这样做效率低下。实际上，考虑一个例子，其中f ( X ) 和f *（X）总是输出从 1 到 100 的均匀概率分布。推测执行每 100 次只会接受一次f * 的猜测。这显然效率低下——f和f * 是相同的函数，我们总是可以接受来自f * 的样本！相反，推测抽样提供了一种根据f ( X ) 和f *（X ）概率地接受或放弃f * 的猜测的方法，从而保证最优性以及相同的输出分布。推测抽样可能在其他环境中很有用，例如在强化学习或物理模拟中（详情请参阅论文）。
推测解码
LLM 不会生成单个下一个标记，而是生成一个概率分布，从中我们可以对下一个标记进行采样（例如，在文本“最著名的电影导演是”之后，LLM 可能会以 70% 的概率生成标记“Steven”，以 30% 的概率生成标记“Quentin”）。这意味着直接应用推测执行来生成 LLM 的输出效率非常低。推测解码利用推测采样来克服这个问题。有了它，我们可以保证尽管成本较低，但生成的样本与朴素解码产生的样本来自完全相同的概率分布。请注意，在贪婪解码的特殊情况下，我们总是对单个最可能的标记进行采样，推测执行可以有效地应用于 LLM 推理，正如我们在前作中所展示的那样。
推测解码是将推测采样应用于自回归模型（如 Transformer）的推理。在这种情况下，f ( X ) 和g ( Y ) 将是同一个函数，以序列为输入，输出由一个 token 扩展的序列的分布。因此，推测解码使我们能够高效地并行计算 token 及其后面的 token，同时保持相同的分布（请注意，推测解码可以并行生成两个以上的 token，请参阅论文）。
要应用推测解码，剩下的就是对解码函数进行快速近似。上面的观察 1 表明，小模型可能在许多较简单的标记上表现良好。事实上，我们在论文中表明，使用现有的现成的较小模型或简单的启发式方法在实践中效果很好。例如，当应用推测解码来加速翻译任务的 11B 参数T5-XXL 模型，并使用较小的 60M 参数 T5-small 作为猜测机制时，我们的速度提高了约 3 倍。
迈向更高效的人工智能
我们看到推测解码已在整个行业中得到采用，并且报告了一些显著的性能提升。这种广泛采用得益于使用推测解码范式的众多富有洞察力和有效技术，这些技术通常与其他新方法结合使用。例如，在分布式设置中展示大规模模型的有效性，使用 多个草稿猜测而不是一个，将目标模型中的知识提炼到草稿模型中，让草稿模型使用目标模型的一部分，对草稿和目标使用单个模型，或一起验证所有草稿标记。该方法还应用于图像和语音生成等领域。
随着 LLM 的使用越来越广泛，对更高效推理的需求也变得越来越重要。我们期待看到更多利用推测解码和其他现有范式的想法，以及全新的方法。
致谢
这项工作是与 Asaf Aharoni、Avinatan Hassidim 和 Danny Vainstein 密切合作的成果。此外，我们还要衷心感谢您对 YaGuang Li、Blake Hechtman、Tao Wang、Toby Boyd、Nathan Lintz、Phil Chen、Nir Shabat、Jayant Madhavan、Aliaksei Severyn、Jakub Adamek、Jonathan Mallinson、Zhifeng Chen、Yoel Drori、Mariano Schain、Charlie Chen、Noam Velan、Nitish Kulkarni 的评论、帮助、富有洞察力的讨论、宝贵的反馈和支持。西德哈斯·穆德加尔、萨沙·戈德斯坦、纳达夫·谢尔曼、皮拉尔·曼琼、费尔南多·佩雷拉、埃亚尔·塞加利斯、埃亚尔·莫拉德、丹尼·瓦列夫斯基、丹尼尔·鲁门、瓦莱丽·尼加德、史蒂夫·贝克、斯里尼瓦桑（奇努）·文卡塔查里、赫玛·布达拉朱、孙子腾、阿南达·瑟莎·苏雷什、伊丽莎白·哈蒙·里德、杰夫·迪恩、普拉巴卡Raghavan、James Manyika 和团队在 Google Research、Google Deepmind 和 Google Search 中。