A walk through of the DeltaNet family of linear attention variants

A note on notation: this article defaults to bra-ket notation because (in my quantum-inspired opinion) it makes the shapes in this derivation very clear. The Math notation switch above rewrites every equation using conventional bold vectors and explicit transposes instead. In bra-ket mode, ∣ q ⟩ \lvert q\rangle ∣ q ⟩ is a column vector, ⟨ k ∣ \langle k\rvert ⟨ k ∣ is a row vector, ⟨ k ∣ q ⟩ \langle k\rvert q\rangle ⟨ k ∣ q ⟩ is a number, and ∣ v ⟩ ⟨ k ∣ \lvert v\rangle\langle k\rvert ∣ v ⟩ ⟨ k ∣ is a matrix. Vectors face right by default, while keys face left when written into the linear-attention state. We work with one causal attention head and real-valued vectors, assume DeltaNet’s keys are normalized, and let the state map from key space to value space.

Modern linear attention variants are complex, and a upon first glance it is not so easy to see what they are designed to achieve. For reference here is the state update equation for Kimi Delta Attention (KDA):

t ​ = S t − 1 ​ Diag ( α t ​ ) ∣ v ^ t ⟩ = S ~ t ∣ k t ⟩ \lvert\widehat v_t\rangle = \widetilde S_t\lvert k_t\rangle ∣ v t ​ ⟩ = S t ​ ∣ k t ​ ⟩ ∣ e t ⟩ = β t ( ∣ v t ⟩ − ∣ v ^ t ⟩ ) \lvert e_t\rangle = \beta_t \left( \lvert v_t\rangle-\lvert\widehat v_t\rangle \right) ∣ e t ​ ⟩ = β t ​ ( ∣ v t ​ ⟩ − ∣ v t ​ ⟩ ) S t = S ~ t + ∣ e t ⟩ ⟨ k t ∣ S_t = \widetilde S_t+\lvert e_t\rangle\langle k_t\rvert S t ​ = S t ​ + ∣ e t ​ ⟩ ⟨ k t ​ ∣ ∣ o t ⟩ = S t ( d k − 1 / 2 ∣ q t ⟩ ) \lvert o_t\rangle = S_t\left(d_k^{-1/2}\lvert q_t\rangle\right) ∣ o t ​ ⟩ = S t ​ ( d k − 1/2 ​ ∣ q t ​ ⟩ ) The reason they are so difficult to understand is that this is the latest in a family of linear attention variants that have been developed over the last few years and the complexity of them has inevitably ballooned such that from the outside the latest variants appear inaccessible.

In this post we are going to walk through the DeltaNet family of linear attention variants, two of which are used by the latest Qwen and Kimi model families, and show how you might have arrived at the same equations by asserting simple things about your hidden state.

softmax attention → linear attention → DeltaNet → Gated DeltaNet → KDA

Only after deriving KDA will we turn to the recurrent and chunkwise Triton programs that execute it.

For a query at token t t t , ordinary causal softmax attention is

Every attention weight is a scalar. It measures the similarity between one key and one query, then softmax turns all of the scores for that query into a distribution. The output is a weighted sum of value vectors.

Over a sequence of length T T T , there are T 2 T^2 T 2 key-query pairs. During autoregressive inference we can cache the keys and values instead of recomputing them, but the cache still grows with the sequence and every new query still has to inspect the entire history.

The obstacle to rearranging this computation is the softmax. Its denominator depends jointly on the current query and every earlier key. So, for the moment, remove it.

For clarity, absorb the constant scale s s s into the query. The deliberately bare version of attention is then

The scalar inner product can move to the right:

Everything that depends on the past can now be collected into one matrix of a fixed size V × K V \times K V × K :

and attention becomes a recurrent write followed by a read:

is the whole trick. The outer product is a matrix; the inner product is a number. We no longer store every past key and value. We store their summed outer products in the fixed-size state S t S_t S t ​ .

This is linear in sequence length rather than quadratic: scan the tokens once, updating the same d v × d k d_v\times d_k d v ​ × d k ​ state at every step. We have paid for that efficiency by discarding softmax’s normalization and selectivity. More sophisticated linear-attention methods use feature maps and normalizers, but this unadorned form exposes the memory problem that motivates DeltaNet.

Suppose we write a pair ∣ v t ⟩ ⟨ k t ∣ \lvert v_t\rangle\langle k_t\rvert ∣ v t ​ ⟩ ⟨ k t ​ ∣ and immediately query the new state with that same key:

⟨ k t ​ ∣ k t ​ ⟩ ​ ​ = S t − 1 ​ ∣ k t ​ ⟩ + ∣ v t ​ ⟩ . ​ The write does not make the memory return ∣ v t ⟩ \lvert v_t\rangle ∣ v t ​ ⟩ . It adds ∣ v t ⟩ \lvert v_t\rangle ∣ v t ​ ⟩ to whatever the memory already returned.

If the old state already produced the correct value, the additive write makes the new state produce twice that value. More generally, keys are not mutually orthogonal, so every write can interfere with previous writes. Linear attention has given us a compact associative memory, but its update behaves like += when what we want is closer to = .

DeltaNet replaces the unconditional linear-attention write with a delta-rule correction. There are two useful ways to derive it.

Before writing token t t t , ask the memory what it currently associates with the new key:

t ​ ⟩ = S t − 1 ​ ∣ k t ​ ⟩ . If we want the memory to return ∣ v t ⟩ \lvert v_t\rangle ∣ v t ​ ⟩ , we should not add the whole value. We should add only the difference:

t ​ ⟩ . Introduce a learned write strength β t ∈ [ 0 , 1 ] \beta_t\in[0,1] β t ​ ∈ [ 0 , 1 ] and define

Then write this error at the current key:

When β t = 1 \beta_t=1 β t ​ = 1 , the result is exactly ∣ v t ⟩ \lvert v_t\rangle ∣ v t ​ ⟩ . Smaller β t \beta_t β t ​ moves the old prediction partway towards the target.

The correction is also local in key space. For any query ∣ x ⟩ \lvert x\rangle ∣ x ⟩ orthogonal to the current key,

⟨ k t ​ ∣ x ⟩ ​ ​ = 0. So the rank-one write changes the response in the selected key direction while leaving every orthogonal direction alone.

The same update falls out of an online learning objective. Treat the current key-value pair as one training example for the linear map S S S :

Its gradient with respect to the state is

This is visibly an outer product: a value-space prediction error times the key bra at which that error was observed. Take one gradient-descent step of size β t \beta_t β t ​ from S t − 1 S_{t-1} S t − 1 ​ :

This is exactly the update we got by requiring immediate reconstruction. The two interpretations are the same:

Expanding the error exposes DeltaNet as a structured state transition plus a new input:

For a unit key, I − β t ∣ k t ⟩ ⟨ k t ∣ I-\beta_t\lvert k_t\rangle\langle k_t\rvert I − β t ​ ∣ k t ​ ⟩ ⟨ k t ​ ∣ has eigenvalue 1 − β t 1-\beta_t 1 − β t ​ in the current key direction and eigenvalue 1 1 1 in every orthogonal direction. It removes the old association along the current key before adding the new one.

DeltaNet fixes the write. It does not yet fix the lifetime of the state.

The linear state compresses the whole history into one matrix. A read

cannot choose to skip an individual old token after that token has been folded into S t S_t S t ​ . Every stored direction that overlaps the query contributes. The delta rule can correct the state around the current key, but stale information in other directions remains available and can distort future reads.

We therefore need a way to forget the old state before using it. Let α t ∈ [ 0 , 1 ] \alpha_t\in[0,1] α t ​ ∈ [ 0 , 1 ] be a learned scalar retention gate:

t ​ = α t ​ S t − 1 ​ . Run the same delta rule against this gated state:

t ​ ∣ v t ​ ⟩ ∣ e t ​ ⟩ S t ​ ​ = α t ​ S t − 1 ​ , = S t ​ ∣ k t ​ ⟩ , = β t ​ ( ∣ v t ​ ⟩ − ∣ v t ​ ⟩ ) , = S t ​ + ∣ e t ​ ⟩ ⟨ k t ​ ∣ , ​ ​ forget , predict , correct , write . ​ ​ This is Gated DeltaNet . The order matters: forget first, predict from the retained state, then correct that prediction. If we predicted before forgetting, the error would describe a different memory from the one we update.

The delta rule gives targeted replacement; the scalar gate gives global erasure. They solve different problems and are complementary.

But α t \alpha_t α t ​ still makes one decision for the entire matrix. The model must retain or forget every key channel at the same rate.

Kimi Delta Attention replaces Gated DeltaNet’s scalar retention with a vector α t ∈ [ 0 , 1 ] d k \alpha_t\in[0,1]^{d_k} α t ​ ∈ [ 0 , 1 ] d k ​ . Put the vector on the diagonal:

Our state maps keys to values, so the key channels are the columns of S S S . Right-multiplication applies a different retention factor to every one:

t ​ = S t − 1 ​ D t ​ . Everything else is the delta rule we have already derived:

t ​ ∣ v t ​ ⟩ ∣ e t ​ ⟩ S t ​ ∣ o t ​ ⟩ ​ = S t − 1 ​ D t ​ , = S t ​ ∣ k t ​ ⟩ , = β t ​ ( ∣ v t ​ ⟩ − ∣ v t ​ ⟩ ) , = S t ​ + ∣ e t ​ ⟩ ⟨ k t ​ ∣ , = S t ​ ( s ∣ q t ​ ⟩) , s = d k − 1/2 ​ , ​ ​ forget each key channel , predict , correct , write , read . ​ ​ That is KDA. Compared with Gated DeltaNet, the conceptual change is only the promotion

The effect is substantial: one channel can be cleared while another is retained.

D t ​ ( I − β t ​ ∣ k t ​ ⟩ ⟨ k t ​ ∣ ) ​ ​ + β t ​ ∣ v t ​ ⟩ ⟨ k t ​ ∣ . ​ The key-space transition is

So A t A_t A t ​ is a diagonal matrix minus a rank-one matrix: a diagonal-plus-low-rank , or DPLR, transition. “DPLR” describes the d k × d k d_k\times d_k d k ​ × d k ​ transition acting on key space. The memory state itself is still the d v × d k d_v\times d_k d v ​ × d k ​ matrix S t S_t S t ​ .

The full journey can now be summarized compactly:

The implementation usually stores g t = log ⁡ α t g_t=\log\alpha_t g t ​ = lo g α t ​ with g t ≤ 0 g_t\leq0 g t ​ ≤ 0 , then obtains the retention factors as exp ⁡ ( g t ) \exp(g_t) exp ( g t ​ ) . In the transposed d k × d v d_k\times d_v d k ​ × d v ​ layout used by the reference code, the recurrence is only five lines:

state = state * g_t.exp().unsqueeze( - 1 ) prediction = einsum( "bhkv,bhk->bhv" , state, k_t) residual = beta_t.unsqueeze( - 1 ) * (v_t - prediction) state = state + einsum( "bhk,bhv->bhkv" , k_t, residual) output = einsum( "bhk,bhkv->bhv" , q_t * scale, state) See the official naive_recurrent_kda reference.

The recurrence above is the natural implementation for autoregressive decode. KDA has two principal execution regimes:

The recurrent Triton launch uses one program per sequence, value head, and 32-wide value tile:

BK = triton.next_power_of_2(K) BV = 32 grid = (triton.cdiv(V, BV ) * N * HV ,) See the fused_recurrent_kda_fwd launch code.

BK covers the key dimension in the normal supported configuration. Each program owns a [BK, BV] tile of the implementation’s transposed state and loops over tokens in order. Different value tiles, heads, and sequences run independently.

The kernel is almost a literal transcription of the recurrence:

state *= tl.exp(g_t[:, None ]) prediction = tl.sum(state * k_t[:, None ], axis = 0 ) residual = beta_t * (v_t - prediction) state += k_t[:, None ] * residual[ None , :] out_t = tl.sum(state * (q_t * SCALE )[:, None ], axis = 0 ) The prediction and read are reductions; the write is an outer product. This is excellent for decode, where only one new token is available at a time. It is less attractive for training and long prefill because these vector operations do not become the large matrix multiplications on which tensor cores are most efficient.

That motivates a second view of exactly the same recurrence.

Chunkwise KDA processes C C C tokens together. It must produce exactly the same states and outputs as the token-by-token recurrence, but it reorganizes the work into matrix products.

For each chunk c c c , we need two results: