A Self-Pruning Transformer: Extreme KV-Cache Compression with Universal Attention

# A Self-Pruning Transformer: Extreme KV-Cache Compression with Universal Attention
Source: [https://arxiv.org/html/2610.09051](https://arxiv.org/html/2610.09051)
Haochen ShenAffiliation:IBM Research, USA SSAIL Lab, University of Illinois Urbana\-ChampaignEmail:[hshen14@illinois\.edu](mailto:)Ahan GuptaAffiliation:IBM Research, USA SSAIL Lab, University of Illinois Urbana\-ChampaignEmail:[ag82@illinois\.edu](mailto:)Derrick LiuYu Chin Fabian LimMudhakar SrivatsaRaghu K\. GantiMinjia ZhangAffiliation:IBM Research, USA SSAIL Lab, University of Illinois Urbana\-ChampaignNaigang Wang

###### Abstract

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,adaptivepruning criterion, removing tokens that contribute least to attention computation\. Experimentally, Universal Attention achieves state\-of\-the\-art10×10\\timescompression on natural language and synthetic task data, whileimprovingdownstream performance compared to both state\-of\-the\-art baselines and unpruned oracles\. It further demonstrates superior long\-context generalization with unprecedented25×25\\timescompression at length 16k\.

## 1Introduction

Large Language Models \(LLMs\) power many of the applications we use today\([OpenAI, 2025](https://arxiv.org/html/2610.09051#bib.bib38);[Anthropic, 2026](https://arxiv.org/html/2610.09051#bib.bib39);[Google DeepMind, 2026](https://arxiv.org/html/2610.09051#bib.bib40);[xAI, 2026](https://arxiv.org/html/2610.09051#bib.bib41)\)\. However, they predominantly employ softmax\-attention layers\([Vaswani et al\., 2017](https://arxiv.org/html/2610.09051#bib.bib1)\), which compute pairwise dot\-products betweenalltokens in the input sequence\. This entails preservingeverytoken, foreachlayer, as the “KV\-cache” during inference time, consuming gigabytes \(if not terabytes\) of device memory, often several times larger than the model itself\. This acts as a significant obstacle to deploying LLMs efficiently at scale\.

To effectively reduce KV\-cache sizes, recent research has proposed the addition of a “forget\-gate” to the softmax self\-attention formulation, enabling unimportant key/value tokens from the past to be deleted\. For example, forgetting transformers\([Lin et al\., 2025a](https://arxiv.org/html/2610.09051#bib.bib2);[Lin et al\., 2025b](https://arxiv.org/html/2610.09051#bib.bib3)\)\(FoX\) introduces and applies a scalar data\-dependent decay term for each new token entering the cache, accumulating over time as an additive bias to the attention matrix\. This effectively zeros out certain query\-key products after softmax normalization, allowing the corresponding key/value pair to be safely evicted\. However, these forget\-gates are relatively inexpressive, yielding simple patterns that devolve to soft sliding\-windows, which can only evict limited key/value entries before downstream model performance degrades\.

In this work, we introduce*Universal Attention*\(UA\), a unifying framework that combines first\- and second\-order statistics of stored keys to produce uniquely expressive forget gates\. Some components of our framework align with and generalize decay mechanisms from popular Linear Attention models, such as DeltaNet\([Yang et al\., 2024b](https://arxiv.org/html/2610.09051#bib.bib16);[Yang et al\., 2025a](https://arxiv.org/html/2610.09051#bib.bib9)\)& Mamba\([Gu and Dao, 2023](https://arxiv.org/html/2610.09051#bib.bib7);[Dao and Gu, 2024](https://arxiv.org/html/2610.09051#bib.bib8)\)\. But while these models rely on a constant\-sizedd×\\timesdstate that stores at mostdlinearly independent key/value associations, we combine these and novel mechanisms into a single pruning technique for the full KV\-cache of SoftMax attention, while preserving RoPE embeddings\([Su et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib28)\)\. We additionally implement hardware\-aware custom UA kernels for high\-throughput training\. The resulting UA architecture is efficient and expressive, achieving over 10×\\timesKV\-cache compression on natural language and synthetic task data, for a 1B transformer model trained on 50B tokens\. This memory overhead is lower than state\-of\-the\-art pruning methods and evenO⁡\(1\)O\(1\)linear attention architectures, despiteimproveddownstream performance, even compared to unpruned oracle baselines\. Cache pruning in UA is also adaptive and input\-dependent, generalizing to sequence lengths not observed during training, with unprecedented25×25\\timescompression in long contexts of 16k\.

## 2Related Work

Shrinking the KV cache is a topical problem with several orthogonal approaches\. These include dedicated architectures with sublinear cache growth at inference time, techniques to reduce KV\-cache access without reducing the overall size, post\-training pruning techniques for vanilla architectures, and combined architectures that use trained forget gates to decay KV cache entries\. UA sits at the intersection of these approaches\.

Sub\-quadratic architectures:Linear attention models collapse the sequence length into a constant,ds​t​a​t​e×\{\}\_\{state\}\\timesdstate\-sized, recurrent hidden state[Katharopoulos et al\. \(2020\)](https://arxiv.org/html/2610.09051#bib.bib43)\. However, when input sequence lengths exceeddstate, retrieval from the hidden state becomes noisy\. To mitigate this issue, models in this family impose a recency bias by decaying past key/value associations\([Gu and Dao, 2023](https://arxiv.org/html/2610.09051#bib.bib7);[Dao and Gu, 2024](https://arxiv.org/html/2610.09051#bib.bib8);[Yang et al\., 2025a](https://arxiv.org/html/2610.09051#bib.bib9)\)\. While still highly expressive, these models experience considerable drops in quality for complex long\-context tasks compared to Softmax self\-attention\([Jelassi et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib45);[Ben\-Kish et al\., 2025](https://arxiv.org/html/2610.09051#bib.bib44)\)\.

In addition, despite their O\(1\) memory in the limit, such models can fail to reduce overhead in practice: for typical model configurations, linear attention consumesgreatermemory for sequences shorter than 1k tokens \- a sizable portion of the usual training length \(see Appendix[B](https://arxiv.org/html/2610.09051#A2)\)\. Fixes aimed at improving model capacity and performance only exacerbate the issue: hybrid models interleaving full attention with linear attention still have the same break\-even point, but become O\(L\) in the limit, reducing best\-case memory savings to 8\-10×\\timesfor typical settings where1/81/8to1/101/10of layers are switched to Softmax\([NVIDIA et al\., 2025](https://arxiv.org/html/2610.09051#bib.bib32);[Lieber et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib46);[Team et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib47)\)\. Log\-Linear Attention\([Guo et al\., 2026](https://arxiv.org/html/2610.09051#bib.bib4)\)maintains multiple state buffers, pushing out the break\-even point even further\. UA instead applies an expressive decay mechanism to Softmax self\-attention directly, avoiding the memory lower\-bound imposed by a constant\-sized recurrent state\.

Sparse attention:[Zaheer et al\. \(2020\)](https://arxiv.org/html/2610.09051#bib.bib13);[Child et al\. \(2019\)](https://arxiv.org/html/2610.09051#bib.bib14);[Yuan et al\. \(2025\)](https://arxiv.org/html/2610.09051#bib.bib6)compute only a subset of the attention matrix\. These techniques range from computing static patterns \(such as sliding\-windows around the diagonal\) to dynamic patterns based on custom data\-dependent heuristics\. Simple patterns like sliding\-window\-attention reduce KV\-cache sizes but degrade long\-context capability unless interleaved with more expressive layers\([Jiang et al\., 2023](https://arxiv.org/html/2610.09051#bib.bib48);[Beltagy et al\., 2020](https://arxiv.org/html/2610.09051#bib.bib50);[Team et al\., 2025](https://arxiv.org/html/2610.09051#bib.bib49)\),while content\-dependent dynamic sparsity provides a more expressive alternative\([Zhang et al\., 2025a](https://arxiv.org/html/2610.09051#bib.bib51);[Shi et al\., 2025](https://arxiv.org/html/2610.09051#bib.bib58)\)\. These dynamic sparse attention patterns are either implemented by modifying full\-attention masks to zero out unnecessary query\-key dot\-products, or by selectively reading queries and keys from memory to compute a partial attention matrix\. Neither implementation reduces KV\-cache sizes\.

KV\-cache eviction:Methods such as H2O and SnapKV\([Zhang et al\., 2023](https://arxiv.org/html/2610.09051#bib.bib11);[Li et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib10)\)use custom heuristics during the prefill & decode stages of LLM inference to determine which key/value tokens to evict from the KV\-cache\.More recent methods refine the eviction criterion under the same fixed\-budget setting: CAOTE\([Goel et al\., 2025](https://arxiv.org/html/2610.09051#bib.bib61)\)combines attention scores with value vectors to minimize the attention\-output error induced by eviction, while KeyDiff\([Park et al\., 2026](https://arxiv.org/html/2610.09051#bib.bib57)\)avoids attention scores and instead retains geometrically distinctive keys based on key\-space similarity\.These strategies compress KV\-caches to a preset target \(e\.g\., evicting 90%\\%of tokens\), while UA’s adaptive pruning dynamically determines which key/value tokens are dispensable\. As a result, UA compresses much further at long sequence length, in addition to improving performance more generally\.

Data\-dependent forget gates:A range of gate formulations have been added to a variety of models, such as LSTMs\([Hochreiter and Schmidhuber, 1997](https://arxiv.org/html/2610.09051#bib.bib52)\)and transformers\([Vaswani et al\., 2017](https://arxiv.org/html/2610.09051#bib.bib1)\), for a variety of purposes\. But from a cache compression perspective, many approaches devolve into soft\-sliding\-window formulations\([Lin et al\., 2025a](https://arxiv.org/html/2610.09051#bib.bib2);[Lin et al\., 2025b](https://arxiv.org/html/2610.09051#bib.bib3);[Hochreiter and Schmidhuber, 1997](https://arxiv.org/html/2610.09051#bib.bib52)\)\. More expressive mechanisms such as the delta rule\([Yang et al\., 2024b](https://arxiv.org/html/2610.09051#bib.bib16);[Yang et al\., 2025a](https://arxiv.org/html/2610.09051#bib.bib9);[Yang et al\., 2025b](https://arxiv.org/html/2610.09051#bib.bib5)\)preclude the use of well\-established elements of attention, such as SoftMax or RoPE embeddings\([Su et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib28)\)\. UA combines multiple forget gates into an additive bias to the standard attention matrix, allowing us to leverage rich positional encoding via RoPE, targeted retrieval via SoftMax, and more sophisticated decay patterns than sliding\-window \(see Section[5\.3](https://arxiv.org/html/2610.09051#S5.SS3)\), all at the same time\.

## 3Background

This section introduces the notation and necessary background relevant to UA’s design\.

Softmax multi\-head self\-attention:Standard softmax multi\-head self\-attention takes an input sequence of tokens\(𝐪i\)i=1L,\(𝐤i\)i=1L,\(𝐯i\)i=1L∈ℝn×d\(\\mathbf\{q\}\_\{i\}\)\_\{i=1\}^\{L\},\(\\mathbf\{k\}\_\{i\}\)\_\{i=1\}^\{L\},\(\\mathbf\{v\}\_\{i\}\)\_\{i=1\}^\{L\}\\in\\mathbb\{R\}^\{n\\times d\}, whereLL,nn&ddare the sequence length, number of heads, and head dimension, respectively and computes:

Softmax​\(𝐐𝐊⊤\+𝐌\)​𝐕,\\text\{Softmax\}\(\\mathbf\{Q\}\\mathbf\{K\}^\{\\top\}\+\\mathbf\{M\}\)\\mathbf\{V\},\(1\)where𝐌∈ℝL×L\\mathbf\{M\}\\in\\mathbb\{R\}^\{L\\times L\}is the causal attention mask broadcasted across heads,𝐐,𝐊,𝐕∈ℝn×L×d\\mathbf\{Q\},\\mathbf\{K\},\\mathbf\{V\}\\in\\mathbb\{R\}^\{n\\times L\\times d\}are tensors formed by stacking𝐪i\\mathbf\{q\}\_\{i\},𝐤i\\mathbf\{k\}\_\{i\}&𝐯i\\mathbf\{v\}\_\{i\}, respectively\. During inference, new tokens need to attend to all previous tokens in𝐊\\mathbf\{K\}&𝐕\\mathbf\{V\}, resulting in per\-layer KV\-caches that consume considerable memory\.

Data dependent forget\-gates:To reduce large KV\-cache sizes, forget\-gates decay past key/value tokens, deleting sufficiently decayed entries from the KV\-cache\. In particular, FoX\([Lin et al\., 2025a](https://arxiv.org/html/2610.09051#bib.bib2);[Lin et al\., 2025b](https://arxiv.org/html/2610.09051#bib.bib3)\)expresses a forget\-gate as an additive bias to the attention matrix\. Using terminology from\([Lin et al\., 2025a](https://arxiv.org/html/2610.09051#bib.bib2)\), FoX computes a gate:ft=σ⁡\(𝐱t​𝐰\+b\)∈ℝf\_\{t\}=\\sigma\(\\mathbf\{x\}\_\{t\}\\mathbf\{w\}\+b\)\\in\\mathbb\{R\}for each token in the input sequence and computes:

Softmax​\(𝐐𝐊⊤\+𝐃\)​𝐕,\\text\{Softmax\}\(\\mathbf\{Q\}\\mathbf\{K\}^\{\\top\}\+\\mathbf\{D\}\)\\mathbf\{V\},\(2\)where𝐃i​j=∑l=j\+1ilog⁡\(fl\)\\mathbf\{D\}\_\{ij\}=\\sum\_\{l=j\+1\}^\{i\}\\log\(f\_\{l\}\)\. These decay\-values decrease monotonically from the main\-diagonal of the attention matrix, ensuring that cache entries with near\-zero gates stay near\-zero in perpetuity\. Sufficiently decayed keys can thus be safely deleted from the KV\-cache\. However, because KV\-cache entries decay in lockstep, this decay mechanism resembles a sliding\-window, an inexpressive KV\-cache eviction strategy\. UA generalizes and incorporates several such decay mechanisms, to achieve more sophisticated behavior\.

## 4Universal Attention

We propose Universal Attention \(UA\), a novel framework and architecture that incorporates context\-aware decay into softmax attention, inspired by modern linear attention methods and the structure of the attention matrix itself\. The result is a model that learns adaptive and dynamic sparsity patterns balancing efficiency with expressivity\.

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 applyingL2L\_\{2\}\-normalization to keys, and a low\-rank sigmoid gate to the attention output\. \(Right\) Construction of𝚲\\mathbf\{\\Lambda\}from three decay mechanisms: similarity\-basedfsimf\_\{\\text\{sim\}\}, query\-awarefqf\_\{\\text\{q\}\}, and key\-awarefkf\_\{\\text\{k\}\}\. These are combined via a geometric mean, transformed to log\-space, and accumulated causally to produce the decay matrix\.Figure[1](https://arxiv.org/html/2610.09051#S4.F1)depicts a UA layer, incorporating decay into softmax attention as a bias term\. However, rather than relying on a single decay mechanism, UA combines multiple context\-aware patterns via their geometric mean to form a unified decay rate value for each pair of positions\. We choose the geometric mean over arithmetic because it preserves the veto property: if any single mechanism deems that a token should not decay, the combined decay is driven to zero\.

Formally, consider the input to a UA layer as𝐗∈ℝL×dm​o​d​e​l\\mathbf\{X\}\\in\\mathbb\{R\}^\{L\\times d\_\{model\}\}, with queries, keys, and values𝐐,𝐊,𝐕∈ℝn×L×d\\mathbf\{Q\},\\mathbf\{K\},\\mathbf\{V\}\\in\\mathbb\{R\}^\{n\\times L\\times d\}, where keys are L2\-normalized\. Universal Attention computes:

Softmax​\(𝐐𝐊⊤\+𝚲\)​𝐕,𝚲i​j:=\{∑k=j\+1i𝐀k​ji≥j−∞i<j,\\displaystyle\\text\{Softmax\}\\left\(\\mathbf\{Q\}\\mathbf\{K\}^\{\\top\}\+\\mathbf\{\\Lambda\}\\right\)\\mathbf\{V\},\\quad\\mathbf\{\\Lambda\}\_\{ij\}:=\\begin\{cases\}\\displaystyle\\sum\_\{k=j\+1\}^\{i\}\\mathbf\{A\}\_\{kj\}&i\\geq j\\\\ \-\\infty&i<j,\\end\{cases\}\(3\)where𝚲i​j\\mathbf\{\\Lambda\}\_\{ij\}represents the accumulated log\-decay from keyjjto queryii: each columnjjof𝐀\\mathbf\{A\}tracks how a single key decays across subsequent positions, and𝚲i​j\\mathbf\{\\Lambda\}\_\{ij\}sums this decay causally, with−∞\-\\inftyenforcing causal masking through the softmax\. We operate in log\-space so that multiplicative decay composes additively\. The matrix𝐀\\mathbf\{A\}aggregates the decay functions inℱ\\mathcal\{F\}via their geometric mean:

𝐀=log⁡\(𝟏−\(∏f∈ℱf⁡\(𝐗\)\)1\|ℱ\|\),\\displaystyle\\mathbf\{A\}=\\log\\\!\\left\(\\mathbf\{1\}\-\\left\(\\prod\_\{f\\in\\mathcal\{F\}\}f\(\\mathbf\{X\}\)\\right\)^\{\\frac\{1\}\{\|\\mathcal\{F\}\|\}\}\\right\),\(4\)where∏\\proddenotes the elementwise product and each context\-aware decay mechanismf:ℝL×dm​o​d​e​l→\[0,1\]L×Lf\\colon\\mathbb\{R\}^\{L\\times d\_\{model\}\}\\to\[0,1\]^\{L\\times L\}produces a pairwise decay matrix\. Below, we introduce three representative mechanisms, linked to the first\- and second\-order statistics of stored keys\.

### 4\.1Decay Components

Our representative decay mechanism familyℱ:=\{f\(q\),f\(k\),f\(sim\)\}\\mathcal\{F\}:=\\left\\\{f^\{\(\\text\{q\}\)\},f^\{\(\\text\{k\}\)\},f^\{\(\\text\{sim\}\)\}\\right\\\}includes \(1\) a Mamba\-style query\-aware forget gatef\(q\)f^\{\(\\text\{q\}\)\}, \(2\) a novel key\-aware forget gatef\(k\)f^\{\(\\text\{k\}\)\}, and \(3\) a DeltaNet\-style similarity decayf\(sim\)f^\{\(\\text\{sim\}\)\}\. The first two mechanisms capture first\-order statistics aligned to the query and key axes of the attention matrix, respectively, while the third captures second\-order interaction effects between embeddings at current and past time steps\.

Query\-aware forget gatef\(q\)f^\{\(\\text\{q\}\)\}:Similar to Mamba2\([Dao and Gu, 2024](https://arxiv.org/html/2610.09051#bib.bib8)\)and FoX\([Lin et al\., 2025a](https://arxiv.org/html/2610.09051#bib.bib2)\), each query position is assigned a scalar forget\-gate computed as a sigmoid with a learned bias\. Since this decay value depends only on the query position, it is broadcast along the key axis to produce the corresponding pairwise decay matrixf\(q\)f^\{\(\\text\{q\}\)\}:

gi\(q\):=σ⁡\(𝐱i​𝐰\(q\)\+b\(q\)\),f\(q\)​\(𝐗\)i​j:=gi\(q\),\\displaystyle g^\{\(\\text\{q\}\)\}\_\{i\}:=\\sigma\\left\(\\mathbf\{x\}\_\{i\}\\mathbf\{w\}^\{\(\\text\{q\}\)\}\+b^\{\(\\text\{q\}\)\}\\right\),\\quad f^\{\(\\text\{q\}\)\}\(\\mathbf\{X\}\)\_\{ij\}:=g^\{\(\\text\{q\}\)\}\_\{i\},\(5\)where𝐰\(q\)\\mathbf\{w\}^\{\(\\text\{q\}\)\}andb\(q\)b^\{\(\\text\{q\}\)\}are learned parameters\. These gates accumulate causally, so that a key’s importance degrades with the number of intermediate positions visited\. This mechanism can be interpreted as each token advancing the entire cache by a variable\-size time step\.

Key\-aware forget gatef\(k\)f^\{\(\\text\{k\}\)\}:Symmetric tof\(q\)f^\{\(\\text\{q\}\)\}, eachkeyis assigned a scalar forget gate via the same sigmoid\+bias parameterization, which then broadcasts over thequerydimension:

gj\(k\):=σ⁡\(𝐱j​𝐰\(k\)\+b\(k\)\),f\(k\)​\(𝐗\)i​j:=gj\(k\)\.\\displaystyle g^\{\(\\text\{k\}\)\}\_\{j\}:=\\sigma\\left\(\\mathbf\{x\}\_\{j\}\\mathbf\{w\}^\{\(\\text\{k\}\)\}\+b^\{\(\\text\{k\}\)\}\\right\),\\quad f^\{\(\\text\{k\}\)\}\(\\mathbf\{X\}\)\_\{ij\}:=g^\{\(\\text\{k\}\)\}\_\{j\}\.\(6\)Whilef\(q\)f^\{\(\\text\{q\}\)\}represents how much history to forget, our novelf\(k\)f^\{\(\\text\{k\}\)\}lets each key control its own intrinsic decay rate, independent of future updates\. This allows the model to designate tokens that persist through time, or that self\-prune quickly, on a case\-by\-case basis\.

Similarity decayf\(sim\)f^\{\(\\text\{sim\}\)\}:To capture higher\-order interaction effects between stored keys, we derive a pairwise decay mechanism similar to DeltaNet\([Yang et al\., 2024b](https://arxiv.org/html/2610.09051#bib.bib16)\)\. The DeltaNet update recurrence takes the form:

𝐒t=𝐒t−1​\(𝐈−βt​𝐤t​𝐤t⊤\)\+βt​𝐯t​𝐤t⊤\\mathbf\{S\}\_\{t\}=\\mathbf\{S\}\_\{t\-1\}\\big\(\\mathbf\{I\}\-\\beta\_\{t\}\\mathbf\{k\}\_\{t\}\\mathbf\{k\}\_\{t\}^\{\\top\}\\big\)\+\\beta\_\{t\}\\mathbf\{v\}\_\{t\}\\mathbf\{k\}\_\{t\}^\{\\top\}\(7\)
The state𝐒t\\mathbf\{S\}\_\{t\}accumulates key/value outer products with step sizeβt\\beta\_\{t\}, and so we can unroll the prior state𝐒t−1:=∑i=0t−1𝐯i′​𝐤i′⁣⊤\\mathbf\{S\}\_\{t\-1\}:=\\sum\_\{i=0\}^\{t\-1\}\\mathbf\{v\}\_\{i\}^\{\\prime\}\\mathbf\{k\}\_\{i\}^\{\\prime\\top\}, wherevi′,ki′\\textbf\{v\}\_\{i\}^\{\\prime\},\\textbf\{k\}\_\{i\}^\{\\prime\}representvi,ki\\textbf\{v\}\_\{i\},\\textbf\{k\}\_\{i\}with changes absorbed from subsequent state updates\. DeltaNet setski′:=ki\\textbf\{k\}\_\{i\}^\{\\prime\}:=\\textbf\{k\}\_\{i\}and absorbs all changes intovi′\\textbf\{v\}\_\{i\}^\{\\prime\}, but an equivalent formulation setsvi′:=vi\\textbf\{v\}\_\{i\}^\{\\prime\}:=\\textbf\{v\}\_\{i\}and absorbs all changes intoki′\\textbf\{k\}\_\{i\}^\{\\prime\}:

𝐒t−1​\(𝐈−βt​𝐤t​𝐤t⊤\)=∑i=0t−1𝐯i​ki′⁣⊤​\(𝐈−βt​𝐤t​𝐤t⊤\)=∑i=0t−1𝐯i​\(ki′−βt​ki′​kt⊤​kt\)⊤\\mathbf\{S\}\_\{t\-1\}\\big\(\\mathbf\{I\}\-\\beta\_\{t\}\\mathbf\{k\}\_\{t\}\\mathbf\{k\}\_\{t\}^\{\\top\}\\big\)=\\sum\_\{i=0\}^\{t\-1\}\\mathbf\{v\}\_\{i\}\\textbf\{k\}\_\{i\}^\{\\prime\\top\}\\big\(\\mathbf\{I\}\-\\beta\_\{t\}\\mathbf\{k\}\_\{t\}\\mathbf\{k\}\_\{t\}^\{\\top\}\\big\)=\\sum\_\{i=0\}^\{t\-1\}\\mathbf\{v\}\_\{i\}\\big\(\\textbf\{k\}\_\{i\}^\{\\prime\}\-\\beta\_\{t\}\\textbf\{k\}\_\{i\}^\{\\prime\}\\textbf\{k\}\_\{t\}^\{\\top\}\\textbf\{k\}\_\{t\}\\big\)^\{\\top\}\(8\)
ForL2L\_\{2\}\-normalizedk, the right\-hand term inside the summation represents the projection ofki′\\textbf\{k\}\_\{i\}^\{\\prime\}ontokt\\textbf\{k\}\_\{t\}, scaled byβt∈\[0,1\]\\beta\_\{t\}\\in\[0,1\]\. Thus we can interpret the DeltaNet update as apartial orthogonalizationof all prior keys against the newest key entering the cache\. Over successive rounds of orthogonalization, the magnitude ofki′\\textbf\{k\}\_\{i\}^\{\\prime\}decays to zero, slowly when it is not aligned to each newkt\\textbf\{k\}\_\{t\}, and quickly when it is\.

Unfortunately, this decay mechanism is not UA\-compatible as\-is, because dynamically shiftingkboth 1\) interferes with RoPE, and 2\) requires overwriting the entire key cache at each time step, which is computationally prohibitive\. We propose a simple modification that transforms DeltaNet decay into an attention\-bias\-friendly form:

𝐤i′−βt​𝐤i′⁣⊤​𝐤t​𝐤t⇒𝐤i′−βt​𝐤i′⁣⊤​𝐤t​𝐤i′=\(1−βt​𝐤i′⁣⊤​𝐤t\)​ki′\\displaystyle\\mathbf\{k\}\_\{i\}^\{\\prime\}\-\\beta\_\{t\}\\mathbf\{k\}\_\{i\}^\{\\prime\\top\}\\mathbf\{k\}\_\{t\}\\mathbf\{k\}\_\{t\}\\quad\\Rightarrow\\quad\\mathbf\{k\}\_\{i\}^\{\\prime\}\-\\beta\_\{t\}\\mathbf\{k\}\_\{i\}^\{\\prime\\top\}\\mathbf\{k\}\_\{t\}\{\\color\[rgb\]\{0,0,1\}\\mathbf\{k\}\_\{i\}^\{\\prime\}\}\\ =\\ \(1\-\\beta\_\{t\}\\mathbf\{k\}\_\{i\}^\{\\prime\\top\}\\mathbf\{k\}\_\{t\}\)\\textbf\{k\}\_\{i\}^\{\\prime\}\(9\)By replacing𝐤t\\mathbf\{k\}\_\{t\}with𝐤i′\\mathbf\{k\}\_\{i\}^\{\\prime\}in the rightmost term, the update reduces to a multiplicative decay factor determined by key similarity\. The decay factor is thusβt​𝐤i′⁣⊤​𝐤t\\beta\_\{t\}\\mathbf\{k\}\_\{i\}^\{\\prime\\top\}\\mathbf\{k\}\_\{t\}, which serves as the basis forf\(sim\)f^\{\(\\text\{sim\}\)\}\. We letf\(query\)f^\{\(\\text\{query\}\)\}absorbβt\\beta\_\{t\}and definef\(sim\)f^\{\(\\text\{sim\}\)\}based solely on the dot product:

f\(sim\)​\(𝐗\)i​j:=ReLU​\(𝐤i⊤​𝐤j\)2\\displaystyle f^\{\(\\text\{sim\}\)\}\(\\mathbf\{X\}\)\_\{ij\}:=\\mathrm\{ReLU\}\(\\mathbf\{k\}\_\{i\}^\{\\top\}\\mathbf\{k\}\_\{j\}\)^\{2\}\(10\)Since keys areL2L\_\{2\}\-normalized,𝐤i⊤​𝐤j∈\[−1,1\]\\mathbf\{k\}\_\{i\}^\{\\top\}\\mathbf\{k\}\_\{j\}\\in\[\-1,1\], and squared ReLU ensures a gate value in\[0,1\]\[0,1\]as desired\. The result is a conservative decay mechanism biased heavily toward zero, but still capable of eliminating a key/value pair when a new entry is perfectly aligned\.

While we focus on these three representative mechanisms, the framework is extensible: additional decay functions can be incorporated intoℱ\\mathcal\{F\}as new advances emerge in linear attention — hence theuniversalnature of our method\.

### 4\.2Layer Design and Algorithm

Pluggingℱ=\{f\(q\),f\(k\),f\(sim\)\}\\mathcal\{F\}=\\\{f^\{\(\\text\{q\}\)\},f^\{\(\\text\{k\}\)\},f^\{\(\\text\{sim\}\)\}\\\}into the general framework[4](https://arxiv.org/html/2610.09051#S4.E4), the decay matrix becomes:

𝐀i​j=log⁡\(1−\(gi\(q\)⋅gj\(k\)⋅ReLU​\(𝐤i⊤​𝐤j\)2\)13\)\.\\displaystyle\\mathbf\{A\}\_\{ij\}=\\log\\\!\\left\(1\-\\Big\(g^\{\(\\text\{q\}\)\}\_\{i\}\\cdot g^\{\(\\text\{k\}\)\}\_\{j\}\\cdot\\mathrm\{ReLU\}\(\\mathbf\{k\}\_\{i\}^\{\\top\}\\mathbf\{k\}\_\{j\}\)^\{2\}\\Big\)^\{\\frac\{1\}\{3\}\}\\right\)\.\(11\)The accumulated log\-decay𝚲i​j=∑k=j\+1i𝐀k​j\\mathbf\{\\Lambda\}\_\{ij\}=\\sum\_\{k=j\+1\}^\{i\}\\mathbf\{A\}\_\{kj\}can then be added to the pre\-softmax logits of our attention layer, depicted in Fig\.[1](https://arxiv.org/html/2610.09051#S4.F1)\. We also take architectural inspiration from linear attention models, which often apply additional channel\-wise sigmoid gates to the intermediate representation\([Dao and Gu, 2024](https://arxiv.org/html/2610.09051#bib.bib8);[Yang et al\., 2025a](https://arxiv.org/html/2610.09051#bib.bib9);[Poli et al\., 2023](https://arxiv.org/html/2610.09051#bib.bib54)\)\. We adopt analogous gates for attention layers\([Qiu et al\., 2025](https://arxiv.org/html/2610.09051#bib.bib30)\), forming unique per\-channel gates via low\-rank projection\([Zhang et al\., 2025b](https://arxiv.org/html/2610.09051#bib.bib55)\)\. With rank equal to head width, these gates add a single attention head’s worth of extra parameters per layer, a negligible quantity with outsize impact on performance\. By reducing reliance on sink tokens\([Qiu et al\., 2025](https://arxiv.org/html/2610.09051#bib.bib30)\), channel gating allows the decay mechanisms greater freedom to learn expressive patterns\. We ablate this component and compare to both gated and vanilla attention baselines in Section[5\.1](https://arxiv.org/html/2610.09051#S5.SS1)\.

Inference pruning:Beyond shaping attention, UA’s learned decay also provides a natural signal for inference\-time KV\-cache pruning\. The log\-decay matrixΛ\\Lambdais implemented as a bias on attention logits, but it can also be interpreted, through elementwise exponentiation, as a set of token\-wise masks applied to the logits post\-softmax:

Softmax​\(𝐐𝐊⊤\+𝚲\)i​j=e\(𝐐i​𝐊j⊤\+𝚲i​j\)∑k=0je\(𝐐i​𝐊k⊤\+𝚲i​k\)=e𝐐i​𝐊j⊤​exp⁡\(𝚲\)i​j∑k=0je𝐐i​𝐊k⊤​exp⁡\(𝚲\)i​k\\text\{Softmax\}\\left\(\\mathbf\{Q\}\\mathbf\{K\}^\{\\top\}\+\\mathbf\{\\Lambda\}\\right\)\_\{ij\}=\\frac\{\\text\{e\}^\{\(\\mathbf\{Q\}\_\{i\}\\mathbf\{K\}\_\{j\}^\{\\top\}\+\\mathbf\{\\Lambda\}\_\{ij\}\)\}\}\{\\sum\_\{k=0\}^\{j\}\\text\{e\}^\{\(\\mathbf\{Q\}\_\{i\}\\mathbf\{K\}\_\{k\}^\{\\top\}\+\\mathbf\{\\Lambda\}\_\{ik\}\)\}\}=\\frac\{\\text\{e\}^\{\\mathbf\{Q\}\_\{i\}\\mathbf\{K\}\_\{j\}^\{\\top\}\}\\exp\(\\mathbf\{\\Lambda\}\)\_\{ij\}\}\{\\sum\_\{k=0\}^\{j\}\\text\{e\}^\{\\mathbf\{Q\}\_\{i\}\\mathbf\{K\}\_\{k\}^\{\\top\}\}\\exp\(\\mathbf\{\\Lambda\}\)\_\{ik\}\}\(12\)Mask values inexp⁡\(𝚲\)\\exp\(\\mathbf\{\\Lambda\}\)start at 1 and decay to 0 asymptotically, providing a natural pruning criterion: during inference, maintain a scalar biasλi\\lambda\_\{i\}for each key/value pair representing the corresponding entry of𝚲\\mathbf\{\\Lambda\}, continuously updateλi\\lambda\_\{i\}using Eq\.[11](https://arxiv.org/html/2610.09051#S4.E11), and evict any key\-value pair whose effective maskeλie^\{\\lambda\_\{i\}\}falls below a chosen thresholdp∈\[0,1\]p\\in\[0,1\]\. Decay bias updates inAare negative by construction, soλi\\lambda\_\{i\}always decreases, thus a token withλi<log⁡p\\lambda\_\{i\}<\\log pwill not return to importance later\. Although simple, this rule fully exploits the learned decay and lets the model itself decide, on a head\-by\-head basis, which and how many tokens should be retained\. More sophisticated pruning algorithms are left to future work\.

Implementation:We codesign UA to behardware\-aligned– interoperating with GQA\([Ainslie et al\., 2023](https://arxiv.org/html/2610.09051#bib.bib15)\)and FlashAttention\([Dao et al\., 2022](https://arxiv.org/html/2610.09051#bib.bib33);[Dao, 2023](https://arxiv.org/html/2610.09051#bib.bib34);[Shah et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib35);[Zadouri et al\., 2026](https://arxiv.org/html/2610.09051#bib.bib36)\), as current kernels do not support both GQA and float\-valued biases\. We implement UA as two kernels, separately computing \(1\) the decay\-mask, and \(2\) the attention\-output\. Additionally, we apply Activation\-Checkpointing\([Chen et al\., 2016](https://arxiv.org/html/2610.09051#bib.bib53)\)to mitigate storing theO\(L2\)sized decay\-maskper\-layerin device memory \- which would otherwise be required to compute gradients in the backwards pass\. Our custom implementation boosts end\-to\-end training throughput by 1\.13×\\timescompared to PyTorch\-2\. We explain our kernel implementations and end\-to\-end pipeline in detail in Appendix[C](https://arxiv.org/html/2610.09051#A3)\.

## 5Evaluation

Setup:We evaluate Universal Attention in a 1B\-parameter transformer following the Llama3 architecture\([Grattafiori et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib31)\), which uses RoPE\([Su et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib28)\)and Grouped Query Attention\([Ainslie et al\., 2023](https://arxiv.org/html/2610.09051#bib.bib15)\)\. Models are trained for 50B tokens on Dolma v1\.7 data\([Soldaini et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib42)\), reweighted per[Chu et al\. \(2024\)](https://arxiv.org/html/2610.09051#bib.bib37)\. Model and training details and hyperparameters are provided in Appendix[A](https://arxiv.org/html/2610.09051#A1)\. For long\-context inputs, we apply YaRN\([Peng et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib29)\)as a tuning\-free adjustment to RoPE\. We report results at two pruning thresholdsp=\(\.05,\.15\)p=\(\.05,\.15\), to illustrate performance tradeoffs at conservative/aggressive levels of compression\. For UA models, we additionally prune in chunks of 128 tokens during prefill, rather than running the whole input and only pruning afterwards\.

Baselines:We compare against both linear\- and softmax\-attention baselines: Mamba2\([Dao and Gu, 2024](https://arxiv.org/html/2610.09051#bib.bib8)\), Mamba2\-hybrid\([NVIDIA et al\., 2025](https://arxiv.org/html/2610.09051#bib.bib32)\), Gated DeltaNet \(GDN\) and its hybrid\([Yang et al\., 2025a](https://arxiv.org/html/2610.09051#bib.bib9)\), as well as standard Llama3 and gated\-attention Llama3\-G\. All baselines are trained identically to UA\.

For context extension, we apply YaRN\([Peng et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib29)\)only to the Llama3\-based models, where RoPE rescaling is a standard extension recipe\. For Mamba2\-hybrid, we follow Nemotron\-H and keep its native RoPE\-free design in the attention layers\.For GDN and its hybrid, we retain the base RoPE setup in the main comparison to avoid introducing an additional positional intervention\. Appendix[I](https://arxiv.org/html/2610.09051#A9)further evaluates GDN\-H with YaRN at 8K and 16K, where it improves long\-context performance but does not change the overall comparison with UA\.

To compare KV cache compression efficiency, we also apply H2O\([Zhang et al\., 2023](https://arxiv.org/html/2610.09051#bib.bib11)\), SnapKV\([Li et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib10)\), and their pyramid variants\([Yang et al\., 2024a](https://arxiv.org/html/2610.09051#bib.bib12)\)to the vanilla Llama3 model\.We also include KeyDiff\([Park et al\., 2026](https://arxiv.org/html/2610.09051#bib.bib57)\)as an additional baseline using the publicly available kvpress\([Devoto et al\., 2025](https://arxiv.org/html/2610.09051#bib.bib62)\)implementation\.To match the compression rate of UA, we aggressively set these methods to keep only 10% of the KV\-cache\. This strong suite of baselines thus encompasses a range of complementary recent approaches to reducing memory overhead, as well as two full\-size oracles\. Further details are provided in Appendix[A](https://arxiv.org/html/2610.09051#A1)\.

Benchmarks:We evaluate the models’ performance on two sets of benchmarks: information retrieval, including LongBench\-E\([Bai et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib19)\)and RULER\([Hsieh et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib18)\), and common\-sense reasoning tasks, following[Yang et al\. \(2025a\)](https://arxiv.org/html/2610.09051#bib.bib9)\. For each task, we track the size of the stored states and KV\-caches \(“effective cache size”\) using their actual sizes in bytes\.

### 5\.1Downstream Performance

Information retrieval:Table[1](https://arxiv.org/html/2610.09051#S5.T1)reports results on inputs within the training sequence length, together with effective cache size, for RULER \(4K\) and LongBench\-E \(0\-4K\)\. On RULER, UA delivers significant improvement, even over unpruned Llama3 oracle models, claiming the highest and second\-highest average scores at its two pruning thresholds\. This improvement is mostly fueled by gains in tracking and retrieval tasks, though performance on aggregation suffers\. Aggregation is a common bottleneck for KV\-cache eviction methods, since it requires combining information from broadly distributed context\. UA shares this weakness, but unlike SnapKV, H2O,and KeyDiff, its substantial degradation is concentrated in aggregation\. On LongBench\-E, the best average is obtained by Mamba2\-H, yet UA remains highly competitive under a much tighter cache budget:p=0\.15p\{=\}0\.15uses under half the memory of Mamba2\-H, and exceeds both Llama3 and Llama3\-G scores with only a fraction of their cache size\. This favorable comparison to full\-cache Llama3 also persists at 3B model scale, with results in Appendix[F](https://arxiv.org/html/2610.09051#A6)\. Overall, UA offers superior memory\-performance tradeoffs\.

Table 1:RULER \(4k sequence length\) and LongBench\-E \(0\-4k sequence length split\)\.RULER \(4k\)LongBench\-E \(0\-4k\)ModelRetrievalTrackingAggregationQAAvg\.Cache \(MB\)Avg\.Cache \(MB\)Sub\-quadratic modelsMamba231\.00\.818\.120\.025\.02012\.720\.0Mamba2\-H57\.23\.220\.929\.743\.235\.518\.931\.7GDN37\.72\.618\.118\.429\.03214\.932\.0GDN\-H59\.94\.419\.624\.944\.146\.015\.742\.2Quadratic modelsLlama354\.28\.413\.630\.440\.8143\.716\.4113\.3\+ SnapKV21\.713\.47\.525\.619\.514\.412\.911\.3\+ Pyramid SnapKV21\.220\.36\.225\.119\.414\.412\.911\.3\+ H2O4\.20\.210\.719\.97\.314\.413\.011\.3\+ Pyramid H2O4\.30\.29\.720\.07\.214\.412\.611\.3\+ KeyDiff \(K=10%​LK=10\\%L\)33\.824\.710\.224\.328\.014\.415\.211\.3Llama3\-G53\.716\.321\.428\.341\.9143\.717\.9113\.3Universal Attentionp=0\.0568\.436\.45\.729\.550\.321\.417\.920\.2p=0\.1567\.330\.66\.929\.149\.312\.918\.014\.4

Common\-sense reasoning:Table[2](https://arxiv.org/html/2610.09051#S5.T2)presents language modeling and zero\-shot reasoning results at a context length of 4k\. Due to benchmark setup, H2O, SnapKV, and KeyDiff only prune their cachesafterproducing output logits from full\-cache attention, and so are omitted from comparison\. UA achieves competitive\-to\-superior perplexity and reasoning scores, with notable gains in LAMBADA accuracy, PIQA, and ARC\-easy, despite aggressive pruning to only 38\.7 KV\-cache tokens on average\.While the absolute zero\-shot reasoning gains are modest, these results provide strong evidence that UA preserves model quality at substantially lower cache cost: UA matches or slightly exceeds Llama3 in average reasoning accuracy \(50\.74/50\.88 vs\. 50\.28\) while reducing cache from 2\.6 MB to 1\.9/1\.5 MB, and uses roughly 10–13×\\timesless cache than Mamba2 and 17–21×\\timesless than GDN while maintaining comparable or better accuracy\.Note that due to short input sequence length, the Mamba2 and GDN linear attention architectures consume more memory than the corresponding full attention models, and UA compression levels are less dramatic in comparison, though easily still the lowest effective cache size overall\.This demonstrates that the learned decay preserves short\-context quality and improves the memory–quality tradeoff, while potentially acting as an implicit regularizer for perplexity and downstream task performance\.

Table 2:Language modeling perplexity and zero\-shot common\-sense reasoning accuracy\.
### 5\.2Context Length Extension

To better understand how the memory\-performance trade\-off evolves with context length, Figure[2](https://arxiv.org/html/2610.09051#S5.F2)plots average score against effective cache size for RULER \(left\) and LongBench\-E \(right\)\.Each curve traces a model across increasing sequence lengths from 1K to 16K, so both cache size and task difficulty vary along the curve\. Consequently, a lower score at a larger cache size does not imply that additional cache is harmful, but simply reflects evaluation at a longer and more challenging context\. Better configurations at a given context length lie toward the upper\-left,and UA consistently traces the strongest frontier as sequence length increases\.

The two UA variants expose a clear trade\-off within the method itself:p=0\.05p\{=\}0\.05preserves higher scores, whilep=0\.15p\{=\}0\.15achieves more aggressive compression\. Unlike full\-attention baselines such as Llama3 and Llama3\-G, whose cache grows linearly with context length and quickly exceeds the plotting range, UA keeps cache usage under 25 MB even at the longest contexts shown\. Relative to Llama3, UA withp=0\.15p\{=\}0\.15achieves4\.5×4\.5\\times–45\.3×45\.3\\timescache reduction on RULER and7\.9×7\.9\\times–26\.9×26\.9\\timeson LongBench\-E\. The gap widens at longer contexts because UA’s cache plateaus, while Llama3 continues to grow\. By contrast, SnapKV and H2O reduce cache at a larger accuracy cost, while hybrid baselines grow the cache with diminishing returns\. The full breakdown of benchmark scores by task and sequence length is provided in Appendix[I](https://arxiv.org/html/2610.09051#A9)\.

![Refer to caption](https://arxiv.org/html/2610.09051v1/figs/RULER_LongBench.png)Figure 2:Score versus effective cache size trade\-off on RULER \(left\) and LongBench\-E \(right\)\.Each line traces a model across increasing sequence lengths from 1K to 16K; thus, moving along a curve changes both cache size and task difficulty\.The gold star \(★\\bigstar\) denotes the ideal operating point of high score and low overhead\.
### 5\.3Empirical Analysis

We examine UA’s actual learned patterns, uncovering several interesting and spontaneous behaviors\. Fig\.[3](https://arxiv.org/html/2610.09051#S5.F3)\(left\) shows an example learned decay pattern, expressed as the set of accumulated masksexp⁡\(𝚲\)\\exp\{\(\\mathbf\{\\Lambda\}\)\}\. Attention logits are not pictured\. We can isolate contributions from each of our three decay mechanisms: vertical stripes represent tokens with low key\-linked decayf\(k\)f^\{\(\\text\{k\}\)\}, horizontal lines represent selective filtering through similarity\-based decayf\(sim\)f^\{\(\\text\{sim\}\)\}, and progressive query\-linked decayf\(q\)f^\{\(\\text\{q\}\)\}causes entries to fade with distance from the diagonal\. UA’s mixture of mechanisms appears to be expressive\.

![Refer to caption](https://arxiv.org/html/2610.09051v1/figs/ua_empiricals.png)Figure 3:\(Left\) One example head’s learned pattern of accumulated forget gatese​x​p​\(𝚲\)exp\(\\mathbf\{\\Lambda\}\), exhibiting behaviors from multiple decay mechanisms\. \(Middle\) Progressive UA cache size over two different input sequences\. Pruning is highly dynamic and input\-adaptive, with spontaneously learned growth\-and\-reset behaviors\. \(Right\) Evolution of effective cache size over the course of training, with an extreme initial dip and gradual recovery\.Dynamic pruning is illustrated in Fig\.[3](https://arxiv.org/html/2610.09051#S5.F3)\(middle\), plotting average cache size across heads as a function of input length\. While cache length is small relative to input length, the exact size varies, both over time and between different inputs\. The model learns to be naturally parsimonious, accumulating information when it is needed, and pruning aggressively once it is not\. This ability to discard unnecessary information could explain the consistent improvement of UA over the oracle Llama3 baselines in Section[5\.1](https://arxiv.org/html/2610.09051#S5.SS1)\.

This notion of learned frugality is reinforced in Fig[3](https://arxiv.org/html/2610.09051#S5.F3)\(right\), which plots the effective KV\-cache size during training\. Here we measure effective size as the percentage of cache entries aboveppat the final token position / in the final row of the attention matrix\. Caches immediately sparsify at the beginning, then the model gradually re\-incorporates past information as its language modeling capabilities improve\. Eventually, cache sizes stabilize and even fall gently over the last half of training, with the model naturally becoming more efficient\. This suggests that longer training regimes could produce further compression\.

### 5\.4Ablations

UA decay mechanisms:We ablate UA’s decay mechanisms using models trained on 10B tokens instead of 50B, and report averaged RULERS1,S2, &S3scores across sequence lengths\. Table[3](https://arxiv.org/html/2610.09051#S5.T3)\(left\) shows that all components contribute to performance, supporting our “universal” approach of combining multiple decay mechanisms\. Interestingly, removing the most well\-studied, Mamba\-style query\-aware decayf\(q\)f^\{\(\\text\{q\}\)\}causes the smallest degradation, while the other decay mechanisms provide substantial complementary benefits\. Appendix[E](https://arxiv.org/html/2610.09051#A5)provides additional architectural ablations, including each decay mechanism in isolation, Qwen3\-style normalized queries and keys, and comparisons with FoX and sliding\-window attention\.

Pruning threshold:We ablate different pruning thresholdsppon our 50B\-token model in Table[3](https://arxiv.org/html/2610.09051#S5.T3)\(right\), evaluating LongBench\-E with inputs truncated to 16K and perplexity on held\-out natural\-language validation data matching the training distribution\. We report both the average compression ratio across KV heads and the maximum compression achieved on a single sequence\. UA degrades gracefully as compression increases: atp=0\.15p\{=\}0\.15, it achieves21\.4×21\.4\\timesaverage compression on LongBench\-E and14\.9×14\.9\\timeson held\-out validation data, while changing the LongBench\-E score and validation perplexity by less than one point each\. Because the learned decay patterns vary across models and the compression–performance trade\-off changes smoothly rather than exhibiting a natural cutoff, there is no single universally optimal threshold\. We therefore usep=0\.05p\{=\}0\.05andp=0\.15p\{=\}0\.15as representative operating points for gentler and more aggressive pruning, respectively, while allowing the threshold to be selected according to the desired memory–quality trade\-off\.

Table 3:\(Left\): Impact of UA components on downstream retrieval tasks\.\(Right\): Impact of pruning threshold, with performance/compression tradeoff\.

## 6Limitations

Aggregation tasks:Universal Attention is comparatively weaker on aggregation tasks, which often require combining weak evidence distributed across many positions rather than preserving a sparse set of individually salient tokens\. This contrasts with UA’s strengths on retrieval, tracking, and QA, suggesting a trade\-off in the types of long\-range information emphasized by different architectures\. Our hybrid results in Appendix[D](https://arxiv.org/html/2610.09051#A4)further show that this weakness can be mitigated by combining UA with complementary memory mechanisms, making UA suitable as a cache\-efficient attention component within broader long\-context architectures when stronger aggregation capability is required\.

Additional computation:UA introduces additional computation for constructing its learned decay and pruning masks\. Although this enables adaptive, content\-dependent retention, reducing the associated overhead remains an important systems optimization target\.

Training from scratch:We evaluate UA as a novel architecture, but given that it generalizes standard softmax attention, a natural question is whether existing models can be adapted to UA, rather than having to pretrain it completely from scratch\. Early results in Appendix[G](https://arxiv.org/html/2610.09051#A7)suggest that this is possible, but a significant gap remains compared to fully\-pretrained UA\. Closing this gap would greatly increase the practical feasibility of UA’s deployment\.

Inference optimization:While preliminary inference experiments in Appendix[H](https://arxiv.org/html/2610.09051#A8)demonstrate that UA’s reduced cache can translate into lower decode latency, the current implementation does not include a fully fused, inference\-optimized dynamic cache manager\. Efficiently exploiting dynamically sparse KV caches requires specialized sparse kernels and systems support\([Ye et al\., 2025](https://arxiv.org/html/2610.09051#bib.bib56)\), and developing such an implementation remains an important direction for future work\.

## 7Conclusion

We present Universal Attention, a technique that generalizes existing decay mechanisms, introduces a novel decay mechanism, and combines these with RoPE and softmax attention into a single, unified layer\. UA’s expressive decay mechanism provides a natural adaptive pruning criterion, yielding superior memory\-performance tradeoffs, improving downstream performance with lower memory overhead than evenO⁡\(1\)O\(1\)linear attention models\. Our results demonstrate that trainable self\-pruning architectures provide a practical avenue to mitigate large memory overhead during inference\.

## Acknowledgments

We sincerely appreciate the anonymous reviewers\. Their insightful feedback helps significantly improve the quality of the paper\. This research was supported by the National Science Foundation \(NSF\) under Grant No\. 2441601\. The work utilized the Delta and DeltaAI system at the National Center for Supercomputing Applications \(NCSA\) and Jetstream2 at Indiana University through allocation CIS240055 from the Advanced Cyberinfrastructure Coordination Ecosystem: Services & Support \(ACCESS\) program, which is supported by National Science Foundation grants \#2138259, \#2138286, \#2138307, \#2137603, and \#2138296\. The Delta advanced computing resource is a collaborative effort between the University of Illinois Urbana\-Champaign and NCSA, supported by the NSF \(award OAC 2005572\) and the State of Illinois\. UIUC SSAIL Lab is supported by research funding and gift from IBM, Amazon, Google, and AMD\.

## References

- Ainslieet al\.\(2023\)J\. Ainslie, J\. Lee\-Thorp, M\. de Jong, Y\. Zemlyanskiy, F\. Lebron, and S\. SanghaiGQA: training generalized multi\-query transformer models from multi\-head checkpoints\.InProceedings of the 2023 Conference on Empirical Methods in Natural Language Processing,H\. Bouamor, J\. Pino, and K\. Bali \(Eds\.\),Singapore,pp\.4895–4901\.External Links:[Link](https://aclanthology.org/2023.emnlp-main.298/),[Document](https://dx.doi.org/10.18653/v1/2023.emnlp-main.298)Cited by:[Appendix B](https://arxiv.org/html/2610.09051#A2.SS0.SSS0.Px3),[§4\.2](https://arxiv.org/html/2610.09051#S4.SS2.p3.1),[§5](https://arxiv.org/html/2610.09051#S5.p1.1)\.
- Anthropic \(2026\)AnthropicModels overview\.Note:Claude API Docs\. Accessed: 2026\-03\-31External Links:[Link](https://docs.anthropic.com/en/docs/about-claude/models)Cited by:[§1](https://arxiv.org/html/2610.09051#S1.p1.1)\.
- Baiet al\.\(2024\)Y\. Bai, X\. Lv, J\. Zhang, H\. Lyu, J\. Tang, Z\. Huang, Z\. Du, X\. Liu, A\. Zeng, L\. Hou, Y\. Dong, J\. Tang, and J\. LiLongBench: a bilingual, multitask benchmark for long context understanding\.InProceedings of the 62nd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),Bangkok, Thailand,pp\.3119–3137\.External Links:[Link](https://aclanthology.org/2024.acl-long.172),[Document](https://dx.doi.org/10.18653/v1/2024.acl-long.172)Cited by:[Appendix A](https://arxiv.org/html/2610.09051#A1.SS0.SSS0.Px3.p1.1),[§5](https://arxiv.org/html/2610.09051#S5.p5.1)\.
- Beltagyet al\.\(2020\)I\. Beltagy, M\. E\. Peters, and A\. CohanLongformer: the long\-document transformer\.External Links:2004\.05150,[Link](https://arxiv.org/abs/2004.05150)Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p4.1)\.
- Ben\-Kishet al\.\(2025\)A\. Ben\-Kish, I\. Zimerman, S\. Abu\-Hussein, N\. Cohen, A\. Globerson, L\. Wolf, and R\. GiryesDeciMamba: exploring the length extrapolation potential of mamba\.InThe Thirteenth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=iWSl5Zyjjw)Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p2.1)\.
- Bisket al\.\(2020\)Y\. Bisk, R\. Zellers, R\. L\. Bras, J\. Gao, and Y\. ChoiPIQA: reasoning about physical commonsense in natural language\.InThirty\-Fourth AAAI Conference on Artificial Intelligence,Cited by:[Appendix A](https://arxiv.org/html/2610.09051#A1.SS0.SSS0.Px3.p1.1)\.
- Chenet al\.\(2016\)T\. Chen, B\. Xu, C\. Zhang, and C\. GuestrinTraining deep nets with sublinear memory cost\.ArXivabs/1604\.06174\.External Links:[Link](https://api.semanticscholar.org/CorpusID:15865278)Cited by:[§4\.2](https://arxiv.org/html/2610.09051#S4.SS2.p3.1)\.
- Childet al\.\(2019\)R\. Child, S\. Gray, A\. Radford, and I\. SutskeverGenerating Long Sequences with Sparse Transformers\.arXiv e\-prints,pp\.arXiv:1904\.10509\.External Links:[Document](https://dx.doi.org/10.48550/arXiv.1904.10509),1904\.10509Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p4.1)\.
- Chuet al\.\(2024\)L\. Chu, D\. Kumari, T\. Dao, A\. Gu, R\. Ganti, D\. Agrawal, M\. Srivatsa, D\. Wertheimer, Y\. C\. F\. Lim, A\. Viros, N\. Gonzalez, T\. HoangTrong, O\. Arviv, Y\. Perlitz, M\. Shmueli, H\. Shen, M\. Zhang, G\. Goodhart, N\. Wang, N\. Hill, J\. Rosenkranz, C\. Liu, A\. Hoque, C\. Yang, S\. Sharma, A\. Uong, J\. Gala, S\. Zawad, and R\. GordonBamba: inference\-efficient hybrid mamba2 model\.Note:[https://huggingface\.co/blog/bamba\#bamba\-inference\-efficient\-hybrid\-mamba2\-model\-%F0%9F%90%8D](https://huggingface.co/blog/bamba#bamba-inference-efficient-hybrid-mamba2-model-%F0%9F%90%8D)Accessed: 2025\-05\-06Cited by:[Appendix A](https://arxiv.org/html/2610.09051#A1.SS0.SSS0.Px1.p2.1),[§5](https://arxiv.org/html/2610.09051#S5.p1.1)\.
- Clarket al\.\(2018\)P\. Clark, I\. Cowhey, O\. Etzioni, T\. Khot, A\. Sabharwal, C\. Schoenick, and O\. TafjordThink you have solved question answering? try arc, the ai2 reasoning challenge\.ArXivabs/1803\.05457\.Cited by:[Appendix A](https://arxiv.org/html/2610.09051#A1.SS0.SSS0.Px3.p1.1)\.
- Daoet al\.\(2022\)T\. Dao, D\. Y\. Fu, S\. Ermon, A\. Rudra, and C\. RéFlashAttention: fast and memory\-efficient exact attention with io\-awareness\.External Links:2205\.14135,[Link](https://arxiv.org/abs/2205.14135)Cited by:[§4\.2](https://arxiv.org/html/2610.09051#S4.SS2.p3.1)\.
- Dao and Gu \(2024\)T\. Dao and A\. GuTransformers are SSMs: generalized models and efficient algorithms through structured state space duality\.InInternational Conference on Machine Learning \(ICML\),Cited by:[Appendix B](https://arxiv.org/html/2610.09051#A2.SS0.SSS0.Px1),[§1](https://arxiv.org/html/2610.09051#S1.p3.1),[§2](https://arxiv.org/html/2610.09051#S2.p2.1),[§4\.1](https://arxiv.org/html/2610.09051#S4.SS1.p2.1),[§4\.2](https://arxiv.org/html/2610.09051#S4.SS2.p1.2),[§5](https://arxiv.org/html/2610.09051#S5.p2.1)\.
- Dao \(2023\)T\. DaoFlashAttention\-2: faster attention with better parallelism and work partitioning\.External Links:2307\.08691,[Link](https://arxiv.org/abs/2307.08691)Cited by:[§4\.2](https://arxiv.org/html/2610.09051#S4.SS2.p3.1)\.
- Devotoet al\.\(2025\)A\. Devoto, M\. Jeblick, and S\. JégouExpected attention: kv cache compression by estimating attention from future queries distribution\.arXiv preprint arXiv:2510\.00636\.External Links:[Link](https://arxiv.org/abs/2510.00636)Cited by:[§5](https://arxiv.org/html/2610.09051#S5.p4.1.2)\.
- Galeet al\.\(2020\)T\. Gale, M\. Zaharia, C\. Young, and E\. ElsenSparse gpu kernels for deep learning\.External Links:2006\.10901,[Link](https://arxiv.org/abs/2006.10901)Cited by:[Appendix H](https://arxiv.org/html/2610.09051#A8.p1.1)\.
- Gaoet al\.\(2024\)L\. Gao, J\. Tow, B\. Abbasi, S\. Biderman, S\. Black, A\. DiPofi, C\. Foster, L\. Golding, J\. Hsu, A\. Le Noac’h, H\. Li, K\. McDonell, N\. Muennighoff, C\. Ociepa, J\. Phang, L\. Reynolds, H\. Schoelkopf, A\. Skowron, L\. Sutawika, E\. Tang, A\. Thite, B\. Wang, K\. Wang, and A\. ZouThe language model evaluation harness\.Zenodo\.External Links:[Document](https://dx.doi.org/10.5281/zenodo.12608602),[Link](https://zenodo.org/records/12608602)Cited by:[Appendix A](https://arxiv.org/html/2610.09051#A1.SS0.SSS0.Px3.p1.1)\.
- Goelet al\.\(2025\)R\. Goel, J\. Park, M\. Gagrani, D\. Jones, M\. Morse, H\. Langston, M\. Lee, and C\. LottCAOTE: kv cache selection for llms via attention output error\-based token eviction\.External Links:2504\.14051,[Link](https://arxiv.org/abs/2504.14051)Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p5.1.3)\.
- Google DeepMind \(2026\)Google DeepMindGemini 3\.Note:Accessed: 2026\-03\-31External Links:[Link](https://deepmind.google/models/gemini/)Cited by:[§1](https://arxiv.org/html/2610.09051#S1.p1.1)\.
- Grattafioriet al\.\(2024\)A\. Grattafiori, A\. Dubey, A\. Jauhri, A\. Pandey, A\. Kadian, A\. Al\-Dahle, A\. Letman, A\. Mathur, A\. Schelten, A\. Vaughan, A\. Yang, A\. Fan, A\. Goyal, A\. Hartshorn, A\. Yang, A\. Mitra, A\. Sravankumar, A\. Korenev, A\. Hinsvark, A\. Rao, A\. Zhang, A\. Rodriguez, A\. Gregerson, A\. Spataru, B\. Roziere, B\. Biron, B\. Tang, B\. Chern, C\. Caucheteux, C\. Nayak, C\. Bi, C\. Marra, C\. McConnell, C\. Keller, C\. Touret, C\. Wu, C\. Wong, C\. C\. Ferrer, C\. Nikolaidis, D\. Allonsius, D\. Song, D\. Pintz, D\. Livshits, D\. Wyatt, D\. Esiobu, D\. Choudhary, D\. Mahajan, D\. Garcia\-Olano, D\. Perino, D\. Hupkes, E\. Lakomkin, E\. AlBadawy, E\. Lobanova, E\. Dinan, E\. M\. Smith, F\. Radenovic, F\. Guzmán, F\. Zhang, G\. Synnaeve, G\. Lee, G\. L\. Anderson, G\. Thattai, G\. Nail, G\. Mialon, G\. Pang, G\. Cucurell, H\. Nguyen, H\. Korevaar, H\. Xu, H\. Touvron, I\. Zarov, I\. A\. Ibarra, I\. Kloumann, I\. Misra, I\. Evtimov, J\. Zhang, J\. Copet, J\. Lee, J\. Geffert, J\. Vranes, J\. Park, J\. Mahadeokar, J\. Shah, J\. van der Linde, J\. Billock, J\. Hong, J\. Lee, J\. Fu, J\. Chi, J\. Huang, J\. Liu, J\. Wang, J\. Yu, J\. Bitton, J\. Spisak, J\. Park, J\. Rocca, J\. Johnstun, J\. Saxe, J\. Jia, K\. V\. Alwala, K\. Prasad, K\. Upasani, K\. Plawiak, K\. Li, K\. Heafield, K\. Stone, K\. El\-Arini, K\. Iyer, K\. Malik, K\. Chiu, K\. Bhalla, K\. Lakhotia, L\. Rantala\-Yeary, L\. van der Maaten, L\. Chen, L\. Tan, L\. Jenkins, L\. Martin, L\. Madaan, L\. Malo, L\. Blecher, L\. Landzaat, L\. de Oliveira, M\. Muzzi, M\. Pasupuleti, M\. Singh, M\. Paluri, M\. Kardas, M\. Tsimpoukelli, M\. Oldham, M\. Rita, M\. Pavlova, M\. Kambadur, M\. Lewis, M\. Si, M\. K\. Singh, M\. Hassan, N\. Goyal, N\. Torabi, N\. Bashlykov, N\. Bogoychev, N\. Chatterji, N\. Zhang, O\. Duchenne, O\. Çelebi, P\. Alrassy, P\. Zhang, P\. Li, P\. Vasic, P\. Weng, P\. Bhargava, P\. Dubal, P\. Krishnan, P\. S\. Koura, P\. Xu, Q\. He, Q\. Dong, R\. Srinivasan, R\. Ganapathy, R\. Calderer, R\. S\. Cabral, R\. Stojnic, R\. Raileanu, R\. Maheswari, R\. Girdhar, R\. Patel, R\. Sauvestre, R\. Polidoro, R\. Sumbaly, R\. Taylor, R\. Silva, R\. Hou, R\. Wang, S\. Hosseini, S\. Chennabasappa, S\. Singh, S\. Bell, S\. S\. Kim, S\. Edunov, S\. Nie, S\. Narang, S\. Raparthy, S\. Shen, S\. Wan, S\. Bhosale, S\. Zhang, S\. Vandenhende, S\. Batra, S\. Whitman, S\. Sootla, S\. Collot, S\. Gururangan, S\. Borodinsky, T\. Herman, T\. Fowler, T\. Sheasha, T\. Georgiou, T\. Scialom, T\. Speckbacher, T\. Mihaylov, T\. Xiao, U\. Karn, V\. Goswami, V\. Gupta, V\. Ramanathan, V\. Kerkez, V\. Gonguet, V\. Do, V\. Vogeti, V\. Albiero, V\. Petrovic, W\. Chu, W\. Xiong, W\. Fu, W\. Meers, X\. Martinet, X\. Wang, X\. Wang, X\. E\. Tan, X\. Xia, X\. Xie, X\. Jia, X\. Wang, Y\. Goldschlag, Y\. Gaur, Y\. Babaei, Y\. Wen, Y\. Song, Y\. Zhang, Y\. Li, Y\. Mao, Z\. D\. Coudert, Z\. Yan, Z\. Chen, Z\. Papakipos, A\. Singh, A\. Srivastava, A\. Jain, A\. Kelsey, A\. Shajnfeld, A\. Gangidi, A\. Victoria, A\. Goldstand, A\. Menon, A\. Sharma, A\. Boesenberg, A\. Baevski, A\. Feinstein, A\. Kallet, A\. Sangani, A\. Teo, A\. Yunus, A\. Lupu, A\. Alvarado, A\. Caples, A\. Gu, A\. Ho, A\. Poulton, A\. Ryan, A\. Ramchandani, A\. Dong, A\. Franco, A\. Goyal, A\. Saraf, A\. Chowdhury, A\. Gabriel, A\. Bharambe, A\. Eisenman, A\. Yazdan, B\. James, B\. Maurer, B\. Leonhardi, B\. Huang, B\. Loyd, B\. D\. Paola, B\. Paranjape, B\. Liu, B\. Wu, B\. Ni, B\. Hancock, B\. Wasti, B\. Spence, B\. Stojkovic, B\. Gamido, B\. Montalvo, C\. Parker, C\. Burton, C\. Mejia, C\. Liu, C\. Wang, C\. Kim, C\. Zhou, C\. Hu, C\. Chu, C\. Cai, C\. Tindal, C\. Feichtenhofer, C\. Gao, D\. Civin, D\. Beaty, D\. Kreymer, D\. Li, D\. Adkins, D\. Xu, D\. Testuggine, D\. David, D\. Parikh, D\. Liskovich, D\. Foss, D\. Wang, D\. Le, D\. Holland, E\. Dowling, E\. Jamil, E\. Montgomery, E\. Presani, E\. Hahn, E\. Wood, E\. Le, E\. Brinkman, E\. Arcaute, E\. Dunbar, E\. Smothers, F\. Sun, F\. Kreuk, F\. Tian, F\. Kokkinos, F\. Ozgenel, F\. Caggioni, F\. Kanayet, F\. Seide, G\. M\. Florez, G\. Schwarz, G\. Badeer, G\. Swee, G\. Halpern, G\. Herman, G\. Sizov, Guangyi, Zhang, G\. Lakshminarayanan, H\. Inan, H\. Shojanazeri, H\. Zou, H\. Wang, H\. Zha, H\. Habeeb, H\. Rudolph, H\. Suk, H\. Aspegren, H\. Goldman, H\. Zhan, I\. Damlaj, I\. Molybog, I\. Tufanov, I\. Leontiadis, I\. Veliche, I\. Gat, J\. Weissman, J\. Geboski, J\. Kohli, J\. Lam, J\. Asher, J\. Gaya, J\. Marcus, J\. Tang, J\. Chan, J\. Zhen, J\. Reizenstein, J\. Teboul, J\. Zhong, J\. Jin, J\. Yang, J\. Cummings, J\. Carvill, J\. Shepard, J\. McPhie, J\. Torres, J\. Ginsburg, J\. Wang, K\. Wu, K\. H\. U, K\. Saxena, K\. Khandelwal, K\. Zand, K\. Matosich, K\. Veeraraghavan, K\. Michelena, K\. Li, K\. Jagadeesh, K\. Huang, K\. Chawla, K\. Huang, L\. Chen, L\. Garg, L\. A, L\. Silva, L\. Bell, L\. Zhang, L\. Guo, L\. Yu, L\. Moshkovich, L\. Wehrstedt, M\. Khabsa, M\. Avalani, M\. Bhatt, M\. Mankus, M\. Hasson, M\. Lennie, M\. Reso, M\. Groshev, M\. Naumov, M\. Lathi, M\. Keneally, M\. Liu, M\. L\. Seltzer, M\. Valko, M\. Restrepo, M\. Patel, M\. Vyatskov, M\. Samvelyan, M\. Clark, M\. Macey, M\. Wang, M\. J\. Hermoso, M\. Metanat, M\. Rastegari, M\. Bansal, N\. Santhanam, N\. Parks, N\. White, N\. Bawa, N\. Singhal, N\. Egebo, N\. Usunier, N\. Mehta, N\. P\. Laptev, N\. Dong, N\. Cheng, O\. Chernoguz, O\. Hart, O\. Salpekar, O\. Kalinli, P\. Kent, P\. Parekh, P\. Saab, P\. Balaji, P\. Rittner, P\. Bontrager, P\. Roux, P\. Dollar, P\. Zvyagina, P\. Ratanchandani, P\. Yuvraj, Q\. Liang, R\. Alao, R\. Rodriguez, R\. Ayub, R\. Murthy, R\. Nayani, R\. Mitra, R\. Parthasarathy, R\. Li, R\. Hogan, R\. Battey, R\. Wang, R\. Howes, R\. Rinott, S\. Mehta, S\. Siby, S\. J\. Bondu, S\. Datta, S\. Chugh, S\. Hunt, S\. Dhillon, S\. Sidorov, S\. Pan, S\. Mahajan, S\. Verma, S\. Yamamoto, S\. Ramaswamy, S\. Lindsay, S\. Lindsay, S\. Feng, S\. Lin, S\. C\. Zha, S\. Patil, S\. Shankar, S\. Zhang, S\. Zhang, S\. Wang, S\. Agarwal, S\. Sajuyigbe, S\. Chintala, S\. Max, S\. Chen, S\. Kehoe, S\. Satterfield, S\. Govindaprasad, S\. Gupta, S\. Deng, S\. Cho, S\. Virk, S\. Subramanian, S\. Choudhury, S\. Goldman, T\. Remez, T\. Glaser, T\. Best, T\. Koehler, T\. Robinson, T\. Li, T\. Zhang, T\. Matthews, T\. Chou, T\. Shaked, V\. Vontimitta, V\. Ajayi, V\. Montanez, V\. Mohan, V\. S\. Kumar, V\. Mangla, V\. Ionescu, V\. Poenaru, V\. T\. Mihailescu, V\. Ivanov, W\. Li, W\. Wang, W\. Jiang, W\. Bouaziz, W\. Constable, X\. Tang, X\. Wu, X\. Wang, X\. Wu, X\. Gao, Y\. Kleinman, Y\. Chen, Y\. Hu, Y\. Jia, Y\. Qi, Y\. Li, Y\. Zhang, Y\. Zhang, Y\. Adi, Y\. Nam, Yu, Wang, Y\. Zhao, Y\. Hao, Y\. Qian, Y\. Li, Y\. He, Z\. Rait, Z\. DeVito, Z\. Rosnbrick, Z\. Wen, Z\. Yang, Z\. Zhao, and Z\. MaThe llama 3 herd of models\.External Links:2407\.21783,[Link](https://arxiv.org/abs/2407.21783)Cited by:[Appendix A](https://arxiv.org/html/2610.09051#A1.SS0.SSS0.Px2.p1.1),[§5](https://arxiv.org/html/2610.09051#S5.p1.1)\.
- Gu and Dao \(2023\)A\. Gu and T\. DaoMamba: linear\-time sequence modeling with selective state spaces\.arXiv preprint arXiv:2312\.00752\.Cited by:[§1](https://arxiv.org/html/2610.09051#S1.p3.1),[§2](https://arxiv.org/html/2610.09051#S2.p2.1)\.
- Guoet al\.\(2026\)H\. Guo, S\. Yang, T\. Goel, E\. P\. Xing, T\. Dao, and Y\. KimLog\-linear attention\.InThe Fourteenth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=mOJgZWkXKW)Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p3.1)\.
- Hochreiter and Schmidhuber \(1997\)S\. Hochreiter and J\. SchmidhuberLong short\-term memory\.Neural Comput\.9\(8\),pp\.1735–1780\.External Links:ISSN 0899\-7667,[Link](https://doi.org/10.1162/neco.1997.9.8.1735),[Document](https://dx.doi.org/10.1162/neco.1997.9.8.1735)Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p6.1)\.
- Hsiehet al\.\(2024\)C\. Hsieh, S\. Sun, S\. Kriman, S\. Acharya, D\. Rekesh, F\. Jia, Y\. Zhang, and B\. GinsburgRULER: what’s the real context size of your long\-context language models?\.arXiv preprint arXiv:2404\.06654\.Cited by:[Appendix A](https://arxiv.org/html/2610.09051#A1.SS0.SSS0.Px3.p1.1),[§5](https://arxiv.org/html/2610.09051#S5.p5.1)\.
- Jelassiet al\.\(2024\)S\. Jelassi, D\. Brandfonbrener, S\. M\. Kakade, and E\. MalachRepeat after me: transformers are better than state space models at copying\.InProceedings of the 41st International Conference on Machine Learning,ICML’24\.Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p2.1)\.
- Jianget al\.\(2023\)A\. Q\. Jiang, A\. Sablayrolles, A\. Mensch, C\. Bamford, D\. S\. Chaplot, D\. de las Casas, F\. Bressand, G\. Lengyel, G\. Lample, L\. Saulnier, L\. R\. Lavaud, M\. Lachaux, P\. Stock, T\. L\. Scao, T\. Lavril, T\. Wang, T\. Lacroix, and W\. E\. SayedMistral 7b\.External Links:2310\.06825,[Link](https://arxiv.org/abs/2310.06825)Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p4.1)\.
- Katharopouloset al\.\(2020\)A\. Katharopoulos, A\. Vyas, N\. Pappas, and F\. FleuretTransformers are rnns: fast autoregressive transformers with linear attention\.InProceedings of the 37th International Conference on Machine Learning,ICML’20\.Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p2.1)\.
- Liet al\.\(2024\)Y\. Li, Y\. Huang, B\. Yang, B\. Venkitesh, A\. Locatelli, H\. Ye, T\. Cai, P\. Lewis, and D\. ChenSnapKV: llm knows what you are looking for before generation\.arXiv preprint arXiv:2404\.14469\.Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p5.1),[§5](https://arxiv.org/html/2610.09051#S5.p4.1)\.
- Lieberet al\.\(2024\)O\. Lieber, B\. Lenz, H\. Bata, G\. Cohen, J\. Osin, I\. Dalmedigos, E\. Safahi, S\. Meirom, Y\. Belinkov, S\. Shalev\-Shwartz, O\. Abend, R\. Alon, T\. Asida, A\. Bergman, R\. Glozman, M\. Gokhman, A\. Manevich, N\. Ratner, N\. Rozen, E\. Shwartz, M\. Zusman, and Y\. ShohamJamba: a hybrid transformer\-mamba language model\.External Links:2403\.19887,[Link](https://arxiv.org/abs/2403.19887)Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p3.1)\.
- Linet al\.\(2025a\)Z\. Lin, E\. Nikishin, X\. O\. He, and A\. CourvilleForgetting transformer: softmax attention with a forget gate\.InThe Thirteenth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=q2Lnyegkr8)Cited by:[§1](https://arxiv.org/html/2610.09051#S1.p2.1),[§2](https://arxiv.org/html/2610.09051#S2.p6.1),[§3](https://arxiv.org/html/2610.09051#S3.p3.1),[§4\.1](https://arxiv.org/html/2610.09051#S4.SS1.p2.1)\.
- Linet al\.\(2025b\)Z\. Lin, J\. Obando\-Ceron, X\. O\. He, and A\. CourvilleAdaptive computation pruning for the forgetting transformer\.InSecond Conference on Language Modeling,External Links:[Link](https://openreview.net/forum?id=xNj14CY5S1)Cited by:[§1](https://arxiv.org/html/2610.09051#S1.p2.1),[§2](https://arxiv.org/html/2610.09051#S2.p6.1),[§3](https://arxiv.org/html/2610.09051#S3.p3.1)\.
- Liuet al\.\(2026\)J\. Liu, H\. Pu, and Z\. ZhangAsyncSparse: accelerating sparse matrix\-matrix multiplication on asynchronous gpu architectures\.External Links:2604\.17834,[Link](https://arxiv.org/abs/2604.17834)Cited by:[Appendix H](https://arxiv.org/html/2610.09051#A8.p1.1)\.
- Merityet al\.\(2016\)S\. Merity, C\. Xiong, J\. Bradbury, and R\. SocherPointer sentinel mixture models\.External Links:1609\.07843Cited by:[Appendix A](https://arxiv.org/html/2610.09051#A1.SS0.SSS0.Px3.p1.1)\.
- NVIDIAet al\.\(2025\)NVIDIA, :, A\. Blakeman, A\. Basant, A\. Khattar, A\. Renduchintala, A\. Bercovich, A\. Ficek, A\. Bjorlin, A\. Taghibakhshi, A\. S\. Deshmukh, A\. S\. Mahabaleshwarkar, A\. Tao, A\. Shors, A\. Aithal, A\. Poojary, A\. Dattagupta, B\. Buddharaju, B\. Chen, B\. Ginsburg, B\. Wang, B\. Norick, B\. Butterfield, B\. Catanzaro, C\. del Mundo, C\. Dong, C\. Harvey, C\. Parisien, D\. Su, D\. Korzekwa, D\. Yin, D\. Gitman, D\. Mosallanezhad, D\. Narayanan, D\. Fridman, D\. Rekesh, D\. Ma, D\. Pykhtar, D\. Ahn, D\. Riach, D\. Stosic, E\. Long, E\. Segal, E\. Evans, E\. Chung, E\. Galinkin, E\. Bakhturina, E\. Dobrowolska, F\. Jia, F\. Liu, G\. Prasad, G\. Shen, G\. Liu, G\. Chen, H\. Qian, H\. Ngo, H\. Liu, H\. Li, I\. Gitman, I\. Karmanov, I\. Moshkov, I\. Golan, J\. Kautz, J\. P\. Scowcroft, J\. Casper, J\. Seppanen, J\. Lu, J\. Sewall, J\. Zeng, J\. You, J\. Zhang, J\. Zhang, J\. Huang, J\. Xue, J\. Huang, J\. Conway, J\. Kamalu, J\. Barker, J\. Cohen, J\. Jennings, J\. Parmar, K\. Sapra, K\. Briski, K\. Chumachenko, K\. Luna, K\. Santhanam, K\. Kong, K\. Sivamani, K\. Pawelec, K\. Anik, K\. Li, L\. McAfee, L\. Derczynski, L\. Pavao, L\. Vega, L\. Voegtle, M\. Bala, M\. R\. de Melo, M\. N\. Sreedhar, M\. Chochowski, M\. Kliegl, M\. Stepniewska\-Dziubinska, M\. Le, M\. Novikov, M\. Samadi, M\. Andersch, M\. Evans, M\. Martinez, M\. Chrzanowski, M\. Ranzinger, M\. Blaz, M\. Smelyanskiy, M\. Fawzy, M\. Shoeybi, M\. Patwary, N\. Lee, N\. Tajbakhsh, N\. Xu, O\. Rybakov, O\. Kuchaiev, O\. Delalleau, O\. Nitski, P\. Chadha, P\. Shamis, P\. Micikevicius, P\. Molchanov, P\. Dykas, P\. Fischer, P\. Aquilanti, P\. Bialecki, P\. Varshney, P\. Gundecha, P\. Tredak, R\. Karimi, R\. Kandu, R\. El\-Yaniv, R\. Joshi, R\. Waleffe, R\. Zhang, S\. Kavanaugh, S\. Jain, S\. Kriman, S\. Lym, S\. Satheesh, S\. Muralidharan, S\. Narenthiran, S\. Anandaraj, S\. Bak, S\. Kashirsky, S\. Han, S\. Acharya, S\. Ghosh, S\. T\. Sreenivas, S\. Clay, S\. Thomas, S\. Prabhumoye, S\. Pachori, S\. Toshniwal, S\. Prayaga, S\. Jain, S\. Das, S\. Kierat, S\. Majumdar, S\. Han, S\. Singhal, S\. Niverty, S\. Alborghetti, S\. Panguluri, S\. Bhendigeri, S\. N\. Akter, S\. Migacz, T\. Shiri, T\. Kong, T\. Roman, T\. Ronen, T\. Saar, T\. Konuk, T\. Rintamaki, T\. Poon, U\. De, V\. Noroozi, V\. Singh, V\. Korthikanti, V\. Kurin, W\. U\. Ahmad, W\. Du, W\. Ping, W\. Dai, W\. Byeon, X\. Ren, Y\. Xu, Y\. Choi, Y\. Zhang, Y\. Lin, Y\. Suhara, Z\. Yu, Z\. Li, Z\. Li, Z\. Zhu, Z\. Yang, and Z\. ChenNemotron\-h: a family of accurate and efficient hybrid mamba\-transformer models\.External Links:2504\.03624,[Link](https://arxiv.org/abs/2504.03624)Cited by:[2nd item](https://arxiv.org/html/2610.09051#A1.I1.i2.p1.1),[§2](https://arxiv.org/html/2610.09051#S2.p3.1),[§5](https://arxiv.org/html/2610.09051#S5.p2.1)\.
- OpenAI \(2025\)OpenAIIntroducing gpt\-5\.Note:Accessed: 2026\-03\-31External Links:[Link](https://openai.com/index/introducing-gpt-5/)Cited by:[§1](https://arxiv.org/html/2610.09051#S1.p1.1)\.
- Papernoet al\.\(2016\)D\. Paperno, G\. Kruszewski, A\. Lazaridou, N\. Q\. Pham, R\. Bernardi, S\. Pezzelle, M\. Baroni, G\. Boleda, and R\. FernándezThe LAMBADA dataset: word prediction requiring a broad discourse context\.InProceedings of the 54th Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),K\. Erk and N\. A\. Smith \(Eds\.\),Berlin, Germany,pp\.1525–1534\.External Links:[Link](https://aclanthology.org/P16-1144/),[Document](https://dx.doi.org/10.18653/v1/P16-1144)Cited by:[Appendix A](https://arxiv.org/html/2610.09051#A1.SS0.SSS0.Px3.p1.1)\.
- Parket al\.\(2026\)J\. Park, D\. Jones, M\. J\. Morse, R\. Goel, M\. Lee, and C\. LottKeyDiff: key similarity\-based kv cache eviction for long\-context llm inference in resource\-constrained environments\.External Links:2504\.15364,[Link](https://arxiv.org/abs/2504.15364)Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p5.1.3),[§5](https://arxiv.org/html/2610.09051#S5.p4.1.2)\.
- Penget al\.\(2024\)B\. Peng, J\. Quesnelle, H\. Fan, and E\. ShippoleYaRN: efficient context window extension of large language models\.InThe Twelfth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=wHBfxhZu1u)Cited by:[§5](https://arxiv.org/html/2610.09051#S5.p1.1),[§5](https://arxiv.org/html/2610.09051#S5.p3.1)\.
- Poliet al\.\(2023\)M\. Poli, S\. Massaroli, E\. Nguyen, D\. Y\. Fu, T\. Dao, S\. Baccus, Y\. Bengio, S\. Ermon, and C\. RéHyena hierarchy: towards larger convolutional language models\.InProceedings of the 40th International Conference on Machine Learning,ICML’23\.Cited by:[§4\.2](https://arxiv.org/html/2610.09051#S4.SS2.p1.2)\.
- Qiuet al\.\(2025\)Z\. Qiu, Z\. Wang, B\. Zheng, Z\. Huang, K\. Wen, S\. Yang, R\. Men, L\. Yu, F\. Huang, S\. Huang, D\. Liu, J\. Zhou, and J\. LinGated attention for large language models: non\-linearity, sparsity, and attention\-sink\-free\.External Links:2505\.06708,[Link](https://arxiv.org/abs/2505.06708)Cited by:[§4\.2](https://arxiv.org/html/2610.09051#S4.SS2.p1.2)\.
- Sakaguchiet al\.\(2019\)K\. Sakaguchi, R\. L\. Bras, C\. Bhagavatula, and Y\. ChoiWinoGrande: an adversarial winograd schema challenge at scale\.arXiv preprint arXiv:1907\.10641\.Cited by:[Appendix A](https://arxiv.org/html/2610.09051#A1.SS0.SSS0.Px3.p1.1)\.
- Sapet al\.\(2019\)M\. Sap, H\. Rashkin, D\. Chen, R\. Le Bras, and Y\. ChoiSocial iqa: commonsense reasoning about social interactions\.InProceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing \(EMNLP\-IJCNLP\),pp\.4463–4473\.Cited by:[Appendix A](https://arxiv.org/html/2610.09051#A1.SS0.SSS0.Px3.p1.1)\.
- Shahet al\.\(2024\)J\. Shah, G\. Bikshandi, Y\. Zhang, V\. Thakkar, P\. Ramani, and T\. DaoFlashAttention\-3: fast and accurate attention with asynchrony and low\-precision\.External Links:2407\.08608,[Link](https://arxiv.org/abs/2407.08608)Cited by:[§4\.2](https://arxiv.org/html/2610.09051#S4.SS2.p3.1)\.
- Shiet al\.\(2025\)J\. Shi, Y\. Wu, Y\. Peng, B\. Wu, L\. Wang, G\. Liu, and Y\. LuoTrainable dynamic mask sparse attention\.External Links:2508\.02124,[Link](https://arxiv.org/abs/2508.02124)Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p4.1.2)\.
- Soldainiet al\.\(2024\)L\. Soldaini, R\. Kinney, A\. Bhagia, D\. Schwenk, D\. Atkinson, R\. Authur, B\. Bogin, K\. Chandu, J\. Dumas, Y\. Elazar, V\. Hofmann, A\. H\. Jha, S\. Kumar, L\. Lucy, X\. Lyu, N\. Lambert, I\. Magnusson, J\. Morrison, N\. Muennighoff, A\. Naik, C\. Nam, M\. E\. Peters, A\. Ravichander, K\. Richardson, Z\. Shen, E\. Strubell, N\. Subramani, O\. Tafjord, P\. Walsh, L\. Zettlemoyer, N\. A\. Smith, H\. Hajishirzi, I\. Beltagy, D\. Groeneveld, J\. Dodge, and K\. LoDolma: An Open Corpus of Three Trillion Tokens for Language Model Pretraining Research\.arXiv preprint\.External Links:[Link](https://arxiv.org/abs/2402.00159)Cited by:[Appendix A](https://arxiv.org/html/2610.09051#A1.SS0.SSS0.Px1.p2.1),[§5](https://arxiv.org/html/2610.09051#S5.p1.1)\.
- Suet al\.\(2024\)J\. Su, M\. Ahmed, Y\. Lu, S\. Pan, W\. Bo, and Y\. LiuRoFormer: enhanced transformer with rotary position embedding\.Neurocomput\.568\(C\)\.External Links:ISSN 0925\-2312,[Link](https://doi.org/10.1016/j.neucom.2023.127063),[Document](https://dx.doi.org/10.1016/j.neucom.2023.127063)Cited by:[§1](https://arxiv.org/html/2610.09051#S1.p3.1),[§2](https://arxiv.org/html/2610.09051#S2.p6.1),[§5](https://arxiv.org/html/2610.09051#S5.p1.1)\.
- Teamet al\.\(2025\)G\. Team, A\. Kamath, J\. Ferret, S\. Pathak, N\. Vieillard, R\. Merhej, S\. Perrin, T\. Matejovicova, A\. Ramé, M\. Rivière, L\. Rouillard, T\. Mesnard, G\. Cideron, J\. Grill, S\. Ramos, E\. Yvinec, M\. Casbon, E\. Pot, I\. Penchev, G\. Liu, F\. Visin, K\. Kenealy, L\. Beyer, X\. Zhai, A\. Tsitsulin, R\. Busa\-Fekete, A\. Feng, N\. Sachdeva, B\. Coleman, Y\. Gao, B\. Mustafa, I\. Barr, E\. Parisotto, D\. Tian, M\. Eyal, C\. Cherry, J\. Peter, D\. Sinopalnikov, S\. Bhupatiraju, R\. Agarwal, M\. Kazemi, D\. Malkin, R\. Kumar, D\. Vilar, I\. Brusilovsky, J\. Luo, A\. Steiner, A\. Friesen, A\. Sharma, A\. Sharma, A\. M\. Gilady, A\. Goedeckemeyer, A\. Saade, A\. Feng, A\. Kolesnikov, A\. Bendebury, A\. Abdagic, A\. Vadi, A\. György, A\. S\. Pinto, A\. Das, A\. Bapna, A\. Miech, A\. Yang, A\. Paterson, A\. Shenoy, A\. Chakrabarti, B\. Piot, B\. Wu, B\. Shahriari, B\. Petrini, C\. Chen, C\. L\. Lan, C\. A\. Choquette\-Choo, C\. Carey, C\. Brick, D\. Deutsch, D\. Eisenbud, D\. Cattle, D\. Cheng, D\. Paparas, D\. S\. Sreepathihalli, D\. Reid, D\. Tran, D\. Zelle, E\. Noland, E\. Huizenga, E\. Kharitonov, F\. Liu, G\. Amirkhanyan, G\. Cameron, H\. Hashemi, H\. Klimczak\-Plucińska, H\. Singh, H\. Mehta, H\. T\. Lehri, H\. Hazimeh, I\. Ballantyne, I\. Szpektor, I\. Nardini, J\. Pouget\-Abadie, J\. Chan, J\. Stanton, J\. Wieting, J\. Lai, J\. Orbay, J\. Fernandez, J\. Newlan, J\. Ji, J\. Singh, K\. Black, K\. Yu, K\. Hui, K\. Vodrahalli, K\. Greff, L\. Qiu, M\. Valentine, M\. Coelho, M\. Ritter, M\. Hoffman, M\. Watson, M\. Chaturvedi, M\. Moynihan, M\. Ma, N\. Babar, N\. Noy, N\. Byrd, N\. Roy, N\. Momchev, N\. Chauhan, N\. Sachdeva, O\. Bunyan, P\. Botarda, P\. Caron, P\. K\. Rubenstein, P\. Culliton, P\. Schmid, P\. G\. Sessa, P\. Xu, P\. Stanczyk, P\. Tafti, R\. Shivanna, R\. Wu, R\. Pan, R\. Rokni, R\. Willoughby, R\. Vallu, R\. Mullins, S\. Jerome, S\. Smoot, S\. Girgin, S\. Iqbal, S\. Reddy, S\. Sheth, S\. Põder, S\. Bhatnagar, S\. R\. Panyam, S\. Eiger, S\. Zhang, T\. Liu, T\. Yacovone, T\. Liechty, U\. Kalra, U\. Evci, V\. Misra, V\. Roseberry, V\. Feinberg, V\. Kolesnikov, W\. Han, W\. Kwon, X\. Chen, Y\. Chow, Y\. Zhu, Z\. Wei, Z\. Egyed, V\. Cotruta, M\. Giang, P\. Kirk, A\. Rao, K\. Black, N\. Babar, J\. Lo, E\. Moreira, L\. G\. Martins, O\. Sanseviero, L\. Gonzalez, Z\. Gleicher, T\. Warkentin, V\. Mirrokni, E\. Senter, E\. Collins, J\. Barral, Z\. Ghahramani, R\. Hadsell, Y\. Matias, D\. Sculley, S\. Petrov, N\. Fiedel, N\. Shazeer, O\. Vinyals, J\. Dean, D\. Hassabis, K\. Kavukcuoglu, C\. Farabet, E\. Buchatskaya, J\. Alayrac, R\. Anil, Dmitry, Lepikhin, S\. Borgeaud, O\. Bachem, A\. Joulin, A\. Andreev, C\. Hardin, R\. Dadashi, and L\. HussenotGemma 3 technical report\.External Links:2503\.19786,[Link](https://arxiv.org/abs/2503.19786)Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p4.1)\.
- Teamet al\.\(2024\)J\. Team, B\. Lenz, A\. Arazi, A\. Bergman, A\. Manevich, B\. Peleg, B\. Aviram, C\. Almagor, C\. Fridman, D\. Padnos, D\. Gissin, D\. Jannai, D\. Muhlgay, D\. Zimberg, E\. M\. Gerber, E\. Dolev, E\. Krakovsky, E\. Safahi, E\. Schwartz, G\. Cohen, G\. Shachaf, H\. Rozenblum, H\. Bata, I\. Blass, I\. Magar, I\. Dalmedigos, J\. Osin, J\. Fadlon, M\. Rozman, M\. Danos, M\. Gokhman, M\. Zusman, N\. Gidron, N\. Ratner, N\. Gat, N\. Rozen, O\. Fried, O\. Leshno, O\. Antverg, O\. Abend, O\. Lieber, O\. Dagan, O\. Cohavi, R\. Alon, R\. Belson, R\. Cohen, R\. Gilad, R\. Glozman, S\. Lev, S\. Meirom, T\. Delbari, T\. Ness, T\. Asida, T\. B\. Gal, T\. Braude, U\. Pumerantz, Y\. Cohen, Y\. Belinkov, Y\. Globerson, Y\. P\. Levy, and Y\. ShohamJamba\-1\.5: hybrid transformer\-mamba models at scale\.External Links:2408\.12570,[Link](https://arxiv.org/abs/2408.12570)Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p3.1)\.
- Vaswaniet al\.\(2017\)A\. Vaswani, N\. Shazeer, N\. Parmar, J\. Uszkoreit, L\. Jones, A\. N\. Gomez, Ł\. Kaiser, and I\. PolosukhinAttention is all you need\.InAdvances in Neural Information Processing Systems,Vol\.30,pp\.\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf)Cited by:[§1](https://arxiv.org/html/2610.09051#S1.p1.1),[§2](https://arxiv.org/html/2610.09051#S2.p6.1)\.
- Wanget al\.\(2019\)A\. Wang, Y\. Pruksachatkun, N\. Nangia, A\. Singh, J\. Michael, F\. Hill, O\. Levy, and S\. BowmanSuperGLUE: a stickier benchmark for general\-purpose language understanding systems\.InAdvances in Neural Information Processing Systems,H\. Wallach, H\. Larochelle, A\. Beygelzimer, F\. d'Alché\-Buc, E\. Fox, and R\. Garnett \(Eds\.\),Vol\.32,pp\.\.External Links:[Link](https://proceedings.neurips.cc/paper/2019/file/4496bf24afe7fab6f046bf4923da8de6-Paper.pdf)Cited by:[Appendix A](https://arxiv.org/html/2610.09051#A1.SS0.SSS0.Px3.p1.1)\.
- xAI \(2026\)xAIOverview — xai docs\.Note:Documents Grok 4\.20 as the newest flagship model\. Accessed: 2026\-03\-31External Links:[Link](https://docs.x.ai/overview)Cited by:[§1](https://arxiv.org/html/2610.09051#S1.p1.1)\.
- Yanget al\.\(2024a\)D\. Yang, X\. Han, Y\. Gao, Y\. Hu, S\. Zhang, and H\. ZhaoPyramidInfer: pyramid KV cache compression for high\-throughput LLM inference\.InFindings of the Association for Computational Linguistics: ACL 2024,L\. Ku, A\. Martins, and V\. Srikumar \(Eds\.\),Bangkok, Thailand,pp\.3258–3270\.External Links:[Link](https://aclanthology.org/2024.findings-acl.195/),[Document](https://dx.doi.org/10.18653/v1/2024.findings-acl.195)Cited by:[§5](https://arxiv.org/html/2610.09051#S5.p4.1)\.
- Yanget al\.\(2025a\)S\. Yang, J\. Kautz, and A\. HatamizadehGated delta networks: improving mamba2 with delta rule\.InThe Thirteenth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=r8H7xhYPwz)Cited by:[Appendix B](https://arxiv.org/html/2610.09051#A2.SS0.SSS0.Px2),[§1](https://arxiv.org/html/2610.09051#S1.p3.1),[§2](https://arxiv.org/html/2610.09051#S2.p2.1),[§2](https://arxiv.org/html/2610.09051#S2.p6.1),[§4\.2](https://arxiv.org/html/2610.09051#S4.SS2.p1.2),[§5](https://arxiv.org/html/2610.09051#S5.p2.1),[§5](https://arxiv.org/html/2610.09051#S5.p5.1)\.
- Yanget al\.\(2025b\)S\. Yang, Y\. Shen, K\. Wen, S\. Tan, M\. Mishra, L\. Ren, R\. Panda, and Y\. KimPaTH attention: position encoding via accumulating householder transformations\.InThe Thirty\-ninth Annual Conference on Neural Information Processing Systems,External Links:[Link](https://openreview.net/forum?id=ZBlHEeSvKd)Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p6.1)\.
- Yanget al\.\(2024b\)S\. Yang, B\. Wang, Y\. Zhang, Y\. Shen, and Y\. KimParallelizing linear transformers with the delta rule over sequence length\.InThe Thirty\-eighth Annual Conference on Neural Information Processing Systems,External Links:[Link](https://openreview.net/forum?id=y8Rm4VNRPH)Cited by:[§1](https://arxiv.org/html/2610.09051#S1.p3.1),[§2](https://arxiv.org/html/2610.09051#S2.p6.1),[§4\.1](https://arxiv.org/html/2610.09051#S4.SS1.p4.1)\.
- Yeet al\.\(2025\)Z\. Ye, L\. Chen, R\. Lai, W\. Lin, Y\. Zhang, S\. Wang, T\. Chen, B\. Kasikci, V\. Grover, A\. Krishnamurthy, and L\. CezeFlashInfer: efficient and customizable attention engine for llm inference serving\.External Links:2501\.01005,[Link](https://arxiv.org/abs/2501.01005)Cited by:[§6](https://arxiv.org/html/2610.09051#S6.p4.1)\.
- Yuanet al\.\(2025\)J\. Yuan, H\. Gao, D\. Dai, J\. Luo, L\. Zhao, Z\. Zhang, Z\. Xie, Y\. Wei, L\. Wang, Z\. Xiao, Y\. Wang, C\. Ruan, M\. Zhang, W\. Liang, and W\. ZengNative sparse attention: hardware\-aligned and natively trainable sparse attention\.InProceedings of the 63rd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),W\. Che, J\. Nabende, E\. Shutova, and M\. T\. Pilehvar \(Eds\.\),Vienna, Austria,pp\.23078–23097\.External Links:[Link](https://aclanthology.org/2025.acl-long.1126/),[Document](https://dx.doi.org/10.18653/v1/2025.acl-long.1126),ISBN 979\-8\-89176\-251\-0Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p4.1)\.
- Zadouriet al\.\(2026\)T\. Zadouri, M\. Hoehnerbach, J\. Shah, T\. Liu, V\. Thakkar, and T\. DaoFlashAttention\-4: algorithm and kernel pipelining co\-design for asymmetric hardware scaling\.External Links:2603\.05451,[Link](https://arxiv.org/abs/2603.05451)Cited by:[§4\.2](https://arxiv.org/html/2610.09051#S4.SS2.p3.1)\.
- Zaheeret al\.\(2020\)M\. Zaheer, G\. Guruganesh, K\. A\. Dubey, J\. Ainslie, C\. Alberti, S\. Ontanon, P\. Pham, A\. Ravula, Q\. Wang, L\. Yang,et al\.Big bird: transformers for longer sequences\.Advances in Neural Information Processing Systems33\.Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p4.1)\.
- Zellerset al\.\(2019\)R\. Zellers, A\. Holtzman, Y\. Bisk, A\. Farhadi, and Y\. ChoiHellaSwag: can a machine really finish your sentence?\.InProceedings of the 57th Annual Meeting of the Association for Computational Linguistics,Cited by:[Appendix A](https://arxiv.org/html/2610.09051#A1.SS0.SSS0.Px3.p1.1)\.
- Zhanget al\.\(2025a\)H\. Zhang, H\. Fan, K\. Sha, Y\. Huang, and Y\. FengDAM: dynamic attention mask for long\-context large language model inference acceleration\.InFindings of the Association for Computational Linguistics: ACL 2025,W\. Che, J\. Nabende, E\. Shutova, and M\. T\. Pilehvar \(Eds\.\),Vienna, Austria,pp\.4663–4676\.External Links:[Link](https://aclanthology.org/2025.findings-acl.242/),[Document](https://dx.doi.org/10.18653/v1/2025.findings-acl.242),ISBN 979\-8\-89176\-256\-5Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p4.1.2)\.
- Zhanget al\.\(2025b\)Y\. Zhang, Z\. Lin, X\. Yao, J\. Hu, F\. Meng, C\. Liu, X\. Men, S\. Yang, Z\. Li, W\. Li, E\. Lu, W\. Liu, Y\. Chen, W\. Xu, L\. Yu, Y\. Wang, Y\. Fan, L\. Zhong, E\. Yuan, D\. Zhang, Y\. Zhang, Y\. T\. Liu, H\. Wang, S\. Fang, W\. He, S\. Liu, Y\. Li, J\. Su, J\. Qiu, B\. Pang, J\. Yan, Z\. Jiang, W\. Huang, B\. Yin, J\. You, C\. Wei, Z\. Wang, C\. Hong, Y\. Chen, G\. Chen, Y\. Wang, H\. Zheng, F\. Wang, Y\. Liu, M\. Dong, Z\. Zhang, S\. Pan, W\. Wu, Y\. Wu, L\. Guan, J\. Tao, G\. Fu, X\. Xu, Y\. Wang, G\. Lai, Y\. Wu, X\. Zhou, Z\. Yang, and Y\. DuKimi linear: an expressive, efficient attention architecture\.External Links:2510\.26692Cited by:[§4\.2](https://arxiv.org/html/2610.09051#S4.SS2.p1.2)\.
- Zhanget al\.\(2023\)Z\. Zhang, Y\. Sheng, T\. Zhou, T\. Chen, L\. Zheng, R\. Cai, Z\. Song, Y\. Tian, C\. Ré, C\. Barrett, Z\. ”\. Wang, and B\. ChenH2O: heavy\-hitter oracle for efficient generative inference of large language models\.InAdvances in Neural Information Processing Systems,A\. Oh, T\. Naumann, A\. Globerson, K\. Saenko, M\. Hardt, and S\. Levine \(Eds\.\),Vol\.36,pp\.34661–34710\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/2023/file/6ceefa7b15572587b78ecfcebb2827f8-Paper-Conference.pdf)Cited by:[§2](https://arxiv.org/html/2610.09051#S2.p5.1),[§5](https://arxiv.org/html/2610.09051#S5.p4.1)\.

## Appendix ATraining, Model, and Evaluation Details

#### Training:

All models are trained across 16 NVIDIA A100s in parallel\. We employ a learning rate of3​e−43e\-4, with warmup over 2k steps and cosine decay\. Optimizer is AdamW withλ=\(\.9,\.95\)\\lambda=\(\.9,\.95\)and weight decay0\.10\.1\.

The models presented in Section[5\.1](https://arxiv.org/html/2610.09051#S5.SS1),[5\.2](https://arxiv.org/html/2610.09051#S5.SS2), and[5\.3](https://arxiv.org/html/2610.09051#S5.SS3)are trained on 50B tokens, while the models for ablation in Section[5\.4](https://arxiv.org/html/2610.09051#S5.SS4)are trained on 10B tokens\. The training data is drawn from Dolma v1\.7\([Soldaini et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib42)\)using the data\-mixture reweighting adopted by[Chu et al\. \(2024\)](https://arxiv.org/html/2610.09051#bib.bib37), with a context length of 4096 and 0\.5M tokens per minibatch\.

#### Model:

We use the Llama\-3 tokenizer\([Grattafiori et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib31)\)\(vocabulary size 128,256\) for all models\. Unless noted below, all models have 32 layers withdmodel=1280d\_\{\\mathrm\{model\}\}=1280\.

- •Base architectures: - –Universal Attention / Llama\-3:nheads=16n\_\{\\mathrm\{heads\}\}\{=\}16,nkv=4n\_\{\\mathrm\{kv\}\}\{=\}4,dff=4096d\_\{\\mathrm\{ff\}\}\{=\}4096,θRoPE=×105\\theta\_\{\\mathrm\{RoPE\}\}\{=\}5\\\!\\times\\\!10^\{5\}\. Context extension uses YaRN with\(βslow,βfast\)=\(0\.7,35\)\(\\beta\_\{\\mathrm\{slow\}\},\\beta\_\{\\mathrm\{fast\}\}\)=\(0\.7,35\)\. - –Mamba2:dstate=128d\_\{\\mathrm\{state\}\}\{=\}128,dconv=4d\_\{\\mathrm\{conv\}\}\{=\}4,dff=3072d\_\{\\mathrm\{ff\}\}\{=\}3072\. - –Gated DeltaNet:nheads=4n\_\{\\mathrm\{heads\}\}\{=\}4,dhead=256d\_\{\\mathrm\{head\}\}\{=\}256,dff=3072d\_\{\\mathrm\{ff\}\}\{=\}3072\. For Mamba2 and Gated DeltaNet,dffd\_\{\\mathrm\{ff\}\}is reduced from 4096 to keep total size≈1\{\\approx\}1B\.
- •Hybrid variants \(GDN\-H, Mamba2\-H\): In both hybrids, layers 4, 12, 20, and 28 are replaced with grouped\-query attention \(nheads=16n\_\{\\mathrm\{heads\}\}\{=\}16,nkv=4n\_\{\\mathrm\{kv\}\}\{=\}4\)\. GDN\-H retainsθRoPE=×105\\theta\_\{\\mathrm\{RoPE\}\}\{=\}5\\\!\\times\\\!10^\{5\}; Mamba2\-H uses no positional encoding in its attention layers following[NVIDIA et al\. \(2025\)](https://arxiv.org/html/2610.09051#bib.bib32)\.
- •KV\-cache compression baselines: - –SnapKV:sparsity ratior=0\.1r\{=\}0\.1,window\_size=\\,\{=\}\\,32,kernel\_size=\\,\{=\}\\,7, max pooling\. - –H2O:rheavy=0\.05r\_\{\\mathrm\{heavy\}\}\{=\}0\.05,rrecent=0\.05r\_\{\\mathrm\{recent\}\}\{=\}0\.05\.

#### Evaluation:

We utilizeEleutherAI\\lm\-evaluation\-harness\([Gao et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib17)\)for evaluating on LongBench\-E\([Bai et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib19)\)and the common\-sense reasoning benchmarks: WikiText\([Merity et al\., 2016](https://arxiv.org/html/2610.09051#bib.bib20)\), LAMBADA\([Paperno et al\., 2016](https://arxiv.org/html/2610.09051#bib.bib21)\), PIQA\([Bisk et al\., 2020](https://arxiv.org/html/2610.09051#bib.bib22)\), HellaSwag\([Zellers et al\., 2019](https://arxiv.org/html/2610.09051#bib.bib23)\), WinoGrande\([Sakaguchi et al\., 2019](https://arxiv.org/html/2610.09051#bib.bib24)\), ARC\-Easy and ARC\-Challenge\([Clark et al\., 2018](https://arxiv.org/html/2610.09051#bib.bib25)\), Social IQa\([Sap et al\., 2019](https://arxiv.org/html/2610.09051#bib.bib26)\), and BoolQ\([Wang et al\., 2019](https://arxiv.org/html/2610.09051#bib.bib27)\)\. For RULER\([Hsieh et al\., 2024](https://arxiv.org/html/2610.09051#bib.bib18)\), we utilize the official GitHub repository111[https://github\.com/NVIDIA/RULER](https://github.com/NVIDIA/RULER), following the defaultsynthetictask set\.

## Appendix BCompute State Size & KV\-Cache Size

Linear attention models are known for their constant state size, which does not grow as the input sequence length grows\. However, at short sequence length, linear attention can often have larger state than a softmax attention with similar number of parameters\. Here we present the steps we used to compute the state size in our Mamba2 model and Gated DeltaNet model, as well as the KV\-cache size of the GQA model\.

#### Mamba2\([Dao and Gu, 2024](https://arxiv.org/html/2610.09051#bib.bib8)\):

Following the official implementation of Mamba2222[https://github\.com/state\-spaces/mamba](https://github.com/state-spaces/mamba), we take the following parameters:

- •dmodeld\_\{\\text\{model\}\}
- •expand, 2 by default,
- •dstated\_\{\\text\{state\}\}, 128 by default\.

The state’s shape is\(nheads,dstate,dhead\)\(n\_\{\\text\{heads\}\},d\_\{\\text\{state\}\},d\_\{\\text\{head\}\}\), where the total number of elements is

sMamba=\\displaystyle s\_\{\\text\{Mamba\}\}=nheads×dstate×dhead\\displaystyle\\ n\_\{\\text\{heads\}\}\\times d\_\{\\text\{state\}\}\\times d\_\{\\text\{head\}\}=\\displaystyle=dmodel×expanddhead×dstate×dhead\\displaystyle\\ \\frac\{d\_\{\\text\{model\}\}\\times\\texttt\{expand\}\}\{d\_\{\\text\{head\}\}\}\\times d\_\{\\text\{state\}\}\\times d\_\{\\text\{head\}\}=\\displaystyle=expand×dmodel×dstate\.\\displaystyle\\ \\texttt\{expand\}\\times d\_\{\\text\{model\}\}\\times d\_\{\\text\{state\}\}\.

#### Gated DeltaNet\([Yang et al\., 2025a](https://arxiv.org/html/2610.09051#bib.bib9)\):

Following the official implementation listed inflash\-linear\-attention333[https://github\.com/fla\-org/flash\-linear\-attention](https://github.com/fla-org/flash-linear-attention), we take the following parameters:

- •dheadd\_\{\\text\{head\}\}, 256 by default,
- •expand\_v, 2 by default,
- •nheadsn\_\{\\text\{heads\}\}, 6 by default\.

The state’s shape is\(num\_heads,v\_head\_dim,head\_dim\)\(\\texttt\{num\\\_heads\},\\texttt\{v\\\_head\\\_dim\},\\texttt\{head\\\_dim\}\), where the total number of elements is

sGDN=\\displaystyle s\_\{\\text\{GDN\}\}=nheads×dhead, v×dhead\\displaystyle\\ n\_\{\\text\{heads\}\}\\times d\_\{\\text\{head, v\}\}\\times d\_\{\\text\{head\}\}=\\displaystyle=expand\_v×nheads×dhead2\.\\displaystyle\\ \\texttt\{expand\\\_v\}\\times n\_\{\\text\{heads\}\}\\times d\_\{\\text\{head\}\}^\{2\}\.

#### Grouped\-Query Attention\([Ainslie et al\., 2023](https://arxiv.org/html/2610.09051#bib.bib15)\):

Each entry of the GQA’s KV cache contains a key and a value, where both of them share the size ofnheads, kv×dheadn\_\{\\text\{heads, kv\}\}\\times d\_\{\\text\{head\}\}\. Thus, the total size of the KV\-cache after seeing a sequence of lengthllis

sGQA=2×nheads, kv×dhead×l\.\\displaystyle s\_\{\\text\{GQA\}\}=2\\times n\_\{\\text\{heads, kv\}\}\\times d\_\{\\text\{head\}\}\\times l\.
To compare the actual memory footprint, we can unify them using the actual size of the stored states\. Thus, for a sequence of lengthll, a hybrid model consisting ofnMamban\_\{\\text\{Mamba\}\}Mamba2 layers,nGDNn\_\{\\text\{GDN\}\}Gated DeltaNet layers, andnGQAn\_\{\\text\{GQA\}\}grouped\-query attention layers, has the following size of saved tensors:

nMamba×sMamba\+nGDN×sGDN\+nGQA×sGQA\.\\displaystyle n\_\{\\text\{Mamba\}\}\\times s\_\{\\text\{Mamba\}\}\+n\_\{\\text\{GDN\}\}\\times s\_\{\\text\{GDN\}\}\+n\_\{\\text\{GQA\}\}\\times s\_\{\\text\{GQA\}\}\.Plugging\-in the actual numbers from our setting as listed in Appendix[A](https://arxiv.org/html/2610.09051#A1), we got the following formula for the size of the stored states and KV\-caches for each of our model configurations, as listed in Table[4](https://arxiv.org/html/2610.09051#A2.T4)\. Here we store all states and KV\-caches in bf16\.

Table 4:Size of the stored states and KV\-caches, w\.r\.t\. sequence lengthll\.

## Appendix CKernel Design

We detail a memory\-efficient Universal Attention layer in listing[1](https://arxiv.org/html/2610.09051#LST1)alongside the triton kernels for both the forward and backwards passes in listings[2](https://arxiv.org/html/2610.09051#LST2)&[3](https://arxiv.org/html/2610.09051#LST3), respectively\.

Decay:The decay mask kernels launch one kernel for both the forward and backwards pass\. Thread\-blocks tile over rows, computing the decay output \(in the forward pass\), or gradients \(in the backwards pass\) for a chunk of columns\.

Memory\-efficient attention:The memory\-efficient attention kernel launches one kernel in the forward pass and two in the backwards pass\. In the forward pass, the attention kernel launches 2D thread\-blocks with a block computing one chunk \(consisting of multiple rows\) of the output tensor\. In the backwards pass, we first launch a preprocessing kernel that computeso⊙\\odotdOsimilar to FlashAttention\. Then we launch the main backwards kernel that computes the gradients of the queries keys and values\. The two functionsmem\_efficient\_bwd\_dkvandmem\_efficient\_bwd\_dqare two subroutines of the main kernel and follow logic similar to FlashAttention\-2\. However, prior to computing softmax normalized values, the decay mask is applied to intermediate attention outputs to correctly compute normalized attention scores\. If GQA is applied, our kernel mimics a broadcast on the fly by indexing the correct kv\-head of the decay\-mask\.

1classUniversalAttention:

2

3

4defforward\(qgate,kgate,Q,K,V\):

5decay=decay\_fwd\(qgate,kgate,K\)

6o=mem\_efficient\_fwd\(Q,K,V,decay\)

7deldecay

8returno

9

10defbackward\(dO\):

11decay=decay\_fwd\(qgate,kgate,K\)

12dq,dk,dv,ddecay=mem\_efficient\_bwd\(Q,K,V,decay,dO\)

13dqgate,dkgate,dkv2=decay\_bwd\(ddecay\)

14returndqgate,dkgate,dq,dk,dv

Listing 1:Pseudocode for a memory\-efficient UA layer\. A collection of kernels compute the decay\-mask and attention output, additively applying the decay mask to the attention matrix in a fused online softmax kernel, avoiding expensive I/O to HBM\. Activation\-Checkpointing additionally avoids saving the quadratic decay mask per\-layer\.1

2defdecay\_fwd\(src,dest,K\):

3b,n,f=tl\.program\_id\(0\),tl\.program\_id\(1\),tl\.program\_id\(2\)

4osum,srcsum,destsum=zeros\(by,bx\),zeros\(by,bx\),zeros\(by,bx\)

5forkinrange\(l//by\+1\):

6k,kt=K\[b,n,k\*by:k\*\(by\+1\),:\],K\[b,n,f\*bx:k\*\(bx\+1\),:\]\.T

7srcsum\+=src\[b,n,k\*bx:k\*\(bx\+1\)\]

8destsum\+=dest\[b,n,k\*by:k\*\(by\+1\)\]

9Γ\\Gammac=relu\(dot\(k,kt\)\)\*\*2\+srcsum\+destsum

10osum\+=log\(1\-Γ\\Gammac\)

11decay\[b,n,k\*by:k\*\(by\+1\),k\*bx:k\*\(bx\+1\)\]=osum

12returndecay

13

14defmemory\_efficient\_fwd\(Q,K,V,decay\):

15i,n=tl\.program\_id\(0\),tl\.program\_id\(1\)

16li,mi=\[0\]\*by,\[\-inf\]\*by

17qi=Q\[n,i\*by:\(i\+1\)\*by,:\]

18oi=zeros\(\(by,d\)\)

19forjinrange\(l//bx\+1\):

20s=tl\.matmul\(qi,k\[n%kvh,j\*bx:\(j\+1\)bx,:\]\.T\)\+\\

21decay\[n%kvh,i\*by:\(i\+1\)by,j\*bx:\(j\+1\)bx\]

22mij=tl\.maximum\(mi,tl\.max\(s,1\)\)

23s\-=mij

24p=tl\.exp\(s\)

25lij=tl\.sum\(s,axis=\-1\)

26alpha=tl\.exp\(mi\-mij\)

27li,mi=li\*alpha\+lij,mij

28oi=oi\*alpha\+\\

29tl\.matmul\(p,v\[j\*bx:\(j\+1\)bx,:\]\)

30store\(oi\)

Listing 2:Pseudocode for the decay forward and memory efficient forward kernels illustrating the per thread\-block work\.1defdecay\_bwd\(ddecay,src,dest,K\):

2b,n,f=tl\.program\_id\(0\),tl\.program\_id\(1\),tl\.program\_id\(2\)

3dki,dsrci,suff=zeros\(by,d\),zeros\(by\),0

4forjinrev\_range\(l//bx\+1\):

5ki=K\[b,n,f\*by:\(f\+1\)\*by,:\]

6kj=K\[b,n,j\*bx:\(j\+1\)\*bx,:\]

7Γ\\Gamma=tl\.dot\(ki\*sqrt\(src\),\(kj\*sqrt\(dest\)\)\.T\)

8Γ\\Gammac=relu\(Γ\\Gamma\)\*\*\(2/3\)

9dd=rev\_cumsum\(ddecay\[b,n,f\*by:\(f\+1\)\*by,j\*bx:\(j\+1\)\*bx\]\)\+suff

10suff\+=tl\.sum\(dd,\-1\)

11dd=where\(Γ\\Gamma\>0,\-dd/\(1\-Γ\\Gammac\)\*\(2/3\)\*relu\(Γ\\Gamma\)\*\*\(\-1/3\),0\)

12dki\+=tl\.dot\(dd,kj\*sqrt\(dest\)\)

13dsrci\+=tl\.sum\(dd\*Γ\\Gamma,\-1\)/\(2\*src\)

14dk\_i\[b,n,f\*by:\(f\+1\)\*by,:\]=dki\*sqrt\(src\)

15dsrc\[b,n,f\*by:\(f\+1\)\*by\]=dsrci

16

17returndk\_i\+dk\_j,dsrc,ddest

18

19

20defmem\_efficient\_bwd\_dkv\(Q,K,V,decay,dO\):

21j,n=tl\.program\_id\(0\),tl\.program\_id\(1\)

22kj=K\[n%kvh,j\*bx:\(j\+1\)\*bx,:\]

23vj=V\[n%kvh,j\*bx:\(j\+1\)\*bx,:\]

24dkj,dvj=zeros\(bx,d\),zeros\(bx,d\)

25foriinrange\(l//by\+1\):

26qi=Q\[n,i\*by:\(i\+1\)\*by,:\]

27aff=decay\[n%kvh,i\*by:\(i\+1\)\*by,j\*bx:\(j\+1\)\*bx\]

28s=tl\.dot\(kj,qi\.T\)\+aff\.T

29p=tl\.exp\(s\-M\[n,i\*by:\(i\+1\)\*by\]\)

30dvj\+=tl\.dot\(p,dO\[n,i\*by:\(i\+1\)\*by,:\]\)

31dp=tl\.dot\(vj,dO\[n,i\*by:\(i\+1\)\*by,:\]\.T\)

32ds=p\*\(dp\-D\[n,i\*by:\(i\+1\)\*by\]\)

33dkj\+=tl\.dot\(ds,qi\)

34DK\[n%kvh,j\*bx:\(j\+1\)\*bx,:\]=dkj

35DV\[n%kvh,j\*bx:\(j\+1\)\*bx,:\]=dvj

36

37defmem\_efficient\_bwd\_dq\(Q,K,V,decay,dO\):

38i,n=tl\.program\_id\(0\),tl\.program\_id\(1\)

39qi=Q\[n,i\*by:\(i\+1\)\*by,:\]

40doi=dO\[n,i\*by:\(i\+1\)\*by,:\]

41mi=M\[n,i\*by:\(i\+1\)\*by\]

42dqi=zeros\(by,d\)

43forjinrange\(l//bx\+1\):

44kj=K\[n%kvh,j\*bx:\(j\+1\)\*bx,:\]

45vj=V\[n%kvh,j\*bx:\(j\+1\)\*bx,:\]

46aff=decay\[n%kvh,i\*by:\(i\+1\)\*by,j\*bx:\(j\+1\)\*bx\]

47s=tl\.dot\(qi,kj\.T\)\+aff

48p=tl\.exp\(s\-mi\)

49dp=tl\.dot\(doi,vj\.T\)

50ds=p\*\(dp\-D\[n,i\*by:\(i\+1\)\*by\]\)

51ddecay\[n%kvh,i\*by:\(i\+1\)\*by,j\*bx:\(j\+1\)\*bx\]=ds

52dqi\+=tl\.dot\(ds,kj\)

53DQ\[n,i\*by:\(i\+1\)\*by,:\]=dqi

Listing 3:Pseudocode for the decay backward and memory\-efficient decay backward kernels illustrating the per thread\-block work\.
## Appendix DHybrid Model Results

We further evaluate UA in hybrid architectures that combine it with complementary sequence\-modeling mechanisms\. We consider two variants:UA\-Softmax, which replaces one out of every eight UA layers with standard full attention, andUA\-Mamba, which alternates UA and Mamba layers with an even split\. All models follow the same training recipe as our main experiments\. To isolate the architectural effects of hybridization, we evaluate these models without KV\-cache pruning\.

Table[5](https://arxiv.org/html/2610.09051#A4.T5)shows that UA remains effective when incorporated into both hybrid architectures\. UA\-Softmax achieves performance close to pure UA, with RULER/LongBench\-E averages of 48\.8/17\.6 compared with 48\.5/18\.1 for UA\. UA\-Mamba is slightly weaker overall at 47\.9/15\.9, but reduces the effective cache relative to the full\-attention hybrid\.

The hybrids also exhibit different capability trade\-offs\. In particular, introducing full\-attention or Mamba layers substantially improves aggregation performance relative to pure UA, while reducing tracking performance\. This suggests that UA can serve as a cache\-efficient, high\-resolution attention component within broader hybrid architectures, with full\-attention or recurrent layers providing complementary memory behavior when stronger aggregation capability is desired\.

Table 5:Hybrid UA model performance, compared against pure Llama3 and UA models, with no eviction\.
## Appendix EAdditional Ablation Results

We further evaluate the robustness of UA to alternative architectural choices and compare against individual decay mechanisms and sliding\-window attention\. Table[6](https://arxiv.org/html/2610.09051#A5.T6)reports averaged RULERS1,S2, andS3scores across context lengths for 1B models trained on 10B tokens\.

UA remains effective when adopting Qwen3\-style normalized queries and keys, with slightly improved performance at 16K\. We also evaluate each decay mechanism in isolation\. In particular, using only the query\-aware decayf\(q\)f^\{\(\\text\{q\}\)\}corresponds to the FoX\-style design\. While this variant performs reasonably at 4K, its performance degrades substantially with context length, reaching 40\.0 at 16K compared with 58\.4 for full UA\. The key\-aware and similarity\-aware mechanisms likewise underperform the complete model when used alone, indicating that the three decay signals provide complementary retention behavior\. Sliding\-window attention is competitive only at the shortest context length and degrades sharply as context increases\. Overall, these results show that UA is robust to alternative query/key normalization choices and that combining multiple decay mechanisms is important for robust long\-context performance\.

Table 6:Additional ablations on model design choices\. 1B models, trained on 10B tokens\.
## Appendix F3B Model Results

To evaluate whether UA’s benefits persist at larger model scale, we train 3B\-parameter Llama3\-style softmax and UA models using the same 50B\-token training recipe as in our main experiments\. Table[7](https://arxiv.org/html/2610.09051#A6.T7)reports RULER and LongBench\-E results together with effective cache size\.

At 3B scale, UA continues to improve both model quality and cache efficiency\. On RULER at 4K context, UA achieves an average score of 52\.6 compared with 48\.8 for Llama3\-3B, whilep=0\.05p\{=\}0\.05reduces the effective cache from 341\.8 MB to 46\.6 MB, corresponding to a7\.3×7\.3\\timesreduction\. On LongBench\-E \(0–4K\), UA reaches 23\.5 compared with 21\.5 for Llama3\-3B while reducing cache from 271\.9 MB to 41\.5 MB, a6\.5×6\.5\\timesreduction\. These results indicate that UA’s quality–memory advantages extend beyond the 1B setting\.

Table 7:Llama3\-3B and UA\-3B model performance, trained on 50B tokens\.
## Appendix GPost\-Training Adaptation

Although UA is designed as a distinct attention architecture and is primarily evaluated through pretraining from scratch, we also investigate whether it can be introduced into an existing pretrained Transformer checkpoint through continued training\.

Starting from our pretrained Llama3 model, we convert its attention layers to UA while retaining the pretrained Llama3 weights, randomly initialize the UA\-specific parameters—including the low\-rank attention\-gate projections and decay parameters—and continue training for an additional 10B tokens\. Table[8](https://arxiv.org/html/2610.09051#A7.T8)reports the resulting RULER and LongBench\-E performance together with effective cache size\.

Table 8:Post\-training adaptation of a pretrained Llama3\-1B checkpoint to UA\.Post\-training adaptation substantially improves over the original Llama3 baseline while enabling significant cache compression\. Atp=0\.05p\{=\}0\.05, Llama3\-UA improves the RULER average from 40\.8 to 44\.3 while reducing the cache from 142\.4 MB to 38\.6 MB; on LongBench\-E, it improves from 16\.4 to 17\.2 while reducing the cache from 113\.3 MB to 41\.3 MB\. The adapted model remains below UA trained from scratch, which reaches a RULER average of 50\.3 with a 21\.4 MB cache at the same pruning threshold\. These results suggest that UA can be incorporated into existing checkpoints through continued training, although training from scratch currently yields the strongest performance–memory trade\-off\.

## Appendix HLatency Analysis

We evaluate the end\-to\-end inference\-time efficiency of UA\. An inference\-optimized dynamic cache manager is not implemented in the current system, as efficiently operating on dynamically pruned KV caches requires sparse matrix–vector computation and specialized sparse kernels, which remain active areas of systems research\([Liu et al\., 2026](https://arxiv.org/html/2610.09051#bib.bib59);[Gale et al\., 2020](https://arxiv.org/html/2610.09051#bib.bib60)\)\. Nevertheless, the current implementation already provides measurable benefits during training and decoding\.

For prefill, our custom UA kernels improve end\-to\-end training throughput by approximately 1\.13×\\times; since prefill exercises the same forward attention path, this suggests that kernel optimization can also reduce TTFT, although we do not report a matched inference\-prefill baseline here\. For decoding, we evaluate a naive UA inference implementation with and without KV\-cache pruning\. As shown in Table[9](https://arxiv.org/html/2610.09051#A8.T9), pruning yields approximately 1\.1×\\times–1\.4×\\timesper\-token decode speedup \(11–29% lower latency\), with larger benefits at longer prompt lengths\.

Table 9:Prefill latency, per\-step decode latency with and without KV\-cache pruning, and post\-eviction cache size at different sequence lengths\.
## Appendix IDetailed RULER results for Figure[2](https://arxiv.org/html/2610.09051#S5.F2)

We report the complete RULER results underlying Figure[2](https://arxiv.org/html/2610.09051#S5.F2), including per\-task scores for all models at context lengths from 1K to 16K and their corresponding effective cache sizes\. In addition to the configurations shown in the main comparison, we include GDN\-H with YaRN at 8K and 16K for a consistent long\-context comparison, as well as full\-cache UA \(p=0p\{=\}0\) to separate the effects of the UA architecture from those of cache pruning\. Applying YaRN substantially improves GDN\-H at long context, particularly on retrieval, but UA withp=0\.05p\{=\}0\.05still achieves a better overall RULER score with a smaller effective cache at both 8K and 16K\.

Table 10:Detailed RULER \(1k\) scores, with effective cache size\.Table 11:Detailed RULER \(2k\) scores, with effective cache size\.Table 12:Detailed RULER \(4k\) scores, with effective cache size\.Table 13:Detailed RULER \(8k\) scores, with effective cache size\.Table 14:Detailed RULER \(16k\) scores, with effective cache size\.