クロネッカー積アテンション（KPA）：キーと値をコンパクトな因子のクロネッカー積の和で表し、KVキャッシュを最大10倍圧縮、別の回転ブランチも不要
本論文は、キーと値をトークンに依存するコンパクトな行列因子のクロネッカー積の和で書くKPAを提案する。テンソル積アテンション（TPA）を厳密に一般化し、TPAのランク1のベクトル外積は退化した場合にあたる。この形に合わせた構造化回転演算子がコンパクトな因子だけを回すため、MLAのような分離したRoPEブランチは要らない。約1.25億〜15億パラメータの4規模と9つのベンチマークで、KPAはTPA、MHA、MQA、GQA、MLAと同等かそれ以上で、KV圧縮は6.4〜10倍。モデルはすべて著者が一から訓練した小型のものだ。
大規模モデルの推論で、ますます目立つ壁がKVキャッシュです。重複計算を避けるため、モデルは各層の各トークンのキーと値をすべて保存し、このキャッシュは系列長、バッチサイズ、層の数に比例して増えます。コンテキストが数十万トークンに達し、スループットのためにバッチを大きくすると、KVキャッシュがGPUメモリの大半を占め、重みや活性化の余地がほとんどなくなり、扱えるコンテキストの長さ、バッチ、展開の規模が制限されます。既存の対策はそれぞれ何かを犠牲にしています。マルチクエリアテンション（MQA）とグループクエリアテンション（GQA）は、複数のヘッドでキーと値を共有するので、表現の多様性を失います。マルチヘッド潜在アテンション（MLA）は圧縮した潜在表現をキャッシュしますが、別の分離した位置エンコーディングの仕組みが要ります。テンソル積アテンション（TPA）は低ランクのベクトル外積で分解しますが、厳しすぎる低ランク構造に縛られかねません。論文2610.12657が提案するクロネッカー積アテンション（KPA）は、この状況で登場します。
KPAの中心の考えは、完全なキーと値の行列をキャッシュせず、各トークンのキーと値を、そのトークンに依存するコンパクトな行列因子のクロネッカー積の和で書くことです。クロネッカー積は2つの小さな行列から大きな行列を作るので、ヘッドの段階とヘッド内の特徴の段階の両方に構造を置けます。TPAのベクトル外積はそのうち一方しか扱えません。論文は、実証研究でアテンションのヘッドが、構文、意味、位置などの別々の役割に分かれることが示されているため、よいKV分解は両方の段階で同時に働くべきだと述べます。そしてTPAのランク1のベクトル外積分解がKPAの「退化したクロネッカー分割」で現れることを示し、KPAがTPAを厳密に一般化し、退化しない行列因子ならより豊かなブロック構造を表せると述べます。デコード時には因子だけをキャッシュし、圧縮率はヘッド次元の分け方に依存し、因子が非対称であるほど圧縮は強くなります。この「クロネッカーのヘッド分割」は設計段階で調整できるハイパーパラメータで、他の構造を変えずに圧縮と精度のトレードオフを動かせます。
位置エンコーディングはもう一つの要です。RoPEは完全なキーに作用しますが、KPAは因子しかキャッシュしないので、回転をどこに置くかに答える必要があります。MLAの答えは分離したRoPEの枝を別に持つことです。論文は、クロネッカー表現に合わせた構造化回転演算子を示します。RoPEの相対位置の性質を保ちつつコンパクトな因子だけを回すため、分離した枝は要りません。論文は、そうした枝を必要としないのに、KPAがMLAを上回ると述べています。
実験では、著者らは約1.25億、3.53億、7.73億、15億パラメータの4つの規模でLlama風のモデル（SwiGLUとRoPE、nanoGPTが土台）を、FineWeb-Eduの約500億トークンで訓練しました。キーと値のランクはRK=RV=2、クエリのランクはRQ=6です。lm-evaluation-harnessで、ARC-Easy、ARC-Challenge、BoolQ、HellaSwag、OpenbookQA、PIQA、WinoGrande、MMLU、SciQの9つのゼロショットのベンチマークを評価しました。論文によれば、同じ訓練条件でKPAはTPA、MHA、MQA、GQA、MLAと同等かそれ以上で、KVキャッシュを最大10倍圧縮します。ある規模では、KPAはTPAと同じ6.4倍の圧縮で最良の平均を得て、精度が1.49ポイント高く、最大の向上はARC-Easyです。別の構成は8倍圧縮に達します。約6.7倍の圧縮では他の基準を上回り、ARC-Easy、BoolQ、SciQで差がはっきりします。15億では9項目中8項目でリードし、1.25億ではすべてのKPA構成がTPAの46.21%を上回ります。ヘッド段階の因子を大きくした非対称の分割は、対称の分割とTPAの両方を上回ります。
この論文をどう見るか。長所は、考え方が明快で、厳密な一般化を証明し、圧縮と回転位置エンコーディングの共存という実際の問題に答え、圧縮率を調整できることです。留保すべき点も明確です。第一に、モデルはすべて著者が一から訓練した小さなもので、最大でも15億パラメータ、約500億トークンであり、既存の大型モデルに適用したわけではありません。第二に、評価はゼロショットの精度が中心です。論文自身が結びで、効率的な分解デコードのカーネルを作ることと、より長いコンテキストでKPAを評価することを自然な次の一歩としており、デコード速度、遅延、長いコンテキストでの挙動のデータはまだありません。第三に、結果は自己報告で、独立した再現は見ていません。KVキャッシュのコストが気になるなら、追う価値のある設計として扱い、カーネルと長いコンテキストの結果が出てからGQAやMLAの代わりになるかを判断するのがよいでしょう。
Sources
FAQ
KPAはどのようにKVキャッシュを圧縮しますか。
通常のマルチヘッドアテンションは、各トークンの完全なキーと値をキャッシュします。KPAは、各トークンのキーと値を、そのトークンの隠れ状態から計算されるコンパクトな行列因子のクロネッカー積の和として書き、デコード時には因子だけをキャッシュします。クロネッカー積はヘッドの段階とヘッド内の特徴の段階の両方に構造を与えるため、ベクトル外積を使うTPAより豊かなブロック構造を表せます。論文は、TPAのランク1の外積がKPAの退化した場合であることを示し、KPAが厳密により一般的だとしています。
位置エンコーディングはどう扱いますか。
MLAは圧縮した潜在表現に回転位置埋め込みを使うために、別の分離したRoPEの枝が必要です。KPAは、クロネッカー表現に合わせた構造化回転演算子を与えます。RoPEの相対位置の性質を保ちつつ、コンパクトな因子だけを回すため、分離した枝は不要です。論文は、この点でKPAがその枝なしでもMLAを上回ると述べています。
結果はどうで、留保すべき点は何ですか。
著者らは約1.25億、3.53億、7.73億、15億パラメータのLlama風モデルを、FineWeb-Eduの約500億トークンで一から訓練し、ARC、BoolQ、HellaSwag、PIQA、MMLUなど9つのゼロショットのベンチマークで比較しました。同じ訓練条件でKPAが基準と同等かそれ以上で、6.4倍圧縮でTPAより平均が1.49ポイント高く、最大10倍圧縮、15億では9項目中8項目でリードするといいます。留保点は、モデルがすべて小さく一から訓練したもので、論文自身が効率的な分解デコードのカーネルとより長いコンテキストでの評価を次の課題に挙げているため、実際のデコード速度と長いコンテキストでの挙動はまだ示されていないことです。