SpecLA: 線形注意モデルのための効率的な投機的デコード
本論文は、状態を持つ線形注意モデルのための投機的デコードランタイムSpecLAを提案する。トポロジ認識カーネルを用いたチェーンとツリーの検証、検証中に生成されたコンパクト因子を保存して受理状態を復元、信頼度枝刈りとターゲット調整済みEAGLEスタイルのドラフターを使用する。NVIDIA H100上の公開GDN-1.3Bターゲットで、自己回帰デコードと比較して最大1.70倍のエンドツーエンド高速化を達成。
線形注意モデル(Mamba、RWKVなど)は、リカレント状態を導入することでTransformerで増大するKVキャッシュを置き換え、長いシーケンス推論時のメモリ使用量を大幅に削減する。しかし、これらのモデルの自己回帰デコードは依然としてトークンを1つずつ処理する必要があり、推論速度が逐次計算に制限される。投機的デコード(Speculative Decoding)は、複数のドラフトトークンを並行検証することで自己回帰デコードを高速化する手法だが、既存のシステムはすべてTransformerのKVキャッシュ向けに設計されており、状態依存を持つ線形注意モデルには直接適用できない。主な課題は、検証プロセスがリカレントな依存関係に従わなければならず、受け入れられた軌跡のみを更新すること、そしてドラフターが無駄な検証作業を引き起こさないように有用な候補を生成することである。
これらの課題に対処するため、研究チーム(Zhibin Wangら7名)は2026年7月に投稿された論文でSpecLAを提案した。SpecLAは、トポロジ認識カーネルを用いてチェーンおよびツリー状の候補を効率的に検証し、検証中に生成されたコンパクトな因子表現を保存して受け入れられた状態を復元する。さらに、信頼度枝刈りとターゲットに合わせたEAGLEスタイルのドラフターを組み合わせ、検証器に高品質な候補を提供する。実験はNVIDIA H100 GPU上で、公開されているGDN-1.3Bモデルをターゲットとして実施された。結果、SpecLAは生成品質を犠牲にすることなく、自己回帰デコード比で最大1.70倍のエンドツーエンド高速化を達成した。この研究は、線形注意モデルの効率的な推論に新たな道を開き、リアルタイム対話型アプリケーション(チャットボットやコード生成など)への展開を促進することが期待される。
SpecLAの革新性は、トポロジ認識検証カーネルによりチェーンやツリーの依存関係を効率的に処理できる点にあり、これは従来の投機的デコードシステムでは扱えなかったものである。また、検証中にコンパクト因子を保存することで、受け入れ後の状態回復を高速化し、信頼度枝刈りによって低品質な草案をフィルタリングして検証資源を有効活用する。本論文はcs.CLに分類され、今後の大規模モデルや異なるハードウェアへの拡張が期待される。