# Muir：线性注意力中的敏感度分配门控与匹配滤波读出 - **作者**：Tushaar Gangavarapu、Gil Tabak、Clemens JS Schaefer - **原文版本**：[arXiv:2611.08376v1](https://arxiv.org/abs/2611.08376v1)，2026-11-08 获取；[本地 PDF](../source/paper.pdf) - **发表**：截至 2026-11-08 未查到正式会议或期刊发表记录；Google Research 条目将其列为 arXiv (2026)（[官方记录](https://research.google/pubs/muir-linear-attention-with-sensitivity-allocated-gating-and-matched-filter-readout/)）。 - **代码**：未查到独立公开代码仓库；作者在论文附录 J 提供了 JAX 版 token mixer 与分块核参考实现（[原文](https://arxiv.org/abs/2611.08376v1)，PDF pp. 62–65）。 ## 论文总览 Muir 研究的是一个看似矛盾的问题：线性注意力把任意长的历史压入固定大小的矩阵状态，因此推理状态不随上下文增长；但有限状态容量会让不同键相互干扰，模型需要在“忘得快以抑制干扰”和“记得久以保留旧信息”之间取舍。作者以 Gated DeltaNet-2（GDN-2）为基线，指出它的门控参数分配、状态读出和 TPU 分块训练各有一个瓶颈。 论文的论证主线是： 1. **写入门比擦除门更值得分配容量。** 单步状态对写门的敏感度不依赖当前键的稀疏度；擦除门的作用则同时乘上键坐标平方和已有状态行能量。经过 SiLU 和单位范数化的键只把大部分能量放在少数通道，因此擦除门可用更低秩投影。 2. **不同衰减速度造成状态行能量失衡。** 慢衰减行累积历史写入，快速衰减行主要保留近期信息。无权读出时，高能量行会压过低能量行；单靠增加写时擦除来压噪声，又会更快抹掉远处的绑定。 3. **读出时按通道重加权可以分担擦除。** 静态读出增益适合单一检索时距；动态增益可以随查询变化，在不同检索时距间切换。读出抑制干扰后，写时擦除可以降低，状态有效记忆跨度随之增加。 4. **低秩门控节省的参数用于加宽 FFN。** 主模型在每层参数预算不变时，将 FFN 宽度从 GDN-2 的 5,888 增至 7,680（在 $d=2{,}048$ 时增加 30.4%），同时容纳一个查询读出门。 5. **边界枢轴分块核让大块 TPU 计算既稳定又并行。** 它避免对衰减前缀做除法，避免物化完整的三维衰减张量，并把三角求解改成分治式块矩阵运算。 论文覆盖 166M–1.34B 参数的模型，并报告语言建模、常识题、合成和真实检索、长度外推、吞吐以及机制探针。强项是理论解释、机制诊断与架构和硬件实现之间的对应；主要取舍是静态读出与动态读出并非同一配置的最优选择：前者更适合单一远距查询，后者更适合多跳查询和连续长流。 ## 结构地图 | 原文章节 | 内容 | PDF 页码与定位 | |---|---|---| | 1. Introduction | 线性状态容量、GDN-2 的参数和读出问题、TPU 分块瓶颈 | pp. 1–3，Fig. 1 p. 2 | | 2. Preliminaries | 通道门控 Delta Rule、Chunkwise Parallel Form、TPU 硬件背景 | pp. 3–5 | | 3. Channel-Wise Gated Delta Memory Dynamics | 写/擦敏感度、稳态行能量、匹配滤波、合成实验 | pp. 5–9，Fig. 2 p. 9 | | 4. Muir | 五门融合低秩投影、读出增益、Boundary-Pivoted Chunkwise Kernel | pp. 10–14，Fig. 3 p. 10、Fig. 4 p. 13 | | 5. Experiments | 预训练配置、语言建模、检索、长度外推、吞吐、数值稳定、消融 | pp. 14–21，Figs. 5–9 pp. 16–20，Table 2 p. 21 | | 6. Mechanistic Analysis | 门敏感度、行能量、读出增益、擦除干预、多跳诊断 | pp. 21–24，Fig. 10 p. 22、Figs. 11–12 p. 23、Fig. 13 p. 24 | | 7. Conclusion | 总结、静态/动态增益取舍与实现限制 | pp. 24–25 | | Appendices | 证明、合成任务、核分析、配置、完整结果和参考实现 | proofs pp. 40–43；Tables 6–23 pp. 44–61；JAX 参考实现 pp. 62–65 | ## 一：动机 ### 固定状态带来吞吐，也带来干扰 对每个注意力头，线性注意力用 $S_t\in\mathbb{R}^{d_k\times d_v}$ 汇总过去的键值关联。它避免 Transformer 随序列长度增长的 KV cache，但不同 token 必须共享有限数量的键方向。键重叠时，读出一个值也会带出其他写入，Delta Rule 因此在新写入前先减去当前键方向上的旧预测。 GDN-2 将这套更新细化为逐通道衰减、擦除和写入门。它为衰减 $f$ 与输出 $g$ 使用秩 128 投影，却为擦除 $b$ 和写入 $w$ 使用完整的 $d\times d$ 投影。与 Transformer 的每层 $12d^2$ 非嵌入参数预算对齐时，门控挤压了 FFN：在 $d=2{,}048$ 时，GDN-2 的 FFN 宽度是 5,888，而 Transformer 是 8,192。 作者提出三个相互关联的缺口： - **门控容量没有按状态敏感度分配。** 擦除和写入都用 dense 投影，尽管擦除仅沿当前键激活的方向影响状态。 - **读出与状态行的噪声规模不匹配。** 通道衰减使慢行能量远大于快行；无权查询会让慢行主导输出。若只靠更强擦除来抑制这类干扰，就会缩短所有通道的有效记忆。 - **通道衰减使大块并行训练数值困难。** 传统分解需要除以极小的累计衰减，或者构造 $C\times C\times d_k$ 张量、顺序求解三角系统；这些操作在 TPU 的慢速向量单元上成为瓶颈。 ## 二：方法 ### 2.1 GDN-2 状态更新与 Muir 读出 每头的键、查询经 SiLU 后作 $\ell_2$ 归一化，值保持未归一化。令 $\hat S_t=\operatorname{Diag}(\alpha_t)S_{t-1}$、$\hat k_t=b_t\odot k_t$、$\hat v_t=w_t\odot v_t$，GDN-2 的更新为 $$ S_t = \hat S_t-k_t^\top(\hat k_t\hat S_t)+k_t^\top\hat v_t = \bigl(I-k_t^\top(b_t\odot k_t)\bigr)\operatorname{Diag}(\alpha_t)S_{t-1} +k_t^\top(w_t\odot v_t), \qquad o_t=q_tS_t. $$ 这里 $k_t,q_t\in\mathbb{R}^{1\times d_k}$，$v_t\in\mathbb{R}^{1\times d_v}$；$\alpha_t\in(0,1)^{d_k}$ 控制每个状态行的保留率，$b_t\in(0,2)^{d_k}$ 控制沿当前键的擦除，$w_t\in(0,1)^{d_v}$ 控制每个值通道写入量。输出门 $g_t$ 在每头 RMSNorm 后逐通道调制。 Muir 保留同一递归，但读出时增加查询增益： $$ o_t=(\gamma_t\odot q_t)S_t,\qquad \gamma_t=\gamma^{\mathrm{stat}}\odot 2\sigma(z_t^q). $$ $\gamma^{\mathrm{stat}}$ 是可学习的静态通道轮廓，$2\sigma(z_t^q)\in(0,2)$ 是依赖当前 token 的动态调制；初始化时两者的乘积为 1。它只缩放查询，不缩放写状态用的键，因此不破坏 $k_t$ 的单位范数。 ### 2.2 理论分析：写入敏感度、行能量与匹配滤波 **写门和擦除门进入状态更新的方式不对称。** 对任意扰动 $\delta b,\delta w$，命题 3.1 给出 $$ D_bS_t[\delta b]=-k_t^\top\bigl((\delta b\odot k_t)\hat S_t\bigr),\qquad D_wS_t[\delta w]=k_t^\top(\delta w\odot v_t). $$ 在各维独立、单位方差的各向同性扰动下， $$ \mathbb{E}\|D_bS_t\|_F^2=\sum_j k_{t,j}^2\|\hat S_{t,j,:}\|_2^2, \qquad \mathbb{E}\|D_wS_t\|_F^2=\|v_t\|_2^2. $$ 因此写门每一步都能改变输入值；擦除门只有在“该行已有能量”且“当前键使用了该行”时才明显影响状态。训练检查点中键的有效参与通道数仅约 6/128（Muir 为 5.98，GDN-2 为 7.43），测得写门对状态的敏感度约为擦除门的 13.3 倍（166M，Table 3 / Fig. 10，PDF pp. 22–23）。作者还把擦除键 $b\odot k$ 分解为沿 $k$ 的分量和正交残差 $s$；该残差形成非正规擦除变换，增大转移矩阵谱范数。低秩擦除投影可减少这种无效的斜向变化，同时通道偏置仍能保留不同的基础擦除率（Sec. 3.1、Appendix D.1，PDF pp. 5–6、40–41）。 **稳态行能量由注入强度和有效时间常数共同决定。** 在作者设定的平稳 i.i.d. 输入流中，记 $c_j=\mathbb{E}[k_{t,j}^2]$、$\lambda_j=\mathbb{E}[\alpha_{t,j}^2]$、$\bar b_j=\mathbb{E}[b_{t,j}]$、$\sigma_v^2=\mathbb{E}\|\hat v_t\|_2^2$。在行近似不相关及二阶矩闭合后，稳态行能量近似为 $$ m_j^\ast=\mathbb{E}\|S_{t,j,:}\|_2^2 \approx \frac{c_j\sigma_v^2}{1-\lambda_j+2\lambda_j\bar b_jc_j} =c_j\sigma_v^2\,\tau_{\mathrm{eff},j}, \quad \tau_{\mathrm{eff},j}=\frac{1}{1-\lambda_j+2\lambda_j\bar b_jc_j}. $$ 快衰减由 $1-\lambda_j$ 主导；慢衰减时，状态行会积累更多值写入，除非通过擦除额外缩短其时间常数。该公式中的 key 四阶矩先有精确递推，之后再用 $\mathbb{E}[k_i^2k_j^2]\approx c_ic_j$ 和 $\mathbb{E}[b_j^2]\approx\bar b_j^2$ 做闭合；所以它是有明确假设的近似定律，不是对任意数据流的恒等式（Proposition 3.4、Appendix D.2，PDF pp. 6、41–42）。 **匹配滤波在查询时权衡目标信号与干扰。** 若要在时距 $\ell$ 后取回曾写入的 $v_{t-\ell}$，每个通道的目标信号存活率记作 $\bar a_{\ell,j}$，干扰行方差为 $m_j^\ast$。在目标信号与跨写入干扰满足文中独立/不相关条件时，期望 MSE 为 $$ \mathcal L_\ell(\gamma) =\sigma_v^2\left(1-\sum_j\gamma_j c_j\bar a_{\ell,j}\right)^2 +\sum_j\gamma_j^2c_jm_j^\ast. $$ 最优读出方向因此满足 $\gamma_j^\ast(\ell)\propto \bar a_{\ell,j}/m_j^\ast$：只按行范数归一化会过度放大目标信号已消失的快行；最优权重同时看“剩余目标信号”和“行噪声”。但 $\bar a_{\ell,j}$ 随时距改变，不同 $\ell$ 的最优向量一般不共线，一个静态增益不能对所有查询距离都最优（Proposition 3.6、Remark 3.7，PDF pp. 7–8）。 对均值为 $H$ 的几何检索时距分布，作者得到静态滤波器 $$ \gamma_j^{\mathrm{stat}}\propto \frac{g_H(\lambda_j)}{m_j^\ast}, \qquad g_H(\lambda_j)= \frac{\sqrt{\lambda_j}}{1+(H-1)(1-\sqrt{\lambda_j})}. $$ 这会产生介于不加权和全白化之间的分数幂关系 $\gamma_j\propto \rho_j^{-p}$。在训练模型中拟合到的中位指数为 $p=0.35$。理论上，读出滤波减少高能行的干扰后，可以降低写时擦除率，使各通道的 $\tau_{\mathrm{eff},j}$ 增大；这是“读出替代部分擦除”的机制（Corollary 3.8、Proposition 3.10，PDF pp. 8、42–43）。 ![图 2：受控合成记忆实验](figures/fig2_toys.png) 图 2 的三个面板依次比较静态增益与擦除强度、短/长时距混合任务，以及学习任务中的门秩。单头合成任务有 64 个键通道；在 $H=27$ 时，完全逆 RMS 白化比单位增益更差，而半幂和理论滤波降低 MSE，并将最优擦除 $\mu^\ast$ 从 0.60 降到 0.50。拟合滤波与静态最小二乘 oracle 基本重合。在另一项 20-seed 的端到端两头任务中，低秩擦除门在所测秩上改善长时距召回；写门降到秩 8 以下则损害表现，秩 16 接近 dense，秩 32 超过 dense。图和参数见 Sec. 3.2 / Fig. 2，PDF pp. 8–9；完整数值见 Appendix E / Tables 6–7，pp. 44–45。 ### 2.3 五门融合投影与参数再分配 Muir 将衰减 $f$、擦除 $b$、输出 $g$、写入 $w$、查询增益 $q$ 五路门控合并为一个共享降维投影，再用各自的上投影和通道偏置产生 logits： $$ h_t=x_tW_\downarrow,\quad \dim(h_t)=R_{\mathrm{tot}}=512,\qquad z_t^m=h_t^mW_{m\uparrow}+c_m,\quad m\in\{f,b,g,w,q\}. $$ 主模型的秩分配是 $(r_f,r_b,r_g,r_w,r_q)=(64,64,128,128,128)$。低敏感度的衰减和擦除门用秩 64；直接调节输出、写入和值读出的三路门用秩 128。总秩 512 刚好填满两个 $256$ 列 MXU tile，且分组方式也对齐 TPU 的 128-lane VMEM tile，减少窄矩阵乘法的填零浪费。消融不使用动态 $q$ 门时，将这 128 秩转给写门，$r_w=256$，保持总秩不变。 在每层 $12d^2$ 参数预算下，Muir 的门控投影参数量为 $2dR_{\mathrm{tot}}$，并设 $d_{\mathrm{ffn}}=4d-R_{\mathrm{tot}}$。在 $d=2{,}048$ 时，GDN-2 的门投影为 9.44M 参数，Muir 为 2.10M；节省的 7.34M/层把 FFN 从 5,888 扩至 7,680（+30.4%），保持总参数相同。若固定原 FFN 宽度，门控压缩可让整层非嵌入参数与线性 FLOPs 下降 14.6%。(Sec. 4.1 / Fig. 3，PDF pp. 10–11；配置见 Appendix G / Table 11，p. 49。) ![图 3：Muir 模块、五门秩分配和参数预算](figures/fig3_architecture.png) 左侧是残差块与 Muir token mixer；右上是按硬件 tile 对齐的秩分配；右下显示等参数条件下将门投影节省量投入 FFN。作者报告的是模型设计和参数/FLOP 预算，不是单独证明 FFN 加宽能解释所有检索提升。 ### 2.4 Boundary-Pivoted Chunkwise Kernel 训练时，逐 token 更新状态的算术强度低。把长度为 $T$ 的序列拆成 $C=256$ token 的块，可以把块内运算改成矩阵乘法。定义相对衰减 $\pi_{t,s}=\exp(\Lambda_t-\Lambda_s)$，则块内 Delta Rule 化为求解单位下三角系统 $TZ=[\hat V,\hat K_{\odot\pi}]$，其中 $T=I+L_{\mathrm{strict}}$；求得 $U,Y$ 后即可用矩阵乘法算块输出和下一块状态。 直接实现有两个问题：用 $\pi_t/\pi_s$ 拆解会在快衰减通道上令分母下溢、倒数溢出；逐对计算虽稳定，却需要 $C\times C\times d_k$ 张量并由向量单元处理；顺序前代也受慢向量单元限制。 Muir 对三角矩阵做分治块逆。每个长度为 8 的叶块直接从非正的 log-decay 差计算衰减，并用有限 Neumann 级数求逆；之后反复合并相邻块。对跨左右子块的键对，选左块末端为枢轴 $p$，将相对衰减拆成 $$ \pi_{t,s}=\exp(\Lambda_t-\Lambda_p)\odot\exp(\Lambda_p-\Lambda_s), \qquad s\le p