AI・機械学習
上級

ドラフトモデル(Speculative Decoding)(ドラフトモデル(スペキュレイティブデコーディング))

ドラフトモデル(Draft Model)は、投機的デコーディングにおいて候補トークン列を高速に生成する小型・軽量モデルである。ターゲットモデルの出力分布を近似することが求められ、モデル蒸留・層スキップ・自己ドラフトなど多様な構築手法が存在する。ドラフトモデルの品質(採択率)と推論コストのバランスが投機的デコーディング全体のスピードアップを決定する最重要因子である。

0 回閲覧
0 いいね

ドラフトモデルの役割と要件

投機的デコーディング(Speculative Decoding)において、ドラフトモデル(Draft Model)はターゲットモデルの出力を「推測」する役割を担う。ドラフトモデルに求められる要件は、一見矛盾する2つの特性の両立である。

  1. 高速性: ターゲットモデルよりも大幅に高速に推論できること。ドラフト生成のレイテンシがターゲットモデルの1回のフォワードパスよりも十分小さくなければ、投機的デコーディングのオーバーヘッドが利得を上回ってしまう
  2. 高精度: ターゲットモデルの出力分布をできるだけ忠実に近似すること。ドラフトの品質が高いほど採択率αが上がり、1回の検証ステップで多くのトークンが採択される

この2つのトレードオフが、ドラフトモデル設計の根本的な課題である。

評価軸理想的な特性現実的な制約
パラメータ数ターゲットの1/10〜1/5小さすぎると採択率が低下
推論レイテンシターゲットの1/10以下GPUメモリに2モデル分必要
分布一致度KLダイバージェンス最小タスク・ドメイン依存
語彙互換性ターゲットと同一語彙異なるtokenizerの場合は変換が必要
メモリ使用量最小限KVキャッシュも含めたGPUメモリ管理

独立ドラフトモデル方式

最も基本的なアプローチは、ターゲットモデルと同じアーキテクチャファミリーの小型モデルをドラフトモデルとして使用する方法である。例えば、以下のような組み合わせが実績を持つ。

ターゲットモデルドラフトモデルパラメータ比典型的採択率
Llama 2 70BLlama 2 7B1:100.70〜0.80
Llama 3 70BLlama 3 8B1:8.750.72〜0.82
CodeLlama 34BCodeLlama 7B1:4.860.75〜0.85
Vicuna 33BVicuna 7B1:4.710.68〜0.78
GPT-4GPT-3.5-turbo非公開非公開(推定0.7〜0.8)
Mixtral 8x22BMixtral 8x7BMoE構造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)はこの問題にアプローチした手法の一つである。