
你可能不知道,当下最火的长文本大模型加速方案里,藏着一个让工程师们头疼了很久的悖论。
这个悖论是:你想训练一个"选择器"去挑出最重要的文本片段,但挑选这个动作本身,恰恰切断了让选择器变聪明的那根线。
先说清楚为什么这事儿这么重要。现在的大语言模型处理长文本时,每生成一个新词,都要把前面所有词的信息都过一遍,这个过程叫注意力计算。文本越长,这个计算量就越大,而且是平方级增长。一万字的文档和十万字的文档,计算成本可能差一百倍。
**注意力稀疏化*:一种给大模型"减负"的技术,核心思路是不再让每个新词都去看前面所有的词,而是只挑一小部分最相关的历史内容来看,从而大幅降低计算量。
这就像你在图书馆写论文,参考文献浩如烟海,你不可能把每一本书都从头翻到尾。聪明的做法是先找出最相关的几十页,集中精力读透它们。问题是,谁来决定哪几十页是"最相关"的?如果找的人眼光不准,挑出来的都是无关紧要的内容,那后面读得再认真也白搭。
一个训练不动的"选择器"
现在市面上做注意力稀疏化的方法,大多依赖一个轻量级的"选择器"模块,给每个文本片段打分,然后按分数高低选出前K个,这个K就是所谓的"注意力预算",也就是你愿意花多少计算资源去关注历史内容。
**Top-K选择*:一种排序后取前K个最高分选项的操作,因为它是"离散"的(选中就是选中,没选中就是没选中,中间没有过渡状态),所以没法用常规的梯度下降方法去优化。
问题就出在这个"离散"上。想象你在教一个新员工筛选简历,你希望他学会判断谁最适合这个岗位。但如果你只告诉他"选中"或"不选中"这两种结果,却从不告诉他为什么某份简历比另一份好那么一点点,他永远学不会精细的判断力,只能靠瞎猜和运气积累经验。
在数学上,这意味着语言模型的最终损失函数没法通过反向传播,把有用的梯度信号送回选择器那里去更新它。选择器就像一个瞎子摸象的裁判,它打出的分数和模型最终预测得好不好之间,压根没有直接的因果关系可以学习。
于是现有的方法想出了一个绕路的办法,用"蒸馏"来训练选择器。具体做法是让选择器去模仿原始密集注意力模型的分布,也就是说,教它去学习"原来的大模型在生成这个词时,主要看了前面哪些内容"。
这个办法乍一看很合理,但仔细想想会发现一个致命漏洞。
**排序错位*:论文里提出的一个核心概念,指选择器学到的"重要性排序"和"这些内容对模型最终预测结果的实际影响"之间存在的不一致。
这个错位到底体现在哪儿?打个比方,你让一个学生模仿老教授阅读论文时的眼球轨迹,教授的目光在哪些段落停留久,学生就跟着记住哪些段落重要。但教授阅读时的注意力分布,是在他能读遍全文、没有任何篇幅限制的情况下形成的。而学生只被允许精读其中一小部分,教授看得多但用得少的段落,可能恰恰是学生这有限篇幅里最不该浪费时间的地方。蒸馏学到的是"教授看哪里",而不是"在只能读少数几段的情况下,哪几段对答对题目最关键"。
这种错位还有两个更具体的副作用。第一,教师模型是逐层单独指导的,没法照顾到不同层之间的配合关系,可能某一层漏掉的信息刚好被另一层补上了,但蒸馏看不到这种互补性。第二,蒸馏只盯着注意力权重本身,却忽略了被关注的内容矩阵(也就是实际携带信息的那部分)对最终结果的影响,选中了正确的位置,不代表用好了那里面的信息。
SAS:把梯度硬生生"焊"回选择器身上
论文提出的方法叫SAS,思路说起来其实挺朴素的:与其教选择器去模仿别人的判断,不如让它直接从最终答案对不对这件事上学习。
具体怎么做?关键一步是把选择器打出的连续分数,直接塞进注意力计算的softmax公式里,以对数形式加在原始的注意力得分上。这样一来,梯度就能顺着这条路径,从最终的语言建模损失一路传回选择器,不再需要绕道模仿别人。
这个设计说起来简单,但真正落地却需要四个精心打磨的细节,缺一个都不行。研究团队通过大量对比实验,摸索出了这四条关键经验。
第一个细节是门控的位置:要放在softmax里面,而不是外面。
论文把这两种方式分别叫做"内部门控"和"外部门控"。外部门控是先算出正常的注意力权重分布,再拿选择器的分数去给结果做加权缩放。内部门控则是把选择器的分数当成一个"偏置项",直接加到softmax之前的原始得分上,这样它会参与到整个归一化的竞争过程里去。
这两者的差别,好比两种给餐厅打分的方式。外部门控相当于先让所有顾客各自打分出一份完整的满意度排行榜,然后你再拿着自己的喜好权重去乘一遍,但排行榜本身早已定型,你的偏好没法影响它的形成过程。内部门控则是让你的偏好从一开始就参与到排名的博弈里,你喜欢的餐厅会挤掉你不喜欢的餐厅的名次,而不仅仅是在最后加个滤镜。论文的梯度推导也证实了这一点:内部门控能提供"哪个块比哪个块更该被关注"的相对信号,外部门控只能在固定权重上小修小补,学不到这种此消彬长的博弈关系。实验数据也很直接,内部门控设计训练一个epoch后能达到54.4分的准确率,外部门控只有41.6分,差了整整13个百分点。
第二个细节是激活函数要用softmax做归一化,而不是sigmoid或者干脆不加处理。
**门控激活函数*:把选择器打出的原始分数转换成正式门控值的数学函数,论文比较了三种做法,归一化的softmax、独立判断的sigmoid、以及完全不处理直接注入。
为什么归一化这么关键?因为论文里设计了一个巧妙的机制:永远保留"当前正在处理的这一块内容",给它一个固定不变的门控值1,而历史内容的门控值则需要和这个基准做比较。如果用softmax,所有历史块的分数会被放在同一个"总量"里做竞争性分配,这个总量会反过来校准历史内容整体上该占多大的注意力份额,相对于永远稳赢的当前块而言。但如果用sigmoid,每个历史块的门控值是独立算出来的,互相之间没有博弈关系。论文的实验观察到一个很有意思的现象,训练久了以后,sigmoid门控会集体逐渐饱和,都往1这个方向靠拢,这就等于所有历史块都变得和当前块差不多重要,选择器实际上失去了区分度,退化成了"什么都要看一眼"的无门控注意力。原始分数直接注入的做法也有类似问题,但表现是相反的,分数会集体塌缩向0,方差越来越小,同样是丧失了区分能力。
这就像评委打分,如果每个评委各自打分互不参照(sigmoid),久而久之大家会趋同,变成谁都给差不多的高分,失去区分度。而如果强制所有评委必须在总分100分里分配名额(softmax),那么谁给谁打高分就必然意味着别人分数被挤压,这种此消彬长的压力会一直维持选手之间清晰的排名差距。
第三个细节是要保留连续的分数,而不是训练时就直接用二值化的硬掩码。
研究团队尝试过一种叫"直通估计器"(一种让离散操作在反向传播时假装是连续操作的技巧)的方案,前向传播时依然做硬性的Top-K筛选,只是反向传播时假装梯度可以直接穿过去。但这种做法有个数学上的隐患:被剔除的那些块,它们的注意力权重分母,只由被选中的那些块组成,一旦某个没被选中的块实际得分反而更高,它对应的梯度就会失控地暴涨,因为公式里的指数运算是没有上界的。软门控则天然把所有块都纳入同一个归一化分母里,权重永远被限制在1以内,梯度更稳定。实验里也验证了这点,离散方案训练100步时准确率一度冲到46.5,看着挺唬人,但后面反而掉到46.0,不如连续方案最终收敛到的54.4分,说明前期看起来猛烈的梯度信号其实是噪声,不是有效信息。
第四个细节则纯粹是为了效率考虑:训练时只更新被选中的块,而不是所有的候选块都参与梯度计算。
论文管这个叫"训练范围"。全范围训练里,每个未被选中的块都能拿到自己独立的梯度反馈,告诉它"你这次内容不够好所以没选上"。而稀疏范围训练里,未被选中的块拿不到这种针对自身内容的直接反馈,只能通过被选中块的间接归一化效应,被动地跟着调整一点点。这就好比公司年底评优,全范围训练相当于每个没获奖的员工都收到一份详细反馈,告诉他具体哪里做得不够;稀疏范围训练则是只有获奖者被公开表扬,没获奖的人只能从获奖者被表扬这件事上隐约猜测自己差在哪儿,猜得没那么准,但胜在成本低很多。实验证明,虽然稀疏范围一开始收敛慢一些,但训练足够久之后,两者最终效果几乎打平,54.8分对54.4分。考虑到训练成本能省下不少,论文最终选择了稀疏范围作为默认方案。
Triton内核:给长文本训练搭一座桥
光有设计思路还不够,真要把这套方法用到几万字的长文本训练上,还得解决一个纯工程问题。
**FlashAttention*:一种通过分块流式处理来加速并节省内存的注意力计算方式,被广泛用作各类大模型高效注意力实现的基础。
如果照着最直白的写法,先把带有门控偏置的完整注意力矩阵算出来,再做后续处理,长文本场景下这个矩阵会大到内存根本装不下。团队专门开发了一个融合了门控注入逻辑的Triton内核(Triton是一种编写高性能GPU计算代码的编程语言),把选择器的对数门控值直接融合进FlashAttention原本分块滚动扫描的计算流程里,扫到哪个块就顺手把门控偏置加进去、顺手判断这个块该不该被跳过,全程不需要把完整的注意力矩阵摊开存在内存里。反向传播时同样采用这种分块重算的策略,把每个被选中块内所有token的梯度累加起来,得到这个块对应的门控梯度。
这套内核设计使得SAS的训练成本变得可控,训练时可以处理长达32768个token的序列。
实验结果:数字不会说谎
理论说得再漂亮,最终还是要看跑分。研究团队在Qwen3的4B、8B、14B三种规模的模型上做了系统性对比,测试涵盖数学推理、通用知识问答、长文本理解和智能体任务四个维度。
在推理类任务上,SAS在几乎所有预算和模型规模下都超过了对照组SeerAttention-R(一种同样使用AttnGate选择器但通过蒸馏训练的强基线方法)。差距最悬殊的地方出现在预算特别紧张的时候,比如在2048 token预算下的AIME24竞赛数学题测试中,4B模型上SAS比SeerAttention-R高出13个百分点。
| 预算 | 方法 | MATH500(4B) | GPQA-Diamond(4B) | AIME24(4B) | AIME25(4B) |
| 1024 | SeerAttn-R | 84.67 | 39.84 | - | - |
| 1024 | SAS | **90.65** | **50.41** | - | - |
| 2048 | SeerAttn-R | 91.85 | 49.94 | 55.83 | 45.16 |
| 2048 | SAS | **93.47** | **54.86** | **68.85** | **56.38** |
| 4096 | SeerAttn-R | 94.10 | 55.40 | 69.32 | 58.59 |
| 4096 | SAS | **93.67** | **55.11** | **71.72** | **59.97** |
在预算4096的情况下,SAS在AIME24上的成绩(71.72)甚至超过了完全不做稀疏化的满血密集注意力模型(71.25),这意味着只用一小部分计算量,反而比全量计算做得更好。
长文本理解任务上,一个特别值得说的现象是,选择器完全是用数学题训练出来的,却被直接拿去测试和数学毫不相关的长文档理解任务,结果依然全面超过对照方法,说明这种从最终目标反推出来的选择能力,比死记硬背某个具体分布更容易迁移到陌生场景。在预算2048、14B模型的情况下,处理超过8000字的长文档时,SAS的表现比对照方法高出2.4分。
智能体任务上,涉及多轮工具调用的BFCL测试中,SAS在所有模型规模和预算设置下都领先,4B模型在2048预算下领先3.5分。在更复杂的VitaBench测试里,预算4096时SAS在外卖配送、到店消费、在线旅行三个真实业务场景下大多数指标都优于对照方法,非常接近满血密集注意力的表现。
为什么SAS学到的选择更好用
跑分好只是结果,团队还想搞清楚背后的原因,这部分分析读起来挺有意思。
一个反直觉的发现是,单看每一层,SAS选中的内容覆盖的注意力权重总量反而比蒸馏方法更少。按理说覆盖的权重越多越"忠实于原模型",可SAS明明分数更好,这是为什么?
答案藏在跨层协作里。团队把每个方法在所有层选中的块取并集,再去对比这个并集和满血密集注意力模型实际关注内容的重合程度,发现SAS的并集重合度明显更高。这说明SAS虽然在单层上显得"保守",但不同层之间学会了分工合作,你负责看这部分,我负责看那部分,组合起来反而覆盖得更全,而蒸馏方法是逐层独立模仿,层与层之间没有这种默契。
这就像一个团队做项目调研,蒸馏方法相当于让每个人单独按照同一份模板去覆盖尽可能多的资料,结果大家覆盖的内容高度重叠,遗漏的盲区也差不多;而端到端训练出来的分工方式,更像是团队成员事先商量好各自主攻不同方向,单个人负责的范围看起来窄了,但因为分工明确不重复劳动,整个团队实际触达的信息面反而更广。
另一个有意思的发现是关于生成效率的。研究团队观察到,用SAS训练出来的模型在做推理任务时,生成的答案长度更短,被截断(也就是没在规定长度内答完)的概率也更低。这在AIME等高难度数学题上体现得尤其明显。这背后的逻辑不难理解,如果模型没能拿到关键的上下文信息,它往往会陷入漫无目的的长篇推理,反复兜圈子试图找到答案,而拿到了真正有用的上下文,模型能更快地锁定思路,一步到位。
在实际部署效率方面,团队把这套方法接入了SGLang推理框架进行了端到端测速。在单条请求、批量为1的情况下,8000字左右的短文本两种方法速度相当,但拉长到51.2万字的超长文本后,SAS的解码延迟只有密集注意力的不到五分之一,快了5.6倍。如果把批量提高到8,在6.4万字长度下加速比可以冲到13倍左右。
不过团队也做了一个细致的成本拆解,发现随着文本长度继续变长,选择过程本身(给所有候选块打分、再做Top-K排序)反而会成为新的瓶颈。在51.2万字的极端场景下,选择阶段能占到整个解码步骤耗时的90%,真正的注意力计算反而只占一小部分。这提示了一个未来值得优化的方向,当文本长得离谱时,"怎么选"本身也得想办法加速。
这套方法还能用在预训练阶段吗
团队还额外做了一个验证性实验,看看这套端到端训练思路能不能不局限于"后训练"阶段,而是直接融入模型从头训练的过程。
他们在OLMo3-7B这个开源基座模型上进行了大约500亿token规模的继续预训练,这次不再冻结主干模型,而是让主干网络和选择器一起联合优化。结果显示,SAS在通用知识、数学、代码等多项下游任务上取得了参与对比模型中最好的平均分43.28,非常接近完全不做稀疏化的原始密集模型(43.88),并且在长文本理解测试LongBench上,SAS达到了并列最高的30.0分,明显超过密集基座模型的29.0分,尤其是在超过8000字的长输入场景中优势最为突出。
这个结果提供了一个初步但重要的信号,端到端的稀疏训练思路,不仅能用来给已经训练好的模型"打补丁",似乎也有潜力从一开始就融入模型的正式训练流程里。
写在后面
读完这篇论文,最触动我的其实不是那几个百分点的提升,而是那个"排序错位"的洞察本身。
我们平时做很多事情,都习惯了用"模仿最好的做法"去替代"直接优化目标",这是一种很自然的思维捷径,因为模仿的目标是清晰可见的,而直接优化往往涉及一个不可导、卡壳、无法反向传播的中间环节。这篇论文提醒我,绕过这个卡壳环节固然省事,但绕过去的代价,是学到的东西可能压根就不是你真正想要的。
论文里那个关于门控饱和的图(图3)也让我印象很深。你会看到,如果激活函数选错了,不管是sigmoid还是直接注入原始分数,训练久了以后所有历史块的门控值都会趋同,要么全都逼近1,要么全都塌缩到0附近,方差越来越小。这其实是一种"隐性的退化",模型看起来还在训练、损失还在下降,但选择器实际上已经丧失了区分能力,变成了一个摆设。这种退化不会在损失曲线上直接报警,你得专门去看门控值的分布才能发现问题。这提醒我一个更普遍的道理,有时候一个系统看起来在正常运转,指标也在变好,但某个关键部件可能早已经悄悄失去了它本该发挥的作用。
论文里还留了一个没解决的问题,在RULER这种专门测试超长上下文检索能力的基准上,SAS虽然比对照方法有进步,但随着文本长度拉到十几万字,性能依然会显著下滑,离密集注意力还有明显差距。团队把这个归因于"块摘要"这种压缩表示方式,天然难以保留那种散落在长文本某个角落的细粒度信号,比如那种"大海捞针"式的关键线索。这算是留给后来者的一道题,怎么在保持效率的前提下,让选择器不丢失那些藏得很深的细节。
Q&A
Q1:SAS和SeerAttention-R这两种注意力稀疏化方法最本质的区别是什么?
A:两者用的选择器结构完全相同,唯一区别在于训练方式。SeerAttention-R通过蒸馏让选择器模仿原始密集注意力模型的分布,SAS则把选择器分数直接注入注意力计算,让语言建模损失通过标准反向传播直接更新选择器,不再需要模仿老师模型。
Q2:SAS在计算资源紧张的情况下表现如何?
A:SAS在低预算场景下优势最明显。比如在2048 token预算下,Qwen3-4B模型的AIME24竞赛数学测试中,SAS比对照方法SeerAttention-R高出13个百分点,说明预算越紧张,准确的内容排序就越关键。
Q3:SAS这种方法目前还有哪些局限性?
A:在超长文本检索测试RULER上,随着文本长度增长到十几万字,SAS性能依然会明显下滑,和满血密集注意力还有差距。研究团队认为这可能与块摘要压缩方式难以保留细粒度的"大海捞针"式线索有关,是未来需要优化的方向。







