ユニバーサル・アシスト生成: 任意のアシスタントモデルで高速デコード
TL;DR
Hugging Face と Intel Labs は、Universal Assisted Generation（UAG）という手法を開発しました。この手法は、対象モデルと同じトークナイザーを共有していなくても、任意の小型言語モデルをスペキュレーティブデコードのアシスタントとして使用できるようにします。これにより、従来は互換性のある小型バリアントが存在しなかったモデルでも、推論速度が 1.5 倍から 2.0 倍向上します。
アシスト生成におけるトークナイザー依存の課題
アシスト生成（スペキュレーティブデコード）は、小型のアシスタントモデルがトークン列を予測し、より大きな対象モデルがそれを単一のフォワードパスで検証することで LLM の推論を高速化します。これによりレイテンシが大幅に削減されますが、従来のアシスト生成では対象モデルとアシスタントモデルが同じトークナイザーを共有している必要があります（すなわち、同一ファミリーのモデルである必要があります）。
この依存関係はボトルネックとなります。多くの高性能モデルは、対象モデルの 50〜100 倍程度に小さくなる十分に小型のバージョンが存在せず、実質的な速度向上が得られません。例えば、gemma-2-9b には 2B のバリアントしかなく、CodeLlama-13b にはこの目的に使用できる小型バージョンがありません。
Universal Assisted Generation（UAG）メカニズム
Universal Assisted Generation は、双方向トークナイザー変換プロセスを実装することでトークナイザーの制約を取り除きます。これにより、対象モデルと任意のアシスタントモデルを組み合わせることが可能になり、たとえば小型の vicuna-68m を使用して gemma-2-9b を高速化できます。
双方向トークナイザー変換
- Assistant to Target（アシスタント → 対象）: アシスタントモデルがトークン列を生成すると、そのトークンはテキストに変換され、続いて対象モデルのトークナイザーで再トークナイズされます。
- Target to Assistant（対象 → アシスタント）: 対象モデルがトークンを検証した後、得られた対象トークンはアシスタント側のトークン形式に変換され、次のイテレーションのためにアシスタントモデルのコンテキストに追加されます。
語彙不一致の処理
再エンコード時の精度を確保するため、UAG は新たに生成されたシーケンスの前に過去トークンのコンテキストウィンドウを付加します。この全シーケンスを対象フォーマットに再エンコードし、最新の対象トークンと整列させて新トークンの正確な挿入位置を決定します。さらに、対象からアシスタントへの変換時には、語彙が一致しないトークンはアシスタントモデルのキー・バリュー（KV）キャッシュから破棄され、データの整合性が保たれます。
パフォーマンスベンチマーク
UAG は、同一ファミリーの適切なアシスタントが存在しないモデルに対して、顕著なレイテンシ改善をもたらします。以下のタスクで観測された速度向上率を示します。
| 対象モデル | アシスタントモデル | タスク | 速度向上 |
|---|---|---|---|
CodeLlama-13b-Instruct-hf |
tiny_starcoder_py |
コード生成 | 1.90x |
Mixtral-8x22B-Instruct-v0.1 |
vicuna-68m |
要約 | 1.52x |
gemma-2-9b |
vicuna-68m |
要約 | 1.76x |
Mixtral-8x22B-Instruct-v0.1 |
Qwen2-0.5B-Instruct |
長文要約 | 1.78x |
Llama-3.1-70B |
Qwen2-0.5B-Instruct |
長文要約 | 1.78x |
Phi-3-medium-128k-instruct |
Qwen2-0.5B-Instruct |
長文要約 | 1.91x |
実験は、モデルサイズに応じて単一の A6000 GPU から 4 台の A100 GPU までのハードウェア上で、ランダムに選択した 100 件のサンプルを使用して実施しました。
実装と使用方法
Universal Assisted Generation は 🤗 Transformers ライブラリに 4.46.0 以降で統合されています。ユーザーは tokenizer と assistant_tokenizer の両方を generate() メソッドに渡すことで UAG を利用できます。
from transformers import AutoModelForCausalLM, AutoTokenizer
prompt = "Alice and Bob"
checkpoint = "google/gemma-2-9b"
assistant_checkpoint = "double7/vicuna-68m"
assistant_tokenizer = AutoTokenizer.from_pretrained(assistant_checkpoint)
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer(prompt, return_tensors="pt")
model = AutoModelForCausalLM.from_pretrained(checkpoint)
assistant_model = AutoModelForCausalLM.from_pretrained(assistant_checkpoint)
outputs = model.generate(**inputs, assistant_model=assistant_model, tokenizer=tokenizer, assistant_tokenizer=assistant_tokenizer)
現在の制限と今後の課題
UAG は現在、do_sample=True の場合に多項サンプリングをサポートしています。スペキュレーティブサンプリングとは異なり、多項サンプリングは対象モデルがアシスタントと同じトークンをサンプリングしなかった場合に自動的にそのトークンを破棄するため、同一トークナイザーを共有するアシスタントに比べてスループットが低下することがあります。今後のアップデートではスペキュレーティブサンプリングのサポートを追加し、UAG を 🤗 Transformers のパイプラインに直接統合して、よりシームレスなユーザー体験を提供する予定です。