Softmax attention as a sum of RNNs

Attention
$t;$ $s\le t$	query position; an earlier key position
$q_t,$ $k_s;$ $d$	query and key vectors; their dimension (the head dimension)
$v_s;$ $e$	value vector; its dimension
$o_t$	attention output at position $t$
$\phi,$ $\psi$	feature maps of linear attention, applied to queries and keys
$S_t,$ $z_t$	linear-attention state and normalizer
Taylor expansion
$x^{\otimes n}$	$n$ -fold Kronecker power of $x\in\mathbb R^d$ , all $d^n$ products of $n$ coordinates
$H^{(n)}_t$	order- $n$ state, $\sum_{s\le t}k_s^{\otimes n}v_s^\top$ , of size $d^n\times e$
Cottention and 2Mamba
$N$	sequence or context length in tokens
$m$	learned per-head scalar in Cottention's scaling $N^{\sigma(m)}$
$\sigma$	logistic sigmoid
$dt$	Mamba-2's per-token step size, which also scales the values
$u_t$	learned linear function of token $t$ that sets its decay
$a_t$	2Mamba decay at token $t$ , in $(0,1)$
$w_{ts}$	2Mamba's unnormalized weight of key $s$ for query $t$

Context. Causal softmax attention $o_t=\sum_{s\le t}e^{q_t^\top k_s}v_s\big/\sum_{s\le t}e^{q_t^\top k_s}$ : position $t$ averages the values $v_s\in\mathbb R^e$ of positions $s\le t$ , weighted through its query $q_t\in\mathbb R^d$ and their keys $k_s\in\mathbb R^d$ ; the usual $1/\sqrt d$ is absorbed into $q_t$ makes each output a weighted average of all earlier values, so a decoder keeps every past key and value (the KV cachethe stored keys and values of all earlier tokens, which grows by one key and one value per token per head) and training cost grows with the square of the sequence length. Linear attentionattention with $e^{q^\top k}$ replaced by $\phi(q)^\top\psi(k)$ , so that the sum over the past can be accumulated in a fixed-size state (Katharopoulos et al., 2020) replaces $e^{q^\top k}$ by a product of feature mapsfunctions $\phi,\psi:\mathbb R^d\to\mathbb R^d$ applied to the query and to the key separately, such as ReLU or $\mathrm{elu}(x)+1$ $\phi(q)^\top\psi(k)$ ; the sum over the past then passes through one $d\times e$ matrix that is updated token by token, so the layer is an RNNrecurrent neural network: a layer that carries a fixed-size state from token to token, so memory at inference does not grow with the context with constant memory at inference. In language modeling it is less accurate than softmax at the same size, and why has been unclear. Mongaras and Larson (2025) expand the exponential in its Taylor series and show that the numerator of softmax attention is itself a sum of recurrences, one for each power of $q^\top k$ , with linear attention as the first-order term. Cottentioncosine attention: linear attention whose score is the cosine similarity of query and key, from the same group (Mongaras et al., 2024) is the group's earlier linear attention, and 2Mambaa linear-time attention layer: a simplified Mamba-2 whose query-key score is squared, giving a second-order state (Mongaras and Larson, 2026) is the follow-up: Mamba-2a linear-time sequence layer (Dao and Gu, 2024) that can be written as linear attention with a data-dependent decay of the past, plus several extra components reduced to the parts that matter, with a second-order state.

Softmax attention is linear attention summed over every order of interaction between query and key coordinates. The order- $n$ term runs on a state of $d^n\times e$ numbers (about $d^n/n!$ rows once repeated products are merged), so a fixed-size recurrent state keeps only finitely many terms, and the question becomes how many matter. In models of about 300M parameters, the normalizing denominator can be replaced by any vector norm with no loss, terms up to the tenth power reproduce softmax, and each added order moves linear attention toward it. 2Mamba keeps the second-order term, with a data-dependent decay of the past, and nearly matches softmax's test loss at 300M and 700M parameters while its memory stays fixed, below that of a KV cache beyond about a thousand tokens per head. All of this evidence comes from models of at most 2B parameters, most of them trained on a few billion tokens.

Setting. Fix one head. Linear attention computes $o_t=S_t^\top\phi(q_t)\big/z_t^\top\phi(q_t)$ from a state $S_t=S_{t-1}+\psi(k_t)v_t^\top$ of $d\times e$ numbers and a normalizer $z_t=z_{t-1}+\psi(k_t)$ , whatever the length of the past. To see which parts of softmax attention such a state can hold, use two facts: $e^x=\sum_n x^n/n!$ , and the $n$ -th power of an inner product is the inner product of Kronecker powers $x^{\otimes n}=x\otimes\cdots\otimes x\in\mathbb R^{d^n}$ , the vector of all products $x_{i_1}x_{i_2}\cdots x_{i_n}$ of $n$ coordinates of $x$ , $(q^\top k)^n=(q^{\otimes n})^\top k^{\otimes n}$ . Substituting both and exchanging the sums over $n$ and $s$ writes the numerator as a weighted sum of recurrences, each with the form of linear attention, where $H^{(n)}_t$ the order- $n$ state $\sum_{s\le t}k_s^{\otimes n}v_s^\top$ , a $d^n\times e$ matrix that is updated by one rank-one term per token is the order- $n$ state: $\begin{aligned}&\sum_{s\le t}e^{q_t^\top k_s}v_s\\&\quad=\sum_{n=0}^{\infty}\frac{1}{n!}\underbrace{\big(H^{(n)}_t\big)^{\top}q_t^{\otimes n}}_{\text{order-}n\text{ readout}},\\&H^{(n)}_t=\underbrace{H^{(n)}_{t-1}+k_t^{\otimes n}v_t^\top}_{d^n\times e\ \text{state}}.\end{aligned}$ The $n=0$ state is the running sum of the values, $n=1$ is linear attention with $\phi=\psi$ the identity, and $n=2$ collects every product $k_{s,i}k_{s,j}$ of two key coordinates. The denominator is the same sum with each $v_s$ replaced by 1.

Feature maps into $\mathbb R^d$ can reshape the first-order term but cannot add the higher orders, which multiply query coordinates with each other; to carry the order- $n$ term exactly, a state needs one row per distinct product of $n$ coordinates, $\binom{d+n-1}{n}$ rows ( $d(d+1)/2$ at $n=2$ , about $d^n/n!$ when $d$ is much larger than $n$ ), so the size of the state is what limits linear attention. The weights $1/n!$ make high orders small but not negligible, and the experiments (Steps 2 and 3) say that at this scale the exponential matters and the exact denominator does not. For a linear-time layer, the 2026 ablations give a short recipe. Of Mamba-2's many components, a data-dependent decay and a width-2 convolution of the inputs carry most of its accuracy (Step 5); squaring the score then adds the second order (Step 6). This needs a state about $(d+1)/2$ times larger than linear attention's, about 32 times at $d=64$ , which still beats a KV cache once the context passes about a thousand tokens. Exponentiating the score instead of squaring it gives softmax attention with a decay, slightly better than softmax, but it needs a KV cache again.

Step 1

The numerator of softmax attention is exactly a sum of recurrences. The identity $(q^\top k)^n=(q^{\otimes n})^\top k^{\otimes n}$ comes from writing $(\sum_i q_ik_i)^n$ as a sum over all index tuples $(i_1,\dots,i_n)$ of $\prod_j q_{i_j}\prod_j k_{i_j}$ . The exponential series converges absolutely for every score, so the sums over $n$ and $s\le t$ can be exchanged, and the query factors out of the sum over $s$ ; what remains inside is $H^{(n)}_t$ , a state of the same kind linear attention keeps, one per order. The same algebra with the sum running over all $s$ gives the bidirectional case. Two consequences matter later. A truncation of the series to its first-order term, as in linear attention with identity features, can make scores negative, which is why linear attention needs positive feature maps or a norm on its output. A single even power keeps scores nonnegative for any vectors: $(q^\top k)^2\ge0$ , and the truncation $1+x+x^2/2=((x+1)^2+1)/2$ is positive everywhere (my check). 2Mamba uses the square.

Step 2

The denominator works as a normalizer, and any vector norm does the job. Llama-2-styledecoder-only transformers with the Llama 2 layout: RMS norm before each block, rotary position embeddings and gated MLPs models of about 300M parameters (20 layers, head dimension 64) were trained at sequence length 1,024 for 100,000 steps of 36 sequences, about 3.7B tokens, on each of three web and book corpora (the Pilean open English text dataset of about 825 GiB drawn from many domains, such as code, papers and articles, SlimPajamaa cleaned and deduplicated 627B-token web text dataset and FineWeba cleaned and deduplicated web text dataset of about 15 trillion tokens from Common Crawl snapshots), with softmax replaced by its numerator followed by either a gatea per-position factor between 0 and 1 that multiplies the output, as the output gate of an LSTM does or a vector norm. The gate version divides by the sequence length, clamps the scores at 5 and clips gradients, and still shows recoverable loss spikes and slightly worse loss. The norm version, $o_t$ set to the numerator divided by its L2 norm, follows the softmax loss curve on all three datasets, and on FineWeb also at 2B parameters and at sequence length 4,096; RMS norma vector divided by the root mean square of its entries, then scaled coordinate-wise by a learned gain and layer norma vector shifted to mean zero, divided by its standard deviation, then scaled and shifted coordinate-wise by learned parameters do as well. With the L2 norm the output points in the same direction as softmax's, since both divide the same vector by a positive number, and only its length differs (my observation), so the match says the length carried little information here. Two ablations rule out the alternatives: detaching the denominator from the gradient hurts, and keeping the norm but replacing the exponential by a ReLU feature map is much worse, so the exponential is the part that cannot be dropped.

Step 3

Each added order moves linear attention toward softmax, and ten orders reproduce it. Summing the powers 0 to $n$ of $q^\top k$ with weights $1/n!$ gives a training-loss curve that falls steadily toward softmax's as $n$ grows and matches it at $n=10$ . Applying the same added orders to the scores of ReLU, $\mathrm{elu}+1$ and cosine linear attention, $(\phi(q)^\top\psi(k))^n$ , improves the first two but never reaches softmax even at $n=10$ ; the authors attribute this to the feature maps restricting the vectors before the product. Cosine attention gains nothing, consistent with its score lying in $[-1,1]$ , where higher powers stay small. These runs test what function the layer computes and say nothing about an efficient model: an order-10 state at $d=64$ would have $64^{10}$ rows, so the high-order runs must use the quadratic attention-matrix form, and only order 2 is practical as a recurrence.

Step 4 (Cottention)

Cosine attention is a usable linear attention when its scale is learned. Normalizing each query and key to unit length makes the score their cosine similarity $q^\top k/(\|q\|\|k\|)$ , the cosine of the angle between the two vectors, in $[-1,1]$ , and the score factors as in linear attention with no feature map, so the layer runs as an RNN with a $d\times e$ state per head. Without a denominator a row of scores can sum to as much as $N$ , the sequence length, which made training unstable; dividing by $N^{\sigma(m)}$ with a learned per-head scalar $m$ a learned scalar per head, initialized to 0.5, so the divisor $N^{\sigma(m)}$ ranges between 1 and $N$ ; here $\sigma$ is the logistic sigmoid stabilizes it, and $m$ decreases during training. A BERTa bidirectional transformer encoder pretrained by predicting masked-out words model with cosine attention averages 81.8 on GLUEa standard suite of English language-understanding tasks, such as entailment and sentiment, scored on fine-tuning against 83.1 for the same model with softmax, and causal models at 300M and 1.2B parameters reach perplexitythe exponential of the mean next-token cross-entropy; lower is better similar to softmax's (shown as curves, with no table). Memory at inference is constant, and a custom CUDA kernela hand-written GPU program, here one that fuses the masked cumulative sums of causal linear attention makes training memory linear in $N$ . By the expansion of the 2025 paper, cosine attention is a first-order method with bounded scores, consistent with its loss staying above softmax's in that paper's runs.

Step 5 (2Mamba)

In Mamba-2, the decay and a short convolution carry the accuracy. The 2026 paper puts each Mamba-2 component, one at a time, into plain linear attention (300M parameters, FineWeb, sequence length 2,048, about 6.6B tokens) and reports test loss. Plain linear attention with a norm on its output has test loss 3.13. A causal convolution of width 2 over the queries, keys and values gives 2.93; wider windows add little (2.90 at width 4). A decay of the past, the A-maska factor $\prod_{r=s+1}^{t}a_r$ on the score of key $s$ at query $t$ , with a per-token, per-head decay $a_t\in(0,1)$ computed from the token, gives 2.88 when the decay is $a_t=\exp(-\mathrm{softplus}(u_t))$ with $u_t$ a learned linear function of the token and softplus $\mathrm{softplus}(x)=\log(1+e^x)$ , a smooth positive version of $\max(x,0)$ , against 2.98 for Mamba-2's original form, which ties the decay to a step size $dt$ that also scales the values. The other parts (output gate, skip term, value step size) each move the loss by 0.07 or less. Combined, convolution and softplus decay give 2.77, and adding the value step size gives 2.76; this is Mamba-2S, the simplified base. Adding the remaining components back changes nothing or hurts (the output gate: 2.78).

Step 6 (2Mamba)

Squaring the score nearly reaches softmax with a fixed state. From Mamba-2S, 2Mamba removes the output norm, squares the score and normalizes by the sum of scores. With weights

w_{ts}

w_{ts}=(q_t^\top k_s)^2\prod_{r=s+1}^{t}a_r

, the squared score of key

s

for query

t

times the decay accumulated since

s

and

q,k,v

taken after the width-2 convolution, the output is

o_t=\sum_{s\le t}w_{ts}\htmlData{gr=gn4}{v_s}/\sum_{s\le t}w_{ts}

, and since

(q^\top k)^2=(q^{\otimes 2})^\top k^{\otimes 2}

it runs as a recurrence. One decoding step for one head, in plain PyTorch, with value dimension

e=d

(

a

computed from the current token; the paper keeps only the

d(d+1)/2

distinct products, with multinomial weights):

# F = torch.nn.functional; q, k, v come after the width-2 convolution
# state H: (d*d, d) and z: (d*d,), both start at zero
q2 = torch.outer(q, q).flatten()     # all d*d products q_i q_j
k2 = torch.outer(k, k).flatten()
a = torch.exp(-F.softplus(x @ w_a))  # decay in (0, 1); x is the token's input
H = a * H + torch.outer(k2, v)
z = a * z + k2
o = (q2 @ H) / (q2 @ z)              # q2 @ k2 = (q @ k)**2 >= 0

I checked numerically that this recurrence equals the quadratic form. At 300M and 700M parameters and training lengths 2,048, 4,096 and 8,192, its test loss nearly matches softmax's and lies below Mamba-2's (curves only, no table); the value step size, which made 700M training diverge in the low-precision kernels, is left out of the final algorithm. Per head the state holds

d(d+1)^2/2+3d

numbers (numerator, denominator and one cached token for the convolution) against

2Nd

for a KV cache, so 2Mamba uses less memory once

N\gt(d+1)^2/4+3/2

why: the state is fixed while the KV cache grows by

2d

numbers per token; the two are equal at

2Nd=d(d+1)^2/2+3d

, about 1,058 tokens at

d=64

. On a needle-in-a-haystacka test that hides a password in a long stretch of irrelevant text and asks the model to repeat it, measuring retrieval from far back in the context test with contexts up to about 16,000 characters, inside its 8,192-token training length, a model trained for 400,000 steps (batches of 64 sequences of 8,192 tokens) retrieves slightly better than softmax and much better than Mamba-2. Exponentiating the score instead, 2Mamba-E, is slightly better than softmax at 700M but needs a KV cache; it is softmax attention with a decay mask and convolution, close to the Forgetting Transformersoftmax attention whose scores are multiplied by a learned per-token forget gate accumulated over the gap, with log gate

\log\sigma(x)

(Lin et al., 2025), whose log decay

\log\sigma(x)=-\mathrm{softplus}(-x)

has the same form.

Technical notes The paper's phrase that linear attention is a first-order approximation means the

n=1

term alone with its own normalizer; a first-order Taylor truncation of the numerator would also keep the

n=0

running sum of values. Its statement that feature maps can only manipulate one recurrent chain holds for maps into

\mathbb R^d

; a map into a larger space can include higher orders, as the second-order Taylor features of Based (Arora et al., 2024) do, and 2Mamba's squared score is such a map with a decay. The loss comparisons in the three papers are shown as curves without repeated seeds or error bars, with few numbers in tables, and none goes beyond 2B parameters. Most runs see a few billion training tokens; the 2Mamba runs at length 8,192 see about 26B, and the 400,000-step run about 210B (my counts from the stated batch sizes and steps).

Terms

Softmax attention: $o_t=\sum_{s\le t}e^{q_t^\top k_s}v_s\big/\sum_{s\le t}e^{q_t^\top k_s}$ ; needs every past key and value
Linear attention: $o_t=S_t^\top\phi(q_t)/z_t^\top\phi(q_t)$ , $S_t=S_{t-1}+\psi(k_t)v_t^\top$ : an RNN with a $d\times e$ state
Kronecker power identity: $(q^\top k)^n=(q^{\otimes n})^\top k^{\otimes n}$ , with $x^{\otimes n}\in\mathbb R^{d^n}$
Order- $n$ state: $H^{(n)}_t=\sum_{s\le t}k_s^{\otimes n}v_s^\top$ ; softmax's numerator is $\sum_n (H^{(n)}_t)^\top q_t^{\otimes n}/n!$
Norm in place of the denominator: numerator divided by its L2 norm; matched softmax at 300M and 2B parameters
Cottention: cosine-similarity scores, divided by $N^{\sigma(m)}$ ; constant memory, GLUE 81.8 against 83.1
Mamba-2S: linear attention + width-2 convolution + softplus decay + value step size; test loss 2.76 against 3.13 for plain linear attention
2Mamba: weights $(q_t^\top k_s)^2\prod_{r=s+1}^t a_r$ , normalized by their sum; state $d(d+1)^2/2+3d$ per head
Memory crossover: 2Mamba's state is smaller than a KV cache once $N\gt(d+1)^2/4+3/2$ , about 1,058 at $d=64$
2Mamba-E: exponentiated score with decay; slightly better than softmax, needs a KV cache; close to the Forgetting Transformer

Quiz

By the Taylor expansion of softmax attention, what keeps a linear-attention layer from representing it exactly?
Why does 2Mamba square the query-key score rather than use a positive feature map such as ReLU on queries and keys?
Cottention normalizes queries and keys to unit length and has no softmax denominator. What keeps its training stable?

Sources

Gabriel Mongaras and Eric C. Larson. On the Expressiveness of Softmax Attention: A Recurrent Neural Network Perspective. Transactions on Machine Learning Research, 2025.
Gabriel Mongaras, Trevor Dohm, and Eric C. Larson. Cottention: Linear Transformers With Cosine Attention. arXiv preprint, 2024.
Gabriel Mongaras and Eric C. Larson. 2Mamba2Furious: Linear in Complexity, Competitive in Accuracy. arXiv preprint, 2026.