DropKV leads tested eviction methods with a small token cache
In plain words
Language models keep information about earlier words in a cache so they can use it as they continue processing a long passage.
To make that cache smaller, DropKV estimates how much removing each saved token would change the model’s attention output, then removes tokens with the smallest estimated effect.
On RULER, a test of long-context tasks, DropKV had higher average accuracy than the other tested removal methods across three models at a 128K-token context and a 5% cache budget. That result applies to those models and settings.
The research
A cache can be made smaller by permanently removing saved token information, but choosing several tokens together to minimize the change in a model’s attention output is computationally hard. Can a simple per-token score make that choice tractable without losing too much task performance?
DropKV scores cached tokens by their estimated effect on attention output, then removes the least disruptive ones to preserve task performance with a limited cache.
DropKV scores each token by its estimated effect on attention output, adds squared scores across recent queries, then removes the lowest. This replaces a difficult joint choice with per-token scoring.
At a 5% cache budget, its average scores exceeded competing eviction methods on RULER’s long-context tasks and 16 English and code LongBench tasks across three models. That supports its performance in these tested settings.
The theoretical bound applies only when attention weights are small and the cone condition holds; that condition failed for some sampled triples. The evaluations cover three models and the specified tasks and cache budgets, and the Triton acceleration is for prefill scoring, not decode-time scoring.
Analysis
Deep
- Original paper authors
- Aozhong Zhang, Selcuk Gurses, Yanxia Deng, Naigang Wang, Chi-Chun Liu
- Submitted to arXiv
- License
- CC BY 4.0