A Self-Pruning Transformer: Extreme KV-Cache Compression with Universal Attention
Davis Wertheimer · Haochen Shen · Ahan Gupta · Derrick Liu · Yu Chin Fabian Lim · Mudhakar Srivatsa · Raghu K. Ganti · Minjia Zhang · Naigang Wang
Universal Attention unifies softmax attention with linear attention decay mechanisms. By leveraging this learned decay as a natural pruning criterion, it achieves state-of-the-art KV-cache compression and long-context gains with high efficiency.
摘要
The large KV-cache size of modern LLMs creates a barrier to efficient deployment. Recent work has explored replacing attention layers' RoPE positional embeddings with alternative decay-based mechanisms, which can then be used to prune KV-cache during inference. However, these decay functions have limited expressivity, and in practice devolve into sliding-window-like eviction patterns. In this work, we propose a unifying framework for complementary and novel decay mechanisms, capturing complex key statistics and interactions while preserving expressive RoPE embeddings and Softmax attention. The resulting *Universal Attention* is a highly expressive and end-to-end trainable architecture, whose composite decay mechanism acts as a natural, *adaptive* pruning criterion, removing tokens that contribute least to attention computation. Experimentally, Universal Attention achieves state-of-the-art $10\times$ compression on natural language and synthetic task data, while *improving* downstream performance compared to both state-of-the-art baselines and unpruned oracles. It further demonstrates superior long-context generalization with unprecedented $25\times$ compression at length 16k.