4 · 训练：让稀疏结构真正可学

稀疏拓扑减少了可交互边，也改变了优化问题。固定窗口不给远距离依赖直接梯度路径；离散 top-k 在选择边界处不可导；训练期使用稠密注意力、推理期突然裁剪，又会造成分布偏移。模型需要在训练中学会使用有限连接，架构定义本身不会自动产生这项能力。

需要同时观察路由器与主干。路由器必须学会召回关键块，又不能把所有 query 挤到少数热点；主干必须适应有限邻域，并学会经由全局 token、跨层传播或周期性稠密层搬运信息。常见手段包括从短到长的上下文课程、从稠密到稀疏的退火、辅助召回损失、负载均衡正则、教师分布蒸馏，以及对未选边的抽样探索。每种手段都会增加计算或状态，节省量要按完整训练 step 重算。

1. 结构改变以后，优化目标也变了

稠密 attention 中，主任务损失可以通过 softmax 权重传到每个允许访问的 key。hard top-k 把候选外的边直接移出计算图，路由器的排序错误会阻断这些梯度。若仅依赖入选边的主任务梯度，模型可能过早集中到少数容易命中的局部块，远程候选长期得不到训练。soft routing 允许所有候选保留非零权重，梯度连续，却无法直接获得 hard routing 的计算节省；直通估计在前向执行离散选择、反向使用连续近似，估计带有偏差，但能把部署结构放进训练。

辅助监督从哪里来，是训练设计的核心。稠密教师可以提供每个 query 的高权重位置或输出分布，路由器学习复现教师候选；如果教师本身只在较短上下文上可靠，监督也会继承长度边界。对比学习可把相关 query—块作为正样本，把相似但无关的块作为困难负样本，前提是数据能够标出证据位置。只用最终答案监督时，模型可能通过局部统计得到正确答案，路由器未必真正召回长程证据。

联合目标常写成

\mathcal{L}=\mathcal{L}_{\mathrm{LM}}+\lambda_r\mathcal{L}_{\mathrm{route}}+\lambda_b\mathcal{L}_{\mathrm{balance}}.

$\mathcal{L}$ 是语言模型目标， $\mathcal{L}$ {\mathrm{route}} $L_{route}$ 约束候选召回或教师分布， $\mathcal{L}_{\mathrm{balance}}$ 抑制路由热点， $\lambda_r$ 与 $\lambda_b$ 控制两个辅助目标的权重。增大 $\lambda_r$ 可能提高教师候选召回，却不保证最终生成质量同步上升；增大 $\lambda_b$ 能让设备负载更均匀，也可能迫使 query 离开最相关的块。训练曲线需要同时报告三项损失、路由利用率和主任务指标，不能只展示总损失。

1.1. 从稠密模型迁移到稀疏模型

从零开始预训练允许模型自始至终适应稀疏图，但成本高，且很难与同预算稠密基线严格比较。把已有稠密 checkpoint 转为稀疏结构更常见：初期保留较大窗口或较多候选，让稀疏输出接近原模型；随后逐步缩小预算，并继续训练路由器和主干。若一步切到很小的候选集，attention 输出分布、残差幅度和梯度尺度会同时变化，短期损失突升不一定能靠少量继续训练恢复。

转换哪些层也影响稳定性。浅层表示局部形态和词法关系，规则窗口通常能够覆盖；深层更可能需要跨段整合，却并不意味着只把深层设为全局就一定最好。周期性稠密层给信息提供确定的全局交换点，代价是这些层仍保持平方成本。训练时可先只替换部分层，通过消融观察质量与吞吐，再扩展到更多层。层级配置必须与推理实现一致，否则训练所得收益会被部署时的额外裁剪改变。

长上下文课程同时改变序列长度和负样本数量。长度增加后，每个 query 面对更多相似块，固定 top-k 的检索难度会上升；若同步增加 $k$ ，训练计算也会增长。合理的课程会分别控制长度、候选预算和证据距离，让模型逐步见到更难的选择，而不是把三项同时放大后无法判断损失来自哪里。采样数据还要包含真正依赖远处信息的 token，否则大部分梯度仍由局部可预测内容贡献。

2. 稳定性、负载与分布式训练

路由器输出在训练早期容易抖动。相邻 step 对同一模式选择不同块，会让主干看到不断变化的输入分布；top-k 边界附近的小分数差异还可能在低精度下翻转排序。温度退火、候选过采样、指数滑动平均或路由一致性损失可以减少波动，但都会改变探索范围。过早稳定可能把初期错误固化，长期高熵又会妨碍缓存和 kernel 优化。

批内负载是模型目标与系统目标相交的地方。设一批 query 对块 $r$ 的选择次数为 $c_r$ ，理想执行希望 $c_r$ 不要过度集中；语义相关性却可能天然集中在共享前缀或任务说明。负载均衡应约束设备无法承受的热点，而非机械追求所有块完全均匀。分布式训练需要记录 all-to-all 的消息大小、设备间 token 数差异与等待时间，并把这些量和路由损失共同调节。

选择结果还参与反向传播的状态管理。训练需要保存候选索引，或在 backward 中以确定性方式重算；若随机路由、dropout 或索引更新使两次选择不一致，梯度会对应另一张计算图。activation checkpointing 只能在随机状态和路由参数都可复现时安全重算。对动态 Sparse Attention 来说，索引张量本身也是 activation，其容量与通信必须计入训练开销。

选择器训练沿着监督信号进入这一问题：稠密 attention 怎样产生 oracle，hard top-k 截断了哪条梯度，预算又怎样进入目标函数。监督与联合训练进一步拆解 DSA、NSA 与 MoBA 的训练数据流，给出召回、概率质量覆盖和训练—推理一致性的验收方法。评价不能只看最终 loss：还要记录路由熵、块利用率、关键证据召回、跨层路径和不同长度下的退化。一个在短序列上收敛的稀疏模型，不会仅凭可扩展的 mask 自动获得长程推理能力。

训练章承接动态路由与 KV 选择定义的选择对象，把候选规则变成可优化参数；随后交给内核检查训练和部署是否执行相同拓扑，再由评测验证长度、质量和系统收益。相邻章节之间最容易出现的断层，是训练论文只报告模型指标，内核实现却为规则化而改动候选粒度。两边的块大小、top-k 和 fallback 必须逐项对应。

失效分析至少覆盖四种情况：路由塌缩到局部或少数热点，训练长度外候选召回下降，辅助损失改善但语言模型质量下降，以及部署 kernel 对候选做截断后出现训练—推理偏差。每种情况都要从数据、目标、结构和实现四个方向定位，避免把所有问题归结为路由器容量不足。

2.1. 导航

稀疏长上下文微调讨论另一种训练用途：把稀疏图当作full attention微调的近似。LongLoRA给出S²-Attn的完整索引、复杂度、位置处理、参数消融与复现路径，部署时恢复原始full attention。

可微近似与直通估计、负载均衡、长度课程、稠密教师蒸馏、训练稳定性和推理拓扑对齐共同决定训练路径。公式会明确梯度流向，实验会固定 token 预算与实际计算量。

完整训练方案需要回答：哪些参数从已有模型继承，路由监督从哪里产生，离散选择怎样传梯度，长度和预算怎样调度，负载约束作用于哪个统计量，部署时哪些规则必须保持一致。训练曲线还要区分主任务、召回、负载和实现偏差，单个总损失不足以判断稳定性。

复现实验还应保存路由随机种子、索引刷新周期、候选并列时的排序规则以及每层实际保留比例。这些细节会改变梯度路径和设备负载，缺失时即使损失函数相同，也难以判断两次训练为何得到不同的稳定性与长程召回。

3. 稀疏训练改变了梯度图

设第 $\ell$ 层路由器参数为 $\phi$ ，主干参数为 $\theta$ ，硬集合为 $C_i=\operatorname{TopK}(r_\phi(q_i,K))$ 。稀疏输出

o_i=\sum_{j\in C_i}\alpha_{ij}v_j

对这一 query 的直接 attention 路径，梯度只经过入选的 key 和 value。未入选位置仍可能经其他 query、压缩分支或残差路径收到梯度，主干的共享参数也会从其他位置更新。对 $\phi$ 而言，集合成员在排序不变的小邻域内保持不动，普通反向传播无法告诉路由器「第 $k+1$ 名其实应该入选」。若路由分数还参与入选项的连续加权，这部分权重可以求导；离散索引本身仍不可导。

若使用连续门 $g_{ij}=\sigma(r_{ij}/\tau)$ ，可把输出改写为

o_i=\frac{\sum_jg_{ij}\exp(s_{ij})v_j} {\sum_jg_{ij}\exp(s_{ij})}.

温度 $\tau>0$ 时，有限分数对应的门值都非零，完整计算通常仍接近稠密。其导数为 $g_{ij}(1-g_{ij})/\tau$ ：降低温度会使梯度集中在分数接近零的位置，远离零的位置则逐渐饱和。这是以零为阈值的门， $\tau\to0$ 时保留正分数项，数量未必等于 $k$ 。若部署使用固定 top-k，还需使用相应的容量约束或 top-k 松弛，并直接比较软门与硬路由输出。

3.1. 教师信号应该监督哪一个量

稠密教师提供 logits、softmax 概率和最终 attention 输出，分别对应三种监督。排序损失鼓励相关 key 排在前面；分布损失让候选内概率接近教师；输出蒸馏直接约束聚合后的向量。可写成

\mathcal L_{\text{route}}= \lambda_{\text{rank}}\mathcal L_{\text{rank}} +\lambda_{\text{mass}}\mathcal L_{\text{mass}} +\lambda_{\text{out}}\|o_i^{\text{dense}}-o_i^{\text{sparse}}\|_2^2.

这里各项还要说明梯度如何抵达选择器。设教师概率为 $p_{ij}$ ，硬候选质量 $M_i=\sum_{j\in C_i}p_{ij}$ 可用于评测，但排序不变时， $1-M_i$ 对仅生成离散索引的选择器没有梯度。可微质量目标可取 $\mathcal L_{\text{mass}}=1-\sum_j g_{ij}p_{ij}$ ，并配合候选数量约束，否则把所有门打开就能降低损失。输出蒸馏若只经过 hard top-k，也只能训练入选路径中的连续参数；要更新索引选择，还需代理梯度或额外排序监督。

排序标签容易预计算，却会把边界两侧很接近的分数硬分正负；概率质量能区分重要程度，但需要教师 softmax；输出蒸馏允许漏掉与其他 value 冗余的边，却可能让不同错误互相抵消。训练早期可依靠教师排序，随后提高主任务和输出项比重，使选择器适应正在变化的主干表示。

教师本身也受长度影响。固定长度生成的 top- $k$ 标签，在更长序列上面对更多高分干扰项；同一个 $k$ 覆盖的概率质量会下降。蒸馏数据应覆盖目标长度分布，并按距离、query 熵和教师尾部质量采样。否则平均 loss 主要由短序列和尖锐 query 决定，恰好绕开最难的长程选择。

4. 三分支结构怎样分担风险

NSA 将压缩、选择和局部窗口放在同一 attention 结构中。可以把输出抽象为

o_i=g_i^{\text{cmp}}o_i^{\text{cmp}}+g_i^{\text{sel}}o_i^{\text{sel}}+g_i^{\text{win}}o_i^{\text{win}},

三个标量门 $g_i^c\in[0,1]$ 由当前位置输入特征经过 MLP 与 sigmoid 得到，独立计算，不要求和为 1。压缩分支提供低分辨率全局覆盖，选择分支读取少量原始块，窗口分支保留邻近细节。三条分支使用同一个 query，却有不同的 KV 表示，各自计算 attention 后再门控相加。选择分支漏掉远处细节时，压缩分支仍可能携带概要；压缩把两个实体混合时，精确选择可以恢复其中一个原始块。NSA 原文 §3.2给出了这一融合公式。

这种互补关系要靠分支消融验证。只关闭选择分支可测精确远距边的贡献，只关闭压缩分支可测全局低分辨率通道，只关闭窗口可测局部稳定性。若完整模型优于各单支，但任意一支关闭都几乎不变，可能存在冗余；若某一支承担绝大多数梯度，其他分支可能训练不足。还应记录门控前后的各分支输出范数，避免把门值很小的名义分支算作有效通道。

4.1. MoBA 的路由为何像专家选择

MoBA 把 KV 块视为候选专家，query 选择少数块执行 attention。设块 $b$ 的代表为 $\bar k_b$ ，路由分数为 $r_{ib}=q_i^\top\bar k_b$ 。若相关信号只出现在块内一个 token，简单均值会把它缩小约块大小的倍数，而无关 token 的波动仍进入代表。块越大，计算越规则，稀有信号的信噪比却可能下降；块越小，路由更精细，top-k 和 kernel 调度成本上升。

因果训练还带来边界效应。当前 query 所在块通常只能读取块内过去部分，历史完整块则可以全部读取。把当前块强制加入候选能保护局部连续性，却让不同位置拥有不同的有效候选预算。报告 top- $k$ 时要说明是否包含当前块、局部块是否额外计算，以及不足一个完整块的尾部怎样处理。

MoBA 的块均值与 query 点积不引入独立的路由器参数；固定主干后，没有额外路由器可单独预热。带可学习 indexer 的方法才可以先冻结主干、训练 indexer，再联合解冻。两者的训练对象不同。

MoBA 能在 full 与 sparse attention 之间切换。原文 §3.2比较了先用 MoBA 训练 90% token、再用 full attention 训练余下 10% 的混合配方，还研究了部分层保留 full attention 的做法。预算退火则是另一种实验设计：逐步改变选中块数时，应记录总训练 token 与实际执行边数，才能区分拓扑变化和计算预算的影响。

5. 长度课程与预算课程要解耦

设第 $u$ 个训练阶段的序列长度为 $n_u$ ，每个 query 的精确候选数为 $k_u$ 。若 $n_u$ 增大而 $k_u$ 固定，保留比例下降，候选竞争通常更激烈；每个 query 的 attention 成本近似固定，全序列边数约为 $n_uk_u$ ，随长度线性增加。若 $k_u/n_u=\rho$ 固定，每个 query 的候选数随长度增加，全序列边数约为 $\rho n_u^2$ ，仍是平方级。因果起始位置和局部保底会改变具体计数，但不改变这里的渐近阶。两种课程分别检验固定单步预算下的外推与固定保留比例下的建模能力。

可以先在固定长度上收紧 $k$ ，再固定 $k$ 增加长度，从而区分预算压力与长度外推。证据距离也应独立控制：把无关 padding 拉长不能代替真正远距依赖；将目标事实移动到不同距离，才能观察路由器是否依靠绝对位置或局部捷径。训练 batch 中还要保留一定比例的短样本，避免所有梯度都被长样本的少量困难 token 主导。

5.1. 训练稳定性的可观测量

路由熵只描述分数分布，还需要四组统计共同判断。第一组是教师 top- $k$ 召回和概率质量覆盖；第二组是不同 head、层、样本之间的候选重合；第三组是每个物理块被选择次数的分位数；第四组是路由更新前后主任务梯度范数。候选突然改变时，梯度尖峰可能出现在主干而非路由器，单看辅助损失会漏掉。

令连续两次参数更新后的候选为 $C_i^{u}$ 与 $C_i^{u+1}$ ，路由翻转率可写成

F_i^u=1-\frac{|C_i^u\cap C_i^{u+1}|}{|C_i^u\cup C_i^{u+1}|}.

训练早期较高翻转率表示探索，后期若仍然很高，主干持续面对变化的计算图；翻转率迅速降到零且召回不佳，则可能过早固化。把翻转率与质量、负载和温度曲线对齐，能判断稳定化策略是在减少无意义抖动，还是冻结错误路由。