Browse this library

Chinese source article

4.1 · 稀疏训练与适配

解释稀疏结构怎样进入优化过程,以及原生训练与长上下文微调为何需要不同证据

View original Markdown ↗

稀疏注意力可以在三个时刻进入模型:从预训练开始就参与每次前向和反向;在已有稠密模型上训练一个选择器,再让主模型适应稀疏路径;只在长上下文微调阶段使用便宜的近似图,部署时恢复原始 full attention。三种做法都会出现“训练时用了稀疏注意力”,学习对象却不同。第一种让表示与稀疏拓扑共同形成,第二种要处理稠密到稀疏的分布转换,第三种主要降低微调成本,不能直接证明稀疏图适合最终推理。

训练问题的核心不只是选择出哪些边,而是监督信号如何抵达选择器。hard top-kk 对离散索引不可直接求导;若先用稠密 attention 产生 oracle,选择器可以学习重现高价值候选,但稠密 teacher 的成本仍存在于训练阶段;若使用可微近似或多分支 attention,梯度可以流入路由,却可能让训练时的软选择与推理时的硬预算不一致。每种路线都要说明哪部分参数从任务 loss 学习,哪部分依赖额外标签或蒸馏。

从稠密 oracle 学什么

给定某层某头的完整 attention,可以把高分 token、概率质量覆盖或对输出最有影响的块作为候选标签。top-kk 位置标签最直接,却把第 kk 与第 k+1k+1 名之间很小的分数差放大成离散边界;概率质量目标更平滑,但选择器最终仍需满足固定物理预算。使用最终任务 loss 作为监督更接近目标,信用分配也更长,容易让选择器学到数据捷径。

DSA 一类两阶段方案先在 dense attention 下训练 indexer,再切换到 token 级稀疏 attention 继续训练。第一阶段让索引器看到可靠 oracle,第二阶段让主模型适应真实漏召回和 hard top-kk。转换时机、稀疏预算退火与学习率会影响稳定性:过早切换会把尚未可靠的路由噪声注入主干,过晚切换则让模型缺少在目标拓扑下重组表示的训练量。

NSA 将压缩、选择和滑动窗口作为并行分支放进原生训练。局部分支提供稳定近邻,压缩分支聚合全局概要,选择分支读取重要块。这样的残差式组合降低单一选择器失误的破坏性,但每条分支都有成本,训练和推理必须使用一致的预算与 kernel 才能兑现效率。MoBA 用块摘要路由到少量 KV blocks,并保留当前因果块;块级门控更适合连续执行,页内无关 token 也会一起参与。

训练—推理一致性

常见偏差来自四处。训练时可能用较短序列,选择器从未见过百万 token 下的候选竞争;训练时使用 soft 权重,推理时只保留 hard top-kk;训练时用稠密 oracle 或全局统计,部署时只能访问便宜摘要;训练 kernel 与 Decode kernel 的分块规则不同,让同一预算产生不同候选。只要其中一处变化,最终质量就不能只用训练 loss 推断。

选择器还会与表示共同演化。主模型为了被正确检索,可能把关键信息写成 indexer 更容易辨认的方向;indexer 又根据当前表示更新。这种共适应是原生训练的优势,也是迁移难点。冻结主干只训 indexer 能减少扰动,却限制表示对稀疏瓶颈的适应;全部联合训练能提高上限,也更容易出现路由塌缩或少数块过载。

训练日志应包含路由熵、候选利用率、不同 head 的重复率、关键证据 recall、保留 attention mass 和梯度范数。最终 loss 相近时,这些量可以提前暴露选择器只依赖局部位置、所有 query 选择同一批块,或某些层完全退化为窗口。把预算从小到大画成质量曲线,也比只挑一个效果最好的点更能说明机制。

长上下文微调中的近似图

LongLoRA 的 S²-Attn 服务于另一种需求:已有模型希望用有限显存完成长上下文微调。它把 heads 分组,并对部分组做位移,使每组在局部块内计算时仍能跨越原始块边界。训练阶段的稀疏图降低显存和计算,部署时可以恢复 full attention。这里验证的是近似训练信号能否让模型学会更长位置范围,不是稀疏推理能否保持同样输出。

这种适配还涉及 RoPE 扩展、LoRA 参数位置、embedding 与 normalization 是否训练,以及长数据构造。若只比较 attention mask 而没有固定这些因素,很难把收益归因于 S²-Attn。训练长度扩展也不等于有效上下文等比例增长,必须用变量追踪、多证据聚合和长度外推任务检查。

原生稀疏训练与稀疏微调可以共享部分工具,例如块 mask、稀疏 kernel 和长度 curriculum;评价目标仍应分开。前者关心目标拓扑下从训练到 serving 的完整效率和质量,后者关心以较低资源获得长上下文适配,并允许推理回到稠密路径。

怎样建立有说服力的对照

原生训练应固定模型参数量、训练 token、数据顺序、上下文分布和优化预算,比较 dense、相同局部窗口、oracle 路由与真实路由。两阶段转换还要报告切换前后 loss 跳变、恢复所需 token 和不同预算下的结果。若稀疏模型额外使用更多数据或更长训练,最终分数不能单独归因于结构。

微调路线要固定 base checkpoint、训练样本、位置扩展方法、可训练参数和总 tokens,比较 full attention、窗口、S²-Attn 等图。部署评测分别运行训练时稀疏图与恢复后的 full attention,才能知道收益来自参数学习还是推理图本身。显存、steps/s、最终有效上下文和短上下文退化都应报告。

选择器训练沿 oracle、hard top-kk、联合训练与机制指标展开,LongLoRA给出 S²-Attn、位置扩展和微调复现。两篇放在同一路线,是为了比较稀疏结构进入优化的不同方式;它们的成功标准并不相同,不能用一条“训练更省”的结论互相替代。

预算也属于训练分布

选择器在固定 top-k 下训练,部署时临时减半预算,改变的不只是计算量。候选边界、路由熵和残差分支承担的压力都会变化。更稳妥的做法是在训练中覆盖一组预算,或为目标预算分别校准;评测则画出预算曲线,观察模型是在预算下降时平滑退化,还是越过某个点后突然崩溃。突变往往意味着任务所需证据数超过容量,或选择器依赖某个固定冗余。

长度也是训练分布的一部分。短序列下 top-k 占比很高,选择问题接近稠密;长度增加后,相同 k 要在更多干扰项中检索,路由难度和证据竞争同时上升。因此应同时比较固定 k、固定比例和分层预算。固定比例会让成本随长度增长,固定 k 则让质量压力持续增加,两种曲线回答不同问题。

实验还要排除“参数或数据补偿结构缺陷”。稀疏模型若使用更多训练 token、更长 curriculum 或额外 teacher,最终质量提升可能来自预算增加。最有说服力的结果包括等数据等计算对照,以及在相同最终系统成本下允许各方法选择最佳训练配方的对照。前者隔离机制,后者回答实际资源配置,两者一起比单一榜单更完整。