AI・機械学習
上級

ドラフトモデル(ドラフトモデル)

投機的デコーディングにおいて、ターゲットモデルに先行して複数トークンの候補を高速生成する小型言語モデル。ターゲットモデルの1/10〜1/100のパラメータ数で、同一のtokenizerを共有し、類似した確率分布を出力することが求められる。蒸留(Knowledge Distillation)で訓練されたモデルが最も高い受理率を達成する。

0 回閲覧
0 いいね

ドラフトモデルとは

ドラフトモデル(Draft Model)は、投機的デコーディング(Speculative Decoding)において「下書き」となるトークン列を高速に生成する小型の言語モデルである。ターゲットモデル(本番用の大型モデル)が1トークン生成するのにかかる時間で、ドラフトモデルは4〜8トークンを生成できる。生成されたドラフトトークンはターゲットモデルによって一括検証され、確率的に受理または却下される。

ドラフトモデルの品質(ターゲットモデルとの分布類似度)が投機的デコーディング全体の性能を決定づけるため、適切なドラフトモデルの選択と訓練は最重要の設計判断となる。

ドラフトモデルの種類

種類代表例パラメータ比受理率追加メモリ特徴
独立小型モデルLlama 3.1 8B(→70B用)1/980〜90%16GB同一ファミリーで最高の受理率
蒸留専用モデルVicuna 68M(→13B用)1/19070〜80%0.3GB専用訓練で軽量化
n-gram(モデル不要)vLLM ngram mode040〜60%ほぼ0コンテキスト内パターン再利用
Self-Draft(追加ヘッド)Medusa / Eagle追加1〜3%75〜90%0.1〜0.5GBターゲットモデルに小型ヘッド装着
層スキップLayerSkip(Meta)同一60〜75%0ターゲットの前半層のみで推論

独立小型モデルの設計指針

必須条件

  1. Tokenizer の完全一致: ターゲットモデルと同一の語彙テーブル(vocab)とトークン化アルゴリズム(BPE/SentencePiece)を使用する。語彙が異なるとトークンIDの対応が取れず、検証フェーズが機能しない
  2. 推論速度の十分な差: ドラフトモデルの1トークン生成時間がターゲットモデルの1/K以下(K=投機トークン数)でないと高速化しない。目安としてパラメータ比1/5〜1/50
  3. 分布の類似度: KLダイバージェンス D_KL(p||q) が小さいほど受理率が高い。同一ファミリーの小型モデルが最適

訓練方法

ドラフトモデルの最も効果的な訓練方法は、ターゲットモデルからの知識蒸留(Knowledge Distillation)である。

  • ロジット蒸留: ターゲットモデルの出力ロジット分布を教師信号として小型モデルを訓練。temperature=2〜5でソフトラベルを生成し、KL損失で最適化
  • on-policy蒸留: ドラフトモデル自身が生成したテキストに対してターゲットモデルのロジットを取得し、実際のデコーディング時に近い分布で訓練
  • フィーチャー蒸留: 中間層の隠れ状態も教師信号に使用。層の対応関係(layer mapping)を定義して追加損失を加える

サイズの選び方

ターゲットサイズ推奨ドラフトサイズ理由
405B8B〜14B405Bのフォワードパスが十分遅いため、やや大きいドラフトでも受理率向上のメリットが大きい
70B1B〜8Bバランス型。8Bなら受理率90%近く到達
13B〜27B0.5B〜2Bターゲットが比較的高速なため、ドラフトは極力軽量に
7B〜8Bngram or Self-Draftターゲット自体が小型のため、外部ドラフトモデルのオーバーヘッドが相対的に大きい

n-gramドラフト(モデル不要方式)

n-gramドラフトは追加モデルなしで投機的デコーディングを実現する方式である。既に生成済みのコンテキスト内で出現したn-gram(n連続トークン)パターンをルックアップテーブルに保存し、現在の末尾トークン列と一致するパターンがあれば、その続きをドラフトとして提案する。

動作例

コンテキストに「東京都渋谷区」というトークン列が既に出現している場合、現在の生成位置で「東京都」が出現すると、n-gramテーブルから「渋谷区」をドラフトとして提案する。

メリットとデメリット

  • メリット: 追加メモリ不要、追加モデルのロード不要、任意のモデルに即座に適用可能
  • デメリット: 受理率が低い(40〜60%)、繰り返しの多いテキスト以外では効果が限定的、コンテキストが短い場合は効果なし

Self-Draft方式

Medusa

Meta Researchが2024年に発表した手法。ターゲットモデルの最終隠れ状態に複数の小型予測ヘッド(MLP 1〜2層)を追加し、各ヘッドが異なる将来位置のトークンを同時に予測する。

  • ヘッド数: 通常3〜5個(将来3〜5トークンを同時予測)
  • 追加パラメータ: ターゲットモデルの1〜3%
  • 訓練: ターゲットモデルは凍結し、追加ヘッドのみをファインチューニング(数時間〜1日)
  • 受理率: 75〜85%

Eagle(EAGLE-2)

Microsoftが2024年に発表した手法。ターゲットモデルの最終隠れ状態と直前の埋め込みを入力とする自己回帰ヘッドを追加し、ドラフトトークンを逐次生成する。

  • 追加パラメータ: ターゲットモデルの約0.5%
  • 訓練: 同上(ターゲット凍結、ヘッドのみ訓練)
  • 受理率: 80〜90%(Medusaより高い傾向)
  • 特徴: 木構造探索(Tree Attention)との組み合わせでさらに高速化

パフォーマンスチューニング

投機トークン数(K)の最適化

投機トークン数Kは受理率αとのトレードオフ。最適なKは以下の式で近似される。

最適K ≈ -1 / ln(α)

受理率α最適K期待生成トークン/ステップ
0.99〜109.0
0.84〜54.0
0.732.6
0.621.8
0.51〜21.3

動的K調整

最新の実装では、受理率をリアルタイムに監視し、Kを動的に調整する手法が採用されている。受理率が高い区間ではKを増やし、低い区間では減らすことで、全体の効率を5〜10%改善できる。

よくある質問(FAQ)

Q1: ドラフトモデルはターゲットモデルと同じファミリーでなければなりませんか? A: 必須ではないが、同一ファミリーが最も高い受理率を達成する。異なるファミリーの場合、tokenizerの一致が必要で、受理率は10〜20ポイント低下する傾向がある。

Q2: ドラフトモデルのファインチューニングは必要ですか? A: 汎用的な利用ではそのまま使用可能。特定ドメイン(医療、法律、コードなど)でターゲットモデルをファインチューンしている場合、ドラフトモデルも同じデータでファインチューンすると受理率が10〜15ポイント向上する。

Q3: ドラフトモデルを複数使うことはできますか? A: 可能。Cascade Speculative Decoding(2024年提案)では、複数サイズのドラフトモデルを連鎖させる。例: 0.5B → 2B → 70Bの3段階。最初の段階で明らかに異なるトークンを除外し、中間モデルで精度を上げる。

まとめ

  • ドラフトモデルは投機的デコーディングの「下書き生成」を担う小型言語モデル
  • 同一ファミリーの小型モデル(パラメータ比1/5〜1/50)が最も高い受理率を達成
  • n-gram方式は追加メモリ不要だが受理率は低い(40〜60%)
  • Self-Draft方式(Medusa/Eagle)はターゲットモデルに小型ヘッドを追加する軽量アプローチ
  • ドラフトモデルの品質がシステム全体の性能を決定するため、適切な選択と蒸留訓練が重要