通用辅助生成：使用任意助手模型实现更快解码

TL;DR

Hugging Face 与 Intel Labs 开发了通用辅助生成（UAG），这项技术使得任何小型语言模型都可以作为推测解码的助手使用，无论它是否与目标模型共享相同的分词器。这为之前缺乏兼容小规模变体的模型提供了 1.5 倍至 2.0 倍的推理加速。

辅助生成中分词器依赖的挑战

辅助生成（或称推测解码）通过使用小型助手模型预测一系列 token，然后让更大的目标模型在一次前向传播中进行验证，从而加速大语言模型的推理。虽然这显著降低了延迟，但标准的辅助生成要求目标模型和助手模型共享相同的分词器（即它们必须来自同一模型系列）。

这种依赖导致了瓶颈，因为许多高性能模型缺乏足够小的版本（通常比目标模型小 50‑100 倍），无法提供显著的加速。例如，gemma-2-9b 只有 2B 的变体，而 CodeLlama-13b 并没有可用于此目的的更小版本。

通用辅助生成（UAG）机制

通用辅助生成通过实现双向分词器翻译过程，消除了分词器的限制。这使得目标模型可以与任何助手模型配对，例如使用体积极小的 vicuna-68m 来加速 gemma-2-9b。

双向分词器翻译

  1. Assistant to Target（助手 → 目标）: 当助手模型生成一系列 token 时，这些 token 会被转换为文本，然后使用目标模型的分词器重新分词。
  2. Target to Assistant（目标 → 助手）: 在目标模型验证这些 token 后，得到的目标 token 会被转换回助手模型的 token 格式，并追加到助手模型的上下文中，以供下一轮迭代使用。

处理词汇表差异

为确保重新编码的准确性，UAG 在新生成的序列前预置了先前 token 的上下文窗口。整个序列会重新编码为目标模型的格式，并与最新的目标 token 对齐，以确定新 token 的精确插入位置。此外，在从目标模型翻译回助手模型时，任何不匹配的 token 都会从助手模型的键值（KV）缓存中丢弃，以保持数据完整性。

性能基准

UAG 为缺乏合适同系列助手的模型提供了显著的延迟改进。以下是在不同任务中观察到的加速效果：

目标模型 助手模型 任务 加速比
CodeLlama-13b-Instruct-hf tiny_starcoder_py 代码生成 1.90x
Mixtral-8x22B-Instruct-v0.1 vicuna-68m 摘要 1.52x
gemma-2-9b vicuna-68m 摘要 1.76x
Mixtral-8x22B-Instruct-v0.1 Qwen2-0.5B-Instruct 长上下文摘要 1.78x
Llama-3.1-70B Qwen2-0.5B-Instruct 长上下文摘要 1.78x
Phi-3-medium-128k-instruct Qwen2-0.5B-Instruct 长上下文摘要 1.91x

实验在 100 个随机选取的示例上进行，使用的硬件范围从单个 A6000 GPU 到四个 A100 GPU，具体取决于模型规模。

实现与使用

通用辅助生成已集成到 🤗 Transformers 库中，起始于 4.46.0 版本。用户可以通过向 generate() 方法传入 tokenizer 和 assistant_tokenizer 来实现 UAG。

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)

当前局限性与未来工作

UAG 目前在 do_sample=True 时支持多项式采样。与推测采样不同，多项式采样会在目标模型未采样到与助手相同的 token 时自动拒绝该 token，这可能导致相较于共享相同分词器的助手模型吞吐量降低。未来的更新计划引入对推测采样的支持，并将 UAG 直接集成到 🤗 Transformers 流水线中，以提供更简化的用户体验。

Sources