Diffusion LLMの推論を高速化するPolestar──KVキャッシュ再利用とトークンコミットを最適化

論文概要

Diffusion LLM (dLLM) の推論における効率化手法「Polestar」が提案された。dLLM は双方向アテンション機構を用いるため KV キャッシュの再利用が困難であり、また静的な信頼度閾値による並列デコードは精度低下を招く課題がある。Polestar は、デコード過程でトークン表現が変化する「ドリフト」現象を統一的な信号として利用することで、KV キャッシュのスパースな更新（Polestar-Cache）と、トークンの早期コミット（Polestar-Commit）を同時に実現する。

関連研究

既存の推論最適化手法は、主に KV キャッシュの圧縮や静的な並列デコード戦略に依存していた。しかし、これらは dLLM 特有の双方向アテンションによる表現の動的な変化を考慮できていない。本研究は、この表現ドリフトを明示的に観測・利用することで、トレーニング不要の推論フレームワークを実現した点で先行研究と一線を画す。

新規性と貢献

Polestar の最大の新規性は、トークンの表現ドリフトを「KV キャッシュの陳腐化判定」と「トークンコミットの信頼性判定」という二つの異なる課題を解決するための共通信号として定義したことにある。これにより、推論の精度とスループットのトレードオフを改善し、Pareto フロンティアを更新した。

提案手法の詳細

Polestar は以下の二つのコンポーネントで構成される。

• Polestar-Cache: ドリフト量を監視し、KV キャッシュの陳腐化を検知。必要な箇所のみを更新することで、キャッシュの効率的な再利用を可能にする。
• Polestar-Commit: 急激なドリフトイベントを検知し、コミット準備が整ったトークンを特定。これにより、信頼性を担保した並列デコードを実現する。

評価・考察

数学およびコーディングのベンチマークにおいて、Polestar は既存手法と比較して最大 10.73% の精度向上と、最大 3.7 倍のスループットを実現した。また、フォワードパスあたり 3.67 トークンという高いデコード並列度を達成している。

応用例と今後の展望

本手法は、dLLM を用いた推論サービスを展開する企業において、ハードウェアリソースのコストを削減しつつ生成品質を維持する実務的な手段となる。特に、推論負荷が高い数学的推論やコード生成タスクにおいて、数千〜数万リクエスト規模のサービス運用時に劇的なコスト改善が期待できる。今後の課題は、ドリフト検知の計算オーバーヘッドをさらに削減し、より多様なモデルアーキテクチャへの汎用性を検証することである。

結論

Polestar は、dLLM の推論効率を制約していた二つの主要課題を、トークン表現ドリフトという単一の指標で解決した。トレーニング不要かつ高い性能向上を実現しており、実用的な推論最適化手法として重要な一歩である。

注釈

• dLLM (Diffusion LLM): 拡散モデルの枠組みを言語モデルの推論に応用したモデル群。
• KV キャッシュ: 推論時に以前の計算結果を保持し、再計算を避けるためのメモリ領域。
• トークンコミット: デコードされたトークンを最終的な出力として確定させるプロセス。