2.4 · 训练原生稀疏: 让模型在稀疏图上学习

稀疏 attention 有两种完全不同的使用方式. 第一种先训练全注意力模型, 部署时再删除一部分连接、KV 或 attention blocks. 它能直接使用已有 checkpoint, 也要承受训练函数与推理函数之间的差异. 第二种把稀疏计算图放进优化过程, 模型从训练中就知道哪些历史会被压缩、哪些候选会被 top-k 截断、哪些局部位置始终可见. 训练原生稀疏的核心是「模型在部署时使用的稀疏图上学习」, 起点可以是随机初始化, 也可以是已有全注意力 checkpoint.

Native Sparse Attention(NSA)是这条路线最完整的例子. 它把压缩、选择和滑窗做成三条分别归一化的attention分支, 从前向、反向一直覆盖 Prefill 与 Decode. MoBA 从另一侧进入: 每个 query 路由到少量 KV blocks, 结构不增加 attention 参数, 因而可以在 full attention 与 MoBA 之间切换. DSA 又采用两阶段转换: 先保留 dense attention 训练 Lightning Indexer, 再让主模型在 token 级稀疏 attention 上继续训练. 三者都训练稀疏结构, 候选来源、梯度路径和硬件粒度并不相同.

判断一种训练原生稀疏方案, 要同时追踪五件事: 稀疏图何时进入优化过程, 多路 attention 怎样共同更新, hard top-k 截断哪些梯度, token 预算怎样落到 kernel 与 KV 读取, 以及从头预训练和继续训练各自承担多少适应成本. NSA 解析列出了具体 block 参数、论文实验和 kernel 流程; MoBA 解析继续拆解块均值路由、因果约束与 full/sparse 混合; NSA、MoBA 与 DSA 则可以放进同一组符号中比较训练路线.

1. 分界落在计算图与优化过程

1.1. 后插稀疏改变已经学好的函数

令全注意力在位置 $t$ 的输出为

o_t^{\mathrm{full}}= \sum_{s\in V_t}\alpha_{t,s}v_s, \qquad \alpha_{t,s}= \frac{\exp(q_t^\top k_s/\sqrt d)} {\sum_{j\in V_t}\exp(q_t^\top k_j/\sqrt d)}. \tag{1}

$V_t$ 是因果可见历史. 部署阶段选一个集合 $S_t\subset V_t$ , 再把 softmax 分母限制到 $S_t$ , 输出变成

o_t^{\mathrm{sparse}}= \sum_{s\in S_t}\widehat\alpha_{t,s}v_s, \qquad \widehat\alpha_{t,s}= \frac{\exp(q_t^\top k_s/\sqrt d)} {\sum_{j\in S_t}\exp(q_t^\top k_j/\sqrt d)}. \tag{2}

即使 $S_t$ 覆盖了大部分原始 attention mass, 分母变化也会重新缩放所有保留位置. 一个低概率 token 若是随机字符串、变量定义或工具返回值的唯一来源, 漏选造成的影响远高于它在平均 attention 中的占比. 全注意力训练也没有压力把信息写入确定会保留的窗口、压缩状态或全局 token.

后插路线依旧有价值. KV eviction、Quest、MInference 一类方法能在没有大规模训练预算时改造现有模型, 其中一些还保留精确 attention 计算, 只近似候选集合. 它们节省的是指定部署阶段, 训练原生路线则把结构限制带入参数学习. 两类方法应分别比较 checkpoint 迁移成本、训练节省、推理节省和质量恢复成本.

1.2. 「原生」要求四个阶段共享结构契约

一个稀疏层可以写成候选构造器 $R_\theta$ 与主 attention $A_\phi$ 的组合:

S_t=R_\theta(q_t,K_{:t}),\qquad o_t=A_\phi(q_t,K_{S_t},V_{S_t}). \tag{3}

评估训练原生稀疏时, 可以沿前向、反向和部署检查三个环节. 前向确实在 $S_t$ 上计算主 attention, 训练 loss 观察到候选遗漏后的输出; 反向不依赖完整 $t\times t$ attention map 才能完成; 部署时的候选粒度、预算、因果规则与训练一致. Prefill 和 Decode 可以使用不同 kernel, 数学上的 $R_\theta$ 与 $A_\phi$ 仍应一致. 「使用同一结构」也不要求每个训练 step 都保持固定稀疏率. 课程训练可以先用 full attention, 再逐步减小候选预算; 某些层可以保留 full attention; 短序列可以退化到 dense kernel. 关键在于模型经历过准备部署的稀疏图, 并且最终优化阶段覆盖部署配置. MoBA 的 full/sparse 切换与 DSA 的 dense warm-up 都符合这一点.

反向算子的存在是硬条件. 前向只读取 $k$ 个 blocks, backward 若重新物化所有 QK 分数, 预训练仍支付平方成本. NSA 为 selection attention 设计 forward 与 backward kernel, 并把连续 KV blocks 与 GQA group放进同一执行约束. 训练速度要报告完整 step time、峰值显存和重算开销, 单独的 forward latency 无法代表原生训练收益.

NSA、MoBA 与 DSA 位于三种位置。 NSA 从架构层面定义三路输出. 压缩分支覆盖全局低分辨率历史, selection 分支回读少量连续原始 blocks, sliding-window 分支承担局部模式, 三路通过可学习 gate 融合. 压缩 attention 的概率还生成 selection importance. 它同时改变 attention 表示、候选生成和 kernel, 适合从预训练开始共同学习. MoBA 保留普通 Q/K/V 参数, 用每个 KV block 的 mean-pooled key 给 query 打分, top-k 后在选中 blocks 内运行精确 attention, 当前 block 强制进入并使用 causal mask. 官方实现明确要求继续训练, 无法直接把预训练模型换成 MoBA 后立即获得可靠质量. 它又能把 $k$ 扩到全部 blocks 退化为 full attention, 因而适合渐进转换和 full/sparse 混合训练.

DSA 在 MLA 主 attention 之外增加低维 Lightning Indexer. Dense warm-up 阶段保留完整主 attention, 冻结主模型, 用聚合后的主 attention 分布监督 indexer. Sparse training 阶段启用 top-k token selection, 再优化整个主模型适应稀疏图. Indexer 输入从主计算图 detach, indexer 只接受自己的 KL loss, 主模型只接受语言模型 loss. 这是一种训练原生的转换方案, 其主 attention 与 selector 使用两条明确分开的梯度路径. 三条路线可以按「稀疏从哪里进入训练」区分. NSA 让稀疏成为 attention layer 的原始结构; MoBA 保留参数形状, 在 attention mask 与 query-to-block assignment 上进入; DSA 保留 MLA 主干, 用新增 indexer 改变主 attention 的 token集合. 因而「是否增加参数」与「是否训练原生」没有必然关系. MoBA 不增加参数但需要继续训练, DSA 增加 selector并做两阶段训练, NSA 连分支投影都从预训练中学习.

路线	候选单位	稀疏训练信号	与 full attention 的切换	主要固定通路
NSA	连续 blocks	LM loss 经三路输出, coarse probability 同时用于 route	需要架构级转换	compression 与 window
MoBA	KV blocks	LM loss 适应 hard routed attention	$k$ 覆盖全部 blocks 时恢复 full	当前 block
DSA	单个 tokens	Indexer KL + 主模型 LM loss	dense warm-up 后进入 sparse stage	由 main MLA 与训练配方决定

训练成本的比较也要统一 token口径. 从头 sparse预训练节省大量长序列 attention, 但若为了达到同等质量多训练更多 tokens, 总 FLOPs可能回升. 继续训练路线保留已有 checkpoint价值, 转换阶段本身仍是成本. 报告应同时列出起始 checkpoint、转换 tokens、最大长度、有效训练 FLOPs和最终质量, 单看每 step加速会忽略适应成本.

支持集是模型的一部分。对位置 $t$ , 稀疏 attention 先从合法历史 $V_t$ 中选出支持集 $S_t\subseteq V_t$ , 再在 $S_t$ 上计算主 attention. 写成

S_t=\operatorname{Select}(x_{\le t};\theta_r), \qquad y_t=\operatorname{Attn}(q_t,K_{S_t},V_{S_t};\theta_a).

$\theta_r$ 是路由参数, $\theta_a$ 是主 attention 参数. 稠密模型固定 $S_t=V_t$ ;训练原生稀疏把 $S_t$ 也纳入学习过程. 支持集决定主模型能看到哪些变量, 它既是离散结构, 也是梯度传播的边界. 如果目标位置 $s^$ 没进入 $S_t$ , 主 attention 对 $K_{s^$ },V_{s^} $K_{s^{*}}, V_{s^{*}}$ 的直接梯度为零. 训练信号只能通过路由辅助损失、连续代理或其他层的信息路径告诉 selector「这里漏了」. 若这些通道也没有覆盖 $s^$ $s^{*}$ , 当前错误不会自动教会路由器把它选进来. 这与普通参数初值不好不同: 未选位置连局部导数入口都不存在.

因此可学习性至少包含三层. 表达层要求存在某组路由参数能把必要位置放入预算;信号层要求漏选时有梯度或探索机会;优化层要求训练动态不会过早收缩到局部捷径. 三层中任一失败, 模型都可能拥有理论容量却学不到正确稀疏图. 预算 $k$ 还定义了结构上限. 若某一步需要同时读取 $m>k$ 个相互独立的 values, 单个 sparse head无法全部纳入支持. 多 heads、coarse summary、跨层积累或更大 $k$ 可以补偿, 但支持容量不会由训练技巧突破. 教师attention top-k在预算 $k$ 下仍失败, 只能说明这一候选规则未能完成任务. 它最大化保留的教师概率, 未必保留任务最需要的值; 还要比较其他候选、证据覆盖和传播路径, 才能区分预算、候选与主模型问题.

离散 top-k 为什么没有普通梯度。令 selector score 为 $z_s$ , hard top-k mask 为

m_s=\mathbf1[s\in\operatorname{TopK}(z,k)].

只要 score 的微小扰动没有改变排序边界, $m_s$ 完全不变, 对 $z_s$ 的导数为零. 第 $k$ 与第 $k+1$ 名交换时 mask跳变, 该点又不可导. 因而从主任务 loss经过 hard mask反传, 常规链式法则无法给未跨边界的 selector score提供有效梯度. 已选 token 的主 attention score仍有梯度, 但这只优化集合内权重. 若 selector 与主 attention共享 query/key参数, 集合内梯度会间接改变未来 selector score;它没有直接说明哪个集合外 token应该进入. 共享参数提供弱耦合, 不能代替明确路由监督. soft relaxation把 mask换成连续权重, 例如

\pi_s=\frac{\exp(z_s/\tau)}{\sum_j\exp(z_j/\tau)}.

温度 $\tau$ 大时分布平滑, 非零目标或代理信号可以作用于各位置, 训练计算接近稠密;温度变小时接近离散选择, 梯度集中在少量边界位置并可能饱和. 训练用全量 soft权重、推理用hard top-k还存在支持不一致, 模型可能依赖许多小权重的合计贡献. straight-through estimator前向使用 hard mask, 反向假装经过某个连续函数. 它的前向执行真实hard mask, 反向使用有偏代理梯度. 代理梯度是否指向更好的支持集没有普遍保证, 需要通过 oracle、探索和任务曲线验证. Gumbel噪声等随机松弛可以采样离散候选, 方差和温度调度又成为训练变量.

另一条路线用教师分布训练 selector. 稠密 attention、高预算路由或 coarse branch产生目标 $p_s$ , selector最小化 KL、排序或覆盖损失. 这样未选位置只要教师给质量就有梯度. 教师计算增加训练成本, 可以只在部分层、部分 batch或短片段启用. 教师自身若偏向局部捷径, 学生会继承.

原生稀疏不等于从第一步就用最硬预算。从头训练时 hidden与selector都未形成, 立即使用很小 $k$ 会让早期随机集合决定哪些参数得到梯度. 某些位置偶然被选后优化更快, 又更容易继续入选, 形成自增强. 训练可以从较大预算、较高温度或更多随机探索开始, 随表示稳定逐步收紧. 设训练步 $u$ 的预算为 $k(u)$ , 一个简单日程从 $k_0$ 降到目标 $k_1$ . 大预算阶段提高必要位置进入支持的概率, 主模型先学会使用长程信息;收紧阶段迫使 selector保留最有价值部分. 若收紧太晚, 模型长期依赖稠密支持, 目标稀疏图适应不足;太早则回到随机锁定.

预算退火与学习率相互作用. 支持结构变化会突变每层输入分布, 在大学习率阶段可能持续重排, 在学习率很小时又来不及适应新预算. 可在预算变化后保留一段稳定训练, 观察路由 recall和主 loss重新收敛. 单纯按总 token百分比设日程没有跨模型通用值. 探索也可保留固定比例随机候选. 令支持由 exploitation top-k 与 exploration随机 $r$ 个位置组成. 随机位置为集合外参数提供被主任务看见的机会, 同时消耗预算并加入噪声. $r$ 随训练下降, 类似从探索转向利用. 对唯一远距证据, 均匀随机命中概率约为 $r/N$ , 长上下文下可能很低, 分块或教师引导探索更有效.

原生训练的关键是目标结构贯穿足够长训练, 让主模型与路由器共同适应. 允许早期软化、宽预算和教师辅助不违背这一点;只在训练末尾临时套一个稀疏 mask, 适应 token不足, 才更接近后插转换.

2. 压缩、选择与滑窗怎样共同训练

2.1. 三个分支各自归一化再融合

设分支集合 $\mathcal C={\mathrm{cmp},\mathrm{slc},\mathrm{win}}$ , 每条分支拥有自己的 K/V 投影与可见集合 $S_t^c$ . NSA三路共享当前query, 此处统一记作 $q_t$ . 分支输出为

o_t^c=\operatorname{Attn} \left(q_t,K_{S_t^c}^c,V_{S_t^c}^c\right), \qquad c\in\mathcal C. \tag{4}

NSA 使用 MLP 与 sigmoid 产生分支 gate $g_t^c$ , 最终输出写成

o_t=\sum_{c\in\mathcal C}g_t^c o_t^c. \tag{5}

每条分支有自己的 softmax 分母, sigmoid gate也不要求三路权重之和为1. 同一原始位置可能通过压缩表示、selected raw KV 和 window raw KV 多次影响输出, 三个概率不能直接拼成一张统一 attention map. 把三路 K/V 合并后做一次 softmax 会改变式 (5), 训练目标和 kernel也随之改变. 独立 K/V 投影给三个分支分配不同表示空间. Window 分支可以专注邻近词法和局部状态; compression 分支学习把一个 block 可用的信息聚合进单个 entry; selection 分支保留精细 token. 共享全部投影能减少参数, 也会让局部梯度与远距压缩梯度在同一 K/V 上竞争. NSA 选择独立投影来降低这种干扰.

2.2. Coarse 分支既输出内容也产生路由信号

令 compression block 长度为 $l$ , stride 为 $d$ , 从 $r=0$ 开始编号的第 $r$ 个完整重叠block 的表示为

\widetilde k_r^{\mathrm{cmp}}= \phi_K(x_{rd+1:rd+l}),\qquad \widetilde v_r^{\mathrm{cmp}}= \phi_V(x_{rd+1:rd+l}). \tag{6}

$\phi_K,\phi_V$ 是带 block 内位置编码的可学习映射. 采用从 $r=0$ 开始、仅纳入完整窗口的教学计数, 到位置 $t$ 时可用compressed entries数为

N_{\mathrm{cmp}}(t)= \max\left(0,\left\lfloor\frac{t-l}{d}\right\rfloor+1\right). \tag{7}

Compression attention 对这些 entries 做完整低分辨率覆盖. 它直接参与式 (5), 因此语言模型梯度会训练压缩表示保存可用于输出的信息. 同一 attention 的概率

p_{t,r}^{\mathrm{cmp}}= \operatorname{Softmax}_r \left(\frac{q_t^\top \widetilde k_r^{\mathrm{cmp}}}{\sqrt{d_k}}\right) \tag{8}

还承担 selection 分支的 importance source. 这让 coarse 表示同时服务「产生输出」和「指出原始细节位于哪里」. 两项职责可能冲突: 对生成有用的压缩 Value 不一定对应最佳路由 key, 训练时应分别观察 compression 输出质量与候选 recall. Compression blocks 与 selection blocks 的划分可以不同. 用一个固定稀疏映射矩阵 $A_{b,r}$ 表示 compressed entry $r$ 对 selection block $b$ 的空间覆盖, GQA group 内有 $H$ 个 query heads, 则 block importance 可以统一写成

u_{t,b}=\sum_{h=1}^{H}\sum_r A_{b,r}p_{t,r}^{\mathrm{cmp},(h)}. \tag{9}

NSA 论文给出了由 block length、stride 与重叠关系展开的具体求和. 式 (9)保留其核心: coarse probability 先按空间关系投到原始 blocks, 再跨 GQA heads 聚合, 组内共享一套 top- $n$ .

Fine分支的直接梯度与选择边界. 令 $I_t=\operatorname{TopN}(u_t,n)$ , selection block 长度为 $l'$ , 则 fine 分支读取

S_t^{\mathrm{slc}}= \bigcup_{b\in I_t} \{bl'+1,\ldots,\min((b+1)l',t)\}. \tag{10}

在一次前向中固定 $I_t$ , selected attention 与普通 attention 相同. 设上游梯度为 $g_t=\partial\mathcal L/\partial o_t^{\mathrm{slc}}$ , 对选中位置 $s$ 的 logit $z_{t,s}$ 有

\frac{\partial\mathcal L}{\partial z_{t,s}} =\alpha_{t,s}\,g_t^\top \left(v_s-o_t^{\mathrm{slc}}\right), \qquad s\in S_t^{\mathrm{slc}}. \tag{11}

未选位置不参与 fine attention, 对这条路径的 Q/K/V 梯度为零. Hard top- $n$ 的集合变化也没有普通导数, 式 (11)不会告诉一个落榜 block「差一点就入选」. NSA 的 $u_{t,b}$ 来自可微的 compression attention, 这条 coarse 输出路径持续训练 query 与 compressed K/V; selection 集合本身仍是离散操作. 这条梯度边界可以写得更明确. Selection 集合在一个不跨越排序边界的小邻域内保持不变, 因此

\frac{\partial o_t^{\mathrm{slc}}}{\partial u_{t,b}}=0 \quad\text{a.e.}, \tag{12}

其中 a.e. 表示除 top- $n$ 交换边界外几乎处处. 在边界上集合发生离散跳变, 普通导数也不存在. $p^{\mathrm{cmp}}$ 收到的实际梯度主要来自 coarse output, 它作为 route score 的用途不会自动添加一条平滑梯度. 这解释了为什么候选 recall必须单独测量, LM loss下降不能直接证明 selector排序正确. 做一个两位置手算. Selected set含位置 $1,2$ , softmax权重为 $(0.75,0.25)$ , Value为 $v_1=(2,0),v_2=(0,2)$ , 则输出 $o=(1.5,0.5)$ . 令上游梯度 $g=(1,-1)$ . 式 (11)给出第一个 logit梯度 $0.75,(1,-1)\cdot((2,0)-(1.5,0.5))=0.75$ , 第二个为 $0.25,(1,-1)\cdot((0,2)-(1.5,0.5))=-0.75$ . 排名第三但未选的位置即使 Value更适合当前任务, fine路径梯度仍为零. 若把 hard top- $n$ 换成 soft relaxation、Gumbel或 straight-through estimator, selector能获得近似梯度, 训练图与部署图也会出现新差异. Soft routing训练时读取所有 blocks, 无法获得原生稀疏训练节省; straight-through前向 hard、反向 soft, 梯度带有估计偏差; 随机采样还会增加 kernel shape波动. 采用哪种方案要同时报告 forward实际稀疏率与 backward是否扫描全历史.

候选训练可以额外加入 teacher loss, 但这属于可选设计, 不能与 NSA 原论文的语言模型训练混为一项. 若有 full attention 教师分布 $p_{t,s}^{$ , 可以定义 block target $P_{t,b}^{$ }=\sum_{s\in b}p_{t,s}^{*} $P_{t, b}^{*} = \sum_{s \in b} p_{t, s}^{*}$ , 再加入

\mathcal L_{\mathrm{route}}= \sum_tD_{KL} \left(P_t^{*}\Vert\operatorname{Softmax}(u_t)\right). \tag{13}

DSA 的 indexer warm-up 正是这类显式教师对齐, 只是其候选粒度为 token, target 来自主 attention heads 的聚合. NSA 依靠 compression 路径与 LM loss 联合适配, 两者提供给 selector 的监督方式不同.

Window 分支抑制局部捷径对长程分支的干扰。 Window 集合为 $S_t^{\mathrm{win}}={\max(1,t-w+1),\ldots,t}$ . 局部模式往往收敛更快, 若 compression 与 selection 和局部 token 混在同一 softmax 中, 大量近邻高分会压低远距路径的梯度. NSA 把最近 $w$ 个 token 放进独立分支, 让 coarse 与 fine 分支继续收到自己的归一化信号. 分支 gate 的梯度也很直接. 令 $g_t^c=\sigma(a_t^c)$ , 上游梯度为 $r_t=\partial\mathcal L/\partial o_t$ , 则

\frac{\partial\mathcal L}{\partial a_t^c} =r_t^\top o_t^c\, g_t^c(1-g_t^c). \tag{14}

某条分支的输出尺度更大时, gate 容易优先沿它下降; gate 接近 $0$ 或 $1$ 后, sigmoid 饱和又会降低更新. 训练日志要记录分支输出范数、gate 分位数与分支梯度, 只看平均 gate 不能判断某条路径是否在少量长程样本中发挥作用. Window 与 selected blocks 可能重叠, 两条分支仍独立计算. 这部分重复是稳定训练换来的成本. 若实现为了省算直接从 selection 中删除 window blocks, fine 分支的 softmax 分母和训练输入都会变化. 可以把「允许重叠」和「去重」作为结构消融, 不能在部署 kernel 中悄悄改动.

路由目标分别约束排序、覆盖与成本。分布蒸馏令教师概率为 $p_s$ , selector概率为 $q_s$ , 使用

L_{KL}=\sum_{s\in V_t}p_s\log\frac{p_s}{q_s}.

它利用全历史概率, 对每个位置都有信号. 教师 attention若多 heads、多层, 还要决定聚合方式. 概率平均偏向共同关注, 最大值或并集目标强调任一 head需要的位置. 学生预算有限时无法同时覆盖所有教师尖峰, 聚合规则会塑造稀疏图. 排序目标直接比较正负位置. 对教师高质量位置 $s^+$ 与低质量位置 $s^-$ , hinge loss可写为

L_{rank}=\max(0,\mu-z_{s^+}+z_{s^-}).

它要求正例至少高出 margin $\mu$ , 不要求拟合完整概率尺度. 负例如何采样很重要: 随机远离边界的负例太容易, 当前 top-k附近的 hard negatives更能稳定支持集. 只用教师 top-k作正例又会忽略集合外仍有一定质量的位置. 覆盖目标关注选中集合承载的教师质量. 若 soft mask为 $m_s\in[0,1]$ , 可用

L_{cov}=-\sum_s p_sm_s.

hard mask下这就是最大化教师 mass recall, 梯度仍需连续代理或策略估计. 覆盖比集合 recall更贴近输出, 唯一低概率证据可能在教师尚未识别时被忽略. 下游任务 loss用于补足教师度量与真实行为之间的差距. 成本目标约束选择量:

L_{cost}=\lambda\sum_sm_s

或约束超过预算的部分. 固定 top-k已让数量等于 $k$ , 成本差异仍来自 block/page布局与负载. 若目标硬件按块执行, 正则应惩罚非空块或页面数, 只惩罚 token数可能学出高度分散的支持. 训练目标与实际执行粒度要一致. 四类损失回答不同问题: KL拟合教师分布, ranking稳定边界, coverage保留重要质量, cost控制资源. 把它们简单相加需要标度校准. 主任务 loss常按 token平均, 路由 loss可能按边或query求和;序列长度变化会改变相对权重. 先明确归一化分母, 再调系数.

负载均衡与证据覆盖不是同一目标。 MoBA式 block experts中, router把 query分给少量 KV blocks. 若少数 blocks持续热门, 对应设备或线程承担更多工作. 负载均衡损失鼓励各 block接收相近流量, 提升并行利用率. 语义分布本来就可能不均匀, 例如开头说明、最近窗口或固定文档标题确实被频繁访问. 强行均匀会牺牲路由正确性. 设 block $b$ 接收的query比例为 $f_b$ , 平均router概率为 $p_b$ . 常见均衡项与 $\sum_bf_bp_b$ 或方差相关. 它控制总体使用, 不保证每个query覆盖必要证据. 一个router可以均匀地把queries送到错误blocks, 均衡指标仍很好. 容量上限又与均衡不同. 每个 block或专家最多接收 $C_b$ 个queries, 超出的 token要丢弃、重路由或进入保底分支. overflow率直接改变实际支持. 训练只计算理想 top-k、推理受capacity截断, 就会出现支持不一致. 训练中要模拟相同容量与冲突解决规则.

覆盖也有 head维度. 多个 query heads共享候选时, 候选集合要覆盖各 heads的重要位置. 若用平均attention mass训练, 专门复制数字或代码符号的少数 head会被主流 heads淹没. 可优化最差 head覆盖、分位数或对专门化 heads加权, 相应候选并集会扩大. 局部 window分支会天然吸收大量高频近邻依赖, 让全局router专注远程位置. 如果教师分布包含局部 mass, selector预算可能再次被近邻占满, 与 window重复. 训练可在给定 window覆盖后, 用剩余教师质量监督远程分支. 局部 token依然重要, 只是已经由window承担, 远程预算应减少重复边.

探索不足会形成支持塌缩。训练初期某些blocks因位置偏置、初始化或数据频率略占优势, 被选次数更多, 获得更多主任务梯度, 表示变得更有用, 之后继续被选. 冷门blocks很少进入支持, keys和router都缺少学习信号. 这种正反馈可以在总体 loss正常下降时发生. 塌缩有几种形态. 位置塌缩长期只看最近窗口或开头 sinks;内容塌缩偏向高频标点、标题和模板;head塌缩让多个router heads选择相同集合;样本塌缩对不同query给几乎固定mask. 每类都要用不同统计识别:位置直方图、token类型分布、head间Jaccard和query间集合熵. router熵高也不保证探索有效. 概率分布可以很平, hard top-k仍因微小固定偏置选择相同边界;加入噪声后才会更换支持. 反过来, 熵低可能表示模型对正确唯一证据很自信. 熵是诊断量, 不能单独作为目标最大化.

随机探索应遵守因果与分块结构. Prefill位置 $t$ 只能从 $V_t$ 采样;packed batch不能跨文档;block采样要避免用未来 token统计决定早期query. 若探索单位与推理单位不同, 训练学到的支持分布也会错位. 推理按block选, 训练随机token边不能完整模拟. 一种保守策略保留确定性局部窗口和少量随机远程blocks, 再让学习router分配剩余预算. 局部通道保证基础语言建模, 随机通道提供新证据曝光, 学习通道逐步接管. 随训练推进减少随机份额, 观察远程 recall是否保持. 如果随机份额一撤就下降, selector尚未学会稳定替代.

支持抖动与支持僵化都可能伤害训练。第 $k$ 与第 $k+1$ 名score接近时, 很小参数更新会让位置频繁交换. 主 attention每步看到不同图, 梯度方差上升. 这种支持抖动在量化、分布式top-k tie和低精度累加下更明显. 可以监控相邻训练step在固定probe batch上的Jaccard与margin. 提高margin、降低router学习率或使用短期候选平滑能减少抖动. 平滑过强又会让支持僵化, 新证据难以进入. 训练目标是让有明确偏好的query稳定, 对不确定query保留探索, 不是让所有mask固定. EMA teacher可提供较慢变化的路由目标. 学生仍按当前参数执行, 教师分布用于蒸馏. 它降低目标漂移, 可能滞后于主模型新学到的表示. teacher更新率需要与预算日程和模型学习速度匹配.

支持变化还会改变不同参数的有效batch size. 某个历史key只有被许多queries选中时才积累稳定梯度, 冷门key更新稀疏. 对路由参数统计被选频次与梯度方差, 能发现一部分状态长期没有训练. 仅看全模型梯度范数会掩盖这种不均匀.

3. 从主任务梯度到训练配方

3.1. 总损失要区分主模型与路由器

一个通用训练目标可以写成

\mathcal L= \mathcal L_{\mathrm{LM}} +\lambda_r\mathcal L_{\mathrm{route}} +\lambda_b\mathcal L_{\mathrm{balance}} +\lambda_g\mathcal L_{\mathrm{gate}}. \tag{15}

$\mathcal L_{\mathrm{LM}}$ 优化最终 token 预测; $\mathcal L_{\mathrm{route}}$ 对齐 teacher 或提高候选质量; $\mathcal L_{\mathrm{balance}}$ 约束 block 热点与设备负载; $\mathcal L_{\mathrm{gate}}$ 防止分支过早关闭. 具体方法未必包含后三项. NSA 的核心训练结果来自端到端 LM 目标与三路结构; MoBA 使用 parameter-less gate, 没有独立 router 参数; DSA 明确加入 indexer KL loss. 梯度归属决定各项能优化什么. DSA sparse stage 将 indexer input detach, $\mathcal L_{\mathrm{route}}$ 只更新 indexer, $\mathcal L_{\mathrm{LM}}$ 只更新主模型. 这样避免 indexer loss 扭曲主 hidden state, 也让 LM loss 无法直接修正 selector. NSA 中 compression attention 同时参与输出和 importance, LM loss 会更新 coarse query/K/V 与 gate, hard top- $n$ 仍隔断集合边界. MoBA 的 block score由 Q 与 mean-pooled K 计算, 选中 blocks 内的 LM 梯度能更新 Q/K/V, 未选 block 的路由边界依旧没有连续梯度. 预算通常作为结构约束直接进入 $n,k,w,l'$ 等超参数, 比对总 loss 加一个「稀疏率惩罚」更容易控制 kernel shape. 动态预算可以为每个 query 预测 $n_t$ , 训练时还需约束 batch 内最大值、分位数和设备负载. 平均预算满足目标, 少数 query 激活大量 blocks 时, 真实 step time 仍由最慢 row 决定.

3.2. 从头预训练能让表示完整适应稀疏图

从随机初始化训练 NSA, 所有层一开始就面对 compressed、selected 与 window 三路约束. 模型可以把长期状态写入压缩器容易保存的表示, 让 fine route 的 query 与 coarse importance 共同成形, 也能在局部窗口外建立逐层传播路径. 论文实验在 27B 总参数、3B 激活参数的 MoE 骨干上预训练, 再做长上下文继续训练; 这一结果对应特定模型、数据与 GQA 配置. 从头训练的风险落在早期路由. 随机 query 与 compressed keys 产生的 top- $n$ 接近随机, fine 分支可能长期读不到有用远距 token. Window 与 compression 路径提供稳定信号, 还可以在初期增大 $n$ 、提高 coarse gate 或周期性插入 full layers, 再逐步收紧预算. 每次预算变化都要观察 loss 是否出现跳变、route entropy 是否塌缩以及长距离 recall 是否恢复.

训练长度分布同样重要. 大部分样本短于 window 时, 模型几乎没有动力训练 compression 与 selection. 平均 LM loss 会被大量短位置支配, 长序列末端的退化很难看见. 需要同时报告 trailing loss、按 query-evidence 距离分桶的 loss, 并确保足够样本真正跨过 $w$ 和 $nl'$ .

继续训练适合已有 checkpoint, 但要留足适应量。 MoBA 与普通 attention 使用同一组 Q/K/V 参数, $k$ 覆盖全部历史 blocks 时可恢复 full attention. 官方论文利用这一性质测试两阶段训练: 大部分 token 使用 MoBA, 后段切回 full attention; 还在 SFT 中保留若干 full attention 层, 缓解 prompt loss mask 导致的稀疏梯度传播问题. 这些实验说明切换可以平滑, 具体比例仍要随模型和任务重新验证. DSA 给出更明确的转换流程. Dense warm-up 保留完整主 attention, 冻结主模型, 用主 attention heads 聚合后的 target distribution训练 Lightning Indexer. 报告中的 warm-up 为 $1000$ steps、每步 $16$ 条 $128$ K 序列. 随后 sparse stage 每个 query 选择 $2048$ 个 KV tokens, 优化主模型与 indexer共 $15000$ steps; indexer 继续在选中集合上做 KL 对齐, 主模型用 LM loss 适应稀疏输入.

两阶段 KL 的可见范围不同. Warm-up 的 teacher覆盖整个合法历史, 能直接惩罚 indexer漏掉高质量 token. Sparse stage的对齐限制在当前 selected set内, 更侧重让已选 token的相对 score接近主 attention; 集合外 token无法通过这项受限 KL重新进入. Dense warm-up提供了初始全局排序, sparse stage再保持入选集合内的校准, 主模型则围绕实际 top-k结果调整表示. 转换过程需要三条曲线. 第一条是 full-history teacher上的 indexer mass recall, 检查 selector是否仍能找到主 attention需要的位置; 第二条是 sparse LM loss与 full reference的差距, 检查主模型适应程度; 第三条是 indexer KL, 检查打分校准. KL持续下降而 full-history recall下降时, 受限集合内拟合正在掩盖候选遗漏.

NSA 增加压缩器、三套 K/V 投影与 gate, 从 dense checkpoint 转换时没有 MoBA 那样的参数同构. 可以从原 Q/K/V 复制或分支初始化, 再用 dense teacher做输出蒸馏与 candidate supervision, 但需要重新验证分支平衡. 很短的 SFT 无法替代长上下文继续预训练, 因为 selector 需要看到多种距离、结构和证据密度. 一份可执行的转换日程可以分四段. 第一段保留 full attention, 初始化新分支并做输出蒸馏; 第二段以较大候选预算启用 sparse forward, 让主模型适应漏边; 第三段逐步收紧 $n$ 或 $k$ , 同时增加长样本比例; 第四段固定部署预算完成长上下文继续训练与 SFT. 每段都保存 checkpoint, 发生 loss跳变时回到上一个预算, 比一次性从 full切到目标稀疏率更容易定位问题.

预算课程还要避免只按 step变化. 模型在短序列上即使 $n$ 很小也可能覆盖大部分历史, 到长序列才真正稀疏. 可以按实际可见比例 $\rho_t=|S_t|/|V_t|$ 或 evidence距离控制课程, 确保训练早期和后期都包含跨窗口样本. 对混合长度 batch, 记录每个长度桶的 $\rho$ 分布比记录全局平均更准确.

SFT 与 RL 会重新改变稀疏梯度分布。 SFT 常把 prompt token 的 LM loss mask 掉, 梯度只从 response 回传. Full attention 能让 response query直接连接全部 prompt, block sparse route 若漏掉某段 prompt, 那段表示在多层中收到的监督会更弱. MoBA 论文用高层 full attention缓解这一现象, 也提示训练原生稀疏不能只在预训练阶段验收. RL 的 rollout 由推理引擎生成, 训练框架再重放序列. 两边必须对 candidate selection、top-k tie、量化和因果 mask保持一致. 同一 token 在 rollout 与训练中看到不同 blocks, log probability 差异会混入 off-policy error. 保存 route indices可以提高一致性, 代价是额外轨迹存储; 重新计算则要求两个 runtime 数值对齐. 长 CoT 还会改变预算压力. Prompt 证据在生成后段距离越来越远, 固定 window 无法覆盖, selection 要持续找回早期状态. 训练集若只有短 response, selector 可能只学到 Prefill 内检索, 到长 rollout 后段召回下降. RL 验收要按生成位置记录候选覆盖、gate 与 TPOT, 不能只比较最终 reward.

训练与推理必须共享支持语义。训练期 selector若在全精度 score上取 token top-k, 推理期改成量化 block top-k, 两者预算同为 $k$ 也不是同一支持. block会带入组内额外 token, 量化会改变排序, page对齐与capacity截断还会继续修改集合. 模型部署前要用推理真实规则做足够长适配. 支持一致性可以逐项列出: 合法集合 $V_t$ 、选择单位、预算、tie-breaking、局部/window合并、去重、容量溢出、精度和fallback. 任一项变化都可能改变 $S_t$ . 训练可用慢参考实现, 数学语义必须与推理kernel相同;kernel优化不能偷偷改变边界规则. Prefill与Decode也可能使用不同支持实现. Prefill成块处理许多queries, Decode逐 token;如果训练只覆盖Prefill block mask, 逐步Decode的候选更新与state追加可能未学到. 在训练或验证中用逐 token路径重放同一序列, 与并行因果实现对齐, 能检查支持和logits一致性.

dropout会让训练支持更丰富, 推理关闭后分布收缩. 对attention probability dropout问题较小, 对router score或mask dropout会直接改变可见图. 若它被用作探索, 需要确认关闭后 learned selector仍有足够覆盖. 可在训练末段降低router dropout, 让模型适应确定性支持. 长度一致性同样关键. 训练长度 $N_{train}$ 内每个query可从有限历史选择, 推理到更长 $N$ 后, 候选竞争数增加, top-k比例下降, score极值分布变化. 固定 $k$ 在长度外更稀疏. 训练应包含目标长度或使用能控制长度分布变化的目标, 不能只验证kernel能运行.

路由器与主模型的学习速度要协调。 selector学得太快, 会在主 attention尚未形成前根据早期噪声固定支持;学得太慢, 主模型长期在随机或粗糙图上优化, 学会绕过远程分支. 可以为router使用独立学习率、warmup与梯度裁剪, 目标是让支持随表示改善而更新, 又不在每步剧烈跳变. 冻结方案能帮助定位. 先用teacher或高预算训练主模型, 暂时冻结selector;再固定主模型训练selector模仿oracle;随后联合微调. 分阶段降低耦合难度, 最终质量可能受早期图限制. 从头联合训练更原生, 优化相互依赖更强. 两种配方回答不同问题, 应在相同token预算下比较.

梯度尺度也不一致. 主loss按每个token一个标量, selector可能为每个query-key对产生大量logits. 若路由loss对所有边求和, 长度翻倍会显著增大梯度;若按合法边平均, 单个关键正例信号被大量负例稀释. 常见做法按query归一化, 再对正负样本或教师mass加权. 共享Q/K参数时, 主attention梯度与router辅助梯度可能冲突. 一个方向有利于候选排序, 另一个有利于候选内表示. 可以测梯度余弦, 冲突明显时分开投影、使用adapter或调整损失权重. 参数共享节省容量并传递信号, 也把两个目标绑定在同一空间. 训练日志应把主loss、路由loss、coverage、margin、support churn与分支梯度分开. 总loss下降无法说明selector改善, 因为主模型可能越来越依赖window捷径. 路由recall提高也不保证任务变好, 教师目标可能与最终主attention不同. 多条曲线要在同一probe batch上对应.

SFT 会改变支持分布。预训练文本的依赖多由局部语言统计主导, SFT指令包含问题、上下文、格式约束与答案, 远程证据位置和输出结构不同. router在预训练阶段形成的局部偏好可能不适合长文问答. SFT需要继续更新selector, 或至少验证冻结路由是否覆盖指令数据所需证据. 短SFT样本还会使长程分支缺少梯度. 即使底座支持128K, SFT大多只有几千 token, selector会在局部样本上继续优化, 产生遗忘. 混入长上下文SFT、保持预训练路由蒸馏或降低router学习率, 可以保护长程能力. 哪种做法有效取决于SFT数据量和支持漂移.

答案only loss让context token本身没有直接next-token监督, 路由仍通过答案位置反传到被读取证据. 若答案很短, 每个样本提供的router query很少. 在问题理解和中间推理位置加入监督可增信号, 也会改变SFT目标. 至少要统计有效router queries数量, 不能只看总context tokens. 格式模板容易形成捷径. selector持续选择开头system prompt、末尾问题和最近token, loss也能下降;正文中段证据被忽略. 随机化证据位置、加入无关段落和相同模板下的不同答案, 能迫使支持依赖内容. 固定模板上的高准确率不足以证明动态路由.

RL 的稀疏奖励更延迟。 RL中一条长轨迹只在结果或若干步骤得到奖励, selector在早期漏掉证据与最终失败相隔很远. 策略梯度方差已经较高, 离散支持再增加一层信用分配. 若router动作被视为策略的一部分, 需要记录其log-prob并纳入重要性比率;若router确定性地由hidden产生, 梯度主要通过模型输出策略间接到达. 奖励可能鼓励更长推理, 同时增加Decode steps与历史长度. 固定 $k$ 下支持比例进一步下降, selector要处理不断增长的轨迹. agent任务还会插入observation, 这些片段在后续步骤的重要性可能突变. RL数据分布与预训练文本差别很大.

成本可以进入reward, 例如惩罚选中blocks或full attention调用. 若惩罚过强, 策略会选择便宜但信息不足的图;奖励稀疏时, 成本信号每步都有, 更容易主导优化. 应先在固定质量边界下调成本系数, 并监控任务分桶, 避免平均reward用短任务收益遮住长任务退化. off-policy数据还存在支持不一致. 旧策略生成轨迹时的hidden与selector不同, 当前模型重放同一tokens会选择新集合. 如果训练只存token不存路由状态, 无法精确重现行为策略计算图. 对常规语言模型更新可以重新前向, 对把router视为显式动作的算法则要保存或重建其概率.

RL阶段预算若与部署一致, 模型有机会适应真实成本;为了训练吞吐临时用更大预算, 部署收紧会再次产生结构偏移. 可以在部分rollout或更新batch使用目标预算, 测质量-成本曲线. 稀疏支持属于策略执行的一部分, 不能只在RL结束后替换.

课程学习要覆盖距离、干扰与操作类型。只增加序列长度不是完整课程. 路由难度还取决于候选数量、相似干扰、证据位置和记忆操作. 单一needle、无干扰时, selector可能靠表面模式找到证据;加入多个相似keys、覆盖更新与多跳关系后, 才能训练细致支持. 距离课程从窗口内、刚越过窗口到多个窗口外, 检查模型何时启用远程分支. 干扰课程增加相似但错误的blocks, 训练margin. 操作课程区分复制原文、取最新值、累计统计和多证据组合. 四类任务需要不同支持, 能防止router把「远」简单等同于「重要」.

预算课程与任务课程联动. 简单阶段用大预算会让模型不必精确选择, 难阶段突然收紧容易崩溃. 可以先在大预算下学会任务, 再逐步缩小并增加干扰;也可以固定目标预算、从少候选到多候选增长长度. 两条路径分别控制支持容量与搜索难度. 课程转换处要看旧能力. 学会远程检索后加入大量局部数据, router可能重新塌缩;学会最新值覆盖后加入计数, 过强遗忘会伤害累计. 混合replay与分桶验证能发现灾难性遗忘. 平均loss连续下降并不保证每种记忆操作保持.

4. 稀疏预算必须同时满足数学和硬件

4.1. 三路总预算可以直接手算

NSA 在位置 $t$ 读取的逻辑条目数近似为

N_{\mathrm{NSA}}(t)= N_{\mathrm{cmp}}(t)+nl'+w. \tag{16}

取论文训练配置 $l=32,d=16,l'=64,n=16,w=512$ , 当 $t=65536$ 时, $N_{\mathrm{cmp}}=4095$ , 按零起点完整窗口计数, selected raw tokens为 $1024$ , window 为 $512$ , 总计 $5631$ , 约占 full history 的 $8.59%$ . 三路各自计算 attention, 重叠位置也会重复, 所以这个比例表示逻辑 keys 数, 还不能直接当成 kernel 加速比. 整个长度 $T$ 的训练前向成本可近似写成

C_{\mathrm{NSA}}(T) \propto \sum_{t=1}^{T} \left(\frac{t}{d}+nl'+w\right)d_h =O\left(\frac{T^2d_h}{d} +T(nl'+w)d_h\right). \tag{17}

压缩分支仍随历史增长, 系数由 stride $d$ 缩小; selection 与 window 在预算固定后对每个 query 为常数. 计算 compressed entries 的 MLP、importance 映射、top- $n$ 和 gate 还要另计. Decode 的常驻状态则包含三套分支 cache 与压缩表示, 它的字节比例和式 (17)不同. MoBA 的成本结构更简单. 长度 $t$ 、block size $B$ 、top- $k$ blocks 时, router 扫过约 $t/B$ 个 block keys, 主 attention 读取最多 $kB$ 个 raw tokens, 当前 block包含在 $k$ 内. $t=32768,B=512,k=3$ 时, 主 attention 读取最多 $1536$ tokens, 即 full history 的 $4.69%$ , 另加 $64$ 个 block summaries 的路由. Block 越细, route candidates 增加; block 越粗, 每次选中带入更多无关邻居. MoBA 把当前 causal block与历史 non-causal blocks分别送入变长 FlashAttention, 再用 online softmax合并两部分. 合并需要保存每部分的 row max与归一化和. 若当前块输出为 $(o_c,m_c,l_c)$ , 历史块输出为 $(o_h,m_h,l_h)$ , 统一最大值 $m=\max(m_c,m_h)$ , 则

o=\frac{ e^{m_c-m}l_co_c+e^{m_h-m}l_ho_h} {e^{m_c-m}l_c+e^{m_h-m}l_h}. \tag{18}

直接把 $o_c,o_h$ 平均会丢失两个分区的 softmax分母, 与在候选并集上做一次 attention不等价. 当前块强制入选还解决了 causality: block mean可能包含 query之后的 token, 它不能用于决定当前块是否可见, 当前块内部必须继续使用 triangular mask. 计算稀疏也不自动压缩常驻 cache. Selection 分支为了让未来 query任意选择历史 block, 通常仍要保存全长 raw K/V; window分支只保存最近 $w$ ; compression分支保存约 $t/d$ 个 entries. 假设 $H_{kv}=4,d_k=192,d_v=128$ , BF16下每 token一套 K/V为 $4(192+128)\times2=2560$ bytes. $t=65536$ 时, selection cache约 $160$ MiB, window cache约 $1.25$ MiB, compression entries约 $10$ MiB, 都是单层逻辑值. 三路独立投影的总 cache可能超过一层 full attention, NSA主要节省 attention计算与每步实际读取, cache容量必须单独核算.

这也影响「训练原生」的系统目标. 长 Prefill与训练关注少算 QK/AV, Decode关注少读实际候选 KV; 高并发服务还关心常驻容量. 一个结构可以在前两项获益, 第三项持平或变差. 若要同时压 cache, 还需跨层共享、量化、淘汰或固定状态机制, 这些改动会引入新的训练约束.