TriAttention

Definition:

TriAttentionは,RoPE適用前の空間[pre-RoPE space]において Q/Kベクトル が安定した非ゼロ中心付近に集中するという性質[Q/K concentration]を利用し,その中心から導出される三角級数[trigonometric series]によって各キーの重要度を推定することにより,LLMの長文推論におけるKVキャッシュをオンザフライで圧縮する手法である[Mao et al., 2026].クエリ・キーの実際の値をリアルタイムに参照することなく,事前キャリブレーションで得た分布中心のみから注意ロジットの距離依存性を三角関数の和として近似・予測できることを理論的根拠とする.

背景と問題設定

LLMにおける長文推論[extended reasoning]では, デコード中に生成されるKVキャッシュがメモリの深刻なボトルネックとなる.このため, 重要度の低いKVペアを逐次的に削除するKVキャッシュ圧縮が研究されている.

既存の主要手法[SnapKV, R-KVなど]は, 直近のpost-RoPEクエリの注意スコアを用いてキーの重要度を推定する.しかしRoPE[Rotary Position Embedding]はクエリベクトルを位置に応じて回転させるため, 代表的なクエリとして使えるウィンドウが非常に狭く, 重要なキーを見逃して永続的に削除してしまうという問題があった.これが長文推論の不安定性の主因とされている.

核心的発見:Q/K concentration

本論文の出発点は, pre-RoPE空間[RoPE適用前の空間]における新たな観察である.著者らは, ほぼ全てのアテンションヘッドにおいて, pre-RoPEのQおよびKベクトルが固定の非ゼロ中心付近に強く集中[concentrated]しており, この集中は位置や入力コンテキストを跨いで安定していることを発見した.この性質をQ/K concentrationと呼ぶ.

RoPEは以下のように, 位置$t$に対してベクトルを周波数ごとに回転させる:

\[\tilde{q}_t = \mathrm{RoPE}[q_t, t], \quad \tilde{k}_s = \mathrm{RoPE}[k_s, s]\]

注意ロジットは内積 $\tilde{q}_t^\top \tilde{k}_s$ であるが, pre-RoPEの$q$と$k$が各周波数帯において中心$\mu_q, \mu_k$の周りに集中しているとき, これを三角級数[trigonometric series]として展開できる:

\[\tilde{q}_t^\top \tilde{k}_s \approx \sum_f A_f \cos\bigl[\omega_f [t - s] + \phi_f\bigr]\]

ここで$\omega_f$は周波数, $A_f$と$\phi_f$はQ/Kの中心ベクトル$\mu_q, \mu_k$から決まる振幅と位相である.すなわち注意ロジットはQ-K間の相対距離$[t-s]$の三角関数の和として近似できる.これによりクエリ・キーの実際の値を用いなくても, 各ヘッドが「どの距離のキーを好むか」という距離嗜好[distance preference]が中心ベクトルから予測可能となる.

TriAttentionの定義とスコアリング

TriAttentionは, 上記の洞察に基づき, キー$k_s$の重要度スコアを以下の二成分で推定する.

[1]三角級数スコア $S_{\mathrm{trig}}$
オフラインのキャリブレーションでQの分布中心$\mu_q$を事前に計算しておき, 推論時には各キーの位置$s$に対して三角級数から距離嗜好を評価することで重要度を推定する.距離嗜好に合わない位置のキーには低スコアが与えられ, 圧縮の対象となる.

[2]ノルムベーススコア $S_{\mathrm{norm}}$
Q/Kが比較的分散しているヘッド[少数派]に対する補完信号として, 各周波数帯の期待クエリ寄与によって重み付けされたKベクトルのノルムを用いる.

[3]適応的重み付け
二成分の混合比は, 集中度の指標である平均結果長[Mean Resultant Length; $R$]によって自動的に調整される:

\[S = R \cdot S_{\mathrm{trig}} + [1 - R] \cdot S_{\mathrm{norm}}\]

$R$が高い[集中度が高い]ヘッドでは$S_{\mathrm{trig}}$が支配的となり, $R$が低いヘッドでは$S_{\mathrm{norm}}$が補完する.

実験結果

AIME25[32Kトークン生成]での評価では, KVキャッシュ予算2048トークンの条件においてSnapKVやR-KVが精度を約半分に落とすのに対し, TriAttentionはFull Attentionの精度に匹敵する結果を示した.同じ精度水準[40.8%]での比較では2.5倍のスループット向上, 精度を一致させた条件では10.7倍のKVメモリ削減を達成している.またMATH 500では1,405トークン/秒[Full Attentionの223トークン/秒比で6.3倍]のピークスループットを記録した.実用的な側面では, 通常ではメモリ不足でFull Attentionが実行不可能な24GB GPUへの32Bモデル[OpenClaw]の展開を可能にしている.

手法の特徴まとめ

TriAttentionの本質的な特徴は, クエリ・キーの実際の値をリアルタイムに参照せず, 事前キャリブレーションで得た分布中心のみを用いて重要度を推定できる点にある.これにより推論オーバーヘッドを最小化しつつ, 位置回転の影響を受けないpre-RoPE空間の安定した構造を活用できる.また特定のモデルアーキテクチャに依存せず, Qwen3・Qwen2.5・Llama3ベースの複数モデルで有効性が確認されている.

参考文献


2026-04-12.

Mathematics is the language with which God has written the universe.





















追記型アーキテクチャ Deep Agents LiteRT-LM Transformerの万能近似能力 LightMem MindsDB