打破词表壁垒: 如何利用任意起草模型加速 LLM

大语言模型 (LLM) 的推理速度仍然是部署生成式 AI 的主要瓶颈之一。无论是运行聊天机器人、代码助手还是摘要工具，逐个生成 token 的成本和延迟都令人望而却步。

为了解决这个问题，社区已广泛采用投机采样 (Speculative Decoding，SD) 。 这项技术利用一个更小、更快的“起草 (drafter) ”模型来猜测即将出现的 token，然后由更大的“目标 (target) ”模型并行验证。一旦奏效，效果就像魔术一样: 你可以在保持相同输出质量的同时显著提高速度。

然而，投机采样一直存在一个主要的“陷阱”: 词表约束 (The Vocabulary Constraint) 。

传统上，起草模型和目标模型必须使用完全相同的语言——字面意义上的。它们需要共享相同的分词器 (tokenizer) 和词表 (vocabulary) 。这意味着你无法使用 Qwen 系列的高效模型来加速 Llama 模型，也无法用 Phi 模型来加速 Mixtral 模型。如果词表不匹配，标准的 SD 算法就会失效。

在论文 Accelerating LLM Inference with Lossless Speculative Decoding Algorithms for Heterogeneous Vocabularies 中，来自魏茨曼科学研究所、英特尔实验室和 d-Matrix 的研究人员提出了一套突破性的算法来消除这一约束。他们的工作使得任何现成的模型都可以作为任何目标模型的起草者，无论它们的词表是否相同。

本文将剖析这些算法的工作原理，解释为什么处理“异构 (heterogeneous) ”词表如此困难，以及这些新方法如何实现高达 2.8 倍的加速。

背景: 为什么词表很重要

要理解这项创新，我们首先需要了解标准投机采样存在的问题。

LLM 并不处理原始文本；它们处理的是被称为 Token ID 的数字。分词器根据特定的词表将文本转换为这些 ID 的序列。

  • 模型 A (例如 Llama-3) 可能将单词 “apple” 映射为 ID 5043。
  • 模型 B (例如 Qwen-2) 可能将单词 “apple” 映射为 ID 8912。

标准投机采样的工作方式是直接将 Token ID 从起草模型传递给目标模型。如果你尝试使用模型 B 为模型 A 起草，模型 B 发送 ID 8912。模型 A 收到 8912 后，在其自己的字典中查找，结果发现它对应的是完全不同的内容 (或者根本不存在) 。验证随即失败，投机过程也就白费了。

也就是直到现在，唯一的解决方案是从头开始训练一个特定的起草模型来匹配目标的词表——这是一个昂贵且耗时的过程。研究人员提出了三种算法来弥合这一差距，实现了异构投机采样 (heterogeneous speculative decoding) 。

我们将探讨三种提出的方法:

  1. SLEM: 一种稳健的、基于文本的对齐方法。
  2. TLI: 一种对 token 分布的数学调整方法。
  3. SLRS: 一种基于字符串级概率的理论方法。

方法 1: 字符串级精确匹配 (SLEM)

引入的第一个也是最稳健的算法是 算法 2: 字符串级精确匹配 (SLEM) 。 该方法旨在成为 LLM 推理的“万能适配器”。

纯文本的逻辑

由于两个模型之间的 Token ID 不兼容，SLEM 使用了它们都能理解的一种媒介: 纯文本 。

过程不再是传递 Token ID，而是如下所示:

  1. 起草 (Drafting) : 起草模型使用自己的词表生成一系列 token。
  2. 解码 (Decoding) : 这些起草 token 立即被解码为原始文本字符串。
  3. 重新分词 (Re-tokenizing) : 使用目标模型的分词器对该原始文本字符串进行编码。
  4. 验证 (Verification) : 目标模型针对自己的预测来验证这些新 token。

这听起来很简单，但它解决了根本的不兼容问题。起草模型可以用自己的 token 进行“思考”，而目标模型则接收自己的 token 作为输入。

非单射分词器的挑战

研究人员遇到了一个微妙但关键的工程障碍: 非单射分词器 (Non-Injective Tokenizers) 。

理想情况下，分词应该是可逆的。如果你将文本编码为 token 再解码回来，应该得到完全相同的文本。然而，许多现代分词器应用了“规范化”规则。例如:

  • 将多个空格替换为单个空格。
  • 文本小写化。
  • 规范化重音字符 (例如 ‘é’ 变为 ’e’) 。

在异构设置中，这会产生不匹配。起草模型生成文本，但当该文本为目标模型重新分词时，规范化规则可能会对其进行细微修改。如果算法只是盲目地比较生成的 Token ID，它们可能会拒绝有效的草稿，原因仅仅是分词器本身导致的细微格式差异，而不是模型的逻辑错误。

解决方案: SLEM 实现了一种复杂的匹配机制。它双向转换 token，并在已接受的目标 token 和提议的起草 token 之间搜索最长的匹配文本段。它有效地在文本空间中对齐序列，减轻了不同分词规则引入的噪声。

这种稳健性使得 SLEM 在实践中非常有效，并成为了 Hugging Face Transformers 库中异构 SD 的默认算法。

方法 2: Token 级交集 (TLI)

虽然 SLEM 依赖于将 token 转换为文本再转换回来, 算法 4: Token 级交集 (TLI) 则严格停留在概率分布领域。

词表的交集

即使两个模型拥有不同的词表，它们通常也共享很大一部分 token (如常用词、标点符号和字母字符) 。

请看下表 (论文中的表 8) ，它显示了流行模型的词表大小。虽然大小差异巨大，但它们很少是不相交的。

图表 8: 广泛使用的目标模型和起草模型的词表大小。

标准 SD 算法之所以失败，是因为起草模型可能会提议一个目标模型词表中根本不存在的 token。

TLI 策略

TLI 通过修改起草模型的概率分布来改进这一点。它强制起草模型在涉及目标模型时“安分守己”。

  1. 识别交集 : 在两个词表中都存在的 token 集合。
  2. 重新归一化 : 算法获取起草模型的概率分布，并将交集之外的任何 token 的概率设为零。然后，它按比例放大剩余 token 的概率，使其总和为 1。

这样一来，起草模型在数学上就被阻止了建议目标模型根本无法生成的 token。

研究人员证明 (定理 4.1) ，TLI 保证了比朴素的“并集”策略 (即简单地将词表视为合并) 更高的预期接受率。通过将起草模型的精力集中在有效的候选者上，TLI 减少了浪费并提高了速度。

方法 3: 字符串级拒绝采样 (SLRS)

第三种方法, 算法 3 (SLRS) , 在理论上是最具雄心的。标准 SD 执行逐 token 验证，而 SLRS 试图执行字符串级的拒绝采样 。

概念

在标准的 token 级验证中，我们检查 \(P_{target}(token) \geq P_{draft}(token)\)。 在 SLRS 中，研究人员试图检查 \(P_{target}(string) \geq P_{draft}(string)\)。

这是一个强有力的想法，因为一个“字符串” (比如一个单词) 在目标模型中可能表示为一个 token，而在起草模型中可能表示为三个 token。在这里进行逐 token 比较是不可能的，但比较整个字符串出现的概率在理论上应该是可行的。

计算瓶颈

虽然在数学上是合理的并且是“无损的” (保留了精确的目标分布) ，但 SLRS 面临一个巨大的障碍: 计算字符串的概率非常昂贵。

要知道起草模型生成字符串 “hello” 的概率，你必须将起草模型构建 “hello” 的所有可能方式的概率相加。

研究人员使用 Qwen2-7B-Instruct 词表展示了这种复杂性。如下面的图 1 所示，即使像 “hello” 这样简单的单词也可以由 14 种不同的 token 组合形成 (例如 “h”+“ello”, “hel”+“lo”, “hello”) 。

图 1: 左: 来自 Qwen2-7B-Instruct 词表的 14 种有效 token 组合，它们连接后可形成字符串 ‘hello’。右: 所有这些组合的树状可视化。14 个对勾标记表示有效的组合，即可视化树中的叶子节点。在此示例中，计算算法 3 中的 \\(\\psi(t)\\) 需要起草模型进行 16 次前向传递，这是树中非叶子节点的数量加一。

为了准确计算拒绝采样标准 (\(\psi(t)\)) ，算法必须计算这棵树的每个分支的概率。对于 “hello”，这需要起草模型进行 16 次前向传递——仅仅是为了验证一个单词！

随着 token 变长，这种复杂性呈指数级爆炸。下图 (图 2) 显示了 token 长度与组合数量之间的关系。

图 2: 来自 Qwen2-7B-Instruct 词表的 150,000 个选定 token 在不同 token 长度下的组合数量。我们可以看到组合数量随 token 长度呈指数级增长。

由于这种“组合爆炸”，论文得出的结论是，虽然 SLRS 是一个重要的理论贡献，但它可能仅对词表非常小 (如字节级模型) 的起草模型实用，因为那里的组合数量仍然在可控范围内。

算法比较

论文提供了这些方法之间预期接受率的有用比较。

  • Alg 5 (标准 SD) : 对于异构词表未定义 (失败) 。
  • Alg 1 (朴素并集) : 可行，但次优。
  • Alg 2 (SLEM) : 使用精确匹配。兼容性高，接受率中等。
  • Alg 4 (TLI) : 数学交集。接受率高于 Alg 1。
  • Alg 3 (SLRS) : 字符串验证的理论上限，但计算量大。

表 3: 给定异构词表时所有投机方法的预期接受率。如定理 4.1 证明，算法 1 的预期接受率总是小于或等于算法 4 的预期接受率。

实证结果: 真的有效吗？

研究人员在高端硬件 (NVIDIA H100) 和消费级硬件 (RTX 6000) 上，使用流行的数据集如 CNN/DailyMail (摘要) 、Humaneval (编码) 和 SCROLLS (长上下文) 测试了这些算法。

SLEM 性能

SLEM (算法 2) 在各项测试中表现最为稳健。在不存在同系列起草模型的情况下，SLEM 解锁了显著的加速。

例如，使用 CodeLlama-13b 作为目标，使用 Tiny Starcoder (完全不同的系列) 作为起草者，他们在编码任务上实现了 2.79 倍的加速 。

表 1: 广泛使用的模型、任务和硬件设置下算法 2 (SLEM) 与自回归解码 (AR) 的基准对比。结果表明，SLEM 将吞吐量比 AR 提高了高达 2.8 倍。

该表突出了该方法的灵活性。你可以将 Mixtral 模型与 Qwen 起草者或 Vicuna 起草者配对。这种混合搭配的能力允许工程师为特定的硬件找到绝对最快的起草者，而不必局限于目标模型的“小弟”。

TLI 性能

TLI (算法 4) 也表现出了实实在在的改进，虽然在峰值场景下通常略低于 SLEM，但它仍然是一种非常高效的 token 级方法。

表 2: 广泛使用的模型、任务和硬件设置下算法 4 (TLI) 与自回归解码 (AR) 的基准对比。结果表明，TLI 将吞吐量比 AR 提高了高达 1.7 倍。

关键是，TLI 在标准基准测试中实现了高达 1.7 倍的加速 。 当字符串级解码的开销不可取，且词表重叠足够时，它可以作为一种轻量级的替代方案。

对 AI 推理未来的启示

这篇论文的贡献不仅仅是更快的基准测试数据。通过解决异构词表问题，作者从根本上改变了 LLM 推理的经济性。

1. 速度的民主化: 以前，只有有资源训练自定义起草模型的组织才能充分优化其推理栈。现在，任何人都可以下载最先进的目标模型 (如 Llama-3) 并将其与 Hugging Face Hub 上任何可用的轻量级模型配对，从而获得即时加速。

2. 开箱即用的效率: SLEM 和 TLI 集成到 Hugging Face transformers 库中意味着公众现在就可以利用这些成果。论文指出，Hugging Face 的维护者独立评估了这些方法，并因其有效性将 SLEM 设为异构投机的默认算法。

3. 模型架构独立性: 我们不再受限于“模型家族”。推理模型可以由编码模型起草。多语言模型可以由专门的英语模型起草。这为寻找特定任务的“完美搭档”开辟了新的研究领域。

结论

标准投机采样是 LLM 延迟优化的一次飞跃，但它被词表匹配的要求所束缚。Timor 等人的这项工作打破了这些束缚。

通过 SLEM , 我们有了一种稳健的、基于文本的方法来处理分词器的混乱现实。通过 TLI , 我们有了一种统计上合理的方法来利用词表重叠。通过 SLRS , 我们有了一个字符串级验证的理论框架，即使目前的计算限制约束了它。

对于学生和从业者来说，结论很明确: 部署和加速大语言模型的限制正在消失。我们现在可以将模型视为模块化组件，进行混合搭配，以实现速度和准确性的最佳平衡。

这些方法的完整基准测试和代码已在开源的 Hugging Face Transformers 库中提供。