Linear attention

The problem with attention

Let’s start with a review of standard multi-head attention. For a single layer and head, we compute $\mathbf{Q}, \mathbf{K}, \mathbf{V} \in \mathbb{R}^{n_{\text{seq}} \times d_{\text{head}}}$:

$$\mathbf{Q} = \begin{bmatrix} — & \mathbf{q}_1^{\top} & — \\ & \vdots & \\ — & \mathbf{q}_{n_{\text{seq}}}^{\top} & — \\ \end{bmatrix}, \quad \mathbf{K} = \begin{bmatrix} — & \mathbf{k}_1^{\top} & — \\ & \vdots & \\ — & \mathbf{k}_{n_{\text{seq}}}^{\top} & — \\ \end{bmatrix}, \quad \mathbf{V} = \begin{bmatrix} — & \mathbf{v}_1^{\top} & — \\ & \vdots & \\ — & \mathbf{v}_{n_{\text{seq}}}^{\top} & — \\ \end{bmatrix}.$$

We can then compute the output as:

$$\mathbf{O} = \text{softmax}(\mathbf{Q}\mathbf{K}^{\top} \odot \mathbf{M} ) \mathbf{V},$$

where $\mathbf{M} \in \mathbb{R}^{n_{\text{seq}} \times n_{\text{seq}}}$ is the causal mask.¹

The big issue with vanilla attention is that it is quadratic in sequence length. The computational complexity (number of floating point operations) of the matrix multiplication $\mathbf{Q}\mathbf{K}^{\top}$ is $O(n_{\text{seq}}^2 \cdot d_{\text{head}})$, which is quadratic in the sequence length. And the memory complexity of storing $\mathbf{Q}\mathbf{K}^{\top}$ is $O(n_{\text{seq}}^2)$, which is also quadratic in the sequence length.² This becomes problematic when $n_{\text{seq}}$ is large (e.g., think about what happens when sequence length is 1 million tokens).

A refresher on the computational complexity of matrix multiplication.

Suppose we have two matrices $\mathbf{A} \in \mathbb{R}^{m \times n}$ and $\mathbf{B} \in \mathbb{R}^{n \times k}$.

The resulting matrix $\mathbf{C} = \mathbf{A}\mathbf{B} \in \mathbb{R}^{m \times k}$ has $m \cdot k$ entries. Each entry $c_{ij}$ can be computed as the dot product of the $i^{\text{th}}$ row of $\mathbf{A}$ and the $j^{\text{th}}$ column of $\mathbf{B}$, which consists of $n$ multiplications and $n-1$ additions.

Thus, the complexity of the matrix multiplication is $O(m \cdot n \cdot k)$ floating point operations.

Linear attention

Let’s consider the standard attention formulation,³ but without the softmax operation:

$$\mathbf{O} = (\mathbf{Q}\mathbf{K}^{\top})\mathbf{V}.$$

Without the softmax operation, attention becomes a linear operation – it’s just a bunch of matrix multiplications. This means that we take advantage of the associative property of matrix multiplication, and write:

$$\mathbf{O} = \mathbf{Q} (\mathbf{K}^{\top} \mathbf{V}).$$

With this order of operations, the complexity is no longer quadratic in sequence length. The matrix $\mathbf{K}^{\top} \mathbf{V} \in \mathbb{R}^{d_{\text{head}} \times d_{\text{head}}}$ can be computed using $O(d_{\text{head}}^2 \cdot n_{\text{seq}})$ FLOPs, and $\mathbf{Q} (\mathbf{K}^{\top} \mathbf{V}) \in \mathbb{R}^{n_{\text{seq}} \times d_{\text{head}}}$ can also be computed using $O(d_{\text{head}}^2 \cdot n_{\text{seq}})$ FLOPs. So the total complexity is $O(d_{\text{head}}^2 \cdot n_{\text{seq}})$, which is linear in the sequence length.

One thing we’ve conveniently ignored is the causal mask $\mathbf{M} \in \mathbb{R}^{n_{\text{seq}} \times n_{\text{seq}}}$. When we write linear attention[1]Transformers are RNNs: Fast autoregressive transformers with linear attention [link]
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Proceedings of the 37th International Conference on Machine Learning. 2020. with the causal mask, associativity no longer helps us, as we must apply an element-wise masking operation to the attention scores:

$$\mathbf{O} = (\mathbf{Q}\mathbf{K}^{\top} \odot \mathbf{M}) \mathbf{V}.$$

In this parallel view (i.e., when we compute the output for all token positions at once), the causal mask therefore prevents us from computing the output with linear dependence on the sequence length, and we are still stuck with quadratic dependence on the sequence length. We will see, though, that in the recurrent view (i.e., when we compute one token at a time), linear attention is very nice. Luckily, there is a way to balance between the two views via something called the chunkwise parallel form of linear attention, which we’ll describe later.

The recurrent view

As we’ve written it above, linear attention can be viewed as a parallel computation – the output is computed for all token positions at once via matrix multiplications.

But with linear attention, it’s useful to think about it as a recurrent computation – as some computation that is performed sequentially, one token position at a time (as in autoregressive inference). In this recurrent view, we think of the $\mathbf{K}^{\top} \mathbf{V} \in \mathbb{R}^{d_{\text{head}} \times d_{\text{head}}}$ as a memory matrix, that is built up as we read the sequence.

First, we note that $\mathbf{K}^{\top} \mathbf{V}$ can be written as a sum of outer products:

$$\mathbf{K}^{\top} \mathbf{V} = \sum_{i=1}^{n_{\text{seq}}} \mathbf{k}_i \mathbf{v}_i^{\top}.$$

With this view, we can think of the memory matrix being updated with a single key-value pair $(\mathbf{k}_i, \mathbf{v}_i)$ for each token position $i$.

Let’s now define $\mathbf{S}_t$ to be the memory matrix at time $t$ (e.g., up to and including token position $t$):

$$\mathbf{S}_t = \sum_{i=1}^{t} \mathbf{k}_i \mathbf{v}_i^{\top}.$$

It is easy to define a recurrence relation for $\mathbf{S}_t$:

$$\mathbf{S}_{t} = \mathbf{S}_{t-1} + \mathbf{k}_{t} \mathbf{v}_{t}^{\top}.$$

That’s basically all there is to it. With the recurrent view, we maintain a constant-size memory matrix $\mathbf{S}_t$ that gets updated with a new key-value pair at each time step. Note that, once the memory matrix has been built up, additional tokens can be generated with constant time complexity ($O(d_{\text{head}}^2)$) and constant memory complexity (also $O(d_{\text{head}}^2)$), without dependence on the sequence length.

Chunkwise parallel form of linear attention

As we’ve seen so far, when viewed as recurrent computation, linear attention is very nice; however, when viewed as parallel computation, we still have quadratic dependence on sequence length (due to the causal mask).

Chunkwise parallel form[2]Gated linear attention transformers with hardware-efficient training [link]
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Proceedings of the 41st International Conference on Machine Learning. 2024. is a way to balance between the two views, and to make parallel computation not entirely quadratic in sequence length.

The idea is to break up the sequence into chunks, each of length $n_{\text{chunk}}$. Within each chunk (“intra-chunk”), we use the parallel view and use quadratic attention, incurring a cost of $O(n_{\text{chunk}}^2 \cdot d_{\text{head}})$. Between chunks (“inter-chunk”), we use the recurrent view and use linear attention, incurring a cost of $O(d_{\text{head}}^2 \cdot n_{\text{chunk}})$.

Let’s first split the input sequence of length $n_{\text{seq}}$ into $\frac{n_{\text{seq}}}{n_{\text{chunk}}}$ non-overlapping chunks. We’ll use subscript $[\cdot]$ to denote chunk-level quantities. For chunk $i$, we have:

• $\mathbf{Q}_{[i]}, \mathbf{K}_{[i]}, \mathbf{V}_{[i]} \in \mathbb{R}^{n_{\text{chunk}} \times d_{\text{head}}}$: query, key, and value vectors for chunk $i$,
• $\mathbf{S}_{[i]} \in \mathbb{R}^{d_{\text{head}} \times d_{\text{head}}}$: memory matrix after processing chunks $0, 1, \ldots, i-1$.

Inter-chunk memory matrix recurrence

For the inter-chunk recurrence, we compute the current chunk’s memory matrix from the previous chunk’s memory matrix:

$$\mathbf{S}_{[i]} = \mathbf{S}_{[i-1]} + \mathbf{K}_{[i-1]}^{\top} \mathbf{V}_{[i-1]},$$

where $\mathbf{K}_{[i-1]}^{\top} \mathbf{V}_{[i-1]} = \sum_{j \in \text{chunk } i-1} \mathbf{k}_j \mathbf{v}_j^{\top}$ is the sum of outer products for all tokens in chunk $i-1$.

Computing the output for a chunk

For each chunk $i$, we compute the output as:

$$\mathbf{O}_{[i]} = \underbrace{\mathbf{Q}_{[i]} \mathbf{S}_{[i]}}_{\text{inter-chunk}} + \underbrace{(\mathbf{Q}_{[i]} \mathbf{K}_{[i]}^{\top} \odot \mathbf{M}) \mathbf{V}_{[i]}}_{\text{intra-chunk}},$$

where $\mathbf{M} \in \mathbb{R}^{n_{\text{chunk}} \times n_{\text{chunk}}}$ is the causal mask applied within the chunk.

The first term (inter-chunk) allows each token in chunk $i$ to attend to all tokens in previous chunks via the memory matrix $\mathbf{S}_{[i]}$. The second term (intra-chunk) allows each token to attend to previous tokens within the same chunk using standard attention.

Complexity of chunkwise parallel form

Let’s analyze the complexity for processing a single chunk:

• Inter-chunk memory matrix recurrence computation: we’re multiplying a $d_{\text{head}} \times n_{\text{chunk}}$ matrix by a $n_{\text{chunk}} \times d_{\text{head}}$ matrix, for a total cost of $O(n_{\text{chunk}} \cdot d_{\text{head}}^2)$.
• Computing the output for a chunk:
  • Inter-chunk output: we’re multiplying a $n_{\text{chunk}} \times d_{\text{head}}$ matrix by a $d_{\text{head}} \times d_{\text{head}}$ matrix, for a total cost of $O(n_{\text{chunk}} \cdot d_{\text{head}}^2)$.
  • Intra-chunk output: this is the standard quadratic attention computation, which is $O(n_{\text{chunk}}^2 \cdot d_{\text{head}})$.

The total complexity per chunk is therefore $O(n_{\text{chunk}}^2 \cdot d_{\text{head}} + n_{\text{chunk}} \cdot d_{\text{head}}^2)$.

For the entire sequence with $\frac{n_{\text{seq}}}{n_{\text{chunk}}}$ chunks, the total complexity is:

$$O\left(\frac{n_{\text{seq}}}{n_{\text{chunk}}} (n_{\text{chunk}}^2 \cdot d_{\text{head}} + n_{\text{chunk}} \cdot d_{\text{head}}^2)\right) = O(n_{\text{seq}} \cdot n_{\text{chunk}} \cdot d_{\text{head}} + n_{\text{seq}} \cdot d_{\text{head}}^2).$$

The chunk size $n_{\text{chunk}}$ is effectively a knob to control the tradeoff between parallelism and recurrence:

• When $n_{\text{chunk}} = n_{\text{seq}}$, we have the fully parallel form (quadratic complexity, maximally parallel);
• When $n_{\text{chunk}} = 1$, we have the fully recurrent form (linear complexity, maximally sequential);
• Intermediate values balance parallelism and complexity.

Gated delta networks

The gated delta network[3]Gated delta networks: Improving Mamba2 with delta rule [link]
Songlin Yang, Jan Kautz, and Ali Hatamizadeh. International Conference on Learning Representations. 2025. is a variant of linear attention that makes two additions: (1) a decay gate, and (2) a delta rule.

Decay gate

The intuition behind the decay gate is to gradually forget old information as we move forward in the sequence. This can be implemented by multiplying the memory matrix by a (data-dependent) decay factor $\alpha_t \in (0, 1)$:

$$\mathbf{S}_t = \alpha_t \mathbf{S}_{t-1} + \mathbf{k}_t \mathbf{v}_t^{\top}.$$

Delta rule

The intuition behind the delta rule[4]Linear transformers are secretly fast weight programmers [link]
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Proceedings of the 38th International Conference on Machine Learning. 2021. is this: when we update the memory matrix with some key-value pair $(\mathbf{k}_t, \mathbf{v}_t)$, there may already be some conflicting information corresponding to the same key $\mathbf{k}_t$. In addition to naively adding $\mathbf{k}_t \mathbf{v}_t^{\top}$ to the memory matrix, we can “clear out” the old information, and update the entry to correspond to a mix of the old and new information. Notice that we can retrieve the old conflicting value by searching our memory with key $\mathbf{k}_t$: $(\mathbf{v}_t^{\text{old}})^{\top} = \mathbf{k}_t^{\top} \mathbf{S}_{t-1}$. We can then update the memory matrix to clear out the old information, and then update the entry to correspond to a mix of the old and new information:

$$\begin{aligned} \mathbf{S}_t &= \mathbf{S}_{t-1} \;\;\underbrace{-\,\mathbf{k}_t (\mathbf{v}_t^{\text{old}})^{\top}}_{\text{delete old key-value}} \;+\; \underbrace{\mathbf{k}_t\!\left(\beta_t \mathbf{v}_t + (1-\beta_t)\,\mathbf{v}_t^{\text{old}}\right)^{\top}}_{\text{write new (blended) key-value}} \\ &= \mathbf{S}_{t-1} - \mathbf{k}_t \mathbf{k}_t^{\top} \mathbf{S}_{t-1} + \mathbf{k}_t(\beta_t \mathbf{v}_t^{\top} + (1 - \beta_t) \mathbf{k}_t^{\top}\mathbf{S}_{t-1}) \\ &= (\mathbf{I} - \beta_t \mathbf{k}_t \mathbf{k}_t^{\top}) \mathbf{S}_{t-1} + \beta_t \mathbf{k}_t \mathbf{v}_t^{\top}, \end{aligned}$$

where $\beta_t \in (0, 1)$ is a (data-dependent) “writing strength”.

Putting them together

The gated delta network basically just combines the decay gate and the delta rule.

Putting them together, we have:

$$\mathbf{S}_t = (\alpha_t(\mathbf{I} - \beta_t \mathbf{k}_t \mathbf{k}_t^{\top})) \mathbf{S}_{t-1} + \beta_t \mathbf{k}_t \mathbf{v}_t^{\top},$$

where $\alpha_t \in (0, 1)$ is a data-dependent decay factor, and $\beta_t \in (0, 1)$ is a data-dependent writing strength.

References

References cited in the text are listed first, in order of citation; additional references follow, ordered alphabetically.

1. Transformers are RNNs: Fast autoregressive transformers with linear attention [link]
   Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Proceedings of the 37th International Conference on Machine Learning. 2020.
2. Gated linear attention transformers with hardware-efficient training [link]
   Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Proceedings of the 41st International Conference on Machine Learning. 2024.
3. Gated delta networks: Improving Mamba2 with delta rule [link]
   Songlin Yang, Jan Kautz, and Ali Hatamizadeh. International Conference on Learning Representations. 2025.
4. Linear transformers are secretly fast weight programmers [link]
   Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Proceedings of the 38th International Conference on Machine Learning. 2021.
5. FlashAttention: Fast and memory-efficient exact attention with IO-awareness [link]
   Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Advances in Neural Information Processing Systems. 2022.
6. DeltaNet explained (Part I) [link]
   Songlin Yang. 2024.

Footnotes

1. Note that we omit the $\frac{1}{\sqrt{d_{\text{head}}}}$ scaling factor for simplicity. ↩

2. Note that some hardware-aware implementations of attention, such as FlashAttention[5]FlashAttention: Fast and memory-efficient exact attention with IO-awareness [link]
   Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Advances in Neural Information Processing Systems. 2022., can effectively compute attention without materializing the full $\mathbf{Q}\mathbf{K}^{\top} \in \mathbb{R}^{n_{\text{seq}} \times n_{\text{seq}}}$ matrix. ↩

3. We’ll temporarily ignore the causal mask $\mathbf{M}$, but will return to it later. It turns out to be very important. ↩