通用辅助生成：使用任意辅助模型实现更快解码
TL;DR：许多 LLM，如
gemma-2-9b 和 Mixtral-8x22B-Instruct-v0.1，缺乏一个更小的版本可用于辅助生成。在本篇博客文章中，我们介绍了**通用辅助生成**：一种由 Intel Labs 和 Hugging Face 开发的方法，它将辅助生成扩展到**任意模型家族**的小型语言模型 🤯。因此，现在可以通过几乎零开销的方式将**任何**解码器或混合专家模型的推理速度提升 **1.5 倍至 2.0 倍** 🔥🔥🔥。让我们深入了解！
引言
如今，最强大的开源 LLM 通常拥有数十亿到数千亿个参数（你好 Llama-3.1-405B 👋），在生产环境中部署这些庞然大物带来了一系列工程挑战。其中一个挑战是这些大型模型的文本生成速度慢，这促使社区开发了各种技术来加速解码过程。辅助生成，也称为推测解码，是一种非常流行且实用的加速 LLM 推理而不损失准确性的方法。在这篇博客文章中，我们将探讨辅助生成的工作原理，并分享我们将其扩展到 Hugging Face Hub 上14 万个语言模型中任意模型的研究成果 🚀！
辅助生成
辅助生成的核心思想是使用一对模型，即目标模型和辅助模型。辅助模型是目标模型的一个更小、更高效的版本，例如，你可以使用 Llama-3.2-1B 作为更大的 Llama-3.1-70b 目标模型的辅助模型。辅助生成是一个迭代过程。每个循环中，辅助模型自回归地逐个生成一系列 token。然后，目标模型在一次前向传播中验证序列中的所有辅助 token。通过在目标模型的每次前向传播中确认多个 token，而不是每次只生成一个 token，从而实现加速。有关更详细的解释，请参阅原始博客文章。结合最近引入的动态推测策略，辅助生成可将文本生成速度提高 1.5 倍到 3 倍，具体取决于任务和所使用的模型。
辅助生成带来的显著加速伴随着一个显著的缺点：目标模型和辅助模型必须共享相同的分词器，这意味着它们需要来自同一个模型家族。然而，许多广泛使用的模型缺乏足够小巧且准确的小型版本，无法大幅减少延迟。根据我们的经验，当辅助模型比目标模型小 50-100 倍时，通常才能看到有意义的加速。例如，CodeLlama-13b 缺乏一个更小的版本，而 gemma-2-9b 只有一个 2b 变体，它仍然不够小/快，无法实现显著的性能提升。
通用辅助生成
为了缓解这一痛点，Intel Labs 与 Hugging Face 的朋友们共同开发了通用辅助生成（UAG）。UAG 允许选择任意一对目标模型和辅助模型，无论它们的分词器如何。例如，可以将 gemma-2-9b 用作目标模型，而使用微小的 vicuna-68m 作为辅助模型。
我们提出的方法主要思想是双向分词器转换。一旦辅助模型完成一次生成迭代，辅助 token 将被转换为文本，然后使用目标模型的分词器进行分词以生成目标 token。在验证步骤之后，目标 token 同样被转换回辅助 token 格式，然后将其附加到辅助模型的上下文，再开始下一次迭代。
由于辅助分词器和目标分词器使用不同的词汇表，因此有必要处理它们之间的差异。为了准确地重新编码新生成的辅助 token，需要预先添加一个包含几个先前 token 的上下文窗口。然后将整个序列重新编码为目标 token 格式，并与最新的目标 token 对齐，以精确定位新生成的 token 应附加的位置。这个过程在下面的视频中有所说明。
虽然上述视频中未显示，但从目标模型到辅助模型的 token 重编码遵循类似的过程。然而，不匹配的 token 必须从辅助模型的键值 (KV) 缓存中丢弃，以确保数据完整性。
基准测试
下表显示了目标模型与使用不同分词器的辅助模型配对时观察到的延迟改进。
| 目标模型 | 辅助模型 | 数据集 | 任务 | 加速比 |
|---|---|---|---|---|
codellama/CodeLlama-13b-Instruct-hf |
bigcode/tiny_starcoder_py |
openai/humaneval |
代码生成 | 1.90倍 |
mistralai/Mixtral-8x22B-Instruct-v0.1 |
double7/vicuna-68m |
cnn_dailymail |
摘要 | 1.52倍 |
google/gemma-2-9b |
double7/vicuna-68m |
cnn_dailymail |
摘要 | 1.76倍 |
mistralai/Mixtral-8x22B-Instruct-v0.1 |
Qwen/Qwen2-0.5B-Instruct |
tau/scrolls |
长上下文摘要 | 1.78倍 |
meta-llama/Llama-3.1-70B |
Qwen/Qwen2-0.5B-Instruct |
tau/scrolls |
长上下文摘要 | 1.78倍 |
microsoft/Phi-3-medium-128k-instruct |
Qwen/Qwen2-0.5B-Instruct |
tau/scrolls |
长上下文摘要 | 1.91倍 |
请注意，上述目标模型没有适合使用标准辅助生成进行加速的小型变体（小于 10 亿参数）。
每个实验都在 100 个随机选择的示例上进行。使用 Llama 和 Mixtral 目标模型的实验分别使用 2 块和 4 块 A100 GPU。所有其他实验均使用单块 A6000 GPU 运行。
代码
通用辅助生成已集成到 🤗 Transformers 的 4.46.0 版本中。
要使用，将 tokenizer 和 assistant_tokenizer 传递给 generate()
>>> from transformers import AutoModelForCausalLM, AutoTokenizer
>>> prompt = "Alice and Bob"
>>> checkpoint = "google/gemma-2-9b"
>>> assistant_checkpoint = "double7/vicuna-68m"
>>> assistant_tokenizer = AutoTokenizer.from_pretrained(assistant_checkpoint)
>>> tokenizer = AutoTokenizer.from_pretrained(checkpoint)
>>> inputs = tokenizer(prompt, return_tensors="pt")
>>> model = AutoModelForCausalLM.from_pretrained(checkpoint)
>>> assistant_model = AutoModelForCausalLM.from_pretrained(assistant_checkpoint)
>>> outputs = model.generate(**inputs, assistant_model=assistant_model, tokenizer=tokenizer, assistant_tokenizer=assistant_tokenizer)
>>> tokenizer.batch_decode(outputs, skip_special_tokens=True)
['Alice and Bob are sitting in a bar. Alice is drinking a beer and Bob is drinking a']
未来方向
虽然在使用标准辅助生成时，传递 `do_sample=True` 会使用推测性采样算法（论文中的算法 1），但 UAG 目前仅支持多项式采样。在多项式采样中，如果目标模型没有采样到与辅助模型相同的 token，则该 token 会自动被拒绝，而推测性采样则不然。实际上，这意味着 UAG 在 `do_sample=True` 模式下的吞吐量将低于辅助模型具有相同分词器的情况。未来，我们计划为 UAG 添加推测性采样支持。此外，我们打算将 UAG 集成到 🤗 Transformers 管道中，以实现更简洁和流线型的使用方式。
