更快的輔助生成: 動態推測

HuggingFace發表於2024-10-21

⭐ 在這篇部落格文章中，我們將探討 動態推測解碼 ——這是由英特爾實驗室和 Hugging Face 開發的一種新方法，可以加速文字生成高達 2.7 倍，具體取決於任務。從 Transformers🤗 釋出的版本 4.45.0 開始，這種方法是輔助生成的預設模式⭐

推測解碼

推測解碼 技術十分流行，其用於加速大型語言模型的推理過程，與此同時保持其準確性。如下圖所示，推測解碼透過將生成過程分為兩個階段來工作。在第一階段，一個快速但準確性較低的 草稿 模型 (Draft，也稱為助手) 自迴歸地生成一系列標記。在第二階段，一個大型但更準確的 目標 模型 (Target) 對生成的草稿標記進行並行驗證。這個過程允許目標模型在單個前向傳遞中生成多個標記，從而加速自迴歸解碼。推測解碼的成功在很大程度上取決於 推測前瞻 (Speculative Lookahead，下文用 SL 表示)，即草稿模型在每次迭代中生成的標記數量。在實踐中，SL 要麼是一個靜態值，要麼基於啟發式方法，這兩者都不是在推理過程中發揮最大效能的最優選擇。

推測解碼的單次迭代

推測解碼的單次迭代

動態推測解碼

Transformers🤗 庫提供了兩種不同的方法來確定在推理過程中調整草稿 (助手) 標記數量的計劃。基於 Leviathan 等人 的直接方法使用推測前瞻的靜態值，並涉及在每個推測迭代中生成恆定數量的候選標記。另一種 基於啟發式方法的方法 根據當前迭代的接受率調整下一次迭代的候選標記數量。如果所有推測標記都是正確的，則候選標記的數量增加; 否則，數量減少。

我們預計，透過增強最佳化策略來管理生成的草稿標記數量，可以進一步減少延遲。為了測試這個論點，我們利用一個預測器來確定每個推測迭代的最佳推測前瞻值 (SL)。該預測器利用草稿模型自迴歸的生成標記，直到草稿模型和目標模型之間的預測標記出現不一致。該過程在每個推測迭代中重複進行，最終確定每次迭代接受的草稿標記的最佳 (最大) 數量。草稿/目標標記不匹配是透過在零溫度下 Leviathan 等人提出的拒絕抽樣演算法 (rejection sampling algorithm) 來識別的。該預測器透過在每一步生成最大數量的有效草稿標記，並最小化對草稿和目標模型的呼叫次數，實現了推測解碼的全部潛力。我們稱使用該預測器得到 SL 值的推測解碼過程為預知 (orcale) 的推測解碼。

下面的左圖展示了來自 MBPP 資料集的程式碼生成示例中的預知和靜態推測前瞻值在推測迭代中的變化。可以觀察到預知的 SL 值 (橙色條) 存在很高的變化。
靜態 SL 值 (藍色條) 中，生成的草稿標記數量固定為 5，執行了 38 次目標前向傳播和 192 次草稿前向傳播，而預知的 SL 值只執行了 27 次目標前向傳播和 129 次草稿前向傳播 - 減少了很多。右圖展示了整個 Alpaca 資料集中的預知和靜態推測前瞻值。

在 MBPP 的一個例子上的預知和靜態推測前瞻值 (SL)。

在 MBPP 的一個例子上的預知和靜態推測前瞻值 (SL)。

在整個 Alpaca 資料集上平均的預知 SL 值。
在整個 Alpaca 資料集上平均的預知 SL 值。

上面的兩個圖表展示了預知推測前瞻值的多變性，這說明靜態的推測解碼可能使次優的。

為了更接近預知的推測解碼並獲得額外的加速，我們開發了一種簡單的方法來在每次迭代中動態調整推測前瞻值。在生成每個草稿令牌後，我們確定草稿模型是否應繼續生成下一個令牌或切換到目標模型進行驗證。這個決定基於草稿模型對其預測的信心，透過 logits 的 softmax 估計。如果草稿模型對當前令牌預測的信心低於預定義的閾值，即 assistant_confidence_threshold ，它將在該迭代中停止令牌生成過程，即使尚未達到最大推測令牌數 num_assistant_tokens 。一旦停止，當前迭代中生成的草稿令牌將被髮送到目標模型進行驗證。

基準測試

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

目標模型 草稿模型 任務型別 加速比 - 啟發式策略 加速比 - 動態策略
facebook/opt-6.7b facebook/opt-125m summarization 1.82x 2.71x
facebook/opt-6.7b facebook/opt-125m open-ended generation 1.23x 1.59x
Salesforce/codegen-6B-mono Salesforce/codegen-350M-mono code generation (python) 0.89x 1.09x
google/flan-t5-xl google/flan-t5-small summarization 1.18x 1.31x
meta-llama/Llama-3.1-8B meta-llama/Llama-3.2-1B summarization 1.00x 1.52x
meta-llama/Llama-3.1-8B meta-llama/Llama-3.2-1B open-ended generation 1.00x 1.18x
meta-llama/Llama-3.1-8B meta-llama/Llama-3.2-1B code generation (python) 1.09x 1.15x
  • 表格中的結果反映了貪婪解碼 (temperature = 0)。在使用取樣 (temperature > 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 等等都將有資格進行推測解碼。敬請關注！

參考資料

  • Dynamic Speculation Lookahead Accelerates Speculative Decoding of Large Language Models。

在這篇論文中，我們介紹了 DISCO，一種動態推測前瞻最佳化方法，利用分類器決定草稿模型是否應該繼續生成下一個標記，還是暫停，並切換到目標模型進行驗證，而不是僅僅使用對預測機率的簡單閾值。

  • Assisted Generation: a new direction toward low-latency text generation
  • Fast Inference from Transformers via Speculative Decoding

原文連結: https://hf.co/blog/dynamic_speculation_lookahead

原文作者: Jonathan Mamou, Oren Pereg, Joao Gante, Lewis Tunstall, Daniel Korat, Nadav Timor, Moshe Wasserblat

譯者: Zipxuan

相關文章