Gated DeltaNet:给线性注意力装上一块会忘、会改的记忆
Gated DeltaNet: 给线性注意力装上一块会忘、会改的记忆
Gated DeltaNet:给线性注意力装上一块会忘、会改的记忆
论文原文:Gated Delta Networks: Improving Mamba2 with Delta Rule
arXiv:https://arxiv.org/abs/2412.06464
发布日期:2024-12-09
作者:Songlin Yang、Jan Kautz、Ali Hatamizadeh(MIT CSAIL、NVIDIA)
你给模型一篇两万字的会议记录,里面夹着一句“服务器迁移口令是 7391”。读到末尾时,你问它口令是什么。标准 Transformer 可以回头查看前文,代价是上下文越长,计算与显存压力越大。线性模型走了另一条路,它不保留一整摞原始纸张,而是边读边把内容压进一块固定大小的记忆。
Transformer:用注意力机制直接比较当前 token 与上下文中的其他 token,精确检索能力强,但标准自注意力的计算量随序列长度近似二次增长。
线性注意力(Linear Attention):把注意力改写成可递推的状态更新,使推理时间随序列长度线性增长,并把历史压缩进固定大小的状态。
问题随之出现:那块记忆既要长期保存口令,又要在话题转换时快速清掉旧信息,还得允许后来一句话改正前面的记录。只会记,它迟早被塞满;只会忘,它会把仍然有用的东西一起擦掉;只会局部修改,它又来不及应对大规模的上下文切换。
2024 年 12 月,Songlin Yang、Jan Kautz 和 Ali Hatamizadeh 把两种已经存在、却长期分开使用的机制放进同一条更新公式:Mamba2 式的门控遗忘负责“大扫除”,DeltaNet 式的 delta rule 负责“定点改写”。这个组合被命名为 Gated DeltaNet,并在 ICLR 2025 发表。
这篇论文真正解决的,不是把某个榜单再推高一点,而是让固定容量的神经记忆第一次同时拥有快速清空与精确改写两种动作。
一块固定大小的记忆,为什么会在长文本里失控
要理解 Gated DeltaNet,先得把线性注意力从“便宜版注意力”的印象里解放出来。它更接近一种可学习的联想记忆:每读入一个 token,模型产生 key、value 和 query。key 像索引,value 像要保存的内容,query 则在需要输出时发起查询。历史被累积进矩阵状态 $S_t$,输出由 $S_t q_t$ 计算。
联想记忆(Associative Memory):把 key 与 value 建立绑定,之后用相似 key 查询对应 value 的记忆结构。
状态矩阵(State Matrix):线性注意力在递推过程中维护的固定大小矩阵,可视为被压缩后的上下文记忆。
最朴素的更新是:
$$ S_t = S_{t-1} + v_t k_tᵀ $$
它看起来很干净。每个新键值对都被加进状态,旧信息也没有被主动删除。麻烦藏在“固定大小”四个字里。状态矩阵能容纳的相互正交键值对数量受维度约束,当序列持续增长,不同 key 会越来越难保持独立,多个记忆开始叠在相同方向上。查询一个 key 时,别的 value 也会漏进结果。
记忆碰撞(Memory Collision):固定容量状态中,不同键值关联相互干扰,导致查询时无法准确分离目标信息。
这和把整家公司一年的票据都塞进一个只有二十格的收纳盒很像。起初每个格子对应一个项目,找票据很快。项目增加到第五十个后,你只能把多个项目叠放在同一格。如果不清理,也不重排,某次找“项目 A 的发票”时,很可能把项目 B 的付款单一起抽出来。收纳盒没有变小,问题是可区分的位置已经耗尽。
Mamba2 的办法是给整个状态乘一个数据依赖的衰减系数 $α_t$:
$$ S_t = α_t S_{t-1} + v_t k_tᵀ, α_t ∈ (0,1) $$
门控衰减(Gated Decay):由当前输入决定遗忘比例,把历史状态整体缩小,从而释放记忆容量。
当 $α_t$ 接近 0,过去几乎被清空;接近 1,过去大体保留。这个机制能在话题改变时迅速腾出空间,却有一个明显代价:它对所有旧记忆一视同仁。模型如果只想忘掉旧客户的邮寄地址,也会同时削弱仍需保留的合同金额、付款日期和联系人。
生活里的对应场景是一间安装了总调光器的仓库。你发现角落里有一箱过期文件,于是把整间仓库的灯调暗。过期文件确实没那么显眼了,可正在核对的合同也一起看不清。如果没有总调光器,仓库会一直亮着,旧信息持续占据注意;只有总调光器,又无法只处理那一箱文件。收益是快速降低全局干扰,代价是有用内容也同步衰减。
DeltaNet 采取另一种动作。它先用当前 key 从旧状态读出 $S_{t-1}k_t$,再计算新 value 与旧 value 的误差,只沿这个 key 对应的方向做替换:
$$ S_t = S_{t-1}(I-β_t k_tk_tᵀ) + β_t v_tk_tᵀ $$
Delta Rule:根据目标 value 与当前预测 value 的差异更新记忆,只修改与当前 key 相关的部分。
Householder 型转移矩阵:形如 $I-βkkᵀ$ 的低秩状态转移,它能沿特定方向修改状态,同时尽量保留其他方向。
这像酒店前台修改一位客人的房间号。前台先按姓名找到旧记录,再把“1208”改成“1516”,其他客人的记录不动。如果没有这一步定点替换,只能不断追加新记录,系统会同时保留两个房间号;如果每次修改都清空整张住客表,更新虽然彻底,却把其他人的信息也毁了。
Delta rule 在合成检索任务上很强,因为它擅长精确绑定和修订。但每一步通常只处理一个 key 方向。当上下文从法律合同突然切换到烹饪教程,旧主题留下的大量关联无法在一两步内清走。它像有一支很精准的橡皮擦,却一次只能擦掉一个字。
论文用 S-NIAH 三组任务把这两个缺陷拆开测量。
S-NIAH(Single Needle-In-A-Haystack):在长上下文中埋入一个键值对,再测试模型能否在末尾根据 key 找回 value 的合成检索任务。
S-NIAH-1 的背景高度重复,要保存的信息很少,主要考验长期保留。8K 长度时,DeltaNet 仍有 98.8% 准确率,Mamba2 只剩 30.4%,说明持续衰减会伤害远距离记忆。S-NIAH-2 和 S-NIAH-3 混入真实文章背景,状态需要过滤大量干扰。此时没有清理能力的 DeltaNet 在 8K 的数字检索上跌到 14.4%,Gated DeltaNet 达到 29.6%;在 4K 的 UUID 检索上,DeltaNet 为 22.4%,Gated DeltaNet 为 27.6%。
| 模型 | S-NIAH-1 8K | S-NIAH-2 4K | S-NIAH-2 8K | S-NIAH-3 4K |
|---|---|---|---|---|
| DeltaNet | 98.8 | 18.6 | 14.4 | 22.4 |
| Mamba2 | 30.4 | 56.2 | 17.0 | 4.6 |
| Gated DeltaNet | 91.8 | 92.2 | 29.6 | 27.6 |
论文的长度外推图提供了另一种观察角度。模型只在 4K 序列上训练,却被放到最长 20K 的 GovReport、QMSum、NarrativeQA、Qasper、CodeParrot 和 PG19 上测试。图中的纵轴是困惑度,越低越好。纯递归模型里,Gated DeltaNet 的整体曲线更稳,但没有在每个数据集、每个长度上都赢。作者明确称结果是 mixed,这比“长上下文已经解决”的宣传更接近事实。
图里最能说明问题的地方是 8K 之后的走势。Mamba2 在 NarrativeQA 等任务上开始明显恶化,DeltaNet 在部分任务上也因干扰积累而上扬;Gated DeltaNet 及其混合版本多数保持在更低区间。它没有让固定状态变成无限记忆,只是在有限容量内做了更聪明的资源调度。
还可以把这组曲线理解成一次容量压力测试。4K 以内,各种模型都能依靠训练分布内的模式维持表现;越过训练长度后,状态更新规则的偏好才逐渐显形。持续衰减的模型会遗失早期证据,缺少清理的模型会积累干扰,带双重控制的模型则有机会根据内容选择保留、清空或覆盖。曲线没有直接显示每次门控发生在哪里,却把不同记忆策略长期累积后的结果呈现出来。
根本矛盾是:有限状态既要稳定保存少量高价值关联,又要及时释放被低价值信息占据的方向。只优化“记住”或只优化“忘掉”,都会在另一类长文本任务上付出代价。
把全局擦除和定点改写放进同一条公式
Gated DeltaNet 的核心公式只有一行:
$$ S_t = S_{t-1}(α_t(I-β_t k_tk_tᵀ)) + β_t v_tk_tᵀ $$
先看 $α_t$。它控制整个历史状态还剩多少,提供快速遗忘。再看 $β_t$。它控制当前 key 方向上擦除旧 value、写入新 value 的强度。两个动作相乘后,模型可以在同一步里先缩小旧状态,再精确修改一个关联。
Gated Delta Rule:把全局衰减门 $α_t$ 与定向 delta 更新 $β_t$ 合并的状态更新规则。
这个式子的两个极端很有解释力。若 $α_t$ 接近 0,历史状态被快速清空,适合章节切换、对话重置或任务边界。若 $α_t$ 接近 1,它退化为接近纯 DeltaNet 的精确更新,适合长期保留大部分内容,只修正某个键值关联。模型学习的是两者之间连续的动作空间。
可以把它看成经营一家生鲜店。每天闭店时,店长先决定全店要淘汰多少旧库存,这对应 $α_t$;随后针对“牛奶”这个品类核对货架数量,撤掉过期批次并补入新货,这对应 $β_t$ 与 key 定向更新。如果没有全店淘汰,换季后会留下大量无人购买的库存;如果只有全店统一打折清仓,畅销商品也会被误伤;如果没有品类级核对,同一款牛奶的新旧批次会混在一起。组合动作让店长既能应对换季,也能修正单品记录。
论文还给了一个更深的解释:状态 $S_t$ 可以视为一组在推理过程中变化的权重。当前 key 是输入,value 是监督目标,模型每读一个 token,就对“用 $S_tk_t$ 预测 $v_t$”这个小型回归问题做一步在线梯度下降。Delta rule 对应误差驱动的更新,$β_t$ 像自适应学习率;门控衰减 $α_t$ 则像对这组快速权重施加数据依赖的权重衰减。
Fast Weight(快速权重):在处理当前序列时持续更新、用于临时存储信息的权重或状态,与训练完成后固定的慢速参数相对。
在线学习(Online Learning):数据逐步到来时立即更新模型状态,不等待收集完整数据集后再统一训练。
这层解释很有价值,因为它把“记忆模块”从一个黑箱变成了微型学习器。标准前向传播通常像拿着固定工具处理输入,Gated DeltaNet 则在读上下文时临时改造一件工具。每个 token 都既是被处理的数据,也是一次小规模的训练样本。
换成日常场景,这像同声传译员边听会议边维护一张术语表。主办方说“Project Lighthouse 指新版,不再指旧版”,译员立即把该词条的旧释义替换掉;会议从财务议题切到医学议题时,她会整体降低上一套术语的优先级。如果不做在线修订,同一个项目名的新旧含义会冲突;如果不降低旧领域词汇的权重,医学段落里可能不断冒出财务解释。临时术语表的价值正在于会随现场变化。
公式简洁不等于模型块简单。论文中的 query、key、value 都经过线性投影、短卷积和 SiLU 激活,query 与 key 还做 L2 归一化。$α$ 与 $β$ 由独立线性投影生成,输出经过归一化、门控和最终投影。消融实验表明,去掉短卷积后平均困惑度从 27.35 变为 28.95,平均常识准确率从 47.26% 降到 46.16%;去掉输出门后,准确率降到 45.46%。L2 归一化也明显优于 L1 归一化。
短卷积(Short Convolution):在很小的局部窗口内混合相邻 token,为递归记忆补充局部顺序模式。
L2 归一化:把向量缩放到单位欧氏长度,限制 key 与 query 的尺度,减少状态更新中的数值不稳定。
图 1 左侧是两种混合架构。H1 交替放置 Gated DeltaNet 与滑动窗口注意力;H2 再加入 Mamba2。右侧展示单个 Gated DeltaNet 块,几条输入路径最终汇入 gated delta rule,然后经过规范化与输出门。
这张图揭示了论文很务实的一面。作者没有声称一种递归状态可以取代全部注意力。线性递归擅长持续更新状态,滑动窗口注意力擅长比较附近 token,标准注意力擅长从较大范围精确取回原文。混合模型把不同记忆机制安排在不同层,让它们各自承担更合适的工作。
滑动窗口注意力(SWA):每个 token 只关注附近固定窗口内的 token,把全局二次计算缩小为局部计算。
这像医院分诊。导诊台持续维护“当前有哪些病人、病情是否变化”的全局状态,专科医生精查眼前一位病人的近期检查,档案室在需要时取回原始报告。如果把所有工作都交给导诊台,它的白板会装不下细节;如果每次都翻全部档案,速度会慢得无法接诊;如果只看最近一张报告,又可能漏掉长期病史。混合架构的意义不是折中得模糊,而是把不同计算原语分工。
这条公式解决的是记忆动作的表达力问题,混合架构解决的是固定状态无法承担所有检索形态的问题。前者让一块记忆会管理自己,后者承认一块记忆仍有边界。
公式更强以后,为什么训练速度没有被吞掉
递归公式的一个老问题是训练时难以并行。推理可以从 $S_{t-1}$ 算到 $S_t$,但若训练也严格逐 token 运行,GPU 大量计算单元会闲着。Transformer 之所以强势,一部分原因就是矩阵乘法能把许多 token 同时铺到 GPU 上。
Gated DeltaNet 的算法贡献,是把带门控的 delta recurrence 改写成分块并行形式。序列先被切成多个 chunk,块与块之间传递最终状态,块内部的大部分运算被组织成矩阵乘法。作者在 DeltaNet 已有的 WY representation 基础上,把累计衰减 $γ$ 吸收到 query、key、状态与低秩更新中,得到可在现代 tensor core 上高效运行的算法。
Chunkwise Parallelism(分块并行):把长序列切成若干块,块间保留递归依赖,块内改写为可并行矩阵运算。
WY Representation:把一串 Householder 型矩阵乘积压缩成适合矩阵乘法的形式,避免逐个显式相乘。
Tensor Core:GPU 中为低精度矩阵乘加优化的专用计算单元,是大模型训练高吞吐的主要硬件基础之一。
类比一下:一家面包店原本要求一个师傅从称粉、和面、发酵到烘烤完整做完一条面包,再开始下一条。递归关系很清楚,产能却低。分块并行把工作改成烤盘批次,烤盘之间传递上一批留下的酵种状态,同一烤盘里的称量和整形可以多人并行。如果完全取消批次间状态,发酵风味会断掉;如果坚持一条一条做,设备利用率又太差。算法要同时保留依赖和批处理。
图 3 用单张 H100、1.3B 参数模型比较训练吞吐。Gated DeltaNet 大致维持在每秒 45K tokens 左右,与 DeltaNet 接近,比 Mamba2 慢约 2K 到 3K tokens/s。混合版 H1 在不同序列长度下约 52K 到 55K tokens/s,原因是它加入了高度优化的 2K 滑动窗口注意力,并通过层间组合获得更高硬件利用率。标准 Transformer++ 在 2K 时约 55K tokens/s,序列增长到 16K 后跌至约 26K;H1 的曲线则基本保持平坦。
这组结果必须谨慎读。论文比较的是训练吞吐,不能直接推导端到端生成速度;图中 batch size 会随序列长度反向调整,纵轴测的是每秒处理 token 数;不同实现的 kernel 成熟度也会影响数字。论文能支持的结论是:加入 gated delta rule 没有让 DeltaNet 的并行训练优势消失,混合架构在该实验设置下保持了强吞吐。
性能部分采用了相当严格的同条件比较。主要 1.3B 模型都在 FineWeb-Edu 的 100B tokens 上训练,训练长度 4K,使用相同 Llama2 tokenizer、AdamW 配置和批量规模。混合模型的滑动窗口为 2K。这个设计减少了“某模型只是吃了更多数据或用了更大词表”的干扰。
困惑度(Perplexity):语言模型对测试文本的不确定程度,数值越低通常表示预测越好。
Zero-shot:不为目标任务额外训练,直接用预训练模型完成评测。
在语言建模与常识推理上,纯递归模型中的 Gated DeltaNet 平均准确率 55.32%,高于 Mamba2 的 54.89% 和 DeltaNet 的 52.14%;混合版 H1 达到 56.40%。在六个真实检索任务上,纯 Gated DeltaNet 平均 30.6,略高于 Mamba2 的 29.8 和 DeltaNet 的 26.2;H2 达到 40.1,高于 Transformer++ 的 37.0 与 Samba 的 37.3。LongBench 的 14 项平均分里,Gated DeltaNet 为 16.6,H1 为 17.8,H2 为 18.4。
| 模型 | 常识平均准确率 | 真实检索平均分 | LongBench 平均分 |
|---|---|---|---|
| Mamba2 | 54.89 | 29.8 | 13.5 |
| DeltaNet | 52.14 | 26.2 | 13.6 |
| Gated DeltaNet | 55.32 | 30.6 | 16.6 |
| Transformer++ | 52.25 | 37.0 | 11.0 |
| Samba | 54.00 | 37.3 | 15.9 |
| Gated DeltaNet-H1 | 56.40 | 39.0 | 17.8 |
| Gated DeltaNet-H2 | 56.18 | 40.1 | 18.4 |
这些数字还揭示一个不太舒服的事实:纯递归模型依然没有补平真实检索差距。Gated DeltaNet 的 30.6 高于其他纯递归模型,却仍低于 Transformer++ 的 37.0。真正超过 Transformer 的是混合版。固定状态能提升容量利用率,但“把全文压进一个有限矩阵”和“保留可直接寻址的原文”仍是两种能力。
再看 S-NIAH,也不能只挑最好看的列。Gated DeltaNet 在 S-NIAH-1 的 1K、2K 上分别为 98.4% 和 88.4%,低于 Mamba2 的 99.2% 和 98.8%;到了 4K、8K,它才明显优于 Mamba2。它在 S-NIAH-2 的 8K 也只有 29.6%,离可靠检索还很远。论文证明了组合机制更好,没有证明固定状态的容量限制已经消失。
实际工程里还要留意代码发布边界。官方仓库提供训练与 kernel 实现,但没有发布论文中的预训练权重;仓库说明训练代码的数据加载器来自 SlimPajama,而论文主实验使用 FineWeb-Edu。作者称两者趋势接近,可复现实验仍需在数据、评测 prompt 和权重转换上做额外工作。真实检索任务对 prompt 很敏感,官方仓库明确提醒不要只用通用 lm-eval-harness 任务名代替论文评测脚本。
Kernel:直接实现核心张量运算的底层程序,算法公式相同,kernel 设计不同也可能带来显著速度差异。
这里的日常类比是赛车测试。论文证明新发动机在同一赛道、同一轮胎和同等油量下更快,同时没有明显增加维修时间。但这不代表它在城市拥堵、雨天和长途耐久赛里自动胜出。如果不控制赛道和油量,成绩没有可比性;如果只测实验赛道,工程结论又不能无限外推。严谨的读法,是接受它在受控条件下的优势,同时保留部署层面的未知数。
硬件高效并不是公式的附属包装。一个记忆规则如果只能顺序计算,就很难进入大模型训练;一个漂亮 kernel 如果对应的记忆动作太弱,也只是更快地遗忘或碰撞。论文的价值来自算法表达力与并行实现同时成立。
资料与边界
正文实验数字来自 ICLR 2025 论文,实现与复现边界参考 NVIDIA 官方仓库。文章把训练吞吐、语言建模、合成检索和真实检索分开解释,未把任何一项替代为“通用长上下文能力”。
写在后面
这篇论文发表于 2024 年,距今约一年半,它的影响已经越过实验室。2025 年的 Kimi Linear 用更细粒度门控扩展 Gated DeltaNet,在 1M context 下报告最高 75% KV cache 节省与 6 倍解码吞吐;2026 年 Olmo Hybrid 用 Gated DeltaNet 替换 75% 注意力层,在 MMLU 上以少 49% 的训练 token 达到同等准确率;同年的 Gated DeltaNet-2 又把“擦除”和“写入”拆成通道级独立门。真正延续下来的主线,是把有限记忆的管理动作越做越细。
如果未来模型只能带着一只固定大小的箱子走过百万 token,那么决定它能否保持清醒的,也许不是箱子还能扩大多少,而是它能否知道哪一张纸该扔、哪一行该改、哪一份原件必须留到最后。