輔助生成：一種低延遲文字生成的新方向

大型語言模型如今炙手可熱，許多公司投入大量資源來擴充套件它們並解鎖新功能。然而，作為注意力持續時間不斷縮短的人類，我們也不喜歡它們緩慢的響應時間。延遲對於良好的使用者體驗至關重要，因此儘管小型模型的質量較低（例如在程式碼補全中），也常常被使用。

為什麼文字生成如此緩慢？是什麼阻止了您在不破產的情況下部署低延遲大型語言模型？在這篇博文中，我們將重新審視自迴歸文字生成的瓶頸，並引入一種新的解碼方法來解決延遲問題。您將看到，透過使用我們的新方法——輔助生成，您可以在商用硬體上將延遲降低多達 10 倍！

理解文字生成延遲

現代文字生成的核心很容易理解。讓我們看看其核心部分，即機器學習模型。它的輸入包含一個文字序列，其中包括迄今為止生成的文字，以及可能存在的其他模型特定元件（例如，Whisper 也有音訊輸入）。模型接收輸入並執行前向傳播：輸入被饋送到模型並沿其層順序傳遞，直到預測出下一個 token 的未歸一化對數機率（也稱為 logits）。一個 token 可以是整個單詞、子詞，甚至單個字元，這取決於模型。GPT-2 圖解是您深入瞭解文字生成這部分的絕佳參考。

模型的前向傳播會為您提供下一個 token 的 logits，您可以自由操作它們（例如，將不想要的單詞或序列的機率設定為 0）。文字生成的下一步是從這些 logits 中選擇下一個 token。常見的策略包括選擇最可能的 token（稱為貪婪解碼）或從其分佈中取樣（也稱為多項式取樣）。透過迭代地將模型前向傳播與下一個 token 選擇相結合，即可實現文字生成。這個解釋只是解碼方法的冰山一角；有關深入探討，請參閱我們關於文字生成的部落格文章。

從上面的描述中，文字生成中的延遲瓶頸很清楚：對於大型模型，執行一次模型前向傳播很慢，您可能需要在序列中進行數百次。但讓我們深入探討一下：為什麼前向傳播很慢？前向傳播通常由矩陣乘法主導，在快速檢視相應的維基百科部分後，您可以看出記憶體頻寬是此操作的限制（例如，從 GPU RAM 到 GPU 計算核心）。換句話說，前向傳播的瓶頸來自於將模型層權重載入到裝置的計算核心，而不是執行計算本身。

目前，您有三種主要途徑可以探索，以最大限度地利用文字生成，所有這些都旨在解決模型前向傳播的效能問題。首先是特定於硬體的模型最佳化。例如，您的裝置可能與Flash Attention相容，它透過重新排序操作來加速注意力層，或者與INT8 量化相容，後者可以減小模型權重的大小。

其次，當您知道會收到併發文字生成請求時，您可以批次處理輸入，並在延遲懲罰很小的情況下大幅提高吞吐量。現在，載入到裝置中的模型層權重可以並行用於多行輸入，這意味著您將以大致相同的記憶體頻寬負擔獲得更多 token。批次處理的缺點是您需要額外的裝置記憶體（或者將記憶體解除安裝到其他地方）——在此範圍的末端，您可以看到像FlexGen這樣的專案，它們以犧牲延遲為代價來最佳化吞吐量。

# Example showcasing the impact of batched generation. Measurement device: RTX3090
from transformers import AutoModelForCausalLM, AutoTokenizer
import time

tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2").to("cuda")
inputs = tokenizer(["Hello world"], return_tensors="pt").to("cuda")

def print_tokens_per_second(batch_size):
    new_tokens = 100
    cumulative_time = 0

    # warmup
    model.generate(
        **inputs, do_sample=True, max_new_tokens=new_tokens, num_return_sequences=batch_size
    )

    for _ in range(10):
        start = time.time()
        model.generate(
            **inputs, do_sample=True, max_new_tokens=new_tokens, num_return_sequences=batch_size
        )
        cumulative_time += time.time() - start
    print(f"Tokens per second: {new_tokens * batch_size * 10 / cumulative_time:.1f}")

print_tokens_per_second(1)   # Tokens per second: 418.3
print_tokens_per_second(64)  # Tokens per second: 16266.2 (~39x more tokens per second)

最後，如果您有多個裝置可用，您可以使用張量並行來分發工作負載並獲得更低的延遲。使用張量並行，您可以將記憶體頻寬負擔分散到多個裝置上，但現在除了執行多個裝置的金錢成本之外，您還必須考慮裝置間通訊瓶頸。其好處在很大程度上取決於模型大小：對於可以輕鬆放入單個消費級裝置的模型，好處非常有限。根據這篇DeepSpeed 部落格文章的結果，您可以看到可以將一個 17B 引數模型分散到 4 個 GPU 上，以將延遲降低 1.5 倍（圖 7）。

這三種類型的改進可以協同使用，從而產生高吞吐量解決方案。然而，在應用特定硬體最佳化之後，減少延遲的選擇有限——而且現有選擇成本高昂。讓我們解決這個問題！

語言解碼器前向傳播，再探

您上面讀到，每次模型前向傳播都會生成下一個 token 的 logits，但這實際上是一個不完整的描述。在文字生成過程中，典型的迭代是模型接收最新生成的 token 作為輸入，加上所有其他先前輸入的快取內部計算，返回下一個 token 的 logits。快取用於避免冗餘計算，從而加快前向傳播速度，但這不是強制性的（並且可以部分使用）。當停用快取時，輸入包含迄今為止生成的整個 token 序列，輸出包含對應於序列中所有位置的下一個 token 的 logits！位置 N 的 logits 對應於如果輸入由前 N 個 token 組成（忽略序列中所有後續 token）時下一個 token 的分佈。在貪婪解碼的特定情況下，如果您將生成的序列作為輸入並對結果 logits 應用 argmax 運算子，您將獲得生成的序列。

from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")

inputs = tok(["The"], return_tensors="pt")
generated = model.generate(**inputs, do_sample=False, max_new_tokens=10)
forward_confirmation = model(generated).logits.argmax(-1)

# We exclude the opposing tips from each sequence: the forward pass returns
# the logits for the next token, so it is shifted by one position.
print(generated[0, 1:].tolist() == forward_confirmation[0, :-1].tolist())  # True

這意味著您可以將模型前向傳播用於不同的目的：除了輸入一些 token 來預測下一個 token，您還可以將一個序列傳遞給模型，並再次檢查模型是否會生成相同的序列（或其一部分）。

讓我們考慮一下，您可以使用一個神奇的無延遲預言模型，該模型可以為任何給定輸入生成與您的模型相同的序列。為了便於討論，它不能直接使用，它僅限於作為您生成過程的助手。使用上面描述的屬性，您可以使用這個助手模型來獲取候選輸出 token，然後透過您的模型進行前向傳播以確認它們確實是正確的。在這種烏托邦式的情況下，文字生成的延遲將從 O(n) 減少到 O(1)，其中 n 是生成的 token 數量。對於長時間生成，我們談論的是幾個數量級。

更接近現實一步，假設助手模型失去了其預言屬性。現在它是一個無延遲模型，根據您的模型，它會錯誤地預測一些候選 token。由於任務的自迴歸性質，一旦助手錯誤地預測了一個 token，所有後續的候選 token 都必須失效。然而，這並不能阻止您在用您的模型更正錯誤的 token 後再次查詢助手，並迭代地重複此過程。即使助手預測錯誤了幾個 token，文字生成的延遲也比其原始形式低一個數量級。

顯然，不存在無延遲的輔助模型。然而，相對容易找到一個模型，其文字生成輸出近似於另一個模型——相同架構的較小版本，以類似方式訓練，通常符合此屬性。此外，當模型大小差異顯著時，使用較小模型作為助手的成本在考慮到跳過一些前向傳播的好處後，變得微不足道！您現在瞭解了輔助生成的核心。

使用輔助生成進行貪婪解碼

輔助生成是一種平衡行為。您希望助手快速生成候選序列，同時儘可能準確。如果助手質量差，您將承擔使用助手模型的成本，而幾乎沒有任何好處。另一方面，最佳化候選序列的質量可能意味著使用緩慢的助手，從而導致總體速度變慢。雖然我們無法為您自動選擇助手模型，但我們添加了一個額外要求和啟發式方法，以確保與助手在一起的時間保持在可控範圍內。

首先，要求是——助手必須與您的模型具有完全相同的分詞器。如果未滿足此要求，則必須新增昂貴的 token 解碼和重新編碼步驟。此外，這些額外步驟必須在 CPU 上進行，這反過來可能需要緩慢的裝置間資料傳輸。助手的快速使用對於輔助生成的好處顯現至關重要。

最後，啟發式方法。到目前為止，您可能已經注意到電影《盜夢空間》與輔助生成之間的相似之處——畢竟，您正在文字生成內部執行文字生成。每個候選 token 都會有一個輔助模型的前向傳播，我們知道前向傳播是昂貴的。雖然您無法提前知道輔助模型將正確識別多少個 token，但您可以跟蹤此資訊並使用它來限制向輔助模型請求的候選 token 數量——輸出的某些部分比其他部分更容易預測。

總而言之，以下是我們輔助生成迴圈的原始實現（程式碼）

1. 使用貪婪解碼透過輔助模型生成一定數量的候選 token，生成 candidates。第一次呼叫輔助生成時，生成的候選 token 數量初始化為 5。
2. 使用我們的模型，對 candidates 進行前向傳播，獲得 logits。
3. 使用 token 選擇方法（貪婪搜尋的 .argmax() 或採樣的 .multinomial()）從 logits 中獲取 next_tokens。
4. 將 next_tokens 與 candidates 進行比較，並獲取匹配的 token 數量。請記住，此比較必須從左到右進行因果關係：在第一次不匹配之後，所有候選者都將失效。
5. 使用匹配數量來切分內容並丟棄與未確認候選 token 相關的變數。本質上，在 next_tokens 中，保留匹配的 token 以及第一個分歧的 token（我們的模型從有效的候選子序列生成）。
6. 調整下一次迭代中要生成的候選 token 的數量——如果所有 token 都匹配，我們的原始啟發式方法將其增加 2，否則減少 1。

我們以無憂的方式設計了 🤗 Transformers 中的 API。您只需在新引數 assistant_model 中傳遞輔助模型，即可獲得延遲收益！在釋出此部落格文章時，輔助生成僅限於批處理大小為 1。

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)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
# ['Alice and Bob are sitting in a bar. Alice is drinking a beer and Bob is drinking a']

額外的內部複雜性是否值得？讓我們看一下貪婪解碼情況下的延遲資料（取樣結果在下一節中），考慮到批處理大小為 1。這些結果直接來自 🤗 Transformers，沒有任何額外最佳化，因此您應該能夠在您的設定中重現它們。

粗略地看一下收集到的數字，我們看到輔助生成可以在不同的設定中顯著減少延遲，但它並非萬能藥——您應該在將其應用於您的用例之前進行基準測試。我們可以得出結論，輔助生成：

1. 🤏 需要訪問一個至少比您的模型小一個數量級的輔助模型（差異越大越好）；
2. 🚀 在存在 INT8 的情況下，當模型適合 GPU 記憶體時，可將速度提高至 3 倍；否則，可提高至 2 倍；
3. 🤯 如果您使用的模型不適合您的 GPU 並且依賴記憶體解除安裝，您可以看到高達 10 倍的加速；
4. 📄 在輸入型任務中表現出色，例如自動語音識別或摘要。

使用輔助生成進行取樣

貪婪解碼適用於輸入依賴型任務（自動語音識別、翻譯、摘要等）或尋求事實知識的任務。需要高度創造性的開放式任務，例如語言模型作為聊天機器人的大多數用途，應改為使用取樣。輔助生成自然是為貪婪解碼設計的，但這並不意味著您不能將輔助生成與多項式取樣一起使用！

從下一個 token 的機率分佈中抽取樣本將導致我們的貪婪助手更頻繁地失敗，從而降低其延遲收益。然而，我們可以使用大多數基於取樣的應用程式中存在的溫度係數來控制下一個 token 的機率分佈的尖銳程度。在一個極端，當溫度接近 0 時，取樣將近似於貪婪解碼，偏向最可能的 token。在另一個極端，當溫度設定為遠大於 1 的值時，取樣將是混沌的，從均勻分佈中抽取。因此，低溫對您的助手模型更有利，保留了輔助生成的大部分延遲收益，如下所示。

為什麼不親自體驗一下，感受一下輔助生成呢？

未來方向

輔助生成表明，現代文字生成策略已成熟，可進行最佳化。理解它目前是記憶體受限問題，而不是計算受限問題，使我們能夠應用簡單的啟發式方法，以最大限度地利用可用記憶體頻寬，從而緩解瓶頸。我們相信，對輔助模型使用的進一步完善將使我們獲得更大的延遲降低——例如，如果我們請求輔助模型生成多個候選續寫，我們可能能夠跳過更多的前向傳播。自然，釋出高質量的小型模型以用作輔助模型，對於實現和放大收益至關重要。

最初發布在我們的 🤗 Transformers 庫中，與 .generate() 函式一起使用，我們期望在整個 Hugging Face 生態系統中提供它。其實現也完全開源，因此，如果您正在進行文字生成並且不使用我們的工具，請隨意將其用作參考。

最後，輔助生成重新提出了文字生成中的一個關鍵問題。該領域一直在發展，其約束是所有新 token 都是給定模型固定計算量的結果。在純自迴歸模式下，每個同構前向傳播產生一個 token。這篇部落格文章強化了不應如此的觀點：生成的輸出的大部分子部分也可以由大小僅為一部分的模型同樣生成。為此，我們將需要新的模型架構和解碼方法——我們很高興看到未來會怎樣！

相關工作

在這篇部落格文章最初發布後，我注意到其他工作也探索了相同的核心原理（使用前向傳播來驗證更長的續寫）。特別是，請檢視以下工作：

• 分塊並行解碼，由 Google Brain
• 推測性取樣，由 DeepMind

引用

@misc {gante2023assisted,
    author       = { {Joao Gante} },
    title        = { Assisted Generation: a new direction toward low-latency text generation },
    year         = 2023,
    url          = { https://huggingface.co/blog/assisted-generation },
    doi          = { 10.57967/hf/0638 },
    publisher    = { Hugging Face Blog }
}

致謝

我要感謝 Sylvain Gugger、Nicolas Patry 和 Lewis Tunstall 提供了許多寶貴的建議來改進這篇部落格文章。最後，感謝 Chunte Lee 為我們網頁設計了精美的封面。