ツリーベース投機的デコーディング（ツリーベーストウキテキデコーディング）
ツリーベース投機的デコーディング（Tree-based Speculative Decoding）は、ドラフトモデルが単一のトークン列ではなく複数の候補列をツリー（木）構造で同時に生成し、ターゲットモデルがツリーアテンション（Tree Attention）を用いて全候補を一括検証する投機的デコーディングの拡張手法である。SpecInfer、Sequoia、Medusa（ツリーモード）が代表的であり、線形投機と比較して採択トークン数を大幅に増加させることで3〜5倍のスピードアップを実現する。
ツリーベース投機的デコーディングとは
標準的な投機的デコーディングでは、ドラフトモデルが単一の候補トークン列（線形チェーン）を生成し、ターゲットモデルが先頭から順に検証する。この線形方式では、途中のトークンが棄却されると後続の全トークンが無駄になるため、投機長γを伸ばしても利得が飽和しやすい。
ツリーベース投機的デコーディングは、この限界を克服するために候補を木構造で展開する。各ノードが1つのトークンに対応し、ルートから葉までの各パスが1つの候補列を表す。ターゲットモデルは、ツリー全体を1回のフォワードパスで処理するツリーアテンション（Tree Attention）メカニズムにより、全候補パスを同時に検証する。
線形投機とツリー投機の比較は以下の通りである。
| 特性 | 線形投機 | ツリー投機 |
|---|---|---|
| 候補構造 | 単一チェーン（1パス） | ツリー（複数パス） |
| 投機トークン数 | γ個 | ツリーのノード数（10〜64個） |
| 検証効率 | 棄却点以降すべて無駄 | 棄却されたサブツリーのみ無駄 |
| 期待採択トークン数 | (1-α^(γ+1))/(1-α) | 線形より20〜80%多い |
| フォワードパスコスト | γ+1 トークンのバッチ | ノード数分のバッチ（ツリーアテンション） |
| 実装複雑度 | 低 | 高（ツリーアテンション、KVキャッシュ管理） |
ツリーアテンション（Tree Attention）の仕組み
ツリー構造の候補を1回のフォワードパスで処理するためには、標準的なcausal attention maskを拡張したツリーアテンション機構が必要である。通常のcausal maskでは、位置iのトークンは位置0〜iの全トークンに注意を向けるが、ツリーアテンションでは各ノードがルートからそのノードまでのパス上の祖先ノードにのみ注意を向ける。
具体的には、ツリーアテンションマスクMは以下のように構築される。
- M[i][j] = 1 : ノードjがノードiの祖先（またはi自身）である場合
- M[i][j] = 0 : それ以外
この非連続なアテンションパターンを効率的に実装するために、いくつかの最適化手法が提案されている。
| 実装手法 | 説明 | フレームワーク |
|---|---|---|
| カスタムCUDAカーネル | ツリーマスクを直接GPUカーネルに実装 | SpecInfer |
| FlashAttentionの拡張 | Block-sparseマスクとしてツリー構造をエンコード | Sequoia |
| パディング方式 | ツリーを線形に展開しパディングでマスク | Medusa（初期実装） |
| Paged KVキャッシュ統合 | vLLMのPagedAttention上にツリーを構築 | vLLM |
FlashAttention互換のツリーアテンション実装は、特にSequoiaプロジェクトで進展しており、標準的なFlashAttention-2のブロックスパースマスク機能を活用して効率的なGPUカーネル実行を実現している。
SpecInfer: ツリー投機の先駆者
SpecInfer（Miao et al., 2024）は、ツリーベース投機的デコーディングの先駆的研究であり、以下の3つの革新を導入した。
- マルチドラフトモデル: 複数の小型ドラフトモデル（Small Speculative Models, SSMs）を使用して、多様な候補ツリーを構築する。各SSMは異なる初期化やファインチューニングを持ち、候補の多様性を確保する
- トークンツリー構築: 各SSMが独立に候補列を生成し、それらをマージしてトークンツリーを構築する。共通のプレフィックスを持つ候補はツリー上で共有される
- ツリーベース並列検証: ターゲットモデルがツリー全体を1回のフォワードパスで処理し、各ノードの採択・棄却をルートから深さ優先順に決定する
SpecInferの実験結果では、Llama-2-70Bに対してLlama-2-7Bを2〜4個のSSMとして使用した構成で、標準的な線形投機デコーディングの1.5〜2倍、通常の自己回帰生成の3〜4倍のスピードアップが報告されている。
Sequoia: ハードウェア適応型動的ツリー
Sequoia（Chen et al., 2024）は、SpecInferの固定ツリー構造を動的に最適化するフレームワークであり、ハードウェア特性（GPU演算能力、メモリ帯域幅、バッチサイズ）に応じて最適なツリートポロジーを自動決定する。
Sequoiaの主要な技術的貢献は以下の通りである。
- 動的ツリートポロジー最適化: 推論環境（GPUスペック、バッチサイズ、シーケンス長）に基づいて、ツリーの幅（branching factor）と深さ（depth）を動的に決定する。A100では広いツリー、RTX 4090ではやや深いツリーが最適となる傾向がある
- Robust Speculation: ドラフトモデルの不確実性を考慮した投機戦略。エントロピーが高い位置ではより多くの分岐を生成し、低い位置では分岐を抑制することで、ツリーのノード数を効率的に配分する
- Offloading対応: モデルパラメータの一部をCPUメモリやディスクにオフロードする環境でも、ツリー投機の利点を維持する最適化
| ハードウェア | 最適ツリー構成 | スピードアップ |
|---|---|---|
| A100 (80GB) | 深さ6、幅4 | 3.5〜4.5倍 |
| H100 (80GB) | 深さ5、幅5 | 4〜5倍 |
| RTX 4090 | 深さ7、幅3 | 2.5〜3.5倍 |
| CPU offloading | 深さ8、幅2 | 6〜10倍（ベースが遅いため相対的に大きい） |
ツリー構造の最適化戦略
ツリーのトポロジー（形状）は投機的デコーディングの効率に直結する。主要な最適化戦略を以下に示す。
Top-k分岐
各ノードでドラフトモデルの上位k個のトークンを分岐として展開する。kの値はノードの深さに応じて減少させるのが一般的である（浅い層ではk=4〜5、深い層ではk=2〜3）。
エントロピーベース分岐
ドラフトモデルの出力エントロピーに基づいて分岐数を動的に調整する。低エントロピー（高確信）のノードでは分岐を抑制し、高エントロピー（低確信）のノードではより多くの代替候補を生成する。
バジェット制約最適化
ツリーの総ノード数にバジェット制約（例: 最大64ノード）を設け、制約下でツリー全体の期待採択トークン数を最大化するように分岐を割り当てる。この問題は動的計画法（DP）で解くことができ、Sequoiaはこの最適化を実行時に高速に解くアルゴリズムを実装している。
| 最適化戦略 | メリット | デメリット |
|---|---|---|
| Top-k分岐 | 実装が単純 | 不確実なノードでの分岐不足 |
| エントロピーベース | 適応的なリソース配分 | エントロピー計算のオーバーヘッド |
| バジェット制約 | 理論的最適解 | 最適化の計算コスト |
KVキャッシュ管理の課題
ツリーベース投機的デコーディングでは、KVキャッシュの管理が線形方式よりも大幅に複雑になる。ツリーの各ノードに対応するKVエントリが必要であり、検証後に採択されたパス以外のエントリをすべて削除（プルーニング）する必要がある。
主な課題と解決策は以下の通りである。
- メモリ効率: ツリーのノード数分のKVキャッシュが一時的に必要。vLLMのPagedAttentionでは、ツリーの各パスが共有プレフィックスのページを共有できるため、メモリオーバーヘッドを抑制できる
- プルーニング効率: 検証後に不採択パスのKVエントリを効率的に削除する必要がある。連続メモリ上での削除は断片化を引き起こすため、ページベースの管理が有利
- ロールバック: 再試行（retry）の際に、正しい位置までKVキャッシュをロールバックする操作が必要
よくある質問（FAQ）
Q1: ツリー投機は線形投機より常に優れていますか？
必ずしもそうではない。ツリーアテンションのフォワードパスは線形方式よりも計算コストが高く、ノード数が多いほどオーバーヘッドが増大する。採択率αが非常に高い（0.90以上）場合は線形投機でも十分な利得が得られるため、ツリーの複雑さに見合わないことがある。ツリー投機が特に有効なのは、採択率が中程度（0.5〜0.8）で分岐の多様性が利得につながるケースである。
Q2: ツリーのノード数にはどのような制限がありますか？
主な制限はGPUメモリ（KVキャッシュの一時的なオーバーヘッド）とフォワードパスの計算コストである。実用的には16〜64ノードが一般的であり、Sequoiaの実験ではA100で64ノード、RTX 4090で32ノードが効率的な上限とされている。ノード数を増やしてもフォワードパスのコスト増加により、ある点から実効スピードアップは飽和する。
Q3: ツリー投機はどのフレームワークで利用できますか？
vLLM、TensorRT-LLM、SGLangの最新バージョンでツリーベース投機的デコーディングがサポートされている。vLLMでは --spec-decode-draft-model-tp-size と --num-speculative-tokens の設定でツリー投機が有効化される。Medusaのツリーモードも広く利用可能である。Hugging Face TransformersのAssisted Generationは現時点では線形方式のみをサポートしている。
Q4: SpecInferとSequoiaの主な違いは何ですか？
SpecInferは複数の独立したドラフトモデル（SSMs）からツリーを構築するのに対し、Sequoiaは単一のドラフトモデルからハードウェア最適化されたツリーを構築する。Sequoiaはツリートポロジーの動的最適化とオフロード対応を持つ点でより実用的であり、より幅広いハードウェア環境で安定した高速化を実現する。