LeanK: Learnable K Cache Channel Pruning for Efficient Decoding

Large language models (LLMs) enable long-context tasks but face efficiency challenges due to the growing key-value (KV) cache. LeanK proposes a novel KV cache pruning method that leverages the static sparsity of QK vectors along the channel (head_dim) dimension. A static, channel-wise pruning mask is derived through a novel two-stage training process. Experimental results demonstrate up to 70% reduction in K cache and 16–18% reduction in V cache, with minimal impact on end-to-end model performance.
This method introduces a new direction for KV cache compression that is orthogonal to existing approaches such as token-level compression and quantization. In addition, LeanK observes that the channel-wise norm distribution of QK vectors in each attention head correlates with the head’s retrieval capacity. This insight opens up promising avenues for enhancing the long-context understanding abilities of LLMs.