低遅延テキスト生成のための Hugging Face アシスト生成

TL;DR

Hugging Face は Assisted Generation と呼ばれる新しいデコード手法を導入しました。これは、より小さく高速なアシスタントモデルが候補トークンを提案し、より大きなモデルが単一のフォワードパスでそれらを検証することで、一般的なハードウェア上でレイテンシを最大 10 倍まで削減でき、特にメモリオフロードを使用する場合に効果的です。

テキスト生成レイテンシのボトルネック

テキスト生成のレイテンシは、主に計算量の問題ではなくメモリ帯域幅の問題です。標準的な自己回帰フォワードパスでは、GPU の RAM から GPU 計算コアへモデル層の重みをロードする際にボトルネックが発生します。大規模モデルは完全な応答を生成するために数百回の順次フォワードパスが必要になるため、このメモリ転送のオーバーヘッドが支配的なコストとなります。

既存の最適化手法（Flash Attention、INT8 量子化、バッチング（スループット向上のため）や Tensor Parallelism（メモリ帯域幅の分散））は存在しますが、これらはしばしば高額なコストやレイテンシのトレードオフを伴います。Assisted Generation は、プライマリモデルが必要とするフォワードパスの総数を減らすことでこの問題に対処します。

Assisted Generation の仕組み

Assisted Generation は言語デコーダの特性を利用します。単一のフォワードパスで次のトークンだけでなく、トークン列全体を検証できるのです。キャッシュを使用せずにトークン列をモデルに入力すると、モデルはその列の各位置に対するロジットを返します。

Assisted Generation ループ

このプロセスは小さなアシスタントモデルと大きなプライマリモデルの間で循環します。

候補生成: 小さなアシスタントモデルが貪欲デコードを用いて短い候補トークン列（最初は 5 トークンのウィンドウ）を生成します。
検証: プライマリモデルがこれらの候補に対して単一のフォワードパスを実行し、各位置のロジットを取得します。
比較: プライマリモデルの予測トークンを左から右へ、アシスタントの候補と比較します。
修正: 最初の不一致がアシスタントの失敗箇所を示します。その不一致点のプライマリモデルの予測トークンを採用し、以降のアシスタント候補は破棄します。
調整: ヒューリスティックにより次回の候補数を調整します。すべてのトークンが一致した場合はウィンドウを 2 増やし、いずれかが不一致だった場合は 1 減らします。

アシスタントモデルの要件

効率を保つため、アシスタントモデルは次の 2 条件を満たす必要があります。

トークナイザーの共有: アシスタントはプライマリモデルと全く同じトークナイザーを使用し、CPU ベースのデコードや再エンコードの高コストを回避します。
サイズ差: アシスタントはプライマリモデルに対して少なくとも 1 桁（10 倍）以上小さく、生成時間がプライマリモデルのフォワードパスに比べて無視できる程度であることが求められます。

パフォーマンスと能力

Assisted Generation はハードウェア構成やモデル設定に応じてさまざまな速度向上を提供します。

メモリオフロード: モデルが GPU メモリに収まらずオフロードを利用する場合、最大 10 倍 の速度向上が得られます。
GPU 常駐モデル: モデルが GPU メモリに収まる場合、標準設定で 2 倍、INT8 量子化時で 3 倍 の速度向上が期待できます。
タスク適合性: 入力に強く依存するタスク（自動音声認識（ASR）、翻訳、要約など）で特に効果的です。

サンプリングとの互換性

Assisted Generation は貪欲デコード向けに設計されていますが、マルチノミアルサンプリングでも使用可能です。ただし、アシスタントの有効性は temperature 設定に依存します。temperature が低い（0 に近い）ほど貪欲デコードに近く、レイテンシ削減効果が保たれます。一方、temperature が高いとランダム性が増し、アシスタントが失敗しやすくなり速度向上が減少します。

今後の方向性と関連研究

Hugging Face は、テキスト生成の将来はトークンあたりの固定計算コストから脱却する可能性があると示唆しています。テキストの複雑さに応じて、シーケンスの異なる部分をサイズの異なるモデルで生成するようなアーキテクチャが考えられます。

このアプローチは、Google Brain の Blockwise Parallel Decoding や DeepMind の Speculative Sampling といった、フォワードパスで長い続きのテキストを検証する手法と核心的な原則を共有しています。