長い推論のKVキャッシュを減らす
LLMは計算済みのキーと値をKVキャッシュに保存し、次のトークンを生成するときに再利用します。推論が長くなるほどキャッシュも増えます。過去の状態を削除すればメモリを節約できますが、解答に必要な情報まで失うことがあります。
CASKは解答に必要な状態を残し、繰り返される中間推論をまとめる手法です。圧縮後も元のモデルの推論をどこまで維持できるかを調べました。
重要度スコアの限界
TriAttentionなど、トークンに重要度を付けて一部を残す手法を調べました。スコアの計算を変えても、実際に残るトークンがほとんど変わらない条件がありました。スコアを精密にするだけでは、推論に必要な情報を十分に保てませんでした。
そこで生成中の状態を役割で分けました。解答の形成や次の推論を支える状態はprotected coreとして保持します。繰り返しの探索や再確認から生じる似た状態はmergeable scratchに分け、ここだけを統合します。
統合では各状態の寄与を反映した代表キーと値を計算します。後続トークンとの関連性を考慮した距離で統合対象を選び、失われる情報も確認します。coreは統合対象に含めません。
入力と生成の圧縮
長い入力文書は、生成が始まる前にキャッシュ予算の大部分を使います。生成中の推論を圧縮する仕組みがあっても、動作させる余地がありません。CASKでは処理を二段階に分けました。
- 入力区間のキャッシュを一部削除し、生成に使う空間を確保します。
- 生成区間ではcoreを保持し、scratchだけを選択的に統合します。
実験方法と結果
推論実験にはQwen3-8Bを使い、H100 PCIe上でbfloat16・SDPA設定で実行しました。AIME24とAIME25からそれぞれ6例を比較しました。
まず非圧縮モデルの出力を取得し、同じトークン列を圧縮モデルへ一つずつ入力します。次のトークンの一致率、元のトークンに与える確率、最初に予測がずれる位置を記録しました。
| データセット | 手法 | 次トークン一致率 | 平均NLL |
|---|---|---|---|
| AIME24 | TriAttention | 88.2% | 0.383 |
| AIME24 | CASK | 90.7% | 0.268 |
| AIME25 | TriAttention | 89.1% | 0.357 |
| AIME25 | CASK | 90.3% | 0.313 |
予算384では、両データセットでCASKの一致率が高く、平均NLLが低くなりました。一致率は、予算512のTriAttentionによるAIME24の89.4%、AIME25の89.9%も上回りました。
長い入力で残った課題
LongBenchのmulti_newsとvcsumでは生成区間の統合が実際に動作し、元のモデルの出力に近づきました。qasperの改善は主に入力区間の処理によるものでした。
qmsumとgov_reportでは、入力だけで予算を使い切る問題が残りました。2wikimqaも予算384ではCASKの一致率が低くなりました。
論文と実装
CASK: Core-Aware Selective KV Compression for Reasoning Traces
Buseong Kim · Heejun Gwon
にarXivで公開しました。arXiv:2604.10900。
公開リポジトリにはCASKとTriAttentionの実装、実験用スクリプト、測定結果、論文の図を収録しています。