Haoran Zhang

Linear Attention, Visualized

By

In July 2026, Moonshot released Kimi K3, a 2T-parameter model with a 1M-token context window. K3 uses Moonshot's linear-attention variant, Kimi Delta Attention (KDA). Moonshot reports that it trails only Claude Fable 5 and GPT-5.6 Sol across its evaluation suite.

NVIDIA had already made a similar choice in Nemotron 3, which combines softmax attention with Mamba-2 recurrent layers. KDA and Mamba-2 use different recurrences. Both offer linear-time sequence mixing during training and fixed-state decoding, although this efficiency comes with some loss of expressiveness.

In this post, we visualize the path from self-attention to the linear-attention variants used in recent models.

Self-attention recap

Given hidden states $H$, three learned projections produce queries, keys, and values:

$$Q=HW_Q,\qquad K=HW_K,\qquad V=HW_V,$$

where $Q,K\in\mathbb R^{n\times d_k}$ and $V,O\in\mathbb R^{n\times d_v}$, with $n$ denoting the sequence length.

$$O = \operatorname{Attention}(Q,K,V) = \operatorname{softmax}\!\left( \frac{QK^{\mathsf T}}{\sqrt{d_k}}+M_{\mathrm{add}} \right)V, \qquad M_{\mathrm{add}} = \begin{bmatrix} 0&-\infty&-\infty\\ 0&0&-\infty\\ 0&0&0 \end{bmatrix}.$$

Each query $q_t$ represents a linear functional on the key space:

$$\phi_{q_t}:\mathbb R^{d_k}\to\mathbb R, \qquad k_j\mapsto\langle q_t,k_j\rangle.$$

Evaluating it at $k_j$ produces the scalar score $\langle q_t,k_j\rangle$. Softmax normalizes these scores into weights, and $o_t$ is the weighted sum of the corresponding values.

Self-attention is quadratic

The attention operation uses two matrix multiplications over the sequence:

$$\begin{aligned} QK^{\mathsf T}:\quad (n\times d)(d\times n)&\longrightarrow(n\times n), &\approx 2n^2d\ \text{FLOPs},\\ \operatorname{softmax}(\cdots)V:\quad (n\times n)(n\times d)&\longrightarrow(n\times d), &\approx 2n^2d\ \text{FLOPs}. \end{aligned}$$

Masking and softmax add $\mathcal O(n^2)$ work, so the attention operation costs

$$4n^2d+\mathcal O(n^2)=\mathcal O(n^2d),$$

which is quadratic in the sequence length.

A direct implementation stores the $n\times n$ score matrix. FlashAttention avoids materializing this matrix in off-chip memory, but the computation remains quadratic.

Naïve linear attention

The quadratic cost comes from multiplying $Q$ by $K^{\mathsf T}$ first, which produces an $n\times n$ score matrix. If we remove softmax, matrix multiplication is associative, so we can change the order of computation:

$$O=(QK^{\mathsf T})V=Q(K^{\mathsf T}V).$$

The left-hand order compares every query with every key before mixing the values. The right-hand order first summarizes all key–value pairs in the $d\times d$ matrix $K^{\mathsf T}V$, then applies each query to that summary. It never constructs an $n\times n$ matrix.

The two matrix multiplications now cost

$$\begin{aligned} K^{\mathsf T}V:\quad (d\times n)(n\times d)&\longrightarrow(d\times d), &\approx 2nd^2\ \text{FLOPs},\\ Q(K^{\mathsf T}V):\quad (n\times d)(d\times d)&\longrightarrow(n\times d), &\approx 2nd^2\ \text{FLOPs}. \end{aligned}$$

Together, they require approximately $4nd^2=\mathcal O(nd^2)$ work, which is linear rather than quadratic in the sequence length.

This computation is not causal: the shared summary $K^{\mathsf T}V$ contains key–value pairs from the entire sequence, so every query can read the future. The causal version is

$$O=\bigl(QK^{\mathsf T}\odot M\bigr)V, \qquad M= \begin{bmatrix} 1&0&0\\ 1&1&0\\ 1&1&1 \end{bmatrix}, \qquad \odot:\ \text{Hadamard product}.$$

At first glance, the causal mask appears to prevent direct reassociation of the full matrix product. However, row $t$ of $O$ depends only on $q_t$, $K_{1:t}$, and $V_{1:t}$, as shown below.

Once the masked columns are removed, row $t$ becomes an ordinary matrix product over a prefix:

$$o_t =q_tK_{1:t}^{\mathsf T}V_{1:t} =q_t\!\left(K_{1:t}^{\mathsf T}V_{1:t}\right).$$

We call the product of the key and value prefixes the state $S_t$:

$$S_t :=K_{1:t}^{\mathsf T}V_{1:t} \in\mathbb R^{d\times d}.$$

$S_t$ is a fixed-size linear map from key space to value space. Moving from one token to the next adds exactly one outer product, so we can build this state recurrently instead of recomputing the entire prefix:

$$S_0=0, \qquad S_t=S_{t-1}+k_t^{\mathsf T}v_t, \qquad o_t=q_tS_t.$$

Each recurrent step costs $\mathcal O(d^2)$, so processing $n$ tokens costs $\mathcal O(nd^2)$. The cost is linear in the sequence length, hence the name linear attention.

The recurrent form reduces the arithmetic complexity, but it uses GPUs inefficiently:

  1. Large activation memory. A direct autograd implementation saves every $d\times d$ state $S_t$ for the backward pass. The $n$ saved states require $\mathcal O(nd^2)$ memory.
  2. Sequential token updates. Computing $S_t$ requires $S_{t-1}$. This dependency prevents the GPU from updating different tokens in parallel.
  3. Poor Tensor Core utilization. Tensor Cores accelerate tiled matrix multiply-accumulate operations. The outer product $k_t^{\mathsf T}v_t$ has shape $(d\times1)(1\times d)$, so its reduction dimension is only one. This shape does not fill the Tensor Core tiles along the reduction dimension.

Chunkwise algorithm

To solve these problems, we introduce the chunkwise algorithm. It uses the parallel form $\bigl(QK^{\mathsf T}\odot M\bigr)V$ within each chunk and the recurrent form $S_{\mathrm{out}}=S_{\mathrm{in}}+K^{\mathsf T}V$ between chunks.

Assume $C$ divides $n$. Split the sequence into $n/C$ chunks of $C$ consecutive tokens. Let $Q_r,K_r,V_r\in\mathbb R^{C\times d}$ denote chunk $r$.

Block the causal matrix

Let $A=QK^{\mathsf T}\odot M$. Partition $A$ into $C\times C$ blocks and partition $V$ and $O$ into $C\times d$ blocks. Let $A_{r,s}$ denote block $(r,s)$ of $A$. In block row $r$, blocks to the left contain earlier chunks, the diagonal block contains the current chunk, and blocks to the right contain future tokens.

A n×n in C×C blocks
×
V C×d blocks
=
O C×d blocks
history local future

The output for chunk $r$ splits into history and local contributions:

$$O_r =\underbrace{\sum_{s<r}A_{r,s}V_s}_{\text{history}} +\underbrace{A_{r,r}V_r}_{\text{local}}.$$

The local block keeps a $C\times C$ causal mask because tokens in the chunk have different prefix lengths. The history blocks need no mask because every token in chunk $r$ may read every token in an earlier chunk.

Compress the history

Let $K_{<r}$ and $V_{<r}$ contain all chunks before chunk $r$. Associativity compresses the dense history:

$$\sum_{s<r}A_{r,s}V_s =Q_rK_{<r}^{\mathsf T}V_{<r} =Q_rS_r, \qquad S_r:=K_{<r}^{\mathsf T}V_{<r}.$$

$S_r$ contains all key–value updates before chunk $r$. Let $M_c$ denote the $C\times C$ causal mask inside the chunk. The chunk output is

$$O_r =\underbrace{Q_rS_r}_{\text{history}} +\underbrace{\bigl(Q_rK_r^{\mathsf T}\odot M_c\bigr)V_r}_{\text{local}}.$$

Update the state

After computing $O_r$, add the current chunk to the history:

$$S_{r+1}=S_r+K_r^{\mathsf T}V_r.$$

The matrix product $K_r^{\mathsf T}V_r$ batches $C$ rank-one updates into a dense GEMM. Its reduction dimension is $C$ rather than $1$, so it uses Tensor Cores efficiently.

For the recurrent state, autograd now saves one boundary state per chunk. These $n/C$ states require $\mathcal O((n/C)d^2)$ memory. Chunk boundaries still form a sequential chain, but the chain now contains $n/C$ steps instead of $n$.

Cost

Each chunk requires

$$\mathcal O\!\left( \underbrace{C^2d}_{\text{local}} +\underbrace{Cd^2}_{\text{state update}} \right).$$

Across $n/C$ chunks, the total cost is

$$\mathcal O\!\left(nCd+nd^2\right).$$

Larger chunks increase the local $nCd$ term but shorten the sequential state chain. $C=1$ recovers the recurrent form, while $C=n$ recovers the full parallel form. Every valid $C$ provides exact results. FLA uses $C=64$, while FlashKDA uses $C=16$.

RetNet: scalar decay

In Naïve linear attention, each outer-product write $k_t^{\mathsf T}v_t$ enters the state with coefficient one and never decays. New and old writes therefore receive the same age-independent coefficient. RetNet introduces a recency bias by shrinking the existing state before adding the current write:

$$S_0=0, \qquad S_t=\gamma S_{t-1}+k_t^{\mathsf T}v_t, \qquad o_t=q_tS_t, \qquad 0<\gamma\le 1.$$

$\gamma$ is a fixed scalar, so it scales the entire state uniformly. The figure follows one key–value write through the recurrence:

A write from $s$ steps earlier has weight $\gamma^s$. The weight decreases exponentially with age.

Parallel form

To compute the recurrence in parallel, note that a write at position $j$ undergoes $t-j$ decay steps before query $q_t$ reads it. Its coefficient is therefore $\gamma^{t-j}$. Collect these coefficients in a decayed causal matrix:

$$\Gamma[t,j] = \begin{cases} \gamma^{t-j}, & j\le t,\\ 0, & j>t. \end{cases}$$

The diagonal is one because the current write has not decayed. Moving one column to the left adds one decay step.

The parallel form replaces the binary causal mask $M$ with $\Gamma$:

$$O=\bigl(QK^{\mathsf T}\odot\Gamma\bigr)V.$$

Chunkwise form

The chunkwise form again separates incoming history from interactions inside the chunk. Let $S_r$ denote the state immediately before chunk $r$, and let $\ell=1,\ldots,C$ index positions in the current chunk.

Decay incoming history

Before $q_{r,\ell}$ (the query at local position $\ell$ in chunk $r$) reads $S_r$, the state crosses $\ell$ decay steps. Define the boundary-to-query factors

$$\lambda_\ell:=\gamma^\ell, \qquad \lambda=(\gamma,\gamma^2,\ldots,\gamma^C)^{\mathsf T}.$$

The first illustration applies one entry of $\lambda$ to each row of $Q_rS_r$. (Hover a row of $Q_r$ to trace its decay factor and the resulting history row.)

We can move the row scale to $Q_r$ and define the decayed query:

$$\overleftarrow{Q}_r := \operatorname{diag}\!\left(\gamma,\gamma^2,\ldots,\gamma^C\right)Q_r.$$

The left arrow denotes left-boundary decay: row $\ell$ includes the factor $\gamma^\ell$ from the chunk boundary to $q_{r,\ell}$.

This gives the history term

$$O_r^{\mathrm{hist}}=\overleftarrow{Q}_rS_r.$$

Compute the local block

Partitioning the full parallel computation into $C\times C$ blocks exposes the same history and local split:

(QKT ⊙ Γ)decayed scores · n×n → C×C blocks
×
VC×d blocks
=
OC×d blocks
diagonal = local
below diagonal = history
above diagonal = future

A write at local position $m$ crosses $\ell-m$ decay steps before $q_{r,\ell}$ reads it. The local decay matrix is

$$\Gamma_C[\ell,m] = \begin{cases} \gamma^{\ell-m}, & m\le\ell,\\ 0, & m>\ell. \end{cases}$$

The local term is

$$O_r^{\mathrm{local}} =\bigl(Q_rK_r^{\mathsf T}\odot\Gamma_C\bigr)V_r.$$

Update the state

We now have the chunk output $O_r$. Next, we construct the state passed to the next chunk. The incoming state crosses all $C$ steps, while a write at local position $\ell$ crosses only $C-\ell$ steps. Define the write-to-boundary factors

$$\rho_\ell:=\gamma^{C-\ell}, \qquad \rho=(\gamma^{C-1},\gamma^{C-2},\ldots,1)^{\mathsf T}.$$

Because token rows of $K_r$ become columns of $K_r^{\mathsf T}$, $\rho$ scales columns in the state update. (Hover a column of $\Delta S_r$ to trace the decayed keys and the corresponding column of $V_r$.)

Write that column scaling as a diagonal matrix and absorb it into the keys:

$$\overrightarrow{K}_r := \operatorname{diag}\!\left(\gamma^{C-1},\gamma^{C-2},\ldots,1\right)K_r.$$

The right arrow denotes right-boundary decay: key row $\ell$ includes the factor $\gamma^{C-\ell}$ from its position to the chunk's right boundary.

Complete chunkwise form

Combining the history and local terms gives the complete chunkwise equations:

$$\begin{aligned} O_r &=\underbrace{\overleftarrow{Q}_rS_r}_{\text{history}} +\underbrace{\bigl(Q_rK_r^{\mathsf T}\odot\Gamma_C\bigr)V_r}_{\text{local}},\\ S_{r+1} &=\gamma^CS_r+\overrightarrow{K}_r^{\mathsf T}V_r. \end{aligned}$$

The arrows identify the relevant boundary. Queries absorb decay from the left boundary, while keys absorb decay toward the right boundary. The incoming history state decays from the left boundary to each query position; each key–value write made inside the chunk decays from its position to the right boundary. This exact regrouping keeps the same $\mathcal O(nCd+nd^2)$ asymptotic cost as the Naïve chunkwise algorithm.

Mamba-2: input-dependent scalar decay

RetNet applies the same decay $\gamma$ at every position. Mamba-2 replaces it with a token-dependent scalar $a_t$. Its recurrent form is

$$S_t=a_tS_{t-1}+k_t^{\mathsf T}v_t, \qquad o_t=q_tS_t, \qquad 0<a_t<1.$$

Each token supplies one scalar for the entire state. A value near one preserves history; a value near zero nearly clears it.

Parallel form

The outer-product write $k_j^{\mathsf T}v_j$ enters the state at position $j$. To reach $S_t$, subsequent transitions multiply it by $a_{j+1},a_{j+2},\ldots,a_t$. Its survival coefficient is

$$a_{j+1}a_{j+2}\cdots a_t.$$

Collect these survival coefficients in the decay matrix $\Gamma$. Entry $\Gamma[t,j]$ scales the score between $q_t$ and $k_j$. (Hover an entry of $\Gamma$ to trace the corresponding query, key, value, and output.)

For an inclusive interval, define

$$\lambda[i:j]:=\prod_{s=i}^{j}a_s, \qquad \text{empty product }(i>j)\text{ returns }1.$$

The decay matrix becomes

$$\Gamma[t,j] = \begin{cases} \lambda[j+1:t], & j\le t,\\ 0, & j>t. \end{cases}$$

The parallel equation keeps the same form as RetNet:

$$O=\bigl(QK^{\mathsf T}\odot\Gamma\bigr)V.$$

Chunkwise form

The chunkwise derivation follows RetNet. Each fixed power of $\gamma$ is replaced by a product of token-dependent gates.

Within chunk $r$, let $a_{r,\ell}$ denote the gate at local position $\ell$, and define

$$\lambda_r[i:j]:=\prod_{s=i}^{j}a_{r,s}.$$

Decay incoming history

Before $q_{r,\ell}$ reads the incoming state $S_r$, the recurrence multiplies $S_r$ by $\lambda_r[1:\ell]$. As in RetNet, absorb this scalar into the query row:

$$\overleftarrow q_{r,\ell} :=\lambda_r[1:\ell]q_{r,\ell}, \qquad O_r^{\mathrm{hist}}=\overleftarrow Q_rS_r.$$

The left arrow again denotes decay from the chunk's left boundary to each query position.

Compute the local block

The write $k_{r,m}^{\mathsf T}v_{r,m}$ is multiplied by $\lambda_r[m+1:\ell]$ before it contributes to the output at local position $\ell$. Therefore,

$$\Gamma_{C,r}[\ell,m] = \begin{cases} \lambda_r[m+1:\ell], & m\le\ell,\\ 0, & m>\ell. \end{cases}$$

The local output is

$$O_r^{\mathrm{local}} =\bigl(Q_rK_r^{\mathsf T}\odot\Gamma_{C,r}\bigr)V_r.$$

Update the state

We now have $O_r$. To construct the state for the next chunk, multiply the incoming state by $\lambda_r[1:C]$. The write at local position $\ell$ is multiplied only by the subsequent gates, whose product is $\lambda_r[\ell+1:C]$. Absorb this suffix product into its key:

$$\overrightarrow k_{r,\ell} :=\lambda_r[\ell+1:C]k_{r,\ell}.$$

The right arrow denotes decay from each write position to the chunk's right boundary. The state update is

$$S_{r+1} =\lambda_r[1:C]S_r+\overrightarrow K_r^{\mathsf T}V_r.$$

Complete chunkwise form

Combining the output and state equations gives

$$\begin{aligned} O_r &=\underbrace{\overleftarrow Q_rS_r}_{\text{history}} +\underbrace{\bigl(Q_rK_r^{\mathsf T}\odot\Gamma_{C,r}\bigr)V_r}_{\text{local}},\\ S_{r+1} &=\lambda_r[1:C]S_r+\overrightarrow K_r^{\mathsf T}V_r. \end{aligned}$$

Mamba-2 therefore keeps RetNet's scalar state transition but lets each token choose how much history survives. The gate remains scalar, so it scales every entry of the state equally.

DeltaNet: correcting memory collisions

Mamba-2 decays the entire history at once. Sometimes we instead want to correct the association addressed by the current key without uniformly decaying the entire memory. DeltaNet provides this more selective update.

Why additive writes collide

Treat $S$ as a key–value associative memory. For each association $(k_i,v_i)$, we would like to retrieve $v_i$ by reading the memory with $k_i$:

$$k_iS\approx v_i.$$

However, for a unit-norm key $k_i$, the read contains

$$\begin{aligned} k_iS &=\sum_{j=1}^{m}(k_i k_j^{\mathsf T})v_j\\ &=v_i+\underbrace{\sum_{j\ne i}(k_i k_j^{\mathsf T})v_j}_{\text{cross-talk}}. \end{aligned}$$

The first term is the intended value. The other terms come from keys that overlap with $k_i$. Orthogonal keys eliminate this interference, but learned keys generally are not mutually orthogonal. A raw outer-product write can therefore collide with associations already stored in $S$.

DeltaNet first asks what value the current state associates with $k_t$:

$$\widehat v_t=k_tS_{t-1},$$

The residual $v_t-\widehat v_t$ points from this prediction toward the desired value $v_t$. DeltaNet writes a step in that direction at $k_t$:

$$S_t =S_{t-1}+\beta_t k_t^{\mathsf T}(v_t-\widehat v_t), \qquad 0<\beta_t\le 1.$$

An intuitive example

Consider a simplified memory with two associations. The labels below stand for key and value vectors. Assume the Alice and Bob keys are orthogonal and unit norm, so a write at Bob's key does not change the value read at Alice's key.

Current memory
Aliceat school
Bobat home
+
New association
Bobat the office

Reading the state with Bob's key first returns “at home.” The target is “at the office,” so DeltaNet writes the difference between those two value vectors at Bob's key. Reading Bob after the update gives

$$k_{\mathrm{Bob}}S_{\mathrm{new}} =(1-\beta)\,v_{\mathrm{home}} +\beta\,v_{\mathrm{office}}.$$

If $\beta=1$, the new value replaces the old prediction exactly. A smaller $\beta$ makes a partial correction.

More generally, when $\lVert k_t\rVert_2=1$,

$$k_tS_t=(1-\beta_t)\widehat v_t+\beta_t v_t.$$

The unit-norm condition matters. Without it, reading the update back with $k_t$ introduces an additional factor $\lVert k_t\rVert_2^2$.

The example explains one token update. We now derive the recurrent and full-sequence matrix forms.

Recurrent form

Define the scaled correction

$$u_t:=\beta_t(v_t-\widehat v_t).$$

The complete token recurrence is

$$\begin{aligned} \widehat v_t&=k_tS_{t-1},\\ u_t&=\beta_t(v_t-\widehat v_t),\\ S_t&=S_{t-1}+k_t^{\mathsf T}u_t,\\ o_t&=q_tS_t. \end{aligned}$$

Once $u_t$ is known, the state has the same additive form as Naïve linear attention:

$$S_t=\sum_{i=1}^{t}k_i^{\mathsf T}u_i.$$

Unlike the raw values $v_t$, the corrections $u_t$ are causally coupled: each $u_t$ depends on earlier corrections through $S_{t-1}$. We therefore cannot form all rows of $U$ independently as in Naïve linear attention. We first collect this dependency into a triangular system and solve for $U$; only then can we reuse the linear-attention readout.

The triangular correction system

Substitute

$$S_{t-1}=\sum_{i<t}k_i^{\mathsf T}u_i$$

into the prediction. Then

$$u_t =\beta_t\left(v_t-\sum_{i<t}(k_tk_i^{\mathsf T})u_i\right).$$

For the first three tokens,

$$\begin{aligned} u_1&=\beta_1v_1,\\ u_2&=\beta_2\left(v_2-(k_2k_1^{\mathsf T})u_1\right),\\ u_3&=\beta_3\left(v_3-(k_3k_1^{\mathsf T})u_1-(k_3k_2^{\mathsf T})u_2\right). \end{aligned}$$

Only earlier corrections appear on the right. Define the strictly lower-triangular matrix

$$L[t,i] =\begin{cases} k_tk_i^{\mathsf T},&i<t,\\ 0,&i\ge t. \end{cases}$$

Let $M$ be the inclusive causal mask and $M^-:=M-I$ its strictly lower-triangular part. Then

$$L=(KK^{\mathsf T})\odot M^-.$$

Stack $u_t$ into $U\in\mathbb R^{n\times d}$ and define $B:=\operatorname{diag}(\beta)$. The row recurrences become (as you should verify):

$$\begin{aligned} U&=B(V-LU),\\ (I+BL)U&=BV. \end{aligned}$$

Equivalently,

$$U=(I+BL)^{-1}BV.$$

Solving for $U$

We do not compute $(I+BL)^{-1}$ explicitly. Forming a dense $n\times n$ inverse costs $\mathcal O(n^3)$.

Set $A:=I+BL$ and $R:=BV$. Since $A$ is unit lower triangular, forward substitution solves

$$AU=R$$

one row at a time. At row $t$, the earlier rows of $U$ are already known, and $u_t$ is the only new row. (Hover a row to follow the known inputs and current unknown.)

For $U\in\mathbb R^{n\times d}$, row $t$ costs $\mathcal O(td)$ and the full solve costs $\mathcal O(n^2d)$.

Parallel output

After solving for $U$, the state is a sum of key–correction outer products. The causal output therefore reuses the Naïve linear-attention equation with $U$ in place of $V$:

$$\begin{aligned} O_{\mathrm{naïve}}&=(QK^{\mathsf T}\odot M)V,\\ O_{\mathrm{DeltaNet}}&=(QK^{\mathsf T}\odot M)U, \qquad U=[I+B(KK^{\mathsf T}\odot M^-)]^{-1}BV. \end{aligned}$$

Chunkwise algorithm

The full $n$-token triangular system still requires quadratic work. Divide the sequence into chunks of $C$ tokens instead. Each chunk starts from its incoming boundary state. Local dependencies then fit in a $C\times C$ system, while a $d\times d$ state carries history between chunks. This regrouping is exact.

One chunk as an affine recurrence

Let $r$ index chunks and $\ell\in{1,\ldots,C}$ index tokens in the current chunk. $S_r$ is the state before the chunk, and $S_{r+1}$ is the state after it. Rewrite the DeltaNet update as

$$\begin{aligned} A_{r,\ell} &:=I-\beta_{r,\ell}k_{r,\ell}^{\mathsf T}k_{r,\ell},\\ X_{r,\ell} &:=\beta_{r,\ell}k_{r,\ell}^{\mathsf T}v_{r,\ell},\\ S_r^{(\ell)} &=A_{r,\ell}S_r^{(\ell-1)}+X_{r,\ell}, \qquad S_r^{(0)}=S_r,\quad S_r^{(C)}=S_{r+1}. \end{aligned}$$

For example, a chunk with $C=2$ contains two local transitions:

Use $\Lambda_r[i:j]$ for the ordered transition product over local positions $i$ through $j$ in chunk $r$. Applying the lemma to the incoming state $S_r$ gives

$$\begin{aligned} S_{r+1} &=\Lambda_r[1:C]S_r+H_r,\\ H_r &:=\sum_{i=1}^{C}\Lambda_r[i+1:C]X_{r,i}. \end{aligned}$$

The first term carries the incoming state across the chunk. $H_r$ is the within-chunk convolution sum of local writes. Computing either expression directly would require many dense $d\times d$ transition products.

Compact the transition product with WY

For DeltaNet, the two vectors in each DPLR correction are $k_{r,\ell}$ and $\beta_{r,\ell}k_{r,\ell}$. We absorb $\beta_{r,\ell}$ into the triangular factor below. WY then gives

$$\Lambda_r[1:C]=I-K_r^{\mathsf T}W_r.$$

To construct $W_r$, maintain the partial form

$$\Lambda_r[1:\ell] = I-\sum_{i=1}^{\ell}k_{r,i}^{\mathsf T}w_{r,i}.$$

The first three rows of $W_r$ follow by substituting one transition at a time:

$$\begin{aligned} w_{r,1} &=\beta_{r,1}k_{r,1},\\ w_{r,2} &=\beta_{r,2}\!\left( k_{r,2}-(k_{r,2}k_{r,1}^{\mathsf T})w_{r,1} \right),\\ w_{r,3} &=\beta_{r,3}\!\left( k_{r,3} -(k_{r,3}k_{r,1}^{\mathsf T})w_{r,1} -(k_{r,3}k_{r,2}^{\mathsf T})w_{r,2} \right). \end{aligned}$$

The general row is

$$w_{r,\ell} = \beta_{r,\ell}\!\left( k_{r,\ell} -\sum_{i<\ell} (k_{r,\ell}k_{r,i}^{\mathsf T})w_{r,i} \right).$$

Let $M_c$ be the inclusive $C\times C$ causal mask, $M_c^-:=M_c-I$, and

$$B_r:=\operatorname{diag}(\beta_r), \qquad L_r:=(K_rK_r^{\mathsf T})\odot M_c^-.$$

As in the full-sequence parallel form, stacking these row recurrences produces:

$$(I+B_rL_r)W_r=B_rK_r, \qquad \Lambda_r[1:C]=I-K_r^{\mathsf T}W_r.$$

Compact the local writes

The local convolution sum $H_r$ is the state produced by the tokens in chunk $r$ when the incoming state is zero. The full-sequence parallel form already showed that DeltaNet writes its state as $K^{\mathsf T}U$. Applying the same result within the chunk gives

$$H_r=K_r^{\mathsf T}U_r.$$

The first three rows of $U_r$ are (as you may verify)

$$\begin{aligned} u_{r,1} &=\beta_{r,1}v_{r,1},\\ u_{r,2} &=\beta_{r,2}\!\left( v_{r,2}-(k_{r,2}k_{r,1}^{\mathsf T})u_{r,1} \right),\\ u_{r,3} &=\beta_{r,3}\!\left( v_{r,3} -(k_{r,3}k_{r,1}^{\mathsf T})u_{r,1} -(k_{r,3}k_{r,2}^{\mathsf T})u_{r,2} \right). \end{aligned}$$

The general row is

$$u_{r,\ell} = \beta_{r,\ell}\!\left( v_{r,\ell} -\sum_{i<\ell} (k_{r,\ell}k_{r,i}^{\mathsf T})u_{r,i} \right).$$

This is exactly the full-sequence correction recurrence, applied only within chunk $r$. Stacking the rows gives:

$$(I+B_rL_r)U_r=B_rV_r.$$

Correct the incoming state

Substitute the compact forms of $\Lambda_r[1:C]$ and $H_r$ into the unrolled recurrence:

$$\begin{aligned} S_{r+1} &=(I-K_r^{\mathsf T}W_r)S_r+K_r^{\mathsf T}U_r\\ &=S_r+K_r^{\mathsf T}(U_r-W_rS_r). \end{aligned}$$

Define the actual correction rows as

$$Z_r=U_r-W_rS_r.$$

$U_r$ contains the corrections obtained with an empty incoming state. $W_rS_r$ is the transformed contribution of the predictions from $S_r$. Their difference gives the corrections produced by the actual chunk:

$$S_{r+1}=S_r+K_r^{\mathsf T}Z_r.$$

The update now has the same additive form as Naïve linear attention. We can therefore reuse its chunkwise output formula with $Z_r$ in place of $V_r$:

$$O_r = Q_rS_r + (Q_rK_r^{\mathsf T}\odot M_c)Z_r.$$
Complete per-chunk algorithm

For each chunk $r$,

$$\begin{aligned} L_r &=(K_rK_r^{\mathsf T})\odot M_c^-,\\ (I+B_rL_r)T_r &=B_r,\\ W_r &=T_rK_r,\\ U_r &=T_rV_r,\\ Z_r &=U_r-W_rS_r,\\ O_r &=Q_rS_r+(Q_rK_r^{\mathsf T}\odot M_c)Z_r,\\ S_{r+1} &=S_r+K_r^{\mathsf T}Z_r. \end{aligned}$$ 
<p>The local quantities $L_r,T_r,W_r,$ and $U_r$ can be computed for all chunks in parallel. The boundary pass then threads $S_r$ through the chunks to form $Z_r$, $O_r$, and $S_{r+1}$.</p>

Gated DeltaNet (GDN)

Gated DeltaNet (GDN) extends DeltaNet by adding a token-dependent scalar decay to the state transition:

$$\begin{aligned} \text{DeltaNet:}\qquad S_t &=(I-\beta_tk_t^{\mathsf T}k_t)S_{t-1} +\beta_tk_t^{\mathsf T}v_t,\\ \text{GDN:}\qquad S_t &=a_t(I-\beta_tk_t^{\mathsf T}k_t)S_{t-1} +\beta_tk_t^{\mathsf T}v_t. \end{aligned}$$

GDN adds scalar decay to DeltaNet in the same way that Mamba-2 adds scalar decay to Naïve linear attention. DeltaNet also has the same chunkwise form as Naïve linear attention after replacing $V_r$ with $Z_r$.

Method Chunkwise readout Chunkwise state update
Naïve $O_r=Q_rS_r+(Q_rK_r^{\mathsf T}\odot M_c)V_r$ $S_{r+1}=S_r+K_r^{\mathsf T}V_r$
Mamba-2 $O_r=\overleftarrow Q_rS_r+(Q_rK_r^{\mathsf T}\odot\Gamma_{C,r})V_r$ $S_{r+1}=\lambda_r[1:C]S_r+\overrightarrow K_r^{\mathsf T}V_r$
DeltaNet $O_r=Q_rS_r+(Q_rK_r^{\mathsf T}\odot M_c)Z_r$ $S_{r+1}=S_r+K_r^{\mathsf T}Z_r$
GDN $O_r=\overleftarrow Q_rS_r+(Q_rK_r^{\mathsf T}\odot\Gamma_{C,r})Z_r$ $S_{r+1}=\lambda_r[1:C]S_r+\overrightarrow K_r^{\mathsf T}Z_r$
Complete per-chunk algorithm

For each chunk $r$, define the decay products and the correction gate:

$$\begin{aligned} \lambda_r[i:j] &:=\prod_{s=i}^{j}a_{r,s}, \qquad \lambda_r[i:j]=1\ \text{when }i>j,\\ \Gamma_{C,r}[\ell,i] &:= \begin{cases} \lambda_r[i+1:\ell],&i\le\ell,\\ 0,&i>\ell, \end{cases}\\ B_r &:=\operatorname{diag}(\beta_{r,1},\ldots,\beta_{r,C}). \end{aligned}$$
$$\begin{aligned} \overleftarrow q_{r,\ell} &:=\lambda_r[1:\ell]q_{r,\ell},& \overleftarrow k_{r,\ell} &:=\lambda_r[1:\ell]k_{r,\ell},& \overrightarrow k_{r,\ell} &:=\lambda_r[\ell+1:C]k_{r,\ell}. \end{aligned}$$
$$\begin{aligned} L_r &=(K_rK_r^{\mathsf T}\odot\Gamma_{C,r})\odot M_c^-,\\ (I+B_rL_r)T_r &=B_r,\\ U_r &=T_rV_r,\\ \overleftarrow W_r &=T_r\overleftarrow K_r,\\ Z_r &=U_r-\overleftarrow W_rS_r,\\ O_r &=\overleftarrow Q_rS_r +(Q_rK_r^{\mathsf T}\odot\Gamma_{C,r})Z_r,\\ S_{r+1} &=\lambda_r[1:C]S_r+\overrightarrow K_r^{\mathsf T}Z_r. \end{aligned}$$

Gated Delta Product (GDP)

Gated Delta Product (GDP) replaces GDN's single generalized Householder transition per token with an ordered product of $n_h$ transitions.

Each transition has its own key, value, and correction gate, all projected from the same input token. The transitions act along different learned key directions, increasing computation and expressiveness without changing the shape of $S_t$.

For substep $j$, define

$$A_{t,j}:=I-\beta_{t,j}k_{t,j}^{\mathsf T}k_{t,j}, \qquad X_{t,j}:=\beta_{t,j}k_{t,j}^{\mathsf T}v_{t,j}.$$

For example, with $n_h=3$, the token update is

$$\begin{aligned} S_{t,0} &=a_tS_{t-1},\\ S_{t,1} &=A_{t,1}S_{t,0}+X_{t,1},\\ S_{t,2} &=A_{t,2}S_{t,1}+X_{t,2},\\ S_{t,3} &=A_{t,3}S_{t,2}+X_{t,3},\\ S_t &:=S_{t,3}, \qquad o_t=q_tS_t. \end{aligned}$$

The decay $a_t$ is applied before the first substep, so the previous state decays only once. The query reads the state only after all three updates have finished.

GDP as GDN over a virtual sequence

GDP can be written as an ordinary GDN recurrence over a longer sequence. Treat its $n_h$ substeps as consecutive positions in that sequence. We call the original positions real tokens and the expanded positions virtual steps.

For $n_h=3$, one real token expands into three virtual steps. Two real tokens therefore become six:

real token $t=1$real token $t=2$
virtual step $m$123456
substep $(t,j)$$(1,1)$$(1,2)$$(1,3)$$(2,1)$$(2,2)$$(2,3)$
$\tilde k_m$$k_{1,1}$$k_{1,2}$$k_{1,3}$$k_{2,1}$$k_{2,2}$$k_{2,3}$
$\tilde v_m$$v_{1,1}$$v_{1,2}$$v_{1,3}$$v_{2,1}$$v_{2,2}$$v_{2,3}$
$\tilde\beta_m$$\beta_{1,1}$$\beta_{1,2}$$\beta_{1,3}$$\beta_{2,1}$$\beta_{2,2}$$\beta_{2,3}$
$\tilde a_m$$a_1$11$a_2$11
$\tilde q_m$00$q_1$00$q_2$
retained output——$o_1$——$o_2$

The decay appears at the first virtual step of each real token; the remaining steps use $1$. The query appears at the last virtual step; the earlier steps use zero queries because they do not produce outputs.

To write this construction generally, flatten token $t$ and substep $j$ into

$$m=n_h(t-1)+j.$$

Map the GDP inputs to this virtual sequence by setting

$$\begin{aligned} (\tilde k_m,\tilde v_m,\tilde\beta_m) &=(k_{t,j},v_{t,j},\beta_{t,j}),\\ \tilde a_m &= \begin{cases} a_t,&j=1,\\ 1,&j>1, \end{cases}\\ \tilde q_m &= \begin{cases} q_t,&j=n_h,\\ 0,&j<n_h. \end{cases} \end{aligned}$$

Every virtual step then follows the GDN recurrence:

$$\begin{aligned} \tilde S_m &=\tilde a_m (I-\tilde\beta_m\tilde k_m^{\mathsf T}\tilde k_m)\tilde S_{m-1} +\tilde\beta_m\tilde k_m^{\mathsf T}\tilde v_m,\\ S_t &=\tilde S_{tn_h}, \qquad o_t=\tilde q_{tn_h}\tilde S_{tn_h}. \end{aligned}$$

Flattening preserves the update order. We can therefore run the GDN chunkwise algorithm on a virtual sequence of length $nn_h$ and retain every $n_h$-th output. Only the sequence becomes longer; the recurrent state keeps the same shape.

Gated Linear Attention (GLA)

Mamba-2 lets each token choose how much of the previous state survives, but its gate is scalar:

$$S_t=a_tS_{t-1}+k_t^{\mathsf T}v_t.$$

The same $a_t$ scales every row of $S_{t-1}$. It can shorten or extend the lifetime of the memory as a whole, but it cannot preserve one channel while rapidly forgetting another.

Gated Linear Attention (GLA) replaces the scalar gate with a vector

$$\boldsymbol a_t=(a_{t,1},\ldots,a_{t,d})\in(0,1)^d,$$

and uses the recurrence

$$S_t=\operatorname{Diag}(\boldsymbol a_t)S_{t-1}+k_t^{\mathsf T}v_t, \qquad o_t=q_tS_t.$$

The two gates scale the same state differently:

GLA therefore gives each channel its own forgetting rate.

Chunkwise form

Multiplying diagonal decay matrices is equivalent to multiplying their diagonal vectors elementwise:

$$\operatorname{Diag}(\boldsymbol a_{r,i}) \cdots \operatorname{Diag}(\boldsymbol a_{r,j}) = \operatorname{Diag}\left( \boldsymbol a_{r,i}\odot\cdots\odot\boldsymbol a_{r,j} \right).$$

We therefore replace Mamba-2's scalar prefix product with a vector prefix product. Within chunk $r$, define

$$\boldsymbol\lambda_r[i:j] :=\bigodot_{s=i}^{j}\boldsymbol a_{r,s}, \qquad \boldsymbol\lambda_r[i:j]=\boldsymbol 1\ \text{when }i>j.$$

Here, $\bigodot$ denotes a repeated Hadamard (elementwise) product.

After a write enters the state, later tokens decay each of its channels independently:

Incoming state. The incoming state reaches local position $\ell$ through $\boldsymbol\lambda_r[1:\ell]$. Stack these vectors and absorb them into $Q_r$:

$$\overleftarrow q_{r,\ell} := \boldsymbol\lambda_r[1:\ell]\odot q_{r,\ell}, \qquad \overleftarrow Q_r := \begin{bmatrix} \boldsymbol\lambda_r[1:1]\\ \boldsymbol\lambda_r[1:2]\\ \vdots\\ \boldsymbol\lambda_r[1:C] \end{bmatrix} \odot Q_r.$$

Local writes. A write at position $m$ reaches position $\ell$ through $\boldsymbol\lambda_r[m+1:\ell]$. Its local score is therefore

$$A_r[\ell,m] = \begin{cases} q_{r,\ell} \bigl(\boldsymbol\lambda_r[m+1:\ell]\odot k_{r,m}\bigr)^{\mathsf T}, &m\le\ell,\\ 0,&m>\ell. \end{cases}$$

For the state update, absorb the decay from each write to the right boundary into its key:

$$\overrightarrow k_{r,\ell} :=\boldsymbol\lambda_r[\ell+1:C]\odot k_{r,\ell}, \qquad \overrightarrow K_r := \begin{bmatrix} \boldsymbol\lambda_r[2:C]\\ \boldsymbol\lambda_r[3:C]\\ \vdots\\ \boldsymbol\lambda_r[C+1:C] \end{bmatrix} \odot K_r.$$

The complete chunkwise equations are

$$\begin{aligned} O_r &=\overleftarrow Q_rS_r+A_rV_r,\\ S_{r+1} &=\operatorname{Diag}\bigl(\boldsymbol\lambda_r[1:C]\bigr)S_r +\overrightarrow K_r^{\mathsf T}V_r. \end{aligned}$$

Kimi Delta Attention (KDA)

Kimi Delta Attention (KDA) adds GLA's channelwise decay to GDN's delta rule:

$$\begin{aligned} S_t &=(I-\beta_tk_t^{\mathsf T}k_t) \operatorname{Diag}(\boldsymbol a_t)S_{t-1} +\beta_tk_t^{\mathsf T}v_t\\ &=\operatorname{Diag}(\boldsymbol a_t)S_{t-1} +\beta_tk_t^{\mathsf T} \left( v_t-k_t\operatorname{Diag}(\boldsymbol a_t)S_{t-1} \right),\\ o_t&=q_tS_t. \end{aligned}$$

KDA first decays the state channel by channel, then applies the delta correction. If $\boldsymbol a_t=a_t\boldsymbol 1$, it reduces exactly to GDN.

Reuse GLA and GDN

The GLA and GDN sections give most of the chunkwise algorithm. GLA gives the decayed matrices $\overleftarrow Q_r$ and $\overrightarrow K_r$, together with the local score matrix $A_r$. GDN shows that the delta rule replaces the value rows $V_r$ with correction rows $Z_r$. Therefore,

$$\begin{aligned} O_r &=\overleftarrow Q_rS_r+A_rZ_r,\\ S_{r+1} &=\operatorname{Diag}\bigl(\boldsymbol\lambda_r[1:C]\bigr)S_r +\overrightarrow K_r^{\mathsf T}Z_r. \end{aligned}$$

It remains to compute $Z_r$. KDA reuses GDN's correction solve but replaces scalar decay with channelwise decay. In this solve, decay appears in two quantities: $L_r$ and $\overleftarrow K_r$.

KDA keeps this structure and replaces the scalar decay products with channelwise products:

Quantity GDN: scalar decay KDA: channelwise decay
Decay from $i$ through $j$ $\lambda_r[i:j]=\prod_{s=i}^{j}a_{r,s}$ $\boldsymbol\lambda_r[i:j]=\bigodot_{s=i}^{j}\boldsymbol a_{r,s}$
Incoming prediction key $\overleftarrow k_{r,\ell}=\lambda_r[1:\ell]k_{r,\ell}$ $\overleftarrow k_{r,\ell}=\boldsymbol\lambda_r[1:\ell]\odot k_{r,\ell}$
$L_r[\ell,m]$ for $m<\ell$ $\lambda_r[m+1:\ell](k_{r,\ell}k_{r,m}^{\mathsf T})$ $k_{r,\ell}\bigl(\boldsymbol\lambda_r[m+1:\ell]\odot k_{r,m}\bigr)^{\mathsf T}$

Using the new $L_r$ and $\overleftarrow K_r$, reuse GDN's triangular transform. Let $B_r:=\operatorname{diag}(\beta_{r,1},\ldots,\beta_{r,C})$. Then

$$\begin{aligned} (I+B_rL_r)T_r &=B_r,\\ U_r &=T_rV_r,\\ \overleftarrow W_r &=T_r\overleftarrow K_r,\\ Z_r &=U_r-\overleftarrow W_rS_r. \end{aligned}$$
Complete per-chunk algorithm

For each chunk $r$, define the channelwise decay products and correction gate:

$$\begin{aligned} \boldsymbol\lambda_r[i:j] &:=\bigodot_{s=i}^{j}\boldsymbol a_{r,s}, \qquad \boldsymbol\lambda_r[i:j]=\boldsymbol 1\ \text{when }i>j,\\ B_r &:=\operatorname{diag}(\beta_{r,1},\ldots,\beta_{r,C}). \end{aligned}$$
$$\begin{aligned} \overleftarrow q_{r,\ell} &:=\boldsymbol\lambda_r[1:\ell]\odot q_{r,\ell},& \overleftarrow k_{r,\ell} &:=\boldsymbol\lambda_r[1:\ell]\odot k_{r,\ell},& \overrightarrow k_{r,\ell} &:=\boldsymbol\lambda_r[\ell+1:C]\odot k_{r,\ell}. \end{aligned}$$

Stack the transformed rows into $\overleftarrow Q_r$, $\overleftarrow K_r$, and $\overrightarrow K_r$. Define the two causal within-chunk matrices:

$$\begin{aligned} A_r[\ell,m] &:= \begin{cases} q_{r,\ell} \bigl(\boldsymbol\lambda_r[m+1:\ell]\odot k_{r,m}\bigr)^{\mathsf T}, &m\le\ell,\\ 0,&m>\ell, \end{cases}\\ L_r[\ell,m] &:= \begin{cases} k_{r,\ell} \bigl(\boldsymbol\lambda_r[m+1:\ell]\odot k_{r,m}\bigr)^{\mathsf T}, &m<\ell,\\ 0,&m\ge\ell. \end{cases} \end{aligned}$$
$$\begin{aligned} (I+B_rL_r)T_r &=B_r,\\ U_r &=T_rV_r,\\ \overleftarrow W_r &=T_r\overleftarrow K_r,\\ Z_r &=U_r-\overleftarrow W_rS_r,\\ O_r &=\overleftarrow Q_rS_r+A_rZ_r,\\ S_{r+1} &=\operatorname{Diag}\bigl(\boldsymbol\lambda_r[1:C]\bigr)S_r +\overrightarrow K_r^{\mathsf T}Z_r. \end{aligned}$$

Sources and further reading

Transformers are RNNs, RetNet, Gated Linear Attention (GLA), Mamba-2 / SSD, DeltaNet, Gated DeltaNet (GDN), DeltaProduct (GDP), and Kimi Linear / Kimi Delta Attention (KDA)

Blogs