Attention is recurrent

Most introductions to attention, define it as

But this is a bit backwards, it's only the parallel form, used for prefill and training. And it makes the "causal mask" seem like a fundamental thing. But if you rather look from the recursion perspective, the casualness comes automatically, because when you auto-regressively generate a sequence, you can simply not use future tokens, as they do not yet exist. The parallel form is hugely important though, and it is what makes modern inference systems reasonable. Enabling LLMs to re-use shared prefixes efficiently, allowing systems to have huge system prompts and making back and forth conversation efficient. To get an understanding of how different they are the parallel form and decoding mode are, you can see 1000+ of tokens per second in prefill on macbooks with the M-chip, while the decoding part is stuck at 30 TPS, being dependent on the HBM/VRAM speed.¹

When we generate tokens we have a prompt $x_{\leq n}=(x_0, x_1, ..., x_n)$ and we wish to compute $x_{n+1}$ and so on. I.e $p\left(x_{n+1} | x_{\leq n}\right)$.

We have regular attention defined as:

And we have the generation process from prompt $x=(x_0,)$: $$x_0\longrightarrow o_0, h_0 \longrightarrow x_1\longrightarrow o_1,h_1 \longrightarrow x_2 \longrightarrow \cdots$$

With these definitions it turns out that parallel form, i.e how to compute $o_n$ directly for a prompt $x_{\leq n}=(x_0, x_1, ..., x_n)$ is embarrassingly parallel and we can compute all $o_0, o_1, ..., o_n$ simultaniously ( no need to first compute $o_0, o_1, ...$ in a sequential matter), allowing you to simple change $q$ for $Q$, and we suddenly go from GEMV to GEMMs and we see the matrix units go to work.

For recursive decoding, to generate $o_n$ they did not have to recalculate $k_{\leq n-1}$ and $v_{\leq n-1}$ as they were already calculated when $o_{n-1}$ was generated, i.e using the identities:²

Linear attention

Why? You can deepen your Transformers without any big compute increase by interleaving linear layers. It can be used as positional encoding, so the full attention layers can get rid of RoPE. You can scale the context lengths without hitting the quadratic time complexity wall.

Now let's remove the softmax, we can collapse the history into a single time independent state $S\in \mathbb{R}^{d \times d}$.³ Since now we are no longer just keeping the information from all previous tokens, but compressing them into a fixed-sized state, there are more engineering tuning/hacks to improve the compression, just look at the complexity of LSTMs and GRUs...

Starting from the auto-regressive form

we get

This is the key part, as if we define the summation as

we can be smart when calculating the attention for the $n$ th token, $o_n$, we don't have to recalculate the whole sum if we keep track of the last state $S_{n-1}$:

So when we want to calculate $o_n$ in linear attention, instead of keeping $(k_0,v_0),(k_1,v_1),\ldots,(k_n,v_n)$ i.e $(K_{\leq n}, V_{\leq n})$ (the KV-cache), we only need to keep track of the previous state $S_{n-1}$. We can say that for linear transformer $o_n=o_n(S_{n-1}, q_n, k_n, v_n)$, while regular attention $o_n=o_n(K_{\leq n}, V_{\leq n})$. The former arguments are all fixed-size, while the latter arguments grow linearly with $n$.

So generation (and at arrow between $i-1$ and $i$, the previous state $S_{i-1}$ is cached) now looks like

with

Prefill

Now suppose our known prompt is

Sequentially we would compute

Then

The last output

is what eventually gives us

Again, $x_4$ is not part of the prefill. We prefill $x_0,\ldots,x_3$, and the representation at the last known position predicts $x_4$.

Unlike regular attention, there is now an actual recurrence

But luckily the recurrence is just a cumulative sum:

And addition is associative, s o all the prefix states $S_0,S_1,S_2,S_3$ can be computed with a parallel prefix scan. This is the first place where the difference to regular attention becomes interesting.

For regular attention there wasn't really a state recurrence at all. Once $Q,K,V$ were known, every $o_i$ could be computed independently, so prefill was basically just stacking the autoregressive equations into a GEMM.

For linear attention there is a recurrent state, but the recurrence happens to be one of the easiest possible ones:

So plain linear attention is just the simplest associative scan (a cumsum).

Note on non-causal simple linear attention

if we just remove the softmax from non-causal attention we have

which we can regroup as $O=Q(K^TV)$ which complexity is linear in $n$. (https://en.wikipedia.org/wiki/Matrix_chain_multiplication) But when we have a causal pre-fill sequence, we must now use multiplicative mask $L$ which has entries in $\{0, 1\}$ rather than an additive $M$ with $\{-\infty, 0\}$ entries get zerod by the softmax,

And we can sadly not do this re-arrangment as $\odot$ is elementwise multiplication.⁴

Gated linear attention

Now we can make the state forget with parameter $\alpha_n$, so instead of $S_n=S_{n-1}+k_n^Tv_n$ we have

Autoregressively this is still trivial:

If we unroll it,

So an old write at position $i$ survives until position $n$ with weight

This is no longer the simplest assocative scan, as it's now a weighted scan.

Delta attention

Instead of blindly adding the new value, we first ask: what does the current memory already predict for this key? The current prediction is

So rather than writing $v_n$, we write only the error

With a write strength $\beta_n$,

This is still a simple autoregressive update. Given $S_{n-1}$ and the new token $x_n$, compute

then

and finally

Here $u_n$ is the actual corrected value we write into memory. And now something important has changed. For ordinary linear attention,

only depends on token $n$. But for delta attention,

and

depends on the entire previous state. So the writes themselves are now causally coupled. This is why prefill stops being just a cumsum. What happens during prefill?

Prefill

Take again

Assume $S_{-1}=0$.

Then

and

For token $1$,

For token $2$,

And for token $3$,

So unlike normal attention, we cannot just say "compute every row independently".

And unlike vanilla linear attention, we cannot just cumsum independent writes either. The write $u_3$ depends on $u_2$, which depends on $u_1$, which depends on $u_0$. At first this looks inherently sequential. But notice that these dependencies are all linear. Move the previous writes to the left:

and so on.

So for the whole prompt we get one lower-triangular system

where

Therefore

And this is where the inverse in the parallel/chunkwise GDN/KDA formulas comes from.

It is just what happens when we take a causal recurrence

and ask:

instead of solving $u_0$, then $u_1$, then $u_2$, can we solve all of them together?

The answer is a triangular solve.

So now the little hierarchy becomes