推測デコードで Whisper 推論が 2× 速くなる

TL;DR

Hugging Face は、OpenAI の Whisper モデルに推測デコードを適用すると、推論時間が約 2× 短縮されること（例：小規模ベンチマークで 73 秒から 33 秒へ）を示しています。転写品質に損失はなく、元のモデルと全く同じトークン出力が保証されるためです。

推測デコードとは何か？

Leviathan et al.（2022）により提案された推測デコードは、高速アシスタントモデル と 大規模メインモデル を組み合わせます。アシスタントは短い候補トークン列（例：5 トークン）を生成し、メインモデルは単一のフォワードパスでこの列を検証し、最初の不一致までのすべてのトークンを受け入れ、残りを破棄します。修正されたプレフィックスがアシスタントの新たなコンテキストとなり、サイクルが繰り返されます。アシスタントがメインモデルよりはるかに高速に動作するため、全体のデコード速度が向上し、検証ステップにより 同一の最終出力 が保証されます。

主な要件:

  • アシスタントはメインモデルと正確に同じトークナイザー／語彙を共有している必要があります。
  • メインモデルより少なくとも 3× 高速であり、かつ「簡単」なトークンの大多数（≈70‑80 %）を正しく予測できる必要があります。

ベースライン Whisper 推論速度

著者らは、73 サンプルの LibriSpeech バリデーションセット（総計 ≈9 MB）で Whisper large‑v2 のベンチマークを行いました。使用した条件は以下の通りです：

  • float16 精度、FlashAttention（attn_implementation="sdpa"）、および low‑CPU‑mem ローディング。
  • サンプルごとの生成時間を測定。

結果: 合計 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 の assisted generation 戦略を有効にし、推測デコードを実装します。

スピードアップ結果（英語）

  • 推測デコード時間: 合計 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× の速度と ≥ 70‑80 % のトークン一致率を目指す。
  • 特定の言語向けには、大規模 Whisper モデル（例：large‑v3）をメインモデルとしてファインチューニングし、同じアーキテクチャを蒸留して高速アシスタントを作成します。これによりトークン分布が整合し、両モデルの WER が改善されます。

バッチサイズの考慮事項

  • 推測デコードは batch size = 1 のときに最大の効果を発揮します。
  • バッチが大きくなると、バッチ全体の すべての候補 がメインモデルと一致する必要があり、不一致があると早期に破棄され、スピードアップが減少します。
  • 経験的には、バッチサイズ 4 までスピードアップが持続しますが、それ以上になるとオーバーヘッドが利益を上回ります（Distil‑Whisper 論文、Sec. D.3 参照）。

実践的なポイント

  • 推測デコードは既存の Whisper パイプラインの ドロップイン置換 であり、assistant_model 引数を追加するだけです。
  • 転写品質を変えずに 無料の 2× 推論加速 を提供します。
  • このアプローチはモデルに依存せず、同じトークナイザーを共有する適切な高速アシスタントがあれば、任意のトランスフォーマーベースのシーケンスモデルが恩恵を受けられます。

“推測デコードは既存の Whisper パイプラインに最適なドロップイン置換を提供します。無料の 2× スピードアップを実現しながら、同じ精度を維持します。” — Sanchit Gandhi, Hugging Face ブログ (2023‑12‑20)

謝辞: 本稿は Patrick von Platen、Pedro Cuenca、Joao Gante にフィードバックと 🤗 Transformers における assisted‑generation 実装への貢献を感謝します。

Sources