简单好用！北大、普林斯顿联合提出即插即用的大语言模型加速方法

2024-03-27 04:03

©PaperWeekly 原创 · 作者 | 何震宇

单位 | 北京大学博士生

研究方向 | 大语言模型

最近，大语言模型（LLM）生成过程的加速技术，例如投机解码、Medusa（美杜莎）等，都带来了令人印象深刻的速度提升。这些方法通常依赖于将 LLM 与一个小型的草稿模型配对。小型的草稿模型试图在每个解码步骤中以更低的延迟预测多个草稿 token，并让 LLM 并行验证它们，从而减少了 LLM 的解码步数。

然而，获得高质量的草稿 token 仍然是一门艺术：它必须在 LM 更少量的参数和更强大的预测能力之间找到平衡，同时草稿模型必须匹配 LLM 的词汇表；此外，它还应该便于集成到用于部署 LLM 的分布式系统中。

为了应对这些挑战，Medusa 引入了一个有效的微调方法来获得草稿模型，它基于 LLM 本身微调了多个头（Medusa head）来预测草稿 token。然而，对额外微调的要求仍然有诸多不便。

这就引出了一个问题——能否设计一个即插即用的加速方法？一个无需训练或微调新模型即可提供快速生成草稿 token 的方法？

为了实现这个目标，作者引入了 REST: Retrieval-Based Speculative Decoding，一种基于检索的投机解码方法。REST 不使用额外的草稿模型，而是利用数据存储来根据部分输入来检索草稿 token。通过在 CPU 上进行检索，REST 避免了额外的 GPU 负载。由于不需要联合训练或微调，REST 可以立即加速任何预训练语言模型。

在这篇文章中，研究人员将探讨 REST 的灵感和内部工作机制，揭示它如何在没有额外训练模型的情况下实现令人印象深刻的速度提升，通过清晰的解释和生动的例子来展示 REST 的能力。

论文地址：

https://arxiv.org/pdf/2311.08252.pdf（camera-ready版本将在不久后更新）

收录会议：

NAACL 2024

博客地址：

https://sites.google.com/view/rest-llm

代码地址：

https://github.com/FasterDecoding/REST

提出REST的动机

首先，来观察一下 Python 代码的生成过程。例如，当模型生成 “for i” 时，熟悉 Python 的程序员可以高度预测接下来的几个令牌很可能是 “in range(”。

然而，当使用像 CodeLlama 这样的 LLM 时，它必须依次生成 "in"，"range"，和 "(" 这三个 token。这极大地降低代码生成的速度。那么有没有更高效的方法来一次性生成这些常见的短语呢？

答案是有的！这就是 REST 进入画面的地方。

REST: 基于检索的投机解码

如图所示，在推理过程中，语境被当做 query 来从数据存储中检索相关的语境，所有的后面紧跟着的序列，可以当做候选序列。

候选序列中每个的所有前缀都有可能充当语境后面的 token。然而，候选序列的数量可能是很庞大的，由于 GPU 显存大小和运算能力的限制无法一次性喂给 LLM 来验证。

因此，作者在候选序列的基础上建立了一个字典树（前缀树）Trie，Trie 中的每个结点到根节点就形成了候选序列中的一个前缀，而结点的权重即是该前缀在候选序列中出现的次数。

基于 Trie，可以轻易地选出出现次数最多的多个前缀，来当做草稿序列。然而，作者观察到这些草稿序列中，会经常出现重复的前缀，比如图中 in range 和 in range( 都有前缀 in range, in range( 和 in sorted 都有前缀 in，如果直接喂给大模型验证的话，会造成对相同前缀的重复计算，浪费了计算资源。

因此，研究人员观察到出现次数最多的前缀实际上是从权重最大的多个结点构成的子树（subtree）中获得的，通过对该子树进行广度优先遍历（BFS），可以得到一个伪句子（pseudo sequence），由字典树的特性可得草稿序列一定都是这个伪句子的子序列。通过专门的设计自注意力掩码，可以在 LLM 验证的时候还原伪句子中 token 原本的依赖关系（也称作树注意力）。

最后，如果草稿序列中有多个草稿 token 被 LLM 接受，那么 LLM 的推理速度就得到了提高。注意，这里作者采用的接受策略是当且仅当草稿 token 和 LLM 本应该生成的 token 相同时，才接受。因此通过 REST 来加速生成的内容和用原本 LLM 自回归生成的内容是完全一样的。

REST有多快？

在单张 A6000 GPU 上，研究人员对 CodeLlama 7B 和 13B 在 HumanEval 上进行测试，所用的 datastore大小为 27GB，取得了 2.12~2.36 倍的加速效果；对 Vicuna 7B 和 13B 在 MT-Bench 上进行了测试，所用的 datastore 大小为 12GB，取得了 1.62~1.77 倍的加速效果。