Hugging Face Transformers 4.45.0 動的投機的デコード
Hugging Face と Intel Labs は、タスクに応じて最大 2.7 倍のテキスト生成速度向上を実現する動的投機的デコードを開発しました。この手法は、Transformers 4.45.0 のリリースから支援生成のデフォルト動作モードとなります。
動的投機的デコードの仕組み
動的投機的デコードは、標準的な投機的デコードを改良し、"speculation lookahead"（SL）を最適化することで、より高速なドラフトモデルが大規模なターゲットモデルによって並行して検証される前に生成するトークン数を調整します。
従来の手法は静的な SL または前回イテレーションの受容率に基づくヒューリスティックを使用していましたが、動的投機的デコードはアシスタントモデル自身の信頼度を用いて停止時期を決定します。具体的には、システムは各予測トークンのロジットのソフトマックスを監視します。アシスタントモデルの信頼度が事前に定義された assistant_confidence_threshold を下回った場合、そのイテレーションのトークン生成は停止し、シーケンスは検証のためにターゲットモデルへ送られます（たとえ num_assistant_tokens の上限に達していなくても）。
パフォーマンスベンチマーク
RTX 4090 上で greedy デコード（temperature = 0）を使用して実施したベンチマークにより、動的アプローチはさまざまなモデルペアとタスクにおいてヒューリスティックベースの手法を一貫して上回ることが示されました。
| ターゲットモデル | ドラフト（アシスタント）モデル | タスク | ヒューリスティックによるスピードアップ | 動的によるスピードアップ |
|---|---|---|---|---|
facebook/opt-6.7b |
facebook/opt-125m |
要約 | 1.82x | 2.71x |
facebook/opt-6.7b |
facebook/opt-125m |
自由生成 | 1.23x | 1.59x |
Salesforce/codegen-6B-mono |
Salesforce/codegen-350M-mono |
コード生成（python） | 0.89x | 1.09x |
google/flan-t5-xl |
google/flan-t5-small |
要約 | 1.18x | 1.31x |
meta-llama/Llama-3.1-8B |
meta-llama/Llama-3.2-1B |
要約 | 1.00x | 1.52x |
meta-llama/Llama-3.1-8B |
meta-llama/Llama-3.2-1B |
自由生成 | 1.00x | 1.18x |
meta-llama/Llama-3.1-8B |
meta-llama/Llama-3.2-1B |
コード生成（python） | 1.09x | 1.15x |
これらのベンチマークからの主な発見は次のとおりです：
- Llama 3.1/3.2 pairing: 動的アプローチは要約で 1.52 倍のスピードアップを達成しましたが、ヒューリスティックアプローチは有意な向上を示しませんでした。
- Code generation:
codegen-6B-monoに対しては、ヒューリスティックアプローチが実際に速度低下（0.89x）を引き起こしたのに対し、動的アプローチは速度向上（1.09x）を提供しました。
Transformers 4.45.0 における実装
動的投機は Transformers 4.45.0 における支援デコードのデフォルトモードです。assistant_model を generate メソッドに渡すことで利用できます。
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
prompt = "Alice and Bob"
checkpoint = "EleutherAI/pythia-1.4b-deduped"
assistant_checkpoint = "EleutherAI/pythia-160m-deduped"
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer(prompt, return_tensors="pt").to(device)
model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
assistant_model = AutoModelForCausalLM.from_pretrained(assistant_checkpoint).to(device)
outputs = model.generate(**inputs, assistant_model=assistant_model)
設定パラメータ
ユーザーは assistant_model.generation_config 内の以下のパラメータを調整して、特定のデータセットに対するパフォーマンスを最適化できます。
assistant_confidence_threshold: ドラフトモデルが生成を停止する信頼度レベル（ロジットのソフトマックス）。num_assistant_tokens: アシスタントがイテレーションごとに生成できる最大トークン数。num_assistant_tokens_schedule:'dynamic'（デフォルト）、'heuristic'、または'constant'のいずれかに設定して、投機戦略を切り替えます。
理論的根拠：オラクルモデル
動的調整の必要性を検証するために、研究者は「オラクル」を使用しました。これは、ドラフトモデルとターゲットモデル間で不一致が生じるまでトークンを生成し、各イテレーションで最適な speculation lookahead を特定します。MBPP データセットと Alpaca データセットの分析により、最適なドラフトトークン数に大きなばらつきがあることが明らかになり、静的な lookahead 値は推論速度最大化に対して最適でないことが証明されました。