Scissorhandsとは?重要トークンを残してLLMのKVキャッシュを圧縮する技術

Scissorhandsは、将来も参照されやすい重要トークンだけをKVキャッシュに残すことで、LLM推論のメモリを削減する手法です。仕組み、実験結果、長文RAGやLLMサービングへの活用を解説します。

参考文献

Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time

Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, Anshumali Shrivastava

論文を見る

今回の論文

今回取り上げるのは、Zichang Liuらによる論文「Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time」です。2023年5月にarXivで公開された、LLM推論とメモリ最適化の研究です。URLは https://arxiv.org/abs/2305.17118 、研究分野はTransformer推論、KVキャッシュ圧縮、LLMサービングです。

選んだ理由は、モデルを再学習せずに「過去を全部保持する」という推論時の前提を見直すからです。長文RAG、会話履歴、同時リクエストが増えるほどKVキャッシュは運用上の制約になります。どの情報を残すべきかを注意重みから判断する発想は、推論基盤を設計する際にも応用しやすいと考えられます。

どんな技術か

Scissorhandsは、LLMが次のトークンを生成する際に保持するKVキャッシュから、将来も重要になりそうなトークンを選んで残す圧縮手法です。名前どおり、不要と判断した過去トークンのKeyとValueを「切り落とし」、固定のメモリ予算内で生成を続けます。

核になる観察は、ある時点で強く参照されたトークンは、その後の時点でも強く参照されやすいというものです。論文はこれを「重要性の持続仮説」と呼びます。直近の注意重みの履歴を使えば、未来に必要になりそうな情報を完全には分からなくても、低いリスクで選別できるという考え方です。

課題

自己回帰型LLMは、過去の各トークンについて計算したKeyとValueを再利用します。これがKVキャッシュです。生成が長くなるほど、また同時に扱うリクエスト数が多いほど、キャッシュは線形に増えます。論文の試算では、バッチサイズ128・系列長2048のOPT-175Bで、重みが約325GBなのに対しKVキャッシュは約1,152GBです。

既存の量子化は、重みやKVの各数値を少ないビットで保存できますが、「全トークンを保存する」前提のままでは系列長に伴う増加を止められません。単純な先頭切り捨てやスライディングウィンドウは、古いが重要な命令、固有名詞、RAGで取得した根拠を消してしまう恐れがあります。

実際には、長いチャット履歴を持つエージェント、数十ページの文書を読むRAG、GPUを共有して高い同時実行数を求めるAPIで問題になります。品質を保ったまま、どの過去情報を残すかを推論時に決める必要があります。

用語解説

KVキャッシュ
Transformerの注意機構で、過去トークンから作ったKeyとValueを保存した領域です。再計算を避けられる一方、トークン数・層数・バッチ数に応じて増えるため、本論文の削減対象になります。
自己回帰生成
直前までの入力と生成済みトークンを条件に、次の一トークンを順番に出す方式です。各ステップで過去全体を参照するため、KVキャッシュが必要になります。
注意重み(attention weight)
現在のトークンが過去のどのトークンをどれだけ参照するかを表す値です。Scissorhandsは、この値が大きいトークンを将来にも役立つ候補として扱います。
重要性の持続仮説
ある過去ステップで大きな影響を持ったトークンは、将来の生成でも重要であり続ける、という論文の仮説です。未来の注意を直接見られないというキャッシュ削減の難しさを、過去の観測で補います。
パープレキシティ
正解の続きテキストをどれだけ自然に予測できたかを測る言語モデリング指標で、低いほどよい値です。圧縮で生成能力が落ちていないかを確認するために使われます。

技術の仕組み

Scissorhandsは、モデルの重みを変更せず、推論ループのKVキャッシュ管理だけを置き換えます。したがって追加学習なしで既存モデルへ適用する設計です。

注意の偏りと持続性を利用する

ある時刻tの注意出力は、各過去トークンのValueを注意重みα(t, i)で加重平均したものです。注意重みは均一ではなく、少数トークンに集中する傾向があります。さらにOPTモデル、OpenBookQA、WikiTextで調べると、系列前半で重要だったトークン集合と後半で重要だった集合の重なりは、多くの層で95%超でした。

これは「今までほとんど参照されなかったトークン」を捨てても、次の注意出力に与える影響が小さい可能性を示します。ただし、論文はこれを全モデル・全入力で成立する法則とは扱っていません。圧縮は仮説に基づく近似であり、品質評価が前提です。

固定予算で残すトークンを選ぶ

各注意ヘッドに保持できるトークン数の予算Bを設定します。新しいトークンのKVを追加してキャッシュがBを超えたとき、直近wステップの注意重みを見ます。平均的な注意より低い状態が何度も続いたトークンほど、削除候補としてスコアを高くします。

一方、最新rトークンは必ず残します。生成直後のトークンは重要性を判断する観測が足りず、局所的な文法や直近の指示にも必要になりやすいためです。論文の実験ではw=400、r=10を使用しました。削除後は残ったKVだけで通常どおり注意計算を行います。

層ごと・ヘッドごとの予算配分

全層に同じ量を割り当てるのではなく、論文は後段層へ多めの予算を与えます。後段ほど重要性の持続率が少し下がるためです。これは実装上の重要な点で、キャッシュ削減率だけをグローバルに指定しても、品質低下が層に偏る可能性があります。

圧縮処理自体には注意重みを集計する計算が必要です。ただし、毎トークンではなく、一定数を追加してからまとめて削除します。論文では予算の半分に相当する量をまとめて削除する設定を使い、圧縮後の注意計算量削減でオーバーヘッドを相殺する設計にしています。

実験と結果

論文はOPT-6B、13B、30B、66Bを対象に、言語モデリングと5-shot下流タスクで評価しました。実験は4台のNVIDIA A100 40GB GPUで行い、C4の言語モデリング、HellaSwag、MathQA、PIQA、Winograndeを使っています。

最大5倍のKVキャッシュ削減

圧縮率1倍を無圧縮、5倍を元のKV容量の約20%に相当する条件として比較しています。OPT-66Bでは、言語モデリングと下流精度が5倍圧縮まで概ね維持されました。WinograndeとMathQAでは、OPT-66Bの5倍圧縮後も精度維持が報告されています。

ただし小さいモデルほど余裕は小さく、言語モデリングのパープレキシティはOPT-13Bで元のKVの50%程度、OPT-66Bで75%程度まで明確に維持されたと報告されています。すべてのモデル・タスクで「常に5倍」が保証された結果ではありません。モデル規模、入力、評価タスクごとに圧縮率を決めるべきです。

4ビット量子化との併用

HellaSwagで2倍のトークン削減と4ビット量子化を組み合わせたところ、OPT-6Bの精度は無圧縮0.702、Scissorhands 0.706、併用0.704でした。OPT-13Bではいずれも0.720です。系列長方向の削減と、値のビット幅方向の削減は別のレバーであり、組み合わせられることを示しています。

またC4上のOPT-13Bを3倍圧縮した注意重みの比較では、無圧縮との差分は概ねゼロ付近に集中しました。重要トークンの一部を落とす事例は残るものの、注意分布の大部分を近似できたことが品質維持の一因と考えられます。

何に使える?

Scissorhandsは、モデル能力を上げる技術というより、同じGPU予算でより長い文脈または多くの同時リクエストを扱うための技術です。

長文RAGと文書対話

検索文書、会話履歴、生成トークンを長く保持するRAGで有望です。先頭から一律に削る代わりに、注意を集めた文書断片や指示を残せます。特に、回答が長くなるレポート生成や、複数ターンで同一資料を参照する社内アシスタントでメモリ上限を緩められる可能性があります。もっとも、根拠文を必ず保持したい用途では、引用トークンを保護するルールとの併用が必要です。

高同時実行のLLM API

KVメモリがボトルネックのサービングでは、圧縮した分だけバッチ数を増やせます。リクエスト長に応じて予算Bを割り当て、品質低下が許されにくい契約書要約やコード生成には高い予算を与える、といったSLO別の運用が考えられます。

ローカルLLMの長時間エージェント

限られたVRAMでツール実行ログや会話を扱うローカルエージェントにも適しています。要約メモリや外部検索だけに頼らず、モデル内部のKVについても「参照され続ける情報」を残せます。実用時はキャッシュ削減をエージェントの記憶機能と混同せず、必要な事実は外部ストアへ永続化してください。

開発や事業へのヒント

この論文の実務的な価値は、コンテキスト長を単なる入力制限ではなく、情報選別とメモリ配分の問題として捉え直した点にあります。

圧縮率ではなく品質予算を先に決める

小規模プロダクトでは、まず代表的な長文問い合わせを50〜100件用意し、無圧縮を基準に正確性、引用の正しさ、レイテンシ、VRAMを測るとよいです。そのうえで2倍、3倍と圧縮し、どの質問型で誤りが増えるかを見ます。「何倍削れたか」ではなく、許容できる品質低下の範囲で予算Bを決める設計になります。

重要情報を保護する仕組みを加える

Scissorhandsは注意重みに基づく汎用的な削除です。プロダクトでは、システムプロンプト、ユーザーが固定した要件、RAGの引用候補、ツール実行の最新状態などに保護フラグを付け、削除対象から外す拡張が考えられます。これは論文の直接の評価対象ではありませんが、重要性をモデルの注意だけに委ねないための実装上の補完です。

サービング指標にKV使用量を加える

GPU使用率だけでは、長い入力で突然同時実行数が落ちる原因を追いにくいです。リクエストごとのKVトークン数、圧縮発生回数、削除率、応答品質を同じトレースへ記録すると、圧縮設定を改善しやすくなります。Scissorhandsのような技術は、モデル評価と推論基盤の監視をつなぐ領域です。

限界

この手法は、過去に低い注意だったトークンが将来も不要である傾向に依存します。話題が急に戻る会話、長距離の照応、後半で初めて意味を持つ細部では、重要情報を早く捨てる可能性があります。論文の最大実験モデルはOPT-66Bであり、他のアーキテクチャや現在の大規模モデルへの一般化は個別検証が必要です。

また、注意重みを観測・集計できる推論実装が必要です。外部APIのLLMには通常適用できず、vLLMなどキャッシュ管理へ手を入れられる実行基盤が前提になります。圧縮時の集計、各層・各ヘッドの予算配分、量子化との相互作用も実装を複雑にします。

最後に、KVキャッシュはトークンの表現であって監査可能な記憶ではありません。削除後にどの事実を参照できるかを保証する技術ではないため、正確性や根拠提示が必須の業務では、外部検索・引用検証・回帰評価を併用してください。

よくある質問

Q. Scissorhandsはモデルの再学習やファインチューニングが必要ですか?

A. 必要ありません。訓練済みモデルの重みは固定し、推論時にKVキャッシュを予算内へ保つ処理を追加します。ただし注意重みとKVにアクセスできる推論エンジンは必要です。

Q. 単純なスライディングウィンドウと何が違いますか?

A. スライディングウィンドウは古い順に削除します。Scissorhandsは古さではなく、直近の複数ステップで低い注意しか受けなかったかで削除候補を決めます。そのため古くても繰り返し参照されるトークンを残せる可能性があります。

Q. 何倍まで圧縮しても精度は落ちませんか?

A. いいえ。論文ではOPT-66Bの一部タスクで5倍圧縮まで精度維持を示しましたが、小さいモデルや異なる条件では低下し得ます。自分のモデル、入力長、タスクで無圧縮との回帰評価を行ってください。

Q. RAGの根拠文が削除されることはありますか?

A. あります。注意が低いと判定されれば削除候補になります。根拠を必ず残す必要があるなら、引用元トークンを保護する、回答ごとに再検索する、外部メモリから再投入する、といった制御を組み合わせるべきです。

Q. 量子化だけで十分ではありませんか?

A. 量子化は各KV値を小さい表現で保存しますが、保持するトークン数は減りません。Scissorhandsはトークン数を減らすため、長文で増えるメモリに別方向から効きます。論文では4ビット量子化との併用も評価されています。

今日の学び

この論文は、LLM推論で増え続けるKVキャッシュという課題を扱いました。Scissorhandsは、注意を集めたトークンの重要性は持続しやすいという観察を使い、低重要度のKVを固定予算の外へ削除します。

得られるヒントは、長文対応を「より大きなGPUを用意する」だけで考えず、情報の重要度に応じて記憶を割り当てることです。まずは品質評価と保護すべき情報の定義を用意し、KV使用量を運用指標として観測するところから始めるとよいでしょう。