稠密注意力有个结构性浪费:softmax 归一化后,大部分因果分数空间拿到的质量小到可以忽略,但稠密 kernel 仍要为每个 QK tile 走完打分后的整条流水线——归一化、乘 V、写回。分数算完才发现这块区域不重要,但计算已经花掉。BAAI 团队 9 月 26 日挂在 arXiv 的新论文 MassAlloc Attention(MALA)盯上的正是这段被浪费的 post-score 计算,数日冲上 Hugging Face Daily Papers 榜首,累计 763 赞。
QK 全算,后面按质量分配
MALA 的思路分两半。第一半,不砍 QK 打分——每个合法因果交互保留 score 访问,守住注意力分布的完整性。第二半才是重点:打分后的计算不再全员执行,而是按归一化贡献分配。前向时,MALA 借用演化中的 online-softmax normalizer 实时估计各区域质量占比,低贡献区域的后续计算直接跳过;反向复用定稿的 normalizer 推导嵌套保留支撑,只用标准注意力状态。一个统一的 tolerance 参数同时治理训练和推理——跳多少、留多少,两边同一把尺子。
提速多少:反传 3 倍,训练 FLOPs 省 23%
数字相当扎实。8K matched-work 研究里,总 post-score 工作量对齐后,MALA 的平均省略质量是 0.0188%,逐实例最优的 reference-mass oracle 是 0.0182%,几乎贴着理论上限跑。关联记忆任务上,MALA 在 8K 拿到 89.67% 准确率,全量注意力 FullAttn 是 89.97%,差 0.3 个百分点。工程侧,128K 张量并行算子基准里,MALA 把训练前向延迟压到 1/2.2、反向压到 1/3.0,推理解码快 1.6 倍;作者在 Hugging Face 社区帖补充测试环境为 8×H100、TP=8。规模验证:0.6B 到 14B 的 scaling-law 训练中困惑度全程贴近 FullAttn,总 FLOPs 下降;14B 在 32K 上下文训练省 23.1% 总 FLOPs,另有 32B 继续训练模型,知识、推理、长上下文检索与 FullAttn 相当。
和稀疏注意力不是一条路
这条赛道不缺新面孔:LEMA 把 KV 缓存挪出显存,HyQuant 做注意力混合精度量化,Video DeltaNet 给视频生成换混合注意力,都在啃同一个问题。MALA 的差异点是不预设稀疏结构:固定窗口、top-k 要人事先告诉模型稀疏在哪,MALA 把决定权交还给 softmax 统计本身——分布已经告诉你哪里重要,照着分配就行。另一个容易被忽略的细节是反向提速(3.0 倍)大于前向(2.2 倍),这本质是训练场算法,对预训练成本占大头的大厂,23% 的 FLOPs 直接换算成真金白银。
冷水也得泼
论文页 Models citing this paper 为 0,Hugging Face 页面也没放出代码仓库链接,第三方复现暂时无从谈起——目前所有基准数字都是作者自报。tolerance 超参的敏感性、跨架构迁移性,要等工件落地才有答案。同团队还挂了姊妹篇 CoWindow Attention(arXiv:2609.32704),让多个注意力头分担历史访问,14B 训练省 28.5% FLOPs,两篇对照可看清 BAAI 的完整布局。
注意力计算预算的分配权,正从人拍板的固定规则移交给分布自己的统计量。如果 MALA 的 tolerance 机制进主流训练框架,长上下文训练的成本曲线会再弯一次——下一次 128K 级预训练的账单,可能就少一截。
参考:arxiv.org/abs/2609.32712