Hugging Face 輔助生成以降低延遲的文字生成

TL;DR

Hugging Face 推出了 Assisted Generation，一種新型解碼方法，旨在降低自回歸文字生成的延遲。透過使用較小且更快的輔助模型提出候選 token，讓較大的模型在單一次前向傳播中驗證這些候選 token，於一般硬體上（特別是使用記憶體卸載時）可將延遲降低至最高 10 倍。

文字生成延遲的瓶頸

文字生成的延遲主要是記憶體頻寬問題，而非計算問題。在標準的自回歸前向傳播中，瓶頸出現在將模型層權重從 GPU 記憶體載入至 GPU 計算核心的過程。由於大型模型需要數百次連續的前向傳播才能產生完整回應，這種記憶體傳輸開銷成為主要成本。

雖然已有 Flash Attention、INT8 量化、批次處理（提升吞吐量）以及 Tensor Parallelism（分散記憶體頻寬）等優化手段，但它們往往伴隨高昂的金錢成本或在延遲上做出妥協。Assisted Generation 透過減少主要模型所需的前向傳播總次數來解決此問題。

Assisted Generation 的運作原理

Assisted Generation 利用語言解碼器的一個特性：單一次前向傳播即可驗證一整段 token 序列，而不僅僅是預測下一個 token。如果在未使用快取的情況下將一段 token 序列輸入模型，模型會回傳該序列中每個位置的 logits。

Assisted Generation 迴圈

此流程在小型輔助模型與大型主要模型之間循環進行：

  1. 候選產生：小型輔助模型使用貪婪解碼產生一段短的候選 token 序列（起始為 5 token 的視窗）。
  2. 驗證：主要模型對這些候選 token 執行單一次前向傳播，取得每個位置的 logits。
  3. 比較：將主要模型預測的 token 與輔助模型的候選 token 從左至右逐一比較。
  4. 校正：第一個不匹配的地方即為輔助模型失敗的點。保留主要模型在該位置的預測，並丟棄其後所有的輔助候選 token。
  5. 調整：根據比較結果以啟發式方式調整下一輪請求的候選數量：若全部匹配則視窗大小加 2，若有任何不匹配則視窗大小減 1。

輔助模型的需求

為了維持效率，輔助模型必須符合兩項條件：

  • 共享分詞器：輔助模型必須使用與主要模型完全相同的 tokenizer，避免昂貴的 CPU 端解碼與重新編碼步驟。
  • 尺寸差異：輔助模型的規模應至少比主要模型小一個量級，以確保其產生時間相較於主要模型的前向傳播可忽略不計。

效能與能力

Assisted Generation 依硬體與模型配置的不同，可提供不同程度的加速：

  • 記憶體卸載：當模型無法全部放入 GPU 記憶體而需依賴卸載時，可達到最高 10 倍 的加速。
  • GPU 常駐模型：模型能完整放入 GPU 記憶體時，加速幅度介於 2 倍（標準）至 3 倍（搭配 INT8 量化）。
  • 任務適用性：此方法對於以輸入為基礎的任務最為有效，例如自動語音辨識（ASR）、翻譯與摘要。

與抽樣的相容性

雖然設計上是針對貪婪解碼，但 Assisted Generation 亦可與多項式抽樣（multinomial sampling）結合使用。然而，輔助模型的效能取決於 temperature 設定。低 temperature（接近 0）近似貪婪解碼，能保留大部分延遲優勢；而高 temperature 會增加隨機性，使輔助模型更容易失效，從而降低加速效果。

未來方向與相關工作

Hugging Face 認為文字生成的未來可能會脫離「每個 token 固定計算成本」的框架。未來的架構或許會根據文字的複雜度，讓序列的不同部分由不同規模的模型生成。

此概念與其他研究共享核心原理，特別是 Blockwise Parallel Decoding（Google Brain）與 Speculative Sampling（DeepMind），兩者皆利用前向傳播來驗證較長的續寫。

SUMMARY: Hugging Face 引入了 Assisted Generation，一種利用較小的輔助模型預測候選 token，並由較大的模型驗證的解碼方法，可在某些硬體配置下將延遲降低至最高 10 倍。

TITLE: Hugging Face 輔助生成以降低延遲的文字生成

Sources