智客 ZICQ
EN 登录 / 注册
智客前沿 前沿论文 #线性注意力 #长上下文 #状态门控 #Muir #GDN-2

Muir重构线性注意力门控:1.34B模型外推1000万Token

Mr.Xu 的头像

文 / Mr.Xu 智客前沿编译 · 审校

发布时间:

中文阅读 (Chinese) English Version

摘要:Muir 是一种面向长上下文线性注意力模型的新型状态更新架构。其核心洞察是,键激活与单位归一化会把能量集中到少量通道,使擦除门无需全秩参数;写入操作则会影响全部通道,因此 Muir 将多类门控融合为敏感度分配的低秩投影,并以匹配滤波器读出门重新加权不同衰减速度的状态行,抑制慢衰减行累积的干扰。论文还提出并行状态更新算法,缓解宽矩阵单元在大分块下传统分块算法的数值不稳定。实验覆盖166M至1.34B参数:Muir 相比 GDN-2 降低1.8%至2.0%的困惑度,在32K上下文下训练速度为 Transformer 的2.7倍,并可从4K训练长度外推至1000万Token。


从门控冗余到状态读出失衡

Muir 没有继续把线性注意力的状态更新当作一组同权重的门控函数,而是先分析键、值与状态行对更新的敏感度差异。对于门控 Delta Rule 一类模型,状态更新可以抽象为:先衰减旧状态,再根据当前键擦除与旧状态相关的预测,最后写入新值。

GDN-2 将擦除门与写入门解耦,并允许逐通道衰减。Muir 的关键观察是,键向量经过激活和单位归一化后,能量通常集中于少数通道,因此擦除门真正需要调节的有效维度远低于写入门;值写入则会直接缩放全部通道,对状态的影响更加广泛。据此,把所有门都保留为全秩投影会消耗参数,却没有等比例增加状态容量。

另一个问题出现在读出端。慢衰减状态行会持续累积较大的范数,在未加权读出中压过快衰减状态行。模型因此可能被迫用更强的写时擦除抑制干扰,削弱有限状态中的信息保留能力。

敏感度分配的低秩门控

Muir 将衰减、擦除和写入相关门控融合到低秩投影中,由投影结构依据通道敏感度分配表达能力:需要覆盖广泛通道的写入路径保留更高容量,而仅作用于稀疏有效子空间的擦除路径采用更紧凑的参数化。

从计算角度看,若门控秩从完整的 $d$ 降至 $r$,相关投影的计算量可由约 $O(d^2)$ 降至 $O(dr)$,参数与中间激活也会同步减少。省下的参数没有被全部用于增加状态宽度,而是投入更宽的前馈网络。这一取舍的关键不是“低秩一定更好”,而是:只有当矩阵运算足够宽时,较低秩的更新单元才可能成为算力瓶颈;Muir 实际上是在重新分配矩阵单元内部的参数预算。

不过,低秩门控也会形成新的表达瓶颈。如果键能量并非稳定稀疏,或者写入任务需要同时控制大量状态方向,低秩门可能难以表达精确的擦除与写入组合。论文当前证据集中在其线性注意力模型族和既定训练分布上,迁移到其他状态结构仍需验证。

匹配滤波器读出如何处理状态行失衡

Muir 增加匹配滤波器读出门,根据状态与当前查询的匹配关系,对不同状态行重新加权,而不是直接汇总全部状态。其作用类似从含噪状态中选择与查询最相关的分量,抑制慢衰减大范数行对结果的支配。

该机制把“状态写入是否正确”和“状态如何被读取”拆成两个问题:写入端负责控制干扰累积,读出端则主动执行匹配选择。它对多跳和多针检索尤其有意义,因为这类任务需要从长期状态中定位少量相关证据;若高范数噪声行在读出时占主导,底层存储即使完整,答案仍可能被污染。

代价是读出门需要额外的投影与状态交互开销。论文报告 Muir 在训练吞吐上匹配 GDN-2,并更快达到相同的100B-token验证损失,但该结论依赖具体硬件、矩阵单元宽度和内核实现,不能直接外推到所有线性注意力框架。

宽矩阵单元引发数值不稳定

传统线性注意力通常使用 chunkwise 算法,把序列切成若干块,在块内并行计算、块间递推状态。块越大,矩阵单元越容易被充分利用,但累计误差与不稳定的归一化也会放大。Muir 观察到,标准分块算法在其宽矩阵单元所需的大 chunk size 下会出现数值不稳定。

为此,论文设计了并行状态更新机制,以避免依赖易失稳的大块递推路径。其目标不是降低理论复杂度,而是在现代矩阵加速器上同时获得稳定训练与高吞吐。这类优化高度依赖内核实现:结果能否迁移到 AMD GPU、国产加速器或不同 Triton/CUDA 版本,需要检查归约顺序、精度策略和状态合并规则。

实验结果与适用边界

在166M至1.34B参数规模上,Muir 相比 GDN-2 将困惑度降低1.8%至2.0%,相比 Transformer 降低2.5%至3.6%;常识准确率最高提高1.9个百分点,抽取式召回也有所改善。1.34B 模型在多跳和多针召回上均超过两条基线。

长上下文结果更突出:模型从4K训练长度外推至1000万Token,困惑度没有明显漂移;在静态读出条件下,Passkey 检索直到100万Token仍接近100%成功率。训练侧,Muir 达到与 GDN-2 接近的吞吐,收敛到100B-token验证损失的速度快1.2倍,在32K上下文下训练速度为 Transformer 的2.7倍。

这些结果说明线性注意力不一定只能以较弱检索能力换取线性复杂度和常数内存状态。不过,评测仍有三项限制:第一,摘要未提供硬件型号与绝对吞吐,单看加速比不足以判断部署收益;第二,Passkey 和多针任务主要检验精确检索,不能完全代表长文档问答中的语义聚合与推理;第三,10百万Token外推结果来自特定模型与训练配方,尚不能证明任意更长上下文、噪声比例和分布漂移下都保持稳定。

开发者落地建议

对训练团队,Muir 最值得复现的部分不是单独的“匹配滤波”模块,而是三个设计决策的组合:按敏感度压缩擦除门、用查询相关读出校正状态行失衡,以及为宽矩阵单元重写稳定更新内核。

工程上建议优先验证三项指标:端到端训练 token/s 与显存占用、训练 loss 与验证 loss 的漂移差异,以及长上下文检索准确率随上下文长度变化的曲线。还应加入不同衰减率、不同 key 稀疏度和不同状态维度的消融,以确认收益确实来自敏感度分配,而不是更大的前馈网络带来的额外容量。

如果生产目标主要是在有限显存下处理超长上下文,Muir 值得进入候选架构;如果任务更偏短上下文稠密推理,其宽矩阵单元、稳定内核和训练配方是否抵消部署复杂度,需要按实际硬件测算后再决定。


消息来源:Hugging Face Daily Papers (2026-11-05)

—— 完 ——

主题标签: #线性注意力 #长上下文 #状态门控 #Muir #GDN-2

编辑部与事实核查说明:本篇专刊由智客前沿资讯管线根据官方技术公告、学术论文及开源工程文档整理编译,经算法实体核验与人工编辑审校后发布。若发现技术事实纰漏,欢迎依据勘误方针或一键向编辑部发送勘误反馈。

社区整体评论区

正在加载实时智能评论与划词标注…