推測解碼讓 Whisper 推理速度提升 2 倍

TL;DR

Hugging Face 顯示，將推測解碼套用於 OpenAI 的 Whisper 模型，可將推理時間大幅縮短約 2×（例如在小型基準測試中從 73 秒降至 33 秒），且不會降低轉錄品質，因為此方法保證產生與原始模型完全相同的 token 輸出。

什麼是推測解碼？

推測解碼由 Leviathan 等人（2022）提出，將一個 快速助理模型 與一個 較大的主模型 配對。助理模型產生一段短的候選 token 序列（例如五個 token）。主模型隨後在一次前向傳播中驗證此序列，接受直到第一個不匹配為止的所有 token，並丟棄其餘部分。校正後的前綴成為助理模型的新上下文，循環重複。由於助理模型的運算速度遠快於主模型，整體解碼速度提升，而驗證步驟確保 相同的最終輸出 與僅使用主模型時一致。

關鍵需求：

  • 助理模型必須與主模型共享完全相同的 tokenizer/詞彙表。
  • 它的速度應至少比主模型 3× 更快，同時正確預測大多數（≈70‑80 %）的「簡單」token。

Whisper 基線推理速度

作者在 73 個 LibriSpeech 驗證樣本（總計約 9 MB）上對 Whisper large‑v2 進行基準測試，使用：

  • float16 精度、FlashAttention（attn_implementation="sdpa"）以及低 CPU 記憶體載入。
  • 以每個樣本測量產生時間。

結果： 總計 72.99 秒（≈每樣本 1 秒），字錯誤率（WER）為 3.5 %。

將推測解碼套用於 Whisper

助理模型選擇

  • Distil‑Whisper distil‑large‑v2 被選為助理模型。它保留 Whisper 的編碼器，但僅使用 32 層解碼器中的 2 層，提供 6× 的速度優勢，同時在分布外資料上保持在完整模型的 1 % WER 內。
  • 編碼器可以在主模型與助理模型之間共享，因此 VRAM 開銷僅略增（約 8 %）。

實作細節

assistant_model = AutoModelForCausalLM.from_pretrained(
    "distil-whisper/distil-large-v2",
    torch_dtype=torch_dtype,
    low_cpu_mem_usage=True,
    use_safetensors=True,
    attn_implementation="sdpa",
).to(device)

# Generation with assistant
outputs = model.generate(**inputs, assistant_model=assistant_model, **kwargs)

assistant_model 參數會在 🤗 Transformers 中啟用 協助生成 策略，該策略實作推測解碼。

加速結果（英文）

  • 推測解碼時間： 總計 32.70 秒（≈每樣本 0.45 秒）。
  • 加速率： 2.2× 快於基線。
  • WER： 與基線相同，為 3.5 %，證實輸出完全一致。

相同的做法也適用於高階 pipeline API，只需傳入 generate_kwargs={"assistant_model": assistant_model}。

多語言轉錄

Distil‑Whisper 的檢查點僅支援英語，因此在多語言情境下，作者使用最小的多語言 Whisper 檢查點（tiny）作為助理模型。

在 VoxPopuli 的 73 個荷蘭語樣本上進行基準測試：

  • 基線（large‑v2）： 116.5 秒，WER 12.8 %。
  • 推測解碼（assistant = tiny）： 62.1 秒，WER 12.8 %。
  • 加速率： 1.9×。

此方法同時適用於轉錄與翻譯任務，只需在 generate 時提供相應的 language 與 task 參數。

提升效能的策略

選擇助理模型

  • 目標是相較於主模型至少提升 ≥ 3× 的速度，且 token 一致率達 ≥ 70‑80 %。
  • 針對特定語言，可將大型 Whisper 模型（例如 large‑v3）微調作為主模型，並將相同架構蒸餾為快速助理。此舉可對齊 token 分布，提升兩者的 WER。

批次大小考量

  • 推測解碼在 batch size = 1 時可獲得最大效益。
  • 較大的批次需要整批的所有候選都與主模型匹配；若有不匹配則會提前丟棄，導致加速效果下降。
  • 實驗顯示，加速效益可持續至 batch size 4；超過此數量，額外開銷會超過收益（參見 Distil‑Whisper 論文，Sec. D.3）。

實務要點

  • 推測解碼是一個 即插即用 的替代方案，可直接套用於現有的 Whisper 流程：只需加入 assistant_model 參數。
  • 它提供 免費 2× 的推理加速，同時保證轉錄品質不變。
  • 此方法與模型無關；任何基於 transformer 的序列模型，只要有共享相同 tokenizer 的更快助理模型，都能受惠。

「推測解碼提供了完美的即插即用替代方案，可用於現有的 Whisper 流程，因為它提供免費的 2× 加速，同時保持相同的準確度。」 — Sanchit Gandhi, Hugging Face 部落格 (2023‑12‑20)

致謝：本文感謝 Patrick von Platen、Pedro Cuenca 與 Joao Gante 的回饋，以及 🤗 Transformers 中的協助生成實作。

Sources