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.