Muir:采用敏感度分配门控与匹配滤波读出的线性注意力
Muir: Linear Attention with Sensitivity-Allocated Gating and Matched-Filter Readout
精选理由
通过按敏感度分配门控并优化读出,Muir提升线性注意力的语言建模与超长上下文检索能力,同时保持高效训练。
AI 摘要
arXiv:2611.08376v1 公告类型:新论文 摘要:线性注意力将上下文压缩为固定大小的键值状态,从而实现线性时间训练和恒定内存推理。门控 delta 规则模型会衰减该状态,并在写入新值之前擦除其对传入键的预测;Gated DeltaNet-2(GDN-2)通过逐通道衰减以及解耦的满秩擦除门和写入门,将这一机制加以推广。然而,我们发现,状态更新对擦除的敏感度远低于对写入的敏感度:键激活与单位归一化将键的能量集中到稀疏的通道子集上,而值写入会缩放所有通道,因此擦除门可以采用更低秩的参数化方式。此外,衰减缓慢的状态行会累积远大于衰减较快行的范数,从而主导未加权的读出,并迫使模型在写入时执行更强的擦除以抑制干扰。基于这一分析,我们提出 Muir:将所有门融合为按敏感度分配的低秩投影,利用匹配滤波读出门对状态行重新加权,并将节省的参数投入更宽的前馈网络。由于标准分块算法在采用足以充分利用宽矩阵运算单元的大分块尺寸时会出现数值不稳定,我们设计了一种并行更新方法,以实现稳定、高效的训练。在 1.66 亿至 13.4 亿参数规模下,Muir 相较 GDN-2 将困惑度降低了 1.8–2.0%,相较 Transformer 降低了 2.5–3.6%,将常识任务准确率最多提高了 1.9 个百分点,并在抽取式召回上优于 GDN-2。在 13.4 亿参数规模下,Muir 在多跳召回和多针召回上均优于这两个基线,将其 4K 训练长度外推 2,560 倍至 1,000 万 token 的序列流而不发生困惑度漂移,并在使用静态读出的情况下,在长达 100 万 token 的序列中保持接近 100% 的口令检索准确率。Muir 在取得这些提升的同时,保持了与 GDN-2 相当的吞吐量,达到 GDN-2 训练 1,000 亿 token 时的验证损失所需速度为其 1.2 倍,并在 32K 序列长度下达到 Transformer 的 2.7 倍训练速度。