PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.09051v1cs.LG

TL;DR

Softmax attention’s per-layer KV-cache consumes substantial memory, and existing decay gates can be too inexpressive for extensive pruning. Universal Attention combines key-statistic-based decay mechanisms while preserving RoPE and softmax attention. It reports over 10× compression with improved downstream performance, and 25× compression at length 16k.

  • Problem

    Softmax attention requires retaining every token’s keys and values in per-layer KV-caches, while existing forget-gates often devolve into limited sliding-window-like eviction.

    Source: §1 · p. 1

  • Method

    Universal Attention combines first- and second-order key statistics into expressive forget gates and integrates decay with RoPE and softmax attention.

    Source: §1 · p. 1, §7 · p. 10

  • Results

    Across natural-language and synthetic tasks, UA achieves over 10× KV-cache compression while improving downstream performance over state-of-the-art methods and unpruned oracles; at length 16k, compression reaches 25×.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Trainable self-pruning architectures provide a practical avenue to mitigate large inference-memory overhead while retaining RoPE and softmax attention.

    Source: §7 · p. 10

    Early results suggest existing models can be adapted to UA, but a significant performance gap remains compared with fully pretrained UA, constraining deployment feasibility.

    Source: §6 · p. 10

Abstract

from arXiv · show

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, $\textit{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 $\textit{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.

1 Introduction

Large per-layer KV-caches make softmax-attention LLMs costly to deploy, while existing forget-gates often yield sliding-window-like pruning. Universal Attention combines key statistics into expressive gates and reports strong compression alongside improved downstream performance.

1 Introduction

  • Softmax attention requires preserving every token in each layer’s KV-cache during inference, consuming gigabytes or more and obstructing efficient deployment at scale.Source: §1 · p. 1
  • Existing forget-gates can devolve into soft sliding windows, limiting how many key/value entries they can evict before downstream performance degrades.Source: §1 · p. 1
  • Universal Attention combines first- and second-order statistics of stored keys to produce uniquely expressive forget gates.Source: §1 · p. 1
  • Over 10× KV-cache compression on natural-language and synthetic tasks accompanies improved downstream performance over state-of-the-art pruning methods and unpruned oracle baselines; at length 16k, compression reaches 25×.The reported model is a 1B transformer trained on 50B tokens.Source: §1 · p. 2

2 Related Work

Related work spans several orthogonal approaches to reducing KV-cache costs. UA is positioned among them as adaptive pruning that retains softmax attention and RoPE while combining richer decay mechanisms.

Related work overview

  • KV-cache reduction includes sublinear-cache architectures, reduced cache access, post-training pruning, and trained forget gates; UA sits at their intersection.Source: §2 · p. 2

Sub-quadratic architectures

  • Linear-attention models use a constant-sized recurrent state, but retrieval becomes noisy beyond its state dimension and quality can drop on complex long-context tasks.UA instead applies expressive decay directly to softmax attention, avoiding the memory lower bound imposed by a constant-sized recurrent state.Source: §2 · p. 2

Sparse attention

  • Sparse-attention methods compute only part of the attention matrix, but the described implementations do not reduce KV-cache size.Static sliding-window patterns can also degrade long-context capability unless combined with more expressive layers.Source: §2 · p. 2

KV-cache eviction

  • KV-cache eviction methods such as H2O and SnapKV use inference-stage heuristics, while UA dynamically determines which tokens are dispensable rather than pruning to a preset target.The passage contrasts fixed-budget eviction with UA’s adaptive pruning and reports greater compression at long sequence lengths alongside more general performance improvement.Source: §2 · pp. 2–3

Data-dependent forget gates

  • Some data-dependent forget gates devolve into soft sliding windows, while more expressive delta-rule mechanisms preclude standard Softmax or RoPE elements.UA combines multiple gates as an additive bias to the standard attention matrix, retaining RoPE and SoftMax alongside more sophisticated decay patterns.Source: §2 · p. 3

3 Background

The background introduces notation and standard softmax self-attention, then describes data-dependent forget-gates as a way to decay and evict past key/value tokens.

Softmax multi-head self-attention

  • Standard softmax multi-head self-attention operates on an input sequence of token representations, with L, n, and d denoting sequence length, head count, and head dimension.Source: §3 · p. 3
  • The causal attention mask M is broadcast across heads, while Q, K, and V stack the query, key, and value tensors; new tokens attend to prior K and V entries stored per layer.That per-layer storage forms the inference-time KV-cache.Source: §3 · p. 3

Data dependent forget-gates

  • Forget-gates decay past key/value tokens so sufficiently decayed entries can be deleted; FoX expresses its gate as an additive bias to the attention matrix.FoX computes a scalar gate for each input token.Source: §3 · p. 3
  • FoX’s accumulated decay decreases monotonically away from the attention matrix’s main diagonal, but entries decay in lockstep, producing an inexpressive sliding-window-like eviction pattern.The passage says sufficiently decayed keys can then be safely deleted from the cache.Source: §3 · p. 3

4 Universal Attention

Universal Attention augments softmax attention with learned context-aware decay, combining complementary mechanisms into adaptive sparsity patterns. Its decay also guides token-level KV-cache pruning, while custom kernels support efficient training.

4 Universal Attention

  • UA adds context-aware decay to softmax attention and combines multiple decay patterns to learn adaptive, dynamic sparsity.The decay mechanisms are combined by geometric mean to form a unified decay rate for each position pair.Source: §4 · p. 3, Figure 1 · p. 4

4.1 Decay Components

  • Three mechanisms capture query- and key-axis first-order statistics and pairwise second-order interactions; the key-aware gate gives each token its own intrinsic decay rate.That rate is independent of future updates, allowing tokens to persist or self-prune quickly on a case-by-case basis.Source: §4.1 · pp. 4–5
  • The DeltaNet-style similarity mechanism is adapted into a multiplicative, key-similarity-based decay factor rather than dynamically shifting cached keys.The original shifting mechanism interferes with RoPE and would require overwriting the entire key cache at each step.Source: §4.1 · p. 5

4.2 Layer Design and Algorithm

  • At inference, UA evicts a key-value pair when its learned effective mask falls below threshold p; the decreasing decay bias means evicted tokens will not regain importance.The model determines which and how many tokens to retain on a head-by-head basis.Source: §4.2 · p. 6
  • Custom kernels and activation checkpointing support UA’s implementation, which raises end-to-end training throughput by 1.13× versus PyTorch-2.The two kernels separately compute the decay mask and attention output.Source: §4.2 · p. 6

5 Evaluation

In a 1B Llama3-style setup, Universal Attention (UA) is evaluated on retrieval and reasoning alongside attention baselines and cache-pruned Llama3 variants. UA offers favorable memory–performance trade-offs across downstream tasks and longer contexts, while learning input-adaptive pruning patterns and benefiting from complementary decay mechanisms.

5 Evaluation

  • The 1B Llama3-architecture models are trained for 50B tokens; evaluation covers retrieval and common-sense reasoning, with cache-pruned Llama3 baselines and attention-model baselines trained identically to UA.Long-context inputs use YaRN for Llama3-based models, and UA prefill pruning proceeds in chunks of 128 tokens.Source: §5 · pp. 6–7

5.1 Downstream Performance

  • On RULER (4K), UA claims the top two average scores at its pruning thresholds; on LongBench-E (0–4K), p=0.15 uses under half Mamba2-H’s memory and exceeds Llama3 and Llama3-G with a fraction of their cache.Source: §5.1 · p. 7, Table 1 · p. 7

5.2 Context Length Extension

  • As sequence length rises from 1K to 16K, UA traces the strongest score–cache frontier; against Llama3, p=0.15 reduces cache 4.5×–45.3× on RULER and 7.9×–26.9× on LongBench-E.The p=0.05 variant preserves higher scores, while p=0.15 compresses more aggressively; UA cache remains under 25 MB at the longest contexts shown.Source: §5.2 · p. 8, Figure 2 · p. 9

5.3 Empirical Analysis

  • UA learns mixed decay patterns, varies cache size across inputs with growth-and-reset behavior, and sparsifies early in training before gradual recovery as language-modeling capabilities improve.The three illustrated decay mechanisms are key-linked, similarity-based, and query-linked; cache sizes eventually stabilize and gently decline over the latter half of training.Source: §5.3 · pp. 8–9, Figure 3 · p. 9

5.4 Ablations

  • All decay components contribute; removing query-aware f(q) causes the smallest degradation, while p=0.15 yields 21.4× LongBench-E and 14.9× validation compression with less than one-point changes in score and perplexity.No universal threshold is identified; p=0.05 and p=0.15 are representative operating points for gentler and more aggressive pruning.Source: §5.4 · pp. 9–10, Table 3 · p. 10

6 Limitations

Universal Attention has limitations in aggregation performance, computation overhead, training from scratch, and inference optimization. The authors identify hybrid memory mechanisms and specialized inference support as relevant paths forward.

6 Limitations

  • UA is comparatively weaker on aggregation tasks that combine weak evidence across many positions, though complementary memory mechanisms can mitigate this weakness.The authors contrast this with UA’s strengths on retrieval, tracking, and QA.Source: §6 · p. 10
  • UA adds computation to construct learned decay and pruning masks, leaving that overhead as a systems optimization target.This computation enables adaptive, content-dependent retention.Source: §6 · p. 10
  • Adapting existing models to UA appears possible, but early results remain significantly behind fully pretrained UA.The authors identify closing this gap as important for practical deployment.Source: §6 · p. 10
  • Preliminary experiments show lower decode latency with UA’s reduced cache, but the implementation lacks a fused, inference-optimized dynamic cache manager.The authors identify specialized sparse kernels and systems support as future work for efficiently exploiting dynamic sparsity.Source: §6 · p. 10

7 Conclusion

Universal Attention combines generalized and novel decay mechanisms with RoPE and softmax attention in one layer. Its adaptive pruning yields improved memory-performance trade-offs, and the authors present trainable self-pruning as a practical avenue for reducing inference memory overhead.

7 Conclusion

  • UA generalizes existing decay mechanisms and adds a novel one, combining them with RoPE and softmax attention in a unified layer.Source: §7 · p. 10
  • UA improves downstream performance with lower memory overhead than O(1) linear attention models.The authors attribute its memory-performance trade-off to an expressive decay mechanism that provides an adaptive pruning criterion.Source: §7 · p. 10
  • Trainable self-pruning architectures provide a practical avenue to mitigate large memory overhead during inference.Source: §7 · p. 10

A Training, Model, and Evaluation Details

The appendix documents training and evaluation settings, procedures for comparing stored-state and KV-cache sizes, and memory-efficient UA kernel design.

A Training, Model, and Evaluation Details

  • Models use 50B training tokens for Sections 5.1–5.3 and 10B for Section 5.4, with Dolma v1.7 data, context length 4096, and 0.5M tokens per minibatch.Training uses 16 NVIDIA A100s, AdamW, a 3e−4 learning rate, 2k-step warmup, cosine decay, and weight decay 0.1.Source: A Training, Model, and Evaluation… · p. 20
  • The setup specifies Llama-3 tokenization, 32 layers and dmodel = 1280 by default, plus model-specific Mamba2, Gated DeltaNet, and hybrid configurations.Mamba2 and Gated DeltaNet reduce dff to keep total size approximately 1B; hybrid layers replace layers 4, 12, 20, and 28 with grouped-query attention.Source: A Training, Model, and Evaluation… · p. 20
  • Evaluation covers LongBench-E, RULER’s default synthetic task set, and common-sense reasoning benchmarks using the EleutherAI evaluation harness.Source: A Training, Model, and Evaluation… · p. 20

B Compute State Size & KV-Cache Size

  • The appendix compares stored-state and KV-cache sizes as sequence length grows, calculating actual saved-tensor sizes for hybrid models and model configurations.The comparison includes Mamba2 and Gated DeltaNet state shapes, GQA cache size, and bf16 storage.Source: B Compute State Size & KV-Cache Size · pp. 20–21, Table 4 · p. 21

C Kernel Design

  • UA’s memory-efficient kernels compute decay masks and attention outputs, applying the mask within fused online softmax to avoid expensive HBM I/O.Activation checkpointing avoids saving the quadratic decay mask per layer.Source: C Kernel Design · p. 22

D Hybrid Model Results

Hybrid UA models retain strong performance, with UA-Softmax close to pure UA on reported averages. Adding full-attention or Mamba layers improves aggregation but reduces tracking performance.

D Hybrid Model Results

  • UA-Softmax replaces one in every eight UA layers with standard full attention, while UA-Mamba alternates UA and Mamba layers evenly.The models are evaluated without KV-cache pruning to isolate hybridization effects.Source: D Hybrid Model Results · p. 23
  • Introducing full-attention or Mamba layers substantially improves aggregation performance relative to pure UA while reducing tracking performance.The authors suggest these layers can provide complementary memory behavior when stronger aggregation is desired.Source: D Hybrid Model Results · p. 23
  • RULER/LongBench-E averages are 48.8/17.6 for UA-Softmax versus 48.5/18.1 for UA, while UA-Mamba scores 47.9/15.9.Table 5 compares hybrid UA models with pure Llama3 and UA without eviction.Source: Table 5 · p. 23, D Hybrid Model Results · p. 23

E Additional Ablation Results

Additional ablations show that UA’s combined decay mechanisms retain long-context performance across normalization choices, while 3B results extend its quality–memory advantages beyond the 1B setting.

E Additional Ablation Results

  • UA with Qwen3-style normalized queries and keys performs effectively, with slightly improved performance at 16K; the ablations evaluate 1B models trained on 10B tokens.The comparison also includes individual decay mechanisms and sliding-window attention.Source: E Additional Ablation Results · p. 24, Table 6 · p. 25
  • At 16K, query-aware decay alone scores 40.0 versus 58.4 for full UA; key-aware and similarity-aware decay alone also underperform the complete model.The three decay signals provide complementary retention behavior; sliding-window attention is competitive only at the shortest context and degrades sharply as context increases.Source: E Additional Ablation Results · pp. 24–25

F 3B Model Results

  • At 3B, RULER (4K) rises from 48.8 for Llama3-3B to 52.6 for UA, while p=0.05 cuts cache from 341.8 MB to 46.6 MB (7.3×).On LongBench-E (0–4K), UA scores 23.5 versus 21.5 for Llama3-3B, with cache reduced from 271.9 MB to 41.5 MB (6.5×).Source: F 3B Model Results · p. 25, Table 7 · p. 25

G Post-Training Adaptation

The paper converts a pretrained Llama3 checkpoint to UA and continues training. Adaptation improves performance and reduces cache relative to the original model, though training from scratch yields a stronger performance–memory trade-off.

G Post-Training Adaptation

  • At p=0.05, Llama3-UA raises RULER from 40.8 to 44.3 while reducing cache from 142.4 MB to 38.6 MB; scratch UA reaches 50.3 with a 21.4 MB cache.The adapted model remains below UA trained from scratch at the same pruning threshold.Source: G Post-Training Adaptation · pp. 25–26, Table 8 · p. 26
  • On LongBench-E, adapted Llama3-UA improves from 16.4 to 17.2 while reducing cache from 113.3 MB to 41.3 MB.The conversion retains pretrained Llama3 weights, initializes UA-specific parameters randomly, and continues training for 10B tokens.Source: G Post-Training Adaptation · pp. 25–26, Table 8 · p. 26

H Latency Analysis

UA’s current implementation shows decoding speedups with KV-cache pruning, while the paper notes that optimized dynamic-cache inference is not yet implemented. Detailed RULER comparisons report better overall scores with smaller caches for UA at 8K and 16K.

H Latency Analysis

  • The current system lacks an inference-optimized dynamic cache manager; efficient operation on dynamically pruned caches requires sparse matrix–vector computation and specialized sparse kernels.These remain active areas of systems research.Source: H Latency Analysis · p. 26
  • KV-cache pruning yields approximately 1.1×–1.4× per-token decode speedup (11–29% lower latency), with larger benefits at longer prompt lengths.The latency table reports per-step decode latency with and without pruning at different sequence lengths.Source: H Latency Analysis · p. 26, Table 9 · p. 26
  • Custom UA kernels improve end-to-end training throughput by approximately 1.13×; the paper suggests this may also reduce TTFT but reports no matched inference-prefill baseline.Prefill exercises the same forward attention path.Source: H Latency Analysis · p. 26

I Detailed RULER results for Figure 2

Related papers

Loading 2610.09051v1…