LLM推測デコーディング最適化実踐：ドラフトモデルからMedusaマルチヘッド高速化まで

AI与大数据

概要

  • 推測デコーディングはLLM推論高速化の切り牌：2-3倍の高速化、精度の損失ゼロ、2026年には推論サービスの標準機能に
  • 3つの主要な推測戦略：ドラフトモデル、Medusaマルチヘッド、ツリー推測、それぞれ最適なシナリオがある
  • ドラフトモデル選定のポイント：ドラフトモデル速度は目標モデルの5倍以上、受容率70%以上で2倍以上の高速化を実現
  • Medusaマルチヘッドは追加モデル不要：目標モデルに予測ヘッドを追加、学習コストが低くデプロイが簡単
  • 本記事ではvLLM + 推測デコーディングデプロイメントと独自ドラフトモデル学習の実踐を提供

目次

推測デコーディング：LLM推論高速化の切り牌

自己回帰デコーディングのボトルネック

LLM推論の核心的なボトルネックは、自己回帰デコーディングの直列性にある：1つのトークンを生成するたびに完全な順伝播が必要で、GPU利用率が極めて低い。

モデル規模 トークンあたりのレイテンシ GPU利用率 スループット
7B 15ms 25% 40 tok/s
14B 28ms 18% 22 tok/s
70B 80ms 12% 8 tok/s
405B 250ms 5% 2.5 tok/s

推測デコーディングの核心的な考え方

┌──────────────────────────────────────────────────────────────┐
│              推測デコーディング vs 自己回帰デコーディング                             │
│                                                                │
│  自己回帰デコーディング（トークンごとに生成）                                     │
│  ┌──────────────────────────────────────────────────────┐    │
│  │  [BOS] → f() → t1 → f() → t2 → f() → t3 → f() → t4 │    │
│  │  4回の順伝播、4つのトークン                                 │    │
│  └──────────────────────────────────────────────────────┘    │
│                                                                │
│  推測デコーディング（ドラフト＋検証）                                         │
│  ┌──────────────────────────────────────────────────────┐    │
│  │  ドラフトモデル高速生成: t1' t2' t3' t4' (1回の順伝播)         │    │
│  │  目標モデル並列検証: [t1' t2' t3' t4'] → f() (1回の順伝播) │    │
│  │  仮定t1't2't3'正しく、t4'誤り：                        │    │
│  │  受容t1't2't3'、拒否t4'、修正してt4                    │    │
│  │  2回の順伝播、4つのトークン → 2倍の高速化                       │    │
│  └──────────────────────────────────────────────────────┘    │
└──────────────────────────────────────────────────────────────┘

2026年推測デコーディング手法比較

手法 高速化比 精度損失 追加VRAM デプロイ複雑度
ドラフトモデル 2-3倍 ゼロ 30-50% 中
Medusaマルチヘッド 1.5-2.5倍 ゼロ 5-10% 低
Eagleツリー推測 2.5-3.5倍 ゼロ 10-15% 中
自己推測(Self-Spec) 1.5-2倍 ゼロ 0% 低
投機サンプリング+量子化 2-3.5倍 極小 15-25% 中

推測デコーディングの原理と数学的基礎

受容率と高速化比

推測デコーディングの高速化比は、ドラフトトークンの受容率（acceptance rate）に依存する：

理論的高速化比 = γ / (1 + (1-α) × γ)

ここで：
- γ = ドラフト長（1回の推測あたりのトークン数）
- α = 受容率（ドラフトトークンが目標モデルに受容される確率）

例：
- γ=5, α=0.8: 高速化比 = 5 / (1 + 0.2×5) = 5/2 = 2.5×
- γ=5, α=0.9: 高速化比 = 5 / (1 + 0.1×5) = 5/1.5 = 3.3×
- γ=5, α=0.6: 高速化比 = 5 / (1 + 0.4×5) = 5/3 = 1.67×
- γ=3, α=0.8: 高速化比 = 3 / (1 + 0.2×3) = 3/1.6 = 1.875×

受容率と高速化比関係

受容率alpha gamma=3 gamma=5 gamma=7 gamma=10
0.5 1.2× 1.25× 1.27× 1.25×
0.6 1.36× 1.47× 1.52× 1.47×
0.7 1.58× 1.82× 1.94× 1.96×
0.8 1.88× 2.5× 2.92× 3.33×
0.9 2.31× 3.33× 4.38× 5.88×

検証アルゴリズムの実装

import torch
import torch.nn.functional as F

def speculative_decode(
    target_model,
    draft_model,
    input_ids,
    draft_length=5,
    temperature=1.0,
):
    draft_tokens = []
    draft_probs = []
    
    current_ids = input_ids.clone()
    
    for _ in range(draft_length):
        with torch.no_grad():
            outputs = draft_model(current_ids)
            next_token_logits = outputs.logits[:, -1, :]
            next_token_prob = F.softmax(next_token_logits / temperature, dim=-1)
            
        next_token = torch.multinomial(next_token_prob, num_samples=1)
        draft_tokens.append(next_token)
        draft_probs.append(next_token_prob)
        current_ids = torch.cat([current_ids, next_token], dim=-1)
    
    draft_token_ids = torch.cat(draft_tokens, dim=-1)
    
    with torch.no_grad():
        target_outputs = target_model(current_ids)
        target_logits = target_outputs.logits[:, -(draft_length + 1):, :]
        target_probs = F.softmax(target_logits / temperature, dim=-1)
    
    accepted_tokens = []
    n_accepted = 0
    
    for i in range(draft_length):
        draft_token = draft_token_ids[:, i]
        draft_prob = draft_probs[i].gather(1, draft_token.unsqueeze(-1))
        target_prob = target_probs[:, i].gather(1, draft_token.unsqueeze(-1))
        
        accept_prob = torch.min(
            torch.ones_like(target_prob),
            target_prob / (draft_prob + 1e-10),
        )
        
        if torch.rand(1, device=accept_prob.device) < accept_prob:
            accepted_tokens.append(draft_token)
            n_accepted += 1
        else:
            residual_prob = F.relu(target_probs[:, i] - draft_probs[i])
            residual_prob = residual_prob / residual_prob.sum(dim=-1, keepdim=True)
            corrected_token = torch.multinomial(residual_prob, num_samples=1)
            accepted_tokens.append(corrected_token)
            break
    
    if n_accepted == draft_length:
        bonus_prob = target_probs[:, -1]
        bonus_token = torch.multinomial(bonus_prob, num_samples=1)
        accepted_tokens.append(bonus_token)
    
    result = torch.cat(accepted_tokens, dim=-1)
    return result, n_accepted

ドラフトモデル選定と学習

ドラフトモデル選定の原則

原則 要件 理由
速度 ≥5倍目標モデル ドラフト生成がボトルネックにならないこと
受容率 ≥70% 70%未満では2倍の高速化に届かない
サイズ ≤1/5目標モデル VRAMオーバーヘッドを制御可能
語彙 目標モデルと同一 語彙マッピングの損失を回避

主要なドラフトモデルの組み合わせ

目標モデル ドラフトモデル 速度比 受容率 高速化比
LLaMA-3-70B LLaMA-3-8B 6× 75% 2.3×
Qwen2.5-72B Qwen2.5-7B 5.5× 78% 2.5×
DeepSeek-V3 DeepSeek-V2-Lite 7× 72% 2.2×
Mistral-Large Mistral-7B 6× 76% 2.4×
Gemma-2-27B Gemma-2-2B 8× 70% 2.1×

ドラフトモデル学習

from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer
from datasets import load_dataset

class DraftModelTrainer:
    def __init__(self, target_model_path, draft_model_path, output_path):
        self.target = AutoModelForCausalLM.from_pretrained(
            target_model_path, torch_dtype=torch.float16, device_map="auto"
        )
        self.draft = AutoModelForCausalLM.from_pretrained(
            draft_model_path, torch_dtype=torch.float16, device_map="auto"
        )
        self.tokenizer = AutoTokenizer.from_pretrained(target_model_path)
        self.output_path = output_path
    
    def generate_distillation_data(self, prompts, num_samples=50000):
        distill_data = []
        
        for prompt in prompts:
            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.target.device)
            
            with torch.no_grad():
                outputs = self.target.generate(
                    **inputs,
                    max_new_tokens=256,
                    temperature=0.7,
                    do_sample=True,
                    num_return_sequences=5,
                )
            
            for output in outputs:
                text = self.tokenizer.decode(output, skip_special_tokens=True)
                distill_data.append({"text": text})
            
            if len(distill_data) >= num_samples:
                break
        
        return distill_data
    
    def train(self, distill_data, epochs=3, lr=5e-5):
        from trl import SFTTrainer, SFTConfig
        
        sft_config = SFTConfig(
            output_dir=self.output_path,
            num_train_epochs=epochs,
            per_device_train_batch_size=4,
            gradient_accumulation_steps=8,
            learning_rate=lr,
            warmup_ratio=0.03,
            bf16=True,
            logging_steps=100,
            save_strategy="epoch",
            max_seq_length=2048,
        )
        
        trainer = SFTTrainer(
            model=self.draft,
            train_dataset=distill_data,
            args=sft_config,
        )
        
        trainer.train()
        self.draft.save_pretrained(self.output_path)
        return self.draft

ドラフトモデル受容率最適化

最適化手法 受容率向上 方法
蒸餸ファインチューニング +10-15% 目標モデルの出力でドラフトモデルを学習
語彙アライメント +5-8% ドラフトと目標の語彙を一致させる
温度アライメント +3-5% ドラフトと目標のサンプリング温度を一致させる
コンテキストアライメント +5-10% ドラフトモデルに同じsystem promptを使用
組み合わせ最適化 +20-30% 上記すべて

Medusaマルチヘッド推測実踐

Medusaアーキテクチャ

┌──────────────────────────────────────────────────────────────┐
│              Medusaマルチヘッドアーキテクチャ                                    │
│                                                                │
│  オリジナルLLMバックボーン                                            │
│  ┌──────────────────────────────────────────────────────┐    │
│  │  Input → Transformer Layers → Hidden State            │    │
│  └──────────────────────────────────────────────────────┘    │
│                         ↓                                     │
│  オリジナルLM Head (t+1を予測)                                       │
│  ┌──────────┐                                                │
│  │ LM Head  │ → P(t+1)                                      │
│  └──────────┘                                                │
│                                                                │
│  Medusa Head 0 (t+2を予測)                                     │
│  ┌──────────────┐                                            │
│  │ Linear+Softmax│ → P(t+2)                                 │
│  └──────────────┘                                            │
│                                                                │
│  Medusa Head 1 (t+3を予測)                                     │
│  ┌──────────────┐                                            │
│  │ Linear+Softmax│ → P(t+3)                                 │
│  └──────────────┘                                            │
│                                                                │
│  Medusa Head K-1 (t+Kを予測)                                   │
│  ┌──────────────┐                                            │
│  │ Linear+Softmax│ → P(t+K)                                 │
│  └──────────────┘                                            │
└──────────────────────────────────────────────────────────────┘

Medusa実装

import torch
import torch.nn as nn
from transformers import PreTrainedModel

class MedusaHead(nn.Module):
    def __init__(self, hidden_size, vocab_size):
        super().__init__()
        self.linear = nn.Linear(hidden_size, vocab_size, bias=False)
    
    def forward(self, hidden_states):
        return self.linear(hidden_states)

class MedusaModel(nn.Module):
    def __init__(self, base_model, num_heads=4):
        super().__init__()
        self.base_model = base_model
        self.num_heads = num_heads
        
        hidden_size = base_model.config.hidden_size
        vocab_size = base_model.config.vocab_size
        
        self.medusa_heads = nn.ModuleList([
            MedusaHead(hidden_size, vocab_size) 
            for _ in range(num_heads)
        ])
        
        for head in self.medusa_heads:
            nn.init.xavier_uniform_(head.linear.weight)
    
    def forward(self, input_ids, attention_mask=None):
        outputs = self.base_model.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            use_cache=True,
        )
        
        hidden_states = outputs.last_hidden_state
        
        base_logits = self.base_model.lm_head(hidden_states)
        
        medusa_logits = []
        for head in self.medusa_heads:
            medusa_logits.append(head(hidden_states))
        
        return base_logits, medusa_logits, outputs.past_key_values

class MedusaDecoding:
    def __init__(self, model, tokenizer, num_heads=4, top_k=10):
        self.model = model
        self.tokenizer = tokenizer
        self.num_heads = num_heads
        self.top_k = top_k
    
    def generate(self, input_ids, max_new_tokens=256):
        generated = input_ids.clone()
        
        for _ in range(max_new_tokens):
            base_logits, medusa_logits, _ = self.model(generated)
            
            base_prob = torch.softmax(base_logits[:, -1, :], dim=-1)
            next_token = torch.argmax(base_prob, dim=-1, keepdim=True)
            
            candidates = [next_token]
            
            for head_idx in range(self.num_heads):
                head_prob = torch.softmax(
                    medusa_logits[head_idx][:, -1, :], dim=-1
                )
                top_tokens = torch.topk(head_prob, self.top_k)
                candidates.append(top_tokens.indices[:, :1])
            
            candidate_ids = torch.cat(candidates, dim=-1)
            
            with torch.no_grad():
                verify_outputs = self.model.base_model(candidate_ids)
                verify_logits = self.model.base_model.lm_head(
                    verify_outputs.last_hidden_state
                )
            
            accepted = [next_token]
            for i, cand in enumerate(candidates[1:]):
                verify_prob = torch.softmax(verify_logits[:, i, :], dim=-1)
                if verify_prob.gather(1, cand).item() > 0.1:
                    accepted.append(cand)
                else:
                    corrected = torch.argmax(verify_prob, dim=-1, keepdim=True)
                    accepted.append(corrected)
                    break
            
            new_tokens = torch.cat(accepted, dim=-1)
            generated = torch.cat([generated, new_tokens], dim=-1)
            
            if new_tokens.shape[-1] > 1:
                break
        
        return generated

Medusa学習

class MedusaTrainer:
    def __init__(self, medusa_model, lr=1e-4):
        self.model = medusa_model
        self.optimizer = torch.optim.AdamW(
            [p for p in self.model.medusa_heads.parameters()],
            lr=lr,
            weight_decay=0.01,
        )
        self.base_model = medusa_model.base_model
        for param in self.base_model.parameters():
            param.requires_grad = False
    
    def train_step(self, input_ids, labels):
        with torch.no_grad():
            outputs = self.base_model.model(input_ids=input_ids)
            hidden_states = outputs.last_hidden_state
        
        total_loss = 0
        for head_idx, head in enumerate(self.model.medusa_heads):
            shift = head_idx + 1
            head_logits = head(hidden_states[:, :-shift, :])
            head_labels = labels[:, shift:]
            
            loss = F.cross_entropy(
                head_logits.reshape(-1, head_logits.size(-1)),
                head_labels.reshape(-1),
            )
            total_loss += loss
        
        total_loss /= len(self.model.medusa_heads)
        total_loss.backward()
        self.optimizer.step()
        self.optimizer.zero_grad()
        
        return total_loss.item()
Medusa設定 ヘッド数 Top-K 高速化比 追加VRAM
Medusa-2 2 5 1.5× 3%
Medusa-4 4 10 2.0× 5%
Medusa-4 4 20 2.2× 5%
Medusa-8 8 10 2.5× 10%

ツリー推測とEagleアプローチ

Eagleツリー推測の原理

┌──────────────────────────────────────────────────────────────┐
│              Eagleツリー推測                                     │
│                                                                │
│  従来の線形推測:                                                │
│  t1 → t2 → t3 → t4 → t5 (1つのパス)                           │
│                                                                │
│  Eagleツリー推測:                                               │
│              t1                                               │
│            /    \                                             │
│          t2a    t2b                                           │
│         /  \    /  \                                          │
│       t3a t3b t3c t3d                                        │
│                                                                │
│  複数の候補パスを並列検証し、受容率が最も高いパスを選択                      │
│  利点：あるパスが拒否されても、他のパスが受容される可能性がある                  │
└──────────────────────────────────────────────────────────────┘

Eagle実装

class EagleSpeculativeDecoder:
    def __init__(self, target_model, draft_model, tree_width=4, tree_depth=3):
        self.target = target_model
        self.draft = draft_model
        self.tree_width = tree_width
        self.tree_depth = tree_depth
    
    def build_speculation_tree(self, input_ids):
        tree_nodes = [{"ids": input_ids.clone(), "parent": None}]
        leaves = []
        
        for depth in range(self.tree_depth):
            new_nodes = []
            for node in tree_nodes[-self.tree_width ** depth:]:
                with torch.no_grad():
                    outputs = self.draft(node["ids"])
                    probs = torch.softmax(outputs.logits[:, -1, :], dim=-1)
                    top_k = torch.topk(probs, self.tree_width)
                
                for i in range(self.tree_width):
                    child_ids = torch.cat([
                        node["ids"],
                        top_k.indices[:, i:i+1],
                    ], dim=-1)
                    child = {
                        "ids": child_ids,
                        "parent": node,
                        "prob": top_k.values[:, i].item(),
                        "token": top_k.indices[:, i],
                    }
                    new_nodes.append(child)
                    
                    if depth == self.tree_depth - 1:
                        leaves.append(child)
            
            tree_nodes.extend(new_nodes)
        
        return tree_nodes, leaves
    
    def verify_tree(self, tree_nodes, leaves):
        best_path = []
        best_score = 0
        
        for leaf in leaves:
            path = []
            node = leaf
            while node["parent"] is not None:
                path.append(node)
                node = node["parent"]
            path.reverse()
            
            with torch.no_grad():
                full_ids = leaf["ids"]
                outputs = self.target(full_ids)
                target_probs = torch.softmax(outputs.logits, dim=-1)
            
            score = 1.0
            accepted = 0
            for i, step in enumerate(path):
                target_prob = target_probs[:, -(len(path) - i), :].gather(
                    1, step["token"]
                ).item()
                score *= target_prob
                if target_prob > 0.3:
                    accepted += 1
                else:
                    break
            
            if score > best_score:
                best_score = score
                best_path = path[:accepted]
        
        return best_path

ツリー推測 vs 線形推測

手法 候補パス 受容率 高速化比 計算オーバーヘッド
線形gamma=5 1 70% 2.3倍 低
ツリーw=2,d=3 8 82% 2.8倍 中
ツリーw=3,d=3 27 88% 3.2倍 中高
ツリーw=4,d=3 64 92% 3.5倍 高
Eagle(w=4,d=4) 256 95% 3.8倍 高

本番デプロイ：vLLM + 推測デコーディング

vLLM推測デコーディング設定

from vllm import LLM, SamplingParams

llm = LLM(
    model="Qwen/Qwen2.5-72B-Instruct",
    speculative_model="Qwen/Qwen2.5-7B-Instruct",
    num_speculative_tokens=5,
    speculative_max_model_len=4096,
    gpu_memory_utilization=0.9,
    tensor_parallel_size=4,
    trust_remote_code=True,
)

sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=512,
)

outputs = llm.generate(
    ["量子コンピューティングの基本原理を説明してください"],
    sampling_params=sampling_params,
)

for output in outputs:
    print(output.outputs[0].text)

vLLM推測デコーディングパフォーマンス

設定 目標モデル ドラフトモデル スループット(tok/s) 高速化比
推測なし 72B なし 8 1倍
ドラフト推測 72B 7B 18 2.25倍
ドラフト推測 72B 0.5B 22 2.75倍
Medusa-4 72B なし(内蔵) 16 2.0倍
Eagle 72B 7B 25 3.1倍

本番デプロイの注意事項

注意事項 説明 解決策
ドラフトモデルVRAM 追加で30-50%占有 ドラフトモデルをINT8に量子化
受容率の変動 タスクによって受容率の差が大きい ドラフト長を動的に調整
KV Cache管理 拒否されたトークンのKVをクリーンアップ vLLMが自動管理
バッチ処理互換性 推測デコーディングと連続バッチ処理の競合 vLLMが対応済み
初回トークンレイテンシ 推測デコーディングが初回トークンレイテンシを増加 チャットシナリオでは推測を無効化

動的推測戦略

class DynamicSpeculativeConfig:
    def __init__(self):
        self.min_draft_length = 2
        self.max_draft_length = 8
        self.target_acceptance_rate = 0.75
        self.current_draft_length = 5
        self.history = []
    
    def update(self, acceptance_rate):
        self.history.append(acceptance_rate)
        if len(self.history) > 100:
            self.history = self.history[-100:]
        
        avg_rate = sum(self.history) / len(self.history)
        
        if avg_rate > 0.85:
            self.current_draft_length = min(
                self.current_draft_length + 1,
                self.max_draft_length,
            )
        elif avg_rate < 0.65:
            self.current_draft_length = max(
                self.current_draft_length - 1,
                self.min_draft_length,
            )
        
        return self.current_draft_length

まとめと関連情報

重要ポイントの振り返り

  1. 推測デコーディングはLLM推論高速化の無料ランチ：2-3倍の高速化、精度の損失ゼロ
  2. ドラフトモデル：選定のポイントは速度5倍以上＋受容率70%以上、蒸餸ファインチューニングで受容率を20-30%向上可能
  3. Medusaマルチヘッド：追加モデル不要、5%のVRAMオーバーヘッドで2倍の高速化、デプロイが最も簡単
  4. Eagleツリー推測：マルチパス並列検証、3.5×+加速、高スループットシナリオに最適

手法の推奨

シナリオ 推奨手法 高速化比
迅速な導入 Medusa-4 2.0倍
汎用推論 ドラフトモデル(7B→72B) 2.5倍
高スループット Eagleツリー推測 3.5×
VRAM制限 Medusa+量子化ドラフト 2.2倍

推論データのフォーマット変換が必要ですか？JSON→YAMLツールとテキスト差分比較で推測デコーディングのデバッグデータを素早く処理できます。

関連記事

ブラウザローカルツールを無料で試す →

#推测解码#LLM推理加速#Speculative Decoding#Medusa头#大模型推理优化#2026