Modern Linear Attention - From Vanilla Linear Attention to KDA

This series of tutorial is based on

• Blog post: https://sustcsonglin.github.io/blog/2024/deltanet-1/
• Blog post: https://sustcsonglin.github.io/blog/2024/deltanet-2/
• Paper: https://arxiv.org/abs/2406.06484

I strongly recommend reading the original materials if you are good at algorithms and math.

1. Linear Attention

We are all familiar with traditional Softmax Attention, often referred to as Full Attention. It applies Softmax normalization to the attention similarity matrix, giving

$$o_t = \sum_{j=1}^t\frac{\exp{(q_t^Tk_j)}}{\sum_{l=1}^t\exp{(q_t^Tk_l)}}v_j$$

Here, the $q$ and $k$ vectors are row vectors taken from the original $Q$ and $K$ matrices, while the vectors in the equations are written as column vectors. From a vector perspective, the meaning of Attention becomes very clear: the output of the $t$-th token is simply a weighted sum. Since $q^Tk$ is a scalar, we can write

$$o_t = \sum_{j=1}^t w_j v_j$$

In other words, Attention computes a weighted sum of the $v$ vectors using a set of weights. Because of causal modeling, the output at position $t$ can only aggregate the first $t$ vectors. Therefore, if we temporarily ignore how $w_j$ is computed, Attention itself is linear. The real issue lies in the computation of the weights:

$$w_j=\frac{\exp{(q_t^Tk_j)}}{\sum_{l=1}^t\exp{(q_t^Tk_l)}}$$

This expression includes the Softmax function. Intuitively, we exponentiate the similarity between the query vector $q$ and each key vector $k$, then divide by the sum of the exponentiated similarity scores over the first $t$ tokens. This gives us the normalization. Notice that we are still summing over all preceding tokens.

Because of Softmax, we must first compute the query-key similarities, i.e. $\mathrm{softmax}(QK^T)$, which produces a $T\times T$ matrix, where $T$ is the sequence length. The idea behind Linear Attention is very simple: remove the Softmax first, giving

$$o_t = \sum_{j=1}^t (q_t^Tk_j)v_j\quad w_j=q_t^Tk_j$$

There is now no exponential normalization, which means we can apply a few mathematical tricks:

$$\begin{aligned} o_t&=\sum_{j=1}^t (q_t^Tk_j)v_j \\ &=\sum_{j=1}^tv_j(k_j^Tq_t) && q_t^Tk_j\text{ is a scalar} \\ &=\left(\sum_{j=1}^tv_jk_j^T\right)q_t && \text{Associativity} \end{aligned}$$

Define $S_t=\sum_{j=1}^tv_jk_j^T$. Then

$$o_t=S_tq_t\in \mathbb{R}^d$$

and

S_t=S_{t-1}+v_tk_t^T\in \mathbb{R}^{d\times d}$$. This immediately turns Attention into an RNN! The hidden state is $S_t$, which is updated by $v_tk_t^T$, while the output is obtained by mapping the query vector through the state. We now only need to store the state matrix $S$; we no longer need to keep every KV pair around in order to compute Attention. The space complexity drops from $O(Ld)$ to $O(d^2)$. In other words, the memory footprint no longer grows with sequence length. At the same time, the computational complexity changes from $O(L^2d)$ to $O(Ld^2)$. Linear Attention looks great, but simply removing Softmax causes a major performance drop. Although the computation above is equivalent to an RNN with a fixed-size $d\times d$ state, mixing all KV vectors into a single matrix destroys its ability to precisely retrieve a particular historical KV pair. One of the biggest advantages of Softmax Attention is that the entire history of KV pairs is exposed directly to the model, allowing it to freely decide which historical information to use. We can see this clearly by looking at the retrieval process. Suppose a query $q$ wants to retrieve the value associated with some key $k$. In other words, we want the output to be one exact historical value $v_j$: $$o_t = v_j$$. But our computation is $$o_t=\sum_{j=1}^tv_j(k_j^Tq_t)

To retrieve $v_j$ as accurately as possible, the ideal solution would satisfy

$$\begin{cases} k_i^Tq_t=1\quad i=j \\ k_i^Tq_t=0 \quad i\neq j \end{cases}$$

The closest natural choice is $q_t=k_j$. Then, without loss of generality, assume all vectors are normalized:

$$\begin{aligned} o_t&=\sum_{j=1}^tv_j(k_j^Tq_t) \\ &=v_jk_j^Tk_j + \sum_{i\neq j}v_i(k_i^Tk_j)\\ &=v_j + \underbrace{\sum_{i\neq j}v_i(k_i^Tk_j)}_{\text{retrieval error}} \end{aligned}$$

To retrieve $v_j$ exactly, we would need $k_i^Tk_j=0$ for every $i\neq j$, meaning that all key vectors would have to be mutually orthogonal. But in a $d$-dimensional space, we can have at most $d$ mutually orthogonal vectors. Once the sequence becomes longer than that, retrieval error becomes unavoidable.

You might ask: doesn’t Softmax Attention face the same problem? This is exactly where Softmax helps. When Softmax Attention wants to retrieve a particular value, it only needs to make $q_t^Tk_j$ much larger than the other similarity scores. Softmax will naturally concentrate most of the probability mass on the $j$-th weight, making the distribution almost one-hot:

$$\mathrm{softmax}([0.7,0.3,...,\underbrace{15}_{\text{j'th weight}},0.2])\approx [0, 0, ..., 1, 0]$$

So yes, Softmax is actually very useful (confirmed). The idea behind DeltaNet is to address this problem by minimizing retrieval error.

2. The Delta Rule

The idea behind the Delta Rule is simple: it is just gradient descent applied to a linear predictor. Suppose we have a linear predictor

$$\hat y = \sum_i w_i x_i$$

and use mean squared error as the loss function:

$$\mathcal{L}=\frac{1}{2}(y - \hat y)^2$$

Assuming $y$ is the ground-truth target, the gradient with respect to the weights is

$$\begin{aligned} \frac{\partial \mathcal{L}}{\partial w_i}&=\frac{\partial \mathcal{L}}{\partial \hat y}\frac{\partial \hat y}{\partial w_i}\\ &=-(y-\hat y)x_i \end{aligned}$$

By gradient descent,

$$\Delta w_i = -\eta \frac{\partial \mathcal{L}}{\partial w_i} = \eta (y-\hat y)x_i$$

which gives

$$w' = w - \eta (\hat y-y)x_i \tag{1}$$

From this, we can directly write down the core state-matrix update used by DeltaNet:

$$S_t= S_{t-1} - \beta_t (S_{t-1}k_t-v_t)k_t^T \tag{2}$$

Comparing this with Equation (1), the DeltaNet update rule is really performing online gradient descent. The input is $k_t$, the prediction is $S_{t-1}k_t$, the target is $v_t$, and the learning rate is $\beta_t$. In effect, DeltaNet trains its state matrix so that when the input is $k_t$, the corresponding $v_t$ can be retrieved from the state matrix as accurately as possible.

There is another way to understand DeltaNet. Suppose we define $S_{t-1}k_t$ as the previously stored, or “old,” value $v_t^{\text{old}}$. Then

$$\begin{aligned} S_t &= S_{t-1} - \beta_t (S_{t-1}k_t-v_t)k_t^T \\ &=S_{t-1} + \beta_t(v_t-v_t^{\text{old}})k_t^T \\ &=S_{t-1} + (\beta_t v_t - \beta_t v_t^{\text{old}})k_t^T \\ &=S_{t-1} + (\beta_tv_t - \beta_t v_t^{\text{old}}+v_t^{\text{old}}-v_{t}^{\text{old}})k_t^T\\ &=S_{t-1}+(\beta_tv_t+(1-\beta_t)v_t^{\text{old}}-v_t^{\text{old}})k_t^T \end{aligned}$$

Define

$$v_t^{\text{new}} = (1-\beta_t)v_t^{\text{old}} + \beta_tv_t \tag{3}$$

Then

$$\begin{aligned} S_t &= S_{t-1} + (v_t^{\text{new}}-v_t^{\text{old}})k_t^T \\ &=S_{t-1} + v_t^{\text{new}} k_t^T - v_t^{\text{old}}k_t^T \end{aligned}$$

Looking at Equation (3), we can see that the new value is simply a mixture of the old value and the current value. The state-matrix update can therefore be interpreted as erasing the old value from the state matrix and then writing back a mixture of the old and current values. When $\beta_t=0$, the old value is preserved. When $\beta_t=1$, the old value is completely erased and replaced by the current value.

Finally, Linear Attention itself can also be viewed as online gradient descent. Starting from the Linear Attention state update,

$$\begin{aligned} S_t&=S_{t-1}+v_tk_t^T \\ &=S_{t-1} + \eta_t v_tk_t^T && (\eta_t=1)\\ &=S_{t-1} - \eta_t\nabla (-\langle S_{t-1}k_t, v_t \rangle ) \end{aligned}$$

we can define the loss function as

$$\mathcal{L}(S)=-\langle Sk_t, v_t \rangle \quad \text{Dot Product}$$

If the vectors are normalized, this loss reaches its minimum value of $-1$ when the two vectors are maximally aligned. By contrast, the Delta Rule used in DeltaNet comes from an MSE loss:

$$\mathcal{L}(S)=\frac{1}{2}\| Sk_t - v_t \|^2$$

Comparing these two loss functions, we can also say that one difference between DeltaNet and Linear Attention is the objective used for the online update: DeltaNet uses an MSE loss, while Linear Attention corresponds to a linear loss based on the dot product.

3. Chunkwise Parallelism

One reason Attention became so popular is that training can be highly parallelized: almost all of the computation can be expressed as matrix multiplication, which scales extremely well on GPUs. An RNN, on the other hand, is iterative and usually requires a for-loop, so it cannot be fully parallelized. However, once the recurrence contains no nonlinear operation, certain parallel algorithms become possible.

A classic example is the prefix-sum problem. Suppose we want the sum of the first $t$ elements of a sequence:

$$S_t = \sum_{i=1}^t x_t = S_{t-1} + x_t$$

At first glance, this looks difficult to parallelize because the $t$-th prefix sum appears to require the $(t-1)$-th prefix sum first. But addition is associative. Suppose the sequence is $[a,b,c,d]$. What we want is the prefix sum at every position, together with the total sum. For example, the prefix before $a$ is $0$, the prefix before $b$ is $a$, the prefix before $c$ is $a+b$, and so on.

Computing the total sum is easy because of associativity:

$$a + b+c+d=(a+b)+(c+d)$$

The computation can therefore be arranged hierarchically.

We can compute $a+b$ and $c+d$ independently in parallel, then add the two results. This process is called the up-sweep. However, the up-sweep only gives us the total sum $a+b+c+d$. What we actually want is the prefix sum associated with every leaf at the bottom of the tree.

Now define the sum range of a node as

$$S_{l:r}=\sum_{i=l}^{r-1}x_i$$

and define the prefix sum of a node as

$$P_l = \sum_{i=0}^{l-1} x_i$$

which represents the sum of everything before the starting point of that node’s range. Then we can make two observations:

• The prefix of a left child is always equal to the prefix of its parent. For example, the node representing $a+b+c+d$ has prefix $0$, and the node representing $a+b$ also has prefix $0$. The node representing $c+d$ has prefix $a+b$, and its left child $c$ also has prefix $a+b$. This is because a left child always starts at the same position as its parent.

• The prefix of a right child is the parent’s prefix plus the sum range of the left child. For example, the node $c+d$ has prefix $0+a+b$. If the split point is $m$, the left child covers $S_{l:m}$ and the right child covers $S_{m:r}$, so $P_m=P_l + S_{l:m}$.

Therefore, we can propagate the prefix information from the root down the tree. We only need to follow two rules:

• The left child receives the parent’s prefix.
• The right child receives the parent’s prefix plus the left child’s sum range.

Following these rules gives us the prefix sum at every bottom-level leaf. This process is called the down-sweep. Together, the up-sweep and down-sweep compute the prefix sums of the entire sequence. This algorithm is known as the Blelloch Scan.

Can we use a similar parallel scan to compute DeltaNet? Yes. DeltaNet’s recurrence differs slightly from a standard prefix sum. Starting from Equation (2), we can rearrange it as

$$\begin{aligned} S_t &= S_{t-1} - \beta_t (S_{t-1}k_t-v_t)k_t^T \\ &=S_{t-1} - \beta_tS_{t-1}k_tk_t^T+\beta_t v_tk_t^T \\ &=S_{t-1}(I-\beta_t k_tk_t^T)+\beta_t v_tk_t^T \\ &=S_{t-1}M_t +X_t \end{aligned} \tag{4}$$

where

$$M_t=I-\beta_tk_tk_t^T\quad X_t=\beta_t v_tk_t^T$$

Compared with an ordinary prefix sum, each update now also multiplies the previous state by a matrix. Let us start with a simple example:

$$\begin{aligned} &S_1=a_1 \\ &S_2 = a_1M_2 + a_2 \\ &S_3 = a_1(M_3M_2) + (a_2M_3 + a_3) \\ &S_4 = a_1(M_4M_3M_2) + (a_2M_4M_3+a_3M_4 + a_4) \end{aligned}$$

Define the binary operator

$$[M_i,X_i]\oplus [M_j,X_j]=[M_iM_j,M_jX_i+X_j]\tag{5}$$

Then

$$\begin{aligned} &S_1=a_1 \\ &S_2 = a_1M_2 + a_2 = [M_2, a_2] \\ &S_3 = a_1(M_3M_2) + (a_2M_3 + a_3) = [M_2,a_2]\oplus [M_3,a_3] \\ &S_4 = a_1(M_4M_3M_2) + (a_2M_4M_3+a_3M_4 + a_4) = [M_2,a_2]\oplus [M_3,a_3] \oplus [M_4,a_4] \end{aligned}$$

We have now converted the recurrence into another prefix-sum problem. However, to apply a parallel scan, the operator must be associative. So we need to prove that the operator we just defined is associative.

Proof.

Rewrite Equation (4) in matrix-multiplication form:

$$S_t =S_{t-1}M_t +X_t \Rightarrow \begin{bmatrix}S_t & 1\end{bmatrix}=\begin{bmatrix}S_{t-1} & 1\end{bmatrix}\begin{bmatrix}M_t & 0 \\X_t & 1\end{bmatrix}$$

Expanding only the first component gives

$$\begin{aligned}\begin{bmatrix}S_t & 1\end{bmatrix}&=\begin{bmatrix}S_{t-1} & 1\end{bmatrix}\begin{bmatrix}M_t & 0 \\X_t & 1\end{bmatrix}\\&=\begin{bmatrix}S_{t-2} & 1\end{bmatrix}\begin{bmatrix}M_{t-1} & 0 \\X_{t-1} & 1\end{bmatrix}\begin{bmatrix}M_{t} & 0 \\X_{t} & 1\end{bmatrix} \\&=\begin{bmatrix}S_{t-3} & 1\end{bmatrix}\begin{bmatrix}M_{t-2} & 0 \\X_{t-2} & 1\end{bmatrix}\begin{bmatrix}M_{t-1} & 0 \\X_{t-1} & 1\end{bmatrix}\begin{bmatrix}M_{t} & 0 \\X_{t} & 1\end{bmatrix} \\&= \cdots \\&=\prod_t W_t\end{aligned}$$

For any two such matrices $W$, we have

$$\begin{bmatrix}M_{t-1} & 0 \\X_{t-1} & 1\end{bmatrix}\begin{bmatrix}M_{t} & 0 \\X_{t} & 1\end{bmatrix} = \begin{bmatrix}M_{t}M_{t-1} & 0 \\M_tX_{t-1}+X_t & 1\end{bmatrix}_{[:, 0]}=[M_{t-1},X_{t-1}]\oplus[M_t,X_t]$$

The product preserves the same structural form. Now we need to prove

$$\left(\prod_{t=1}^T W_t\right)_{[:, 0]}=\bigoplus_{t=1}^T [M_t, X_t]$$

When $T=1$,

$$(W_1)_{[:,0]}=\begin{bmatrix}M_{1} & 0 \\X_{1} & 1\end{bmatrix}_{[:,0]}=[M_1,X_1]$$

Assume that when $T=N$,

$$(\prod_{t=1}^NW_t)_{[:,0]}=\begin{bmatrix}\tilde M_{N} & 0 \\\tilde X_{N} & 1\end{bmatrix}_{[:,0]}=[\tilde M_N, \tilde X_N]=\bigoplus_{t=1}^N [M_t,X_t]$$

Then for $T=N+1$,

$$\begin{aligned}(\prod_{t=1}^{N+1}W_t)_{[:,0]}&=(W_{N+1}\prod_{t=1}^{N}W_t)_{[:,0]}\\&=\begin{bmatrix}\tilde M_{N} & 0 \\\tilde X_{N} & 1\end{bmatrix}\begin{bmatrix}M_{N+1} & 0 \\X_{N+1} & 1\end{bmatrix}_{[:,0]} \\&=\begin{bmatrix}\tilde M_NM_{N+1} & 0 \\M_{N+1}X_{N+1}+X_{N+1} & 1\end{bmatrix}_{[:,0]} \\&=[\tilde M_NM_{N+1}, M_{N+1}X_{N+1}+X_{N+1}]=\bigoplus_{t=1}^N [M_t,X_t] \oplus [M_{N+1}, X_{N+1}]\\&=\bigoplus_{t=1}^{N+1}[M_t,X_t]\end{aligned}$$

By induction, matrix multiplication is equivalent to the custom operator above. Since matrix multiplication is associative, our custom operator is associative as well.

Q.E.D.

Once associativity has been established, we can use a Blelloch-Scan-like algorithm to accelerate DeltaNet.

However, DeltaNet uses another parallel-scan algorithm called the Brent-Kung Scan. The key difference is that Brent-Kung reuses partial sums that were already computed during the up-sweep.

Suppose we want to compute

$$x_0+x_1+x_2+x_3+x_4+x_5+x_6+x_7$$

First, compute the following pairs independently:

$$p_1=x_0+x_1,p_3=x_2+x_3,p_5=x_4+x_5,p_7=x_6+x_7$$

Then combine them further:

$$p_3\leftarrow p_1+p_3=S_{0:3}$$

$$p_7\leftarrow p_5+p_7=S_{4:7}$$

Finally,

$$p_7\leftarrow=p_3+p_7=S_{0:7}$$

This completes the up-sweep. At this point, we already have several prefix sums:

$$P_1=p_1=S_{0:1},P_3=p_3=S_{0:3},P_7=p_7=S_{0:7}$$

Next, compute the remaining prefix sums in the down-sweep:

$$P_2=P_1+x_2=S_{0:2}$$

$$P_4=P_3+x_4=S_{0:4}$$

$$P_5=P_3+p_5=S_{0:5}$$

$$P_6=P_5+x_6=S_{0:6}$$

Now we have all prefix sums. So the idea behind the Brent-Kung Scan is to use results we already computed to fill in the prefixes that were not explicitly produced by the up-sweep. The computation of DeltaNet follows the same idea: first compute $[M_0,X_0]\oplus [M_1,X_1]$ and $[M_2,X_2]\oplus [M_3,X_3]$, then combine those results, and finally reuse the intermediate results to recover the remaining prefix.

However, DeltaNet does not only involve matrix multiplication by $M$. Repeated multiplication of the $M$ matrices introduces another problem. Consider multiplying two of them:

$$\begin{aligned} M_0M_1&=(I-\beta_0k_0k_0^T)(I-\beta_1k_1k_1^T) \\ &=I-\beta_1k_1k_1^T-\beta_0k_0k_0^T+\beta_0\beta_1k_0(k_0^Tk_1)k_1^T \end{aligned}$$

$$\begin{aligned} M_2M_3&=(I-\beta_2k_2k_2^T)(I-\beta_3k_3k_3^T) \\ &=I-\beta_3k_3k_3^T-\beta_2k_2k_2^T+\beta_2\beta_3k_2(k_2^Tk_3)k_3^T \end{aligned}$$

At the final stage of the up-sweep, we would need to compute

$$(I-\beta_1k_1k_1^T-\beta_0k_0k_0^T+\beta_0\beta_1k_0(k_0^Tk_1)k_1^T)(I-\beta_3k_3k_3^T-\beta_2k_2k_2^T+\beta_2\beta_3k_2(k_2^Tk_3)k_3^T)\tag{6}$$

Let us compare the cost of two ways of performing this computation. The first is to treat each $M$ as a dense matrix. Multiplying two $d\times d$ matrices costs $O(d^3)$. Alternatively, because each $M$ is represented as a sum of rank-1 terms, we can exploit that structure and reduce an individual product to roughly $O(d^2)$ by separately computing the required outer products and inner products before summing them.

The problem is that as the recurrence depth increases, the number of expanded terms also grows. Equation (6) already shows how each multiplication produces more terms, and this expansion grows exponentially with depth. So exploiting the low-rank structure directly does not seem especially attractive; in practice, it may not be much better than simply treating the matrices as dense.

The second problem is space complexity. A parallel scan needs to store intermediate matrices. For a sequence of length $L$, this requires $O(Ld^2)$ memory.

Let us consider a different approach. The state computation of Linear Attention can be written as

$$S_t= \sum_{i=1}^t v_ik_i^T$$

Because this is a sum of vector outer products, we do not need to store every intermediate state. But the computation is completely sequential. A compromise is to introduce intermediate checkpoints. Suppose we split a sequence of length $L$ into chunks of size $C$, giving $n=\lceil L/C \rceil$ chunks. We treat each chunk boundary as a checkpoint, and compute the state matrix at each boundary using matrix multiplication.

Define

$S_{[i]}:=S_{iC}\in \mathbb{R}^{d\times d}$, the state matrix at the beginning of chunk $i$,

$\square_{[i]}=\square_{iC+1:(i+1)C}\in \mathbb{R}^{C\times d}$ for $\square \in \{Q,K,V,O\}$, the matrix formed by the slice belonging to chunk $i$,

and $\square_{[i]}^r=\square_{iC+r}\in \{q,k,v,o,S\}$, the $r$-th element, or the $r$-th state matrix, inside chunk $i$.

Then the Linear Attention computation can be written as

$$S_{[i]}^r=S_{[i]}+\sum_{i=1}^r v_{[i]}^t{k_{[i]}^t}^T \quad \text{starting from the chunk boundary, accumulate the keys and values up to position $r$}$$

$$o_{[i]}^r=S_{[i]}q_{[i]}^r + \sum_{t=1}^r v_{[i]}^t ({k_{i}^t}^Tq_{[i]}^r)\quad \text{the output is obtained from the state matrix and the query $q$}$$

This allows us to write the computation in matrix form:

$$S_{[t+1]}=S_{[t]} + V_{[t]}^TK_{[t]}\quad \text{the state at the start of the next chunk is the current chunk's starting state plus the KV contribution of this chunk}$$

$$O_{[t]}=Q_{[t]}S_{[t]}^T+(Q_{[t]}K_{[t]}^T\odot M)V_{[t]} \quad \text{all outputs in the chunk are computed from the chunk's queries, state, and KV pairs}$$

Here, $M$ is the causal mask. With this formulation, we do not need to store every intermediate state, nor do we need to execute the entire computation sequentially. At the same time, the heavy computation is expressed as matrix multiplication, allowing us to take advantage of Tensor Cores.

Chunkwise parallelism for DeltaNet is slightly more involved. In the previous analysis, every multiplication of the $M$ matrices seemed to increase the number of terms exponentially. However, there is a much more compact representation of these products.

Let us expand consecutive products of $M$ again:

$$P_1 = M_1 = I - \beta_1k_1k_1^T\quad w_1:=\beta_1k_1$$

$$\begin{aligned} P_2 &= P_1M_2 \\ &= I - \beta_1k_1k_1^T-\beta_2k_2k_2^T + \beta_1\beta_2k_1(k_1^Tk_2)k_2^T\\ &=I - w_1k_1^T-\beta_2[k_2-w_1(k_1^Tk_2)]k_2^T \end{aligned}\quad w_2=\beta_2[k_2-w_1(k_1^Tk_2)]$$

$$\begin{aligned} P_3 &= P_2M_3 \\ &=(I - w_1k_1^T-w_2k_2^T)(I-\beta_3k_3k_3^T)\\ &=I-w_1k_1^T-w_2k_2^T-\beta_3k_3k_3^T+\beta_3w_1k_1^Tk_3k_3^T+\beta_3w_2k_2^Tk_3k_3^T \\ &=I - w_1k_1^T - w_2k_2^T-\beta_3[k_3-w_1k_1^Tk_3-w_2k_2^Tk_3]k_3^T \end{aligned}\quad w_3 :=\beta_3[k_3-w_1k_1^Tk_3-w_2k_2^Tk_3]$$

This suggests the pattern

$$w_n:=\beta_n\left[k_n-\sum_{i=1}^{n-1}w_i(k_i^Tk_n)\right]$$

which allows us to write a product of consecutive $M$ matrices as

$$P_n = I - \sum_{i=1}^n w_ik_i^T$$

Let us prove this.

Proof.

When $n=1$,

$P_1 = M_1 = I - w_1k_1^T\quad w_1:=\beta_1k_1$

which is clearly true. Assume the expression holds for $n=N$. Then for $n=N+1$,

$\begin{aligned}P_{N+1}&=P_NM_{N+1} \\&=(I - \sum_{i=1}^N w_ik_i^T)(I-\beta_{N+1}k_{N+1}k_{N+1}^T)\\&=I - \beta_{N+1}k_{N+1}k_{N+1}^T - \sum_{i=1}^N w_ik_i^T + (\sum_{i=1}^N w_ik_i^T)\beta_{N+1}k_{N+1}k_{N+1}^T \\&=I - \sum_{i=1}^N w_ik_i^T - \beta_{N+1}\left[k_{N+1} - \sum_{i=1}^N w_i(k_i^Tk_{N+1})\right]k_{N+1}^T \\&=I - \sum_{i=1}^N w_ik_i^T -w_{N+1}k_{N+1}^T \\&=I - \sum_{i=1}^{N+1}w_ik_i^T\end{aligned}$

By induction,

$P_n = I - \sum_{i=1}^n w_ik_i^T$.

Q.E.D.

Now return to Equation (4):

$$\begin{aligned} S_t &=S_{t-1}(I-\beta_t k_tk_t^T)+\beta_t v_tk_t^T \end{aligned}$$

Notice that

$$S_1=\beta_1v_1k_1^T\quad u_1:=\beta_1v1$$

and

$$\begin{aligned} S_2 &= u_1k_1^T(I-\beta_2k_2k_2^T)+\beta_2v_2k_2^T \\ &=u_1k_1^T + \beta_2[v_2-u_1(k_1^Tk_2)]k_2^T \end{aligned}\quad u_2:=\beta_2[v_2-u_1(k_1^Tk_2)]$$

Similarly,

$$\begin{aligned} S_3 &= (u_1k_1^T+u_2k_2^T)(I-\beta_3k_3k_3^T)+\beta_3v_3k_3^T\\ &=u_1k_1^T+u_2k_2^T+\beta_3[v_3-u_1(k_1^Tk_3)-u_2(k_2^Tk_3)]k_3^T \end{aligned}\quad u_3:=\beta_3[v_3-u_1(k_1^Tk_3)-u_2(k_2^Tk_3)]$$

This suggests

$$u_n=\beta_n\left[v_n - \sum_{i=1}^{n-1}u_i(k_i^Tk_n)\right]$$

and therefore

$$S_t=\sum_{i=1}^tu_tk_t^T$$

Let us prove this as well.

Proof.

When $t=1$,

$$S_1=\beta_1v_1k_1^T=u_1k_1^T$$

Assume the expression holds for $t=N$. Then for $t=N+1$,

$$\begin{aligned}S_{N+1}&=\left(\sum_{t=1}^N u_tk_t^T\right)(I-\beta_{N+1}k_{N+1}k_{N+1}^T)+\beta_{N+1}v_{N+1}k_{N+1}^T \\&=\sum_{t=1}^N u_tk_t^T - \left(\sum_{t=1}^t u_tk_t^T\right)(\beta_{N+1}k_{N+1}k_{N+1}^T) +\beta_{N+1}v_{N+1}k_{N+1}^T\\&=\sum_{t=1}^N u_tk_t^T + \beta_{N+1}[v_{N+1}-\sum_{t=1}^N u_t(k_t^Tk_{N+1})]k_{N+1}^T\\&=\sum_{t=1}^N u_tk_t^T + u_{N+1}k_{N+1}^T \\&=\sum_{t=1}^{N+1}u_tk_t^T\end{aligned}$$

By induction,

S_t=\sum_{i=1}^tu_tk_t^T$$.*Q.E.D.*

We can now derive the chunkwise-parallel form of DeltaNet. Suppose we again split the sequence into chunks of size $C$. Expanding the DeltaNet state recurrence within a chunk gives

$$\begin{aligned} S_{[i]}^r &= S_{[i]}^{r-1}M_{[i]}^r + X_{[i]}^r \\ &=(S_{[i]}^{r-2}M_{[i]}^{r-1}+X_{[i]}^{r-1})M_{[i]}^r + X_{[i]}^r \\ &=S_{[i]}^{r-2}M_{[i]}^{r-1}M_{[i]}^r + X_{[i]}^{r-1}M_{[i]}^r + X_{[i]}^r \\ &=\cdots \\ &=S_{[i]}\prod_{t=1}^r M_{[i]}^t + \sum_{t=1}^rX_{[i]}^t\prod_{s=t+1}^rM_{[i]}^s\\ &=S_{[i]}\underbrace{\prod_{t=1}^r(I-\beta_{[i]}^tk_{[i]}^t {k_{[i]}^t}^T)}_{\text{Product of M}} + \underbrace{\sum_{t=1}^r\left( \beta_{[i]}^tv_{[i]}^t {k_{[i]}^t}^T\prod_{s=t+1}^r(I-\beta_{[i]}^sk_{[i]}^s {k_{[i]}^s}^T) \right)}_{\text{Expansion of S}} \end{aligned}$$

Notice that the second term is exactly the form obtained by expanding the state matrix for $r$ steps: it is a sum of $X$ terms multiplied by subsequent $M$ matrices. Using the two compact expressions derived above, we get

$$S_{[i]}^r=S_{[i]}\left(I-\sum_{t=1}^rw_{[i]}^t {k_{[i]}^t}^T\right) + \sum_{t=1}^ru_{[i]}^t {k_{[i]}^t}^T$$

where

$$w_{[t]}^r = \beta_{[t]}^r\left( k_{[t]}^r - \sum_{i=1}^{r-1}w_{[t]}^i (k_{[t]}^i)^Tk_{[t]}^r \right)$$

and

$$u_{[t]}^r = \beta_{[t]}^r\left( v_{[t]}^r - \sum_{i=1}^{r-1}u_{[t]}^i (k_{[t]}^i)^Tk_{[t]}^r \right)$$

Similarly, the output can be written as

$$\begin{aligned} o_{[i]}^r &= S_{[i]}^rq_{[i]}^r \\ &=S_{[i]}q_{[i]}^r + \sum_{t=1}^r(u_{[i]}^t - S_{[i]}w_{[i]}^t)({k_{[i]}^t}^Tq_{[i]}^r) \end{aligned}$$

In matrix form,

$$S_{[t+1]} = S_{[t]}(I-W_{[i]}K_{[i]}) + U_{[i]}^TK_{[i]} = S_{[i]} + \left(U_{[i]} - W_{[i]}S_{[i]}^T\right)^TK_{[i]}$$

and

$$O_{[i]}=Q_{[i]}S_{[i]}^T + (Q_{[i]}K_{[i]}^T\odot M)(U_{[i]} - W_{[i]}S_{[i]}^T)$$

So DeltaNet can also be accelerated using a chunkwise formulation. Chunkwise parallelism not only lets us exploit fast matrix multiplication, but also avoids storing every intermediate state matrix. It is essentially a trade-off between parallelism and space complexity.

There is still one remaining problem: computing $w_{[t]}^r$ and $u_{[t]}^r$ appears to be sequential. Fortunately, both variables can be solved for directly.

Starting from

$$w_r = \beta_r k_r - \beta_r \sum_{i<r} w_i (k_i^Tk_r)$$

define

A_{ri}=-\beta_r k_i^Tk_r$$, where $A$ is a strictly lower-triangular matrix \((i<r)\). Then $$w_r=\beta_r k_r + \sum_{i<r} A_{ri}w_i

Now write everything in matrix form:

$$K = \begin{bmatrix} k_1^T \\ \vdots \\ k_C^T \end{bmatrix}$$

$$W = \begin{bmatrix} w_1^T \\ \vdots \\ w_C^T \end{bmatrix}$$

$$B = \mathrm{diag}(\beta_1,\dots, \beta_C)$$

Then

$$W=BK+AW \Rightarrow W = (I-A)^{-1}BK$$

Similarly,

$$U=(I-A)^{-1}BV$$

So once again, the computation can be accelerated with matrix multiplication.