Kimi Delta Attention: Derivation and Implementation
Kimi Delta Attention: Derivation and Implementation
Kimi Delta Attention (KDA) provides a mechanism for efficient, fixed-size recurrent memory by combining a delta-rule update for targeted state replacement with per-channel forgetting.
Unlike standard softmax attention, which scales quadratically with sequence length, KDA allows for linear scaling by maintaining a state matrix that summarizes past keys and values. It improves upon previous linear attention variants by ensuring that writing new information does not simply add to existing memory but replaces it based on the prediction error, while allowing the model to forget information on a per-dimension basis.
From Quadratic to Linear Attention
Standard causal softmax attention computes a weighted sum of value vectors based on the similarity between a query and all previous keys. This requires $T^2$ operations for a sequence of length $T$. By removing the softmax normalization, attention can be rewritten as a recurrent process:
$$S_{t} = S_{t - 1} + \mid v_{t} \rangle \langle k_{t} \mid$$ $$\mid o_{t} \rangle = S_{t} \mid q_{t} \rangle$$
In this formulation, the state $S_t$ is a fixed-size matrix ($d_v \times d_k$) that stores the summed outer products of keys and values. While this achieves linear complexity, it introduces a "summation problem": adding a new key-value pair $\mid v_t \rangle \langle k_t \mid$ does not replace the old association for that key but adds to it, potentially leading to interference and value inflation.
DeltaNet: Targeted State Replacement
DeltaNet solves the additive write problem by implementing a delta-rule correction. Instead of writing the full value vector, the system writes only the error between the target value and the current memory's prediction for that key.
The Delta-Rule Derivation
- Predict: The memory is queried with the current key to find the existing association: $\mid \hat{v}t \rangle = S{t-1} \mid k_t \rangle$.
- Correct: The error is calculated as the difference between the actual value and the prediction, scaled by a learned write strength $\beta_t$: $\mid e_t \rangle = \beta_t (\mid v_t \rangle - \mid \hat{v}_t \rangle)$.
- Write: This error is written into the state: $S_t = S_{t-1} + \mid e_t \rangle \langle k_t \mid$.
This approach ensures that if $\beta_t = 1$, the memory will return exactly $\mid v_t \rangle$ when queried with $\mid k_t \rangle$ immediately after the write. Mathematically, this is equivalent to taking one step of gradient descent on a reconstruction loss $\mathcal{L}_t = \frac{1}{2} \parallel S \mid k_t \rangle - \mid v_t \rangle \parallel_2^2$.
Gated DeltaNet and Kimi Delta Attention
While DeltaNet fixes how information is written, it does not address how information is forgotten. Gated DeltaNet introduces a scalar retention gate $\alpha_t \in [0, 1]$ that decays the entire state $S_{t-1}$ before the delta update. However, a single scalar forces all key channels to be forgotten at the same rate.
KDA's Per-Channel Forgetting
Kimi Delta Attention (KDA) promotes the scalar gate $\alpha_t$ to a vector $\alpha_t \in [0, 1]^{d_k}$. This vector is applied as a diagonal matrix $D_t = \text{Diag}(\alpha_t)$ to the state:
$$\left(\overset{\sim}{S}\right){t} = S{t - 1} D_{t}$$
The full KDA recurrence is then:
- Forget: $\left(\overset{\sim}{S}\right){t} = S{t - 1} D_{t}$
- Predict: $\mid \hat{v}t \rangle = \left(\overset{\sim}{S}\right){t} \mid k_t \rangle$
- Correct: $\mid e_t \rangle = \beta_t (\mid v_t \rangle - \mid \hat{v}_t \rangle)$
- Write: $S_t = \left(\overset{\sim}{S}\right)_{t} + \mid e_t \rangle \langle k_t \mid$
- Read: $\mid o_t \rangle = S_t (d_k^{-1/2} \mid q_t \rangle)$
This allows the model to independently retain or clear specific key channels, providing significantly more granular control over memory lifetime.
Implementation: Recurrent vs. Chunkwise Execution
KDA is implemented using two different schedules depending on the hardware utilization needs:
Fused Recurrent Kernel
Used primarily for autoregressive decoding, this Triton kernel processes tokens one by one. It is optimized for low latency, performing reductions for predictions and reads, and outer products for writes. Each program handles a tile of the state matrix, looping over tokens in order.
Chunkwise KDA
For training and long prefill, KDA is reorganized into matrix multiplications to leverage tensor cores. This involves:
- Provisional Errors: Computing initial errors by assuming tokens only see the decayed incoming state.
- Causal Dependencies: Using a triangular solve to resolve the dependencies where token $i$'s error depends on writes from tokens $j < i$ within the same chunk.
- State Fast-Forwarding: Advancing the recurrent state across the chunk boundary using a single matrix product: $S_{c+1} = S_c D_{0:C} + E_c K_c^{end}$.
- Causal Outputs: Computing all token outputs within the chunk by combining the decayed incoming state read with the causal contributions of in-chunk writes.
Summary of Linear Attention Evolution
| Mechanism | State Update | Key Innovation |
|---|---|---|
| Linear Attention | $S + \mid v \rangle \langle k \mid$ | Fixed-size recurrent memory |
| DeltaNet | $S + \beta (\mid v \rangle - S \mid k \rangle) \langle k \mid$ | Targeted replacement via delta rule |
| Gated DeltaNet | $\alpha S$, then delta update | Global state forgetting |
| KDA | $S D$, then delta update | Per-key-channel forgetting |
"The big innovation is not so much the mechanism itself but realizing the parallelize-ability of it... Attention is one of those innovations that came mostly from realizing you had better hardware than everybody else and asking yourself how to use it."