← 全部论文2507.16577v2 / r1
版本与下载

Scaling Linear Attention with Sparse State Expansion:中文研读笔记

  • 作者:Yuqi Pan、Yongqi An、Zheng Li、Yuhong Chou、Ruijie Zhu、Xiaohui Wang、Mingxuan Wang、Jinqiao Wang、Guoqi Li。
  • 版本:arXiv:2507.16577v2,修订于 2025-10-01;本次获取与研读日期:2026-11-02。
  • 原文:本地 PDF;arXiv PDF。下文页码均为从 1 开始的 PDF 页码。
  • 发表:已发表于 ICLR 2026 正式论文集,发表题名为 Scaling Linear Attention Capacity with Sparse State Expansion。下文的方法、公式编号与实验数字以 arXiv v2 为准。
  • 代码:作者已在 ICLR 官方补充材料 ZIP 公开 SSE-H 模型、配置和稀疏内核代码;该实现依赖外部工程组件,不是完整独立训练工程。
  • 配套汇报:英文幻灯片 PDF · 可编辑 TeX。

论文总览

SSE(Sparse State Expansion,稀疏状态扩展)要解决的是线性注意力的记忆如何组织、容量如何扩大。普通线性注意力把历史压进一个固定矩阵,每个 token 都向许多状态行写入信息;带遗忘门时,还会持续衰减这些行。作者把 key 映射解释为“将输入分到记忆行的分类器”,据此提出稀疏写入,并进一步把记忆扩展成多个共享投影参数的分区。每个 token 只读写少量分区,未被选中的分区连衰减也暂停。

这带来两个不同的收益:共享参数允许在近似固定模型参数量下增加状态容量;选择性更新使无关输入较少干扰已有记忆。SSE 的实际结构采用分区间硬 top-k、分区内 softmax 软分配,再加一个始终激活的共享分区。纯 SSE 使用线性递推;SSE-H 每五层 SSE 配一层 softmax attention。

最有说服力的证据是检索改善:2B、100B token 设置下,真实检索平均分从 GLA 的 49.29、GDN 的 53.79 提升到 SSE 的 61.46;混合 SSE-H 为 70.87,Transformer 为 73.00。数学推理方面,SSE-H 在 AIME24/25 达到 64.5/50.2,但同流程 Transformer 为 64.1/52.8,因此受控比较支持“混合线性架构可以保留较强推理能力”。摘要中相对公开小模型的大幅领先,还包含训练流程和数据的差异。(Table 1, p. 10;Table 5, p. 12)

结构地图

原文章节在论证中的作用PDF 页码与主要证据
1 Introduction固定状态的压缩瓶颈与两项贡献p. 2
2 PreliminaryKV cache 与固定矩阵记忆的对比pp. 3–4,式 (1)–(5)
3 Row-Sparse State Update: An Information Classification Perspectivekey 的分类解释、行相似性与硬选择pp. 4–6,Fig. 2–3,式 (6)–(8)
4 SSE: State Expansion under the Row-Sparse Framework分区、双向门控与并行实现pp. 6–9,Fig. 4–6,式 (9)–(14)
5.1 Language Modeling and Retrieval预训练、真实检索、RULER、扩长与转换pp. 10–12,Table 1–4
5.2 Reasoning Ability蒸馏与强化学习后的数学能力p. 12,Table 5
5.3 Ablation Study容量、softmax、门控与共享参数pp. 12–14,Fig. 7–8,Table 6–8
6 Limitations and Future Work更复杂状态转移、路由与扩容训练p. 14
7 Related Work;8 Conclusion与 MoM、状态转移研究和混合架构的关系pp. 14–15
9 Contributions作者贡献与单位p. 16
Appendix A Analyzing State Updating as Information Classification几何解释、遗忘分析、MQARpp. 21–25,Proposition 1–4,Fig. 9
Appendix B Pseudocode for Efficient SSE Implementationsmasking 与 varlenpp. 26–27,Algorithm 1–2
Appendix C Experiment Details配置、12B 转换、感受野、32K 检索pp. 27–28,Table 9–11,Fig. 10

一:动机——参数容量与上下文记忆容量是两件事

softmax attention 保存历史 key/value,单头缓存大小随长度 $t$ 增长,约为 $2td$ 个元素。线性注意力则维护 $\mathbf S_t\in\mathbb R^{c\times d}$:

$$ \mathbf S_t=\mathbf S_{t-1}+\mathbf k_t^\top\mathbf v_t, \qquad \mathbf o_t=\mathbf q_t\mathbf S_t. $$

这里 $c$ 是 key/query 的特征维度,也是记忆行数;$d$ 是 value 维度。固定 $c,d$ 后,每步状态读写成本与历史长度无关。代价是许多历史片段叠加在同一个小矩阵里。按相同精度、相同单头维度计数,$c=128$ 的 $cd$ 状态与 64 个 token 的 $2td$ KV cache 元素数相等;这只是存储量类比,不代表两者能保存等量可检索信息。(Sec. 2;Sec. 4, p. 6)

直接增大 $c$ 会扩大 Q/K 投影,可能同时增加参数和计算。SSE 希望把“学会如何读写”的参数,与“当前上下文能存多少”的状态容量分开。它借鉴稀疏专家的路由形式,但不同分区主要保存不同的动态上下文状态,共享 QKV 权重。相较 MoM,本文着重研究共享参数、softmax key 和读写双向门控共同形成的记忆组织方式。(Sec. 4.1;Sec. 7)

二:方法——先分类,再对选中的记忆推进时间

2.1 统一递推把“管理旧信息”和“写入新信息”分开

原文式 (6) 为:

$$ \mathbf S_t=\mathbf\Lambda_t\mathbf S_{t-1} +\phi(\mathbf k_t)^\top\mathbf v_t. $$

$\mathbf\Lambda_t$ 决定旧记忆如何保留或擦除,$\phi(\mathbf k_t)$ 决定新 value 如何分配到各行。GLA 采用输入相关的对角衰减;GDN 采用更有结构的 delta-rule 转移。SSE 的主要创新落在信息分配及状态分区上,实验中的 SSE 使用 GLA 风格转移。

若 $k_t^j$ 很大,第 $j$ 行就接收较强的 $\mathbf v_t$。于是可以把 $\mathbf x_t\mathbf W_k$ 看作分类打分,把状态行看作类别的聚合容器。这些类别是训练得到的特征区域,并没有人工语义标签。Fig. 2 用最大 key 激活为 value 向量着色,呈现聚类;它支持这种解释,但并未提供监督分类准确率。(Sec. 3.1, pp. 4–5)

2.2 从逐行硬选择到可实现的分区稀疏

概念上的行稀疏版本使用:

$$ \mathbf k_t=\operatorname{softmax} \bigl(\operatorname{top\!\text{-}\!k}(\mathbf x_t\mathbf W_k)\bigr). $$

这里应理解为:保留最大的 $k$ 个 logit,其余屏蔽为 $-\infty$,再对保留项归一化。它会产生真正的零写入。若还要获得“未选中记忆不遗忘”的性质,必须让对应的状态转移为恒等;仅把 key 置零仍会留下 $\mathbf\Lambda_t\mathbf S_{t-1}$ 的衰减。(式 (7)–(8), p. 6;Appendix A)

直接在全部扩展状态行上做 top-k 不利于硬件实现。因此最终 SSE 使用两层决策:先在 $N$ 个分区中硬选 $k$ 个,再在各选中分区的 $c$ 行间用普通 softmax 分配。后一步通常是近零稀疏,不是严格为零,也不是再做一次行内 top-k。(Sec. 4.1, pp. 6–8)

SSE 状态分区结构

原文 Fig. 4,p. 7。上半部是一块固定状态;下半部展示扩展后的多个分区。蓝色表示当前被路由选中的分区,绿色是始终选中的共享分区。读图关键是:空闲分区仍占存储空间,但本步不参与更新。

2.3 SSE 的逐步计算与符号

以下按单头描述;为避免与 value 维度混淆,用 $D$ 表示输入隐藏维度。

符号维度/含义
$\mathbf x_t$$1\times D$,当前隐藏表示
$\mathbf q_t,\mathbf k_t$$1\times c$,读出方向和行写入权重
$\mathbf v_t$$1\times d$,写入内容
$\mathbf S_t^i$$c\times d$,第 $i$ 个路由分区的状态
$\mathbf e_t$$1\times N$,分区路由概率
$\mathcal T_t$当前选中分区集合,大小为 $k$
$\mathbf\Lambda_t$$c\times c$,选中分区的状态转移;实验中为对角门

原文式 (9)–(14) 可连成以下过程:

$$ \mathbf e_t=\operatorname{softmax}(\mathbf x_t\mathbf W_e), \qquad \mathcal T_t=\operatorname{TopKIndices}(\mathbf e_t,k), $$
$$ \mathbf q_t=\mathbf x_t\mathbf W_q,\quad \mathbf v_t=\mathbf x_t\mathbf W_v,\quad \mathbf k_t=\operatorname{softmax}(\mathbf x_t\mathbf W_k), $$
$$ \mathbf S_t^i= \begin{cases} \mathbf\Lambda_t\mathbf S_{t-1}^i +e_t^i\mathbf k_t^\top\mathbf v_t,&i\in\mathcal T_t,\\ \mathbf S_{t-1}^i,&i\notin\mathcal T_t, \end{cases} \qquad \mathbf o_t=\sum_{i\in\mathcal T_t}e_t^i\mathbf q_t\mathbf S_t^i. $$

同一个 $e_t^i$ 在写入与读出两端都出现,使路由权重通过输出误差获得连续梯度。top-k 的离散索引本身不提供常规梯度;可训练的是保留下来的门值及其 softmax 参数。查询仍为线性投影,论文没有像 key 一样强制它非负。读写集合可以不同,默认相同。

一个需要按公式理解的细节:v2 p. 8 把选中行的总“分类强度”写成 1,但按式 (9)–(14),实际为

$$ \sum_{i\in\mathcal T_t}\sum_{j=1}^{c} e_t^i k_t^j =\sum_{i\in\mathcal T_t}e_t^i\leq1. $$

当 $k

2.4 稀疏更新为什么有助于长距离保留

定义 $\widetilde{\mathbf\Lambda}_u^i$:若分区 $i$ 在时刻 $u$ 被选中,取 $\mathbf\Lambda_u$;否则取单位阵。令初始状态为零,将递推展开,可以得到下面的教学推导,暂不写共享分区项:

$$ \mathbf o_t= \sum_{s\leq t}\sum_{i\in\mathcal T_t\cap\mathcal T_s} e_t^ie_s^i\, \mathbf q_t\left(\prod_{u=s+1}^{t}\widetilde{\mathbf\Lambda}_u^i\right) \mathbf k_s^\top\mathbf v_s. $$

乘积按时间递推顺序排列;在论文采用的对角转移下顺序不影响结果。这个式子说明三个机制:历史与当前 token 必须路由到共同分区才会发生直接读写关联;双向门值控制关联强度;遗忘只在该分区被访问时累积。

例如某条记忆经过 100 个 token,普通固定衰减 $a=0.99$ 留下 $0.99^{100}\approx0.366$;若它所在分区期间只激活 10 次,则相应衰减为 $0.99^{10}\approx0.904$。这是解释“局部更新时钟”的教学例子,不是论文实验数据。实际保留还取决于新增写入的干扰及查询是否命中该分区。

2.5 共享分区、负载均衡和容量账目

为了保留连续 token 的局部交互,SSE 加入一个始终激活的共享分区,并对其 Q/K 使用秩 64 的 LoRA,V 继续共享。按 Sec. 5.3 的 $N+1$ 分区描述,n4k1 包含四个路由分区和一个共享分区,每步访问一个路由分区加共享分区。因此 $k/N=1/4$ 是路由部分的激活比例。(p. 8;p. 14;Appendix C, p. 27)

原文行级均衡项为

$$ \mathcal L_{\rm balance}=\alpha\frac{c}{k} \sum_{i=1}^{c}f^i k_t^i, $$

其中 $f^i$ 是选择频率。SSE 改为分区级均衡,$\alpha=0.01$,意在避免大量 token 涌入少量分区、使其他状态难以训练。均衡项在训练时使用,推理时直接路由。

以单头计,路由状态大小是 $Ncd$,计入共享分区约为 $(N+1)cd$;一次读写的主状态运算约为 $(k+1)cd$。这两项必须分开:稀疏访问减少活跃计算,并不消除未激活状态的存储成本。此外,路由投影有约 $DN$ 个参数,top-k、分组和内存访问也有成本,因此“参数和开销近似不变”适用于作者所测配置,不应解释成随 $N$ 增加绝对零成本。(Sec. 4.1–4.2;本段复杂度为按结构计数)

2.6 并行实现保留了分区内部的因果顺序

短序列:masking。 复制 QKV 到分区维,屏蔽不属于该分区的 token,再把分区维合并到 batch/head 维,并行调用线性注意力算子。这样避免产生大量极短的分区子序列;chunk 算子通常在 64、128 等块长上较高效。

长序列:varlen。 按“样本—分区”收集 token,在各分区内保留原时间顺序,为子序列建立新的 cu_seqlens。路由部分共约 $kL$ 条 token 记录,可用一次变长算子调用处理所有分区,最后按原位置聚合输出。共享分区也能融合进去。解码阶段只对选中状态做递推,prefill 则继续用分块计算。(Sec. 4.2, pp. 8–9;Appendix B, pp. 26–27)

实现时有两项不能从示意伪代码里省掉:未选中分区的衰减要变为恒等;分组后的输出要散射回原 token 并进行门控求和。原文 Algorithm 1–2 主要展示调度框架,并未完整展开这两项。

SSE 算子不同实现的耗时

原文 Fig. 6,p. 9。横轴为序列长度,纵轴为前向加反向总时间(ms);蓝、绿、红分别对应逐分区调用、masking、varlen,不同线型对应 $N=4,8,16$。测试使用 A100、8 heads、head dimension 128、$K=1$、cu_seqlens=[0,L/2,L]。短序列 masking 有竞争力,长序列 varlen 更好;固定 $K$ 时 varlen 的三条曲线较接近。该图比较 SSE 算子实现,不能直接读作整模型相对 Transformer 的训练或生成加速比。

2.7 理论支持及其具体适用条件

Appendix A 的论证链为“同类输入形成几何簇 → 跨类混写使行趋同 → 读出区分度下降 → 无差别衰减限制保留跨度”。应将其中条件与实证分开理解。

  1. Proposition 1:同类输入落在同一个锥内。 假设 $\|\mathbf x_s\|^2=d$,并按 $\mathbf x_s\mathbf W_k^i>k_{th}^i$ 分类。定义 $\theta^i=\arccos(k_{th}^i/(\sqrt d\|\mathbf W_k^i\|))$,同类向量两两夹角小于 $2\theta^i$,故内积大于 $d\cos(2\theta^i)$。阈值足够高、不同分类方向足够分离,才提供有意义的类间分离。命题措辞里的反向关系不能脱离这些条件单独使用。(p. 21,式 (17)–(19))
  2. Proposition 2:混写增加行相似性。 推导使用正交 $\mathbf W_v$、非负混合系数和类簇分离,主要展开二类情形。它解释共享噪声为什么可能令行同质化;实际网络的任意 value 投影、任意有符号状态不自动满足全部前提。(pp. 22–23,式 (20)–(34))
  3. Proposition 3:读出区分度的下界。 作者在非负查询条件下讨论 $\min_{\mathbf p,\mathbf q}\cos(\mathbf q\mathbf S,\mathbf p\mathbf S)$。v2 的命题陈述把下界直接写成最小行内积,而证明式 (36)、(39) 含有状态范数归一化,二者不能等同。一个直接的量纲检查是:把所有相同行都设为 $(2,0)$,非零非负查询的输出余弦恒为 1,行内积却为 4,后者不能直接作为下界。实际 SSE 的线性 query 也不保证非负。因此这里提供的是机制解释,不能照抄成实际 SSE 的无条件保证。(pp. 23–24;反例为本次分析)
  4. Proposition 4:减少衰减次数延长保留。 在 $\mathbf q_t,\mathbf k_s$ 全为 1 的简化下,$p_{ts}=\sum_j\prod_{u=s+1}^{t}\alpha_u^j$。若全程每次都衰减,达到固定阈值的跨度受对数界约束;若有通道在整个间隔不更新、衰减恒为 1,就能跨越任意长度的空闲间隔保留其贡献,条件是阈值不超过保留的贡献。论文的“无限感受野”应理解为这种存在性保留情形,而不是所有记忆都永不遗忘。(pp. 24–25,式 (40)–(43))

三:实验——不同训练预算分别回答什么

3.1 模型、预算与比较角色

模型采用 MHA–SwiGLU,并去掉卷积层,以突出注意力机制。600M 模型为 24 层、隐藏维度 1024、8 heads,非嵌入参数 300M;2B 模型为 18 层、隐藏维度 2304、18 heads,非嵌入参数 1.3B。2B 混合模型保留 3 层 softmax attention,其余 15 层为线性层。(Table 9, p. 27;Sec. 5, p. 10)

实验阶段数据/上下文回答的问题
小规模预训练600M/15B token;2B/100B token;最长 8K相同基础训练下,SSE 的语言和检索能力如何?
数据扩展及扩长2B SSE-H 与 Transformer 预训练到 2T token,再用 250B token 扩到 32K优势能否保留到更充分训练和更长上下文?
2B Transformer 转换128K 长上下文训练阶段,约 100B token;评测长度 32K用 SSE 替换部分注意力层后能保留多少检索能力?
推理后训练约 80K 蒸馏样本、5 epochs;GRPO 230 步;每 prompt 采 8 条,最长生成 32K混合架构能否支撑数学推理训练?
状态扩容消融2B/20B token分区数 $N$ 与激活数 $k$ 如何影响召回?
key/参数共享消融600M/15B tokensoftmax key 与共享参数各自贡献多少?
双向门消融2B/100B token读、写两端的路由权重是否都重要?

预训练使用 AdamW,总 batch 为 4M token,线性 warm-up 后学习率 $6\times10^{-4}$,weight decay 0.1,gradient clipping 1.0;扩长时使用 cosine learning-rate decay。12B 转换实验另见结果部分,不能套用 2B 的五比一层结构。

baseline 的角色也不同:GLA 对照相同风格的对角状态转移;GDN 对照更强的 delta-rule 更新;MoM 对照带更多参数的稀疏多记忆方法;GLA-H 对照混合本身的收益;Transformer 提供未压缩历史的参考。MoM 的总参数为 800M、活跃参数 600M,且评测不做 token truncation,因此不能把它简称为与 600M SSE 完全等参数。

3.2 指标如何读

  • Table 1 的 Wiki.↓ 为原表语言建模指标;常识平均分来自 PIQA、HellaSwag、WinoGrande、ARC-e、ARC-c、SIQA 六项,越高越好。
  • FDA、SWDE、SQuAD 被用作真实检索任务,Recall-Avg. 为三项等权平均。v2 未展开三者的评分器细节,以下沿用原文分数口径。
  • RULER 的 S-NIAH 衡量在长上下文中定位单条目标信息;多针任务进一步增加 key、query 或 value 的组合负担。表中为百分制表现,越高越好。
  • AIME24/25 的 avg@32 是把评测集重复采样 32 次后求平均正确率;它不是“32 次中有一次成功”的 pass@32。采样温度为 1.0,top-p 为 0.7。其他数学列沿用原表报告,不能全部改标 avg@32。(Sec. 5.2, p. 12)

四:结果——收益主要体现在哪里

4.1 Sec. 5.1:基础预训练的检索改善大于语言指标变化

以下选择 Table 1 的聚合列,保留不同参数预算。所有常识与检索分数越高越好。

模型与预算Wiki.↓常识平均↑检索平均↑
Transformer,600M/15B token33.4742.2255.95
GLA,同预算43.9741.5318.63
GDN,同预算35.2643.0524.84
SSE-n4k1,同预算35.5242.9131.16
SSE-n4k2,同预算35.3343.0531.86
MoM-n4k1,800M 总/600M 活跃,15B token36.0542.6427.69
MoM-n4k2,同 MoM 预算35.1142.7828.27
Transformer,2B/100B token16.4653.5573.00
GLA,同 2B 预算21.3049.1349.29
GDN,同 2B 预算17.0854.0553.79
SSE-n4k1,同 2B 预算16.9254.5761.46
GLA-H,同 2B 预算17.2753.1567.52
SSE-H-n4k1,同 2B 预算16.4754.4870.87

来源:Table 1, p. 10。2B 纯 SSE 相对 GDN 的检索平均分增加 7.67 个百分点,相对 GLA 增加 12.17 个百分点;与 Transformer 仍差 11.54 个百分点。混合后与 Transformer 的差距缩小到 2.13 个百分点。常识平均分已接近或略高于 Transformer,说明这一阶段的主要短板确实集中在精细检索。

600M 的 SSE-n4k1 相对 MoM-n4k1 以较少总参数取得 31.16 对 27.69 的检索分,但不能从这张表推导所有语言指标都更好。原文叙述中“所有任务都超过 GDN”也有例外:600M 的 SWDE 是 SSE 33.30、GDN 33.48;以表中逐项结果为准。

4.2 Sec. 5.1:8K 单针检索显示压缩误差的差别

2B/100B token 模型S-NIAH-1,8KS-NIAH-2,8KS-NIAH-3,8K
Transformer100.0100.0100.0
GLA87.623.216.2
GDN100.08.29.0
SSE-n4k1100.085.262.2
GLA-H100.0100.070.8
SSE-H-n4k1100.0100.097.4

来源:Table 2, p. 11,零样本;原表还给出 2K、4K。最简单的任务很多模型都满分,较难两类才拉开差距。SSE 在 S-NIAH-3 的 2K/4K/8K 为 100.0/88.2/62.2,仍会随长度下降;SSE-H 对应 100.0/99.4/97.4,表明少量完整注意力层补足了压缩状态的检索损失。

4.3 Sec. 5.1 与 Appendix C:更充分训练后,混合模型整体接近

预训练到 2T token、再用 250B token 扩到 32K 后,Table 3 的十项平均分为 SSE-H 45.6、Transformer 45.1,增加 0.5 个百分点。MMLU 从 52.6 到 54.5,MMLU-Pro 从 24.2 到 26.1,C-Eval 从 55.9 到 59.7;SWDE、SQuAD、DROP、GSM8K 则略低,例如 GSM8K 为 49.2 对 50.6。这是任务间有得有失的整体接近,不是全面占优。(p. 11)

扩长后 RULER 的 32K 列如下:

任务TransformerSSE-H-n4k1
S-NIAH-199.493.6
S-NIAH-270.884.4
S-NIAH-367.490.2
MK-NIAH-155.464.6
MQ-NIAH28.644.6
MV-NIAH18.140.6

来源:Table 11, p. 28。SSE-H 在六类中的五类高于此 Transformer baseline,但最简单单针略低,多 query/value 的绝对分数也仍不高。Table 2 与 Table 11 的训练预算和上下文阶段不同,不能把前者的 Transformer 100 分当作后者的对应控制。

4.4 Sec. 5.1:从已有 Transformer 转换,比从零训练更贴近迁移需求

在 2B Transformer 的 128K 长上下文训练阶段替换部分层、继续训练约 100B token,32K 真实检索结果为:

模型QuestQampariTable Query平均
原 Transformer49.171.560.660.4
转换 GLA-H25.656.648.243.5
转换 SSE-H-n4k135.670.650.052.1

来源:Table 4, p. 12。SSE-H 比 GLA-H 高 8.6 个百分点,但比原 Transformer 低 8.3 个百分点。Qampari 保留较好,Quest 仍损失 13.5 个百分点;“支持转换”应结合这种保真度代价理解。

4.5 Sec. 5.2:数学推理要区分公开模型排名与受控架构比较

模型AIME24AIME25MATH500OlympiadBenchAMC23
Qwen3-1.7B Thinking48.336.893.4——
DeepSeek-R1-Distill-Qwen-7B55.539.292.8——
作者 Transformer-2B64.152.893.083.392.0
作者 SSE-H-n4k1-2B64.550.292.185.791.4

来源:Table 5, p. 12;AIME 两列为 avg@32。相对 Qwen3-1.7B,SSE-H 的 AIME24/25 分别高 16.2/13.4 个百分点,但公开基线并非相同训练数据与后训练预算。相对作者同流程 Transformer,差值为 +0.4/−2.6 个百分点,MATH500 为 −0.9,OlympiadBench 为 +2.4,AMC23 为 −0.6。合理的核心结论是:SSE-H 能承接蒸馏及 RL,获得与强 Transformer 接近的数学能力。

12B 转换结果也支持这种接近:AIME24 为 SSE-H 74.3 对 Transformer 75.7,AIME25 为 57.3 对 58.7;MATH500 为 96.0 对 95.8,OlympiadBench 为 85.3 对 82.7,AMC23 为 96.3 对 97.2。这里使用 SSE–SWA–SSE–SWA–Softmax 的层序列,在 128K 阶段转换后进行监督蒸馏,包含滑动窗口注意力,不能当作 2B 架构的简单等比例放大。(Table 10, p. 27)

4.6 Sec. 5.3:扩容有效,但不同曲线的计算预算也在变化

状态分区数量与召回

原文 Fig. 7,p. 13;2B 模型,20B 训练 token。横轴为分区数 $N$,纵轴为真实任务平均召回;红线 $k/N=25\%$,蓝线 $k/N=50\%$。两条线均随 $N$ 增长而上升,同一 $N$ 下激活更多分区也更好。这里固定的是激活比例,所以增大 $N$ 同时增大了 $k$ 和每步活跃计算;不能把这张图与 Fig. 6 中固定 $K=1$ 的耗时结果拼成“零计算增量扩容就获得同样提升”。图展示有限配置的趋势,未给出可外推的尺度定律拟合。

原文称 $k/N=1$ 时“退化为 vanilla linear attention”。准确理解是失去跳过分区的稀疏优势;若仍保留不同的 token 路由权重,多分区状态并不由该条件自动代数等价为原来的单一状态。(Sec. 5.3, p. 13;本句为按式 (13)–(14) 的分析)

4.7 Sec. 5.3:softmax key、双向门和参数共享都有独立贡献

softmax key。 600M/15B token 下,SiLU 改为 softmax 后,FDA/SWDE/SQuAD 从 20.87/26.91/24.90 变为 33.67/33.30/26.51,平均从 24.23 到 31.16,即 +6.93 个百分点。(Table 6, p. 13)

双向门。 在 Table 7 自身的 2B/100B token 消融组内:

门值应用位置检索平均常识平均
写入和读出56.6353.63
都不使用51.8753.01
仅写入50.9252.82
仅读出53.9652.91

双向门比仅读出提高 2.67 个百分点检索,比仅写入提高 5.71 个百分点。这里的完整模型分数 56.63 不同于 Table 1 的 61.46;原文没有解释这两组名义相同规模与 token 数下的差别,因此只在 Table 7 内计算消融效应。(p. 14)

共享参数。 600M/15B token 的共享模型检索分 31.16;每个分区独立参数后,总参数增到 890M、非嵌入参数从 300M 增到 580M,检索反而降到 25.95,常识仅从 42.91 到 42.99。这个实验支持“该设置下,增加独立投影参数不能替代更好的记忆组织”,并未建立任意规模下参数容量都不重要的结论。(Table 8, p. 14)

4.8 状态诊断与 MQAR:机制解释有对应观测

不同分类器的状态行相似性

原文 Fig. 3,p. 5。横纵轴对应状态行,颜色表示行间余弦相似性;越蓝相似性越低。线性分类器的状态行高度同质,softmax 和 SSE 显示更鲜明的对角结构,支持“信息分配更有区分度”的解释。

Fig. 5(p. 8)进一步计算归一化奇异值能量熵:

$$ p_i=\frac{\sigma_i^2}{\sum_j\sigma_j^2}, \qquad H=-\frac{1}{\log n_\sigma}\sum_{i=1}^{n_\sigma}p_i\log p_i. $$

这里 $n_\sigma$ 表示奇异值数量,以免与分区数 $N$ 混淆。较高熵表示能量分散在更多奇异方向;SSE 的整体平均高于 GLA,说明状态较少集中在少数方向。这是表示诊断,不是直接以 bit 计量的记忆容量。

Fig. 8(p. 13)把 key 中小于 $10^{-5}$ 的通道视为近零:路由分区的平均比例约 41.6%,共享分区约 4.9%,GLA 接近 0。它说明共享分区趋向密集保存局部信息,路由分区更具选择性;并不表示 softmax 本身产生了 41.6% 的数学零值。

MQAR 上的行稀疏更新

原文 Fig. 9,p. 25。三列分别为 GLA、GSA、MetaLA,上下两行分别在 key-value 与 query 之间加入零 token、随机噪声 token;横轴为模型维度 64–512,纵轴为召回准确率。实线是 top-k 版本,虚线是原模型。实验为两层、两头,$k=d/4$。稀疏版本通常表现更高且随维度增长更稳定,尤其能减轻噪声插入带来的干扰。

Appendix C 的 Fig. 10(p. 28)对 8K 训练的 2B 模型,用最后 128 个 token 和门值阈值 0.001 估算各层感受野,SSE 各层均值普遍高于 GLA;阴影表示标准差。该量直接衡量门衰减路径的有效跨度,与检索准确率是不同测量。

五:结论——稀疏访问把状态容量变成可独立调节的资源

SSE 的贡献可以沿一个完整因果设计链理解:key 的分类视角 → 减少无关写入与无关衰减 → 多分区共享参数扩容 → 用变长并行算子把稀疏结构转化为计算收益。实证中,检索是最突出的改善方向;少量 softmax 层进一步提升精细召回,使 SSE-H 在充分训练和推理后训练后接近同流程 Transformer。

应用上需要同时考察三种资源:模型参数、持久状态存储、每 token 活跃读写计算。增加 $N$ 主要扩大存储;增加 $k$ 则增加活跃计算;SSE-H 的 softmax 层仍有随上下文增长的 KV cache,整模型不再具有纯递推模型的严格常数上下文内存。Fig. 6 的算子结果也尚不能替代端到端吞吐与显存测量。

后续最值得检验的是:在固定活跃计算下怎样选择更大的状态容量;用包含历史状态的路由 $\mathbf e_t=g(t,\mathbf x_t,\mathbf S_{t-1})$ 改善“写入分区与未来查询分区不匹配”;以及把 SSE 与 delta-rule 转移结合。v2 已将 GDN 作为 baseline,但尚未充分评估 SSE 加在 GDN 转移上 的组合。作者还提出从已有 checkpoint 扩容的 up-training 方向;这些是后续工作,不能视作本文已经完成的实验。(Sec. 6, p. 14)