A Brief History of Linear Attention: From Imitation and Innovation to Feeding Back

线性注意力简史：从模仿、创新到反哺 苏剑林 · 2025-06-20 · Machine Learning linear RNN generative model attention

Among Chinese-language blogs, this site was probably one of the earliest to pay attention to linear attention. When I wrote the first related post, "Exploring Linear Attention: Does Attention Really Need a Softmax?", back in 2020, the main topic of discussion was still Softmax attention in the context of BERT. In hindsight, considering linear attention in the BERT era wasn't particularly wise, because training lengths were relatively short back then, and the models were mostly encoders — using linear attention offered essentially no advantage in that setting. I even wrote a post, "Linear Transformers Probably Aren't the Model You're Waiting For", expressing exactly this view.

It wasn't until the arrival of ChatGPT, which forced everyone toward decoder-only generative models, that things changed — this setup aligns extremely well with the RNN form of linear attention. At the same time, the pursuit of ever-longer training lengths made the quadratic complexity bottleneck of Softmax attention increasingly conspicuous. Against this new backdrop, linear attention has become steadily more competitive, to the point where there are now signs of it "feeding back" into Softmax attention itself.

Quadratic Complexity

Let's first introduce some notation:

\begin{equation}\begin{gathered} \boldsymbol{q}_i,\boldsymbol{k}_i,\boldsymbol{v}_i,\boldsymbol{o}_i \in \mathbb{R}^{d\times 1} \\[6pt] \boldsymbol{Q}=[\boldsymbol{q}_1,\boldsymbol{q}_2,\cdots,\boldsymbol{q}_n]^{\top}\in\mathbb{R}^{n\times d} \\[6pt] \boldsymbol{K}=[\boldsymbol{k}_1,\boldsymbol{k}_2,\cdots,\boldsymbol{k}_n]^{\top}\in\mathbb{R}^{n\times d} \\[6pt] \boldsymbol{V}=[\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_n]^{\top}\in\mathbb{R}^{n\times d} \\[6pt] \boldsymbol{O}=[\boldsymbol{o}_1,\boldsymbol{o}_2,\cdots,\boldsymbol{o}_n]^{\top}\in\mathbb{R}^{n\times d} \\[6pt] \end{gathered}\end{equation}

An attention model is, in essence, a mapping $\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}\to \boldsymbol{O}$. This post is mainly concerned with the causal setting, meaning $\boldsymbol{o}_t$ can depend on at most $\boldsymbol{Q}_{[:t]},\boldsymbol{K}_{[:t]},\boldsymbol{V}_{[:t]}$. In principle, the $d$ of $\boldsymbol{Q},\boldsymbol{K}$ need not match the $d$ of $\boldsymbol{V},\boldsymbol{O}$ — GAU and MLA are examples where they differ — but simplifying them to be the same doesn't change the essence of the problem.

Standard Softmax attention usually refers to the attention mechanism proposed in "Attention is All You Need":

\begin{equation}\boldsymbol{O} = \mathop{\text{softmax}}(\boldsymbol{Q}\boldsymbol{K}^{\top} + \log \boldsymbol{M})\boldsymbol{V}\end{equation}

Here the scaling factor $1/\sqrt{d}$ is omitted, since it can always be absorbed into $\boldsymbol{Q},\boldsymbol{K}$. $\mathop{\text{softmax}}$ denotes exponential normalization along the second dimension, and $\boldsymbol{M}\in\mathbb{R}^{n\times n}$ is a lower-triangular matrix called the mask matrix, defined as

\begin{equation}M_{i,j} = \left\{\begin{aligned} &1, &i \geq j \\ &0, &i < j\end{aligned}\right.\end{equation}

$\log\boldsymbol{M}$ means taking $\log$ of each component of $\boldsymbol{M}$, where $\log 0 = -\infty$. Written out in component form, Softmax attention is

\begin{equation}\boldsymbol{o}_t = \frac{\sum_{j=1}^t \exp(\boldsymbol{q}_t^{\top}\boldsymbol{k}_j) \boldsymbol{v}_j}{\sum_{j=1}^t \exp(\boldsymbol{q}_t^{\top}\boldsymbol{k}_j) }\end{equation}

where the denominator mainly serves to maintain numerical stability; also, if we apply RMSNorm to $\boldsymbol{O}$, the denominator gets automatically cancelled out. So the core of Softmax attention is really the numerator part, i.e.

\begin{equation}\boldsymbol{O} = \exp(\boldsymbol{Q}\boldsymbol{K}^{\top} + \log \boldsymbol{M})\boldsymbol{V} = (\exp(\boldsymbol{Q}\boldsymbol{K}^{\top})\odot \boldsymbol{M})\boldsymbol{V}\end{equation}

where $\odot$ is the Hadamard product and $\exp$ denotes elementwise exponentiation. It's easy to see that the denominator is essentially the same expression with $\boldsymbol{V}$ replaced by an all-ones matrix of shape $n\times 1$, which we can add back in if needed. The standard implementation of Softmax attention requires explicitly computing the $n\times n$ matrix $\exp(\boldsymbol{Q}\boldsymbol{K}^{\top})$, so both space and time complexity scale with $n^2$. Flash Attention reduced the memory requirement, but the quadratic time complexity remains unavoidable.

The Original Form

The earliest approaches to linear attention were mainly about imitating and approximating Softmax attention, with the simplest scheme being to just drop $\exp$, giving

\begin{equation}\boldsymbol{O} = (\boldsymbol{Q}\boldsymbol{K}^{\top}\odot \boldsymbol{M})\boldsymbol{V}\label{eq:linear-attn}\end{equation}

For simplicity, let's agree that matrix multiplication has higher precedence than the Hadamard product, so we can save a set of parentheses. Why is this form "linear" attention? To understand this quickly, let's first consider the non-causal version obtained by dropping $\odot \boldsymbol{M}$, in which case we have $\boldsymbol{O} = (\boldsymbol{Q}\boldsymbol{K}^{\top})\boldsymbol{V} = \boldsymbol{Q}(\boldsymbol{K}^{\top}\boldsymbol{V})$. Note that computing $\boldsymbol{K}^{\top}\boldsymbol{V}$ has complexity $\mathcal{O}(nd^2)$, producing a $d\times d$ matrix, and then multiplying by $\boldsymbol{Q}$ also costs $\mathcal{O}(nd^2)$, so the overall complexity scales linearly with $n$.

As for the causal version $\eqref{eq:linear-attn}$, we can understand it via the component form:

\begin{equation}\boldsymbol{o}_t = \sum_{j=1}^t \boldsymbol{v}_j (\boldsymbol{k}_j^{\top} \boldsymbol{q}_t) = \sum_{j=1}^t (\boldsymbol{v}_j \boldsymbol{k}_j^{\top}) \boldsymbol{q}_t = \left(\sum_{j=1}^t \boldsymbol{v}_j \boldsymbol{k}_j^{\top}\right) \boldsymbol{q}_t\end{equation}

If we denote the bracketed part as $\boldsymbol{S}_t$, then we have

\begin{equation}\boldsymbol{o}_t = \boldsymbol{S}_t \boldsymbol{q}_t, \qquad \boldsymbol{S}_t = \boldsymbol{S}_{t-1} + \boldsymbol{v}_t \boldsymbol{k}_t^{\top}\label{eq:linear-attn-rnn}\end{equation}

This shows that the causal form of attention can be written as a linear RNN with $\boldsymbol{S}_t$ as its state, so each step has constant complexity and the total complexity scales linearly with the sequence length $n$. Note that "linear RNN" has appeared here — it's a broader concept, and linear attention is one particular instance of it. Linear RNNs have also had their own independent line of development, such as the LRU and SSM discussed previously, though the most competitive linear architectures nowadays tend to take the form of linear attention.

Early linear attention still bore some very obvious marks of imitating Softmax attention — for instance, adding a denominator to normalize equation $\eqref{eq:linear-attn}$, which in turn requires $\boldsymbol{k}_j^{\top} \boldsymbol{q}_t$ to be non-negative, leading to non-negative activation functions being applied to $\boldsymbol{Q},\boldsymbol{K}$. A series of works represented by Performer and RFA even took approximating $\exp(\boldsymbol{Q}\boldsymbol{K}^{\top})$ as their explicit starting point.

However, later research such as "The Devil in Linear Transformer" found that normalizing along the sequence-length dimension doesn't fully avoid numerical instability — it's actually better to normalize after the fact, as in

\begin{equation}\boldsymbol{O} = \mathop{\text{RMSNorm}}((\boldsymbol{Q}\boldsymbol{K}^{\top}\odot \boldsymbol{M})\boldsymbol{V})\end{equation}

And once normalization is no longer required, adding a non-negative activation function to $\boldsymbol{Q},\boldsymbol{K}$ to ensure $\boldsymbol{k}_j^{\top} \boldsymbol{q}_t$ is non-negative is no longer necessary either. So is there still any point in applying a (not necessarily non-negative) activation function to $\boldsymbol{Q},\boldsymbol{K}$? My view is that adding an activation function is a matter of choice — it's entirely possible that some particular activation function yields better results — but doing so doesn't change the fundamental form of linear attention, so it doesn't affect our discussion here. Besides, existing results suggest that skipping it entirely already works well enough.

Fancy Forget Gates

From equation $\eqref{eq:linear-attn-rnn}$ we can see that current linear attention is essentially a $\mathop{\text{cumsum}}$, meaning it accumulates all historical information with equal weight. It's not hard to imagine that once enough tokens have been accumulated, each token's share of information becomes vanishingly small, so a fixed-size $\boldsymbol{S}_t$ matrix alone can't even accurately reconstruct any single token — intuitively, each token's "memory" becomes increasingly blurred.

To alleviate this problem, RetNet introduced a forgetting effect into linear attention:

\begin{equation}\boldsymbol{o}_t = \boldsymbol{S}_t \boldsymbol{q}_t, \qquad \boldsymbol{S}_t = \gamma\boldsymbol{S}_{t-1} + \boldsymbol{v}_t \boldsymbol{k}_t^{\top}\label{eq:linear-attn-retnet}\end{equation}

where the decay factor $\gamma\in(0,1)$ was set as a constant in RetNet, though some variants make it a trainable parameter, or replace $\gamma$ with a diagonal matrix, and so on — the linear attention used in MiniMax-01 is also of this type. Note that decay factors existed even before RetNet, but they mostly appeared in the form of linear RNNs, such as the LRU and SSM mentioned in the previous section; RetNet appears to be the first to combine this idea with linear attention. Once a decay factor is introduced, the model tends to forget older historical information, thereby at least preserving the resolution of more recent tokens — in plain terms, this is an embodiment of the "recency bias" that aligns well with the nature of language models, which tends to work better in practice.

There's also a noteworthy detail worth mentioning: RetNet also applies RoPE to $\boldsymbol{Q},\boldsymbol{K}$, which effectively generalizes the decay factor to a complex number $\gamma e^{\text{i}\theta}$ — from the LRU perspective, this amounts to considering complex eigenvalues. Although adding positional encoding to an RNN might seem somewhat odd at first glance, some experiments, such as the recent TransXSSM, suggest that adding RoPE to linear attention does have a certain positive effect. Of course, this may depend on the specific model variant and experimental setup.

A simple generalization of equation $\eqref{eq:linear-attn-retnet}$ is to replace $\gamma$ with a function $\gamma_t$ of position $t$, which already appeared in SSM. Later, works such as DFW, Mamba, and Mamba2 generalized this to be input-dependent, giving rise to the "data-dependent decay" line of work — which is by now already quite similar to the "forget gates" of earlier non-linear RNNs like GRU and LSTM, except that to preserve the model's linearity, the forget gate's dependence on the state (e.g., $\boldsymbol{S}_t$) is removed.

Why do we favor linear RNNs? Because linear RNNs generally admit some way of being trained in parallel, which makes them more competitive than Softmax attention — matching or exceeding it in both training and inference efficiency. Among these, the "general solution" for parallelization is to reformulate the problem as a Prefix Sum problem and apply an associative scan; the general idea was briefly introduced in the "Parallelization" section of "Google's New Work Tries to 'Revive' RNNs: Can RNNs Shine Again?".

However, this "general solution" is not GPU-efficient — what GPUs excel at is matrix multiplication, so finding parallel algorithms that make heavy use of matrix multiplication is ideal. In fact, even without full parallelization, simply finding a chunk-by-chunk recurrent formulation that makes full use of matrix multiplication can substantially improve training efficiency. This, in turn, places demands on the model: for instance, only forget gates with an outer-product form can achieve this. A typical counterexample is Mamba, whose forget gate is non-outer-product in form and thus cannot fully exploit GPU performance — which is precisely why Mamba2, GLA, and other variants followed.

Test-Time Training

At this point, linear attention has evolved from initially simply imitating Softmax attention, to incorporating static decay factors and even "data-dependent decay," forming its own distinctive character and proving its worth across a range of tasks. However, most of this progress has come from manual, experience-driven design. This naturally raises the question: Is there some higher-level principle that could guide the design of linear attention, or even of sequence models (token mixers) in general?

TTT (Test Time Training) offers its own answer to this question. It frames the construction of sequence models as an "online learning" problem, and proposes building (not necessarily linear) RNNs using optimizers. Specifically, it treats $\boldsymbol{K},\boldsymbol{V}$ as a corpus pair $(\boldsymbol{k}_1, \boldsymbol{v}_1),(\boldsymbol{k}_2, \boldsymbol{v}_2),\cdots,(\boldsymbol{k}_t, \boldsymbol{v}_t)$, trains a model $\boldsymbol{v} = \boldsymbol{f}(\boldsymbol{S}_t;\boldsymbol{k})$ on this corpus, and finally outputs $\boldsymbol{o}_t = \boldsymbol{f}(\boldsymbol{S}_t;\boldsymbol{q}_t)$, where $\boldsymbol{S}_t$ are the model's parameters — the model's architecture itself is, to a large extent, arbitrary.

What does this have to do with RNNs? It's simple: optimizers like SGD and Adam are, in essence, RNNs over the model parameters! This observation isn't new, in fact — as early as 2017, during the heyday of meta-learning, researchers had already proposed and exploited this idea, though back then the goal was to use an RNN (LSTM) to simulate a better optimizer; see "Optimization as a Model for Few-Shot Learning" for details.

As the saying goes, "what goes around comes around" — years later, TTT flips this around and proposes constructing RNNs via optimizers instead. The process works like this: given current model parameters $\boldsymbol{S}_{t-1}$, the optimizer (SGD) receives new data $(\boldsymbol{k}_t, \boldsymbol{v}_t)$, updates the model parameters to $\boldsymbol{S}_t$ based on this data, and finally returns the prediction $\boldsymbol{f}(\boldsymbol{S}_{t-1};\boldsymbol{q}_t)$ for $\boldsymbol{q}_t$, and so on. Thus, the RNN implemented by TTT can be written uniformly as

\begin{equation}\boldsymbol{o}_t = \boldsymbol{f}(\boldsymbol{S}_t; \boldsymbol{q}_t), \qquad \boldsymbol{S}_t = \boldsymbol{S}_{t-1} - \eta_t\nabla_{\boldsymbol{S}_{t-1}}\mathcal{L}(\boldsymbol{f}(\boldsymbol{S}_{t-1};\boldsymbol{k}_t), \boldsymbol{v}_t)\label{eq:ttt-rnn}\end{equation}

where $\mathcal{L}(\boldsymbol{f}(\boldsymbol{S}_{t-1};\boldsymbol{k}_t), \boldsymbol{v}_t)$ is the loss function for the current data $(\boldsymbol{k}_t, \boldsymbol{v}_t)$ under the current parameters $\boldsymbol{S}_{t-1}$, and $\eta_t$ is the learning rate parameter, which — following the "data-dependent decay" idea from the previous section — can also be made data-dependent. This formulation covers a great many RNN models; for example, equations $\eqref{eq:linear-attn-rnn}$ and $\eqref{eq:linear-attn-retnet}$ are both special cases of it:

$$\begin{array}{c|cc|ccc} \hline & \text{RNN} & \boldsymbol{o}_t & \boldsymbol{f}(\boldsymbol{S};\boldsymbol{k}) & \mathcal{L}(\boldsymbol{f}(\boldsymbol{S};\boldsymbol{k}),\boldsymbol{v}) & \eta_t \\ \hline \eqref{eq:linear-attn-rnn} & \boldsymbol{S}_t = \boldsymbol{S}_{t-1} + \boldsymbol{v}_t \boldsymbol{k}_t^{\top} & \boldsymbol{o}_t = \boldsymbol{S}_t \boldsymbol{q}_t & \boldsymbol{S}\boldsymbol{k} & -\boldsymbol{v}^{\top}(\boldsymbol{S}\boldsymbol{k}) & 1 \\ \eqref{eq:linear-attn-retnet} & \boldsymbol{S}_t = \gamma\boldsymbol{S}_{t-1} + \boldsymbol{v}_t \boldsymbol{k}_t^{\top} & \boldsymbol{o}_t = \boldsymbol{S}_t \boldsymbol{q}_t & \boldsymbol{S}\boldsymbol{k} & -\boldsymbol{v}^{\top}(\boldsymbol{S}\boldsymbol{k}) + \frac{1-\gamma}{2}\Vert\boldsymbol{S}\Vert_F^2 & 1 \\ \hline \end{array}$$

The original TTT paper focused on exploring non-linear RNNs under mini-batch settings. Later, Titans added momentum to TTT's SGD, and later still, "Test-Time Training Done Right" explored large-batch usage of TTT, as well as the "TTT + Muon" combination. Note that TTT only uses the optimizer to construct the RNN — parameters outside the RNN, such as trainable parameters like $\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}$, are still trained using an overall optimizer once the full model has been assembled.

A more thought-provoking question is: why can TTT serve as a "guiding principle" for constructing RNNs? The core goal of an RNN is to effectively compress historical data into a fixed-size state, and model parameters are, by their very nature, fixed in size — training a model is, in some sense, equivalent to compressing the training data into the model's weights. TTT exploits precisely this deep alignment with the goal of RNNs. To put it plainly, if we regard an RNN as a compression task, TTT treats the model $\boldsymbol{f}$ as the "decompressor," its weights as the "compressed archive," the compression algorithm as SGD, and the compression ratio as the loss $\mathcal{L}$.

This means we no longer need to rack our brains designing recurrent formulations by hand — instead, we can focus on designing the model $\boldsymbol{f}$ and the loss $\mathcal{L}$. Whether a given RNN is strong or reliable can then be judged simply by examining the corresponding $\boldsymbol{f}$ and $\mathcal{L}$.

Beyond this, since TTT constructs RNNs via online learning, the resulting RNNs are naturally very well suited to in-context learning (ICL) tasks — this is another advantage of TTT as a "guiding principle." Earlier, "Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers" even went in the reverse direction, explaining GPT's ICL capability by stripping the Softmax out of Softmax attention to obtain linear attention — from today's perspective, what it did was essentially construct the corresponding TTT formulation.

Out with the Old, In with the New

For example, the loss function corresponding to the earliest linear attention is $-\boldsymbol{v}^{\top}(\boldsymbol{S}\boldsymbol{k})$, which at a glance looks like a somewhat unreliable objective, because it is unbounded below — this could cause $\boldsymbol{S}$ to diverge to infinity. By contrast, RetNet added an L2 regularization term to the loss function, avoiding this risk, and from an optimization perspective it also mitigated overfitting, thereby yielding a better RNN.

However, although using the inner product as a loss function is simple and has some justification, it does not directly encourage $\boldsymbol{S}\boldsymbol{k}=\boldsymbol{v}$, so it is not an ideal regression loss. A better objective function should be the squared loss, i.e. $\frac{1}{2}\Vert\boldsymbol{S}\boldsymbol{k} - \boldsymbol{v}\Vert^2$. Substituting this into the TTT formula $\eqref{eq:ttt-rnn}$ gives

\begin{equation}\boldsymbol{o}_t = \boldsymbol{S}_t\boldsymbol{q}_t, \qquad \boldsymbol{S}_t = \boldsymbol{S}_{t-1} - \eta_t \underbrace{(\boldsymbol{S}_{t-1} \boldsymbol{k}_t - \boldsymbol{v}_t)\boldsymbol{k}_t^{\top}}_{\nabla_{\boldsymbol{S}_{t-1}}\frac{1}{2}\Vert\boldsymbol{S}_{t-1}\boldsymbol{k}_t - \boldsymbol{v}_t\Vert^2}\end{equation}

This is DeltaNet, a name that comes from Parallelizing Linear Transformers with the Delta Rule over Sequence Length; it was proposed even earlier in Linear Transformers Are Secretly Fast Weight Programmers. Note that $\eta_t (\boldsymbol{S}_{t-1} \boldsymbol{k}_t - \boldsymbol{v}_t)\boldsymbol{k}_t^{\top} = (\boldsymbol{S}_{t-1} (\sqrt{\eta_t}\boldsymbol{k}_t) - (\sqrt{\eta_t}\boldsymbol{v}_t))(\sqrt{\eta_t}\boldsymbol{k}_t)^{\top}$, which means $\eta_t$ can always be absorbed into the definition of $\boldsymbol{k}_t,\boldsymbol{v}_t$, so in what follows we only consider the case $\eta_t=1$:

\begin{equation}\begin{aligned} \boldsymbol{S}_t =&\, \boldsymbol{S}_{t-1} -(\boldsymbol{S}_{t-1} \boldsymbol{k}_t - \boldsymbol{v}_t)\boldsymbol{k}_t^{\top} \\ =&\, \boldsymbol{S}_{t-1} -(\boldsymbol{S}_{t-1} \boldsymbol{k}_t)\boldsymbol{k}_t^{\top} + \boldsymbol{v}_t\boldsymbol{k}_t^{\top} \\ =&\, \boldsymbol{S}_{t-1} (\boldsymbol{I} - \boldsymbol{k}_t\boldsymbol{k}_t^{\top}) + \boldsymbol{v}_t\boldsymbol{k}_t^{\top} \end{aligned}\label{eq:linear-attn-deltanet}\end{equation}

If needed, we can later replace $\boldsymbol{k}_t,\boldsymbol{v}_t$ with $\sqrt{\eta_t}\boldsymbol{k}_t,\sqrt{\eta_t}\boldsymbol{v}_t$ to recover $\eta_t$. Compared with the earliest form of linear attention $\eqref{eq:linear-attn-rnn}$, the difference in DeltaNet is that before adding $\boldsymbol{v}_t\boldsymbol{k}_t^{\top}$, it first subtracts $(\boldsymbol{S}_{t-1} \boldsymbol{k}_t)\boldsymbol{k}_t^{\top}$, where $\boldsymbol{S}_{t-1} \boldsymbol{k}_t$ can be understood as the prediction of the new input $\boldsymbol{k}_t$ made by the old model $\boldsymbol{S}_{t-1}$.

Intuitively, "subtract first, then add" means first removing the model's old understanding of $\boldsymbol{k}_t$, then supplementing it with new knowledge based on $(\boldsymbol{k}_t,\boldsymbol{v}_t)$, achieving the effect of "clearing out the old to make way for the new." This rule is called the "Delta Rule," which is where the name "Delta" in DeltaNet comes from. The Delta Rule is nothing new — it is also known as Least Mean Square, the Widrow-Hoff Algorithm, etc., and dates back to the 1960s. In fact, there is very little that is entirely new in this field; many modifications can be traced back to some "ancient" piece of work, and current efforts are mostly focused on mining the parts of it that can be made scalable.

It should also be pointed out that, chronologically, DeltaNet came before TTT. Understanding RNNs from the perspective of online learning had already appeared sporadically in some works prior to TTT, but TTT systematically articulated this "guiding principle" and used it to construct new RNN models. So we have presented TTT first here, to make the overall introduction flow more naturally.

Some readers may wonder: is DeltaNet still a linear RNN? The answer is yes. What we mean by a linear RNN is that the recurrence formula depends linearly on the state variable, though its dependence on the input or $\boldsymbol{q},\boldsymbol{k},\boldsymbol{v}$ can be nonlinear (of course, different forms of dependence will have different degrees of parallelization efficiency). From equation $\eqref{eq:linear-attn-deltanet}$ we can see that the right-hand side always contains only the first power of $\boldsymbol{S}_{t-1}$, so it satisfies the definition of linearity.

Matrix Inversion and Generalizations

Earlier we said that the most ideal (i.e. GPU-efficient) parallel algorithm for linear RNNs is one that makes full use of matrix multiplication. To achieve this, let us first rewrite DeltaNet as

\begin{equation}\boldsymbol{S}_t = \boldsymbol{S}_{t-1} + (\boldsymbol{v}_t - \boldsymbol{S}_{t-1} \boldsymbol{k}_t)\boldsymbol{k}_t^{\top}\end{equation}

Let $\boldsymbol{u}_t = \boldsymbol{v}_t - \boldsymbol{S}_{t-1} \boldsymbol{k}_t$, so that $\boldsymbol{S}_t = \boldsymbol{S}_{t-1} + \boldsymbol{u}_t\boldsymbol{k}_t^{\top}$ — in other words, this is just the earliest form of linear attention with $\boldsymbol{V}$ replaced by $\boldsymbol{U}=[\boldsymbol{u}_1,\boldsymbol{u}_2,\cdots,\boldsymbol{u}_n]^{\top}$. Iterating this $t-1$ times, we obtain

\begin{equation}\boldsymbol{S}_{t-1} = \sum_{j=1}^{t-1} \boldsymbol{u}_j\boldsymbol{k}_j^{\top}\qquad\Rightarrow\qquad \boldsymbol{u}_t = \boldsymbol{v}_t - \left(\sum_{j=1}^{t-1} \boldsymbol{u}_j\boldsymbol{k}_j^{\top}\right)\boldsymbol{k}_t = \boldsymbol{v}_t - \sum_{j=1}^{t-1} \boldsymbol{u}_j(\boldsymbol{k}_j^{\top}\boldsymbol{k}_t)\end{equation}

The final equality, written in matrix form, is $\boldsymbol{U} = \boldsymbol{V} - (\boldsymbol{K}\boldsymbol{K}^{\top}\odot \boldsymbol{M}^-)\boldsymbol{U}$, where $\boldsymbol{M}^-=\boldsymbol{M} - \boldsymbol{I}$. This is a linear system of equations whose solution can be written directly as

\begin{equation}\boldsymbol{U} = (\boldsymbol{I} + \underbrace{\boldsymbol{K}\boldsymbol{K}^{\top}\odot \boldsymbol{M}^-}_{\text{denote}\boldsymbol{B}})^{-1}\boldsymbol{V}\end{equation}

Here we encounter $(\boldsymbol{I}+\boldsymbol{B})^{-1}$, the inverse of a $n\times n$ matrix, whose standard complexity is $\mathcal{O}(n^3)$ — even higher than Softmax attention! Fortunately, we don't need the explicit inverse but only $\boldsymbol{U}$, which can be turned into solving the linear system $(\boldsymbol{I}+\boldsymbol{B})\boldsymbol{U}=\boldsymbol{V}$, bringing the complexity down to $\mathcal{O}(n^2)$. Further, by exploiting the fact that $\boldsymbol{I}+\boldsymbol{B}$ is lower triangular and the low-rank structure of $\boldsymbol{B}$, the complexity can be reduced to linear, and once written in terms of block matrix multiplication, it can fully exploit the GPU. For these details, readers will have to consult the original paper — here we're only laying out the main mathematical principles.

After DeltaNet came Gated DeltaNet (GDN), which further introduced a forget gate into DeltaNet — a fairly predictable development. Gated DeltaNet's original way of introducing the gate was

\begin{equation}\boldsymbol{S}_t = \alpha_t \boldsymbol{S}_{t-1} (\boldsymbol{I} - \beta_t\boldsymbol{k}_t\boldsymbol{k}_t^{\top}) + \beta_t\boldsymbol{v}_t\boldsymbol{k}_t^{\top}\label{eq:gdn-orgi}\end{equation}

But in my personal opinion, this formulation actually breaks the Delta Rule explicitly; a better formulation would be, as in Comba, to multiply only the first $\boldsymbol{S}_{t-1}$:

\begin{equation}\boldsymbol{S}_t = \gamma_t\boldsymbol{S}_{t-1} + \eta_t(\boldsymbol{v}_t - \boldsymbol{S}_{t-1}\boldsymbol{k}_t)\boldsymbol{k}_t^{\top}\label{eq:gdn-comba}\end{equation}

This is equivalent to taking the loss function as $\frac{1}{2}\Vert\boldsymbol{S}\boldsymbol{k} - \boldsymbol{v}\Vert^2 + \frac{1-\gamma}{\eta}\Vert\boldsymbol{S}\Vert_F^2$. Of course, mathematically the two formulations are equivalent:

\begin{equation}\alpha_t\boldsymbol{S}_{t-1} (\boldsymbol{I} - \beta_t\boldsymbol{k}_t\boldsymbol{k}_t^{\top}) + \beta_t\boldsymbol{v}_t\boldsymbol{k}_t^{\top} = \alpha_t \boldsymbol{S}_{t-1} + \alpha_t \beta_t (\boldsymbol{v}_t/\alpha_t - \boldsymbol{S}_{t-1}\boldsymbol{k}_t)\boldsymbol{k}_t^{\top}\end{equation}

That is, by setting $\gamma_t = \alpha_t, \eta_t = \alpha_t \beta_t$ and then absorbing $1/\alpha_t$ into $\boldsymbol{v}_t$, we can transform one into the other. So these two forms are mathematically identical; since most values of $\alpha_t$ will be close to 1, they probably don't differ much in capability either (Comba claims $\eqref{eq:gdn-comba}$ is slightly better), except that the latter more intuitively preserves the appearance of the Delta Rule.

Theoretically, Gated DeltaNet can also be written in the form of DeltaNet, because we just need to define $\bar{\alpha}_t = \prod_{j=1}^t \alpha_t$, and then dividing both sides of equation $\eqref{eq:gdn-orgi}$ by $\bar{\alpha}_t$ gives

\begin{equation}\bar{\alpha}_t^{-1}\boldsymbol{S}_t = \bar{\alpha}_{t-1}^{-1}\boldsymbol{S}_{t-1} (\boldsymbol{I} - \beta_t\boldsymbol{k}_t\boldsymbol{k}_t^{\top}) + \beta_t(\bar{\alpha}_t^{-1}\boldsymbol{v}_t)\boldsymbol{k}_t^{\top}\end{equation}

Then, combining this with $\boldsymbol{o}_t = \boldsymbol{S}_t \boldsymbol{q}_t = (\bar{\alpha}_t^{-1}\boldsymbol{S}_t) (\bar{\alpha}_t\boldsymbol{q}_t)$, we find that simply setting $\bar{\alpha}_t\boldsymbol{q}_t,\bar{\alpha}_t^{-1}\boldsymbol{v}_t$ as the new $\boldsymbol{q}_t,\boldsymbol{v}_t$ simplifies things to the form of DeltaNet. However, this result only has theoretical derivation value in certain cases (such as deriving the attention matrix in the next section), because in actual computation, regardless of parameterization, for sufficiently large $t$, either $\bar{\alpha}_t$ or $\bar{\alpha}_t^{-1}$ risks overflowing.

There is another generalization of DeltaNet called DeltaProduct, which expands $\boldsymbol{k},\boldsymbol{v}$ by some multiple before applying DeltaNet or Gated DeltaNet, in an attempt to enhance the model's state-tracking ability. However, in my own aesthetic judgment, rather than expanding by a constant factor as DeltaProduct does, it might be more worthwhile to try a quadratic-complexity RNN as in Chapter on Spacetime: Viewing Attention as a Quadratic-Complexity RNN, and see whether there is a chance to surpass Softmax attention.

Feeding Back in Progress

Speaking of surpassing Softmax attention — as mentioned at the outset, today's linear attention is not only competitive with Softmax attention but has even begun to "feed back" into it. This may seem hard to believe, but on reflection it is not difficult to understand. In a sense, Softmax attention has been regressing over the years — going from MHA, to GQA, to MQA, all in the name of subtraction to compress the KV cache. Linear attention, on the other hand, has no KV cache problem, and so has kept moving in a better direction.

To make this clearer, let's write out the attention mechanisms mentioned earlier in matrix form:

$$\begin{array}{c|c} \hline & \text{formula} \\[4pt] \hline \text{Softmax Attention} & (\exp(\boldsymbol{Q}\boldsymbol{K}^{\top})\odot \boldsymbol{M})\boldsymbol{V} \\[4pt] \text{earliest linear Attention} & (\boldsymbol{Q}\boldsymbol{K}^{\top}\odot \boldsymbol{M})\boldsymbol{V} \\[4pt] \text{after adding forget gate} & (\boldsymbol{Q}\boldsymbol{K}^{\top}\odot \boldsymbol{\Gamma})\boldsymbol{V} \\[4pt] \text{DeltaNet} & (\boldsymbol{Q}\boldsymbol{K}^{\top}\odot \boldsymbol{M})(\boldsymbol{I} + \boldsymbol{K}\boldsymbol{K}^{\top}\odot \boldsymbol{M}^-)^{-1}\boldsymbol{V} \\[4pt] \text{Gated DeltaNet} & \begin{gathered}((\boldsymbol{Q}\boldsymbol{K}^{\top}\odot \boldsymbol{M})(\boldsymbol{I} + \boldsymbol{K}\boldsymbol{K}^{\top}\odot \boldsymbol{M}^-)^{-1}\odot\boldsymbol{\Gamma})\boldsymbol{V} \\ =(\boldsymbol{Q}\boldsymbol{K}^{\top}\odot \boldsymbol{\Gamma})(\boldsymbol{I} + \boldsymbol{K}\boldsymbol{K}^{\top}\odot \boldsymbol{\Gamma}^-)^{-1}\boldsymbol{V}\end{gathered} \\[4pt] \hline \end{array}$$

where

\begin{equation}\Gamma_{i,j} = \left\{\begin{aligned} &\prod_{\tau=j+1}^i \gamma_{\tau}, &i > j \\[6pt] &\qquad 1, &i = j \\[6pt] &\qquad 0, &i < j\end{aligned}\right.\end{equation}

and $\boldsymbol{\Gamma}^- = \boldsymbol{\Gamma} - \boldsymbol{I}$. Seen this way, the form of Softmax attention is still stuck at the stage of the earliest linear attention (which, of course, also demonstrates just how powerful it is). So how do we implement this "feedback"? First we need a way to convert Softmax attention into linear attention, which isn't difficult — as far back as The Path to the Transformer, Part 5: Linear Attention as Infinite-Dimensional Attention we summarized three schemes for converting Softmax attention into infinite-dimensional linear attention.

In short, there exists a mapping $\phi$ that maps $\boldsymbol{Q},\boldsymbol{K}$ from $n\times d$ to $n\times \infty$, satisfying $\exp(\boldsymbol{Q}\boldsymbol{K}^{\top}) = \phi(\boldsymbol{Q})\phi(\boldsymbol{K})^{\top}$ — this is called the "kernel trick." What follows is then simple: we just need to replace $\boldsymbol{Q},\boldsymbol{K}$ in the linear attention formulas in the table above with $\phi(\boldsymbol{Q}),\phi(\boldsymbol{K})$, then find a way to recover $\exp$ and normalize it, and we get a new variant of Softmax attention. For example, substituting into the forget-gate formula gives us

\begin{equation}(\phi(\boldsymbol{Q})\phi(\boldsymbol{K})^{\top}\odot \boldsymbol{\Gamma})\boldsymbol{V} = \exp(\boldsymbol{Q}\boldsymbol{K}^{\top} + \log\boldsymbol{\Gamma})\boldsymbol{V}\end{equation}

If $\gamma_t$ is constant, this is exactly the ALIBI proposed in Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation; and if $\gamma_t$ depends on the input, then it's the FoX proposed in Forgetting Transformer: Softmax Attention with a Forget Gate.

An even more interesting result is DeltaFormer, proposed in Understanding Transformer from the Perspective of Associative Memory. As the name suggests, it is the DeltaNet version of Softmax attention. Replacing $\boldsymbol{Q},\boldsymbol{K}$ in DeltaNet with $\phi(\boldsymbol{Q}),\phi(\boldsymbol{K})$, we get

\begin{equation}\begin{aligned} &\,(\phi(\boldsymbol{Q})\phi(\boldsymbol{K})^{\top}\odot \boldsymbol{M})(\boldsymbol{I} + \phi(\boldsymbol{K})\phi(\boldsymbol{K})^{\top}\odot \boldsymbol{M}^-)^{-1}\boldsymbol{V} \\[8pt] =&\,\underbrace{\exp(\boldsymbol{Q} \boldsymbol{K}^{\top} + \log\boldsymbol{M})}_{\text{denote}\boldsymbol{A}}(\boldsymbol{I} + \underbrace{\exp(\boldsymbol{K} \boldsymbol{K}^{\top} + \log\boldsymbol{M}^-)}_{\text{denote}\boldsymbol{B}})^{-1}\boldsymbol{V} \end{aligned}\end{equation}

If normalization is needed, we simply replace $\exp$ with $\text{softmax}$. Compared with Softmax attention, DeltaFormer changes the original $\boldsymbol{A}\boldsymbol{V}$ into $\boldsymbol{A}(\boldsymbol{I}+\boldsymbol{B})^{-1}\boldsymbol{V}$. Note that

\begin{equation}\begin{aligned} \boldsymbol{A}(\boldsymbol{I}+\boldsymbol{B})^{-1}\boldsymbol{V} =&\, \boldsymbol{A}(\boldsymbol{I}-\boldsymbol{B}+\boldsymbol{B}^2- \boldsymbol{B}^3 + \cdots)\boldsymbol{V} \\ =&\, \boldsymbol{A}(\boldsymbol{V}-\boldsymbol{B}\boldsymbol{V}+\boldsymbol{B}^2\boldsymbol{V}- \boldsymbol{B}^3\boldsymbol{V} + \cdots) \end{aligned}\end{equation}

so DeltaFormer is equivalent to first computing attention multiple times using $\boldsymbol{K},\boldsymbol{K},\boldsymbol{V}$, superposing the results to form a new $\boldsymbol{V}$, and then computing attention once more with $\boldsymbol{Q},\boldsymbol{K}$. This property gives it remarkable effectiveness on multi-hop tasks (such as code). Moreover, this feature of DeltaFormer also means it pairs particularly well with MQA, because only $(\boldsymbol{I}+\boldsymbol{B})^{-1}\boldsymbol{V}$ is involved in this part — $\boldsymbol{K},\boldsymbol{V}$ — and for MQA, $\boldsymbol{K},\boldsymbol{V}$ has only a single head, so the computational cost is significantly lower than with MHA.

However, in my view, this kind of fixed-coefficient superposition may amount to "no free lunch": for example, my own experimental results show that DeltaFormer's language-model loss doesn't change much overall, which means that if the loss on some tasks drops noticeably, the loss on some other tasks must be rising to compensate.

Hardcore Encoding Techniques

Another feedback-related piece of work worth noting is PaTH Attention, from PaTH Attention: Position Encoding via Accumulating Householder Transformations, which feeds DeltaNet back into Softmax attention from the perspective of positional encoding.

In The Path to the Transformer, Part 6: A Completeness Analysis of Rotary Position Embedding we pointed out that for any orthogonal matrix $\boldsymbol{\Omega}$, $\boldsymbol{R}_m = \boldsymbol{\Omega}^m$ is a generalized form of RoPE. Besides rotation matrices, what other easily constructed orthogonal matrices are there? PaTH uses the Householder matrix: let $\boldsymbol{w}$ be any column vector of norm $\sqrt{2}$; then $\boldsymbol{I}-\boldsymbol{w}\boldsymbol{w}^{\top}$ is an orthogonal matrix. We also derived this in An Orthogonal Matrix Transforming One Unit Vector into Another, where the geometric meaning is a mirror reflection.

It's easy to see that this is exactly the same as the $\boldsymbol{I}-\boldsymbol{k}_t\boldsymbol{k}_t^{\top}$ that $\boldsymbol{S}_{t-1}$ is multiplied by in DeltaNet, so PaTH simply borrows this part directly — that is, it gives up the form $\boldsymbol{\Omega}^m$, as well as the constraint that $\boldsymbol{w}$ has norm $\sqrt{2}$, and instead directly expresses positional information as a product of a sequence of $\boldsymbol{I}-\boldsymbol{w}\boldsymbol{w}^{\top}$ terms:

\begin{equation}\boldsymbol{q}_i^{\top}\boldsymbol{k}_j \qquad\to\qquad \boldsymbol{q}_i^{\top}\underbrace{(\boldsymbol{I}-\boldsymbol{w}_i\boldsymbol{w}_i^{\top})(\boldsymbol{I}-\boldsymbol{w}_{i-1}\boldsymbol{w}_{i-1}^{\top})\cdots(\boldsymbol{I}-\boldsymbol{w}_{j+1}\boldsymbol{w}_{j+1}^{\top})}_{\text{denote}\boldsymbol{R}_{i,j}}\boldsymbol{k}_j \end{equation}

Writing $\boldsymbol{R}_{i,j}$ in recursive form gives $\boldsymbol{R}_{i,j} = (\boldsymbol{I}-\boldsymbol{w}_i\boldsymbol{w}_i^{\top})\boldsymbol{R}_{i-1,j},\boldsymbol{R}_{j,j} = \boldsymbol{I}$. Comparing this with DeltaNet's equation $\eqref{eq:linear-attn-deltanet}$, the above is equivalent to having $\boldsymbol{v}_t$ identically zero, except that the initial value $\boldsymbol{S}_0$ is no longer zero. Using the same procedure as in the section "Matrix Inversion and Generalizations", we can obtain

\begin{equation}\boldsymbol{R}_{i,j} = \boldsymbol{I} - \boldsymbol{W}_{[j:i]}^{\top}(\boldsymbol{I} + \boldsymbol{W}_{[j:i]}\boldsymbol{W}_{[j:i]}^{\top}\odot\boldsymbol{M}^-)^{-1}\boldsymbol{W}_{[j:i]}\end{equation}

where $\boldsymbol{W}=[\boldsymbol{w}_1,\boldsymbol{w}_2,\cdots,\boldsymbol{w}_n]^{\top}$, with slicing to be understood in the NumPy sense — e.g. $\boldsymbol{W}_{[j:i]}=[\boldsymbol{w}_{j+1},\boldsymbol{w}_{j+2},\cdots,\boldsymbol{w}_i]^{\top}$ — and slicing taking precedence over transposition. Note that the matrix being inverted is lower triangular, and triangular matrices have an important property: the diagonal elements of the inverse equal the reciprocals of the diagonal elements of the original matrix; and if it's a block-triangular matrix, the diagonal blocks satisfy the same property. So we can write

\begin{equation}(\boldsymbol{I} + \boldsymbol{W}_{[j:i]}\boldsymbol{W}_{[j:i]}^{\top}\odot\boldsymbol{M}^-)^{-1} = (\underbrace{(\boldsymbol{I} + \boldsymbol{W}\boldsymbol{W}^{\top}\odot\boldsymbol{M}^-)^{-1}}_{\text{denote}\boldsymbol{J}})_{[j:i,j:i]}\end{equation}

The transformations that follow may be easier to understand written out in component form:

\begin{equation}\begin{aligned} A_{i,j} =&\, \boldsymbol{q}_i^{\top} \boldsymbol{R}_{i,j} \boldsymbol{k}_j \\[6pt] =&\, \boldsymbol{q}_i^{\top}\boldsymbol{k}_j - \boldsymbol{q}_i^{\top}\boldsymbol{W}_{[j:i]}^{\top}\boldsymbol{J}_{[j:i,j:i]}\boldsymbol{W}_{[j:i]}\boldsymbol{k}_j \\ =&\, \boldsymbol{q}_i^{\top}\boldsymbol{k}_j - \sum_{p=1}^d \sum_{l=j+1}^i \sum_{r=j+1}^i \sum_{s=1}^d Q_{i,p} W_{l,p} J_{l,r} W_{r,s} K_{j,s} \\ =&\, \boldsymbol{q}_i^{\top}\boldsymbol{k}_j - \sum_{p=1}^d \sum_{l=1}^i \sum_{r=j+1}^n \sum_{s=1}^d Q_{i,p} W_{l,p} J_{l,r} W_{r,s} K_{j,s} \\ =&\, \boldsymbol{q}_i^{\top}\boldsymbol{k}_j - \sum_{p=1}^d \sum_{l=1}^n \sum_{r=1}^n \sum_{s=1}^d Q_{i,p} W_{l,p} \chi_{l \leq i} J_{l,r} \chi_{r \geq j+1}W_{r,s} K_{j,s} \\ =&\, \boldsymbol{q}_i^{\top}\boldsymbol{k}_j - \sum_{l=1}^n \sum_{r=1}^n \underbrace{\left(\chi_{l \leq i}\sum_{p=1}^d Q_{i,p} W_{l,p}\right)}_{\boldsymbol{Q}\boldsymbol{W}^{\top}\odot\boldsymbol{M}} J_{l,r} \underbrace{\left(\chi_{r \geq j+1} \sum_{s=1}^d W_{r,s} K_{j,s}\right)}_{\boldsymbol{W}\boldsymbol{K}^{\top}\odot\boldsymbol{M}^-} \\ \end{aligned}\end{equation}

There are several key points here: the somewhat clever step is the fourth equality, which exploits the fact that $\boldsymbol{J}$ is lower triangular, so $l < r$ is automatically zero when $J_{l,r}$; in the fifth equality, $\chi$ is an indicator function, equal to 1 when the subscript condition is satisfied and 0 otherwise; in the sixth equality, when we separately sum over the two parts of $p,s$, the results are $\boldsymbol{Q}\boldsymbol{W}^{\top}$ and $\boldsymbol{W}\boldsymbol{K}^{\top}$, and multiplying by $\chi_{l \leq i}$ corresponds exactly to keeping the lower-triangular part of $\boldsymbol{Q}\boldsymbol{W}^{\top}$ (including the diagonal), while multiplying by $\chi_{r \geq j+1}$ corresponds to keeping the lower-triangular part of $\boldsymbol{W}\boldsymbol{K}^{\top}$ (excluding the diagonal).

At this point, we can write out the entire (pre-softmax) attention matrix:

\begin{equation}\boldsymbol{A} = \boldsymbol{Q}\boldsymbol{K}^{\top}\odot\boldsymbol{M} - (\boldsymbol{Q}\boldsymbol{W}^{\top}\odot\boldsymbol{M})(\boldsymbol{I} + \boldsymbol{W}\boldsymbol{W}^{\top}\odot\boldsymbol{M}^-)^{-1}(\boldsymbol{W}\boldsymbol{K}^{\top}\odot\boldsymbol{M}^-) \label{eq:path-attn}\end{equation}

Are you impressed yet? We're not done. Direct matrix inversion has complexity $\mathcal{O}(n^3)$, which is obviously unacceptable, so we still need to exploit the low-rank structure of $\boldsymbol{W}\boldsymbol{W}^{\top}$ to bring the complexity down to $\mathcal{O}(n^2)$, and then work out the backward pass, and finally implement it efficiently in a Flash-Attention-like style. For these details, readers will have to dig into the original paper — it's a genuinely hardcore piece of work from start to finish.

From the perspective of positional encoding, PaTH is a form of CoPE (Contextual Position Encoding): its "position" is not an index $1,2,3,\cdots$, but a positional signal automatically generated from the context content. Similarly, FoX can also be seen as a contextual version of ALIBI. Context-dependent positional information is a defining feature of current linear-attention research, and quite possibly the main direction through which it will feed back into Softmax attention.

Simplification without End

Let us dig a bit deeper into PaTH — this will not only help us understand PaTH better, but also help us get more familiar with DeltaNet, since the two are themselves highly related. In this section we start from two special cases of PaTH, which will help us better understand the connection between PaTH and DeltaNet.

The first special case is $\boldsymbol{W}=\boldsymbol{K}$; substituting this into $\eqref{eq:path-attn}$ gives

\begin{equation}\begin{aligned} \boldsymbol{A} =&\, (\boldsymbol{Q}\boldsymbol{K}^{\top}\odot\boldsymbol{M})(\boldsymbol{I} - (\boldsymbol{I} + \boldsymbol{K}\boldsymbol{K}^{\top}\odot\boldsymbol{M}^-)^{-1}(\boldsymbol{K}\boldsymbol{K}^{\top}\odot\boldsymbol{M}^-)) \\[6pt] =&\, (\boldsymbol{Q}\boldsymbol{K}^{\top}\odot\boldsymbol{M})(\boldsymbol{I} + \boldsymbol{K}\boldsymbol{K}^{\top}\odot\boldsymbol{M}^-)^{-1} \qquad (\text{note}:\boldsymbol{I} - (\boldsymbol{I} + \boldsymbol{A})^{-1} \boldsymbol{A} = (\boldsymbol{I}+\boldsymbol{A})^{-1}) \end{aligned}\end{equation}

Does this look a bit familiar? This is exactly the attention matrix of DeltaNet! From this special case we can see that the difference between PaTH and DeltaFormer is this: DeltaFormer, based on the kernel trick, adds $\exp$ to DeltaNet's $\boldsymbol{Q}\boldsymbol{K}^{\top}$ and $\boldsymbol{K}\boldsymbol{K}^{\top}$ separately, whereas PaTH adds $\exp$ directly to DeltaNet's attention matrix.

The second special case is to reintroduce the constraint $\Vert\boldsymbol{w}\Vert=\sqrt{2}$, in which case $\boldsymbol{I}-\boldsymbol{w}\boldsymbol{w}^{\top}$ is an orthogonal matrix. Let us introduce

\begin{equation}\begin{aligned} \boldsymbol{R}_i \triangleq&\, (\boldsymbol{I}-\boldsymbol{w}_i\boldsymbol{w}_i^{\top})(\boldsymbol{I}-\boldsymbol{w}_{i-1}\boldsymbol{w}_{i-1}^{\top})\cdots(\boldsymbol{I}-\boldsymbol{w}_1\boldsymbol{w}_1^{\top}) \\ =&\, \boldsymbol{I} - \boldsymbol{W}_{[:i]}^{\top}(\boldsymbol{I} + \boldsymbol{W}_{[:i]}\boldsymbol{W}_{[:i]}^{\top}\odot\boldsymbol{M}^-)^{-1}\boldsymbol{W}_{[:i]} \\ =&\,\boldsymbol{R}_{i,0} \end{aligned}\end{equation}

then $\boldsymbol{R}_{i,j} = \boldsymbol{R}_i \boldsymbol{R}_j^{\top}$. This identity means that, just as with RoPE, we can implement the relative-position version of PaTH via an absolute-position formulation: we simply multiply each $\boldsymbol{q}_i^{\top},\boldsymbol{k}_i^{\top}$ by $\boldsymbol{R}_i$, and then apply the usual Softmax Attention implementation. So what operation is multiplication by $\boldsymbol{R}_i$? Repeating the expansion from the previous section, we get

\begin{equation}\begin{aligned} (\boldsymbol{q}_i^{\top} \boldsymbol{R}_{i})_s =&\, (\boldsymbol{q}_i^{\top} - \boldsymbol{q}_i^{\top}\boldsymbol{W}_{[:i]}^{\top}\boldsymbol{J}_{[:i,:i]}\boldsymbol{W}_{[:i]})_s \\ =&\, Q_{i,s} - \sum_{p=1}^d \sum_{l=1}^i \sum_{r=1}^i Q_{i,p} W_{l,p} J_{l,r} W_{r,s} \\ =&\, Q_{i,s} - \sum_{p=1}^d \sum_{l=1}^i \sum_{r=1}^n Q_{i,p} W_{l,p} J_{l,r} W_{r,s} \\ =&\, Q_{i,s} - \sum_{p=1}^d \sum_{l=1}^n \sum_{r=1}^n \chi_{l\leq i} Q_{i,p} W_{l,p} J_{l,r} W_{r,s} \\ =&\, Q_{i,s} - \sum_{l=1}^n \underbrace{\chi_{l\leq i} \sum_{p=1}^d Q_{i,p} W_{l,p}}_{\boldsymbol{Q}\boldsymbol{W}^{\top}\odot\boldsymbol{M}}\, \underbrace{\sum_{r=1}^n J_{l,r} W_{r,s}}_{\boldsymbol{J}\boldsymbol{W}} \end{aligned}\end{equation}

which in matrix form is

\begin{equation}\boldsymbol{\boldsymbol{Q}} - (\boldsymbol{Q}\boldsymbol{W}^{\top}\odot\boldsymbol{M})(\boldsymbol{I} + \boldsymbol{W}\boldsymbol{W}^{\top}\odot\boldsymbol{M}^-)^{-1}\boldsymbol{W}\end{equation}

Does this look familiar again? In fact, the second part is exactly $\text{DeltaNet}(\boldsymbol{Q},\boldsymbol{W},\boldsymbol{W})$! So in this case, the effect implemented by PaTH is equivalent to

\begin{equation}\mathop{\text{SoftmaxAttention}}(\underbrace{\boldsymbol{Q}-\mathop{\text{DeltaNet}}(\boldsymbol{Q},\boldsymbol{W},\boldsymbol{W})}_{\tilde{\boldsymbol{Q}}},\underbrace{\boldsymbol{K}-\mathop{\text{DeltaNet}}(\boldsymbol{K},\boldsymbol{W},\boldsymbol{W})}_{\tilde{\boldsymbol{K}}},\boldsymbol{V})\end{equation}

that is, using DeltaNet to add positional encoding to $\boldsymbol{Q},\boldsymbol{K}$. Viewed this way, PaTH (under the constraint $\Vert\boldsymbol{w}\Vert=\sqrt{2}$) amounts to a kind of intra-layer mixture of Softmax Attention and DeltaNet. Of course, we could also abandon the derivation above and simply use the formula above even when $\Vert\boldsymbol{w}\Vert\neq\sqrt{2}$ does not hold — this would resemble the approach of Canon Layers, adding positional information to $\boldsymbol{Q},\boldsymbol{K}$ via convolution, except that here the convolution is no longer a short convolution but rather the long convolution that is DeltaNet.

An Unconventional Path

Finally, let's look at another recent linear attention model that is equally worth paying attention to — MesaNet (there is also a broadly similar contemporaneous work, Atlas). The Online Learning perspective of TTT tells us that DeltaNet is essentially using SGD to optimize the objective function $\frac{1}{2}\Vert\boldsymbol{S}\boldsymbol{k} - \boldsymbol{v}\Vert^2$, and if we look carefully, we notice that $\boldsymbol{S}\boldsymbol{k}$ is just a linear function of $\boldsymbol{k}$ — so this is actually just a linear regression problem, and linear regression has a closed-form solution!

\begin{equation}\boldsymbol{S}_t = \boldsymbol{G}_t \boldsymbol{H}_t^{-1},\quad \boldsymbol{G}_t = \sum_{j=1}^t \boldsymbol{v}_j \boldsymbol{k}_j^{\top},\quad \boldsymbol{H}_t = \sum_{j=1}^t \boldsymbol{k}_j \boldsymbol{k}_j^{\top}\end{equation}

MesaNet is built by exploiting exactly this closed-form solution to construct a sequence model. The idea originates from Uncovering mesa-optimization algorithms in Transformers, while an efficient training scheme was worked out in MesaNet: Sequence Modeling by Locally Optimal Test-Time Training. Building on the formula above, MesaNet adds a forget gate to $\boldsymbol{G}_t,\boldsymbol{H}_t$, and then, when solving, adds a diagonal matrix $\boldsymbol{\Lambda}_t$ to avoid non-invertibility. The overall model is

\begin{equation}\boldsymbol{o}_t = \boldsymbol{G}_t (\boldsymbol{H}_t + \boldsymbol{\Lambda}_t)^{-1} \boldsymbol{q}_t,\quad \boldsymbol{G}_t = \gamma_t \boldsymbol{G}_{t-1} + \boldsymbol{v}_t \boldsymbol{k}_t^{\top},\quad\boldsymbol{H}_t = \gamma_t \boldsymbol{H}_{t-1} + \boldsymbol{k}_t \boldsymbol{k}_t^{\top}\end{equation}

Clearly, the complexity of $\boldsymbol{G}_t,\boldsymbol{H}_t$ with respect to sequence length is linear, so the computational complexity of $\boldsymbol{o}_t$ is also linear — hence MesaNet still falls within the scope of linear attention. And because of its closed-form solution, it can basically be guaranteed that in most cases it outperforms DeltaNet, or even Gated DeltaNet. From a signal-processing perspective, the difference between MesaNet and DeltaNet is the same as the difference between Recursive Least Squares and Least Mean Squares.

All of this sounds like nothing but advantages — so why does the author file MesaNet under "unconventional path"? In my view, MesaNet "lives by the closed-form solution, and dies by the closed-form solution": the closed-form solution is what makes it generally superior to DeltaNet, but it also gives a sense of having reached a dead end, because the moment you tweak the setup even slightly, there is almost no hope of finding a closed-form solution again. Looking across the whole history of mathematics, virtually every branch that depended on closed-form solutions has by now fallen out of favor, simply because closed-form solutions are far too rare and far too unrepresentative of the general case.

From an implementation standpoint, the matrix $\boldsymbol{H}_t + \boldsymbol{\Lambda}_t$ that MesaNet needs to invert is not triangular. Although $(\boldsymbol{H}_t + \boldsymbol{\Lambda}_t)^{-1} \boldsymbol{q}_t$ can still be turned into solving a linear system rather than computing an explicit inverse, the fact that it is not triangular still substantially increases the cost of solving it. How to parallelize the computation of all the $(\boldsymbol{H}_t + \boldsymbol{\Lambda}_t)^{-1} \boldsymbol{q}_t$ as cheaply as possible will remain a long-standing difficulty for MesaNet. The paper currently uses the "conjugate gradient method" to find an approximate solution — it works, but it is far from perfect.

Then there is the question of theoretical capability: MesaNet is not strictly superior to DeltaNet either. This is because MesaNet's update rule for $\boldsymbol{G}_t,\boldsymbol{H}_t$ is still just a simple moving average, and its matrix inversion does not involve interaction between tokens, so its capability ceiling is presumably below that of DeltaNet, which has the Delta Rule. Intuitively: MesaNet tries its best to remember everything about all the $\boldsymbol{k},\boldsymbol{v}$'s, and trying to keep "everything" can lead to somewhat blurry memory, whereas DeltaNet's principle is "clear out the old to make room for the new" — and precisely because of this "clearing out," it is able to achieve long-term, precise memory of certain content.

We can also understand this non-optimality through a specific example: so far, every attention variant except MesaNet allows the option of sharing K and V — "allows" meaning it is not necessarily optimal, but it can still be trained to produce non-trivial results. MesaNet, however, cannot do this, because if K and V are identical, MesaNet's $\boldsymbol{S}_t$ becomes the identity matrix at all times.

Overall, MesaNet is a model that is pleasing to look at, but its closed-form solution also adds to its complexity and limits its flexibility, leaving quite a bit of room still to be explored. Readers who want to learn more about building sequence models from linear regression can also read TTR, which discusses in detail sequence models built under various linear-regression objectives.

A Road Still Unfolding

This article has given a brief overview of the development of linear attention and introduced the mathematical principles behind some of the models. Linear attention started out by imitating Softmax Attention, but has gradually developed a character of its own, and has by now become a highly competitive approach to sequence modeling — one that has even, in turn, offered new ideas for the development of Softmax Attention itself. The whole process has been full of both fun and insight.