ReFreeKV: 閾値フリーのKVキャッシュ圧縮に向けて

要旨

LLM推論中のメモリ消費を削減するため、KVキャッシュのプルーニング手法がいくつか提案されている。これらの手法は多くのデータセットで損失のないメモリ削減を実現できるものの、多くの場合、十分に強調されていない条件に依存している。すなわち、最適なパフォーマンスを得るためには、入力やドメインに固有のKVキャッシュ予算の閾値を事前に決定する必要がある。しかし、このような入力依存型の設計は、実世界のシナリオではかなり制限される可能性がある。なぜなら、オープンドメインの入力は多様なドメイン、長さ、難易度にわたっており、閾値選択の明確な境界がないからである。その結果、このような入力依存型閾値への依存は、任意の入力に対して大きな性能劣化を引き起こす根本的な制限となり得る。本研究では、ロバストなKV圧縮のための閾値制約を取り除く新しい目的を提案し、フルキャッシュの性能を維持しながら予算配分を適応的に調整する「閾値フリー」な手法を提唱する。そして、この目的の最初の具体化として、新たな手法ReFreeKVを提案する。多様なコンテキスト長、タスクタイプ、モデルサイズを持つ13のデータセットにわたる広範な実験により、その有効性と効率性を実証する。我々のコードはhttps://github.com/Patrick-Ni/ReFreeKVで公開されている。

English

To reduce memory consumption during LLM inference, a handful of methods have been proposed for KV cache pruning. While these techniques can accomplish lossless memory reduction on many datasets, they often hinge on an under-emphasized condition: an input/domain-specific threshold for KV cache budget needs to be pre-determined to achieve the optimal performance. However, such input-sensitive design may be considerably limited in real-world scenarios, as open-domain inputs span diverse domains, lengths and difficulty levels, without clear boundaries for threshold selection. As a result, the dependence of such input-sensitive threshold can be a fundamental limitation that causes large degradation on arbitrary inputs. In this work, we propose a new objective that lifts the threshold constraints for robust KV compression, advocating for "threshold-free" methods that adaptively adjust budget allocation while preserving full-cache performance. We then propose a novel method, ReFreeKV, serving as the first instantiation of this objective. Extensive experiments across 13 datasets with diverse context lengths, task types, and model sizes demonstrate its efficacy and efficiency. Our code is publicly released at https://github.com/Patrick-Ni/ReFreeKV.