跳转至

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

会议: NeurIPS 2026
arXiv: 2605.22791
代码: https://github.com/NVlabs/GatedDeltaNet-2
领域: LLM效率
关键词: 线性注意力, Delta-rule, 快速权重记忆, 通道解耦门控, 分块并行算法

一句话总结

针对 Delta-rule 循环线性注意力中 Key 侧擦除与 Value 侧写入共用单一标量门控的结构约束,Gated DeltaNet-2 引入通道级擦除门 \(\bm{b}_t\) 与写入门 \(\mathbf{w}_t\) 彻底解耦二者,并推导了吸收通道衰减的非对称 WY 分块并行训练算法与门控感知反向传播,在多键值长文本检索和语言建模上全面超越 Mamba-2、KDA 与 Mamba-3。

研究背景与动机

标准 Transformer 依靠自注意力机制直接访问历史全部 token,在多项自然语言与推理任务中确立了统治地位,但其显存和计算代价随序列长度呈二次方剧烈增长,严重制约了超长上下文微调与高吞吐自回归生成。线性注意力与循环序列模型采用固定尺寸的矩阵状态对历史关联进行增量压缩,实现了序列长度维度的线性时间复杂度与常数级的解码显存占用。然而,由于压缩记忆的容量有限,当序列长度显著超出隐藏维度时,多个键值关联必须叠加在同一状态空间内,导致记忆混叠与灾难性干扰,精确检索历史信息的难度急剧增加。

为了增强固定容量状态的保留与更新能力,近期研究引入了多种自适应记忆控制机制。Mamba-2 采用数据依赖的标量衰减来调节记忆遗忘时限;DeltaNet 引入快速权重机制,在写入新值前先读取并扣除当前键对应的旧记忆,实现针对性覆盖;Gated DeltaNet 将标量全局衰减与 Delta-rule 局部覆盖相结合;Kimi Delta Attention(KDA)进一步在 Key 维度上引入通道级细粒度衰减。然而,在现有 Delta-rule 体系中,主动编辑操作依然绑定在一个全局标量门控 \(\beta_t\) 上:该标量同时决定在 Key 轴上抹除多少旧关联、以及在 Value 轴上吸纳多少新信息。

这种将擦除与写入绑定的标量设计构成了循环状态编辑的根本瓶颈。擦除操作属于 Key 空间,其目的是定位并清除特定子空间中的陈旧读出值;而写入操作属于 Value 空间,其目的是决定候选更新中哪些通道特征应当持久化固化。将这两个在物理意义和特征轴上完全正交的决策强制等同,极大削弱了模型在面对复杂语境时的编辑精度。核心 idea:将 Delta 规则中的记忆擦除与写入操作完全解耦,引入通道级擦除门 \(\bm{b}_t\) 作用于 Key 侧读出、通道级写入门 \(\mathbf{w}_t\) 作用于 Value 侧候选更新,并结合吸收累积衰减的紧凑分块 WY 算法,实现既具细粒度通道选择性又具硬件高效并行的高保真循环注意力。

方法详解

整体框架

Gated DeltaNet-2 的核心是解耦擦除与写入的 Gated Delta Rule-2 算子。在整体模型架构中,模型可作为纯循环网络堆叠,也可与滑动窗口注意力(Sliding-Window Attention, SWA)按交替结构组合构建混合模型(Hybrid):Gated DeltaNet-2 负责将长程上下文压缩为固定大小的矩阵状态,而局部滑动窗口注意力负责处理临近精确比对与短程检索。每个 token mixer 内部接收输入特征并生成查询、键、值以及三路互不干扰的自适应门控信号。

整个前向计算包含输入投影与门控生成、通道衰减与非对称编辑、分块 WY 矩阵转换以及输出门控映射四个阶段:

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列特征 x_t"] --> B["输入投影与短期因果卷积<br/>生成 q_t, k_t, v_t 与门控特征"]
    B --> C["通道级门控解耦投影<br/>生成通道衰减 α_t, 擦除门 b_t, 写入门 w_t"]
    C --> D["Gated Delta Rule-2 状态更新<br/>衰减旧状态并解耦执行 Key 侧擦除与 Value 侧写入"]
    D --> E["分块并行 WY 矩阵求解<br/>将通道衰减吸收入非对称擦除因子,解下三角线性系统"]
    E --> F["RMSNorm 与 SiLU 输出门控映射<br/>输出最终注意力上下文向量 o_t"]

关键设计

1. Gated Delta Rule-2:解耦擦除与写入的通道级在线更新

针对以往方法中单一标量 \(\beta_t\) 强行绑定擦除与写入的局限,Gated Delta Rule-2 分别在 Key 维度和 Value 维度引入独立的通道门控向量 \(\bm{b}_t \in [0, 1]^{d_k}\) 与 \(\mathbf{w}_t \in [0, 1]^{d_v}\),定义有效擦除方向 \(\bm{e}_t = \bm{b}_t \odot \bm{k}_t\) 和有效写入向量 \(\bm{z}_t = \mathbf{w}_t \odot \bm{v}_t\)。在每一步更新时,首先利用对角衰减矩阵 \(\mathbf{D}_t = \operatorname{Diag}(\bm{\alpha}_t)\) 对历史状态实施通道级遗忘,随后沿着 \(\bm{e}_t\) 读出陈旧关联 \(\bm{r}_t = \bar{\mathbf{S}}_t^\top \bm{e}_t\),并在 Key 方向写入差分修正:

\[\mathbf{S}_t = \bigl(\mathbf{I} - \bm{k}_t (\bm{b}_t \odot \bm{k}_t)^\top\bigr) \mathbf{D}_t \mathbf{S}_{t-1} + \bm{k}_t (\mathbf{w}_t \odot \bm{v}_t)^\top\]

从快速权重(Fast-Weight)记忆的在线局部回归视角来看,该更新精确等价于极小化一个带门控残差的目标函数:\(\bm{L}_t(\mathbf{S}) = \|\mathbf{S} - \bar{\mathbf{S}}_t\|_F^2 - 2 \langle \mathbf{S}^\top \bm{k}_t, \bm{z}_t - \bar{\mathbf{S}}_t^\top \bm{e}_t \rangle\)。通过令梯度为零直接得到该更新。当 \(\bm{b}_t = \beta_t \mathbf{1}\) 且 \(\mathbf{w}_t = \beta_t \mathbf{1}\) 时,该规则精确退化为 KDA;当进一步约束衰减为标量 \(\bm{\alpha}_t = \alpha_t \mathbf{1}\) 时,则完全退化为 Gated DeltaNet。这保证了模型既能涵盖先验标量更新子空间,又能学到不同通道特征上非均匀的擦除与保留行为。

2. 吸收累积衰减的分块并行 WY 算法

若直接采用逐 token 递归,在现代加速器(如 GPU Tensor Core)上的训练吞吐将严重受限。为了保持分块并行的硬件效率,本文推导了一种紧凑的 WY 变换形式。将序列切分为长度为 \(C=64\) 的数据块,定义块内累积衰减向量 \(\bm{\gamma}_r = \exp(\sum_{i=1}^r \bm{g}_i)\) 并定义归一化状态 \(\mathbf{S}_r = \operatorname{Diag}(\bm{\gamma}_r) \widehat{\mathbf{S}}_r\)。代入递推式后,通道衰减项神奇地从矩阵乘积转移到了 rank-one 增量因子的左右两侧:

\[\widehat{\mathbf{S}}_r = (\mathbf{I} - \bar{\bm{k}}_r \bar{\bm{e}}_r^\top) \widehat{\mathbf{S}}_{r-1} + \bar{\bm{k}}_r \bm{z}_r^\top, \quad \bar{\bm{k}}_r = \bm{\gamma}_r^{-1} \odot \bm{k}_r, \quad \bar{\bm{e}}_r = \bm{\gamma}_r \odot (\bm{b}_r \odot \bm{k}_r)\]

通过构造严格下三角矩阵 \(\mathbf{T} = \operatorname{tril}(\bar{\mathbf{E}} \bar{\mathbf{K}}^\top, -1)\) 并令 \(\mathbf{A} = (\mathbf{I} + \mathbf{T})^{-1}\),块内所有前向依赖被规约为前向代换求解,从而构建出擦除侧辅助矩阵 \(\mathbf{Y} = \mathbf{A} \bar{\mathbf{E}}\) 和写入侧辅助矩阵 \(\mathbf{U} = \mathbf{A} \mathbf{Z}\)。跨块状态更新与块内输出最终均可由密集的矩阵乘法直接完成:

\[\mathbf{S}_{[n+1]} = \operatorname{Diag}(\bm{\gamma}_C) \mathbf{S}_{[n]} + \mathbf{K}_{\text{tail}}^\top (\mathbf{U} - \mathbf{Y} \mathbf{S}_{[n]})\]
\[\mathbf{O}_{[n]} = \mathbf{Q}_\gamma \mathbf{S}_{[n]} + \mathbf{A}_{qk} (\mathbf{U} - \mathbf{Y} \mathbf{S}_{[n]})\]

该算法保持了与标量 Delta-rule 相同的 \(O(C)\) 矩阵代换与高吞吐 Tensor Core 计算模式,未因引入通道门控而破坏分块并行的理论复杂度。

3. 门控感知的反向传播推导与精准梯度累积

在传统标量门控的 Delta-rule 中,标量因子 \(\beta_r\) 在反向传播时可以被整体提取到点积之外,从而对中间梯度进行标量后缩放(scalar post-scaling)。然而,在 Gated Delta Rule-2 中,由于写入侧 \(\mathbf{Z} = \mathbf{W} \odot \mathbf{V}\) 和擦除侧 \(\bar{\mathbf{E}} = \bm{\gamma} \odot (\mathbf{B} \odot \mathbf{K})\) 分别作用了逐行逐通道不同的对角算子,标量提取数学上不再成立,必须在累积矩阵 \(\mathrm{d}\mathbf{A}\) 的点积位置显式嵌入门控权重:

\[\mathrm{d}\mathbf{A} \mathrel{+}= \mathrm{d}\mathbf{U} \, \mathbf{Z}^\top, \quad \mathrm{d}\mathbf{A} \mathrel{+}= \mathrm{d}\mathbf{Y} \, \bar{\mathbf{E}}^\top\]

结合下三角逆的伴随变换 \(\mathrm{d}\mathbf{T} = -\operatorname{tril}(\mathbf{A}^\top \mathrm{d}\mathbf{A} \mathbf{A}^\top, -1)\),反向传播核函数能够精确高效地将梯度回传至通道擦除门 \(\mathbf{B}\)、写入门 \(\mathbf{W}\) 及累积对数衰减 \(\bm{g}\),确保了通道解耦机制在大规模并行训练中的全精度梯度流动。

损失函数 / 训练策略

模型采用标准自回归语言建模交叉熵损失进行无监督预训练。实验采用 1.3B 参数量架构,在 FineWeb-Edu 语料上训练 100B tokens。优化器使用 AdamW,峰值学习率 \(4 \times 10^{-4}\),权重衰减 0.1,梯度裁剪阈值 1.0,并配合 1B tokens 的预热与余弦退火策略;全局批大小设置为 0.5M tokens,预训练序列长度为 4K tokens。在混合架构中,滑动窗口注意力窗口大小固定为 2K tokens。所有对数衰减 \(\bm{g}_t\) 的累积求和与激活在核函数外使用 fp32 保证数值稳定性,查询与键向量在输入前进行 \(L_2\) 归一化以抑制梯度发散。

实验关键数据

主实验

论文在语言建模困惑度、零样本常识推理评测以及长文本检索评测上与 Mamba-2、Gated DeltaNet、KDA 以及 Mamba-3(SISO 与 MIMO-rank 4)展开全面对比。

下表列出 1.3B 模型在语言建模(WikiText/LAMBADA 困惑度)与常识推理任务上的主要结果:

模型架构 模型类型 Wiki. ppl \(\downarrow\) LMB. ppl \(\downarrow\) PIQA acc \(\uparrow\) Hella. acc \(\uparrow\) Wino. acc \(\uparrow\) ARC-e acc \(\uparrow\) ARC-c acc \(\uparrow\) 平均准确率 \(\uparrow\)
Mamba-2 纯循环 16.79 12.38 72.58 55.51 55.33 70.68 35.26 51.82
Gated DeltaNet 纯循环 16.40 11.89 72.31 56.50 56.75 68.81 35.15 52.07
KDA 纯循环 16.81 11.68 72.09 55.75 55.72 70.83 35.92 52.28
Mamba-3 (SISO) 纯循环 16.30 12.99 72.31 55.58 56.20 70.45 34.56 51.42
Mamba-3 (MIMO) 纯循环 16.45 11.66 72.36 56.49 55.78 72.38 38.07 52.39
Gated DeltaNet-2 纯循环 15.90 11.41 72.80 56.84 57.85 72.43 38.23 53.11
Transformer 注意力 19.22 13.72 70.21 56.12 55.85 69.23 33.84 50.86
Mamba-2 混合 (SWA) 17.46 11.29 71.47 57.52 56.17 70.50 34.73 51.99
Gated DeltaNet 混合 (SWA) 16.00 10.82 70.06 57.50 56.83 70.41 35.15 52.25
KDA 混合 (SWA) 16.01 10.66 71.06 56.89 57.77 71.59 35.07 52.68
Mamba-3 (MIMO) 混合 (SWA) 15.81 10.92 71.98 58.19 57.06 70.54 38.48 52.72
Gated DeltaNet-2 混合 (SWA) 15.62 10.43 72.20 58.46 58.56 71.89 36.69 53.97

在长文本合成针锋检索(RULER 基准)中,Gated DeltaNet-2 在多键值干扰场景(MK-NIAH-1)与单键数值检索(S-NIAH-2)的长上下文区间表现尤为亮眼:

模型架构 循环/混合 S-NIAH-2 (4K) S-NIAH-2 (8K) S-NIAH-3 (2K) MK-NIAH-1 (1K) MK-NIAH-1 (2K) MK-NIAH-1 (4K)
Mamba-2 纯循环 62.6 21.0 38.6 29.0 21.2 21.4
Gated DeltaNet 纯循环 87.2 32.0 54.2 58.0 37.0 27.8
KDA 纯循环 89.0 30.6 63.2 54.0 44.2 28.0
Mamba-3 (MIMO) 纯循环 64.2 27.2 72.4 49.4 19.2 18.0
Gated DeltaNet-2 纯循环 93.0 39.2 89.8 72.6 51.4 37.8
Transformer 注意力 44.2 0.0 94.8 75.6 66.6 38.2
KDA 混合 (SWA) 56.0 23.0 93.4 91.4 84.0 40.4
Mamba-3 (MIMO) 混合 (SWA) 53.0 27.8 98.4 82.4 79.0 46.6
Gated DeltaNet-2 混合 (SWA) 57.9 29.2 99.0 93.0 84.6 48.0

消融实验

论文在纯循环设置下针对门控的通道结构与擦除范围进行了细致消融:

配置 Wiki. ppl \(\downarrow\) LMB. ppl \(\downarrow\) 常识推理平均 \(\uparrow\) S-NIAH-2@4K S-NIAH-3@2K MK-NIAH-1@4K 真实检索平均 \(\uparrow\)
Gated DeltaNet-2 (完整模型) 15.90 11.41 53.11 93.0 89.8 37.8 29.88
仅通道写入门 (标量 \(\bm{b}_t\), 通道 \(\mathbf{w}_t\)) 16.55 11.62 52.45 90.6 71.4 30.6 28.92
仅通道擦除门 (通道 \(\bm{b}_t\), 标量 \(\mathbf{w}_t\)) 16.12 11.50 52.79 92.1 84.6 35.2 29.51
扩展擦除范围 \(\bm{b}_t \in [0, 2]^{d_k}\) 15.95 11.44 53.04 93.1 89.4 37.6 29.81

关键发现

  • 通道擦除门是性能增益的主力:仅保留通道擦除门 \(\bm{b}_t\) 能恢复全模型在语言建模与长文本检索上的绝大部分提升(MK-NIAH-1@4K 达到 35.2 对比标量版的 30.6);而仅保留通道写入门 \(\mathbf{w}_t\) 带来的增益相对有限,这表明在固定容量矩阵记忆中,精准定位并从 Key 轴剔除干扰关联是缓解记忆冲突的核心瓶颈。
  • 长上下文多键检索能力显著:在 RULER MK-NIAH-1 任务中,当上下文扩展至 4K 时,Gated DeltaNet-2 纯循环模型的准确率(37.8%)几乎达到 Mamba-3 MIMO(18.0%)的两倍,大幅超越 KDA(28.0%),展现了非对称矩阵外积更新对高熵多重键值记忆的强大维持能力。
  • 训练吞吐几乎无损:在单张 H100 GPU 上的吞吐实测表明,Gated DeltaNet-2 在长序列下吞吐仅从 38.0 Kt/s 微幅下降至 36.1 Kt/s,保持了极平坦的线性扩展优势,在增加通道自由度的同时仅带来极小的计算常数开销。

亮点与洞察

  • 解耦擦除与写入的正交决策:以往 Delta 规则视擦除与写入为硬币的两面,强行绑定幅度;本文指明擦除属于 Key 空间,而写入属于 Value 空间,通道级解耦在不扩充循环状态尺度的前提下,大幅拓宽了记忆编辑的有效自由度。
  • 吸收通道衰减的代数技巧:通过将非均匀对角累积衰减 \(\bm{\gamma}_r\) 分解吸收进 rank-one 擦除因子 \(\bar{\bm{k}}_r\) 与 \(\bar{\bm{e}}_r\),成功维持了标准三角矩阵求解 \(\mathbf{A} = (\mathbf{I} + \mathbf{T})^{-1}\) 的简洁形式,化解了通道级衰减与分块并行在工程实现上的冲突。
  • 快速权重在线优化的自洽视角:从快速权重局部在线回归出发,严谨推导了 Gated Delta Rule-2 的闭式解,统一了 Mamba-2、DeltaNet、Gated DeltaNet、KDA 和 Mamba-3 的矩阵记忆更新范式。

局限与展望

  • 硬件显存带宽与局部累加开销:由于通道擦除与写入门破坏了标量因子的点积外提性质,反向传播核函数需要驻留更多的梯度累加器,在 Hopper 架构上对 warp 数量存在额外调度约束(需限制在 2 或 4 warps)。
  • 扩展擦除范围收益未饱和:将擦除门扩展至负特征值区间(\([0, 2]^{d_k}\))在当前 1.3B 规模下并未带来统计显著的进一步提升,如何在更大模型规模与更长语料下发挥谱半径扩展的潜力仍待探索。
  • 纯循环架构对局部聚合任务的固有限制:在 NQ 和 DROP 等强依赖密集证据局部抽取的任务上,纯循环版本仍落后于全量注意力,必须结合 SWA 混合模块才能达到最佳综合性能。

相关工作与启发

  • vs KDA (Kimi Delta Attention): KDA 首次引入了 Key 轴通道级衰减,但仍使用标量 \(\beta_t\) 统管擦除与写入;本文在继承通道衰减的基础上,彻底将擦除门与写入门拓展为独立的通道向量,在多键干扰场景下检索精度大幅飞跃。
  • vs Mamba-3: Mamba-3 通过指数梯形离散化与旋转状态空间实现复杂的相干累加,但不包含差分残差擦除机制;Gated DeltaNet-2 证明基于 Delta 规则的主动擦除-重写机制在长程记忆解耦上具有更具针对性的优势。
  • vs Gated DeltaNet: Gated DeltaNet 在标量层级组合衰减与覆盖;本文是其在高维通道自由度上的全面推广,并提供了完整的 Triton 分块并行内核与混合架构工程实践。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次提出通道级解耦擦除与写入的 Gated Delta 机制,并给出了严格的快速权重数学推导。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 1.3B/100B tokens 标准预训练、常识推理、合成长程检索、真实 RAG 评测及完整的门控消融与 Triton 吞吐验证。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完备,从公式推导、分块算法到反向求导与硬件实现细节详实透彻。
  • 价值: ⭐⭐⭐⭐⭐ 为次二次方线性注意力与循环模型克服长程记忆干扰提供了坚实、高效的架构范本。