技術的概要：SOCKET - 疎な注意機構のための SOft Collision Kernel EsTimator

問題定義

大規模言語モデル（LLM）は、長文脈推論中に根本的なスケーラビリティの限界に直面しています。プリフィリング段階では入力シーケンス全体を処理しますが、反復的デコーディング段階では、新たに生成されたトークンがすべて、以前にキャッシュされたトークンに注意を向ける必要があります。文脈長が増大する（2K から数百万トークンへ）につれて、この密な注意機構はメモリに束縛されるようになり、GPU 容量を超えることが多く、高価な CPU-GPU 転送を必要とします。

既存の疎な注意手法は、計算をトークンの部分集合に制限することでこれを緩和しようと試みています。しかし、多くの手法は k-means クラスタリングなどのデータ依存ヒューリスティックに依存しており、高い TTFT（First Token までの時間）コストを発生させ、分布シフトに対する頑健性が欠如しています。他の手法は、二値の衝突信号を割り当てる従来の局所感度ハッシュ（LSH）を利用します。本論文は、この「ハード」な定式化がランキングの安定性には不適切であると主張しています。二値の衝突は部分的な類似性を表現できず、粗くノイズの多いランキングを招き、特に厳格なメモリ制約下で検索品質を劣化させます。

手法：SOCKET

著者らは、LSH をバイナリな候補生成器から、原理的かつ類似性を意識したスコアリングカーネルへと再定義する、データ非依存の疎な注意機構であるSOCKET（SOft Collision Kernel EsTimator）を提案します。

中核メカニズム：ソフト LSH

単一のハッシュテーブルあたりクエリを単一の離散バケットに割り当てる従来の LSH と異なり、SOCKET はソフト LSHを採用します：

1. 事前計算（プリフィリング）： キーはランダム射影を用いて $L$ 個の独立したハッシュテーブルにハッシュされます。各キーはテーブルごとに単一のバケット ID を割り当てられ、値ベクトルのノルムがキャッシュされます。
2. ソフト割り当て（デコーディング）： 新しいクエリが到着すると、単一のバケットに割り当てられるのではなく、各ハッシュテーブル内のすべてのバケットにわたる確率分布を誘起します。これは、クエリのランダム超平面への射影に基づいてソフトハッシュスコアを計算し、温度制御されたソフトマックス関数に通すことで達成されます。
3. スコアリング： 各キーについて、すべての $L$ 個のテーブルにわたってそのキーの特定のバケットに割り当てられた確率質量を集約することで、連続的なスコアが計算されます。このスコアは、キーの値ベクトルのノルムで重み付けされます。
4. 選択： トップ-$k$ のキーがこれらのソフトスコアに基づいて決定論的に選択され、この部分集合に対して正確な注意が計算されます。

理論的基盤

本論文は、ソフト LSH スコアを用いて計算された注意が、ソフトマックス注意（具体的には角カーネル）と密接に関連するカーネルを近似することを示す理論的保証を提供します。

• 安定性： 理論的分析（定理 3）は、ソフトバケッタイゼーションが温度 $\tau \to 0$ になるにつれて消滅する制御可能なバイアス（$\epsilon_\tau$）を導入しつつ、滑らかさを維持することを示しています。
• 相関： 補題 4 と実証結果は、ハードスコアよりもソフトスコアの方が、射影されたクエリ - キー類似性の方向構造をよりよく保持することを示しています。ハード LSH は $\ell_1$ 大きさに依存し、直交領域内の方向情報を破棄するため、有限の $L$ において不安定なランキングをもたらします。ソフト LSH は、小信号領域において射影された座標と線形関係を維持し、真の類似性とのより高い相関と推定量の分散の低下をもたらします。

主要な貢献

1. 安定したランカーとしてのソフト LSH： 著者らは、従来のハード LSH よりもはるかに効果的なランキング関数として機能する、データ非依存のソフト LSH メカニズムを提案します。これは、二値の衝突信号を確率的で類似性を意識した集約に置き換え、アドホックな投票なしに安定したトップ-$k$ 選択を可能にします。
2. カスタムカーネルによる効率性： カスタム CUDA スコアリングカーネルと Flash Decode Triton バックエンドが実装されました。これにより、SOCKET は長文脈におけるデコーディング中に FlashAttention よりも最大1.5 倍の高いスループットを達成し、既存の疎な手法の精度を維持または上回ります。
3. データ非依存設計： クラスタリングベースのアプローチとは異なり、SOCKET は再学習、較正、またはオフライン前処理を必要としません。これにより、より高速な TTFT と分布シフトに対する頑健性が実現されます。
4. 理論的保証： 本論文は、ソフトカウント注意推定量のエンドツーエンドの誤差限界を確立し、誤差をサンプリング分散、有限テーブル近似誤差、ソフトバケッタイゼーションバイアスに分解します。

実験結果

著者らは、LongBench（実世界のタスク）とRULER（合成診断タスク）の 2 つのベンチマークにおいて、Llama-3.1-8B-Instruct、Llama-3.2-1B、およびQwen3モデル（4B、8B、30B）で SOCKET を評価しました。

• 精度： LongBench において、SOCKET は Quest、PQCache、Double Sparsity、MagicPig などのベースラインを一貫して上回りました。例えば、Llama-3.1-8B において 33 倍の疎性で、SOCKET は平均スコア 47.83 を達成し、最も強力なベースライン（Quest）を約 0.84 ポイント上回りました。RULER-32K においては、特に高い疎性レベル（例：50 倍）で先行手法と同等かそれ以上の性能を示しました。
• 効率性： SOCKET は NVIDIA A100 および H200 GPU で顕著なスループット向上を示しました。H200 においては、文脈長が増大するにつれて FlashAttention-2 および FlashAttention-3 を上回り、100K トークンで FlashAttention-2 に対して1.58 倍の高速化、140K トークンで FlashAttention-3 に対して1.58 倍の高速化（33 倍の疎性下）を達成しました。
• メモリオーバーヘッド： この手法は、標準的な KV キャッシュに加えて約 15% の控えめなメモリオーバーヘッド（ハッシュテーブルインデックスと値ノルムの保存）を導入しますが、同等の精度を達成するためにハード LSH が必要とする（より多くのハッシュテーブルを必要とする）オーバーヘッドよりも著しく低いです。

意義と主張

本論文は、SOCKET が既存の疎な注意手法に対する原理的かつデータ非依存の代替案を提供することで、長文脈 LLM 推論の重要なボトルネックに対処すると主張しています。ハードな二値ハッシングからソフトな確率的スコアリングへと移行することで、SOCKET は優れた精度と効率性のトレードオフを達成します。

著者らは、自らのアプローチが以下の点で優れていることを強調しています：

• データ依存の前処理の必要性を排除し、TTFT を削減して展開の柔軟性を向上させます。
• 理論的安定性を提供し、従来の LSH の弱点である限られた数のハッシュテーブルであってもトークンのランキングが一貫して維持されることを保証します。
• 実用的なシステムレベルの利益をもたらし、スコアリングカーネルのアルゴリズム的改善が、モデル性能を犠牲にすることなく、現代の GPU アーキテクチャ上で直接的に顕著なスループット向上に変換されることを実証しています。

この研究は、ソフト LSH を単なる検索ツールではなく、デコーディング段階で密な注意を置き換えることができる堅牢なスコアリングカーネルとして位置づけ、LLM の数百万トークン規模の文脈への効率的なスケーリングを可能にします。