Speculative Decoding: LLMの推論を高速化する

このスキルを使う場合

Speculative Decodingは以下が必要な場合に使用します:

• 品質低下なしで推論を1.5-3.6倍高速化したい
• リアルタイムアプリケーション(チャットボット、コード生成)の遅延を削減したい
• 高容量提供でスループットを最適化したい
• 限定ハードウェアで効率的にデプロイしたい
• モデルアーキテクチャを変更せずにより高速に生成したい

主要技術: ドラフトモデル speculative decoding、Medusa(複数ヘッド)、Lookahead Decoding(Jacobi反復)

論文: Medusa(arXiv 2401.10774)、Lookahead Decoding(ICML 2024)、Speculative Decoding Survey(ACL 2024)

インストール

# 標準的なspeculative decoding (transformers)
pip install transformers accelerate

# Medusa (複数デコーディングヘッド)
git clone https://github.com/FasterDecoding/Medusa
cd Medusa
pip install -e .

# Lookahead Decoding
git clone https://github.com/hao-ai-lab/LookaheadDecoding
cd LookaheadDecoding
pip install -e .

# オプション: vLLM with speculative decoding
pip install vllm

クイックスタート

基本的なSpeculative Decoding (ドラフトモデル)

from transformers import AutoModelForCausalLM, AutoTokenizer

# ターゲットモデルをロード (大規模、遅い)
target_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-70b-hf",
    device_map="auto",
    torch_dtype=torch.float16
)

# ドラフトモデルをロード (小規模、高速)
draft_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    device_map="auto",
    torch_dtype=torch.float16
)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-70b-hf")

# Speculative decodingで生成
prompt = "Explain quantum computing in simple terms:"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

# Transformers 4.36+ はassisted generationに対応
outputs = target_model.generate(
    **inputs,
    assistant_model=draft_model,  # Speculative decodingを有効化
    max_new_tokens=256,
    do_sample=True,
    temperature=0.7,
)

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response)

Medusa (複数デコーディングヘッド)

from medusa.model.medusa_model import MedusaModel

# Medusa拡張モデルをロード
model = MedusaModel.from_pretrained(
    "FasterDecoding/medusa-vicuna-7b-v1.3",  # Medusaヘッド付きで事前学習済み
    torch_dtype=torch.float16,
    device_map="auto"
)

tokenizer = AutoTokenizer.from_pretrained("FasterDecoding/medusa-vicuna-7b-v1.3")

# Medusaで生成 (2-3倍高速化)
prompt = "Write a Python function to calculate fibonacci numbers:"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

outputs = model.medusa_generate(
    **inputs,
    max_new_tokens=256,
    temperature=0.7,
    posterior_threshold=0.09,  # 受け入れ閾値
    posterior_alpha=0.3,       # ツリー構築パラメータ
)

response = tokenizer.decode(outputs[0], skip_special_tokens=True)

Lookahead Decoding (Jacobi反復)

from lookahead.lookahead_decoding import LookaheadDecoding

# モデルをロード
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    torch_dtype=torch.float16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# Lookahead decodingを初期化
lookahead = LookaheadDecoding(
    model=model,
    tokenizer=tokenizer,
    window_size=15,    # Lookaheadウィンドウ (W)
    ngram_size=5,      # N-gramサイズ (N)
    guess_size=5       # 並列推測数
)

# 生成 (1.5-2.3倍高速化)
prompt = "Implement quicksort in Python:"
output = lookahead.generate(prompt, max_new_tokens=256)
print(output)

コアコンセプト

1. Speculative Decoding (ドラフトモデル)

アイデア: 小規模ドラフトモデルで候補トークンを生成し、大規模ターゲットモデルで並列に検証します。

アルゴリズム:

1. ドラフトモデルがK個のトークンを推測的に生成
2. ターゲットモデルがすべてのK個のトークンを並列に評価(単一フォワードパス)
3. ドラフトとターゲットが一致するトークンを受け入れ
4. 最初の不一致で拒否し、そこから続行

def speculative_decode(target_model, draft_model, prompt, K=4):
    """Speculative decodingアルゴリズム."""
    # 1. K個のドラフトトークンを生成
    draft_tokens = draft_model.generate(prompt, max_new_tokens=K)

    # 2. ターゲットモデルが1つのフォワードパスですべてのK個のトークンを評価
    target_logits = target_model(draft_tokens)  # 並列!

    # 3. 確率マッチに基づいて受け入れ/拒否
    accepted = []
    for i in range(K):
        p_draft = softmax(draft_model.logits[i])
        p_target = softmax(target_logits[i])

        # 受け入れ確率
        if random.random() < min(1, p_target[draft_tokens[i]] / p_draft[draft_tokens[i]]):
            accepted.append(draft_tokens[i])
        else:
            break  # 拒否、ターゲットから再サンプリング

    return accepted

パフォーマンス:

• 高速化: 優れたドラフトモデルで1.5-2倍
• 品質低下なし(数学的にターゲットモデルと同等)
• ドラフトモデルがターゲットより5-10倍小さい場合に最適

2. Medusa (複数デコーディングヘッド)

出典: arXiv 2401.10774 (2024)

革新: 既存モデルに複数の予測ヘッドを追加し、別のドラフトモデルなしで将来トークンを予測します。

アーキテクチャ:

Input → Base LLM (frozen) → Hidden State
                                ├→ Head 1 (トークンt+1を予測)
                                ├→ Head 2 (トークンt+2を予測)
                                ├→ Head 3 (トークンt+3を予測)
                                └→ Head 4 (トークンt+4を予測)

学習:

• Medusa-1: ベースLLMをフリーズ、ヘッドのみ学習
  • 2.2倍高速化、損失なし
• Medusa-2: ベースLLM+ヘッドを一緒に微調整
  • 2.3-3.6倍高速化、より高品質

ツリーベースの注意:

# Medusaが候補のツリーを構築
# 例: 各ステップでtop-2で2ステップ先を予測

#         Root
#        /    \
#      T1a    T1b  (ステップ1: 2候補)
#     /  \    / \
#  T2a  T2b T2c T2d  (ステップ2: 合計4候補)

# 単一のフォワードパスがツリー全体を評価!

利点:

• 別のドラフトモデルが不要
• 最小限の学習(ヘッドのみ)
• あらゆるLLMと互換性

3. Lookahead Decoding (Jacobi反復)

出典: ICML 2024

コアアイデア: 自己回帰デコーディングを方程式系を解くこととして再構成し、Jacobi反復を使用して並列に解きます。

数学的定式化:

従来:      y_t = f(x, y_1, ..., y_{t-1})  (順序的)
Jacobi:    y_t^{(k+1)} = f(x, y_1^{(k)}, ..., y_{t-1}^{(k)})  (並列)

2つのブランチ:

1. Lookaheadブランチ: n-gramを並列で生成

   • ウィンドウサイズW: 何ステップ先を見るか
   • N-gramサイズN: 過去に何トークン使用するか

2. 検証ブランチ: 見込みのあるn-gramを検証

   • n-gramを生成されたトークンと照合
   • 最初のトークンが一致すれば受け入れ

class LookaheadDecoding:
    def __init__(self, model, window_size=15, ngram_size=5):
        self.model = model
        self.W = window_size  # Lookaheadウィンドウ
        self.N = ngram_size   # N-gramサイズ

    def generate_step(self, tokens):
        # Lookaheadブランチ: W × N候補を生成
        candidates = {}
        for w in range(1, self.W + 1):
            for n in range(1, self.N + 1):
                # 位置wで開始するn-gramを生成
                ngram = self.generate_ngram(tokens, start=w, length=n)
                candidates[(w, n)] = ngram

        # 検証ブランチ: マッチするn-gramを見つける
        verified = []
        for ngram in candidates.values():
            if ngram[0] == tokens[-1]:  # 最初のトークンが最後の入力と一致
                if self.verify(tokens, ngram):
                    verified.append(ngram)

        # 最長の検証されたn-gramを受け入れ
        return max(verified, key=len) if verified else [self.model.generate_next(tokens)]

パフォーマンス:

• 高速化: 1.5-2.3倍(コード生成では最大3.6倍)
• ドラフトモデルや学習が不要
• あらゆるモデルで既製のまま動作

メソッド比較

メソッド	高速化	学習が必要	ドラフトモデル	品質低下
ドラフトモデル Speculative	1.5-2×	いいえ	はい(外部)	なし
Medusa	2-3.6×	最小限(ヘッドのみ)	いいえ(組み込みヘッド)	なし
Lookahead	1.5-2.3×	なし	いいえ	なし
単純なバッチ処理	1.2-1.5×	いいえ	いいえ	なし

高度なパターン

Medusaヘッドの学習

from medusa.model.medusa_model import MedusaModel
from medusa.model.kv_cache import initialize_past_key_values
import torch.nn as nn

# 1. ベースモデルをロード
base_model = AutoModelForCausalLM.from_pretrained(
    "lmsys/vicuna-7b-v1.3",
    torch_dtype=torch.float16
)

# 2. Medusaヘッドを追加
num_heads = 4
medusa_heads = nn.ModuleList([
    nn.Linear(base_model.config.hidden_size, base_model.config.vocab_size, bias=False)
    for _ in range(num_heads)
])

# 3. 学習ループ (Medusa-1ではベースモデルをフリーズ)
for param in base_model.parameters():
    param.requires_grad = False  # ベースをフリーズ

optimizer = torch.optim.Adam(medusa_heads.parameters(), lr=1e-3)

for batch in dataloader:
    # フォワードパス
    hidden_states = base_model(**batch, output_hidden_states=True).hidden_states[-1]

    # 各ヘッドで将来トークンを予測
    loss = 0
    for i, head in enumerate(medusa_heads):
        logits = head(hidden_states)
        # ターゲット: (i+1)位置シフトされたトークン
        target = batch['input_ids'][:, i+1:]
        loss += F.cross_entropy(logits[:, :-i-1], target)

    # バックプロパゲーション
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

ハイブリッド: Speculative + Medusa

# Medusaをspeculative decodingのドラフトモデルとして使用
draft_medusa = MedusaModel.from_pretrained("medusa-vicuna-7b")
target_model = AutoModelForCausalLM.from_pretrained("vicuna-33b")

# ドラフトがMedusaで複数の候補を生成
draft_tokens = draft_medusa.medusa_generate(prompt, max_new_tokens=5)

# ターゲットが単一フォワードパスで検証
outputs = target_model.generate(
    prompt,
    assistant_model=draft_medusa,  # Medusaをドラフトとして使用
    max_new_tokens=256
)

# メリットを組み合わせ: Medusa速度+大規模モデル品質

最適なドラフトモデルの選択

def select_draft_model(target_model_size, target):
    """Speculative decodingの最適なドラフトモデルを選択."""
    # ルール: ドラフトはターゲットより5-10倍小さいべき
    if target_model_size == "70B":
        return "7B"  # 10倍小さい
    elif target_model_size == "33B":
        return "7B"  # 5倍小さい
    elif target_model_size == "13B":
        return "1B"  # 13倍小さい
    else:
        return None  # ターゲットが小さすぎ、Medusa/Lookaheadを代わりに使用

# 例
draft = select_draft_model("70B", target_model)
# "7B"を返す → Llama-2-70bのドラフトとしてLlama-2-7bを使用

ベストプラクティス

1. 正しいメソッドを選択

# 新規デプロイ → Medusa (最良の全体的高速化、ドラフトモデル不要)
if deploying_new_model:
    use_method = "Medusa"

# 既存デプロイで小規模版が利用可能 → ドラフト speculative
elif have_small_version_of_model:
    use_method = "Draft Model Speculative"

# ゼロ学習/セットアップを望む → Lookahead
elif want_plug_and_play:
    use_method = "Lookahead Decoding"

2. ハイパーパラメータのチューニング

ドラフトモデル Speculative:

# K = speculative トークン数
K = 4  # 良好なデフォルト
K = 2  # 保守的 (より高い受け入れ)
K = 8  # 積極的 (より低い受け入れ、受け入れ時により多く)

# ルール: ドラフトモデルが優れている場合、より大きいK → より多くの高速化

Medusa:

# Posterior閾値 (受け入れ信頼度)
posterior_threshold = 0.09  # 標準 (論文から)
posterior_threshold = 0.05  # より保守的 (遅い、より高品質)
posterior_threshold = 0.15  # より積極的 (より高速、品質低下の可能性)

# ツリー深さ (何ステップ先)
medusa_choices = [[0], [0, 0], [0, 1], [0, 0, 0]]  # 深さ3 (標準)

Lookahead:

# ウィンドウサイズW (lookahead距離)
# N-gramサイズN (生成用コンテキスト)

# 7Bモデル (より多くのリソース)
W, N = 15, 5

# 13Bモデル (中程度)
W, N = 10, 5

# 33B+モデル (限定リソース)
W, N = 7, 5

3. 本番デプロイメント

# vLLM with speculative decoding
from vllm import LLM, SamplingParams

# ドラフトモデルで初期化
llm = LLM(
    model="meta-llama/Llama-2-70b-hf",
    speculative_model="meta-llama/Llama-2-7b-hf",  # ドラフトモデル
    num_speculative_tokens=5,
    use_v2_block_manager=True,
)

# 生成
prompts = ["Tell me about AI:", "Explain quantum physics:"]
sampling_params = SamplingParams(temperature=0.7, max_tokens=256)

outputs = llm.generate(prompts, sampling_params)
for output in outputs:
    print(output.outputs[0].text)

リソース

• Medusa論文: https://arxiv.org/abs/2401.10774
• Medusa GitHub: https://github.com/FasterDecoding/Medusa
• Lookahead Decoding (ICML 2024): https://lmsys.org/blog/2023-11-21-lookahead-decoding/
• Lookahead GitHub: https://github.com/hao-ai-lab/LookaheadDecoding
• Speculative Decoding Survey (ACL 2024): https://aclanthology.org/2024.findings-acl.456.pdf
• 総合サーベイ: https://arxiv.org/abs/2401.07851

関連項目

• references/draft_model.md - ドラフトモデルの選択と学習
• references/medusa.md - Medusaアーキテクチャと学習
• references/lookahead.md - Lookahead decodingの実装詳細