透過動態推測實現更快的輔助生成

⭐ 在這篇博文中，我們將探討動態推測性解碼——一種由 Intel 實驗室和 Hugging Face 開發的新型方法，根據任務的不同，它可以將文字生成速度提高多達 2.7 倍。從 Transformers🤗 4.45.0 版本開始，此方法已成為輔助生成 的預設操作模式 ⭐

推測性解碼

推測性解碼是一種流行的技術，用於在保持準確性的同時加速大型語言模型的推理。如下圖所示，推測性解碼透過將生成過程分為兩個階段來工作。在第一階段，一個快速但準確性較低的草稿模型（又稱助手）自迴歸地生成一個令牌序列。在第二階段，一個大型但準確性更高的目標模型對生成的草稿令牌進行並行驗證。這個過程允許目標模型在單次前向傳遞中生成多個令牌，從而加速自迴歸解碼。推測性解碼的成功很大程度上取決於推測前瞻（SL），即草稿模型在每次迭代中生成的令牌數量。在實踐中，SL 是一個靜態值或基於啟發式，兩者都不是在推理過程中榨取最大效能的最佳選擇。

推測性解碼迭代。

動態推測性解碼

Transformers🤗 提供了兩種不同的方法來確定在推理過程中調整草稿（助手）令牌數量的排程。直接的方法，基於 Leviathan 等人的工作，使用靜態的推測前瞻值，並在每次推測迭代中生成恆定數量的候選令牌。另一種方法是基於啟發式的方法，它根據當前迭代的接受率調整下一次迭代的候選令牌數量。如果所有推測性令牌都正確，候選令牌的數量就會增加；否則，就會減少。

我們預計，一種增強的最佳化策略來管理生成的草稿令牌數量可以進一步減少延遲。為了驗證這一論點，我們使用一個預言機來確定每次推測迭代的最佳推測前瞻值。預言機利用草稿模型自迴歸地生成令牌，直到草稿模型和目標模型的預測令牌之間出現差異。這個過程在每次推測迭代中重複進行，最終確定每次迭代接受的最佳（最大）草稿令牌數量。草稿/目標令牌不匹配是透過 Leviathan 等人引入的零溫度拒絕取樣演算法識別的。這個預言機透過在每一步生成最大數量的有效草稿令牌並最小化對草稿模型和目標模型的呼叫次數，從而充分發揮了推測性解碼的潛力。

下圖中左側的圖表展示了來自 MBPP 資料集的一個程式碼生成示例中，在推測迭代中預言機和靜態推測前瞻值的情況。觀察到預言機推測前瞻值（橙色條）存在高度可變性。靜態推測前瞻（藍色條）中，生成的草稿令牌數量固定為 5，執行了 38 次目標前向傳遞和 192 次草稿前向傳遞，而預言機推測前瞻僅執行了 27 次目標前向傳遞和 129 次草稿前向傳遞——顯著減少。右側的圖表顯示了整個 Alpaca 資料集的預言機和靜態推測前瞻。

一個 MBPP 示例上的預言機和靜態推測前瞻 (SL) 值。

整個 Alpaca 資料集的平均預言機推測前瞻。

兩張圖都顯示了預言機推測前瞻值存在顯著差異，表明靜態推測前瞻可能不是最優的。

為了更接近 Oracle 並獲得額外的加速，我們開發了一種簡單的方法，可以在每次迭代中動態調整推測前瞻值。在生成每個草稿令牌後，我們根據助手模型對其預測的置信度（透過 logits 的 softmax 估計）來決定草稿模型是繼續生成下一個令牌還是切換到目標模型進行驗證。如果助手模型對當前令牌預測的置信度低於預定義的閾值（稱為 assistant_confidence_threshold），它將停止該迭代的令牌生成過程，即使尚未達到最大推測令牌數量 num_assistant_tokens。一旦停止，當前迭代中生成的草稿令牌將傳送到目標模型進行驗證。

基準測試

我們在一系列任務和模型配對上對動態方法和啟發式方法進行了基準測試。動態方法在所有測試中都表現出更好的效能。值得注意的是，使用動態方法並將 Llama3.2-1B 作為 Llama3.1-8B 的助手時，我們觀察到加速高達 1.52 倍，而啟發式方法在相同的設定下沒有顯示出顯著加速。另一個觀察結果是，使用啟發式方法時，codegen-6B-mono 會導致減速，而動態方法則顯示出加速。

• 表中結果反映了貪婪解碼（溫度 = 0）。在使用取樣（溫度 > 0）時也觀察到類似趨勢。

• 所有測試均在 RTX 4090 上進行。

• 我們的基準測試已公開，歡迎所有人評估進一步的改進：https://github.com/gante/huggingface-demos/tree/main/experiments/faster_generation

程式碼

動態推測已整合到 Hugging Face Transformers 庫的 4.45.0 版本中，現在作為輔助解碼的預設操作模式。要使用動態推測進行輔助生成，無需更改程式碼——只需像往常一樣執行程式碼即可

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

prompt = "Alice and Bob"
checkpoint = "EleutherAI/pythia-1.4b-deduped"
assistant_checkpoint = "EleutherAI/pythia-160m-deduped"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer(prompt, return_tensors="pt").to(device)

model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
assistant_model = AutoModelForCausalLM.from_pretrained(assistant_checkpoint).to(device)

outputs = model.generate(**inputs, assistant_model=assistant_model)

預設的動態推測前瞻引數反映了最佳值，但可以透過以下程式碼進行調整，以提高特定模型對或資料集的效能

# confidence threshold
assistant_model.generation_config.assistant_confidence_threshold=0.4

# 'constant' means that num_assistant_tokens stays unchanged during generation
assistant_model.generation_config.num_assistant_tokens_schedule='constant'

# the maximum number of tokens generated by the assistant model.
# after 20 tokens the draft halts even if the confidence is above the threshold
assistant_model.generation_config.num_assistant_tokens=20

要恢復到啟發式或常數（如 Leviathan 等人的工作）方法，只需將 num_assistant_tokens_schedule 分別設定為 'heuristic' 或 'constant'，並按如下所示設定 assistant_confidence_threshold=0 和 num_assistant_tokens=5

# Use 'heuristic' or 'constant' or 'dynamic'
assistant_model.generation_config.num_assistant_tokens_schedule='heuristic'
assistant_model.generation_config.assistant_confidence_threshold=0
assistant_model.generation_config.num_assistant_tokens=5

接下來是什麼？

我們引入了一種更快的輔助生成策略，稱為動態推測性解碼，它優於基於啟發式的方法以及生成固定數量候選令牌的方法。

在即將釋出的部落格文章中，我們將展示一種新的輔助生成方法：將任何目標模型與任何輔助模型結合！這將為加速 Hugging Face Hub 上無數沒有足夠小的輔助變體的模型開啟大門。例如，Phi 3、Gemma 2、CodeLlama 等等都將有資格進行推測性解碼。敬請期待！

參考文獻

• 動態推測前瞻加速大型語言模型的推測性解碼。在本文中，我們引入了 DISCO，一種動態推測前瞻最佳化方法，該方法利用分類器來決定草稿模型是否應繼續生成下一個令牌或暫停，並切換到目標模型進行驗證，而不是使用預測機率上的簡單閾值。
• 輔助生成：低延遲文字生成的新方向
• 透過推測性解碼實現 Transformer 的快速推理

引用

@article{mamou2024accelerating,
  title={Accelerating Speculative Decoding using Dynamic Speculation Length},
  author={Mamou, Jonathan and Pereg, Oren and Korat, Daniel and Berchansky, Moshe and Timor, Nadav and Wasserblat, Moshe and Schwartz, Roy},
  journal={arXiv preprint arXiv:2405.04304},
  year={2024}
}