Chapter 21 · Section 21.2 · 70 min

Linear attention

Linear attention is attention rewritten so that its cost grows in proportion to the length of the sequence, not with its square. Ordinary attention compares every query with every key, so $T$ tokens need $T^2$ comparisons. Linear attention first adds up all the keys and values into one small table, and then each query reads its answer from that table. When the model generates text, the table is a running total that each new token updates. So the model carries a memory of fixed size instead of a cache that grows with every token.

That fixed-size memory makes linear attention a recurrent network, the first full model this chapter builds. A recurrent network reads one token at a time and keeps a fixed-size summary of everything it has read, called its state (§3.3). §21.1 set up the problem this solves. A softmax Transformer stores the key and value vectors of every past token, a store called the KV cache, and the cache grows by a fixed number of bytes with every token. A recurrent network keeps only its state. That section also showed that if the state is updated linearly, the network can still be trained in parallel over the whole sequence.

Here is the trick in numbers. Take scaled dot-product attention (§4.3), with the queries, keys, and values of $T$ tokens stacked as the rows of matrices $\mathbf{Q}$ , $\mathbf{K}$ , and $\mathbf{V}$ . Delete the softmax and what is left is $(\mathbf{Q}\mathbf{K}^\top)\mathbf{V}$ . Now move one pair of parentheses, $\mathbf{Q}(\mathbf{K}^\top\mathbf{V})$ . Matrix multiplication is associative, so the answer is the same, but the cost is not. Take $T = 65{,}536$ tokens and a head width of $d = 128$ . The first order builds a $T\times T$ matrix and costs about $2.2\times10^{12}$ floating-point operations (FLOPs). The second builds a $d\times d$ matrix and costs about $4.3\times10^{9}$ , which is 512 times less. And $\mathbf{K}^\top\mathbf{V}$ is a sum over tokens, so a model generating text can keep it as a running total.

§4.7 and §12.6 previewed this idea without proof. This section proves it and then asks what it costs. We will see why the softmax blocks the reordering, and how a product of feature vectors in its place unblocks it. We will prove that causal attention then becomes exactly a recurrent network whose state is a matrix. Linear attention trains in parallel and decodes in constant time and memory, yet for years it was worse than softmax attention. By the end you will be able to say precisely why, with a derivation for each reason, which is what interviewers ask for.

Attention is a weighted average with a similarity score

Fix one attention head, the unit of a layer that has its own queries, keys, and values (§4.2). Each token $t$ has a query $\mathbf{q}_t \in \R^{d_k}$ , a key $\mathbf{k}_t \in \R^{d_k}$ , and a value $\mathbf{v}_t \in \R^{d_v}$ , where $d_k$ and $d_v$ are the widths of the keys and the values. Stacked as rows, they form $\mathbf{Q}, \mathbf{K} \in \R^{T\times d_k}$ and $\mathbf{V}\in\R^{T\times d_v}$ . Here is the output of softmax attention at token $t$ , without the causal mask for now,

\mathbf{o}_t = \sum_{s=1}^{T}\frac{\exp(\mathbf{q}_t^\top\mathbf{k}_s/\sqrt{d_k})}{\sum_{r=1}^{T}\exp(\mathbf{q}_t^\top\mathbf{k}_r/\sqrt{d_k})}\,\mathbf{v}_s .

In plain words, score the query against every key, turn the scores into weights that are positive and add up to one, and return the weighted average of the values. This is the soft lookup of §4.1.

Nothing in that recipe needs the exponential in particular. It needs a way to score how similar a query is to a key, and a rule that turns the scores into weights that are nonnegative and sum to one. Any nonnegative similarity function, divided by its sum over the keys, does that job. Such a function is called a kernel here. Kernel methods in machine learning use the word in a stricter sense, and the kernels built from feature maps below meet that stricter sense too.

📐 Kernel attention

Let $\kappa:\R^{d_k}\times\R^{d_k}\to[0,\infty)$ be a similarity function, called the kernel. Kernel attention maps $\mathbf{Q}\in\R^{T\times d_k}$ , $\mathbf{K}\in\R^{T\times d_k}$ , and $\mathbf{V}\in\R^{T\times d_v}$ to $\mathbf{O}\in\R^{T\times d_v}$ with rows

\mathbf{o}_t = \frac{\sum_{s}\kappa(\mathbf{q}_t,\mathbf{k}_s)\,\mathbf{v}_s}{\sum_{s}\kappa(\mathbf{q}_t,\mathbf{k}_s)} ,

where the sums run over all $s$ (non-causal) or over $s\le t$ (causal), and the denominator is assumed positive. Softmax attention is the case $\kappa_{\text{sm}}(\mathbf{q},\mathbf{k}) = \exp(\mathbf{q}^\top\mathbf{k}/\sqrt{d_k})$ , which we call the softmax kernel.

In plain words, attention is a weighted average of values, and the kernel decides the weights. Statisticians have used this form since 1964, when Nadaraya and, independently, Watson proposed it as a kernel smoother. A kernel smoother predicts the value at a point by averaging nearby observations, each weighted by a kernel. Tsai et al. (2019) read attention this way, with queries and keys as the points and values as the observations.

Why must $\kappa$ be nonnegative? If some weights could be negative, the denominator could be zero or negative. The output would no longer be an average of the values, and a tiny denominator would blow it up. Keep this in mind, because it decides which feature maps are allowed.

Moving the parentheses

Suppose the similarity of a query and a key can be computed in two steps. First turn each vector into a list of $d_\phi$ numbers, its features, with a fixed function $\phi$ called the feature map. Then take the dot product of the two feature lists,

\kappa(\mathbf{q},\mathbf{k}) = \phi(\mathbf{q})^\top\phi(\mathbf{k}), \qquad \phi:\R^{d_k}\to\R^{d_\phi}.

A tiny example shows the idea. In two dimensions, take $\phi(\mathbf{x}) = (x_1^2,\ \sqrt2\,x_1x_2,\ x_2^2)$ . For $\mathbf{q} = (1, 2)$ and $\mathbf{k} = (3, 1)$ the features are $(1, 2\sqrt2, 4)$ and $(9, 3\sqrt2, 1)$ , and their dot product is $9 + 12 + 4 = 25$ . That is $(\mathbf{q}^\top\mathbf{k})^2 = 5^2$ , and the identity holds for every pair, so this $\phi$ computes the kernel $\kappa(\mathbf{q},\mathbf{k}) = (\mathbf{q}^\top\mathbf{k})^2$ . The kernel is never negative, even though one feature can be. It can be zero, though, which is why the maps used in practice keep every feature strictly positive. Any kernel of this two-step form can be reordered, as follows.

Apply $\phi$ to each row, so $\phi(\mathbf{Q})\in\R^{T\times d_\phi}$ and $\phi(\mathbf{K})\in\R^{T\times d_\phi}$ . In the numerator of kernel attention, the query's features do not depend on $s$ , so they come out of the sum,

\sum_{s}\big(\phi(\mathbf{q}_t)^\top\phi(\mathbf{k}_s)\big)\mathbf{v}_s^\top = \phi(\mathbf{q}_t)^\top\Big(\sum_{s}\phi(\mathbf{k}_s)\mathbf{v}_s^\top\Big) = \phi(\mathbf{q}_t)^\top\big(\phi(\mathbf{K})^\top\mathbf{V}\big).

We wrote the output as a row vector to keep the shapes lined up. The first equality writes each term as $\phi(\mathbf{q}_t)^\top\big(\phi(\mathbf{k}_s)\mathbf{v}_s^\top\big)$ , since matrix products can be regrouped, and then takes the common factor $\phi(\mathbf{q}_t)^\top$ out of the sum. The second uses one more definition. The outer product $\phi(\mathbf{k}_s)\mathbf{v}_s^\top$ is the $d_\phi\times d_v$ matrix whose entry $(i, j)$ is $\phi(\mathbf{k}_s)_i\,v_{s,j}$ . A sum of outer products over the rows of two matrices is their product, here $\phi(\mathbf{K})^\top\mathbf{V}$ . The denominator works the same way with $\mathbf{v}_s$ replaced by the number 1.

📐 Linear attention

Given a feature map $\phi:\R^{d_k}\to\R^{d_\phi}$ applied row by row, non-causal linear attention is

\mathbf{O} = \operatorname{diag}\!\big(\phi(\mathbf{Q})\,(\phi(\mathbf{K})^\top\mathbf{1})\big)^{-1}\,\phi(\mathbf{Q})\big(\phi(\mathbf{K})^\top\mathbf{V}\big) \;\in\; \R^{T\times d_v},

with $\phi(\mathbf{K})^\top\mathbf{V}\in\R^{d_\phi\times d_v}$ and $\phi(\mathbf{K})^\top\mathbf{1}\in\R^{d_\phi}$ computed once and shared by every query ( $\mathbf{1}\in\R^T$ is the all-ones vector). It is kernel attention with $\kappa(\mathbf{q},\mathbf{k}) = \phi(\mathbf{q})^\top\phi(\mathbf{k})$ . It is well defined when every denominator $\phi(\mathbf{q}_t)^\top(\phi(\mathbf{K})^\top\mathbf{1})$ is positive, which a feature map with positive entries guarantees.

In plain words, summarize all keys and values once, into a $d_\phi\times d_v$ table and a $d_\phi$ vector. Then let every query read from the summary instead of from every key.

Why can softmax attention not do this? Its kernel does not factor with finitely many features. The Taylor series $\exp(\mathbf{q}^\top\mathbf{k}) = \sum_{n\ge 0}(\mathbf{q}^\top\mathbf{k})^n/n!$ has a term for every power $n$ . The power $(\mathbf{q}^\top\mathbf{k})^n$ is itself a dot product of features, but it needs $d_k^n$ of them, one for every product of $n$ entries. So the series is a feature map with infinitely many features. No finite map can do the same job, because a finite map caps the rank. With $d_\phi$ features, the matrix of kernel values between $n$ queries and $n$ keys is $\phi(\mathbf{Q})\phi(\mathbf{K})^\top$ , whose rank is at most $d_\phi$ . The softmax kernel produces matrices of every rank, as the box below shows. That is the precise sense in which the softmax blocks the reordering. Everything else in this section is a choice of a finite $\phi$ , and the price of that choice.

🔬 Why no finite feature map gives the softmax kernel

Take $n$ distinct numbers $x_1,\dots,x_n$ , and use $x_i\mathbf{e}$ as both the $i$ -th query and the $i$ -th key, for a fixed unit vector $\mathbf{e}\in\R^{d_k}$ . Let $\mathbf{E}$ be the $n\times n$ matrix with entries $e^{x_ix_j}$ . The Taylor series gives $\mathbf{E} = \sum_{m\ge0}\frac{1}{m!}\mathbf{u}_m\mathbf{u}_m^\top$ with $\mathbf{u}_m = (x_1^m,\dots,x_n^m)$ . Every term is positive semidefinite. The first $n$ vectors $\mathbf{u}_0,\dots,\mathbf{u}_{n-1}$ are the columns of a Vandermonde matrix, which is invertible because the $x_i$ are distinct. So $\mathbf{w}^\top\mathbf{E}\mathbf{w} \ge \sum_{m<n}(\mathbf{u}_m^\top\mathbf{w})^2/m! > 0$ for every $\mathbf{w}\ne\mathbf{0}$ , and $\mathbf{E}$ has full rank $n$ . A feature map with $d_\phi$ features would make this rank at most $d_\phi$ , which fails once $n > d_\phi$ .

Now count the cost. §4.7 counts a product of an $m\times k$ matrix and a $k\times n$ matrix as $2mkn$ FLOPs, one multiply and one add per term. The left order, $(\phi(\mathbf{Q})\phi(\mathbf{K})^\top)\mathbf{V}$ , forms a $T\times T$ matrix in $2T^2d_\phi$ FLOPs and multiplies it by $\mathbf{V}$ in $2T^2d_v$ . The right order, $\phi(\mathbf{Q})(\phi(\mathbf{K})^\top\mathbf{V})$ , forms the $d_\phi\times d_v$ summary in $2Td_\phi d_v$ , and the $T$ queries read it in another $2Td_\phi d_v$ . The normalizer adds only lower-order terms, the $T^2$ additions of the row sums in the left order and $O(Td_\phi)$ in the right order, so we leave it out of the count.

📌 Associativity turns quadratic cost into linear cost

\text{left order } 2T^2(d_\phi + d_v), \qquad \text{right order } 4Td_\phi d_v, \qquad \text{the right order is cheaper exactly when } T > T^\ast = \frac{2d_\phi d_v}{d_\phi + d_v}.

With $d_\phi = d_v = d$ the crossover is $T^\ast = d$ , and the saving is a factor $T/d$ . At $T = 65{,}536$ and $d = 128$ that factor is 512. Memory changes the same way, from a $T\times T$ matrix per head to a $d_\phi\times d_v$ one.

The cell below times both orders in NumPy for a head of width 32, so the predicted crossover is at $T = 32$ . A single wall-clock time on a CPU is noisy, so the cell keeps the best of five runs.

Python · cell 1

The left order's time grows about fourfold or more with each doubling of $T$ , as $T^2$ predicts. At the largest sizes the $T\times T$ matrix no longer fits in the CPU's cache, which adds to it. The right order grows about twofold per doubling, linearly in $T$ , and stays far below a millisecond. The exact times depend on your machine, but the gap between the two orders widens with every doubling, at least as fast as the FLOP ratio $T/d$ . In an interview, write the crossover formula first. It shows that linear attention only pays once the sequence is longer than the head is wide.

Warm-up

Which order to multiply

Write two cost functions for non-causal linear attention. Both return Python ints.

attention_flops(T, d_phi, d_v, order) returns the FLOPs of the two matrix products, ignoring the normalizer and the feature map. With order='left' it counts $(\phi(\mathbf{Q})\phi(\mathbf{K})^\top)\mathbf{V}$ , and with order='right' it counts $\phi(\mathbf{Q})(\phi(\mathbf{K})^\top\mathbf{V})$ . Count an $(m\times k)$ by $(k\times n)$ product as $2mkn$ FLOPs.
crossover_T(d_phi, d_v) returns the smallest positive integer $T$ at which the right order costs strictly fewer FLOPs than the left.

def attention_flops(T, d_phi, d_v, order):
    raise NotImplementedError

def crossover_T(d_phi, d_v):
    raise NotImplementedError

Tests, verified solutions, and AI help, open once you add your API key.

Causal linear attention is a recurrent network

A language model needs the causal version, in which token $t$ may only use the keys of tokens $s\le t$ (§4.4). Repeat the factoring with the sums cut off at $t$ ,

\mathbf{o}_t = \frac{\sum_{s\le t}\big(\phi(\mathbf{q}_t)^\top\phi(\mathbf{k}_s)\big)\mathbf{v}_s}{\sum_{s\le t}\phi(\mathbf{q}_t)^\top\phi(\mathbf{k}_s)} = \frac{\Big(\sum_{s\le t}\mathbf{v}_s\phi(\mathbf{k}_s)^\top\Big)\phi(\mathbf{q}_t)}{\Big(\sum_{s\le t}\phi(\mathbf{k}_s)\Big)^{\!\top}\phi(\mathbf{q}_t)} .

The numerator used $(\phi(\mathbf{q}_t)^\top\phi(\mathbf{k}_s))\mathbf{v}_s = \mathbf{v}_s\big(\phi(\mathbf{k}_s)^\top\phi(\mathbf{q}_t)\big)$ , which moves a scalar to the other side of a vector. Each query now reads a prefix sum, a sum over the tokens up to its own. A prefix sum is a recurrence, because the sum up to $t$ is the sum up to $t-1$ plus one term. Give the two running sums names.

📐 Causal linear attention, recurrent form

With a feature map $\phi:\R^{d_k}\to\R^{d_\phi}$ , initial state $\mathbf{S}_0 = \mathbf{0}\in\R^{d_\phi\times d_v}$ , and initial normalizer $\mathbf{z}_0 = \mathbf{0}\in\R^{d_\phi}$ , for $t = 1,\dots,T$ ,

\mathbf{S}_t = \mathbf{S}_{t-1} + \phi(\mathbf{k}_t)\,\mathbf{v}_t^\top, \qquad \mathbf{z}_t = \mathbf{z}_{t-1} + \phi(\mathbf{k}_t), \qquad \mathbf{o}_t = \frac{\mathbf{S}_t^\top\phi(\mathbf{q}_t)}{\mathbf{z}_t^\top\phi(\mathbf{q}_t)} .

Rows of $\mathbf{S}_t$ index features and columns index value dimensions, so $\mathbf{S}_t = \sum_{s\le t}\phi(\mathbf{k}_s)\mathbf{v}_s^\top$ and $\mathbf{z}_t = \sum_{s\le t}\phi(\mathbf{k}_s)$ . Implementations add a small $\epsilon$ (such as $10^{-6}$ ) to the denominator to guard against division by a tiny number.

In plain words, each new token writes the outer product of its key features and its value into a fixed-size matrix, and adds its key features to a fixed-size vector. Then it reads its output in two steps. It multiplies the matrix by its own query features, and it divides by the vector's dot product with the same query features.

§21.1 studied the scalar recurrence $h_t = a_th_{t-1} + u_tx_t$ . There a decay $a_t$ scales the old state, and the new input $x_t$ is added with weight $u_t$ . Linear attention has the same shape with three differences. The state is a $d_\phi\times d_v$ matrix, there is no decay ( $a_t = 1$ ), and each write is the rank-one matrix $\phi(\mathbf{k}_t)\mathbf{v}_t^\top$ . The update is linear in the state, with no nonlinearity on the path from $\mathbf{S}_{t-1}$ to $\mathbf{S}_t$ . That is exactly the property that let the recurrences of §21.1 be computed in parallel.

Before the general proof, run the recurrence by hand once.

✏️ Linear attention by hand on three tokens

Take $d_k = d_v = 2$ and the feature map $\phi(\mathbf{x}) = \operatorname{elu}(\mathbf{x}) + 1$ of Katharopoulos et al. (2020), applied entry by entry. Here $\operatorname{elu}(x) = x$ for $x > 0$ and $e^x - 1$ otherwise. So $\phi(x) = x + 1$ for $x \ge 0$ and $\phi(x) = e^x$ for $x<0$ . It is always positive, and with nonnegative integer inputs every feature is an integer.

$t$	$\mathbf{q}_t$	$\phi(\mathbf{q}_t)$	$\mathbf{k}_t$	$\phi(\mathbf{k}_t)$	$\mathbf{v}_t$
1	$(0, 0)$	$(1, 1)$	$(1, 0)$	$(2, 1)$	$(1, 0)$
2	$(1, 0)$	$(2, 1)$	$(0, 1)$	$(1, 2)$	$(0, 1)$
3	$(0, 1)$	$(1, 2)$	$(0, 0)$	$(1, 1)$	$(1, 1)$

Step 1. $\mathbf{S}_1 = \phi(\mathbf{k}_1)\mathbf{v}_1^\top = \begin{pmatrix}2 & 0\\ 1 & 0\end{pmatrix}$ and $\mathbf{z}_1 = (2, 1)$ . Then $\mathbf{S}_1^\top\phi(\mathbf{q}_1) = (3, 0)$ and $\mathbf{z}_1^\top\phi(\mathbf{q}_1) = 3$ , so $\mathbf{o}_1 = (1, 0) = \mathbf{v}_1$ . That is as it must be with a single key.

Step 2. $\mathbf{S}_2 = \mathbf{S}_1 + \begin{pmatrix}0 & 1\\ 0 & 2\end{pmatrix} = \begin{pmatrix}2 & 1\\ 1 & 2\end{pmatrix}$ and $\mathbf{z}_2 = (3, 3)$ . With $\phi(\mathbf{q}_2) = (2, 1)$ , $\mathbf{S}_2^\top\phi(\mathbf{q}_2) = (5, 4)$ and $\mathbf{z}_2^\top\phi(\mathbf{q}_2) = 9$ , so $\mathbf{o}_2 = (5/9,\ 4/9)$ .

Step 3. $\mathbf{S}_3 = \mathbf{S}_2 + \begin{pmatrix}1 & 1\\ 1 & 1\end{pmatrix} = \begin{pmatrix}3 & 2\\ 2 & 3\end{pmatrix}$ and $\mathbf{z}_3 = (4, 4)$ . With $\phi(\mathbf{q}_3) = (1, 2)$ , $\mathbf{S}_3^\top\phi(\mathbf{q}_3) = (7, 8)$ and $\mathbf{z}_3^\top\phi(\mathbf{q}_3) = 12$ , so $\mathbf{o}_3 = (7/12,\ 2/3)$ .

The quadratic way. The kernel matrix $\phi(\mathbf{Q})\phi(\mathbf{K})^\top$ has rows $(3, 3, 2)$ , $(5, 4, 3)$ , and $(4, 5, 3)$ . Keep the entries on and below the diagonal. That leaves rows $(3)$ , $(5, 4)$ , and $(4, 5, 3)$ , with sums 3, 9, and 12. Row 2 gives $(5\mathbf{v}_1 + 4\mathbf{v}_2)/9 = (5/9, 4/9)$ , and row 3 gives $(4\mathbf{v}_1 + 5\mathbf{v}_2 + 3\mathbf{v}_3)/12 = (7/12, 8/12)$ . These are the same numbers.

Compare with softmax. Query 2 scores the two keys at $\mathbf{q}_2^\top\mathbf{k}_1 = 1$ and $\mathbf{q}_2^\top\mathbf{k}_2 = 0$ . Softmax (without the $\sqrt{d_k}$ , to keep the numbers round) gives weights $e/(e+1) = 0.73$ and $0.27$ . The elu+1 kernel gives $5/9 = 0.56$ and $0.44$ . The $+1$ inside $\phi$ makes every key look somewhat similar to every query. That is a first hint of the flatness we diagnose later.

The same numbers, computed in parallel

The example computed the outputs two ways, a loop over tokens and a masked matrix, and got the same numbers. The theorem says the two always agree.

∴ Recurrent and parallel forms agree

Let $\mathbf{L}\in\R^{T\times T}$ be the lower-triangular matrix of ones, $\mathbf{L}_{ts} = 1$ for $s\le t$ and 0 otherwise. Let $\mathbf{W} = \big(\phi(\mathbf{Q})\phi(\mathbf{K})^\top\big)\odot\mathbf{L}$ , where $\odot$ is the entrywise product. The outputs of the recurrent form are the rows of

\mathbf{O} = \operatorname{diag}(\mathbf{W}\mathbf{1})^{-1}\,\mathbf{W}\mathbf{V}.

Proof. Row $t$ of $\mathbf{W}\mathbf{V}$ is $\sum_{s}\mathbf{W}_{ts}\mathbf{v}_s^\top = \sum_{s\le t}\big(\phi(\mathbf{q}_t)^\top\phi(\mathbf{k}_s)\big)\mathbf{v}_s^\top$ , because $\mathbf{L}$ zeroes the terms with $s>t$ . Moving the scalar across, this is $\phi(\mathbf{q}_t)^\top\sum_{s\le t}\phi(\mathbf{k}_s)\mathbf{v}_s^\top = \phi(\mathbf{q}_t)^\top\mathbf{S}_t = (\mathbf{S}_t^\top\phi(\mathbf{q}_t))^\top$ . Here $\mathbf{S}_t = \sum_{s\le t}\phi(\mathbf{k}_s)\mathbf{v}_s^\top$ by unrolling the recurrence, which is an induction on $t$ from $\mathbf{S}_0 = \mathbf{0}$ . In the same way, entry $t$ of $\mathbf{W}\mathbf{1}$ is $\sum_{s\le t}\phi(\mathbf{q}_t)^\top\phi(\mathbf{k}_s) = \mathbf{z}_t^\top\phi(\mathbf{q}_t)$ . Dividing row $t$ by entry $t$ gives $\mathbf{o}_t$ . ∎

In plain words, causal linear attention is two things at once. It is ordinary masked attention with the softmax replaced by a kernel and a division by row sums, and it is a recurrence. The mask that makes attention causal and the running sum that makes the recurrence causal are the same object.

§21.1 described every causal layer of this chapter by a sequence-mixing matrix. That is a lower-triangular $T\times T$ matrix $\mathbf{M}$ , computed from the layer's input, that mixes each channel $\mathbf{x}\in\R^T$ as $\mathbf{y} = \mathbf{M}\mathbf{x}$ . Linear attention is the first instance with a matrix state. For each value channel, $\mathbf{M} = \operatorname{diag}(\mathbf{W}\mathbf{1})^{-1}\mathbf{W}$ . The parallel and recurrent forms are that section's matrix mode and recurrent mode. The cell implements both, checks the hand example exactly, and checks agreement on random data.

Python · cell 2

Decoding, training, and what each costs

The two forms have opposite strengths, and real systems use each where it is strong.

Decoding. Generating token $t$ needs only $\mathbf{S}_{t-1}$ and $\mathbf{z}_{t-1}$ . The update costs $d_\phi d_v$ multiply-adds for the outer product, and the read costs another $d_\phi d_v$ . That is about $4d_\phi d_v$ FLOPs per head per token. The memory is $d_\phi d_v + d_\phi$ numbers, whatever $t$ is. Softmax attention instead reads a KV cache that grows by one key and one value per token (§6.5). For $d_\phi = d_v = 128$ the state holds 16,512 numbers. At $2\cdot128$ numbers per token, that is about what this head's KV cache holds after 64 tokens, and the state never grows.

Training. The recurrent form has $T$ steps that must run one after another, which is slow on a GPU. If autograd keeps every intermediate state for the backward pass, it also stores $T d_\phi d_v$ numbers. At $T = 8192$ and $d_\phi = d_v = 128$ in float32, the 4-byte format, that is 512 MiB for one head of one layer. Katharopoulos et al. (2020) avoided the storage by writing the gradients as cumulative sums too. The sums run forward in time for the queries and backward in time for the keys and values, so no intermediate state is stored. The parallel form costs $O(T^2)$ like ordinary attention, but it is all matrix multiplication, which GPUs do best.

The practical answer is the chunkwise algorithm of §21.5. It splits the sequence into chunks of $T_c$ tokens, typically 64 to 256. It uses the parallel form inside each chunk and the recurrence across chunks, and gets most of the benefit of both. It is used for training and for prefill, the one parallel pass over the prompt before decoding starts.

mode	work per head	sequential steps	memory	used for
recurrent	$O(Td_\phi d_v)$	$T$	$O(d_\phi d_v)$ state	decoding
parallel (masked, quadratic)	$O(T^2(d_\phi + d_v))$	$O(1)$ matmuls	$O(T^2)$	short sequences, chunks
chunkwise (§21.5)	$O(TT_c(d_\phi + d_v) + Td_\phi d_v)$	$T/T_c$	$O((T/T_c)d_\phi d_v)$	training and prefill

Why can the causal form not use the non-causal shortcut directly? The non-causal right order computes one summary $\phi(\mathbf{K})^\top\mathbf{V}$ , shared by all queries. Causally, every query needs its own prefix of that summary. That is why the causal form is a recurrence rather than one matrix product, and why training it fast needs the chunkwise algorithm.

Interview

Causal linear attention, recurrent and parallel, plus the decode step

This is the standard interview coding question on linear attention. Implement three functions for one head. Q and K have shape (T, d_k), V has shape (T, d_v), and phi maps an array of shape (..., d_k) to (..., d_phi).

causal_linear_attention_recurrent(Q, K, V, phi, eps=1e-6) loops over time. It must never build a (T, T) array, and the tests disable np.tril, np.triu, and np.tri while it runs.
causal_linear_attention_parallel(Q, K, V, phi, eps=1e-6) builds the masked kernel matrix.
linear_attention_step(S, z, q, k, v, phi, eps=1e-6) is one decode step. It takes the state S of shape (d_phi, d_v) and the normalizer z of shape (d_phi,) after $t-1$ tokens. It also takes one token's q and k of shape (d_k,) and v of shape (d_v,). It returns (o, S_new, z_new) with $\mathbf{o}_t = \mathbf{S}_t^\top\phi(\mathbf{q}_t)/(\mathbf{z}_t^\top\phi(\mathbf{q}_t) + \epsilon)$ , and it must not modify its inputs.

All three add eps to the denominator as written, and each token is included in its own sum.

import numpy as np

def causal_linear_attention_recurrent(Q, K, V, phi, eps=1e-6):
    raise NotImplementedError

def causal_linear_attention_parallel(Q, K, V, phi, eps=1e-6):
    raise NotImplementedError

Tests, verified solutions, and AI help, open once you add your API key.

Choosing the feature map

Everything so far holds for any $\phi$ . The choice of $\phi$ decides three things. It decides how close linear attention comes to softmax attention, how large the state is, and whether training is stable. There are four families, from the simplest to the most faithful to softmax.

Simple positive maps

Katharopoulos et al. (2020) used $\phi(\mathbf{x}) = \operatorname{elu}(\mathbf{x}) + 1$ with $d_\phi = d_k$ . They preferred it to $\operatorname{relu}(x) = \max(x, 0)$ , because ReLU has zero gradient for negative inputs. Both maps are cheap and keep the state as small as it can be. The elu+1 map is positive, and ReLU only nonnegative. Neither approximates the softmax kernel.

One property of these maps matters a lot. A map is positively homogeneous if $\phi(c\mathbf{x}) = c\,\phi(\mathbf{x})$ for every $c > 0$ , and ReLU and the identity both are. For such a map, scaling a query by $c$ multiplies every kernel value by $c$ , and the normalized weights do not change. So a query cannot make its attention sharper by growing longer. A softmax query can, since $\exp(c\,\mathbf{q}^\top\mathbf{k})$ widens the gaps between weights exponentially in $c$ . The elu+1 map is not homogeneous, since it equals $x + 1$ for positive inputs. It still cannot sharpen without limit. For a long query $c\mathbf{q}$ , the positive entries of $\phi(c\mathbf{q})$ grow like $c\,q_i$ and the negative ones shrink to zero, so $\phi(c\mathbf{q})$ points along $\operatorname{relu}(\mathbf{q})$ . The weights then settle at a fixed limit, proportional to $\operatorname{relu}(\mathbf{q})^\top\phi(\mathbf{k}_s)$ , which the $+1$ inside $\phi(\mathbf{k}_s)$ keeps close to uniform.

Random features that estimate the softmax kernel

Can a finite $\phi$ approximate $\exp(\mathbf{q}^\top\mathbf{k})$ instead of replacing it? Choromanski et al. (2020) found a positive feature map that is random, and whose dot product equals the softmax kernel on average over the randomness. Their model is called the Performer. The idea rests on one fact about a standard Gaussian vector $\boldsymbol\omega$ , its moment generating function, the average of $e^{\boldsymbol\omega^\top\mathbf{x}}$ as a function of $\mathbf{x}$ .

Start with one dimension. If $\omega$ is a standard normal number, then $\mathbb{E}[e^{\omega x}] = e^{x^2/2}$ . So at $x = 1$ the average of $e^{\omega}$ is $e^{0.5}\approx 1.65$ . In $d_k$ dimensions, take $\boldsymbol\omega\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{d_k})$ and a fixed $\mathbf{x}\in\R^{d_k}$ . Completing the square in the exponent gives

\mathbb{E}\big[e^{\boldsymbol\omega^\top\mathbf{x}}\big] = \int\frac{e^{-\lVert\boldsymbol\omega\rVert^2/2 + \boldsymbol\omega^\top\mathbf{x}}}{(2\pi)^{d_k/2}}\,d\boldsymbol\omega = e^{\lVert\mathbf{x}\rVert^2/2}\int\frac{e^{-\lVert\boldsymbol\omega - \mathbf{x}\rVert^2/2}}{(2\pi)^{d_k/2}}\,d\boldsymbol\omega = e^{\lVert\mathbf{x}\rVert^2/2}.

The middle step uses $-\frac12\lVert\boldsymbol\omega\rVert^2 + \boldsymbol\omega^\top\mathbf{x} = -\frac12\lVert\boldsymbol\omega - \mathbf{x}\rVert^2 + \frac12\lVert\mathbf{x}\rVert^2$ . The last integrand is the density of $\mathcal{N}(\mathbf{x},\mathbf{I})$ , which integrates to 1 (§0.4 defines the multivariate Gaussian).

So $e^{\boldsymbol\omega^\top\mathbf{x} - \lVert\mathbf{x}\rVert^2/2}$ averages to 1 for every $\mathbf{x}$ . Multiply two such factors, one for $\mathbf{q}$ and one for $\mathbf{k}$ , and the average becomes $e^{\mathbf{q}^\top\mathbf{k}}$ , as the derivation below shows. Each random direction $\boldsymbol\omega_i$ then gives one feature.

∴ Positive random features are unbiased

Let $\boldsymbol\Omega\in\R^{d_\phi\times d_k}$ have independent rows $\boldsymbol\omega_i\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{d_k})$ , and define

\phi(\mathbf{x}) = \frac{e^{-\lVert\mathbf{x}\rVert^2/2}}{\sqrt{d_\phi}}\big(e^{\boldsymbol\omega_1^\top\mathbf{x}},\ \dots,\ e^{\boldsymbol\omega_{d_\phi}^\top\mathbf{x}}\big)\in\R^{d_\phi} .

Then every entry of $\phi(\mathbf{x})$ is positive, and for all $\mathbf{q},\mathbf{k}\in\R^{d_k}$

\mathbb{E}_{\boldsymbol\Omega}\big[\phi(\mathbf{q})^\top\phi(\mathbf{k})\big] = e^{\mathbf{q}^\top\mathbf{k}}, \qquad \frac{\operatorname{Var}\big[\phi(\mathbf{q})^\top\phi(\mathbf{k})\big]}{\big(e^{\mathbf{q}^\top\mathbf{k}}\big)^2} = \frac{e^{\lVert\mathbf{q}+\mathbf{k}\rVert^2} - 1}{d_\phi} .

To include the temperature $1/\sqrt{d_k}$ , apply $\phi$ to $\mathbf{q}/d_k^{1/4}$ and $\mathbf{k}/d_k^{1/4}$ , whose dot product is $\mathbf{q}^\top\mathbf{k}/\sqrt{d_k}$ . The proof is in the derivation below.

∂ Unbiasedness and variance of positive random features

The estimator is an average of $d_\phi$ independent copies of one random number,

\phi(\mathbf{q})^\top\phi(\mathbf{k}) = \frac{1}{d_\phi}\sum_{i=1}^{d_\phi}Y_i, \qquad Y = e^{\boldsymbol\omega^\top\mathbf{q} - \lVert\mathbf{q}\rVert^2/2}\,e^{\boldsymbol\omega^\top\mathbf{k} - \lVert\mathbf{k}\rVert^2/2} = e^{-(\lVert\mathbf{q}\rVert^2 + \lVert\mathbf{k}\rVert^2)/2}\,e^{\boldsymbol\omega^\top(\mathbf{q}+\mathbf{k})} .

Mean. By the Gaussian moment generating function with $\mathbf{x} = \mathbf{q}+\mathbf{k}$ ,

\mathbb{E}[Y] = e^{-(\lVert\mathbf{q}\rVert^2 + \lVert\mathbf{k}\rVert^2)/2}\,e^{\lVert\mathbf{q}+\mathbf{k}\rVert^2/2} = e^{\left(\lVert\mathbf{q}+\mathbf{k}\rVert^2 - \lVert\mathbf{q}\rVert^2 - \lVert\mathbf{k}\rVert^2\right)/2} = e^{\mathbf{q}^\top\mathbf{k}},

using $\lVert\mathbf{q}+\mathbf{k}\rVert^2 = \lVert\mathbf{q}\rVert^2 + 2\mathbf{q}^\top\mathbf{k} + \lVert\mathbf{k}\rVert^2$ . An average of unbiased copies is unbiased.

Second moment. Apply the same fact with $\mathbf{x} = 2(\mathbf{q}+\mathbf{k})$ ,

\mathbb{E}[Y^2] = e^{-(\lVert\mathbf{q}\rVert^2 + \lVert\mathbf{k}\rVert^2)}\,e^{2\lVert\mathbf{q}+\mathbf{k}\rVert^2} = e^{\lVert\mathbf{q}\rVert^2 + \lVert\mathbf{k}\rVert^2 + 4\mathbf{q}^\top\mathbf{k}} = e^{2\mathbf{q}^\top\mathbf{k}}\,e^{\lVert\mathbf{q}+\mathbf{k}\rVert^2} .

So $\operatorname{Var}[Y] = \mathbb{E}[Y^2] - \mathbb{E}[Y]^2 = e^{2\mathbf{q}^\top\mathbf{k}}\big(e^{\lVert\mathbf{q}+\mathbf{k}\rVert^2} - 1\big)$ . Averaging $d_\phi$ independent copies divides the variance by $d_\phi$ . Dividing by the squared mean $e^{2\mathbf{q}^\top\mathbf{k}}$ gives the relative variance in the theorem. This is the mean squared error that Lemma 2 of Choromanski et al. (2020) states for these features. ∎

In plain words, each random direction $\boldsymbol\omega_i$ gives one positive feature. The product of two such features equals the softmax kernel on average. The noise of the estimate shrinks like $1/\sqrt{d_\phi}$ , with a constant that grows exponentially in $\lVert\mathbf{q}+\mathbf{k}\rVert^2$ .

🔬 Why positive, and why orthogonal

The classical random Fourier features of Rahimi and Recht (2007) can also estimate the softmax kernel without bias, using $\sin$ and $\cos$ of $\boldsymbol\omega^\top\mathbf{x}$ . But their estimates can be negative. Their mean squared error also grows without bound as the true kernel value goes to zero, and most entries of a sharp attention matrix are near zero. The mean squared error of positive features goes to zero there instead, in proportion to the squared kernel value (Lemma 2 of Choromanski et al. compares the two). The Performer paper reports that the negative estimates in the denominator make training unstable.

Drawing the rows of $\boldsymbol\Omega$ to be exactly orthogonal keeps the estimator unbiased and provably lowers its mean squared error (their Theorem 2). Exactly orthogonal rows exist only for $d_\phi\le d_k$ , and larger $d_\phi$ stacks several independent orthogonal blocks. The Performer's attention mechanism, called FAVOR+, is these positive orthogonal random features.

Now read the variance formula as an engineer would. Take queries and keys with independent unit-variance entries, as at initialization. Then $\lVert\mathbf{q}+\mathbf{k}\rVert^2\approx 2d_k$ , and after the $d_k^{-1/4}$ scaling of each vector it is about $2\sqrt{d_k}$ . At $d_k = 16$ that is 8, and the relative variance of one kernel value is $(e^8 - 1)/d_\phi\approx 2980/d_\phi$ . At $d_k = 128$ it is about 22.6, and $e^{22.6}\approx 7\times10^{9}$ . No affordable $d_\phi$ estimates such kernel values well. The cell measures the error for 16 queries against 32 keys at $d_k = 16$ , first with the vectors shrunk to half length and then at full length.

Python · cell 3

Why does the error fall more slowly at full length? Each feature is a log-normal random number, mostly small with rare huge values, and the rare values carry most of the variance. An average of $d_\phi$ of them settles into the $1/\sqrt{d_\phi}$ rate only once $d_\phi$ is well above that pair's relative variance per feature, $e^{\lVert\mathbf{q}+\mathbf{k}\rVert^2} - 1$ . Below that, a typical draw misses the rare values and underestimates the kernel. At the measured average of 7.4 that variance is about 1,600, far more for the longest pairs, and those pairs dominate the mean. So random features are a genuinely unbiased estimator, yet in practice they do not keep softmax quality in a language model, and the variance formula says why. An interviewer who asks about Performers wants exactly this answer.

Core

Positive random features for the softmax kernel

Implement the positive random features of the theorem. There is no temperature here, so use the inputs as given.

positive_random_features(X, Omega) takes X of shape (n, d_k) and Omega of shape (d_phi, d_k), whose rows are the random directions. It returns the features, of shape (n, d_phi).
softmax_kernel_estimate(q, k, Omega) takes vectors q and k of shape (d_k,) and returns the scalar estimate $\phi(\mathbf{q})^\top\phi(\mathbf{k})$ of $e^{\mathbf{q}^\top\mathbf{k}}$ .

import numpy as np

def positive_random_features(X, Omega):
    return np.concatenate([np.sin(X @ Omega.T), np.cos(X @ Omega.T)], axis=1)   # trig features

def softmax_kernel_estimate(q, k, Omega):
    raise NotImplementedError

Tests, verified solutions, and AI help, open once you add your API key.

A Taylor feature map

The softmax kernel's trouble is its infinite series, so cut the series off. Keep terms up to second order, $e^s\approx 1 + s + s^2/2$ with $s = \mathbf{q}^\top\mathbf{k}$ . Each term factors exactly into features. The constant is $1\cdot 1$ , and the linear term is $\mathbf{q}^\top\mathbf{k}$ . The square is the dot product of the two outer products,

(\mathbf{q}^\top\mathbf{k})^2 = \sum_{i,j}q_iq_jk_ik_j = \operatorname{vec}(\mathbf{q}\mathbf{q}^\top)^\top\operatorname{vec}(\mathbf{k}\mathbf{k}^\top),

where $\operatorname{vec}$ lists a matrix's entries as one long vector. Take the pair from the start of the section, $\mathbf{q} = (1, 2)$ and $\mathbf{k} = (3, 1)$ , so $s = 5$ and $s^2 = 25$ . The outer products list as $(1, 2, 2, 4)$ and $(9, 3, 3, 1)$ , whose dot product is $9 + 6 + 6 + 4 = 25$ . This is the three-feature map from that example with its two equal middle entries kept apart.

∴ Taylor features

For $\mathbf{x}\in\R^{d}$ let $\phi(\mathbf{x}) = \big(1,\ \mathbf{x},\ \operatorname{vec}(\mathbf{x}\mathbf{x}^\top)/\sqrt2\big)\in\R^{1 + d + d^2}$ . Then for all $\mathbf{q},\mathbf{k}$ , with $s = \mathbf{q}^\top\mathbf{k}$ ,

\phi(\mathbf{q})^\top\phi(\mathbf{k}) = 1 + s + \frac{s^2}{2} = \frac{(s+1)^2 + 1}{2} \;\ge\; \frac12 .

Proof. The three blocks contribute $1$ , $\mathbf{q}^\top\mathbf{k}$ , and $\frac12\sum_{i,j}q_iq_jk_ik_j = \frac12\big(\sum_iq_ik_i\big)^2$ . Completing the square shows the sum is at least $\frac12$ , with equality at $s = -1$ . ∎

In plain words, a second-order Taylor expansion of the exponential is itself an exact kernel with $1 + d + d^2$ features. It is always positive, so it can sit in the denominator safely.

The positivity comes with two caveats. First, the kernel is smallest at $s = -1$ and grows again below it. So a key with score $-3$ gets kernel value $2.5$ , while an orthogonal key, with score 0, gets $1$ . Softmax always gives more weight to a higher score, and this kernel does not. Second, the kernel grows only quadratically. So the ratio of two weights is at most polynomial in the scores, while softmax ratios are exponential.

Even so, Zhang et al. (2024) found that this map keeps two properties that make softmax work. Its weights have low entropy, meaning they concentrate on a few keys. And they rise with the score over the range of scores that trained models produce. Arora et al. (2024) made it the core of a model called Based.

The cost is the feature count. At $d = 128$ it would be $1 + 128 + 16{,}384 = 16{,}513$ features and a state of over two million numbers per head. Based therefore projects queries and keys down to $d = 16$ first, giving $d_\phi = 1 + 16 + 256 = 273$ . It pairs this with a small sliding-window softmax attention, in which each token attends with exact softmax to only the previous 64 or 128 tokens (§12.6). The window handles precise comparisons with nearby tokens.

Core

Taylor features and their cost

Implement taylor_features(X) for X of shape (n, d). It returns shape (n, 1 + d + d*d) in the order of the theorem, so that $\phi(\mathbf{q})^\top\phi(\mathbf{k}) = 1 + s + s^2/2$ exactly. The order is the constant, then the $d$ linear features, then the $d^2$ entries of $\operatorname{vec}(\mathbf{x}\mathbf{x}^\top)/\sqrt2$ in row-major order.

Then implement two memory counts, both returning Python ints.

taylor_state_bytes(d_head, d_v, n_heads, bytes_per=4) is the bytes of one layer's decode state. Here d_head is the width $d$ of the queries and keys that enter the feature map (16 in Based, after the projection), and d_v is the width of the values. Count both $\mathbf{S}_t\in\R^{d_\phi\times d_v}$ and $\mathbf{z}_t\in\R^{d_\phi}$ for every head, with $d_\phi = 1 + d + d^2$ .
kv_cache_tokens(state_bytes, n_kv, d_k, bytes_per=2) is the smallest number of tokens $T$ at which one layer's softmax KV cache takes at least state_bytes. The cache holds $2\,n_{kv}d_k$ numbers per token, a key and a value for each of its $n_{kv}$ key-value heads.

import numpy as np

def taylor_features(X):
    return np.concatenate([np.ones((len(X), 1)), X], axis=1)     # first order only

def taylor_state_bytes(d_head, d_v, n_heads, bytes_per=4):
    raise NotImplementedError

Tests, verified solutions, and AI help, open once you add your API key.

Learned feature maps

If no fixed map matches the softmax, learn one that does. Hedgehog (Zhang et al. 2024) makes $\phi$ a single trainable layer with an exponential activation,

\phi(\mathbf{x}) = \exp(\mathbf{W}^\top\mathbf{x} + \mathbf{b}) = \big(e^{\mathbf{w}_1^\top\mathbf{x} + b_1},\dots,e^{\mathbf{w}_{d'}^\top\mathbf{x} + b_{d'}}\big).

Here $\mathbf{W}\in\mathbb{R}^{d_k\times d'}$ is a weight matrix whose columns are the $\mathbf{w}_i$ , and $\mathbf{b}\in\mathbb{R}^{d'}$ holds one bias per feature. In practice $d' = d_k$ , with $\mathbf{W}$ starting as the identity and $\mathbf{b}$ as zero. The map also appends the mirrored features $e^{-(\mathbf{w}_i^\top\mathbf{x} + b_i)}$ , so $d_\phi = 2d_k$ and the cost stays $O(Td_k^2)$ . Each head of each layer gets its own map, shared by its queries and keys.

How the map is trained depends on the starting point. In a model trained from scratch, it is learned with everything else from the next-token loss. To convert a trained Transformer, Hedgehog first freezes the model and trains only the maps, so that their attention weights match the softmax weights the same queries and keys produce. The loss is the cross-entropy between the two sets of weights. Then it fine-tunes the model with its usual loss. The paper reports that Hedgehog recovers over 99% of Transformer quality, both from scratch and when converting a fine-tuned model. Converting a pretrained 125-million-parameter GPT-2 gave a perplexity of 16.7 on WikiText-103 (perplexity is the exponential of the average next-token loss, §3.1).

LoLCATs (Zhang et al. 2024) scaled this conversion to Llama 3.1 70B and 405B in two steps. First it trains each linear attention to reproduce the output of the softmax layer it replaces, under a mean squared error, a step it calls attention transfer. Then it repairs the remaining error with LoRA, the low-rank fine-tuning of §14.5. The two steps together use as few as 40 million tokens. It also keeps a short sliding window of exact attention in every layer. Converting a pretrained model this way costs a tiny fraction of pretraining a new one, which is why it comes up in design interviews.

The widget compares softmax with the fixed maps above (elu+1, ReLU, random features, and Taylor) on one query that has one clearly relevant key among 32. Choose a feature map with the buttons. The head width is $d_k = 16$ . The sharpness slider multiplies the query, which is what lowering a softmax temperature does. For random features a second slider sets $d_\phi$ .

ReLU does not move at all, because it is homogeneous and sharpness only rescales the query. The elu+1 weight creeps up and stops near uniform, at the limit derived above. Taylor rises with sharpness but levels off far below softmax. In the limit its weights become proportional to $s_j^2$ , the squared scores, which do not concentrate on one key. Random features come closest at mild sharpness and the largest $d_\phi$ . A larger $d_k$ makes them worse, because the scaled query and key grow longer. They fall apart as sharpness grows, which is exactly where softmax concentrates its weight. We drew ten sets of features at $d_k = 16$ , sharpness 1, and $d_\phi = 512$ . They misplaced between 16 and 45 percent of the total weight, about 22 percent in a typical draw (the weight misplaced, under the chart). Their kernel values were off by 37 to 87 percent on average. Here is the comparison in one table.

feature map	$d_\phi$	positive	approximates softmax	used in
identity	$d_k$	no	no	layers with no denominator and an output norm (RetNet)
$\operatorname{elu}+1$	$d_k$	yes	no	Katharopoulos et al. (2020)
ReLU	$d_k$	nonnegative (can be all zero)	no	several early variants
positive random features	any, often 256	yes	unbiased, variance $\propto e^{\lVert\mathbf{q}+\mathbf{k}\rVert^2}/d_\phi$	Performer
Taylor, 2nd order	$1 + d + d^2$	yes	to second order	Based ( $d = 16$ )
learned exponential	chosen	yes	trained to mimic	Hedgehog, LoLCATs
SiLU, no denominator	$d_k$	no	no	MiniMax-01's lightning attention, Mamba-2 (§21.6)

The state is an associative memory

So far $\mathbf{S}_t$ has been bookkeeping. It is more useful to see it as a memory that stores key-value pairs and looks them up by content, which is called an associative memory. To see it, drop the denominator and use the identity feature map with unit-length keys. Writing a pair adds the outer product $\mathbf{k}\mathbf{v}^\top$ to the matrix. Reading with a key $\mathbf{k}_j$ returns $\mathbf{S}^\top\mathbf{k}_j$ .

✏️ Three pairs in a two-dimensional memory

Write two pairs with orthogonal unit keys, $\mathbf{k}_1 = (1, 0)$ with $\mathbf{v}_1 = (1, 2)$ , and $\mathbf{k}_2 = (0, 1)$ with $\mathbf{v}_2 = (4, 0)$ . The state is $\mathbf{S} = \mathbf{k}_1\mathbf{v}_1^\top + \mathbf{k}_2\mathbf{v}_2^\top = \begin{pmatrix}1 & 2\\ 4 & 0\end{pmatrix}$ . Reading with $\mathbf{k}_1$ gives $\mathbf{S}^\top\mathbf{k}_1 = (1, 2) = \mathbf{v}_1$ , and reading with $\mathbf{k}_2$ gives $(4, 0) = \mathbf{v}_2$ . Both come back exactly.

Now write a third pair, $\mathbf{k}_3 = (0.6, 0.8)$ with $\mathbf{v}_3 = (0, 5)$ . Its outer product is $\begin{pmatrix}0 & 3\\ 0 & 4\end{pmatrix}$ , so $\mathbf{S} = \begin{pmatrix}1 & 5\\ 4 & 4\end{pmatrix}$ . Reading with $\mathbf{k}_1$ now gives $(1, 5)$ . That is $\mathbf{v}_1 + 0.6\,\mathbf{v}_3$ , the right value plus $0.6$ of the new one, because $\mathbf{k}_3^\top\mathbf{k}_1 = 0.6$ . Reading with $\mathbf{k}_3$ gives $(3.8, 6.2) = \mathbf{v}_3 + 0.6\,\mathbf{v}_1 + 0.8\,\mathbf{v}_2$ . A two-dimensional key space has room for two orthogonal keys, and the third pair leaks into everything.

In general, after $n$ writes the state is $\mathbf{S} = \sum_{s=1}^{n}\mathbf{k}_s\mathbf{v}_s^\top\in\R^{d_k\times d_v}$ , and reading with stored key $j$ gives

\mathbf{S}^\top\mathbf{k}_j = \sum_{s=1}^{n}\mathbf{v}_s\big(\mathbf{k}_s^\top\mathbf{k}_j\big) = \mathbf{v}_j\lVert\mathbf{k}_j\rVert^2 + \sum_{s\ne j}\big(\mathbf{k}_s^\top\mathbf{k}_j\big)\mathbf{v}_s = \mathbf{v}_j + \underbrace{\sum_{s\ne j}\big(\mathbf{k}_s^\top\mathbf{k}_j\big)\mathbf{v}_s}_{\text{crosstalk } \mathbf{c}_j} .

In plain words, what comes back is the stored value plus crosstalk. Crosstalk is interference from every other pair, each weighted by how similar its key is to the query.

∴ Retrieval from an outer-product memory

Let $\mathbf{k}_1,\dots,\mathbf{k}_n\in\R^{d_k}$ be unit vectors, $\mathbf{v}_1,\dots,\mathbf{v}_n\in\R^{d_v}$ , and $\mathbf{S} = \sum_s\mathbf{k}_s\mathbf{v}_s^\top$ .

If the keys are orthonormal, then $\mathbf{S}^\top\mathbf{k}_j = \mathbf{v}_j$ exactly for every $j$ . Orthonormal keys exist only for $n\le d_k$ .
No $d_k\times d_v$ matrix, however it is written, retrieves more than $d_k$ generic pairs exactly by a linear read. If $\mathbf{S}^\top\mathbf{k}_j = \mathbf{v}_j$ for all $j$ , then $\mathbf{K}\mathbf{S} = \mathbf{V}$ with $\mathbf{K}\in\R^{n\times d_k}$ . So the columns of $\mathbf{V}\in\R^{n\times d_v}$ lie in the column space of $\mathbf{K}$ , which has dimension at most $d_k$ . For $n > d_k$ this fails for values in general position.
Suppose the keys are independent and uniform on the unit sphere. Suppose the values are independent of the keys and of each other, with mean zero and $\mathbb{E}\lVert\mathbf{v}_s\rVert^2 = \sigma^2$ . Then $\mathbb{E}[(\mathbf{k}_s^\top\mathbf{k}_j)^2] = 1/d_k$ for $s\ne j$ , and

\mathbb{E}\,\lVert\mathbf{c}_j\rVert^2 = \frac{n-1}{d_k}\,\sigma^2 .

Proof. Part 1 is the display above with $\mathbf{k}_s^\top\mathbf{k}_j = 0$ for $s\ne j$ , and $\R^{d_k}$ holds at most $d_k$ mutually orthogonal unit vectors. Part 2 is proved in its statement. For part 3, fix $\mathbf{k}_j$ . Since $\mathbf{k}_s$ is uniform on the sphere, by symmetry $\mathbb{E}[\mathbf{k}_s\mathbf{k}_s^\top]$ is a multiple of the identity. Its trace is $\mathbb{E}\lVert\mathbf{k}_s\rVert^2 = 1$ , so it is $\mathbf{I}/d_k$ . Then $\mathbb{E}[(\mathbf{k}_s^\top\mathbf{k}_j)^2] = \mathbf{k}_j^\top\mathbb{E}[\mathbf{k}_s\mathbf{k}_s^\top]\mathbf{k}_j = 1/d_k$ . Expand $\lVert\mathbf{c}_j\rVert^2 = \sum_{s,r\ne j}(\mathbf{k}_s^\top\mathbf{k}_j)(\mathbf{k}_r^\top\mathbf{k}_j)\mathbf{v}_s^\top\mathbf{v}_r$ . The terms with $s\ne r$ have expectation zero, because $\mathbb{E}[\mathbf{v}_s^\top\mathbf{v}_r] = 0$ for independent zero-mean values. Each of the $n-1$ terms with $s = r$ has expectation $\sigma^2/d_k$ . ∎

In plain words, a $d_k\times d_v$ memory holds $d_k$ pairs perfectly if their keys are orthogonal, and no linear read can do better. With random keys the noise grows steadily. Its energy is $(n-1)/d_k$ times the signal's, so the signal-to-noise ratio is about $d_k/(n-1)$ . In high dimension the crosstalk is nearly orthogonal to $\mathbf{v}_j$ . So the retrieved vector has cosine about $1/\sqrt{1 + (n-1)/d_k}$ with the right value, about 0.71 at $n = d_k$ .

Causal linear attention writes one pair per token and never erases anything. A layer with $d_\phi = 128$ that has read 100,000 tokens is a memory loaded nearly eight hundred times past its exact capacity. With a feature map the same argument holds with $d_\phi$ in place of $d_k$ . That is one honest reason to want a large $d_\phi$ , as with Taylor features. Schlag et al. (2021) took another route. Their feature map, called DPFP, raises $d_\phi$ to a multiple of $d_k$ without random sampling, and its features are sparse, so different keys tend to be orthogonal after the map. The cell measures retrieval for orthonormal and for random keys.

Python · cell 4

Orthonormal keys retrieve perfectly up to $n = d$ and cannot exist beyond it, so that column shows nan. Random keys lose fidelity smoothly. The measured cosine tracks $1/\sqrt{1 + (n-1)/d}$ , and the measured crosstalk energy tracks $(n-1)/d$ . An interviewer who asks how much a linear-attention layer can remember wants these two facts, exact recall up to $d_\phi$ pairs and crosstalk that grows as $(n-1)/d_\phi$ beyond it.

The widget lets you watch a memory of width $d = 32$ fill one write at a time. Choose the kind of keys, orthonormal, random, or correlated (all sharing one direction). Then press Write one by one, or drag the number of pairs $n$ .

Interview

A matrix state as an associative memory

Implement an outer-product memory with three functions.

write(K, V) takes keys of shape (n, d_k) and values of shape (n, d_v). It normalizes each key to unit length and returns $\mathbf{S} = \sum_s\hat{\mathbf{k}}_s\mathbf{v}_s^\top$ of shape (d_k, d_v).
read(S, q) normalizes the query of shape (d_k,) and returns $\mathbf{S}^\top\hat{\mathbf{q}}$ of shape (d_v,).
crosstalk_energy(K) normalizes the keys and returns the mean over $j$ of $\sum_{s\ne j}(\hat{\mathbf{k}}_s^\top\hat{\mathbf{k}}_j)^2$ . The theorem predicts about $(n-1)/d_k$ for random keys.

import numpy as np

def write(K, V):
    return K.T @ V

def read(S, q):
    return S @ q

Tests, verified solutions, and AI help, open once you add your API key.

A fixed state against a growing cache, at real scale

How big is the state of a real linear-attention layer, and when does it beat a KV cache? The answer decides how many sequences fit on one GPU, so it is a standard question in a systems interview.

✏️ When the state becomes smaller than the cache

Take a layer with 32 heads and $d_\phi = d_k = d_v = 128$ . Its state is $32\cdot128\cdot128 = 524{,}288$ numbers. The normalizer adds $32\cdot128$ more, which is negligible, and modern layers drop it anyway.

Compare a softmax layer with grouped-query attention, in which several query heads share one key head and one value head (§12.5). With $n_{kv} = 8$ key-value heads of width 128, it stores $2\cdot8\cdot128 = 2{,}048$ numbers per token. The two are equal at $T = 524{,}288/2{,}048 = 256$ tokens.

Now count bytes. The state is kept in float32, because the gap between neighboring numbers grows with their size, so in a 16-bit format small updates to a large running sum round away (§21.5 measures this rounding for cumulative sums). The cache is kept in bf16, the 16-bit format with 2 bytes per number. The state is 2 MiB and the cache 4 KiB per token, so they are equal at 512 tokens. Past that the cache keeps growing and the state does not.

📌 A published instance, MiniMax-01

MiniMax-01 (2025) is a 456-billion-parameter mixture-of-experts model, one that routes each token through a few of many feed-forward networks (§13.1), with 45.9 billion parameters active per token. According to its paper and released configuration, it has 80 layers. One softmax-attention layer follows every seven linear-attention layers, so 70 are linear and 10 softmax. Each layer has 64 heads of width 128, and the softmax layers have 8 KV heads.

Its linear layers hold $70\cdot64\cdot128\cdot128 = 73{,}400{,}320$ state numbers, 280 MiB in float32, at any context length. Its softmax layers hold $10\cdot2\cdot8\cdot128\cdot2 = 40{,}960$ bytes per token in bf16, 40 GiB at $2^{20}$ tokens. Had all 80 layers been softmax attention, the cache would be 320 GiB at that length. The model trained on contexts up to one million tokens.

The cell checks these numbers.

Python · cell 5

Speed follows memory. Katharopoulos et al. (2020) generated CIFAR-10 images, sequences of 3,072 values, 4,462 times faster than a softmax Transformer that recomputed attention over the whole prefix at every step. The honest comparison is a softmax Transformer that caches its keys and values, and against that it was still more than 50 times faster (their appendix). Based (Arora et al. 2024) reached 24 times the generation throughput of FlashAttention-2 when generating 1,024 tokens with 1.3-billion-parameter models. On recall-intensive tasks, where fixed states struggle most, it beat Mamba and other sub-quadratic models of its size by 10.36 accuracy points on average in the revised paper. The first version, whose abstract arXiv still lists, reported 6.22.

The MiniMax paper also gives the clearest statement of the cost. In its scaling experiments, a pure linear-attention model matched softmax attention on most downstream tasks. It fell short on needle-in-a-haystack retrieval, finding one planted fact in a long context. The hybrid with one softmax layer in eight matched or beat softmax attention on retrieval. The next part of this section explains that result.

Why linear attention lags softmax attention

For years after 2020, linear attention trained from scratch was consistently worse than softmax attention at the same size. Four mechanisms account for the gap, and each one points to a fix that the rest of this chapter develops.

It is low rank. The kernel matrix $\phi(\mathbf{Q})\phi(\mathbf{K})^\top$ is a $T\times d_\phi$ matrix times a $d_\phi\times T$ matrix, so its rank is at most $d_\phi$ , however long the sequence. Masking with $\mathbf{L}$ raises the rank of the whole matrix. But every block that lies entirely below the diagonal is still a block of $\phi(\mathbf{Q})\phi(\mathbf{K})^\top$ , of rank at most $d_\phi$ . The softmax matrix $\exp(\mathbf{Q}\mathbf{K}^\top/\sqrt{d_k})$ applies the exponential entry by entry, which destroys low rank, and is typically full rank. A layer whose blocks below the diagonal have low rank is a recurrence with a small state, because the rank is the amount of information the past passes to the future. §21.6 makes that equivalence exact, through what are called semiseparable matrices.

Python · cell 6

The softmax kernel matrix has full rank 256, and its lower-left block full rank 128. The elu+1 matrix and its lower-left block both have rank 16, which is $d_\phi = d_k$ , whatever $T$ is.

It cannot be sharp. Softmax turns a score gap $\Delta$ into a weight ratio $e^\Delta$ , so a query can put nearly all its weight on one key. Homogeneous maps cannot sharpen at all, elu+1 adds a floor that flattens everything, and Taylor features give only polynomial ratios. Zhang et al. (2024) measured two differences. Linear attentions trained from scratch had higher-entropy weights than softmax, and their weights did not rise steadily with the query-key dot product. Restoring these two properties closed most of the quality gap. The widget above shows the flatness directly.

Its memory is finite and it never forgets. By the memory theorem, a $d_\phi\times d_v$ state retrieves at most $d_\phi$ pairs exactly. With random keys its retrieval cosine falls as $1/\sqrt{1 + (n-1)/d_\phi}$ . Causal linear attention adds one pair per token with equal weight forever, so old and irrelevant pairs crowd out new ones. Qin et al. (2022) named a related symptom attention dilution, weights spread evenly over long sequences instead of concentrating on nearby tokens.

Its normalizer is unstable. The denominator $\mathbf{z}_t^\top\phi(\mathbf{q}_t)$ grows roughly linearly in $t$ , so each new token's share of the weight shrinks like $1/t$ . The gradients are the deeper problem. Write the weight of key $s$ as $p_{ts} = \kappa_{ts}/\sum_r\kappa_{tr}$ , with $\kappa_{ts} = \phi(\mathbf{q}_t)^\top\phi(\mathbf{k}_s)$ . Its derivative with respect to the kernel value $\kappa_{tr}$ is $(\delta_{sr} - p_{ts})/\sum_{r'}\kappa_{tr'}$ , where $\delta_{sr}$ is 1 if $s = r$ and 0 otherwise. So the derivative can be as large as one over the denominator. Nothing keeps the denominator away from zero, so the gradient has no upper bound. Softmax weights, as functions of the scores, have derivatives of at most $1/4$ .

Qin et al. (2022) proved this unboundedness and measured its effect. Over training, the relative standard deviation of the gradients was 0.58 for elu+1 and 0.47 for Performer, against 0.25 for softmax. Their model, TransNormer, deletes the denominator and normalizes the output instead, $\mathbf{O} = \operatorname{RMSNorm}(\phi(\mathbf{Q})(\phi(\mathbf{K})^\top\mathbf{V}))$ . RMSNorm divides each output vector by its root mean square and multiplies it by learned gains (§5.2). That brought the number down to 0.20.

Taming the denominator is now standard. MiniMax-01's linear-attention layer, which it calls lightning attention, builds on TransNormer. RetNet, GLA, Mamba-2, and Gated DeltaNet also normalize the output, and none of them depends on the division by $\mathbf{z}_t^\top\phi(\mathbf{q}_t)$ . (RetNet does divide each row of its scores by the absolute value of their sum, clamped below at 1, but only for numerical precision. The normalization after it cancels any such scale.) So these layers no longer need $\phi$ to be positive. MiniMax-01's released code, for example, applies SiLU, $x\,\sigma(x)$ with $\sigma$ the logistic sigmoid (§5.3), to queries, keys, and values. It keeps the state in float32 and multiplies it by a fixed decay per head at every step. It puts an RMSNorm and a sigmoid output gate after the read. A few designs keep a denominator but bound it below, dividing by $\max(\lvert\mathbf{z}_t^\top\mathbf{q}_t\rvert, 1)$ , as the deeper box below shows.

The decay is the other half of the modern recipe. The rest of the chapter adds the two missing abilities in turn. First the state learns to forget, through principled continuous-time decays (§21.3), then through decays chosen by the input (§21.4) and gates (§21.6). Then it learns to correct itself, through the delta rule of §21.6, which overwrites the value at a key instead of adding crosstalk.

🔬 An LSTM whose cell is linear attention

The LSTM of §3.3 keeps a cell that a forget gate scales and an input gate writes to. The matrix LSTM of xLSTM (Beck et al. 2024) makes that cell a matrix and writes key-value outer products into it. With scalar gates $f_t$ and $i_t$ computed from the input,

\mathbf{C}_t = f_t\mathbf{C}_{t-1} + i_t\mathbf{v}_t\mathbf{k}_t^\top, \qquad \mathbf{n}_t = f_t\mathbf{n}_{t-1} + i_t\mathbf{k}_t, \qquad \tilde{\mathbf{h}}_t = \frac{\mathbf{C}_t\mathbf{q}_t}{\max\big(\lvert\mathbf{n}_t^\top\mathbf{q}_t\rvert,\ 1\big)} .

With $f_t = i_t = 1$ and the identity feature map this is causal linear attention, apart from the floor in the denominator. The state is stored transposed, $\mathbf{C}_t = \mathbf{S}_t^\top$ , and $\mathbf{n}_t$ is the normalizer $\mathbf{z}_t$ . The gates add the forgetting that plain linear attention lacks. The floor of 1 in the denominator, which the paper takes from RetNet, keeps the division from blowing up. An output gate then scales $\tilde{\mathbf{h}}_t$ . Because no gate reads the previous hidden state, the update is linear in $\mathbf{C}_{t-1}$ , and the layer trains in parallel like every layer in this chapter.

layer	training cost	decode cost per token	decode memory	exact softmax	recall
softmax attention	$O(T^2d)$	$O(Td)$	KV cache, $O(Td)$	yes	exact, the reference
linear attention, elu+1	$O(Td^2)$	$O(d^2)$	state, $O(d^2)$	no	weak, about $d$ pairs
Performer	$O(Td_\phi d)$	$O(d_\phi d)$	$O(d_\phi d)$	unbiased estimate	weak in practice
Based (Taylor plus a 64 to 128 token window)	$O(Td_\phi d_v)$ plus window	constant	constant	to second order	strong for its state size
normalized linear attention with decay (lightning, RetNet)	$O(Td^2)$ , chunked	$O(d^2)$	$O(d^2)$	no	forgets gracefully, still finite
hybrid, one softmax layer in four to eight (§21.6)	mixed	grows only in softmax layers	small cache	in some layers	close to softmax

In the Based row, $d_\phi$ is the projected Taylor dimension, 273 in the paper.

In PyTorch

There is no linear attention in core PyTorch. Production kernels live in libraries such as flash-linear-attention, which use the chunkwise algorithm of §21.5. The block below writes causal linear attention three ways with plain tensor operations. The first is the masked quadratic form. The second is the cumulative-sum form, which stores every $\mathbf{S}_t$ , so the training-memory problem above shows up as the size of one tensor. The third is the recurrent decode loop. Then the block checks that autograd gives the same gradients through the quadratic form and the loop.

🔥 In PyTorchfor a machine with PyTorch

⚠️ Where linear attention implementations go wrong

Reading before writing leaves token $t$ out of its own output, and $\mathbf{o}_1$ becomes $0/0$ .
Storing $\mathbf{S}_t$ transposed works only when $d_\phi = d_v$ , so test with unequal sizes.
Dividing by a sum over the feature axis instead of over keys breaks the check that constant values give constant outputs.
A feature map with negative entries (identity, SiLU, trigonometric random features) together with the denominator divides by numbers near zero or below it. Either use a positive $\phi$ with the denominator, or drop the denominator and normalize the output.
The cumulative-sum form under autograd stores $T d_\phi d_v$ numbers per head, enough to run out of memory at a few thousand tokens.
Accumulating the state in bf16 loses small updates once the state is large, which is why production kernels keep it in float32.
When converting a pretrained model, dropping the $1/\sqrt{d_k}$ temperature changes the kernel being approximated. For random features the temperature is the $d_k^{-1/4}$ scaling of queries and keys.

🧭 What to remember

Attention is a weighted average with a similarity kernel. When the kernel factors as $\phi(\mathbf{q})^\top\phi(\mathbf{k})$ , associativity turns $(\phi(\mathbf{Q})\phi(\mathbf{K})^\top)\mathbf{V}$ , costing $2T^2(d_\phi + d_v)$ , into $\phi(\mathbf{Q})(\phi(\mathbf{K})^\top\mathbf{V})$ , costing $4Td_\phi d_v$ . The second is cheaper for $T > 2d_\phi d_v/(d_\phi + d_v)$ .
Causal linear attention is a recurrence, $\mathbf{S}_t = \mathbf{S}_{t-1} + \phi(\mathbf{k}_t)\mathbf{v}_t^\top$ , $\mathbf{z}_t = \mathbf{z}_{t-1} + \phi(\mathbf{k}_t)$ , $\mathbf{o}_t = \mathbf{S}_t^\top\phi(\mathbf{q}_t)/\mathbf{z}_t^\top\phi(\mathbf{q}_t)$ . It equals the masked quadratic form $((\phi(\mathbf{Q})\phi(\mathbf{K})^\top)\odot\mathbf{L})\mathbf{V}$ divided by row sums, and decoding costs $O(d_\phi d_v)$ per head in constant memory.
Positive random features estimate $e^{\mathbf{q}^\top\mathbf{k}}$ without bias, with relative variance $(e^{\lVert\mathbf{q}+\mathbf{k}\rVert^2} - 1)/d_\phi$ , which is too large at realistic head sizes. Taylor features give $1 + s + s^2/2$ exactly with $1 + d + d^2$ features. Learned maps (Hedgehog, LoLCATs) mimic softmax and can convert pretrained models.
The state is an outer-product associative memory. It retrieves at most $d_\phi$ pairs exactly, and with random keys the crosstalk energy is $(n-1)/d_\phi$ of the signal.
Linear attention lags softmax because it is low rank, cannot be sharp, has a finite memory that never forgets, and has an unstable normalizer. Modern layers drop the normalizer, normalize the output, and add a decay.
At real scale the state of 32 heads of 128 equals a grouped-query cache at 256 tokens and then stops growing. MiniMax-01 keeps 70 of its 80 layers linear, with one softmax layer in eight for retrieval.
Next, §21.3 gives the state a principled way to forget, a continuous-time linear system whose decay rates set how long it remembers.

Katharopoulos et al. (2020). Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. The kernel view, the elu+1 map, the recurrent form, and the cumulative-sum gradients. Read §3 for the derivation and Appendix C for the comparison with cached softmax.
Choromanski et al. (2020). Rethinking Attention with Performers. Positive orthogonal random features (FAVOR+). Lemma 1 is the unbiasedness result, Lemma 2 the variance comparison with trigonometric features, and Theorem 2 the gain from orthogonality.
Schlag, Irie & Schmidhuber (2021). Linear Transformers Are Secretly Fast Weight Programmers. The fast-weight view, the capacity limit, the DPFP feature map, and the delta-rule update that §21.6 develops.
Qin et al. (2022). The Devil in Linear Transformer. Why the normalizer gives unbounded gradients, attention dilution, and the normalized replacement that most modern layers use.
Zhang et al. (2024). The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry. Spikiness and monotonicity as the missing properties, and learned feature maps that restore them.
Arora et al. (2024). Simple linear attention language models balance the recall-throughput tradeoff. Based, Taylor features with a sliding window, the trade-off between state size and recall, and IO-aware kernels.
Zhang et al. (2024). LoLCATs: On Low-Rank Linearizing of Large Language Models. Attention transfer plus LoRA, converting Llama 3 models up to 405B.
MiniMax (2025). MiniMax-01: Scaling Foundation Models with Lightning Attention. Linear attention at frontier scale, its retrieval weakness, and the one-in-eight softmax hybrid.
Tsai et al. (2019). Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel. Attention read as a kernel smoother, the starting point of this section.
Rahimi & Recht (2007). Random Features for Large-Scale Kernel Machines. The random Fourier features that positive random features improve on.

Made for larger screens

Linear attention

Attention is a weighted average with a similarity score

Moving the parentheses

Which order to multiply

Causal linear attention is a recurrent network

The same numbers, computed in parallel

Decoding, training, and what each costs

Causal linear attention, recurrent and parallel, plus the decode step

Choosing the feature map

Simple positive maps

Random features that estimate the softmax kernel

Positive random features for the softmax kernel

A Taylor feature map

Taylor features and their cost

Learned feature maps

The state is an associative memory

A matrix state as an associative memory

A fixed state against a growing cache, at real scale

Why linear attention lags softmax attention

In PyTorch

Read more

Watch next