A walk through of the DeltaNet family of linear attention variants
A walk through of the DeltaNet family of linear attention variants
DeltaNet 线性注意力变体家族详解
Math notation ⟨k|q⟩ kᵀq 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⟩ is a column vector, ⟨k∣ is a row vector, ⟨k∣q⟩ is a number, and ∣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.
关于数学符号的说明:本文默认采用 bra-ket(狄拉克符号)表示法,因为(以我受量子力学启发的观点来看)这种方式能让推导过程中的形状非常清晰。上方的数学符号切换开关会将所有方程重写为传统的粗体向量和显式转置形式。在 bra-ket 模式下,∣q⟩ 是列向量,⟨k∣ 是行向量,⟨k∣q⟩ 是标量,而 ∣v⟩⟨k∣ 是矩阵。向量默认朝右,而键(key)在写入线性注意力状态时朝左。我们处理的是单因果注意力头和实值向量,假设 DeltaNet 的键已归一化,并让状态从键空间映射到值空间。
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): $S_t = S_{t-1}\operatorname{Diag}(\alpha_t)$ $|\widehat v_t\rangle = S_t|k_t\rangle$ $|e_t\rangle = \beta_t (|v_t\rangle - |\widehat v_t\rangle)$ $S_t = S_t + |e_t\rangle\langle k_t|$ $|o_t\rangle = S_t(d_k^{-1/2}|q_t\rangle)$
现代线性注意力变体非常复杂,初看之下很难理解它们的设计目标。作为参考,以下是 Kimi Delta Attention (KDA) 的状态更新方程: $S_t = S_{t-1}\operatorname{Diag}(\alpha_t)$ $|\widehat v_t\rangle = S_t|k_t\rangle$ $|e_t\rangle = \beta_t (|v_t\rangle - |\widehat v_t\rangle)$ $S_t = S_t + |e_t\rangle\langle k_t|$ $|o_t\rangle = S_t(d_k^{-1/2}|q_t\rangle)$
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. That is the route we will take: 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.
它们之所以难以理解,是因为这是过去几年中开发的一系列线性注意力变体中的最新成果,其复杂性不可避免地膨胀,导致从外部看,这些最新的变体显得难以触及。在本文中,我们将深入探讨 DeltaNet 线性注意力变体家族(其中两个已被最新的 Qwen 和 Kimi 模型家族所采用),并展示如何通过对隐藏状态提出简单的假设,从而推导出相同的方程。这就是我们将采取的路径:Softmax 注意力 → 线性注意力 → DeltaNet → Gated DeltaNet → KDA。只有在推导出 KDA 后,我们才会转向执行它的循环和分块 Triton 程序。
1. Begin with quadratic attention
1. 从二次注意力开始
For a query at token $t$, ordinary causal softmax attention is: $a_{ti} = \frac{ \exp(s\langle k_i| q_t\rangle) }{ \sum_{j\leq t} \exp(s\langle k_j| q_t\rangle) }, \quad s=d_k^{-1/2}$ $|o_t\rangle = \sum_{i\leq t}a_{ti}|v_i\rangle$
对于 token $t$ 处的查询(query),普通的因果 Softmax 注意力为: $a_{ti} = \frac{ \exp(s\langle k_i| q_t\rangle) }{ \sum_{j\leq t} \exp(s\langle k_j| q_t\rangle) }, \quad s=d_k^{-1/2}$ $|o_t\rangle = \sum_{i\leq t}a_{ti}|v_i\rangle$
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$, there are $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.
每个注意力权重都是一个标量。它衡量一个键和一个查询之间的相似度,然后 Softmax 将该查询的所有分数转化为一个分布。输出是值向量的加权和。对于长度为 $T$ 的序列,存在 $T^2$ 个键-查询对。在自回归推理过程中,我们可以缓存键和值而不是重新计算它们,但缓存仍会随序列增长,且每个新查询仍必须检查整个历史记录。重排此计算的障碍在于 Softmax。其分母同时依赖于当前查询和所有之前的键。因此,暂时先将其移除。
1.1 Remove the softmax
1.1 移除 Softmax
For clarity, absorb the constant scale $s$ into the query. The deliberately bare version of attention is then: $|o_t\rangle = \sum_{i\leq t} \langle k_i| q_t\rangle |v_i\rangle$
为了清晰起见,将常数缩放因子 $s$ 吸收到查询中。那么,刻意简化的注意力版本为: $|o_t\rangle = \sum_{i\leq t} \langle k_i| q_t\rangle |v_i\rangle$
The scalar inner product can move to the right: $|o_t\rangle = \sum_{i\leq t} |v_i\rangle \langle k_i| q_t\rangle = \left( \sum_{i\leq t} |v_i\rangle\langle k_i| \right) |q_t\rangle$
标量内积可以移到右侧: $|o_t\rangle = \sum_{i\leq t} |v_i\rangle \langle k_i| q_t\rangle = \left( \sum_{i\leq t} |v_i\rangle\langle k_i| \right) |q_t\rangle$
Everything that depends on the past can now be collected into one matrix of a fixed size $V \times K$: $S_t = \sum_{i\leq t} |v_i\rangle\langle k_i|$ and attention becomes a recurrent write followed by a read: $S_t = S_{t-1} + |v_t\rangle\langle k_t|$ $|o_t\rangle = S_t|q_t\rangle$
所有依赖于过去的信息现在都可以收集到一个固定大小为 $V \times K$ 的矩阵中: $S_t = \sum_{i\leq t} |v_i\rangle\langle k_i|$ 注意力机制变成了循环写入,随后进行读取: $S_t = S_{t-1} + |v_t\rangle\langle k_t|$ $|o_t\rangle = S_t|q_t\rangle$
The identity $(|v\rangle\langle k|)|q\rangle = \langle k|q\rangle|v\rangle$ 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$. This is linear in sequence length rather than quadratic: scan the tokens once, updating the same $d_v \times 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.
恒等式 $(|v\rangle\langle k|)|q\rangle = \langle k|q\rangle|v\rangle$ 就是全部诀窍。外积是一个矩阵;内积是一个数字。我们不再存储每一个过去的键和值。我们将它们的外积之和存储在固定大小的状态 $S_t$ 中。这在序列长度上是线性的,而不是二次的:扫描一次 token,在每一步更新相同的 $d_v \times d_k$ 状态。我们通过放弃 Softmax 的归一化和选择性换取了这种效率。更复杂的线性注意力方法使用特征映射和归一化器,但这种原始形式揭示了促使 DeltaNet 产生的内存问题。
1.2 Addition is not assignment
1.2 加法不是赋值
Suppose we write a pair $|v_t\rangle\langle k_t|$ and immediately query the new state with that same key: $S_t|k_t\rangle = (S_{t-1} + |v_t\rangle\langle k_t|) |k_t\rangle = S_{t-1}|k_t\rangle + |v_t\rangle \underbrace{\langle k_t|k_t\rangle}{1} = S{t-1}|k_t\rangle+|v_t\rangle$
假设我们写入一个对 $|v_t\rangle\langle k_t|$,并立即用同一个键查询新状态: $S_t|k_t\rangle = (S_{t-1} + |v_t\rangle\langle k_t|) |k_t\rangle = S_{t-1}|k_t\rangle + |v_t\rangle \underbrace{\langle k_t|k_t\rangle}{1} = S{t-1}|k_t\rangle+|v_t\rangle$
The write does not make the memory return $|v_t\rangle$. It adds $|v_t\rangle$ to whatever the memory already returned. If the old state already produced the correct value, the additive write makes the new state produce the wrong value.
这次写入并没有让内存返回 $|v_t\rangle$。它将 $|v_t\rangle$ 加到了内存原本返回的内容上。如果旧状态已经产生了正确的值,那么这种加法写入会导致新状态产生错误的值。