SpecLA: Efficient Speculative Decoding for Linear-Attention Models
本文介绍了 SpecLA，一种专为有状态线性注意力模型设计的、高效的投机解码运行时，它通过利用拓扑感知验证、紧凑的状态恢复以及目标对齐的草拟器，实现了相比标准自回归解码高达 1.70 倍的端到端加速。
原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。 阅读完整免责声明
想象一下，你正试图写一个故事，但你有一个非常严苛的编辑，他只允许你一次写一个词。在你写下一个词之前，你必须停下来，检查你目前为止写下的所有笔记，更新你的心理状态，然后写下下一个词。这正是许多强大的 AI 模型目前的工作方式——它们一次生成一个标记（token，即单词的一部分），不断地在快速内存和主存储器之间来回传输大量数据。这个过程很慢，就像一个图书管理员，每当你说话时，都必须走到图书馆后面去查阅一本书。
为了提高速度，科学家们发明了一个技巧，叫做“投机采样”（speculative decoding）。想象一下，不再是一个词一个词地写，而是有一个速度更快、但稍微没那么聪明的助手，为你预判接下来的几个词。然后，你要求那位严苛的编辑一次性检查所有这些预判。如果编辑同意了，你就能在通常写一个词的时间内写出好几个词。这种方法在标准的 AI 模型（Transformer）中效果很好，因为它们保留着一份易于编辑的过往单词列表。但有一种新的、更快的 AI 模型类型（称为线性注意力模型/Linear-Attention），它并不保留列表，而是保留一个单一的、稠密的“摘要状态”（summary state），这个状态随着每个新单词的加入而不断变化。旧的这种预判技巧在这里行不通，因为你无法在不重写整个内容的情况下，从摘要状态中“擦除”一个错误的预判。这就是研究人员试图解决的谜题：如何在不破坏这些新模型独特的记忆方式的前提下，实现预判多个单词的速度？
于是，SpecLA 诞生了，这是一个专门为这些具有状态性的线性注意力模型设计的系统，旨在让这种“预判 ahead”的技巧发挥作用。研究人员发现，仅仅试图将旧的猜测方法强加给这些新模型是会惨败的。如果你尝试逐个检查猜测，你会失去速度优势，因为你仍然需要为每一个猜测都进行一次沉重的摘要状态搬运。如果你尝试像处理批处理一样一次性检查它们，数学计算就会变得复杂且缓慢，因为猜测可能会向不同的方向分支，而模型的记忆并不擅长处理分支。
因此，团队构建了 SpecLA，它充当了这些 AI 模型的智能交通控制器。SpecLA 不会将每一次猜测视为一次独立的行程，而是观察猜测的形状。如果猜测形成一条直线，它会让模型的记忆状态保持在快速处理器芯片上，从而避免前往主存储器的缓慢移动。如果猜测像树木一样分支，它会使用一种特殊的“树状掩码”（tree mask）来同时检查它们，而不会混淆不同的路径。最重要的是，当严苛的编辑对某些猜测说“是”而对另一些说“否”时，SpecLA 不会浪费时间重写整个记忆状态。相反，它会保存处理过程中变化的微小、紧凑的“收据”（称为因子/factors）。一旦决策做出，它就能利用这些收据瞬间更新记忆状态，完全跳过了繁重的计算工作。
结果令人振奋。在使用了名为 GDN-1.3B 的公开模型并在强大的 NVIDIA H100 计算机上进行测试时，SpecLA 使 AI 写作速度比标准的逐词生成法快了高达 1.70 倍。在某些测试中，它快了 1.42 倍，而在另一些测试中则快了 1.06 倍。研究人员还进行了较小的测试，以了解其表现如此出色的原因。他们发现，他们这种检查树状猜测的新“混合”方式，比旧的逐个重放猜测的方式快了 1.80 到 7.11 倍。他们还发现，使用这些紧凑的“收据”来更新记忆，比重放单词的速度快了 2.74 到 4.28 倍，而将最终更新延迟到下一步则又节省了 1.15 到 1.44 倍 的时间。
然而，论文谨慎地指出，这种加速效果取决于“猜测助手”的表现。如果助手做出了太多的错误猜测，系统会把时间花在检查这些猜测上，但最后却被拒绝了，这样速度优势就会消失。研究人员表明，为了让系统良好运行，助手需要足够准确，使得至少有 70% 到 80% 的猜测被接受。如果助手是完美的，理论上速度可以更高，但对于现实世界的助手来说，收益是稳健的，但取决于猜测的质量。论文并未声称这解决了所有 AI 模型的所有问题，但它证明了对于这种特定类型的状态模型，一种量身定制的新方法是必要且有效的。
您所在领域的论文太多了？
获取与您研究关键词匹配的最新论文每日摘要——附技术摘要，使用您的语言。