使用先进的通用辅助生成技术加速LLM解码
⭐ 总结：LLM可以使用推测解码实现更快的推理。然而，它们通常缺乏可作为辅助模型并进一步加速的小型版本。这篇博客文章介绍了UAG-TLI，这是一种新的通用辅助生成（UAG）方法，允许使用任何小型LM提供增强的加速。我们对最先进的LLM进行的实验表明，速度提升高达2.5倍。UAG-TLI方法现已集成到Transformers🤗 4.50.0版本中，作为辅助生成（AG）的一部分，使高级AG更易于访问。🚀
引言
大型语言模型（LLM），如 DeepSeek，正在改变AI应用，从聊天机器人到代码生成。然而，其缓慢的推理速度仍然是一个主要瓶颈。推测解码（SD）已成为一种实用的解决方案，通过一次预测多个令牌来加速文本生成。
传统的SD方法要求辅助模型和目标模型共享相同的词汇表。然而，许多LLM没有可用的更小、更轻量级的版本作为辅助模型。这种限制降低了SD的灵活性，阻碍了其更广泛的应用。
在我们之前的博客文章中，我们介绍了UAG，它缓解了这一痛点，并允许使用任何现成的模型作为辅助模型，无论其词汇表如何。然而，这种方法仅限于贪婪解码。但概率解码对于生成多样化、流畅和连贯的文本至关重要。
在这篇博客文章中，我们介绍了UAG-TLI（UAG-令牌级交集），它是UAG的扩展，支持概率解码（即采样）。这一增强功能使得通过UAG进行的推测解码更具适用性，更易于集成，并进一步扩展了提升LLM推理速度的能力。在博客文章的后半部分，我们将展示使用各种推测解码技术（如DeepSeek-R1）加速LLM的示例代码。
UAG-TLI（令牌级交集）
UAG-TLI 的核心思想是将辅助模型概率分布映射到其词汇表与目标模型词汇表的交集上。换句话说，我们从辅助模型的词汇表中删除不存在于目标模型词汇表中的令牌。这确保了辅助模型只生成来自目标词汇表的令牌，从而消除了词汇表之间的翻译需求，并允许使用传统 SD 论文中引入的拒绝采样方法。
UAG-TLI 的主要优势是：
🔧 支持将模型中心上的任何模型用作辅助模型
🎲 支持使用采样（温度 > 0）
📈 提高推理速度
DeepSeek AI 模型基准测试结果
为了展示这项技术的影响，我们对各种LLM进行了基准测试。基准测试代码可用。下表显示了在不影响准确性的前提下显著的加速。
- 表1显示了对于缺乏共享词汇表的更小变体的模型所实现的加速。这使得UAG-TLI成为在非零温度环境下实现推理加速的首选解决方案。
- 表2显示了具有共享词汇表的较小变体的模型。在许多情况下，UAG-TLI比传统AG（传统SD）更有效。例如，在humaneval数据集上，使用UAG-TLI将vicuna-68m作为gemma-2-9b-it的辅助模型，其性能（1.46倍）优于在传统SD模式下使用gemma-2-2b-it作为辅助模型（1.36倍）。
- 我们注意到，DeepSeek-R1-Distill-Qwen-14B和DeepSeek-R1-Distill-Qwen-32B不共享与DeepSeek-R1-Distill-Qwen-1.5B相同的词汇表；因此，该模型只能使用UAG-TLI作为这些模型的辅助模型。
| 目标 | 硬件 | 数据集 | 方法 | 草稿 | 加速比 |
|---|---|---|---|---|---|
| Mixtral-8x22B-Instruct-v0.1 | 4*H100 NVL | 卷轴 | UAG-TLI | Qwen2.5-0.5B-Instruct | 1.69倍 |
| humaneval | UAG-TLI | vicuna-68m | 1.67倍 | ||
| cnn_dailymail | UAG-TLI | vicuna-68m | 1.53倍 | ||
| phi-4 | 1*H100 NVL | 卷轴 | UAG-TLI | Qwen2.5-0.5B-Instruct | 1.45倍 |
| CodeLlama-13b-Instruct-hf | 1*A6000 | humaneval | UAG-TLI | tiny_starcoder | 1.74倍 |
| DeepSeek-R1-Distill-Qwen-14B | 1*A6000 | 卷轴 | UAG-TLI | vicuna-68m | 1.59倍 |
| cnn_dailymail | UAG-TLI | vicuna-68m | 1.31倍 | ||
| humaneval | UAG-TLI | tiny_starcoder | 1.30倍 |
表1：缺乏共享词汇表的小型变体的目标模型的加速性能。
*与在自回归模式下运行目标模型相比的加速。
| 目标 | 硬件 | 数据集 | 方法 | 草稿 | 加速比 |
|---|---|---|---|---|---|
| DeepSeek-R1-Distill-Qwen-32B | 2*A100 80GB PCIe | 卷轴 | 传统SD | DeepSeek-R1-Distill-Qwen-7B | 2.02倍 |
| UAG-TLI | DeepSeek-R1-Distill-Qwen-1.5B | 2.26倍 | |||
| cnn_dailymail | 传统SD | DeepSeek-R1-Distill-Qwen-7B | 1.37倍 | ||
| UAG-TLI | vicuna-68m | 1.38倍 | |||
| humaneval | 传统SD | DeepSeek-R1-Distill-Qwen-7B | 1.96倍 | ||
| UAG-TLI | DeepSeek-R1-Distill-Qwen-1.5B | 1.70倍 | |||
| gemma-2-9b-it | 1*H100 NVL | 卷轴 | 传统SD | gemma-2-2b-it | 2.49倍 |
| UAG-TLI | vicuna-68m | 2.04倍 | |||
| humaneval | 传统SD | gemma-2-2b-it | 1.36倍 | ||
| UAG-TLI | vicuna-68m | 1.46倍 | |||
| DeepSeek-R1-Distill-Llama-70B | 2*H100 NVL | 卷轴 | 基线 | DeepSeek-R1-Distill-Llama-8B | 1.98倍 |
| UAG-TLI | DeepSeek-R1-Distill-Qwen-1.5B | 1.82倍 | |||
| 2*A100 8GB PCIe | humaneval | 基线 | DeepSeek-R1-Distill-Llama-8B | 2.3倍 | |
| UAG-TLI | tiny_starcoder | 1.44倍 |
表2：比较具有共享词汇表变体的目标模型的传统SD和UAG-TLI。
*与在自回归模式下运行目标模型相比的加速。
我们注意到，当使用DeepSeek-R1-Distill-Qwen-14B作为DeepSeek-R1-Distill-Qwen-32B的草稿模型在单个A100 80GB设备上运行时，由于内存卸载，显示出显著的减速。
现已在Hugging Face Transformers中可用
UAG-TLI 现已在 Hugging Face Transformers 库中提供，并在 do_sample=True 时作为异构（不同词汇表）推测解码的默认选择。请参阅 UniversalSpeculativeDecodingGenerator 类。开发者可以轻松地将这些技术集成到其工作流程中，并立即实现更快的 LLM 推理。
from transformers import pipeline
pipe = pipeline(
"text-generation",
model="google/gemma-2-9b",
assistant_model="double7/vicuna-68m", # This extra line is all that's needed!
torch_dtype="bfloat16"
)
pipe_output = pipe("Your prompt here", max_new_tokens=50, do_sample=True)
print(pipe_output[0]["generated_text"])
参考
引文
@article{timor2025acceleratingllminferencelossless,
title={Accelerating LLM Inference with Lossless Speculative Decoding Algorithms for Heterogeneous Vocabularies},
author={Nadav Timor and Jonathan Mamou and Daniel Korat and Moshe Berchansky and Oren Pereg and Gaurav Jain and Roy Schwartz and Moshe Wasserblat and David Harel},
year={2025},
eprint={2502.05202},
archivePrefix={arXiv},
primaryClass={cs.CL},
url={https://arxiv.org/abs/2502.05202},
}
