Switching Linear Attention：中文精读笔记
论文：Hyun Dong Lee、Xavier Gonzalez、Nicolas Zucchet、E. Kelly Buchanan、Emily B. Fox、Scott W. Linderman，Switching Linear Attention。
原文：arXiv:2609.39034v1；本地 PDF。本笔记研读版本为 v1，获取与核查日期为 2026-10-28。下文页码均指从 1 开始的 PDF 页码。
发表：已发表于 COLM 2026 官方预论文集；会议录用论文名单亦列出此文。预论文集的最终 PDF 已公开，会议 DOI 尚待分配。
代码：作者官方公开实现，仓库包含 SwiLA 层、实验和示例。
论文总览
标准 softmax 注意力能按查询从历史键值对中作非线性检索，但自回归推理需要随序列长度增长的 KV 缓存；DeltaNet 一类线性注意力把历史压进固定大小矩阵，却只能用一张线性映射读取。论文提出 SwiLA（Switching Linear Attention）：每个注意力头保存多个在线更新的线性回归器，让每个输出维度按当前输入及预测误差选择专家。写入时用观察到的值计算专家的后验责任，读取时用查询侧先验混合专家预测。这样，读取映射随输入变化，推理状态大小仍不随历史长度增长。（Sec. 1–3，PDF pp. 1–5）
实验支持一个有条件的结论：在等循环状态容量的软最大注意力拟合任务中，多专家分配通常优于把容量全用于独立头；在需要保持情境的 CDAR 任务中，Temporal SwiLA 用较小状态达到接近满分；374M 参数语言模型的纯循环 SwiLA 变体提升了多个召回指标，但最强常识平均分仍属于 KDA，最强召回平均分仍属于带完整 KV 缓存的 FoX。1.3B 参数单 H100 吞吐测试显示 SwiLA 可用更大的生成批次，同时训练速度约为 GDN 的一半。（Sec. 5，PDF pp. 7–10；App. E，pp. 27–33）
结构地图
| 原文章节 | 在论证中的作用 | PDF 定位 |
|---|---|---|
| 1 Introduction；2 Background on Test-Time Regression | 提出 KV 缓存与固定状态检索能力的取舍，引入测试时回归 | pp. 1–3；图 1 在 p. 3 |
| 3 Switching Linear Attention | 混合回归模型、在线更新、Temporal/Gated 扩展与负载均衡 | pp. 3–5；式 (3.1)–(3.10) |
| 4 Related Work | 与 DeltaNet、Mixture-of-Memories（MoM）比较 | pp. 5–6 |
| 5 Experimental Results | softmax 拟合、CDAR、RegBench、语言模型、吞吐及消融 | pp. 6–10；图 2、表 1–2 |
| 6 Conclusion | 归纳固定状态与非线性递推的收益和代价 | p. 10 |
| Appendices A–E | 推导、时间先验、并行化设想和完整实验协议 | pp. 19–33；图 S1–S5 |
一：动机——把“检索”视为测试时拟合
测试时回归（test-time regression，TTR）把一个序列层拆成两件事：接收至今的键值对 $(k_i,v_i)$ 后，在上下文内形成记忆函数 $m_t$；查询 $q_t$ 时输出 $m_t(q_t)$。理想化目标是
这里 $k_i,v_i,q_t\in\mathbb R^D$，$\mathcal M$ 是记忆函数类。它解释了不同序列层的差别：softmax 注意力类似保存所有样本的非参数核平滑；DeltaNet 用一张矩阵 $W_t$ 做在线线性回归，读出 $W_tq_t$。前者有随上下文长度增长的 KV 缓存，后者的状态固定为每头约 $D^2$ 个数，却难以把同一个键在不同情境下映射到不同值。（Sec. 2、4.1，PDF pp. 2–3、6）
SwiLA 的想法是改变 $\mathcal M$：用 $J$ 张线性回归器及输入依赖的混合权重表示一个非线性检索函数。图 1A 用“单直线、分段直线、弯曲目标”说明表达力差别；图 1B 对比三种记忆组织。图 1C 才是实测的等状态容量拟合结果，不应把 A/B 的教学示意当作实验。（Sec. 2–3、图 1，PDF p. 3）
原文图 1，PDF p. 3。C 面板横轴为循环状态标量数、纵轴为测试集 $R^2$（越高越好）；红点代表该容量下最佳 SwiLA 的头数 $H$ 与专家数 $J$ 分解。$B=1$ 只有 $J=1$，可比较的多专家增益从 $B\ge2$ 开始。
二：方法——按预测误差写入，按查询先验读取
2.1 每个输出维度有自己的专家选择
在一个头内，$W_t\in\mathbb R^{J\times D\times D}$ 存储 $J$ 张回归矩阵。对第 $d$ 个值维度，$z_{td}\in\{1,\ldots,J\}$ 是隐含专家编号；第 $j$ 个专家的系数向量为 $w_{tjd}\in\mathbb R^D$。模型假定
$\pi_{jd}$ 是经 softmax 归一化的专家先验；键、查询、值和路由量都由输入 $x_t$ 的可学习投影产生。不同 $d$ 可选择不同 $j$，所以一条 $D$ 维输出理论上有 $J^D$ 种专家组合，而状态参数只随 $JD^2$ 增长。这个计数指组合可能性，不等于模型必然使用全部组合。（Sec. 3.1、App. B.2，PDF pp. 3、20）
2.2 写入：在线混合回归
给定旧状态 $W_{t-1}$ 与新键值对，专家 $j$ 在维度 $d$ 的预测误差为 $\delta_{tjd}=v_{td}-w_{t-1,jd}^{\top}k_t$。边缘对数似然和后验责任为
梯度 $\nabla_{w_{t-1,jd}}\ell_{td}=r_{tjd}\delta_{tjd}k_t$，因而在线更新是
这里 $\beta_{tjd}\in[0,1]$ 是输入依赖的步长。$\delta$ 告诉模型“该专家错了多少”，$r$ 告诉模型“该样本多大程度属于它”；两者相乘决定写入强度。论文称此为在线 EM：责任对应 E 步，而参数更新实际执行的是对边缘似然的一次随机梯度步，不是把当前所有历史样本重新做完整 M 步。（Sec. 3.2，式 (3.2)–(3.4)，PDF p. 4；App. B.1，p. 20）
2.3 读取：用查询侧先验组合专家
状态更新后，第 $d$ 个输出是
读取用 $\pi_{jd}(q_t)$，而非写入时的后验 $r_{tjd}$：后验需要已知 $v_{td}$ 才能计算。查询侧权重随输入变化，因此即使每个专家内部是线性的，混合后的查询到输出关系也可为非线性。以“Harry”在音乐人和魔法师段落中指向不同属性为教学例子，可以把两个专家想成两套不同的键值关联；实际训练使用软路由，未强制专家各自对应人可解释的语义。（Sec. 3.2–3.3，式 (3.5)，PDF p. 4）
2.4 Temporal、Gated 与负载均衡
Temporal SwiLA 为情境跨 token 持续提供额外分布状态。键侧把前一步后验与当前输入先验混合，查询侧把前一步过滤分布与当前先验混合：
$g$ 由输入经 sigmoid 产生；接近 0 时沿用已有情境，接近 1 时采用当前信号。随后以 $\bar r$ 作键侧责任的先验、$\bar\rho$ 作查询侧读出的先验。这是对 HMM 过滤的可计算近似，不是显式精确 HMM 推断。（Sec. 3.3、App. B.3，PDF pp. 4、20–21）
Gated SwiLA 对旧权重施加输入依赖的衰减 $\alpha$。论文将其解释为高斯权重先验带来的正则化；实际采用“先衰减、再计算责任与误差”的形式：
$r^-_{tjd}$ 同样由衰减后状态算出。先衰减与后衰减在 $J>1$ 时会改变专家责任，故不是可随意互换的实现细节。语言模型实验把门控精度沿输出维共享，相当于每个键维度一个门，与 KDA 的门控粒度一致。（Sec. 3.4，式 (3.6)–(3.8)，PDF p. 5；App. B.4、E.3，pp. 21、30）
为避免“少数专家承担全部写入”的塌缩，作者加入互信息式辅助损失 $\mathcal L_{\rm MI}=-\lambda_{\rm marg}H_{\rm marg}+\lambda_{\rm cond}H_{\rm cond}$：边际熵鼓励总体均衡，条件熵鼓励单 token 有所选择。另一个策略是在训练期给路由 logits 加入可学习强度的高斯噪声；但 374M 语言模型实验关闭了这种噪声。辅助损失同时考虑键侧与查询侧。（Sec. 3.5、App. E.3，PDF pp. 5、30）
2.5 固定状态的收益和并行化代价
若每头键值维度为 $D$，矩阵状态规模随 $HJ D^2$ 增长，却不随上下文长度 $T$ 增长；这说的是推理时循环记忆，不代表整个网络没有其他内存。相对 DeltaNet，责任 $r$ 依赖旧预测误差，因此递推对旧状态是非线性的，现有预训练实验沿时间顺序计算。附录 D 给出平行 Newton 线性化的思路：单专家、单输出维的近似雅可比为 $I-g_{tjd}k_tk_t^\top$，可把 DeltaNet 的分块算法用于每次 Newton 迭代。该近似忽略责任分母带来的跨专家耦合；作者在合成输入上约 10 次迭代收敛，但现有 Triton 实现尚未比顺序核更快。（Sec. 4.1、6，PDF pp. 6、10；App. D，pp. 23–26）
三：实验设计——各组检验什么
| 原文小节 | 检验问题与协议 | 指标与比较对象 |
|---|---|---|
| 5.1 softmax 拟合 | 用固定随机 $Q/K/V$ 的因果 softmax 输出作目标；$D=32$、长度 512，65,536 条训练序列、2,048 条测试序列；比较 $B=H\times J$ 的等矩阵状态预算 | 测试 $R^2$ 越高越好；DeltaNet、可训练 softmax 参考（PDF p. 7；图 1C p. 3、图 S1 p. 27） |
| 5.2 CDAR | 4 个情境共用键集，每个情境有 32 个键值对；长度 1,024、词表 8,192。查询组开头只给一次情境 token，要求保持情境。训练 100,000、测试 3,000 条 | 答案位置准确率越高越好；DeltaNet、DN-MoM、softmax。主图报告四个学习率与三个随机种子的最佳结果（图 2A p. 8；App. E.2 pp. 28–29） |
| 5.2 RegBench | 由概率有限自动机生成 10–20 个字符串，模型在上下文内学习下一个符号；五个 8 层模型、每层等循环状态，测试集 500 个自动机 | 下一 token 准确率越高越好，横轴为训练自动机数；Transformer、DeltaNet、DN-MoM（图 2B p. 8；App. E.2 p. 29） |
| 5.3 语言模型 | 13 个约 374M 参数、24 层、隐藏维 1,024 的模型，用 FineWeb-Edu 的 15B token 训练，序列长度 4,096；循环矩阵状态每层 131,072 个数 | 常识任务平均准确率、困惑度，以及截断到 2K token 的六项召回任务 contains 得分（表 1–2，PDF pp. 9–10；App. E.3 p. 30） |
| 5.4 吞吐；5.5 消融 | 独立的 1.3B 参数单 H100 训练/生成测试；另用 115M 参数、5B 训练 token 比较专家数、时间先验和负载均衡 | 千 token/s、峰值 GiB；或测试困惑度及路由统计（图 S3–S5，PDF pp. 31–33） |
等状态容量比较有明确边界：语言模型中 SwiLA 用 $8\times4\times64\times64=131{,}072$ 个矩阵状态标量，DeltaNet/GDN/KDA 用 $8\times128\times128$；作者通过调整前馈层宽度，使总参数量约为 374M。Hybrid Temporal SwiLA 在 24 层中有 6 个全注意力层，因此它的召回成绩同时反映 KV 缓存的作用。（App. E.3，PDF p. 30）
原文图 2A，PDF p. 8。横轴为循环状态标量数（对数刻度），纵轴为测试答案准确率；主图取学习率和种子搜索后的最佳值。
原文图 2B，PDF p. 8。横轴为训练自动机数，纵轴为留出自动机上的测试准确率；不同颜色对应架构。
四：结果——收益出现在哪些条件下
4.1 逼近 softmax：把状态预算分给专家确有帮助
图 1C / S1 显示在 $B\ge2$ 的可比较预算上，最佳 SwiLA 分解的测试 $R^2$ 均高于同状态容量的 DeltaNet。最小的多专家例子是 $H=1,J=2$ 对 $H=2,J=1$；在 $B=4$ 时最佳为 $H=2,J=2$。这支持“头数和头内专家数提供互补容量”的解释。注意图中的 softmax 参考点能保存历史 KV，状态与计算形式不同；该实验测的是拟合能力，不是证明 SwiLA 与 softmax 等价。（Sec. 5.1，PDF p. 7；图 1C p. 3、图 S1 p. 27）
4.2 CDAR 与 RegBench：时间先验在情境持续时最有价值
CDAR 中，同一个键在不同情境必须指向不同值，而情境 token 与后来要回答的位置相隔很远。Temporal SwiLA 的一层、$H=1,J=4,D=32$、$2^{12}$ 状态配置已达到接近满分；无时间先验的 SwiLA 需要更大容量，图 2A 中一层 softmax 模型未解决该任务，两层才进入接近满分区域。附录图 S2 用跨学习率的平均准确率补充主图的“最佳值”口径，趋势仍突出 Temporal SwiLA。（Sec. 5.2，PDF p. 8；App. E.2，pp. 28–29）
RegBench 在仅 1,000 个训练自动机时，Temporal SwiLA 和 DN-MoM 测试准确率约 0.81、0.79，Transformer 约 0.53；训练规模增至 20,000 后，Transformer 约 0.95、居首，循环模型约 0.92–0.94。这里的优势是低数据样本效率，不是所有数据量都胜出。（Sec. 5.2、图 2B，PDF p. 8）
4.3 374M 语言模型：召回改善较稳，常识平均分差距很小
下表保留原表的平均值。常识列为八项准确率的平均（百分点，越高越好）；召回列为六项 contains 得分的平均（百分点，越高越好），所有召回输入截断为 2K token。它们不是同一指标，不能直接相减。（表 1–2，PDF pp. 9–10）
| 模型 | 常识平均准确率 % ↑ | 召回平均 contains % ↑ | 比较角色 |
|---|---|---|---|
| Transformer++ | 45.3 | 28.3 | 标准全注意力 |
| FoX | 44.6 | 33.7 | 带完整 KV 缓存的注意力变体 |
| DeltaNet | 43.7 | 17.3 | 单矩阵循环基线 |
| Gated DeltaNet | 45.0 | 19.2 | 带衰减的单矩阵基线 |
| KDA | 46.6 | 20.3 | 强门控循环基线 |
| SwiLA | 44.7 | 19.7 | 基础混合回归 |
| Temporal SwiLA | 45.2 | 20.3 | 时间先验 |
| Gated SwiLA | 45.4 | 20.5 | 状态衰减 |
| Gated Temporal SwiLA | 45.4 | 20.6 | 时间先验与衰减，粗体仅表示纯循环模型最佳 |
| Hybrid GDN | 44.3 | 29.5 | 18 层 GDN + 6 层全注意力 |
| Hybrid Temporal SwiLA | 45.1 | 31.8 | 18 层 Temporal SwiLA + 6 层全注意力 |
与 DeltaNet 的 17.3 相比，Gated Temporal SwiLA 的召回平均值 20.6 高 3.3 个百分点；与 GDN 的 19.2 相比高 1.4 个百分点，与 KDA 的 20.3 相比高 0.3 个百分点。混合层版本 31.8 高于 Hybrid GDN 的 29.5、Transformer++ 的 28.3，但仍低于 FoX 的 33.7。相对常识平均分，两个 Gated SwiLA 变体只比 Transformer++ 的 45.3 高 0.1 个百分点；KDA 的 46.6 更好。原表的 WikiText 困惑度亦显示 KDA 23.4 优于 Gated SwiLA 27.8（困惑度越低越好），因此不宜把 SwiLA 写成整体最优。（表 1–2，PDF pp. 9–10）
六项召回任务中也有不同形态：例如 SWDE 上 Temporal SwiLA 为 13.6、KDA 为 13.9；FDA 上 Temporal SwiLA 为 11.6、KDA 为 5.5。混合层版本的 SWDE/FDA 为 32.3/49.1，而 FoX 为 36.4/53.0。平均分总结了方向，任务级分数显示收益并不均匀。（表 2，PDF p. 10）
4.4 吞吐：更大批次带来生成优势，顺序训练仍是成本
单 H100 的 1.3B 模型训练测试固定每次更新 131,072 token，序列长从 2K 到 16K 时相应缩小批次。图 S3 中 SwiLA 约 19–22 千 token/s，GDN 约 40 千 token/s，且 SwiLA 的峰值分配内存约 55 GiB、高于 GDN 的约 42 GiB。生成测试使用 2,048 token 合成提示词和 128 token 贪心解码，计时包含 prefill 与 decode。Transformer 在较大批次 OOM，SwiLA 可运行到批次 128；作者概括其可达到比 Transformer 高约 4–5 倍的端到端吞吐。这个倍数反映单卡能容纳的不同批次，不是相同批次的延迟比较。GDN 在较大批次仍比 SwiLA 快，作者报告 SwiLA 推理约慢 1.25–1.5 倍。（Sec. 5.4，PDF p. 9；App. E.4、图 S3–S4，pp. 31–32）
原文图 S4，PDF p. 32。横轴为批次；上图为全流程 token/s，下图为峰值 GiB；叉号表示该模型在相应批次 OOM。
4.5 消融：先解决专家塌缩，再增加专家数
115M 模型消融中，普通 SwiLA、$J=4$、无负载均衡时测试困惑度为 36.769，加入辅助损失后为 16.807；Temporal SwiLA、$J=4$，不加/加均衡分别为 16.718/16.715。图 S5 中无均衡普通模型的键值侧失活专家比例一度超过 40%，而时间先验或辅助损失均减轻了塌缩。固定状态预算下，Temporal SwiLA 从 4 专家增至 16 专家，困惑度从 16.715 略升至 16.862：专家增加迫使每专家的键值投影变窄。由此，容量该分给多少专家与每个专家的维度，需要共同设计。（Sec. 5.5，PDF pp. 9–10；App. E.5、图 S5，p. 33）
原文图 S5，PDF p. 33。上排是查询侧、下排是键值侧；三列依次为总体使用熵、单 token 路由熵、失活专家比例。图例给出每个消融配置的测试困惑度。
五：结论与可继续研究的问题
SwiLA 的主要贡献是把线性注意力的“固定矩阵记忆”改成随输入软切换的多个回归记忆，并把专家写入权重与预测误差联系起来。Temporal 先验解决情境跨 token 保持，门控解决状态衰减，负载均衡解决专家使用不足。综合证据说明，这种结构在同循环状态预算下改善了若干检索问题；真实语言模型中提升最明显的是召回类任务，而最大常识平均分和最大纯召回平均分分别仍由 KDA 与 FoX 保持。（Sec. 3、5–6，PDF pp. 3–10）
实际取舍有两点。首先，固定状态无法像 KV 缓存那样随历史增长，长上下文可记忆量终有上限；混合层模型若加入全注意力，便不再享有全模型严格固定的记忆规模。其次，非线性专家责任让时间方向的训练计算目前是顺序执行，消耗吞吐与显存。附录 D 的分块 Newton 路线表明存在可探索的并行形式，但当前实验未得到系统加速。后续最直接的问题是如何让这套并行方案在真实预训练上同时保持数值精度和吞吐优势。（Sec. 6，PDF p. 10；App. D、E.4，pp. 23–26、31–32）