让注意力自己决定该算什么：一篇关于MassAlloc Attention的解读

你有没有想过一个问题：当你读一本书，读到第300页的时候，你的大脑真的会去"回顾"第一页写了什么吗？

大概不会。你的注意力早就把大部分早期内容压缩成了模糊的印象，只有那些真正重要的情节点还清晰地留在脑子里，比如凶手是谁、主角的动机是什么。可是，如果你是一台正在处理这本书的机器,现在的AI大模型在处理长文本时，情况却完全不同。

现在的主流大模型，不管你给它读的是300页的书还是30万字的代码库，它都要老老实实地把每一个字和之前所有的字重新算一遍关系。哪怕这个字和当前要处理的内容毫无关系，它也得算。这就是这篇论文要解决的问题，而它给出的答案，简单又聪明：让注意力机制自己决定，哪些计算值得做，哪些可以跳过。

这篇论文来自香港科技大学(广州)和北京智源研究院的团队，名字叫MassAlloc Attention，缩写MALA。

**核心问题：算力被浪费在了"没用"的地方**

要理解这篇论文在解决什么，得先搞清楚现在的注意力机制到底是怎么工作的。

Transformer模型*：一种目前主流大语言模型采用的神经网络架构，其核心机制叫"注意力"，用来判断当前处理的这个词，应该和前面哪些词产生关联。

当模型读到句子里的第10000个字，它需要去看前面9999个字里哪些和它相关，然后计算出一个关联程度，这个过程叫打分。打完分之后，还要做归一化、加权求和，把这些前文信息汇总成一个结果，这个后半段流程，论文里称为"后打分计算"(post-score computation)。

问题就出在这里：模型必须对所有9999个字都完整跑完这个后半段流程，即便某些字和当前处理的内容几乎没有任何关系。学术界早就观察到,注意力权重的分布极不均匀，往往只有局部区域、少数几个"锚点"词、或者一小部分和检索相关的位置才拿到了真正有意义的权重，剩下的绝大多数位置分到的权重几乎可以忽略不计。

FlashAttention*：由斯坦福团队提出的注意力计算加速方法，通过分块处理和在线归一化技术减少了内存读写开销，是目前几乎所有大模型训练和推理都会用到的底层技术。

FlashAttention确实让注意力计算的速度和内存效率大幅提升，但它改进的是"怎么算得更快"，没有改变"该不该算"这件事。换句话说，哪怕某个词的注意力权重小到可以忽略，FlashAttention依然会一丝不苟地把它的softmax更新、数值加载、加权求和这一整套流程走完。这就好比你请了一个动作特别快的搬运工，他确实能用最短时间把货物从A搬到B，但如果这批货物根本不需要搬,他还是会照搬不误。

这就是为什么MALA的作者们会问：既然注意力机制自己已经在打分的时候产生了那些"权重很小"的信号，为什么不让它自己根据这些信号来决定后面要不要接着算？

**MALA的核心设计：打分照做，后续按需**

MALA的第一个关键决定，是绝不省略打分这一步。

每一个在因果关系上合法的位置（也就是当前词能看到的历史位置），都必须先算出QK分数(Query和Key的匹配分数,反映两个位置之间的相关程度)。这一步论文里称为"QK score discovery"。

这个设计初看有点反直觉：你不是想省计算量吗？为什么打分这一步一点都不省？

因为如果连打分都跳过，你就永远不知道跳过的这个位置到底重不重要。这就像考试判卷,你不能说"这道题看起来很难，我们直接给0分不批了"，你必须先看一眼答案，才能判断这道题到底值不值得仔细批。MALA的逻辑是：打分的成本远比后面加权求和的成本低，所以先把所有分数都算出来，再根据分数决定哪些值得投入更多计算。

打分做完之后，才是MALA真正省钱的地方：后打分计算的执行与否，由这个分数在整体分布里占的比重决定。

具体怎么判断"占比"够不够格？这里引入了一个叫容忍度的概念。

容忍度τ*：一个人为设定的阈值参数，用来控制"多小的注意力贡献可以被忽略"，MALA的实验里统一设为1。

论文的做法是，先设定一个基准：如果注意力对所有历史位置一视同仁（也就是均匀分布），那么每个位置应该分到1/L的权重，这里L是当前位置能看到的历史长度。然后用容忍度τ去缩放这个基准，得到一个和位置数量、序列长度都无关的统一门槛：τ/L。

任何一个候选的计算块，如果它里面权重最大的那个位置的贡献比例都低于这个门槛，那么整个块就被跳过，不再执行后面的softmax更新、数值加载和加权求和。

这套逻辑在前向计算（也就是模型做推理或者训练时正向跑一遍）和反向计算（也就是训练时用来算梯度、更新参数的那一遍）里，用的信号略有不同。前向用的是不断累积中的"在线归一化系数"，反向用的是前向算完之后保存下来的"最终归一化系数"。这两者的关系保证了一件事：反向阶段能省掉的计算，一定是前向阶段已经省掉的计算的超集。也就是说，反向可以比前向更激进地跳过一些块，但绝不会跳过前向已经保留的块。

这个设计带来一个数学上可以证明的保证：论文推导出，被忽略的注意力概率质量总量，有一个明确的上界，由容忍度τ决定，公式是τ/(1+τ)。当τ设为1时，这个上界大约是50%,但实际测量到的真实遗漏质量远远低于这个理论上限，通常在0.02%这个量级,这说明理论保证是保守的安全网，而实际运行远比最坏情况乐观。

**为什么这套方案值得信任：从"纸面公平"到"实际验证"**

光有理论上界还不够，论文做了一个特别巧妙的对照实验，用来证明MALA的在线决策到底有多接近"理想情况"。

研究者设计了三种参照方案，全部使用完全"作弊"的信息:也就是提前跑完整个流程之后，拿到的最终、最准确的参照概率质量，用这个信息去反过来决定哪些块该保留。第一种按位置分配，所有输入、所有层、所有注意力头共享同一套规则；第二种按层和头精细化，但同一层同一头内所有输入统一对待；第三种是最细粒度的，按每一个具体输入、每一层、每一个头单独分配。

三种方案的计算总量被严格控制成完全相同，都控制在平均每个查询大约1024个位置的规模。

结果是什么？按位置分配的方案效果最差，遗漏质量高达0.6%；按层头精细化之后好了不少，降到0.17%；而按每个具体输入单独分配的第三种方案，遗漏质量降到了0.0182%，这是理论上能达到的最好水平，因为它用的是"事后已知答案"的信息。

MALA自己呢？它完全没有用到任何"事后信息"，只靠打分那一刻在线计算出来的归一化系数，遗漏质量做到了0.0188%,和那个理论最优值只差了0.0006个百分点。

这个对比说明了什么？它证明了注意力自身在打分过程中产生的信号，已经几乎包含了做出最优决策所需要的全部信息。你不需要额外的路由器、额外的选择器，也不需要提前跑一遍来"偷看答案"，注意力打分本身就是最好的向导。这就好比一个经验丰富的医生，光靠听诊器听几声心跳,就能大致判断出病情的轻重，不需要把病人拉去做全套的CT和血液检测才敢下结论。如果非要等所有检查全做完才敢做初步判断,那这个"判断"的意义就大打折扣了,因为诊断的价值恰恰在于用最少的信息尽快做出足够准的决策。

**跨长度的稳定性：一套容忍度，管住所有场景**

如果MALA只在某个特定长度下表现好，那价值就有限。真正重要的是，同一个容忍度τ=1，能不能在从1千到3万2千token的各种长度下都保持可靠。

论文做了系统性的验证，对比"正常跑MALA"和"关掉筛选、完整跑一遍"这两种情况下，输出结果和梯度的误差有多大。

结果显示，随着上下文长度从1K涨到32K（涨了32倍），前向阶段平均需要执行的计算量从470个位置逐渐涨到1066个位置，反向阶段也从462涨到1058，但超过8K之后就基本稳定在1010到1066这个区间,不再随长度线性增长。

这个现象值得多说一句。按理说,文本越长，历史信息越多,该保留的计算量应该跟着涨才对。但数据显示涨到一定程度后就饱和了，这恰好印证了论文最初的观察:注意力的有效信息始终集中在一个相对固定规模的"核心区域"里,不管背景文本多长,真正被反复"翻查"的关键内容量级是相对稳定的。这就像一本侦探小说不管写了500页还是1000页，真正的线索可能就藏在其中固定的十几个关键情节里，剩下的篇幅越多只是背景铺垫,不需要额外增加"追查线索"的精力。

误差方面，在所有测试长度下，平均遗漏概率质量最高不超过0.0062%,输出结果的相对误差最高不超过0.021%，梯度误差（这个决定了训练是否稳定）最高也只在0.38%左右。这些数字都很小，意味着MALA跑出来的结果，和"完整跑一遍不省任何计算"的黄金标准几乎没有区别。

**长距离关联测试：能不能记住"藏得很深"的信息**

前面讲的都是数值上的精度，但更让人关心的问题是：这种"跳过部分计算"的做法，会不会让模型丢掉一些本该记住的重要关联？

论文设计了一个叫"受控关联召回"的测试。做法是：先在文本开头塞进256组键值对，也就是256组"如果看到A就该联想到B"的绑定关系，然后在后面随机的位置提出问题，看模型能不能准确地把A和B对应上。这个任务的特殊之处在于,它没有额外掺杂干扰噪音,纯粹考验模型能不能在长文本里精准定位到之前记下的绑定关系。

对比的基线方法里，有几种是"固定预算"式的稀疏注意力，比如MoBA、NSA、DSA，它们提前设定好一个位置数量上限（比如1024个槙位），不管内容分布如何，都严格卡在这个数字。

在序列长度涨到8192、模型维度设为512的极限测试里，全量注意力(FullAttn，也就是不做任何裁剪的传统方法)拿到89.97%的准确率，而MALA紧跟其后，拿到89.67%，差距不到0.3个百分点。反观那些固定预算的稀疏方法，DSA只有52.61%，MoBA只有47.25%，NSA更是只剩22.61%。

这个结果的分量值得细品。差不多的计算预算规模下，MALA几乎完整保留了长距离关联能力,而那些"死板"按固定名额分配位置的方法，丢失了将近一半甚至更多的关联信息。原因很直观：固定预算方法是提前决定"我只看这么多个位置"，不管这些位置里到底藏没藏着关键线索；而MALA是让注意力自己去发现哪里藏着线索，然后按需分配。这就好比考试限时找答案，一种做法是不管题目内容,规定自己只能翻10页书,翻完就必须停手；另一种做法是先快速扫一眼整本书的目录和关键词,再决定重点翻哪几页。显然后者更容易找到真正的答案，即便两种做法花的总时间差不多。

**真机效果：延迟和内存双双改善**

理论精度和测试集准确率固然重要，但工程落地才是这类研究最终要接受的考验。论文在8块H100 GPU上，用张量并行(把一个模型拆到多块GPU上分摊计算的一种并行策略)的方式，测试了128K超长上下文下的真实运行表现。

结果是：训练阶段的前向计算延迟降低了2.2倍，反向计算延迟降低了3.0倍;推理阶段的解码延迟降低了1.6倍，相比传统的全量注意力方法。更重要的是，这一切改善的同时，每张GPU的峰值显存占用和全量注意力基本持平，没有增加额外的内存负担。

这个"内存不增加"的结果其实不是理所当然的。很多稀疏注意力方案为了决定"该保留哪些位置"，需要额外存储一个选择掩码，或者维护一份路由表，这些辅助结构本身就要占用显存。而MALA的巧妙之处在于，它所有的决策都直接复用注意力计算本身已经产生的状态（比如归一化系数），不额外存任何选择结果，所以显存开销和传统方法几乎一样。这就像一个真正高效的整理师，收纳东西的时候不需要另外准备一个"分类清单"本子，他脑子里记着规则，动手的时候直接按规则分类，不产生额外的纸面记录负担。

**规模化训练：从6亿参数到140亿参数都适用**

工程效率的提升，如果只在小模型上成立，说服力有限。论文进一步在0.6B到14B参数规模的一系列模型上做了系统的规模化训练实验，用了128块GPU。

结果显示，在4K长度的预训练阶段和32K长度的长文本训练阶段，MALA的困惑度(衡量语言模型预测准确性的核心指标，数值越低说明模型越"确定"该说什么，代表性能越好)曲线几乎和全量注意力方法完全贴合。

具体到140亿参数这个规模，两种方法训练完之后的困惑度差异小于0.001，几乎可以忽略。但计算量上却有实打实的节省：4K预训练阶段节省了2.5%的总训练算力，32K长文本训练阶段节省了23.1%。

为什么长文本阶段节省的比例明显更高？答案其实呼应了前面提到的"计算量饱和"现象:文本越长，传统方法需要处理的历史位置越多，浪费在低贡献区域上的计算也越多，而MALA能省下来的部分正好随之增加。

**模型能力测试：知识、推理、检索都没掉链子**

最后一步，是把这套方法真正训练出来的模型拿去做全面的能力评测，看看它在实际任务上表现如何，而不只是看训练过程中的数字指标。

论文分别评测了14B规模的四种注意力变体(全量注意力、MoBA、DSA、MALA)，以及在32B规模上从同一个基础模型继续训练出来的全量注意力和MALA两个版本。

评测覆盖知识类任务（比如MMLU、常识问答）、推理类任务（比如数学题、逻辑推理）、以及长文本检索任务（RULER基准，在32K原生长度和通过YaRN技术扩展到128K的场景下分别测试）。

数据显示，14B规模上，MALA的知识平均分是72.48，全量注意力是72.32；推理平均分MALA是64.66，全量注意力是64.46。32B规模上，MALA知识分75.75对全量注意力75.62，推理分76.10对75.67。长文本检索方面，14B原生32K下MALA拿89.45分，全量注意力89.42分；扩展到128K后，MALA是65.75分，全量注意力65.84分，差距同样很小。

这些数字放在一起看，传达的信息很清楚：省下来的计算，没有换来能力上的明显折损。这在机器学习领域是一件很难得的事情，通常"更快"和"更强"之间存在权衡，而这里的权衡幅度小到几乎测不出来。

**写在后面**

读完这篇论文，最让我觉得有意思的一点，其实是它对"稀疏"这个词的重新定义。

过去很多稀疏注意力的做法，本质上是"提前决定不看",不管是固定窗口、固定块、还是训练一个路由器去预测该看哪里，这些方法的共同特点是决策发生在真正看到内容之前。而MALA反过来，它坚持先把每一处都完整地看一眼（也就是打分这一步绝不省），只是看完之后再决定要不要深入处理。这种"先扫一眼再决定深挖"的思路，其实更接近人类阅读的真实过程,你翻书的时候眼睛确实扫过了每一行,只是大脑对大部分内容一扫而过，只在关键情节上真正停下来细读。

另一个值得单独说一说的细节，是论文里那个"匹配工作量"的对照实验设计。研究者没有满足于"MALA比基线快"这种笼统的结论，而是专门设计了一组严格控制计算总量完全相等的对照，去证明性能提升不是来自"少干活"，而是来自"把活干在了该干的地方"。这种实验设计的严谨性，其实比结果本身更值得学习:很多时候一个方法看起来更好，未必是因为它更聪明，可能只是因为它悄悄做的活更少。这篇论文把这种可能性提前排除掉了。

这篇论文没有解决的问题，可能在于它依然保留了完整的打分计算,也就是那部分随序列长度平方增长的开销一点没减。如果未来能找到办法，在打分阶段本身也引入类似的自适应机制,那节省的空间可能会更大。但这或许也是一种取舍：打分的信息是决策的基础，一旦在这一步也开始"猜"，整套方案建立的信任基础可能就动摇了。

Q&A

Q1：MassAlloc Attention是什么？

A：MassAlloc Attention（简称MALA）是一种新型注意力机制，它让模型先对所有历史位置完整打分，再根据打分结果自动决定哪些位置值得投入后续计算，从而减少低贡献区域的算力浪费，同时保持和传统全量注意力接近的效果。

Q2：MassAlloc Attention会不会降低模型的准确率？

A：实验显示影响很小。在14B参数规模下，知识、推理任务平均分与全量注意力差距在0.2分以内，32K长文本检索分数几乎持平，128K扩展场景下差距也控制在1分左右，同时训练算力最多能省下23.1%。

Q3：MassAlloc Attention相比其他稀疏注意力方法有什么优势？

A：它不依赖提前设定固定的位置预算，而是让注意力自己根据实际内容分布决定计算量。在长距离关联测试中，MALA准确率接近全量注意力的89.97%，而固定预算的MoBA、DSA、NSA等方法准确率明显更低，最低跌到22.61%。