speculative-decoding

投機的デコーディング：LLM推論の高速化

このスキルを使用するタイミング

次のような場合に投機的デコーディングを使用してください：

• 品質を損なうことなく、推論を1.5〜3.6倍高速化したい
• リアルタイムアプリケーション（チャットボット、コード生成）のレイテンシを削減したい
• 大量のリクエストを処理するサービングのスループットを最適化したい
• 限られたハードウェアで効率的にデプロイしたい
• モデルアーキテクチャを変更せずに生成を高速化したい

主要な手法: ドラフトモデルによる投機的デコーディング、Medusa（複数ヘッド）、Lookahead Decoding（ヤコビ反復）

論文: Medusa（arXiv 2401.10774）、Lookahead Decoding（ICML 2024）、Speculative Decoding Survey（ACL 2024）

インストール

# Standard speculative decoding (transformers)
pip install transformers accelerate

# Medusa (multiple decoding heads)
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 .

# Optional: vLLM with speculative decoding
pip install vllm

クイックスタート

基本的な投機的デコーディング（ドラフトモデル）

from transformers import AutoModelForCausalLM, AutoTokenizer

# Load target model (large, slow)
target_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-70b-hf",
    device_map="auto",
    torch_dtype=torch.float16
)

# Load draft model (small, fast)
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")

# Generate with speculative decoding
prompt = "Explain quantum computing in simple terms:"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

# Transformers 4.36+ supports assisted generation
outputs = target_model.generate(
    **inputs,
    assistant_model=draft_model,  # Enable 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

# Load Medusa-enhanced model
model = MedusaModel.from_pretrained(
    "FasterDecoding/medusa-vicuna-7b-v1.3",  # Pre-trained with Medusa heads
    torch_dtype=torch.float16,
    device_map="auto"
)

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

# Generate with Medusa (2-3× speedup)
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,  # Acceptance threshold
    posterior_alpha=0.3,       # Tree construction parameter
)

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

Lookahead Decoding（ヤコビ反復）

from lookahead.lookahead_decoding import LookaheadDecoding

# Load model
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")

# Initialize lookahead decoding
lookahead = LookaheadDecoding(
    model=model,
    tokenizer=tokenizer,
    window_size=15,    # Lookahead window (W)
    ngram_size=5,      # N-gram size (N)
    guess_size=5       # Number of parallel guesses
)

# Generate (1.5-2.3× speedup)
prompt = "Implement quicksort in Python:"
output = lookahead.generate(prompt, max_new_tokens=256)
print(output)

コアコンセプト

1. 投機的デコーディング（ドラフトモデル）

考え方: 小さなドラフトモデルで候補を生成し、大きなターゲットモデルで並列に検証します。

アルゴリズム:

1. ドラフトモデルが投機的にK個のトークンを生成する
2. ターゲットモデルがK個のトークンすべてを並列に評価する（単一のフォワードパス）
3. ドラフトとターゲットが一致するトークンを受理する
4. 最初の不一致を棄却し、そこから続行する

def speculative_decode(target_model, draft_model, prompt, K=4):
    """Speculative decoding algorithm."""
    # 1. Generate K draft tokens
    draft_tokens = draft_model.generate(prompt, max_new_tokens=K)

    # 2. Target model evaluates all K tokens in one forward pass
    target_logits = target_model(draft_tokens)  # Parallel!

    # 3. Accept/reject based on probability match
    accepted = []
    for i in range(K):
        p_draft = softmax(draft_model.logits[i])
        p_target = softmax(target_logits[i])

        # Acceptance probability
        if random.random() < min(1, p_target[draft_tokens[i]] / p_draft[draft_tokens[i]]):
            accepted.append(draft_tokens[i])
        else:
            break  # Reject, resample from target

    return accepted

性能:

• 高速化: 適切なドラフトモデルで1.5〜2倍
• 品質の損失はゼロ（数学的にターゲットモデルと等価）
• ドラフトモデルがターゲットモデルの5〜10倍小さい場合に最も効果的

2. Medusa（複数のデコーディングヘッド）

出典: arXiv 2401.10774（2024）

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

アーキテクチャ:

Input → Base LLM (frozen) → Hidden State
                                ├→ Head 1 (predicts token t+1)
                                ├→ Head 2 (predicts token t+2)
                                ├→ Head 3 (predicts token t+3)
                                └→ Head 4 (predicts token t+4)

学習:

• Medusa-1: ベースLLMを凍結し、ヘッドのみを学習
  • 2.2倍の高速化、無損失
• Medusa-2: ベースLLMとヘッドを一緒にファインチューニング
  • 2.3〜3.6倍の高速化、より高い品質

ツリーベースのアテンション:

# Medusa constructs tree of candidates
# Example: Predict 2 steps ahead with top-2 per step

#         Root
#        /    \
#      T1a    T1b  (Step 1: 2 candidates)
#     /  \    / \
#  T2a  T2b T2c T2d  (Step 2: 4 candidates total)

# Single forward pass evaluates entire tree!

利点:

• 別個のドラフトモデルが不要
• 最小限の学習（ヘッドのみ）
• 任意のLLMと互換性がある

3. Lookahead Decoding（ヤコビ反復）

出典: ICML 2024

中核的な考え方: 自己回帰デコーディングを連立方程式の求解として再定式化し、ヤコビ反復を用いて並列に解きます。

数学的定式化:

Traditional:  y_t = f(x, y_1, ..., y_{t-1})  (sequential)
Jacobi:       y_t^{(k+1)} = f(x, y_1^{(k)}, ..., y_{t-1}^{(k)})  (parallel)

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 window
        self.N = ngram_size   # N-gram size

    def generate_step(self, tokens):
        # Lookahead branch: Generate W × N candidates
        candidates = {}
        for w in range(1, self.W + 1):
            for n in range(1, self.N + 1):
                # Generate n-gram starting at position w
                ngram = self.generate_ngram(tokens, start=w, length=n)
                candidates[(w, n)] = ngram

        # Verification branch: Find matching n-grams
        verified = []
        for ngram in candidates.values():
            if ngram[0] == tokens[-1]:  # First token matches last input
                if self.verify(tokens, ngram):
                    verified.append(ngram)

        # Accept longest verified n-gram
        return max(verified, key=len) if verified else [self.model.generate_next(tokens)]

性能:

• 高速化: 1.5〜2.3倍（コード生成では最大3.6倍）
• ドラフトモデルも学習も不要
• 任意のモデルでそのまま動作する

手法の比較

手法	高速化	学習の必要性	ドラフトモデル	品質の損失
ドラフトモデル型投機的デコーディング	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. Load base model
base_model = AutoModelForCausalLM.from_pretrained(
    "lmsys/vicuna-7b-v1.3",
    torch_dtype=torch.float16
)

# 2. Add Medusa heads
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. Training loop (freeze base model for Medusa-1)
for param in base_model.parameters():
    param.requires_grad = False  # Freeze base

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

for batch in dataloader:
    # Forward pass
    hidden_states = base_model(**batch, output_hidden_states=True).hidden_states[-1]

    # Predict future tokens with each head
    loss = 0
    for i, head in enumerate(medusa_heads):
        logits = head(hidden_states)
        # Target: tokens shifted by (i+1) positions
        target = batch['input_ids'][:, i+1:]
        loss += F.cross_entropy(logits[:, :-i-1], target)

    # Backward
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

ハイブリッド：投機的デコーディング + Medusa

# Use Medusa as draft model for speculative decoding
draft_medusa = MedusaModel.from_pretrained("medusa-vicuna-7b")
target_model = AutoModelForCausalLM.from_pretrained("vicuna-33b")

# Draft generates multiple candidates with Medusa
draft_tokens = draft_medusa.medusa_generate(prompt, max_new_tokens=5)

# Target verifies in single forward pass
outputs = target_model.generate(
    prompt,
    assistant_model=draft_medusa,  # Use Medusa as draft
    max_new_tokens=256
)

# Combines benefits: Medusa speed + large model quality

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

def select_draft_model(target_model_size, target):
    """Select optimal draft model for speculative decoding."""
    # Rule: Draft should be 5-10× smaller
    if target_model_size == "70B":
        return "7B"  # 10× smaller
    elif target_model_size == "33B":
        return "7B"  # 5× smaller
    elif target_model_size == "13B":
        return "1B"  # 13× smaller
    else:
        return None  # Target too small, use Medusa/Lookahead instead

# Example
draft = select_draft_model("70B", target_model)
# Returns "7B" → Use Llama-2-7b as draft for Llama-2-70b

ベストプラクティス

1. 適切な手法を選択する

# New deployment → Medusa (best overall speedup, no draft model)
if deploying_new_model:
    use_method = "Medusa"

# Existing deployment with small model available → Draft speculative
elif have_small_version_of_model:
    use_method = "Draft Model Speculative"

# Want zero training/setup → Lookahead
elif want_plug_and_play:
    use_method = "Lookahead Decoding"

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

ドラフトモデル型投機的デコーディング:

# K = number of speculative tokens
K = 4  # Good default
K = 2  # Conservative (higher acceptance)
K = 8  # Aggressive (lower acceptance, but more when accepted)

# Rule: Larger K → more speedup IF draft model is good

Medusa:

# Posterior threshold (acceptance confidence)
posterior_threshold = 0.09  # Standard (from paper)
posterior_threshold = 0.05  # More conservative (slower, higher quality)
posterior_threshold = 0.15  # More aggressive (faster, may degrade quality)

# Tree depth (how many steps ahead)
medusa_choices = [[0], [0, 0], [0, 1], [0, 0, 0]]  # Depth 3 (standard)

Lookahead:

# Window size W (lookahead distance)
# N-gram size N (context for generation)

# 7B model (more resources)
W, N = 15, 5

# 13B model (moderate)
W, N = 10, 5

# 33B+ model (limited resources)
W, N = 7, 5

3. 本番環境へのデプロイ

# vLLM with speculative decoding
from vllm import LLM, SamplingParams

# Initialize with draft model
llm = LLM(
    model="meta-llama/Llama-2-70b-hf",
    speculative_model="meta-llama/Llama-2-7b-hf",  # Draft model
    num_speculative_tokens=5,
    use_v2_block_manager=True,
)

# Generate
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 デコーディングの実装詳細