Universal Attention composes decay gates for over 10x KV-cache compression and beats unpruned Llama3 on RULER
Synopsis
The work introduces Universal Attention, which combines query-aware, key-aware, and similarity-based decay mechanisms via a geometric mean into an additive bias on the attention matrix, letting the model adaptively prune its KV-cache while keeping RoPE and Softmax; on a 1B Llama3-style model trained on 50B tokens it achieves over 10x compression, with RULER averages of 50.3/49.3 versus 40.8 for unpruned Llama3 at only 21.4/12.9 MB of cache.
Figure 1: Overview of Universal Attention. ( Left ) Standard softmax attention with a fixed causal mask 𝐌 \mathbf{M} . ( Middle ) UA augments the attention logits 𝐐𝐊 ⊤ \mathbf{QK}^{\top} with a learned, content-aware decay bias 𝚲 \mathbf{\Lambda} , while applying L 2 L_{2} -normalization to keys, and a low-rank sigmoid gate to the attention output. ( Right ) Construction of 𝚲 \mathbf{\Lambda} from three decay mechanisms: similarity-based f sim f_{\text{sim}} , query-aware f q f_{\text{q}} , and key-aware f k f_{\text{k}} . These are combined via a geometric mean, transformed to log-space, and accumulated causally to produce the decay matrix.
arXivInterpretation
A unifying framework composes multiple context-aware decay mechanisms via their geometric mean into a single decay rate added as a bias to Softmax attention, preserving RoPE and Softmax while yielding richer decay patterns than a single forget gate. Prior forget gates such as FoX use a single scalar decay that in practice devolves into a soft sliding window; this work combines query-aware, key-aware, and similarity decay, and notes that the geometric mean preserves a veto property so that if any mechanism deems a token should not decay, the combined decay is driven to zero. The three mechanisms are derived formally, and ablations on 1B models trained on 10B tokens show that removing key-aware decay lowers RULER S1-S3 by 36.20 at 16k and removing similarity decay by 31.76, while query-aware decay alone (FoX-style) reaches only 40.0 at 16k versus 58.4 for full UA.
The learned decay itself serves as an adaptive pruning criterion: at inference a scalar bias per key/value pair is maintained and entries whose effective mask falls below a threshold are evicted, with the bias decreasing monotonically so evicted tokens do not return to importance. H2O, SnapKV, and KeyDiff prune to a preset budget (e.g., keeping 10% of the cache), whereas UA lets the model decide per head which and how many tokens to retain, with compression varying by input and length. On the 50B-token model at threshold 0.15, LongBench-E at 16k shows average compression of 21.38 and maximum single-sequence compression of 113.95, with LongBench-E score and validation perplexity each changing by less than one point; during training the cache first sparsifies rapidly, then re-incorporates past information as language modeling improves, and gently falls over the last half of training.
At 1B and 3B scale UA improves both quality and memory: it exceeds unpruned Llama3 and Llama3-G on RULER while using far less cache than linear-attention and hybrid baselines. Unlike linear attention and hybrid models, UA applies decay directly to Softmax attention, avoiding the memory lower bound imposed by a constant-sized recurrent state; unlike pruning baselines, it retains higher scores under a tighter cache budget. The 1B model reaches RULER (4k) averages of 50.3 (p=0.05, 21.4 MB) and 49.3 (p=0.15, 12.9 MB), versus 40.8 (143.7 MB) for Llama3, 41.9 for Llama3-G, and 44.1 (46.0 MB) for GDN-H; at 3B, RULER averages 52.6 versus 48.8 for Llama3-3B, with cache reduced from 341.8 MB to 46.6 MB.
Long-context generalization: UA's cache plateaus with length and remains usable at 16k, whereas full-attention baselines grow linearly. Evaluated beyond the training length from 1K to 16K, UA stays on the strongest memory-performance frontier and keeps cache under 25 MB. Per-task RULER results in Appendix I show UA at 16k with 22.7/13.6 MB of cache versus 299.0 MB for Llama3; GDN-H with YaRN improves at 8K/16K but UA still achieves a better overall RULER score with a smaller effective cache.
Perspective
The results apply to 1B and 3B Llama3-style models with RoPE and GQA pretrained at 4k context on Dolma v1.7 with the data-mixture reweighting of Chu et al., with long context handled by YaRN. The intended setting is long-context inference that benefits from strong retrieval and tracking and can accept per-head adaptive sparse caches; the hybrid experiments show that combining UA with full-attention or Mamba layers improves aggregation tasks, so UA is best suited as a cache-efficient attention component within broader long-context architectures. The threshold can be chosen for the desired memory-quality trade-off, with 0.05 and 0.15 used as representative gentler and more aggressive operating points.
UA is comparatively weaker on aggregation tasks, which the paper attributes to tasks requiring weak evidence distributed across many positions, and notes hybrids can mitigate this; constructing the decay and pruning masks adds computation, and reducing that overhead remains a systems optimization target. The current implementation lacks a fully fused, inference-optimized dynamic cache manager, and the latency gains in Appendix H come from a naive implementation, so whether dynamically sparse caches pay off in real serving still awaits specialized sparse kernels and systems support. Adapting existing checkpoints through continued training is feasible but remains clearly below training from scratch. In addition, several equations and numeric values are missing from the parsed text, so exact hyperparameters and compression ratios should be checked against the original appendices and tables.