Kimi Delta Attention: From Linear Attention to DPLR Transitions

Kimi Delta Attention (KDA) is a linear attention mechanism that optimizes memory updates and state retention by treating the hidden state as an associative memory. By replacing unconditional additive writes with a delta-rule correction and promoting scalar forgetting to per-channel diagonal decay, KDA achieves a flexible, fixed-size state transition known as a diagonal-plus-low-rank (DPLR) operation.

From Quadratic to Linear Attention

Standard causal softmax attention is computationally expensive because it scales quadratically ($T^2$) with sequence length, as every query must inspect every previous key. The core bottleneck is the softmax denominator, which couples the current query with all previous keys.

By removing the softmax normalization, attention can be rewritten as a recurrent write and read operation. Instead of storing all past keys and values, the model maintains a fixed-size state $S_t$ (a matrix of size $d_v \times d_k$) that stores the summed outer products of keys and values:

State Update: $S_t = S_{t-1} + |v_t\rangle \langle k_t|$
Read Operation: $|o_t\rangle = S_t |q_t\rangle$

While this reduces complexity to linear time, it introduces a memory interference problem: additive writes behave like += rather than =, meaning new information is added to existing associations rather than replacing them.

DeltaNet: Targeted Memory Replacement

DeltaNet solves the interference problem by implementing a delta-rule correction. Instead of adding the full value vector $|v_t\rangle$, the model calculates the prediction error—the difference between the target value and what the current state associates with the new key.

The Delta Rule Derivation

Predict: The model queries the current state: $|\hat{v}t\rangle = S{t-1} |k_t\rangle$.
Correct: It calculates the error: $|e_t\rangle = \beta_t (|v_t\rangle - |\hat{v}_t\rangle)$, where $\beta_t$ is a learned write strength.
Write: The error is written to the state: $S_t = S_{t-1} + |e_t\rangle \langle k_t|$.

This mechanism ensures that if $\beta_t = 1$, the memory is updated such that a subsequent read of $|k_t\rangle$ returns exactly $|v_t\rangle$. Mathematically, this is equivalent to taking one gradient descent step on a reconstruction loss $\mathcal{L}_t = \frac{1}{2} (\parallel S |k_t\rangle - |v_t\rangle\parallel)_2^2$.

Gated DeltaNet and Kimi Delta Attention (KDA)

While DeltaNet improves how information is written, it does not manage how information is forgotten. Stale information in the state can distort future reads.

Gated DeltaNet

Gated DeltaNet introduces a scalar retention gate $\alpha_t \in [0, 1]$ that globally decays the state before the delta update: $$\tilde{S}t = \alpha_t S{t-1}$$ This allows the model to erase the entire state globally, but it forces all key channels to be forgotten at the same rate.

Kimi Delta Attention (KDA)

KDA evolves this by replacing the scalar $\alpha_t$ with a vector $\alpha_t \in [0, 1]^{d_k}$, applied as a diagonal matrix $D_t = \text{Diag}(\alpha_t)$. This allows the model to forget each key channel independently: $$\tilde{S}t = S{t-1} D_t$$

Combining this with the delta rule, the full KDA recurrence is:

Forget: $\tilde{S}t = S{t-1} D_t$
Predict: $|\hat{v}_t\rangle = \tilde{S}_t |k_t\rangle$
Correct: $|e_t\rangle = \beta_t (|v_t\rangle - |\hat{v}_t\rangle)$
Write: $S_t = \tilde{S}_t + |e_t\rangle \langle k_t|$
Read: $|o_t\rangle = S_t (d_k^{-1/2} |q_t\rangle)$

The DPLR Transition

Expanding the KDA update reveals that the state transition is a diagonal-plus-low-rank (DPLR) operation. The transition acting on the key space is: $$A_t = D_t - \beta_t D_t |k_t\rangle \langle k_t| = D_t - |b_t\rangle \langle a_t|$$ where $|b_t\rangle = D_t |k_t\rangle$ and $\langle a_t| = \beta_t \langle k_t|$. This structure is critical for efficient implementation, as it combines a diagonal decay with a rank-one update.

Execution Regimes: Recurrent vs. Chunkwise

KDA is implemented using two different schedules depending on the hardware utilization requirements:

Fused Recurrent Kernel

Used primarily for autoregressive decoding, this kernel processes tokens one by one. It is a literal transcription of the recurrence, performing reductions for predictions/reads and outer products for writes. While low-latency, it is less efficient for training because it does not leverage large matrix multiplications on tensor cores.

Chunkwise KDA

For training and long prefill, KDA is reorganized into chunkwise operations. This involves:

Provisional Errors: Computing initial errors as if tokens only saw the incoming state.
Causal Resolution: Using a triangular solve to resolve dependencies between delta errors within the chunk (since token $i$ depends on writes from tokens $0 – i-1$).
State Fast-Forwarding: Advancing the state across the chunk boundary using a single matrix multiplication: $S_{c+1} = S_c D_{0:C} + E_c K_c^{end}$.
Causal Output Calculation: Computing all outputs in a chunk by combining the decayed incoming state read with the causal contributions of in-chunk writes.

Community Insights

Discussion surrounding KDA highlights the tension between theoretical simplicity and practical implementation. Some contributors noted that while the derivation seems intuitive in retrospect, the jump to these architectures often depends on having the compute resources to experiment with these alternatives to standard attention.

Regarding notation, the use of bra-ket notation was praised by some for clarifying the shapes of the tensors, though others found it unfamiliar. As one user noted:

"The bra-ket notation makes this all very simple/intuitive for me. With 'vectors' I always get confused which is horizontal/vertical... With bra-kets the whole thing was very intuitive!"

Others pointed out that linear attention variants are essentially approximations of quadratic attention, and the primary driver for these innovations is the necessity to bypass the $O(n^2)$ scaling limit to handle larger contexts.