ドラフトモデル(Speculative Decoding)(ドラフトモデル(スペキュレイティブデコーディング))
ドラフトモデル(Draft Model)は、投機的デコーディングにおいて候補トークン列を高速に生成する小型・軽量モデルである。ターゲットモデルの出力分布を近似することが求められ、モデル蒸留・層スキップ・自己ドラフトなど多様な構築手法が存在する。ドラフトモデルの品質(採択率)と推論コストのバランスが投機的デコーディング全体のスピードアップを決定する最重要因子である。
ドラフトモデルの役割と要件
投機的デコーディング(Speculative Decoding)において、ドラフトモデル(Draft Model)はターゲットモデルの出力を「推測」する役割を担う。ドラフトモデルに求められる要件は、一見矛盾する2つの特性の両立である。
- 高速性: ターゲットモデルよりも大幅に高速に推論できること。ドラフト生成のレイテンシがターゲットモデルの1回のフォワードパスよりも十分小さくなければ、投機的デコーディングのオーバーヘッドが利得を上回ってしまう
- 高精度: ターゲットモデルの出力分布をできるだけ忠実に近似すること。ドラフトの品質が高いほど採択率αが上がり、1回の検証ステップで多くのトークンが採択される
この2つのトレードオフが、ドラフトモデル設計の根本的な課題である。
| 評価軸 | 理想的な特性 | 現実的な制約 |
|---|---|---|
| パラメータ数 | ターゲットの1/10〜1/5 | 小さすぎると採択率が低下 |
| 推論レイテンシ | ターゲットの1/10以下 | GPUメモリに2モデル分必要 |
| 分布一致度 | KLダイバージェンス最小 | タスク・ドメイン依存 |
| 語彙互換性 | ターゲットと同一語彙 | 異なるtokenizerの場合は変換が必要 |
| メモリ使用量 | 最小限 | KVキャッシュも含めたGPUメモリ管理 |
独立ドラフトモデル方式
最も基本的なアプローチは、ターゲットモデルと同じアーキテクチャファミリーの小型モデルをドラフトモデルとして使用する方法である。例えば、以下のような組み合わせが実績を持つ。
| ターゲットモデル | ドラフトモデル | パラメータ比 | 典型的採択率 |
|---|---|---|---|
| Llama 2 70B | Llama 2 7B | 1:10 | 0.70〜0.80 |
| Llama 3 70B | Llama 3 8B | 1:8.75 | 0.72〜0.82 |
| CodeLlama 34B | CodeLlama 7B | 1:4.86 | 0.75〜0.85 |
| Vicuna 33B | Vicuna 7B | 1:4.71 | 0.68〜0.78 |
| GPT-4 | GPT-3.5-turbo | 非公開 | 非公開(推定0.7〜0.8) |
| Mixtral 8x22B | Mixtral 8x7B | MoE構造 | 0.65〜0.75 |
この方式のメリットは実装が単純でモデルの入手性が高いことだが、ドラフトモデルのGPUメモリ使用量が追加されるデメリットがある。特にGPUメモリが限られる環境では、量子化ドラフトモデル(GPTQ、AWQ、GGUF)を使用することで、メモリオーバーヘッドを軽減できる。
モデル蒸留によるドラフトモデル構築
ターゲットモデルの出力分布を効率的に近似するため、知識蒸留(Knowledge Distillation)を用いてドラフトモデルを専用に訓練する手法が研究されている。一般的な事前学習済みモデルをそのまま使うよりも、ターゲットモデルに特化した蒸留を行うことで採択率を大幅に向上させることが可能である。
蒸留戦略には以下のバリエーションがある。
- 出力分布蒸留(Output Distribution Distillation): ターゲットモデルの出力ロジット分布をソフトターゲットとして、ドラフトモデルのKLダイバージェンスを最小化する。最も標準的な手法であり、温度パラメータτの調整が重要
- 中間特徴量蒸留(Intermediate Feature Distillation): ターゲットモデルの中間層の隠れ状態も蒸留対象とする。ドラフトモデルの内部表現がターゲットに近づくことで、特に長い文脈での採択率が向上する
- タスク特化蒸留(Task-Specific Distillation): 特定のタスク(コード生成、要約、翻訳など)のデータで蒸留を行い、そのタスクでの採択率を最大化する。汎用性は犠牲になるが、特化ドメインでの効率が大幅に向上する
- オンライン蒸留(Online Distillation): 推論時にターゲットモデルの出力を使ってドラフトモデルをリアルタイムに更新する。デプロイ後も採択率が改善し続ける利点がある
| 蒸留手法 | 採択率向上 | 追加コスト | 実用性 |
|---|---|---|---|
| 出力分布蒸留 | +5〜15% | 中(GPU日数〜週単位) | 高 |
| 中間特徴量蒸留 | +3〜10% | 高(ターゲットのアクティベーション保存必要) | 中 |
| タスク特化蒸留 | +10〜20% | 低〜中 | 高(ドメイン限定) |
| オンライン蒸留 | +2〜8%(累積) | 低(推論時の追加計算) | 研究段階 |
自己ドラフト(Self-Drafting)方式
外部のドラフトモデルを別途用意するのではなく、ターゲットモデル自身の一部をドラフトに利用する自己ドラフト方式が注目されている。追加のGPUメモリをほぼ必要としない点が大きなメリットである。
層スキップ(Layer Skipping / Early Exit)
ターゲットモデルのTransformer層の一部をスキップし、浅い層の出力から直接トークンを予測する手法。例えば、70B(80層)のモデルの最初の20層だけを使ってドラフトを生成し、全80層で検証する。Draft & Verify(Zhang et al., 2023)やLayerSkip(Elhoushi et al., 2024)がこのカテゴリに属する。
自己投機的デコーディング(Self-Speculative Decoding)
ターゲットモデルの推論パスを修正し、一部の層やアテンションヘッドをバイパスすることでドラフト生成を行う。完全な層スキップよりも柔軟であり、適応的にスキップする層を選択できる。
Lookahead Decoding
Jacobi反復(Jacobi Iteration)の考え方を応用し、全位置のトークンを同時に予測・更新するアプローチ。ドラフトモデルを必要とせず、N-gramプールと検証の反復によって並列デコーディングを実現する。ただし、標準的な投機的デコーディングと比較してスピードアップは控えめ(1.5〜2.5倍)である。
動的ドラフト戦略
最新の研究では、固定的なドラフトモデルではなく、推論時に動的にドラフト戦略を調整するアプローチが発展している。
- 適応的投機長(Adaptive Speculation Length): 直近の採択率に基づいて投機長γを動的に調整する。採択率が高い区間ではγを増やし、低い区間ではγを減らすことで、平均的なスループットを最大化する。Hugging Face TransformersのAssisted Generationにデフォルト実装されている
- トークンレベル信頼度によるアーリーストップ: ドラフトモデルの出力確率(エントロピー)を監視し、信頼度が低いトークンが出現した時点でドラフト生成を停止する。不確実な推測を避けることで検証フェーズの無駄を削減する
- マルチドラフトモデル: 複数のドラフトモデルをタスクやコンテキストに応じて切り替える。コード生成にはコード特化ドラフト、自然文には汎用ドラフトを使うなど、タスク適応的な運用が可能
ドラフトモデル選定の実践ガイド
実際にドラフトモデルを選定する際の考慮事項を以下にまとめる。
| 考慮事項 | 推奨 | 理由 |
|---|---|---|
| アーキテクチャ | ターゲットと同一ファミリー | 語彙・トークナイザの互換性が保証される |
| パラメータ数 | ターゲットの1/5〜1/10 | 速度と採択率のバランスが最適 |
| 量子化 | INT4/INT8 量子化推奨 | メモリ節約、採択率への影響は限定的 |
| KVキャッシュ | ターゲットと共有可能な設計 | メモリ効率とロールバック処理の簡素化 |
| 評価指標 | タスク特化のα測定 | 汎用ベンチマークよりデプロイ環境での測定を優先 |
よくある質問(FAQ)
Q1: ドラフトモデルなしで投機的デコーディングは可能ですか?
可能である。Medusaはターゲットモデルに追加のデコーディングヘッドを取り付けることで、外部ドラフトモデルを不要にしている。また、Lookahead DecodingはJacobi反復を使ってドラフトモデルなしでの並列デコーディングを実現する。LayerSkipやSelf-Speculative Decodingも、ターゲットモデルの一部をドラフトとして再利用する方式であり、追加モデルのロードが不要である。
Q2: ドラフトモデルのファインチューニングは必要ですか?
必須ではないが、採択率の向上に大きく寄与する。事前学習済みの小型モデルをそのまま使用した場合の採択率が0.65〜0.75程度であるのに対し、ターゲットモデルへの蒸留を行うことで0.80〜0.90まで向上させた報告がある。特にドメイン特化の利用シーンでは、タスク特化蒸留の効果が顕著である。
Q3: 同じファミリーでないモデルをドラフトに使えますか?
語彙(トークナイザ)が異なる場合は直接的な利用が困難であるが、トークナイザマッピング(Tokenizer Mapping)やデトークナイズ→再トークナイズのパイプラインを構築することで対応可能である。ただし、変換のオーバーヘッドとトークン境界の不一致により効率が低下するため、同一語彙のモデルペアを使用することが強く推奨される。BiLD(Big Little Decoder)はこの問題にアプローチした手法の一つである。