Muir：结合敏感性分配门控与匹配滤波读出的线性注意力 Muir: Linear Attention with Sensitivity-Allocated Gating and Matched-Filter Readout
低秩敏感性门控+匹配滤波读出，让线性注意力全面超越GDN-2与Transformer
前置知识
线性注意力与固定维度联想记忆状态
线性注意力把已见 token 的历史压缩进一个固定大小的矩阵状态 $S_t\in\mathbb{R}^{d_k\times d_v}$（每头通常 $d_k=d_v=128$），以 $S_t=S_{t-1}+k_t^\top v_t$ 的外积累加更新，训练 $O(T)$、推理缓存 $O(1)$，从而绕开 softmax 注意力的 $O(T^2)$ 训练开销与 $O(T)$ 的 KV 缓存；查询时用 $o_t=q_tS_t$ 从状态读出。
Muir 的全部设计（衰减、擦除、写入、读出门控）都作用于这个状态矩阵，理解状态如何被写入和读出是读懂本文的前提。
Delta 规则与 Gated DeltaNet-2（GDN-2）
Delta 规则把状态更新看作对键值重构损失 $\frac{1}{2}\|k_tS-v_t\|_2^2$ 的一步在线梯度下降：先用当前键查询状态并擦除其预测，再写入新值，即 $S_t=S_{t-1}-k_t^\top(k_tS_{t-1})+k_t^\top v_t$，可抑制多键串扰。Gated DeltaNet 加入标量指数衰减；GDN-2 进一步把衰减门 $\alpha_t$、擦除门 $b_t$、写入门 $w_t$ 全部升级为逐通道门向量，且擦除/写入用稠密 $d\times d$ 投影。
本文的基线就是 GDN-2，所有改进（非对称低秩门、读出增益 $\gamma_t$）都建立在其逐通道门控递推之上。
分块并行训练（chunkwise）与三角系统求解
为把逐 token 递推展开成矩阵乘法，分块算法把序列切成长度 $C$ 的块，块内展开后归结为解一个 $C\times C$ 单位下三角系统 $TZ=\hat V$。当衰减逐通道变化时，构造 $T$ 需要把键除以累积衰减前缀 $\pi_s$（会下溢）或物化 $C\times C\times d_k$ 的三维衰减张量，且三角求解本身是串行的。
论文用大量篇幅分析该数值与硬件瓶颈，并给出边界枢轴内核；不理解分块形式就无法理解其 TPU 优化贡献。
匹配滤波器（matched filter）
信号处理经典概念：在加性噪声中检测已知形状的信号时，按各通道的信噪比（信号幅度除以噪声方差）加权求和，可使输出信噪比最大化。本文将其移植到状态读出：各状态行的"信号"是目标键值对的存活率 $\bar a_{\ell,j}$，"噪声"是行能量 $m_j^*$，最优查询增益为 $\gamma_j^*\propto\bar a_{\ell,j}/m_j^*$。
这是本文第二个核心创新（读出增益 $\gamma_t$）的理论来源，直接决定多跳检索与长度外推的收益。
TPU 的 MXU 与 VPU
现代 TPU（如 TPU7x）用 $256\times256$ 的 systolic 阵列（MXU）做矩阵乘，吞吐约 2,307 TFLOP/s，而向量处理器（VPU）吞吐仅约其 1/64；矩阵乘任一维度小于 256 时 MXU 会零填充浪费算力。高效内核必须把尽量多的计算表达成 256 对齐的矩阵乘。
本文的门秩分配（总秩 $R_{tot}=512=2\times256$）、块大小 $C=256$ 等设计都是围绕该硬件特性做的协同设计。
研究动机
Transformer 的 softmax 注意力训练代价随序列长度二次增长、推理 KV 缓存线性增长；线性注意力虽把历史压进固定大小状态解决了这一点，但有限维状态只能容纳有限个正交键方向，键一多就相互串扰。GDN-2 用逐通道的衰减/擦除/写入门缓解串扰，却带来三个被忽视的瓶颈。其一，参数分配失衡：GDN-2 对擦除门 $b_t$ 和写入门 $w_t$ 各用一个稠密 $d\times d$ 矩阵，仅这两个矩阵就消耗 $2d^2$ 参数（占每层 $12d^2$ 非嵌入预算的六分之一），迫使 FFN 宽度从 $4d$ 压缩到 $3d-256$（$d=2048$ 时从 8,192 降到 5,888，缩水超 25%）。其二，读出失衡：逐通道衰减让慢衰减状态行在数百到数千个 token 上累积写入，快衰减行只保留近期输入，状态行范数跨度超过一个数量级；在无加权读出 $o_t=q_tS_t$ 下高范数慢行主导输出，模型只能靠更强的写入期擦除压制干扰，导致旧联想被过早覆盖。其三，训练瓶颈：分块算法需除以累积衰减前缀（快速衰减通道在 $C=256$ 时会下溢到约 $10^{-77}$，float32 都无法表示），否则要物化每头 32 MB 的三维衰减张量并在比 MXU 慢约 64 倍的 VPU 上串行解三角系统，宽矩阵单元利用率极低。
本文的目标是本文要在与 Transformer 和 GDN-2 严格等参数（每层 $12d^2$ 非嵌入参数）的约束下重新设计通道门控 Delta 规则模型，同时解决上述三个瓶颈：第一，用实测的状态更新敏感性指导门控参数分配，把省下的门控参数全部再投资到 FFN，使 FFN 宽度恢复甚至超过 Transformer（$d=2048$ 时从 5,888 加宽到 7,680，+30.4%）；第二，在查询端引入读出增益 $\gamma_t$ 平衡状态行贡献，降低写入期擦除需求、延长记忆时域；第三，设计一个数值稳定、能把 98% 计算路由到 MXU 的大块（$C=256$）分块内核。可量化的目标包括：验证 NLL 在 166M–1.34B 全规模优于 GDN-2 与 Transformer，抽取式检索超过 GDN-2，从 4K 训练长度外推到 10M token 流无困惑度漂移，且吞吐与 GDN-2 持平、在 32K 序列上比 Transformer 快约 2.7 倍。
与已有工作不同的是，本文的独特之处在于把「状态动力学分析」放在架构设计之前：先推导可检验的闭式规律，再据此分配参数，而不是像以往工作那样凭直觉堆门。三个被前人忽视的观察构成切入点。第一，状态更新对写入远比对擦除敏感（训练后模型实测约 13 倍）：SiLU 激活加 $\ell_2$ 归一化把单位键能量集中到极少数活跃通道（128 个通道中参与率仅 5.98–7.43，约 5%），擦除扰动在不活跃通道上几乎无效，因此稠密擦除门是浪费，低秩加偏置就够。第二，逐通道衰减导致的行能量稳态差异可以用闭式定律 $m_j^*=c_j\sigma_v^2\tau_{\mathrm{eff},j}$ 精确描述，而信号处理中的匹配滤波器恰好给出最优读出权重 $\gamma_j\propto\bar a_{\ell,j}/m_j^*$——即「查询期补偿」可以替代「写入期擦除」来压制干扰，这一读出—擦除替换关系此前无人利用。第三，现有 GPU 分块内核（16×16 小块、SRAM 内串行求解）不适用于以 256×256 MXU 为核心、VPU 慢 64 倍的 TPU，需要全新的无除法、无三维衰减张量的分块算法。
核心方法
先给直觉：把状态矩阵想成一排书架，每行（通道）像一个按不同速度褪色的书架——慢衰减行积攒了厚厚的书（大范数），快衰减行只有几本新书。GDN-2 的问题是读的时候不挑行，随便一查都被厚书架淹没，于是只好在写的时候频繁把旧书擦掉（强擦除）来防止噪音堆积。Muir 的做法是两手：(1) 写入端，键本身只落在少数几个架子上，擦除门用低秩投影就够，省下的参数全部给 FFN 扩容；(2) 读出端，给查询装一个「匹配滤波器」$\gamma_t$，按每个书架的信噪比（目标信号存活率/行能量）加权，厚书架被压低、快行被适当放大，干扰在查询期就被消掉，写入期就不必猛擦。技术路线上，论文先推导四组理论：单步门敏感性（Prop 3.1）、稳态行能量闭式定律（Prop 3.4）、滞后相关的匹配滤波读出（Prop 3.6 与 Corollary 3.8）和擦除—读出替换（Prop 3.10），再用合成联想记忆实验验证，然后落地为带第五个门 $\gamma_t$ 的 Muir 架构——五个门共享一个融合低秩投影，秩 $(64,64,128,128,128)$、总秩 512 正好填满两块 256 列 MXU tile，FFN 按 $d_{\mathrm{ffn}}=4d-512$ 加宽；最后为 TPU 设计边界枢轴分块内核（无除法、无三维衰减张量、分治块求逆、98% FLOPs 在 MXU）。
核心创新有两点，且都由理论直接推出而非经验调参。第一是「敏感性分配的门控低秩化」：论文证明（Prop 3.1）状态对写入门扰动的期望敏感度为 $\|v_t\|_2^2$、对擦除门仅为 $\sum_j k_{t,j}^2\|\hat S_{t,j,:}\|_2^2$——擦除要起作用必须同时满足对应状态行已有能量且键在该通道有权重；而 SiLU+单位归一化的键把能量集中在约 $6/128$ 个通道上，所以擦除门与衰减门用秩 64 瓶颈即可，写入门、输出门、读出增益直接调制全维向量需秩 128。这从根上区别于 GDN-2 的对称稠密设计，也区别于 KDA/RWKV-7 用共享标量回避稠密门的做法——Muir 保留逐通道解耦但用理论决定谁压缩多少。第二是「匹配滤波读出替代写入擦除」：检索滞后 $\ell$ 下的最优查询增益为 $\gamma_j^*(\ell)\propto\bar a_{\ell,j}/m_j^*\approx\lambda_j^{\ell/2}/\rho_j^2$；静态版本对 horizon 边缘化后呈分数幂律 $\gamma_j\propto\rho_j^{-p}$（$p\in(0,1)$，实测 $p\approx0.35$），动态版本 $\gamma_t=\gamma_{\mathrm{stat}}\odot 2\sigma(z_{q,t})$ 逐 token 适配滞后。Prop 3.10 进一步证明查询期压干扰可以把均衡擦除率压低（实测 $\bar\mu$ 从 0.782 降到 0.652），所有通道的有效记忆时域 $\tau_{\mathrm{eff},j}$ 同时变长——这解释了 Muir 为什么能把 passkey 检索推到 1M token。
方法步骤详情
完整流程分四步。第一步「投影与门控生成」：token 表示 $x_t\in\mathbb{R}^{1\times d}$ 经 $W_q,W_k,W_v$ 投影，过深度因果卷积与 SiLU，$q_t,k_t$ 归一化到单位球；同时 $x_t$ 过融合下投影 $W_\downarrow\in\mathbb{R}^{d\times512}$ 得到瓶颈 $h_t$，再按门各自由上投影加通道偏置生成 logits：$\log\alpha_t=-\exp(a_{\log})\mathrm{softplus}(z_{f,t})$（衰减）、$b_t=2\sigma(z_{b,t})$（擦除）、$w_t=\sigma(z_{w,t})$（写入）、$g_t=z_{g,t}$（输出）、$\gamma_t=\gamma_{\mathrm{stat}}\odot 2\sigma(z_{q,t})$（读出增益），秩分配 $(r_f,r_b,r_g,r_w,r_q)=(64,64,128,128,128)$，且列顺序对齐 256 列 MXU tile 与 128-lane VMEM 边界。第二步「状态更新」：每头执行 $S_t=(I-k_t^\top(b_t\odot k_t))\mathrm{Diag}(\alpha_t)S_{t-1}+k_t^\top(w_t\odot v_t)$，即先衰减、再沿键方向擦除旧预测、最后写入门控值。第三步「匹配滤波读出」：$o_t=(\gamma_t\odot q_t)S_t$，逐头 RMSNorm 后乘 $\mathrm{SiLU}(g_t)$，拼头过 $W_o$ 回残差流；FFN 按 $d_{\mathrm{ffn}}=4d-512$ 加宽（$d=2048$ 时 7,680）。第四步「TPU 训练内核」：把序列切成 $C=256$ 的块，累积对数衰减用 MXU 乘下三角全 1 矩阵得到；叶子块（$n_0=8$）用 Neumann 级数的平方链求逆；每层合并用边界枢轴把 $\pi_{t,s}=\exp(\Lambda_t-\Lambda_p)\odot\exp(\Lambda_p-\Lambda_s)$ 拆成两个落在 $(0,1]$ 的因子，全部化为矩阵乘；自定义 backward 只保存 $T^{-1}$ 与解 $Z$、合并时即时重算枢轴缩放。最终每 token 每头 1.37M MXU FLOPs（98%）对 29.7K VPU FLOPs（2%），VPU 计算降 24 倍、总 FLOPs 降 12%。
技术新颖性
与已有工作的本质区别体现在三处。其一，参数分配依据不同：GDN-2 对称地给擦除/写入门各配稠密 $d\times d$ 矩阵，KDA 与 RWKV-7 为省参数把键擦除或值写入退化为共享标量；Muir 是第一个用一阶敏感性分析（13 倍写入/擦除不对称、键参与率约 5%）定量证明「擦除可以低秩、写入必须高秩」并据此做非对称秩分配的工作。其二，读出机制不同：query-key 归一化注意力与对角 SSM（如 Mamba-3）给 q、k 双侧加逐通道增益，但 Delta 规则要求 $\|k_t\|_2=1$ 以保证秩一擦除的收缩性（特征值在 $(-1,1)$），因此以往只能无加权读出；Muir 的 $\gamma_t$ 只缩放 $q_t$、不动 $k_t$，首次在收缩性约束内实现了逐通道、逐 token 的状态行再平衡，并给出闭式最优解（匹配滤波/分数幂律 $p\in(0,1)$，明显优于 naive 的完全白化 $p=1$）。其三，内核设计不同：GPU 上 FLA/Kimi 内核用 16 token 子块在 SRAM 内串行解三角系统，适配 16×16 TensorCore；Muir 的边界枢轴算法用分治块求逆加枢轴衰减拆分把 $C=256$ 大块全部化为 MXU 矩阵乘，避免了前缀除法下溢（会使 25% 的头产生 NaN）和三维衰减张量（每头 32 MB），是面向宽 systolic 阵列的全新分块并行形式。
实验结果
主实验在 FineWeb-Edu 100B token 上以等参数设置训练 166M/406M/834M/1.34B 四档模型（8 张 TPU7x、序列长 4K、batch 524,288 token/step）。语言建模：1.34B 时 Muir 验证 NLL 1.9973，比 GDN-2（2.0177）低 20.4 mnats、比 Transformer（2.0335）低 36.2 mnats；四档汇总对 GDN-2 降 18.0–20.4 mnats、对 Transformer 降 25.0–36.2 mnats（困惑度分别降 1.8–2.0% 与 2.5–3.6%）；WikiText 困惑度 16.41（GDN-2 为 16.80），LAMBADA 困惑度 10.00 为全场最低；8 任务零样本常识平均 53.84，比 GDN-2 高 1.38 分（部分指标最高提升 1.9 分）；5-shot TriviaQA 19.32% 显著高于 GDN-2 的 16.48%。消融显示 FFN 扩容（Muir no γ）贡献主要 NLL 增益，而动态读出增益在上下文敏感任务上不可替代（LAMBADA 困惑度 10.38→10.00，TriviaQA 16.52%→19.32%）。检索：RULER 六任务 1.34B 平均（1K–8K）Muir 67.2，远超 GDN-2 的 55.3 与 Transformer 的 58.0；5 跳变量追踪 VT 平均 34.3 对 GDN-2 的 21.5；MV-NIAH 68.9 对 53.8；MQ-NIAH 69.1 对 47.6；Based/JRT 六任务 2K 平均 34.60 对 GDN-2 的 32.08（动态 γ 贡献约 69% 的增益），不过 Transformer（45.92）在短上下文抽取上仍然领先。长度外推：Muir 从 4K 训练长度零样本外推到 10M token 流（2,560 倍），25 MB float32 状态无困惑度漂移，动态 γ 把尾部漂移比静态 γ 降低 3.6 倍；passkey 检索中静态 γ 版本在 1M token 达到 90.0%（零样本）/99.4%（64K cooldown 后），而 GDN-2 在 4K–8K 后崩溃，Dynamic YaRN Transformer 需 2.1 TB KV 缓存且 8K 后发散。效率：Muir 与 GDN-2 吞吐持平（4K 时 179.4 对 179.3 K tokens/s），在 16K/32K 分别比 Transformer 快 1.64×/2.74×，且约 85B token 即达到 GDN-2 100B token 的验证损失（1.2 倍提速）。机理分析闭环验证理论：写入/擦除敏感度比 13.3 倍、键参与率 5.98/128、行能量定律 Spearman 0.960、学到的静态增益幂指数 p=0.35、并行擦除率降 16.6%、时域回看窗 49.3→78.7 token（+59.6%）、多跳根保留率 37%→45%。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 语言建模（FineWeb-Edu 验证 NLL，1.34B） | token NLL（越低越好） | 1.9973 | GDN-2 2.0177；Transformer 2.0335 | 较 GDN-2 −20.4 mnats，较 Transformer −36.2 mnats |
| 语言建模（WikiText 词级困惑度，1.34B） | perplexity | 16.41 | GDN-2 16.80；Transformer 17.17 | 较两基线均更低（约 −2.3%/−4.4%） |
| LAMBADA 词级困惑度（1.34B） | perplexity | 10.00 | GDN-2 10.38；Transformer 10.30 | 全场最低，动态 γ 是关键 |
| 零样本常识推理 8 任务平均（1.34B） | accuracy (%) | 53.84 | GDN-2 52.46；Transformer 52.05 | +1.38 / +1.79 分 |
| 5-shot TriviaQA（1.34B） | accuracy (%) | 19.32 | GDN-2 16.48；Transformer 14.56 | +2.84 分 |
| RULER 合成检索 6 任务平均 1K–8K（1.34B） | accuracy (%) | 67.2 | GDN-2 55.3；Transformer 58.0 | +11.9 / +9.2 分 |
| 5 跳变量追踪 VT 平均 1K–8K（1.34B） | accuracy (%) | 34.3 | GDN-2 21.5；Transformer 8.8 | +12.8 分 |
| Based/JRT 真实抽取 6 任务平均 2K（1.34B） | accuracy (%) | 34.60 | GDN-2 32.08（Transformer 45.92 仍最高） | +2.52 分（仍落后 Transformer 11.3 分） |
| Passkey 检索（RULER S-NIAH-1，1M token，静态 γ 变体） | accuracy (%) | 90.0（零样本）/ 99.4（64K cooldown 后） | GDN-2 在 4K–8K 后崩溃；Transformer ≥128K 超 HBM | 有效检索距离扩大约两个数量级 |
| 训练吞吐（1.34B，32K 序列，8×TPU7x） | K tokens/s | 与 GDN-2 持平（4K：179.4 vs 179.3），为 Transformer 的 2.74× | GDN-2（≈1.0×）；Transformer（1×） | 较 Transformer：16K 1.64×、32K 2.74× |
局限与改进
作者承认两点局限。第一，固定门秩预算下静态与动态读出存在硬性取舍：静态 γ（保留 $r_w=256$ 全写入秩）能把 passkey 检索维持到 1M token，但无法匹配多跳/多值检索的多个滞后（VT 上反而低于无 γ 版本），且在连续长文本上因慢行能量持续累积而在 32K 后漂移（Δtail=+0.017 nats）；动态 $\gamma_t$（$r_w,r_q=128,128$）解决多滞后与漂移，却在 passkey 上于 32K 后逐渐衰减——单一变体无法同时最优。第二，内核层面纯 JAX 实现在 <11K 短上下文上仍慢于手工 Pallas Splash Attention，因为 XLA 会把块间临时量物化到 HBM。我的独立观察：(1) 短上下文真实抽取（Based/JRT 2K 平均 34.60 对 Transformer 45.92）差距明显，FFN 扩容并未弥补循环状态在精细抽取上的劣势；(2) 理论建立在平稳 i.i.d. 流、各向同性扰动与二阶矩闭合等假设上，13 倍敏感度等数字依赖其特定配置（SiLU 键、单位归一化、深度卷积），换键构造方式可能改变结论；(3) 实验规模上限 1.34B/100B token 且只用 FineWeb-Edu，更大规模下最优秩分配与 γ 收益是否保持未知；(4) 秩常数（64/128/512）与块大小 C=256 深度绑定 TPU7x 的 256×256 MXU，向其他硬件移植需重新调参。
独立分析的弱点
独立分析出四个弱点。第一，读出增益的「二选一」困境本质是秩预算分配问题：模型无法在推理时按任务切换静态/动态模式，比如同一个部署好的模型既要做远距 passkey 又要多跳路由就必须取舍。改进方向是让每个头学习 $r_w/r_q$ 的软分配，或按作者建议用解析匹配滤波器（每头预测一个滞后并跟踪期望行能量）替代可学习的稠密 $z_q$，把省下的秩还给写入。第二，匹配滤波的静态成分 $\gamma_{\mathrm{stat}}$ 在 4K 标定，遇到 10M 流时慢行能量远超训练分布，导致 32K 后漂移；改进方向是在线估计行能量 $\hat m_j$ 并做自适应归一化（类似 running RMS），而不是依赖训练时学到的固定 profile。第三，理论—实践耦合脆弱：低秩擦除的前提是键能量集中（参与率约 5%），这依赖 SiLU+卷积+单位归一化的键构造；若未来架构改用其他键生成方式，擦除门可能需要更高秩。改进方向是把秩分配做成可微分正则（如对 $W_b$ 加 nuclear-norm 惩罚），让数据自己决定秩而非手工设定 64/128。第四，评测以合成检索为主，缺少指令跟随、代码、多语言等真实长文档任务的验证，也未报告下游微调后的表现；且 1.34B/100B 规模下动态 γ 的 NLL 增益（−3.3 mnats）相对 FFN 扩容（约 −17 mnats）较小，更大数据量下两者占比可能反转，需要在 ≥7B 规模与真实语料上复验。
未来方向
作者提出四个方向：(1) 把动态读出改为解析式匹配滤波——每头预测一个标量滞后并跟踪期望行能量，理论上可统一远距 passkey、多滞后路由和流稳定性而不牺牲写入秩；(2) 架构级解法：滑窗注意力加静态读出 Muir 的混合结构，局部注意力以常数内存精确复制近处内容，循环状态保留满写入秩专职远距查找；(3) 内核级：把叶子 Neumann 求逆、枢轴合并树和跨块状态扫描融合成单个 Pallas 内核，把临时量留在 VMEM 以缩小短上下文吞吐差距；(4) 把敏感性分配门控与匹配滤波读出推广到矩阵状态的测试时优化模型（TTT、Titans、ATLAS 等）。基于本文成果还可延伸：把闭式行能量定律用作训练监控与诊断工具（预测哪些行会积压能量、何时需要正则）；在混合架构（间隔插入 softmax 层）中验证 $\gamma_t$ 与全注意力的互补性；把匹配滤波思想用于 KV 缓存压缩的逐头重要性加权；以及在 GPU（Triton/FLA 框架）上复现边界枢轴内核，检验方法的硬件无关性。
复现评估
复现评估：论文未提及公开完整训练代码或模型权重，但附录 J 提供了 token mixer 与边界枢轴分块前向/反向的 JAX 参考实现（Listing 1），并指出在 GPU 上由于 $\hat q_t=\gamma_t\odot q_t$ 只是把 GDN-2 递推换了查询，Muir mixer 可直接跑在现成的 FLA/Triton GDN-2 内核上（关闭核内 qk 归一化即可），架构复现门槛低。数据全部公开：FineWeb-Edu 预训练语料、WikiText/LAMBADA/RULER/Based-JRT 评测集、Llama-2 32k tokenizer。算力方面：完整预训练需 8 张 TPU7x、100B token、约 19 万步（batch 524,288），普通实验室难以复现全量；但理论部分极易复现——合成实验（单头 $d_k=d_v=64$）与机理探针只需 CPU/单 GPU 量级，166M 档在单节点 8 卡上可行。综合难度：架构与内核复现为中等（有参考实现且附录 D/F 公式详尽），全量训练结果复现为高（TPU7x 专属内核与 XLA 编译调优不可迁移）。对想试用方法的人，最短路径是在 FLA 框架的 GDN-2 基础上加融合低秩门投影与读出增益 $\gamma_t$。
论文图表
