技術要約: LiLiCorr

問題提起

投機的デコーディング（Speculative Decoding）は、より小さな「ドラフト（下書き）」モデルを使用して将来のトークンを提案し、より大きな「ターゲット」モデルがそれらを並列に検証することで、自己回帰型言語モデルの推論を加速させます。DFlashのような有望なドラフトモデルの一種は、拡散型のアーキテクチャを用いて、単一のフォワードパスで将来のトークンブロック全体を予測します。しかし、これらのモデルは位置ごとの周辺分布（per-position marginals）（特定のポジションにおけるトークンの確率を独立して予測すること）に基づいて学習されています。

その結果、DFlashが生成するトークンは個々としては妥当であっても、**結合的な一貫性（jointly incoherent）**を欠くことがよくあります。各スロットで最も確率の高いトークンを独立して選択すると、ターゲットモデルが早い段階で拒絶してしまうシーケンスになりやすく、受け入れ長（acceptance length）とそれによる加速率が制限されてしまいます。周辺分布から一貫したシーケンスを復元することは組合せ爆発的に困難です。なぜなら、$B$ 個のポジションと語彙サイズ $V$ を持つブロックに対する結合分布は $V^B$ 個の実現形を持ち、明示的なモデリングは計算量的に不可能だからです。既存の解決策は、コストの高いドラフトの追加パス（反復的な洗練）に依存するか、相関の負担をターゲットモデルに転嫁するか（投機の目的を損なう）、あるいはスロットを一つずつ処理する逐次的な自己回帰型補正ネットワークを使用するかのいずれかです。

手法: LiLiCorr

LiLiCorr（Lightweight Likelihood Correlation）は、単一のネットワークパスで、各スロットの周辺分布から一貫したシーケンスを復元することでこの問題に対処します。これは、完全な結合分布をモデリングするのではなく、候補となるラティス（格子）上のスコアリングモデルを学習するものとして問題を扱います。

1. 候補ラティスと表現

単一のトークンを選択する代わりに、LiLiCorrはドラフトの周辺分布から各 $B-1$ 個の投機スロットにおけるトップ-$K$ 個の候補トークンを保持します。

• ノード構築: 各候補トークンは、以下の4つのコンポーネントを合計して形成されるベクトル（ノード）として表現されます。
  1. 凍結されたターゲットモデルのトークンの入力埋め込み。
  2. そのスロットにおけるドラフトの隠れ状態（スロット内の全候補で共有）。
  3. ドラフトの確信度（対数確率、ランクなど）のスカラ特徴量を処理する小さなネットワーク。
  4. スロットインデックスと候補のランクのための学習済みエンコーディング。
• 文脈処理: 軽量な双方向Transformerが、$(B-1) \times K$ のラティス全体を結合して処理します。これにより、すべての候補が他のすべての候補（スロット内の競合する候補および隣接するスロットの候補の両方）と情報を交換できるようになり、ブロック全体の文脈をエンコードできます。

2. In/Out ベクトルによる相関

各候補ノードに対して、LiLiCorrは正規化された2つのベクトル、すなわちinベクトル ($i$) と outベクトル ($o$) を出力します。

• ペアワイズ・スコアリング: 隣接するスロット $s$ と $s+1$ は、$s$ における候補の out ベクトルと $s+1$ における候補の in ベクトルとの間のコサイン類似度を計算することによって結合されます。
• 効率性: この定式化により、相関ロジックをバッチ化された行列演算へと簡略化します。ブロック全体のすべてのペアワーズな相互作用は、単一のフォワードパスの後に並列に計算されます。
• 安定性: コサイン類似度（定数 $\sigma$ でスケール）を使用することで、結合学習中に共有DFlashバックボーンを不安定にする可能性のある、無制限のスコアの大きさを防ぎます。

3. 学習目的

LiLiCorrは、DFlashドラフトと**共同学習（co-trained）**されます。学習損失は以下の3つのコンポーネントで構成されます。

1. DFlash 損失: ドラフトの標準的なスロットごとのクロスエントロピー。
2. オラクル・プレフィックス・クロスエントロピー ($L_{CE}$): 前のスロットのグランドトゥルース（正解）トークンを条件としたとき、グランドトゥルースのトークンを最も高くランク付けするようにLiLiCorrを訓練する、教師強制（teacher-forced）による損失。決定的なのは、グラウンドトゥルースがトップ-$K$ 候補の中に存在している最長の先行スロットの連なり（オラクル・プレフィックス）に対してのみ、監督が行われる点です。
3. ディストラクター・ペナルティ ($L_{dist}$): ターゲットモデルが低くランク付けする候補に対してペナルティを与える、コストを考慮した項。これは、グランドトゥルースに対する期待ターゲットコストを最小化し、具体的にはターゲットモデルが拒絶するような「混乱させるもの（confusers）」を押し下げ、正しいトークンのためのマージンを広げます。

4. デコーディング

推論時には、プロセスは2段階で行われます。

1. 並列スコアリング: 単一のLiLiCorrパスが、すべてのin/outベクトルとすべてのペアワーズなマッチングスコアを並列に計算します。
2. 逐次的グリーディ・ウォーク: ブロックは左から右へと確定されます。各スロットにおいて、直前に確定された候補との一致を最大化する候補を選択します。このグリーディ・ウォークは逐次的ですが、ネットワークの重みには一切触れないため、計算コストは非常に低いです。

主な貢献

1. 新しい定式化: 完全な結合分布の計算量を回避しつつ、候補パス上のスコアリングモデルを学習することで、各スロットの周辺分布から一貫したシーケンスを復元する、尤度ベースのアプローチ。
2. 並列相関: 相関ロジックを効率的で並列化可能な行列演算に集約しました。スロットごとに補正ネットワークを実行する従来のドラフト時アプローチとは異なり、LiLiCorrは単一のパスで全候補ラティスをスコアリングし、その後に安価なグリーディ選択を行います。
3. 最先端の性能: 本手法は、多様な設定において、投機的デコーディングのベンチマークを塗り替え、既存のDFlashドラフトや同時期の研究（Domino, DSpark）を上回る性能を示しました。

実験結果

著者らは、Qwen3-8BおよびQwen3-4Bターゲットモデルを用い、greedyおよびtemperature-1デコーディングの下で、9つのベンチマーク（math, code, chat）について評価を行いました。

• 受け入れ長（Acceptance Length）: バニラのDFlashドラフトと比較して、LiLiCorrはすべてのベンチマークにおいて受け入れ長を**9%から19%**増加させました。
• スループット: LiLiCorrは、テストされた72の設定（ベンチマーク、ターゲットサイズ、デコーディングレジーム、並列度を変えたもの）のうち、70の設定で最高のスループットを達成しました。
  • 最も強力なベースラインであるDominoに対し、平均で最大**5.3%**上回っています。
  • LiLiCorrヘッドは、1ブロックあたりのレイテンシのわずか2.8%を占めており、完全に最適化された場合、Dominoのヘッドよりも約2.4倍高速です。
• 汎用性: 位置エンコーディングを拡張（例：YaRN経由）すれば、訓練範囲より1桁長い入力（最大32Kトークン）に対しても、性能のリードが維持されます。
• アブレーション: ディストラクター・ペナルティは、特に短いプロンプトにおいて極めて重要であり、ベンチマークによって受け入れ長をさらに1.5%から8.9%向上させることが示されました。

意義と主張

本論文は、LiLiこれが、ドラフトトークンの相関を逐次的な自己回帰パスから切り離すことで、投機的デコーディングにおける構造的な改善を実現したと主張しています。すべての相関処理を単一の並列パスで行うことで、ドラフトのオーバーヘッドを最小限に抑えつつ、受け入れ長を最大化します。

著者らは、この手法の範囲が限定的であることを強調しています。つまり、ドラフトがグランドトゥルースを含む候補プール（オラクル・プレフィックス）を提案できる能力に依存しています。もし正しいトークンがトップ-$K$ 候補に含まれていなければ、LiLiCorrはそれを復元できません。しかし、候補プールの制約内においては、一貫した結合性を回復するための非常に効率的なメカニズムを提供しており、高並列・長文脈シナリオを含む幅広いサービング条件下で最も高速なシステムとなります。本研究は、スロット間のより長い範囲の結合を検討することが、プレフィックス受け入れのコストに見合うかどうかという未解決の問いを残しつつも、将来の改善の方向性を示唆しています。