Token Mixer I

GATED DELTA NETWORKS

线性 Transformer 使用基于核函数的点积线性注意力取代传统的 Softmax 注意力,并可将其重新表述为一种具有矩阵值隐状态的线性循环神经网络,从而显著降低推理阶段的内存需求。尽管早期线性 Transformer 在语言建模任务上的表现通常逊于标准 Transformer,但近期研究通过引入类似 LSTM 的数据依赖门控机制,取得显著进展,代表性模型包括 GLA 和 Mamba2

如何在长序列中有效管理信息仍然是线性 Transformer 面临的重要挑战。尤其是在上下文内检索任务中,传统 Transformer 依然具有明显优势

  • 线性 Transformer 可以被解释为实现一种基于外积的键值关联记忆,其形式与张量积表征相似。然而,此类模型所能存储的相互正交键值对数量受到模型维度的限制。当序列长度超过该维度时,不同键值关联之间不可避免地会发生记忆碰撞,从而阻碍模型对目标信息的精确检索
  • Mamba2 通过引入一种简单的门控更新规则来缓解这一局限,其状态更新形式为S_t=\alpha_tS_{t-1}+v_tk_t^\top。在每个时间步,该规则利用动态衰减系数\alpha_t\in(0,1)对所有键值关联进行统一衰减。这种方法并未考虑不同键值关联之间的重要性差异,因此可能导致记忆资源利用效率不足。当模型需要遗忘某个特定的键值关联时,其他所有键值关联也会以相同程度被衰减,使得记忆清除过程缺乏针对性
  • DeltaNet 通过以序列化方式利用新到达的键值对对旧键值对进行软替换,从而实现对记忆内容的选择性更新。该方法在上下文内检索的合成基准测试中表现出色。然而,由于其每次只能修改一个键值关联,模型缺乏快速清除过时信息或无关信息的能力。尤其是在上下文发生切换、需要擦除先前信息时,这一局限更为突出。DeltaNet 在真实世界任务中的表现相对一般,这可能是由于其缺少一种有效且稳健的记忆清除机制

鉴于门控更新规则和 Delta 规则在记忆管理方面具有互补优势,提出门控 Delta 规则,这是一种将二者统一起来的简单而直观的机制。该规则能够实现灵活的记忆控制:当令\alpha_t\rightarrow0时,模型可以迅速清除已有记忆;当令\alpha_t\rightarrow1时,该机制实际上退化为纯 Delta 规则,使模型能够在不影响其他信息的情况下,对特定记忆内容进行选择性更新

剩余关键挑战在于如何以硬件高效的方式实现门控 Delta 规则。已有研究基于 WY 表示提出一种能够并行计算 Delta 规则的高效算法。在此基础上,对该方法进行扩展,将门控项纳入其计算过程。该扩展保留分块并行计算的优势,从而使门控 Delta 规则能够在现代硬件上实现高效训练

🐨 MAMBA2: LINEAR ATTENTION WITH DECAY

在不考虑归一化操作以及查询和键的激活函数时,线性 Transformer 可以表述为如下线性递推形式

S_t=S_{t-1}+v_tk_t^\top\in\mathbb{R}^{d_v\times d_k}, \quad o_t=S_tq_t\in\mathbb{R}^{d_v}

  • 其中d_k和 d_v分别表示查询/键和值的注意力头维度。通过展开上述递推关系,分别得到向量形式和矩阵形式

o_t=\sum_{i=1}^{t}(v_ik_i^\top)q_t =\sum_{i=1}^{t}v_i(k_i^\top q_t)\in\mathbb{R}^{d_v}, \quad O=(QK^\top\odot M)V\in\mathbb{R}^{L\times d_v}

  • 其中L表示序列长度,M\in\mathbb{R}^{L\times L}为因果掩码:当i<j时M_{ij}=0;否则M_{ij}=1

基础的线性注意力在语言建模任务上的性能远逊于标准 Transformer。为解决这一问题,一种常见做法是引入衰减项,以遗忘历史信息。以 Mamba2 为例,其可以表示为如下线性递推形式

S_t=\alpha_tS_{t-1}+v_tk_t^\top, \quad o_t=S_tq_t

  • 其中\alpha_t\in(0,1)是一个随时间步 t变化、依赖于输入数据的标量衰减项。定义累积衰减乘积\gamma_j=\prod_{i=1}^{j}\alpha_i。通过展开上述递推关系,可以分别得到如下向量形式和矩阵并行形式

o_t=\sum_{i=1}^{t}\frac{\gamma_t}{\gamma_i}v_ik_i^\top q_t =\sum_{i=1}^{t}v_i\left(\frac{\gamma_t}{\gamma_i}k_i^\top q_t\right), \quad O=\bigl((QK^\top)\odot\Gamma\bigr)V.

\Gamma_{ij}= \begin{cases} \frac{\gamma_i}{\gamma_j}, & i\ge j\\ 0, & i<j \end{cases}

  • \Gamma\in\mathbb{R}^{L\times L}是一个考虑衰减效应的因果掩码。上述并行形式与递推形式之间的等价关系也被称为状态空间对偶性
  • 类似的递推结构还出现在 Gated RFA、xLSTM 和 Gated RetNet 等多种架构中。当 \gamma_t与输入数据无关时,该形式可退化为 RetNet 和 Lightning Attention

为实现硬件友好的线性时间训练,通常采用分块并行形式。分块并行方法将输入和输出划分为若干大小为C的分块,并依据前一个分块的最终状态以及当前分块的查询、键和值矩阵计算当前分块的输出

  • 定义Q_{[t]}:=q_{tC+1:(t+1)C+1}为第 t个分块的查询矩阵,并定义q_{[t]}^r:=q_{tC+r}为第t个分块内的第r个查询向量。第 t个分块的初始状态定义为S_{[t]}:=S_{[t]}^0=S_{[t-1]}^C。通过对递推关系进行部分展开,可以得到

S_{[t]}^r =S_{[t]}+\sum_{i=1}^{r}v_{[t]}^i{k_{[t]}^i}^{\top} \in\mathbb{R}^{d_v\times d_k}, \quad o_{[t]}^r =S_{[t]}^rq_{[t]}^r =S_{[t]}q_{[t]}^r +\sum_{i=1}^{r}v_{[t]}^i \left({k_{[t]}^i}^{\top}q_{[t]}^r\right) \in\mathbb{R}^{d_v}

S_{[t+1]} =S_{[t]}+V_{[t]}^\top K_{[t]} \in\mathbb{R}^{d_v\times d_k}, \quad O_{[t]} =Q_{[t]}S_{[t]}^\top +\bigl(Q_{[t]}K_{[t]}^\top\odot M\bigr)V_{[t]} \in\mathbb{R}^{C\times d_v}

  • 其中M\in\mathbb{R}^{C\times C}是分块内部的因果掩码。上述方程主要由矩阵乘法运算构成,因此可以充分利用张量核心进行硬件层面的优化。这种分块算法也可以方便地扩展至包含衰减项的线性注意力

S_{[t+1]} =\overrightarrow{S}_{[t]} +V_{[t]}^\top\overrightarrow{K}_{[t]} \in\mathbb{R}^{d_v\times d_k}, \quad O_{[t]} =\overleftarrow{Q}_{[t]}S_{[t]}^\top +\bigl(Q_{[t]}K_{[t]}^\top\odot\Gamma_{[t]}\bigr)V_{[t]} \in\mathbb{R}^{C\times d_v}.

(\Gamma_{[t]})_{ij} = \begin{cases}\gamma_{[t]}^i/\gamma_{[t]}^j,&i\ge j,\\0,&i<j,\end{cases}\qquad \gamma_{[t]}^j = \prod_{\ell=tC+1}^{tC+j}\alpha_\ell.

  • 使用左箭头 \overleftarrow{(\,\cdot\,)}和右箭头\overrightarrow{(\,\cdot\,)},分别表示将变量衰减至每个分块的首位置和末位置

\begin{aligned}
{\color{blue}\overleftarrow{\boldsymbol{q}_{[t]}^r}} &= {\color{blue}\gamma_{[t]}^r} \boldsymbol{q}_{[t]}^r && \text{decaying each vector to the first position of chunk } t \\
{\color{blue}\overrightarrow{\boldsymbol{k}_{[t]}^r}} &= \frac{{\color{blue}\gamma_{[t]}^C}}{{\color{blue}\gamma_{[t]}^r}} \boldsymbol{k}_{[t]}^r && \text{decaying each vector to the last position of chunk } t \\
{\color{blue}\overrightarrow{\mathbf{S}_{[t]}}} &= {\color{blue}\gamma_{[t]}^C} \mathbf{S}_{[t]} && \text{decaying the state matrix over the entire chunk } t
\end{aligned}

首先将线性注意力及其带衰减形式表示为等价的递推与并行矩阵计算,并进一步通过 chunkwise 将逐 token 串行更新改写为块内少量大规模 GEMM、块间仅传递紧凑状态矩阵S_{[t]},从而在保持线性复杂度和计算等价性的同时充分利用 GPU Tensor Core

🦌 DELTA NETWORKS: LINEAR ATTENTION WITH DELTA RULE

Delta 更新规则会动态擦除与当前输入键 k_t相关联的旧值v_t^{\mathrm{old}},写入新值v_t^{\mathrm{new}}。新值由当前输入值与旧值根据写入强\beta_t\in(0,1)进行线性组合得到

\begin{aligned} S_t &=S_{t-1} -\underbrace{(S_{t-1}k_t)}_{v_t^{\mathrm{old}}}k_t^\top +\underbrace{\left(\beta_t v_t+(1-\beta_t)S_{t-1}k_t\right)}_{v_t^{\mathrm{new}}}k_t^\top \\ &=S_{t-1}\left(I-\beta_tk_tk_t^\top\right)+\beta_tv_tk_t^\top . \end{aligned}

  • DeltaNet 实现一种以广义 Householder 矩阵I-\beta_tk_tk_t^\top为状态转移矩阵的一阶线性递推。尽管 DeltaNet 在关联回忆和语言建模任务中表现出色,但由于计算效率较低,其长期以来并未受到广泛关注

Chunkwise parallel form:分块并行与 WY 表示法

\mathbf{S}_{[t]}^r = \mathbf{S}_{[t]} \underbrace{\left( \prod_{i=1}^r \left(\mathbf{I} - \beta_{[t]}^i \boldsymbol{k}_{[t]}^i {\boldsymbol{k}_{[t]}^i}^\top\right) \right)}_{:=\mathbf{P}_{[t]}^r} + \underbrace{\sum_{i=1}^r \left( \beta_{[t]}^i \boldsymbol{v}_{[t]}^i {\boldsymbol{k}_{[t]}^i}^\top \prod_{j=i+1}^r \left( \mathbf{I} - \beta_{[t]}^j \boldsymbol{k}_{[t]}^j {\boldsymbol{k}_{[t]}^j}^\top \right) \right)}_{:=\mathbf{H}_{[t]}^r}

  • 其中P_{[t]}^j涉及广义 Householder 矩阵的累积乘积,这可以通过经典的 WY 表示法进行优化

\mathbf{P}_{[t]}^r = \mathbf{I} - \sum_{i=1}^r \mathbf{w}_{[t]}^i {\boldsymbol{k}_{[t]}^i}^\top \in \mathbb{R}^{d_k \times d_k} \qquad \mathbf{w}_{[t]}^r = \beta_{[t]}^r \left( \boldsymbol{k}_{[t]}^r - \sum_{i=1}^{r-1} \left( \mathbf{w}_{[t]}^i ({\boldsymbol{k}_{[t]}^i}^\top \boldsymbol{k}_{[t]}^r) \right) \right) \in \mathbb{R}^{d_k}

\mathbf{H}_{[t]}^r = \sum_{i=1}^r \mathbf{u}_{[t]}^i {\boldsymbol{k}_{[t]}^i}^\top \in \mathbb{R}^{d_v \times d_k} \qquad \mathbf{u}_{[t]}^r = \beta_{[t]}^r \left( \boldsymbol{v}_{[t]}^r - \sum_{i=1}^{r-1} \left( \mathbf{u}_{[t]}^i ({\boldsymbol{k}_{[t]}^i}^\top \boldsymbol{k}_{[t]}^r) \right) \right) \in \mathbb{R}^{d_v}

  • UT transform 让 W、U 一次性用矩阵运算算出来

\mathbf{P}_{[t]} = \mathbf{I} - \mathbf{W}_{[t]}^\top \mathbf{K}_{[t]} \in \mathbb{R}^{d_k \times d_k}, \qquad \mathbf{H}_{[t]} = \mathbf{U}_{[t]}^\top \mathbf{K}_{[t]} \in \mathbb{R}^{d_v \times d_k}

\begin{aligned}
\mathbf{T}_{[t]} &= \left[ \mathbf{I} + \text{strictLower} \left( \text{diag}(\beta_{[t]}) \mathbf{K}_{[t]} \mathbf{K}_{[t]}^\top \right) \right]^{-1} \text{diag} (\beta_{[t]}) \in \mathbb{R}^{C \times C} \\
\mathbf{W}_{[t]} &= \mathbf{T}_{[t]} \mathbf{K}_{[t]} \in \mathbb{R}^{C \times d_k}, \qquad \mathbf{U}_{[t]} = \mathbf{T}_{[t]} \mathbf{V}_{[t]} \in \mathbb{R}^{C \times d_v}
\end{aligned}

  • 硬件高效的 DeltaNet 分块算法,该算法利用矩阵乘法,从而实现基于张量核心的 GPU 优化

\begin{aligned}
\mathbf{S}_{[t+1]} &= \mathbf{S}_{[t]}\mathbf{P}_{[t]} + \mathbf{H}_{[t]} = \mathbf{S}_{[t]} + \left(\mathbf{U}_{[t]} - \mathbf{W}_{[t]}\mathbf{S}_{[t]}^\top \right)^\top \mathbf{K}_{[t]} &&\in \mathbb{R}^{d_v \times d_k} \\
\mathbf{O}_{[t]} &= \mathbf{Q}_{[t]}\mathbf{S}_{[t]}^\top + \left(\mathbf{Q}_{[t]}\mathbf{K}_{[t]}^\top \odot \mathbf{M} \right) \left(\mathbf{U}_{[t]} - \mathbf{W}_{[t]}\mathbf{S}_{[t]}^\top \right) &&\in \mathbb{R}^{C \times d_v}
\end{aligned}

🦊 GATED DELTA NETWORKS

所提出的门控 Delta 规则简单而有效

S_t=S_{t-1}\left(\alpha_t\left(I-\beta_tk_tk_t^\top\right)\right)+\beta_tv_tk_t^\top.

  • 其中依赖于输入数据的门控项\alpha_t\in(0,1)用于控制状态衰减。该形式统一门控机制与 Delta 规则各自的优势:门控项能够实现自适应记忆管理,而 Delta 更新结构则有助于有效学习键值关联

循环状态更新可以被解释为某个在线学习问题的闭式解。近期的线性 RNN 架构通常会在在线学习目标中引入正则化项,以防止当前状态相对于先前状态发生过大偏移,从而实现信息保留。然而,当状态中积累的信息趋于饱和时,这种保留机制会带来问题。在这种情况下,每个状态都会编码多个信息片段的叠加表示,从而增加精确检索的难度

为解决这一局限,Mamba2 和 Gated DeltaNet 引入自适应缩放因子\alpha_t,以放松正则化项对状态更新的约束,使 S_t与S_{t-1}之间能够发生受控偏移。通过选择性遗忘,该机制实现动态记忆管理,并可用于滤除无关信息

  • 线性注意力和 Mamba2 使用简单的负内积损失-\langle S_tk_t,v_t\rangle,而 Longhorn 则采用表达能力更强的在线回归目标\lVert S_tk_t-v_t\rVert^2,以更有效地建模键值关联。由此得到的 Longhorn 更新规则与 Delta 更新规则高度相似,这表明在上下文内关联回忆任务中,门控 Delta 规则可能优于 Mamba2

从快速权重编程、测试时训练和测试时回归的角度来看,隐状态S可以被解释为一个(快速)权重矩阵,而 Delta 规则则通过测试时随机梯度下降优化在线回归目标L(S_t)=\frac{1}{2}\lVert S_tk_t-v_t\rVert^2

S_{t+1}=S_t-\beta_t\nabla L(S_t)=S_t-\beta_t(S_tk_t-v_t)k_t^\top=S_t\left(I-\beta_tk_tk_t^\top\right)+\beta_tv_tk_t^\top.

  • 其中\beta_t表示自适应学习率。门控 Delta 规则可以被理解为在 SGD 更新中引入自适应权重衰减项\alpha_t,而权重衰减是深度学习中被广泛采用的一项技术。与此同时,Titans 的研究也表明,在 RNN 的测试时 SGD 更新中引入权重衰减机制能够取得良好效果

门控负责过滤和清除无关记忆,Delta 规则负责稳定保留并精确写入关键信息,二者结合能够兼顾长程记忆、记忆管理与复杂模式记忆

  • S-NIAH-1:衰减不利于长期记忆保持。 在重复合成上下文中,模型只需长期保存少量关键信息。DeltaNet 在不同序列长度下均接近满分;Mamba2 因持续衰减历史状态,在序列长度超过 2K 后性能明显下降;Gated DeltaNet 借助 Delta 规则增强记忆保持能力,因此退化幅度较小
  • S-NIAH-2/3:门控机制有助于过滤无关信息。 在真实文章上下文中,模型需要暂存大量潜在相关信息。由于状态容量固定,若缺乏有效的记忆清除机制,信息会不断叠加并产生记忆碰撞。DeltaNet 在长序列上的性能显著下降,而 Mamba2 和 Gated DeltaNet 可通过门控机制选择性遗忘无关内容,因此表现更加稳定
  • S-NIAH-3:Delta 规则有助于复杂模式记忆。 当待记忆的值由简单数字替换为 UUID 后,任务对精确写入和复杂模式记忆提出更高要求。Mamba2 的性能迅速下降,而 Gated DeltaNet 表现更好,说明 Delta 规则具有更强的键值关联学习与精确记忆能力

把 gated delta rule 的逐 token 递推改写成按 chunk 并行的训练算法,让主要计算变成大矩阵乘法 + 少量三角求解,从而吃满 GPU tensor core

  • Chunk 内部分展开:把逐 token 递推改写成块内两项分解

\mathbf{S}_{[t]}^r = \mathbf{S}_{[t]} \underbrace{\left( \prod_{i=1}^r {\color{blue}\alpha_{[t]}^i} \left( \mathbf{I} - \beta_{[t]}^i \boldsymbol{k}_{[t]}^i {\boldsymbol{k}_{[t]}^i}^\top \right) \right)}_{:=\mathbf{F}_{[t]}^r} + \underbrace{\sum_{i=1}^r \left( \beta_{[t]}^i \boldsymbol{v}_{[t]}^i {\boldsymbol{k}_{[t]}^i}^\top \prod_{j=i+1}^r {\color{blue}\alpha_{[t]}^j} \left( \mathbf{I} - \beta_{[t]}^j \boldsymbol{k}_{[t]}^j {\boldsymbol{k}_{[t]}^j}^\top \right) \right)}_{:=\mathbf{G}_{[t]}^r}

  • \mathbf{G}_{[t]}^r的低秩展开:把写入累积写成\sum u_i k_i^\top

\mathbf{F}_{[t]}^r = {\color{blue}\gamma_{[t]}^r} \mathbf{P}_{[t]}^r = {\color{blue}\overleftarrow{\mathbf{P}_{[t]}^r}}

\mathbf{G}_{[t]}^r = \sum_{i=1}^r \frac{{\color{blue}\gamma_{[t]}^r}}{{\color{blue}\gamma_{[t]}^i}} \tilde{\mathbf{u}}_{[t]}^i {\boldsymbol{k}_{[t]}^i}^\top \in \mathbb{R}^{d_v \times d_k} \qquad \tilde{\mathbf{u}}_{[t]}^r = \beta_{[t]}^r \left( \boldsymbol{v}_{[t]}^r - \sum_{i=1}^{r-1} \left( \tilde{\mathbf{u}}_{[t]}^i \left( \frac{{\color{blue}\gamma_{[t]}^r}}{{\color{blue}\gamma_{[t]}^i}} {\boldsymbol{k}_{[t]}^i}^\top \boldsymbol{k}_{[t]}^r \right) \right) \right) \in \mathbb{R}^{d_v}

  • UT transform:把递推升级成块级矩阵一次算

\widetilde{\mathbf{U}}_{[t]} = \left[ \mathbf{I} + \text{strictLower} \left( \text{diag}(\beta_{[t]}) \left( {\color{blue}\Gamma_{[t]}} \odot \mathbf{K}_{[t]} \mathbf{K}_{[t]}^\top \right) \right) \right]^{-1} \text{diag}(\beta_{[t]}) \mathbf{V}_{[t]} \in \mathbb{R}^{C \times d_v}

  • 最终状态更新式:只用块末状态 + 当前块的矩阵运算得到下一块状态

\begin{aligned}
\mathbf{S}_{[t+1]} &= {\color{blue}\overrightarrow{\mathbf{S}_{[t]}}} + \left( \widetilde{\mathbf{U}}_{[t]} - {\color{blue}\overleftarrow{\mathbf{W}_{[t]}}} \mathbf{S}_{[t]}^\top \right)^\top {\color{blue}\overrightarrow{\mathbf{K}_{[t]}}} &&\in \mathbb{R}^{d_v \times d_k} \\
\mathbf{O}_{[t]} &= {\color{blue}\overleftarrow{\mathbf{Q}_{[t]}}} \mathbf{S}_{[t]}^\top + (\mathbf{Q}_{[t]} \mathbf{K}_{[t]}^\top \odot \Gamma_{[t]}) \left( \widetilde{\mathbf{U}}_{[t]} - {\color{blue}\overleftarrow{\mathbf{W}_{[t]}}} \mathbf{S}_{[t]}^\top \right) &&\in \mathbb{R}^{C \times d_v}
\end{aligned}

{\color{blue}\overleftarrow{\boldsymbol{q}_{[t]}^r}} = {\color{blue}\gamma_{[t]}^r} \boldsymbol{q}_{[t]}^r, \quad {\color{blue}\overleftarrow{\mathbf{w}_{[t]}^r}} = {\color{blue}\gamma_{[t]}^r} \mathbf{w}_{[t]}^r, \quad {\color{blue}\overrightarrow{\boldsymbol{k}_{[t]}^r}} = \frac{{\color{blue}\gamma_{[t]}^C}}{{\color{blue}\gamma_{[t]}^r}} \boldsymbol{k}_{[t]}^r, \quad {\color{blue}\overrightarrow{\mathbf{S}_{[t]}}} = {\color{blue}\gamma_{[t]}^C} \mathbf{S}_{[t]}

🐼 Kimi Linear: An Expressive, Efficient Attention Architecture

随着大语言模型逐步演化为能力日益强大的智能体,其推理阶段的计算需求——尤其是在长时程任务与强化学习场景下,正日益成为核心瓶颈

  • 随着研究范式转向 RL 测试时扩展,模型需要在推理过程中处理更长的交互轨迹、工具调用过程以及复杂的决策空间,这进一步暴露标准注意力机制在效率方面的根本性局限
  • Softmax 注意力具有二次时间复杂度,同时其键值缓存会随序列长度线性增长,由此带来显著的计算与内存开销,限制系统吞吐率、上下文长度扩展能力以及实时交互性能

线性注意力降低注意力计算复杂度提供一种具有理论依据的方法。然而,由于表达能力有限,传统线性注意力在语言建模任务上的表现长期逊于 Softmax 注意力,甚至在较短序列上也是如此

  • 近年来,两类关键技术进展显著缩小二者之间的性能差距:一类是门控或衰减机制,另一类是 Delta Rule(增量规则)
  • 纯线性结构仍然受到有限状态容量这一根本约束,因此在长序列建模以及上下文内信息检索方面存在理论上的困难。近年来逐渐出现一类混合架构:在大部分层中采用计算效率更高的线性注意力,同时周期性地插入少量具有全局建模能力的 Softmax 注意力层,以在模型质量与计算效率之间取得折中
  • 然而,已有混合模型往往仅在有限规模下进行验证,或者缺少覆盖多种任务与基准的系统性评估。因而,其核心挑战仍然在于:如何设计一种注意力架构,使其在模型质量上达到甚至超过全注意力,同时在速度和内存开销方面获得显著收益

Kimi Linear,一种面向智能体智能与测试时扩展需求设计的混合线性注意力架构,其目标是在显著提升推理效率的同时不牺牲模型质量

  • Kimi Linear 的核心模块是 Kimi Delta Attention(KDA),这是一种具有硬件高效特性的线性注意力机制,它在 Gated DeltaNet 的基础上引入更加细粒度的门控机制
  • GDN 与 Mamba2 类似,采用的是较为粗粒度的逐头遗忘门,即同一个注意力头中的多个特征维度共享统一的遗忘行为
  • KDA 则引入逐通道门控:每个特征维度都拥有独立的遗忘速率,其思想与 Gated Linear Attention(GLA)类似。这种更细粒度的设计能够更加精确地调控有限状态 RNN 中的信息保留与遗忘过程,从而进一步释放 RNN 风格模型在混合架构中的建模潜力
  • KDA 使用一种特殊形式的对角加低秩矩阵(Diagonal-Plus-Low-Rank,DPLR)对其状态转移动力学进行参数化。基于这一特殊结构,进一步设计一种定制化的分块并行算法。相比通用 DPLR 形式,该算法能够显著降低计算开销,同时仍与经典 Delta Rule 保持一致

在整体架构上,Kimi Linear 按照统一的 3:1 比例交替排列 KDA 层与全注意力层,即每三个 KDA 层配置一个全注意力层。借助这种混合结构,在长序列生成过程中,模型的内存占用和 KV Cache 使用量最高可降低 75%;与此同时,周期性插入的全注意力层仍能够提供全局信息交互能力,从而避免纯线性注意力在长程信息传播上的局限。本文的主要贡献如下

  • Kimi Delta Attention:提出一种新的线性注意力机制,通过更细粒度的门控方式改进 Gated Delta Rule,在提升循环状态记忆管理能力的同时兼顾硬件执行效率
  • Kimi Linear 架构:提出一种按照 3:1 的 KDA 层与全局注意力层比例构建的混合架构,在显著降低内存占用的同时,其模型质量能够超过全注意力基线
  • 大规模、公平的实证验证:在 1.4 万亿 token 的训练规模下进行对照实验,结果表明 Kimi Linear 在短上下文、长上下文以及 RL 风格评估任务上均优于全注意力及其他基线方法

🦉 Linear Attention and the Gated Delta Rule

线性注意力(Linear Attention)作为在线学习:线性注意力维护一个矩阵形式的循环状态,用于持续累积键–值(key–value)关联

\mathbf{S}_t = \mathbf{S}_{t-1} + \boldsymbol{k}_t \boldsymbol{v}_t^\top, \quad \boldsymbol{o}_t = \mathbf{S}_t^\top \boldsymbol{q}_t.

  • 从快速权重的视角来看,S_t可以被视为一种联想记忆,用于存储短暂的键到值的映射关系。上述更新过程可以理解为针对一个无界相关性目标执行梯度下降

\mathcal{L}_t(\mathbf{S}) = -\langle \mathbf{S}^\top \boldsymbol{k}_t, \boldsymbol{v}_t \rangle,

  • 该目标会持续强化最近出现的键–值对,但不包含任何遗忘机制。然而,这样的目标函数并没有提供一个明确的判据来决定哪些记忆应当被删除。随着序列不断增长,历史信息会持续累积,从而导致不同记忆之间的干扰在长上下文场景中不断加剧

DeltaNet 将上述循环更新重新解释为针对一个重构目标进行在线梯度下降,其损失函数定义为

L_t(S)=\frac{1}{2}\left\|S^\top k_t-v_t\right\|^2.

  • 以学习率\beta_t执行一步梯度下降,可以得到

S_t
=
S_{t-1}
-
\beta_t\nabla_S L_t(S_{t-1})
=
(I-\beta_t k_tk_t^\top)S_{t-1}
+
\beta_t k_tv_t^\top.

  • 这一更新规则即经典的 Delta Rule(增量规则)。在这种解释下,S被视为一个可学习的联想记忆矩阵:当新的键–值对(k_t,v_t)到来时,模型并非简单地将其累加进状态,而是持续修正已有记忆,使其逐渐逼近映射k_t\mapsto v_t​
  • 其中,(I-\beta_t k_tk_t^\top)所对应的 rank-1 更新结构可以等价地视为一种广义的 Householder 变换。这种特殊的低秩结构十分适合在硬件上实现高效的分块并行计算。更直观地说,普通线性注意力的更新是直接写入新记忆S_t=S_{t-1}+k_tv_t^\top,而 DeltaNet 则首先根据当前状态对k_t的预测结果S_{t-1}^\top k_t计算误差,然后只写入这一预测误差对应的修正量。因此,它本质上是一种带误差校正的在线记忆更新机制

Gated DeltaNet 将遗忘解释为权重衰减

  • 尽管 DeltaNet 的重构式更新能够提高记忆学习的稳定性,但它仍然存在一个问题:过时的关联关系不会被主动遗忘,而可能被长期保留在状态中。为解决这一问题,Gated DeltaNet 引入一个标量遗忘门\alpha_t\in[0,1],得到如下更新规则

S_t
=
\alpha_t(I-\beta_tk_tk_t^\top)S_{t-1}
+
\beta_tk_tv_t^\top.

  • \alpha_t可以被理解为作用于快速权重S_t上的一种权重衰减。它为模型引入显式的遗忘机制,其作用类似于一种数据依赖的L_2正则化
  • 当\alpha_t接近1时,历史状态基本得到保留;而当\alpha_t较小时,已有记忆会被更强地衰减。因此,模型可以根据当前输入动态控制记忆的生命周期。这种简单但有效的改动,为控制记忆持续时间以及缓解不同历史关联之间的干扰提供一种具有明确理论解释的机制,从而改善模型的训练稳定性和长上下文泛化能力,同时保留 DeltaNet 原有的可并行计算结构

GDN 可以被解释为一种乘法式位置编码。其中,状态转移矩阵并非像传统位置编码那样固定,而是数据依赖且可学习的。因此,与 RoPE 所采用的正交旋转变换相比,GDN 实际上放宽其正交性约束,允许模型学习更加一般化的状态转移动力学

🦊 Kimi Delta Attention: Improving Delta Rule with Fine-grained Gating

KDA 在 GDN 标量衰减机制的基础上进行改进,引入更加细粒度的对角门控矩阵\operatorname{Diag}(\alpha_t),从而能够对记忆衰减和位置感知能力进行更精细的控制

\boxed{\mathbf{S}_t = (\mathbf{I} - \beta_t \boldsymbol{k}_t \boldsymbol{k}_t^\top) \textcolor{#C00000}{\mathrm{Diag}(\boldsymbol{\alpha}_t)} \mathbf{S}_{t-1} + \beta_t \boldsymbol{k}_t \boldsymbol{v}_t^\top \in \mathbb{R}^{d_k \times d_v}; \qquad \boldsymbol{o}_t = \mathbf{S}_t^\top \boldsymbol{q}_t \in \mathbb{R}^{d_v}}

  • 通过将上式中的递推关系部分展开为分块(chunk-wise)形式,可得

\mathbf{S}_{[t]}^r
=\underbrace{\left(\overleftarrow{\prod}_{i=1}^r
\left[(\mathbf I-\beta_{[t]}^i\mathbf k_{[t]}^i\mathbf k_{[t]}^{i\top})
{\color{red}\operatorname{Diag}(\alpha_{[t]}^i)}\right]\right)}_{:=\mathbf P_{[t]}^r}\mathbf S_{[t]}^0
+\underbrace{\sum_{i=1}^r\left(\overleftarrow{\prod}_{j=i+1}^r
\left[(\mathbf I-\beta_{[t]}^j\mathbf k_{[t]}^j\mathbf k_{[t]}^{j\top})
{\color{red}\operatorname{Diag}(\alpha_{[t]}^j)}\right]\right)
\beta_{[t]}^i\mathbf k_{[t]}^i\mathbf v_{[t]}^{i\top}}_{:=\mathbf H_{[t]}^r}

  • WY 表示(WY Representation)通常用于将一系列 rank-1 更新压缩为一个紧凑的单一表示。遵循 Comba 中关于\mathbf{P}的构造形式,从而减少后续计算过程中额外进行一次矩阵求逆的需求

\mathbf{P}_{[t]}^{r}
=
\color{red}{
\operatorname{Diag}\!\left(\gamma_{[t]}^{r}\right)
}
-
\sum_{i=1}^{r}
\color{red}{
\operatorname{Diag}\!\left(\gamma_{[t]}^{i\to r}\right)
}
\mathbf{k}_{[t]}^{i}
\mathbf{w}_{[t]}^{i\top}
\qquad
\mathbf{H}_{[t]}^{r}
=
\sum_{i=1}^{r}
\color{red}{
\operatorname{Diag}\!\left(\gamma_{[t]}^{i\to r}\right)
}
\mathbf{k}_{[t]}^{i}
\mathbf{u}_{[t]}^{i\top}

  • 其中,辅助向量\mathbf{w}_t \in \mathbb{R}^{d_k}和\mathbf{u}_t \in \mathbb{R}^{d_v}通过如下递推关系计算

\mathbf{w}_{[t]}^{r}
=
\beta_{[t]}^{r}
\left(
\color{red}{
\operatorname{Diag}\!\left(\gamma_{[t]}^{r}\right)
}
\mathbf{k}_{[t]}^{r}
-
\sum_{i=1}^{r-1}
\mathbf{w}_{[t]}^{i}
\left(
\mathbf{k}_{[t]}^{i\top}
\color{red}{
\operatorname{Diag}\!\left(\gamma_{[t]}^{i\to r}\right)
}
\mathbf{k}_{[t]}^{r}
\right)
\right)

\mathbf{u}_{[t]}^{r\top}
=
\beta_{[t]}^{r}
\left(
\mathbf{v}_{[t]}^{r\top}
-
\sum_{i=1}^{r-1}
\mathbf{u}_{[t]}^{i\top}
\left(
\mathbf{k}_{[t]}^{i\top}
\color{red}{
\operatorname{Diag}\!\left(\gamma_{[t]}^{i\to r}\right)
}
\mathbf{k}_{[t]}^{r}
\right)
\right)

  • UT 变换(UT transform):采用 UT 变换来减少非矩阵乘法FLOPs

\mathbf{M}_{[t]}
=
\left(
\mathbf{I}
+
\operatorname{StrictTril}
\left(
\operatorname{Diag}\!\left(\beta_{[t]}\right)
\left(
\color{red}{
\mathbf{\Gamma}_{[t]}^{1\to C}
}
\odot
\mathbf{K}_{[t]}
\right)
\left(
\frac{
\mathbf{K}_{[t]}
}{
\color{red}{
\mathbf{\Gamma}_{[t]}^{1\to C}
}
}
\right)^{\top}
\right)
\right)^{-1}
\operatorname{Diag}\!\left(\beta_{[t]}\right)

\mathbf{W}_{[t]}
=
\mathbf{M}_{[t]}
\left(
\color{red}{
\mathbf{\Gamma}_{[t]}^{1\to C}
}
\odot
\mathbf{K}_{[t]}
\right),
\qquad
\mathbf{U}_{[t]}
=
\mathbf{M}_{[t]}
\mathbf{V}_{[t]}

  • 下三角矩阵的逆可以通过高斯消元中的前向代入,以逐行迭代的方式高效计算。等价地,以矩阵形式表示时,可以按照分块方式对状态进行更新

\mathbf{S}_{[t+1]}
=
\color{red}{
\operatorname{Diag}\!\left(\gamma_{[t]}^{C}\right)
}
\mathbf{S}_{[t]}
+
\left(
\color{red}{
\mathbf{\Gamma}_{[t]}^{i\to C}
}
\odot
\mathbf{K}_{[t]}
\right)^{\top}
\left(
\mathbf{U}_{[t]}
-
\mathbf{W}_{[t]}
\mathbf{S}_{[t]}
\right)
\in
\mathbb{R}^{d_k\times d_v}

  • 在输出阶段,采用块间递归与块内并行相结合的策略,以最大化矩阵乘法吞吐率,从而充分发挥 Tensor Core 的计算潜力

\mathbf{O}_{[t]}
=
\underbrace{
\left(
\color{red}{\mathbf{\Gamma}_{[t]}^{1\to C}}
\odot
\mathbf{Q}_{[t]}
\right)
\mathbf{S}_{[t]}
}_{\text{inter chunk}}
+
\underbrace{
\operatorname{Tril}
\left(
\left(
\color{red}{\mathbf{\Gamma}_{[t]}^{1\to C}}
\odot
\mathbf{Q}_{[t]}
\right)
\left(
\frac{
\mathbf{K}_{[t]}
}{
\color{red}{\mathbf{\Gamma}_{[t]}^{1\to C}}
}
\right)^{\top}
\right)
}_{\text{intra chunk}}
\underbrace{
\left(
\mathbf{U}_{[t]}
-
\mathbf{W}_{[t]}\mathbf{S}_{[t]}
\right)
}_{\text{``pseudo''-value term}}
\in
\mathbb{R}^{C\times d_v}

KDA 和通用 DPLR 一样,都支持细粒度的逐通道衰减,但这种机制容易在 chunk 计算中的除法操作里产生数值精度问题。GLA 通常通过 log 域和 FP32 二级分块来保证稳定性,但会明显拖慢计算。KDA 将 DPLR 中的两个低秩向量都约束为与 k_t相关,使结构更简单,从而减少二级分块和矩阵乘法次数,更充分利用半精度 Tensor Core

🐼 The Kimi Linear Model Architecture

Kimi Linear 的主体架构沿用 Moonlight 的设计。在细粒度门控之外,还引入若干组件,以进一步增强 Kimi Linear 的表达能力

神经参数化(Neural Parameterization):设x_t \in \mathbb{R}^{d}表示第t个 token 的输入表示,则对于每个注意力头h,KDA 的输入计算如下

q_t^h, k_t^h
=
\operatorname{L2Norm}
\left(
\operatorname{Swish}
\left(
\operatorname{ShortConv}
\left(
W_{q/k}^h x_t
\right)
\right)
\right)
\in
\mathbb{R}^{d_k}

v_t^h
=
\operatorname{Swish}
\left(
\operatorname{ShortConv}
\left(
W_v^h x_t
\right)
\right)
\in
\mathbb{R}^{d_v}

\color{red}{\alpha_t^h}
=
f
\left(
W_\alpha^\uparrow
W_\alpha^\downarrow
x_t
\right)
\in
[0,1]^{d_k}

\beta_t^h
=
\operatorname{Sigmoid}
\left(
W_\beta^h x_t
\right)
\in
[0,1]

  • 其中d_k和d_v分别表示 key 和 value 的 head dimension,在所有实验中二者均被设置为 128,对于上述输入表示,首先应用 ShortConv,随后使用 Swish 激活函数。对于q和k的表示,进一步采用 L2Norm 进行归一化,以确保特征值的稳定性
  • 逐通道衰减参数\color{red}{\alpha_t^h}通过一个低秩投影进行参数化,即采用W_\alpha^\downarrow和W_\alpha^\uparrow,其秩被设置为与 head dimension 相同;同时采用与 GDN 和 Mamba 中类似的衰减函数f(\cdot)​
  • 在通过W_o \in \mathbb{R}^{d\times d}进行输出投影之前,采用逐头 RMSNorm,以及一个数据依赖的门控机制,其参数化形式为

o_t
=
W_o
\left(
\operatorname{Sigmoid}
\left(
W_g^\uparrow W_g^\downarrow x_t
\right)
\odot
\operatorname{RMSNorm}
\left(
\operatorname{KDA}
\left(
q_t,
k_t,
v_t,
\color{red}{\alpha_t},
\beta_t
\right)
\right)
\right)

  • 其中,输出门采用与遗忘门类似的低秩参数化形式;与此同时,该设计能够维持与全秩门控相当的性能,并缓解 Attention Sink 问题

混合模型架构:长上下文检索仍然是纯线性注意力的主要瓶颈,因此将 KDA 与少量完整的全局注意力层,即 Full MLA 层进行混合

  • 对于 Kimi Linear,选择逐层(layerwise)混合方式,即交替使用完整的不同类型层,而不是采用逐头(headwise)混合方式,即在同一层内部混合不同类型的注意力头。这一选择主要是因为逐层混合具有更简单的基础设施实现以及更稳定的训练过程
  • 实验结果表明,采用统一的 3:1 比例,即每 3 个 KDA 层之后配置 1 个完整 MLA 层,可以获得最佳的模型质量与吞吐率折中

MLA 层不使用位置编码(No Position Encoding, NoPE):在 Kimi Linear 中,对所有完整注意力层均采用 NoPE。这一设计将编码位置信息以及近因偏置的全部职责交由 KDA 层承担

  • KDA 被确立为模型中主要的位置感知算子,其作用类似于、甚至可以认为强于短卷积或 SWA 等辅助组件,将全局 NoPE 注意力与专门的位置感知机制相结合,可以获得具有竞争力的长上下文性能
  • NoPE 在实际应用中具有多方面优势,尤其适用于 MLA。首先,NoPE 使得 MLA 在推理阶段能够转换为高效的纯 Multi-Query Attention。其次,NoPE 简化长上下文训练,因为它无需再对 RoPE 的相关参数进行调整,例如调整频率基数,或者采用 YaRN 等上下文扩展方法

🦉 Experiments Insight

  • 细粒度衰减显著提升有限状态记忆的选择性与稳定性:相较于 GDN 的标量衰减,KDA 的逐通道衰减能够更精确地控制信息保留与遗忘。在 Palindrome、MQAR 和 Stack 等记忆密集型任务上,KDA 随序列长度增长始终表现最优,并在检索任务上表现出更快的收敛速度,说明 fine-grained decay 与 delta-rule 更新具有互补作用
  • 线性注意力与全局注意力之间存在明确的最优混合比例,而非全注意力越多越好:实验表明,3:1 的 KDA/MLA 比例在训练损失、泛化性能与推理开销之间取得最佳平衡。进一步增加 KDA 比例会损害验证性能,而增加 MLA 比例则带来额外推理成本。这表明少量全局注意力足以弥补线性状态在全局检索上的局限
  • 局部建模与输出调制对 KDA 的整体性能具有不可忽视的贡献:移除 output gate 会导致性能下降,而 Sigmoid gate 明显优于 Swish gate;同时,去除 ShortConv 也会带来验证性能退化。这说明 Kimi Linear 的性能提升并非仅来源于新的 recurrent update,局部卷积所提供的短程归纳偏置,以及输出门控对信息流的动态调制,同样是架构有效性的关键组成部分
  • 位置建模与全局内容检索的功能解耦有利于长上下文泛化:Kimi Linear 使用 KDA 承担主要的位置感知功能,而 MLA 层采用 NoPE。相较于 RoPE 版本,该设计在短上下文性能基本不受影响的情况下显著改善长上下文表现,说明避免在全局注意力层中引入过强的位置偏置,有助于模型向更长序列外推
  • KDA 的优势不仅体现为推理效率,还表现为更优的训练与 RL 扩展特性:Kimi Linear 在 compute-optimal scaling 下获得约 1.16× 的计算效率,同时在 RL 阶段表现出比 MLA 更快的训练和测试性能提升

🦊 Kimi Delta Attention as learnable position embeddings

标准 Transformer 注意力机制在设计上对输入序列的顺序是不可知的,因此必须显式引入位置编码。在各种位置编码方法中,RoPE 凭借其有效性,已经成为现代大语言模型中事实上的标准方案。类似 RoPE 的乘法式位置编码机制,可以通过如下广义注意力形式进行分析

s_{t,i}=\mathbf q_t^\top\mathbf R_t^\top\mathbf R_i\mathbf k_i
=\mathbf q_t^\top\left(\prod_{j=i+1}^t(\mathbf R_j^\top\mathbf R_{j-1})\right)\mathbf k_i

  • 其中,第t个 query \mathbf{q}_t与第i个 key \mathbf{k}_i之间的位置关系,通过累积矩阵乘积体现。RoPE 将变换矩阵\mathbf{R}_j定义为由d_k/2个二维旋转矩阵构成的块对角矩阵,其中第k个二维旋转块为

\mathbf{R}_j^k
=
\begin{pmatrix}
\cos(j\theta_k) & -\sin(j\theta_k) \\
\sin(j\theta_k) & \cos(j\theta_k)
\end{pmatrix}

  • 其对应的每二维角频率为\theta_k。利用旋转矩阵的性质,即\mathbf{R}_{i-t}=\mathbf{R}_t^\top\mathbf{R}_i,绝对位置信息\mathbf{R}_t和\mathbf{R}_i可以分别作用于\mathbf{q}_t和\mathbf{k}_i,随后被转换为相对位置t-i的编码

\prod_{j=i+1}^{t}\left((\mathbf R_j^k)^\top\mathbf R_{j-1}^k\right)
=(\mathbf R_t^k)^\top\mathbf R_i^k
=\begin{pmatrix}
\cos((i-t)\theta_k)&-\sin((i-t)\theta_k)\\
\sin((i-t)\theta_k)&\cos((i-t)\theta_k)
\end{pmatrix}

  • 进一步地,采用 gated delta rule 的线性注意力也可以写成与之可比较的形式

\mathbf o_t=\sum_{i=1}^t\beta_i\left(\mathbf q_t^\top
\left(\overleftarrow{\prod}_{j=i+1}^t\left[(\mathbf I-\beta_j\mathbf k_j\mathbf k_j^\top){\color{red}\mathbf A_j}\right]\right)
\mathbf k_i\right)\mathbf v_i,\qquad \mathbf S_0=\mathbf 0

  • GDN 可以被解释为一种乘法式位置编码:其状态转移矩阵由数据动态决定,因此放宽 RoPE 所施加的正交性约束,并可能具备更强的表达能力。这一性质也为缓解 RoPE 已知的外推问题提供一种潜在解决方案,因为 RoPE 的固定频率可能使模型过拟合于训练阶段所见的上下文长度。近期一些工作采用 partial RoPE 等折中方案,甚至完全取消显式位置编码,即 NoPE
  • 鉴于 GDN 在功能上承担类似 RoPE 的角色,在模型的全局完整注意力层中采用 NoPE,使位置信息能够由所提出的 KDA 动态建模
  • RoPE 的一个重要优势在于其细粒度位置编码能力:它为不同的二维特征对分配不同的旋转频率,其作用类似于沿特征维度执行非均匀傅里叶变换。相比之下,标准 GDN 对每个 attention head 仅采用单一标量 decay,缺少这种逐维度的多样性。因此,进一步提出带有可学习逐通道门控的 KDA

(Gated) DeltaNet 可以推广至表达能力更强的 Diagonal-Plus-Low-Rank(DPLR) 结构,其形式定义为\mathbf{D}-\mathbf{a}_t\mathbf{b}_t^\top。类似结构也曾在 S4 等模型中得到研究,其中采用静态 DPLR 形式作为状态转移矩阵。在实际计算中,该矩阵通常会被联合对角化到复数域,从而使其表达能力受到对角变换的限制

  • 虽然 DPLR 结构能够引入更丰富的模型交互,并可能通过其 key-value 更新规则增强信息召回能力,但它也存在一个显著局限:计算成本较高且并行化能力较弱。这使得 DPLR 在大规模训练或实时推理场景下天然较慢,因此如何同时保持参数效率与计算效率成为一个重要的设计挑战
  • 为解决这一问题,KDA 引入 DPLR 的一种受约束变体

\mathbf{S}_t
=
\left(
\color{red}{\operatorname{Diag}(\boldsymbol{\alpha}_t)}
-
\beta_t
\mathbf{k}_t
\mathbf{k}_t^\top
\color{red}{\operatorname{Diag}(\boldsymbol{\alpha}_t)}
\right)
\mathbf{S}_{t-1}
+
\beta_t
\mathbf{k}_t
\mathbf{v}_t^\top

\mathbf{S}_t
=
\left(
\mathbf{D}
-
\mathbf{a}_t
\mathbf{b}_t^\top
\right)
\mathbf{S}_{t-1}
+
\beta_t\mathbf{k}_t
\mathbf{v}_t^\top,
\qquad
\mathrm{s.t.},
\quad
\mathbf{D}
=
\color{red}{\operatorname{Diag}(\boldsymbol{\alpha}_t)},
\quad
\mathbf{a}_t
=
\beta_t\mathbf{k}_t,
\quad
\mathbf{b}_t
=
\mathbf{k}_t\odot\boldsymbol{\alpha}_t

  • 通过将\color{red}{\boldsymbol{\alpha}_t}提取出来,KDA 可以像 GLA 一样对状态\mathbf{S}实现细粒度的乘法衰减,随后再执行类似 DeltaNet 的 Householder-style transformation,从而实现高效的状态更新

🦁 Related Works

线性注意力的演进,本质上是从降低复杂度逐步转向提高有限状态记忆的可控性。最早的线性注意力主要利用矩阵乘法的结合律,将显式的 T×T 注意力矩阵转化为固定大小的矩阵状态,从而使计算复杂度关于序列长度呈线性增长

  • 这种简单形式缺少 Softmax 注意力所具有的选择能力:历史信息不断被压缩进同一个有限状态,容易产生记忆干扰。因此,后续工作的重点并非继续降低理论复杂度,而是不断增强状态转移机制:从 RetNet 的数据无关标量衰减,到 Mamba2 的数据依赖标量衰减,再到 GLA 的逐通道对角门控
  • 研究重点逐渐从能否压缩历史转向应该以什么粒度决定哪些历史被保留、哪些被遗忘。KDA 延续这一方向,并将细粒度逐通道衰减与 Delta Rule 结合,使有限状态既具有选择性遗忘能力,也具有基于误差的记忆修正能力

从快速权重记忆的视角看,Delta Rule 改变的不只是门控形式,而是线性注意力的写入机制

  • 论文将线性注意力的循环状态视为一个低容量的联想记忆:慢权重负责学习何时存储、更新和遗忘信息,而快速权重状态则在线维护当前上下文中的键值关联。普通门控或衰减主要解决“旧记忆应该衰减多少”的问题,但并未从根本上改变新信息的写入方式
  • GDN/KDA 中的 gated delta rule 则进一步引入秩 1 纠错更新,使状态能够根据当前键值映射误差对已有记忆进行修正。因此,这里实际上包含两个相对独立的维度:门控决定记忆能够保留多久,Delta Rule 决定记忆应该如何被修正
  • KDA 的意义就在于同时增强这两个维度:逐通道衰减提升遗忘的选择性,而 Delta 更新提升记忆写入与纠错能力

线性注意力与稀疏注意力的根本区别,是压缩历史与保留历史后选择性访问两种不同的记忆范式

  • 稀疏注意力并没有真正消除历史表示,而是通过滑动窗口、路由、聚类、top-k 分块选择等方式,仅计算其中一部分 token。因此,它在精确检索和细粒度选择上通常更具优势,因为原始 KV 信息仍然可以访问;代价则是仍需维护随序列长度增长的 KV Cache,同时动态选择本身也可能引入不可忽略的额外开销
  • NSA、MoBA、DSA 等近期方法开始强调软硬件协同设计,本质上就是为让稀疏选择真正转化为 kernel 层面的实际加速。线性注意力则采用完全不同的策略:直接将历史压缩进固定大小的状态,因此在解码阶段具有更好的内存扩展特性,但也不可避免地形成信息瓶颈
  • 模型并非保存并检索完整历史,而是学习一个足以支持后续预测的压缩状态。因此,两者的核心权衡可以概括为:稀疏注意力通过保留信息换取检索精度,线性注意力通过压缩信息换取状态存储与解码效率

混合注意力的作用并非简单折中计算复杂度,而是显式弥补有限状态模型在精确随机访问上的结构性缺陷

  • 论文明确指出,即使加入门控和 Delta Rule,纯线性注意力在精确复制和极长上下文中的精确检索上仍然存在天然劣势。因此,混合架构的逻辑并不是简单地加入部分全注意力以提高性能,而是让两类机制承担不同职责:KDA 负责大部分层中的状态压缩、持续记忆和高效更新;全注意力则周期性提供对完整历史的全局访问与精确检索能力
  • 在具体实现上,混合方式又可以分为层内混合与层间混合。前者在同一层内混合不同类型的注意力头或 token 混合机制,表达更加灵活,但会引入多条计算路径,并增加缓存管理和分布式并行的复杂度;因此 Kimi Linear 选择固定比例的层间混合,使每层保持统一的计算形式,并采用 3:1 的 KDA/全注意力周期。这体现出明显的算法与系统协同设计:模型结构不仅服务于建模能力,也主动适配现有的并行、缓存与推理基础设施

NoPE 的选择进一步体现 Kimi Linear 对位置建模和内容检索的功能解耦

  • 如果混合架构中的线性注意力与全注意力都携带较强的位置偏置,不同层就可能承担重复甚至不一致的位置建模职责;与此同时,RoPE 的基频参数对上下文扩展较为敏感
  • KDA 的数据依赖状态转移负责位置与近因信息建模,而全注意力层采用 NoPE,主要负责全局内容交互与精确检索
  • 将位置感知从显式、固定频率的位置编码,转化为由循环状态动力学动态产生的位置归纳偏置。全注意力层无需再通过调整 RoPE 基频、YaRN 等方法适配更长上下文,也降低混合架构在上下文扩展过程中的位置编码耦合

🐧 Gated DeltaNet-2: Decoupling Erase and Write in Linear Attention

线性递归注意力以固定大小的递归状态替代显式注意力矩阵,将序列混合转化为线性时间复杂度的递归计算,其内存开销不再随上下文长度增长。这一设计的吸引力十分明确,但其约束同样明显:递归状态本质上是一种压缩的键值记忆,而在长上下文条件下,大量关联关系必须共享同一个有限容量的状态空间,从而使精确检索变得困难。近期工作通过增强递归机制对信息保留行为的控制能力,持续改善这一记忆瓶颈

  • Mamba-2 使用数据依赖的衰减机制来调节记忆时间跨度
  • DeltaNet 将简单的加性写入替换为 delta rule,使模型能够有针对性地覆盖由当前 key 所寻址的关联信息
  • Gated DeltaNet 进一步将 delta rule 与可学习的衰减门结合,使状态同时具备全局遗忘和定向编辑能力
  • Kimi Delta Attention 则在衰减机制上进一步细化,在 key 维度引入逐通道遗忘机制
  • Mamba-3 沿状态空间模型路线继续推进,通过指数—梯形离散化、复值状态转移,以及多输入多输出形式,构建能力更强且效率更高的递归机制
  • 这些进展显著推动递归线性模型的发展,同时也使 delta-rule 记忆中尚未解决的瓶颈更加清晰:当前的主动编辑机制仍使用同一个标量门,同时控制旧信息的擦除与新信息的写入

Gated DeltaNet-2,一种在 delta rule 中将擦除(erase)与写入(write)操作解耦的递归注意力层

  • 将二者绑定到同一个标量上,本质上是一种建模约束,因为擦除和写入实际上作用于状态的不同轴。擦除属于 key 侧操作,其作用是决定旧有读出结果中的哪些坐标应被移除;而写入属于 value 侧操作,其作用是决定输入 value 中的哪些坐标应真正写入并保存在状态中
  • Gated DeltaNet-2 保留 KDA 的逐通道衰减机制,但将原先绑定的标量 delta gate 替换为两个相互独立的门控:一个作用于 key 轴的逐通道 erase gate,以及一个作用于 value 轴的逐通道 write gate。由此,模型可以通过 decay 清除大范围的上下文信息,通过 erase 移除特定的陈旧关联,并通过 write 仅写入那些真正需要长期保留的 value 通道
  • 当 erase gate 与 write gate 被约束为同一个标量时,Gated DeltaNet-2 退化为 KDA;若进一步将 decay 也约束为标量,则可恢复为 Gated DeltaNet
  • 通过将累积的逐通道衰减吸收到秩一的 erase 因子中,该递归形式可以被改写为紧凑的 WY 表示,并保持与现有高效 delta-rule kernel 相同的高层分块计算结构

🐢 Preliminary

作为递归状态的线性注意力:考虑单个注意力头\boldsymbol{q}_t, \boldsymbol{k}_t \in \mathbb{R}^{d_k},\boldsymbol{v}_t \in \mathbb{R}^{d_v},并省略层索引。递归线性注意力层维护一个矩阵状态S_t \in \mathbb{R}^{d_k \times d_v},并利用 query 从该状态中读取信息

  • 将该递归关系在长度为 L 的序列上展开,可以得到熟悉的因果矩阵形式

  • 其中M为因果掩码。状态 S_t的大小与序列长度 L无关,而并行形式则通过矩阵乘法替代逐 token 的递归计算
  • 每一个外积都会被累加到状态中,却没有任何内容被显式移除。旧的关联会一直保留在状态中,直到后续写入通过叠加的方式间接覆盖它们

Chunkwise 形式:高效的线性递归层在训练过程中通常采用 chunkwise 调度方式。将序列划分为大小为 C的若干 chunk。令 Q_{[n]}, K_{[n]}, V_{[n]}分别表示 query、key 和 value 块,并令 S_{[n]}表示该 chunk 开始时的状态。对递归关系进行部分展开,可得

  • 递归仅存在于不同 chunk 之间,而一个 chunk 内所有 token 之间的交互都可以表示为稠密矩阵乘法。当 C固定时,这一方法仍然保持关于序列长度的线性复杂度,同时能够很好地映射到 Tensor Core 上执行

Mamba-2 在每次写入之前引入一个数据依赖的标量衰减因子

  • 这一衰减机制赋予模型一种全局遗忘操作。若定义 \gamma_t = \prod_{i=1}^{t}\alpha_i,那么在时间步 t读取更早时间步 i写入的信息时,对应的缩放因子为\gamma_t/\gamma_i。由此可以得到一个考虑衰减的注意力掩码,同时仍然保持式 (3) 中的 chunkwise 计算结构

DeltaNet 则采用另一种方式,为状态引入主动编辑操作。在写入 v_t之前,模型首先读取当前与 k_t相关联的 value,并将其从状态中减去。对于标量步长 \beta_t \in [0,1],其更新形式为

  • 当 \lVert k_t\rVert_2=1 时,矩阵 k_tk_t^{\top}是一个投影矩阵。因此,当 \beta_t=1时,该操作会覆盖 key (k_t) 所对应的关联;而当 \beta_t=0时,该关联保持不变
  • 从 fast-weight 的视角来看,上式等价于针对如下局部回归损失执行一步在线梯度更新\frac{1}{2}
\left\|
S^{\top}k_t-v_t
\right\|_2^2.

Gated DeltaNet 将上述两种操作结合起来,decay 对整个状态进行统一清除,而 delta rule 则对某个特定关联进行定向编辑。这种分工十分有效,但其中两个门控在每个注意力头内都仍然只是标量

KDA 从 decay 一侧进一步改进这一机制,将标量 \alpha_t替换为逐通道向量 \boldsymbol{\alpha}_t \in (0,1]^{d_k}。定义D_t = \operatorname{Diag}(\boldsymbol{\alpha}_t),则其更新可以写为

  • KDA 允许每一个 key 通道按照各自独立的速率进行衰减,同时保留 DeltaNet 中基于 WY 表示的高效 chunkwise 算法
  • 然而,主动门控 \beta_t仍然只是一个标量。它同时控制两件事情:从当前读取方向中擦除多少旧内容,以及写入多少新的 value 信息
  • Gated DeltaNet-2 正是从这一尚未解除的耦合关系出发,对其进一步进行改进

🐿️ Decoupling erase and write

KDA 通过将衰减扩展为逐通道形式改进 Gated DeltaNet,但式(7)中的标量\beta_t仍然同时承担着两项未必一致的决策

  • 一项决策位于键侧,决定当前读取内容中的哪些坐标应被擦除;另一项决策位于值侧,决定候选值中的哪些坐标应被写入。将这两项决策统一为一个标量,是对更新形式的限制,而非 Delta 规则本身的要求

Gated DeltaNet-2 通过 Gated Delta Rule-2 将这两项决策分离。令

  • 其中,b_t\in[0,1]^{d_k}为擦除门,\mathbf{w}_t\in[0,1]^{d_v}为写入门。擦除门对用于读取旧内容的键坐标进行加权,而写入门对待插入的值坐标进行加权。令 \mathbf{D}_t=\operatorname{Diag}(\boldsymbol{\alpha}_t)。先施加衰减,再执行当前的编辑操作,得到

  • 将式(10)称为 Gated Delta Rule-2。输出为o_t=\mathbf{S}_t^{\top}q_t。擦除矩阵的左因子仍为k_t,从而保留 Delta 规则的写入方向。右因子变为b_t\odot k_t,使读取方向具有通道选择性。写入项变为k_tz_t^{\top},使值更新具有通道选择性

当b_t=\beta_t\mathbf{1}_{d_k}且\mathbf{w}_t=\beta_t\mathbf{1}_{d_v}时,Gated Delta Rule-2 精确退化为 KDA。进一步令\boldsymbol{\alpha}_t=\alpha_t\mathbf{1}_{d_k},则退化为 Gated DeltaNet。因此,该模型将已有的标量门控更新保留为参数绑定的子空间,同时允许在擦除与写入需要不同通道结构时,在这些子空间之外进行学习

  • 该层通过对 token 表示进行独立投影,生成两个门

  • 对数衰减沿用 Gated DeltaNet 的参数化形式

  • 在实际实现中,该衰减激活在传入计算内核之前以 FP32 精度计算,从而避免累积对数衰减中的精度损失

🐹 Fast-weight update perspective

将 Gated Delta Rule-2 解释为对快权重记忆状态的在线更新。状态\mathbf{S}_t存储短暂的键值关联。对于每个 token ,模型首先构造衰减后的状态\bar{\mathbf{S}}_t=\mathbf{D}_t\mathbf{S}_{t-1},随后沿门控擦除方向e_t读取旧内容,并写入朝向门控值目标z_t的修正量。更形式化地,式(9)是以下局部在线优化问题的解

  • 第一项使新状态保持接近衰减后的记忆。第二项施加关联编辑,其残差比较门控写入目标z_t与沿e_t从\bar{\mathbf{S}}_t中读取的内容。由于

  • 使目标函数取得最小值的解为

对于表中的 Mamba-3 一行,采用 SISO 指数—梯形递推。令

  • \mathbf{R}_{1:s}是复数 SSM 视角下依赖数据的累积旋转,且在序列起始处省略前一 token 对应的项。MIMO 版本将每次秩一写入替换为沿 MIMO 秩维度的求和,同时保持相同的在线形式不变
  • Mamba-3 通过指数—梯形输入规则和依赖数据的旋转增强这种写入的表达能力,但不会从状态中减去当前读取的内容
  • Gated DeltaNet 和 KDA 则执行残差式 Delta 编辑。KDA 将衰减由标量形式改为逐通道形式,同时保留标量门控的残差\beta_t(v_t-\bar{\mathbf{S}}_t^{\top}k_t)
  • Gated DeltaNet-2 则将残差本身改为z_t-\bar{\mathbf{S}}_t^{\top}e_t=\mathbf{w}_t\odot v_t-\left(\mathbf{D}_t\mathbf{S}_{t-1}\right)^{\top}(b_t\odot k_t),从而将用于擦除的坐标与用于写入的坐标解耦

🐙 Chunkwise parallel training

Gated Delta Rule-2 保留与 KDA 相同的分块结构。考虑一个块,并省略块索引。令

  • 通过\mathbf{S}_r=\operatorname{Diag}(\gamma_r)\hat{\mathbf{S}}_r定义衰减归一化状态\hat{\mathbf{S}}_r。令\hat{\mathbf{S}}_0=\mathbf{S}_{[n]},式(10)变为纯粹的非对称 Delta 递推,这种归一化是实现高效计算形式的关键。逐通道衰减被吸收到每个秩一擦除项的两个因子中,而更新仍保持为形如\mathbf{I}-\bar{k}_r\bar{e}_r^{\top}的矩阵之积

令\mathbf{B}\in\mathbb{R}^{C\times d_k}和\mathbf{W}\in\mathbb{R}^{C\times d_v}的各行分别为b_r^{\top}和\mathbf{w}_r^{\top}。为使矩阵表示更加紧凑,令\gamma\in\mathbb{R}^{C\times d_k}的各行为\gamma_r^{\top}。令\bar{\mathbf{K}}\in\mathbb{R}^{C\times d_k}、\bar{\mathbf{E}}\in\mathbb{R}^{C\times d_k} 和\mathbf{Z}\in\mathbb{R}^{C\times d_v}的各行分别为\bar{k}_r^{\top}、\bar{e}_r^{\top}和z_r^{\top}。等价地

  • 这里,\mathbf{Y}是擦除侧辅助量,\mathbf{U}是写入侧辅助量。由于\mathbf{T}是对角元为零的三角矩阵,\mathbf{A}可通过每个块内的小规模前向代入求得。于是,块末状态为

  • 其中,\mathbf{K}_{\mathrm{tail}}的第r行为(\gamma_C/\gamma_r)\odot k_r。输出块为

  • 其中,\mathbf{Q}_{\gamma}的第r行为\gamma_r\odot q_r,并且

  • 擦除门通过\bar{\mathbf{E}}的第r行以\gamma_r\odot(b_r\odot k_r)的形式引入。写入门通过\mathbf{Z}的第r行以\mathbf{w}_r\odot v_r的形式引入。其余计算是在固定大小的块上进行三角求解和稠密矩阵乘法。采用 WY 变换,并使用融合的 Triton 内核实现这些方程

反向传播采用与前向传播相同的分解方式。梯度首先流经输出方程和块间状态递推,两者都仅涉及\mathbf{A}_{qk}、\mathbf{K}_{\mathrm{tail}}、\mathbf{Y}和\mathbf{U}。唯一新增的梯度计算是通过式(21)和式(20)的向量—雅可比积

  • 写入侧在值通道上包含一个对角门控,而擦除侧在键通道上包含另一个不同的对角门控。因此,门控因子必须保留在梯度累积的位置。令\mathbf{B}和\mathbf{W}的各行分别为b_r^{\top}和\mathbf{w}_r^{\top},并令 \gamma 表示将累积衰减向量按行堆叠得到的矩阵。则

  • 求逆操作本身具有标准的三角矩阵向量—雅可比积

  • 在此基础上,通过常规的逐元素乘积和反向累积求和,即可得到关于\mathbf{B}、\mathbf{W}、\mathbf{K}、\mathbf{V}以及累积衰减的梯度。这种门控感知的梯度累积,是训练 Gated Delta Rule-2 所需的主要数学改动。其余反向传播内核保持与 KDA 相同的矩阵形状,并可复用相同的状态与输出向量—雅可比积结构