ZigZagkv: Dynamic KV Cache Compression for Long-context Modeling based on Layer Uncertainty
This addresses memory efficiency for long-context modeling in LLMs, offering an incremental improvement over existing uniform compression methods.
The paper tackles the problem of out-of-memory issues in large language models due to growing KV caches during long-context inference by proposing a dynamic compression method based on layer uncertainty, which reduces memory usage to ~20% compared to full KV inference while maintaining nearly lossless performance.
Large Language models (LLMs) have become a research hotspot. To accelerate the inference of LLMs, storing computed caches in memory has become the standard technique. However, as the inference length increases, growing KV caches might lead to out-of-memory issues. Many existing methods address this issue through KV cache compression, primarily by preserving key tokens throughout all layers to reduce information loss. Most of them allocate a uniform budget size for each layer to retain. However, we observe that the minimum budget sizes needed to retain essential information vary across layers and models based on the perspectives of attention and hidden state output. Building on this observation, this paper proposes a simple yet effective KV cache compression method that leverages layer uncertainty to allocate budget size for each layer. Experimental results show that the proposed method can reduce memory usage of the KV caches to only $\sim$20\% when compared to Full KV inference while achieving nearly lossless performance.