Cache compression improves large language model memory use and accuracy
KV-Kaizen: Learning Context-Adaptive Cache Compression Choices
Artificial Intelligence
Summary
Large language models (LLMs) keep track of recent words using memory called KV caches, which grow big and slow down the model. The authors found a way to shrink this memory without throwing away any stored words by choosing the best compression method for each part of the model based on the current text. Their method, called KV-Kaizen, mixes three ways to reduce memory without hurting the model’s accuracy. This approach works better than just removing old words and can make big models run faster while keeping their performance.
What this means in practice
- •For ai infrastructure engineers: Reduce memory usage of large language model inference caches to speed up decoding and lower cost without losing output quality.
- •For ai platform developers: Build LLM-serving systems that adaptively compress caches per model layer to maintain accuracy while supporting longer contexts.
Authors
Joao Monteiro, Louis Béthune, Anastasiia Filippova, Sonia Laguna, David Grangier, Marco Cuturi
Abstract
As the context size of text processed with an LLM grows, the size of KV caches can outstrip the memory allocated for the original model weights. This impacts LLM throughput negatively, since decoding is memory-bound and decode cost grows with cache size. Recent work alleviates this bottleneck by discarding the least relevant tokens. Eviction introduces a tension, since a one-off decision to discard content may prove detrimental later. Instead, we focus on alternative choices that can lead to cache compression without evicting tokens. We achieve this by learning a selector that is able to produce, based on context, a per-layer cache configuration towards an overall compression budget. The selector operates along three axes: sharing one cache across layers (depth), caching at fewer bits (precision), or truncating the low-rank latent cache representations (rank). We call the resulting method KV-Kaizen, for the many small per-layer choices it compounds. We observe that these interventions taken independently and uniformly over all layers limit achievable compression because they degrade accuracy. Crucially, composing them locally and adaptively to the context can instead preserve accuracy while achieving large memory savings. At inference, the selector runs once, before pre-fill. In evaluations on instruction following and reasoning tasks, our selectors reach the Pareto frontier of accuracy against cache size, against learning-free and post-hoc baselines. On long-context tasks, KV-Kaizen improves on eviction and can be composed with it, reaching a 32x smaller decode-time cache on a 14B model while preserving accuracy. A 4x cache size reduction incurs no accuracy degradation from 7B parameters up, and a compressed model is more accurate than a smaller uncompressed one with the same cache size. Together, these findings support pre-training large models and compressing them only afterwards.