通用輔助生成：使用任何助理模型加速解碼
TL;DR
Hugging Face 與 Intel Labs 開發了通用輔助生成（Universal Assisted Generation，UAG），這項技術允許任何小型語言模型作為助理模型用於投機解碼，無需與目標模型共享相同的分詞器。此舉可為先前缺乏相容小規模變體的模型帶來 1.5 倍至 2.0 倍的推理加速。
輔助生成中分詞器依賴的挑戰
輔助生成（亦稱投機解碼）透過使用小型助理模型預測一系列 token，然後讓較大的目標模型在一次前向傳播中驗證這些 token，從而加速 LLM 推理。雖然此方法顯著降低延遲，但標準的輔助生成要求目標模型與助理模型共享相同的分詞器（即必須屬於同一模型家族）。
此依賴性成為瓶頸，因為許多高效能模型缺乏足夠小的版本（通常比目標模型小 50‑100 倍）以提供實質的加速。例如，gemma-2-9b 只有 2B 變體，而 CodeLlama-13b 則沒有更小的版本可供使用。
通用輔助生成（UAG）機制
通用輔助生成透過實作雙向分詞器翻譯流程，移除分詞器限制。這使得目標模型可以與任何助理模型配對，例如使用極小的 vicuna-68m 來加速 gemma-2-9b。
雙向分詞器翻譯
- 助理 → 目標： 當助理模型產生一系列 token 時，先將這些 token 轉換為文字，然後使用目標模型的分詞器重新分詞。
- 目標 → 助理： 目標模型驗證 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 時使用多項式抽樣（multinomial sampling）。與投機抽樣不同，多項式抽樣會在目標模型未抽到與助理模型相同的 token 時自動拒絕該 token，這可能導致相較於共享同一分詞器的助理模型較低的吞吐量。未來的更新計畫加入投機抽樣的支援，並將 UAG 直接整合至 🤗 Transformers pipeline，以提供更流暢的使用者體驗。