0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

鍵を捨てて値だけ残すKeyless Attention、KVキャッシュを半分にする設計

0
Posted at

長い文脈を扱うLLMを動かしていると、GPUメモリの大部分がモデルの重みではなくKVキャッシュに食われる場面に何度もぶつかる。生成トークンが伸びるほど、過去の全トークン分の「キー(K)」と「ベクトル値(V)」を保持し続けるからだ。ここで素朴な疑問が湧く。KとVはきれいに半々でメモリを使うのに、キャッシュから引き出して実際に使うのはVだけ。Kは「どのトークンに注目するか」を決める道案内に使われるだけで、答えそのものではない。この道案内のためだけに、キャッシュの半分を払い続けるのは妥当なのか。

Keyless Attention という論文は、この問いに「キーのキャッシュは要らない」と正面から答える。標準的なQKVアテンションからキー射影を丸ごと消し、Vだけをキャッシュする設計で、KVキャッシュのメモリをちょうど半分にする。

  • Keyless Attention: Value-Space Routing and Value-Only Caching for Efficient Transformers (Xin Gao, Xingming Xu / arXiv 2606.21848, v2: 2026-07-31)

🔑 キーを消すと計算はどう変わるのか

まず標準のアテンションを思い出しておく。クエリQとキーKの内積でスコアを作り、softmaxで重みにして、値Vを混ぜる。

標準:   softmax( Q Kᵀ / √d ) V     ただし Q = X·W_Q,  K = X·W_K,  V = X·W_V

キャッシュに残すのはKとV。ここでKは「Qが各トークンをどれだけ引くか」を決める住所ラベルの役割で、答えの中身ではない。図書館で言えば、本(V)ごとに付いている検索用のインデックスカード(K)にあたる。

Keyless Attentionは、このインデックスカードを捨てる。代わりに、クエリを値の空間へ翻訳する 値ルーティング射影 W_R を一枚差し込み、スコアをQと「Vそのもの」の間で取る。

Keyless: softmax( X·W_Q·W_R·(X·W_V)ᵀ / √d ) X·W_V

肝は、推論時に W_Q·W_R を事前に1つの行列へ畳み込めること。つまり実行時の追加計算はゼロで、キーの射影計算(X·W_K)がまるごと消える。だからメモリが半分になるだけでなく、デコード時のスループットもベースラインを上回る。キャッシュに置くのはVだけ、という一点だけが変わる。

キーをただ捨てるのは以前からある発想(KとVを共有する手法など)だが、論文はそれらを「キーが担っていた道案内の役割を明示的に置き換えていない」と批判する。W_R はまさにその役割の受け皿として入っている点が新しい。

消したのに精度が落ちない理由 🤔

普通なら表現力を削れば品質は下がる。ところが実験では、5モデル中4モデルでKeylessの方がパープレキシティが良い。

モデル 標準QKV Keyless
GPT-2 (557M) 33.50 33.26
Pythia (410M) 40.99 39.22
Qwen2 (1.5B) 34.38 33.79
Llama 3.2 (1B) 39.09 38.59
GPT-2 (280M) 33.71 33.84

下流タスク(GPT-2 557M)でもHellaSwagやStoryClozeなど5問中4問で上回った。

著者が挙げる理由が面白い。W_R と W_V は学習中に勾配が絡み合う(gradient entanglement)ため、ルーティングがコーパス固有の共起パターンへ過剰に特化しにくくなる。要は、キーを消したことが結果的に暗黙の正則化として効いている、という説明だ。実際、ベスト到達後の検証ロスの悪化が標準より緩やかで、過学習しにくい挙動が観測されている。個人的にはこの「制約が汎化を助ける」筋の主張は、パラメータを増やす方向ばかりに慣れた身には気持ちのいい逆張りに映る。

GQAやMLAと何が違うのか

KVキャッシュ削減はここ数年の主戦場で、既に複数の系譜がある。位置づけを整理するとこうなる。

手法 キャッシュするもの キー削減の考え方
MHA(素の実装) K と V(全ヘッド) 削減しない
GQA / MQA K と V(ヘッド間で共有) ヘッドをまとめて縮める
MLA(DeepSeek系) K・Vを圧縮した低ランク潜在 圧縮して1本にする
Keyless Attention V のみ キー射影自体を消す

注目すべきは、Keylessが「ヘッドをまとめる」GQAと直交していること。論文はMHAとGQAの両方で有効性を確認しており、GQAで縮めたうえで、さらにキーを消して半減させる合わせ技が原理上は成り立つ。KVキャッシュ削減の手札がまた一枚、しかも既存手法と積み重ねられる形で増えたと見るのが正確だ。

どこまで信じてよいか

過度な期待は禁物だ。検証されたモデルは最大でもQwen2 1.5B、データはWikiText-103の3000万トークン、学習はA100一枚という小規模な設定にとどまる。7Bや13B、まして本番規模のマルチヘッド構成で同じ50%削減と品質維持が成り立つかは、この論文だけでは判断できない。著者自身、等価性を保証する定理の存在条件がマルチヘッドで常に成り立つわけではないと認めており、収束が標準よりやや遅い(エポックあたりの学習時間が増える)ことも明記している。

それでも、推論時に追加コストゼロでキャッシュを確実に半分にできるという主張は、長文脈サービングのコスト構造を直接叩く。KVキャッシュ最適化を量子化や圧縮で削ってきたチームにとって、「そもそもキーを持たない」という発想は試す価値のある補助線だ。まずは手元の中規模モデルで、W_R を一枚足してキー射影を外し、パープレキシティとデコード速度が論文通りに動くかを再現するところから確かめたい。実装が公開されれば、GQAとの合わせ技での挙動が次の見どころになる。

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?