KVキャッシュ淘汰は賢くなくてもよい？Random Attentionの提案

長い推論チェーンは、大規模言語モデルが複雑な問題を解くうえで役立つ。しかし、生成されたトークンが増え続けると、各トークンのKeyとValueを保存するKVキャッシュが大きなメモリ負担になる。従来の多くの圧縮手法は、後で役立つ可能性をトークンごとに推定し、スコアの高いものを残すという考え方に基づいている。Random Attentionは、この選択スコアが本当に必要なのかを問い直した。

提案手法は単純だ。プロンプトは常に保持し、それ以外のキャッシュを各注意ヘッド内で一様ランダムに淘汰する。重要度の計算やランキングは行わない。4つのモデルと6つの推論タスクによる評価では、強力な既存淘汰手法に匹敵した。さらにvLLMでの実装では、その手法より32〜43%高いスループットを達成した。

主なポイント

• プロンプトを守ることが重要。 制御実験では、手法間の性能差の多くが、選択シグナルによってプロンプト関連情報を保持できたかどうかに起因することが示された。
• 推論テキストには冗長性がある。 モデルは推論を続ける際、必要な情報を言い換えたり再記述したりする。そのため、古いトークンが消えても、後続のテキストに代替情報が残る場合がある。
• 注意ヘッド間にも冗長性がある。 各ヘッドは推論トレースを独自に表現する。最も重要な1個を正確に選ばなくても、関連情報のコピーが十分残れば推論を継続できる。
• スコア計算を省ける。 重要度推定やソートをなくすことで、圧縮処理の計算量と実装の複雑さを抑えられる。

この結果は、すべてのトークンが同じ価値を持つことや、ランダム淘汰があらゆる文脈で安全であることを意味しない。むしろ、長い推論タスクでは生成トレースそのものが冗長であり、複雑な選択器の追加コストに見合う改善が得られない場合がある、という示唆である。

推論基盤にとっては、キャッシュ全体を精密に並べ替える前に、失われると影響が大きいプロンプトを保護し、残りはモデルの繰り返し表現と多ヘッド構造に任せる設計が考えられる。より長いコンテキストや異なるタスクでも成立するかは今後の検証課題だが、単純なランダム方式を基準にして、複雑な淘汰器の実益を測るきっかけになる研究である。

出典：Hugging Face Daily Papers