
你有没有想过,一个公司如果真要做到信息透明,是不是必须让每个员工都知道所有客户的所有历史记录?
大部分管理者的直觉会说:不需要。让财务部门记住财务的事,让客服部门记住客服的事,只要信息在整个公司范围内都能被查到,就够了。没人要求每个人脑子里塞进全公司的账本。
但现在最火的大语言模型,处理超长文本的时候,干的恰恰是那件"每个人都记住全部"的事。这篇来自香港科技大学(广州)和北京智源人工智能研究院的论文,就是在问一个看起来简单却被忽略了很久的问题:真的有必要让每一个注意力头都盯着全部历史吗?
核心问题:越长越贵,贵在哪
要理解这篇论文在解决什么麻烦,得先说清楚大模型是怎么"记住"上下文的。
Transformer模型*:目前主流大语言模型的基础架构,靠"注意力机制"让模型在生成每个词的时候,去参考前面所有已经出现过的词。
这套机制里有个叫FullAttn的东西。
FullAttn(全注意力)*:模型里的每一个注意力头,在处理当前位置的词时,都要和前面所有历史词计算一遍相关性。
问题就出在这个"每一个头都要"上。现在的大模型动辄要处理十几万甚至几十万个token的上下文,如果每个注意力头都要把这几十万个历史位置全部扫一遍,计算量会随着文本长度的平方增长。文本长度翻一倍,计算量翻四倍。这不是线性的麻烦,是滚雪球式的麻烦。
FlashAttention这类工程优化,解决的是"怎么把这个计算做得更快、内存读写更省",但它没有改变一件事:每个头依然要看完整的历史。就像一个仓库管理系统,哪怕你把货架排列得再高效,只要规定每个员工进门都要把整个仓库的货物清点一遍,效率的天花板早就被这条规定锁死了。
论文作者们的直觉来自一个真实的观察:不同的注意力头其实并不是在干同一件事。有的头专门盯着最近几个词,有的头喜欢死死记住开头的信息,有的头则对中间某段历史特别敏感。既然头和头本来就分工不同,那为什么还要强迫每个头都拥有访问全部历史的能力?
这就是这篇论文的起点,一个叫做CoWindow Attention(CoWA)的新设计。
CoWA的设计:把历史分给不同的头去看
CoWA的核心思路可以用一句话概括:让每个注意力头只负责一部分历史,但让所有头合起来依然能覆盖全部历史。
论文管这个性质叫"集体覆盖"(collective coverage)。这里要理解一个前提概念。
KV头*:在分组查询注意力(GQA)机制里,多个查询头会共享同一组键值(Key-Value)表示,这组共享的键值单元就叫KV头,是分配注意力窗口的基本单位。
具体怎么分?CoWA给每个KV头分配三种窗口。第一种是"近邻窗口",也就是最近的一段文本,所有头都共享,保证谁都不会忘记刚发生的事。第二种是"前缀锚点窗口",专门盯着文本最开头的一部分,同样所有头共享,这是因为研究发现很多模型会把序列开头当作一种稳定的参照点。第三种才是重点:"远程窗口",这部分不共享,而是把剩下的历史距离切成若干段,一段分给一个头。
所有头的远程窗口拼起来,正好覆盖了近邻窗口和锚点窗口之外的全部历史。没有一个位置被落下,但也没有哪个位置被所有头重复看一遍。
这就像小区物业安排保安巡逻。如果每个保安都被要求把整个小区从头到尾巡一遍,看起来最保险,但十个保安做的是十份重复劳动,真正出事的时候反而因为路线太长而反应慢。合理的做法是让每个保安负责一段固定区域,大家的巡逻范围拼起来覆盖整个小区,谁也不重复谁的路线。业主要找保安,肯定能在某个区域找到人,但不需要十个保安同时挤在同一段路上。如果不这样分工,保安数量增加带来的只是重复劳动的叠加,而不是覆盖范围的扩大,这正是FullAttn在长文本上遇到的瓶颈。
论文里还有一个细节值得注意:这套窗口划分规则不是靠模型自己学出来的,也不需要额外训练一个"路由器"去决定哪个头该看哪段历史。
路由器/学习型索引器*:一些稀疏注意力方法(如MoBA、DSA)会训练一个小型神经网络,动态决定每个查询应该关注哪些历史token,这个网络就是路由器或索引器。
CoWA完全绕开了这一步,窗口的划分只由token的位置和KV头的编号决定,是一条写死的规则,训练和推理时用的是同一套规则,不会出现训练时一个样子、推理时又变一个样子的情况。这种"位置决定一切"的做法,代价是牺牲了一点灵活性,换来的好处是执行起来极其规整,不需要额外的排序、筛选、路由计算,这些恰恰是别的稀疏注意力方法里最烧算力的部分。
实验证据:拆开看,到底是"窗口小"重要还是"分工"重要
光有设计思路还不够,论文做了一个很讲究的对照实验,来确认到底是什么在起作用。
研究者们担心一件事:CoWA表现好,会不会只是因为它的窗口设置得合适,跟"不同头看不同历史段"这个分工思想没什么关系?
为了把这两个因素拆开,他们做了一个"窗口宽度完全相同,但分工程度不同"的对比。在8个KV头的设置下,让1个、2个、4个、8个不同的远程窗口,重复分配给这8个头,窗口宽度全部保持一致。
结果很说明问题。
只用1种窗口、重复分给全部8个头时,8K长度下的召回准确率只有21.32%;当窗口种类增加到2种时准确率是32.92%;4种窗口时是52.32%;等到8个头各自拿到一个独一无二的窗口,也就是完整的CoWA设计,准确率一下跳到了89.73%,几乎追平了FullAttn的89.97%。
这组数字排开来看很有冲击力:从21.32%到89.73%,中间隔了68个百分点,而唯一变化的因素就是"重复的窗口"换成了"互补的窗口",窗口本身的宽度、位置样式全都没变。这说明真正起作用的不是窗口设计得多精巧,而是"让不同的头去看不同的历史"这个分工思路本身。
如果只是简单粗暴地缩小每个头的注意力范围(比如只保留最近的窗口,不给任何远程窗口),效果会更差。论文里对比的滑动窗口注意力(SWA,只保留近邻窗口,完全没有远程访问能力)在8K长度下的准确率只有5.34%。这也印证了一件事:光是"变稀疏"不够,必须"稀疏地分工合作"才行。
更大范围的检验:一场公平的记忆力考试
窗口对照实验证明了原理,但一个模型能不能打,还得看真刀真枪的对比。论文设计了一个叫"关联召回"的测试任务。
关联召回(associative recall)*:给模型看很多随机配对的键值信息,之后随机抽查,看模型能不能准确从上下文里找出对应的值,这个任务专门检验模型是不是真的"记得住"具体细节,而不是靠语义蒙对答案。
之所以用随机配对,是因为如果用有语义关联的内容,模型可能靠"猜"蒙对答案,随机配对堵死了这条捷径,逼着模型老老实实依赖注意力机制去查找。
研究者把FullAttn、SWA、Seer、InfLLMv2、NSA、MoBA、DSA和CoWA这八种注意力方案放在一起比,控制每种稀疏方法能看到的token预算大致相同,让比较公平。测试的序列长度从1024一路拉到8192,模型维度从64到512不等。
结果在短序列、小预算够用的情况下,大家表现都差不多。但一旦序列拉长到8192,预算变得相对紧张的时候,差距就显现出来了。CoWA在这个长度下拿到89.73%的准确率,几乎贴着FullAttn的89.97%在走;DSA和MoBA分别是53.71%和50.12%;NSA是25.23%;剩下的方法基本都掉到了10%左右。
这个差距意味着什么?意味着在长文本场景下,如果你用的是那些准确率只有10%左右的方法,模型基本上等于把大部分远程信息给忘了,只有CoWA还能保持接近满分的水准。而且论文特意提到,单纯加大模型维度并不能弥补这个差距,说明问题出在注意力机制本身的设计上,不是靠堆参数能糊弄过去的。
系统层面的收益:省下来的不只是理论数字
模型质量保住了,接下来要看的是这套设计到底能不能换来实实在在的速度和内存收益。
研究者在8张H100 GPU上,用张量并行的方式测试了不同注意力方案在128K序列长度下的表现。
张量并行(Tensor Parallelism, TP)*:把模型的计算任务拆分到多张显卡上同时进行,是训练和部署大模型时常用的分布式计算手段。
结果显示,CoWA把训练阶段的前向计算延迟降低到FullAttn的1/7.4,反向传播延迟降低到1/8.6,推理阶段的解码延迟降低到1/3.0。内存方面,训练阶段CoWA的峰值内存占用和FullAttn基本持平,但在解码阶段,CoWA的峰值内存只有FullAttn的1/7.6。
值得一提的是,作为对照的MoBA在128K长度下的反向传播直接内存溢出,没能跑完。这是因为MoBA虽然也做了稀疏化,但它额外需要做区块池化、路由打分、TopK筛选、数据重排等一系列操作,这些"辅助工作"本身也在吃内存和算力。DSA同样需要一个专门的轻量索引器给所有位置打分,再挑出重要的部分。相比之下,CoWA不需要任何打分和筛选的过程,它的可见范围直接由位置和KV头编号算出来,这个计算几乎不花时间。
这就好比两家快递公司,一家每次派件前都要先开个会讨论把这个包裹交给哪个快递员最合适,另一家直接按照固定的邮编分区规则派件,谁负责哪个片区提前写死在规则里。前者看起来更"智能",但开会本身就是成本,包裹越多,开会耗费的时间也越多;后者省掉了这个决策环节,规则简单粗暴,但胜在稳定高效。如果不提前定好分区规则,指望每次都临时决策最优分配,公司规模一大,光是决策的开销就能拖垮整个系统的效率,这正是MoBA和DSA相比CoWA多出来的负担。
从小模型到大模型:这套设计能不能扛得住规模化
一个方法在小规模实验里好用,不代表放到真正的大模型训练里也好用。论文做了一次从0.6B到14B参数、跨越五个模型规模的完整扩展律实验。
扩展律(scaling law)*:研究模型性能如何随参数量、数据量、计算量等因素变化的规律,是判断一个方法能不能"做大"的重要依据。
在预训练和长文本训练两个阶段,CoWA的困惑度曲线都紧紧贴着FullAttn。
困惑度(perplexity)*:衡量语言模型预测下一个词准确程度的指标,数值越低说明模型预测得越准。
到14B参数规模时,CoWA在4K长度预训练阶段的困惑度和FullAttn几乎完全一致,在32K长文本训练阶段两者的差距小于0.01,但与此同时,CoWA在预训练阶段省下了3.1%的总训练算力,在长文本训练阶段省下了28.5%。
28.5%这个数字值得多说一句。长文本训练之所以省得更多,是因为FullAttn的计算量会随着序列长度的增长而平方级膨胀,序列越长,FullAttn付出的代价越夸张,而CoWA每个头只需要处理固定预算内的token,这个预算不会随着序列长度失控增长。序列越长,CoWA相对FullAttn省下来的比例就越可观。
最后一步,研究者把14B规模训练出来的四种注意力方案模型,以及另外通过继续训练得到的32B模型,都拿到知识问答、推理能力、长文本检索这几类标准测试集上跑了一遍。14B模型上,CoWA在知识类任务上拿到72.70分,FullAttn是72.32分;推理类任务CoWA是64.87分,FullAttn是64.46分;在32K原生长度的检索测试RULER上,CoWA是89.13分,FullAttn是89.42分;在通过YaRN技术外推到128K长度后,CoWA反而拿到了66.60分,比FullAttn的65.84分还高一点。32B模型上的结果也呈现出类似的势均力敌。这说明CoWA这套稀疏设计不是那种"小打小闹能用,一放大就露馅"的技巧,而是真的能撑住规模化训练的考验。
写在后面
读完这篇论文,有一个细节让我反复想了很久:论文标题就直接点破了核心观点,全因果覆盖是一种"集体属性",而不是"每个头的属性"。这句话表面上听起来像句废话,但仔细想想,它其实在挑战一个我们习以为常的默认假设:既然模型需要访问全部历史信息,那是不是就应该让每个计算单元都拥有这个能力?
这篇论文用实验数据告诉我们,答案是否定的。能力的完整性和执行者的完整性从来不是一回事,一个团队能做到样样精通,不需要团队里每个人都样样精通。这个道理放在组织管理、分布式系统设计里都不新鲜,但把它系统性地应用到Transformer的注意力头设计上,而且用严谨的对照实验把这个原理的贡献单独剥离出来验证,这一点做得很扎实。
论文里那个"窗口宽度完全一致,只改变分工方式"的对照实验,是我觉得整篇论文最见功力的地方。很多论文提出一个新架构后,效果好了就直接归功于整个架构,但很少有人愿意花力气去拆解,效果的提升到底是来自哪个具体的设计选择。CoWA的作者们显然是认真地在回答"为什么这样做有效",而不只是"这样做有效"。
论文里也提到一句话,暗示了当前方案还没走到终点:训练阶段用的是"等宽度"分配远程窗口,而"等面积"分配方案被明确列为未来工作。这意味着现在的窗口划分可能还不是最优的,只是一个已经验证有效但仍有改进空间的起点。另外,解码阶段如何进一步压缩显存里驻留的KV缓存,论文里也承认还没做,留了个明确的口子。
一个开放的问题是:如果注意力头的分工不是按照固定的位置规则划分,而是让模型自己去学习动态调整这个分工方式,会不会进一步提升效果?还是说,正是因为放弃了这种"聪明"的动态性,换来了执行效率上的巨大红利?这大概是这类位置规则型稀疏注意力和内容自适应型稀疏注意力之间,一个还没有定论的取舍。
Q&A
Q1:CoWindow Attention是什么?
A:CoWindow Attention(CoWA)是一种新的注意力机制设计,让每个KV头只负责历史文本的一部分区域,但所有头合起来能覆盖完整的历史信息,从而在不牺牲模型效果的前提下大幅降低长文本处理的计算和内存开销。
Q2:CoWA和FullAttn相比效果差多少?
A:在多项测试中差距很小。8K长度关联召回测试里CoWA是89.73%,FullAttn是89.97%;14B模型知识类任务CoWA还略高于FullAttn,说明CoWA基本能达到和全注意力相当的效果。
Q3:CoWA能带来多大的效率提升?
A:在128K序列长度、8张GPU的测试中,CoWA把训练前向延迟降到FullAttn的1/7.4,反向延迟降到1/8.6,解码延迟降到1/3.0,解码阶段内存占用只有FullAttn的1/7.6。
好文章,需要你的鼓励
GeoVerse融合几何潜在扩散模型与Wan2.2视频生成先验,并引入全局空间记忆机制,实现了跨视角一致、生成质量高且推理速度快的新视角合成,在多个数据集上大幅超越现有方法GLD。
本文解读TT-VidT,一种通过分离外观锚点与动作令牌、结合Diff Compression训练目标的视频自监督方法。在24种架构目标组合的受控实验中,仅TT3D与Diff Compression搭配才能进入强动作敏感区,且计算量比同类方法降低近一半。
Q&A Q1:Imprint Reader是什么? A:Imprint Reader是上海人工智能实验室团队开发的模型,通过SMaRT训练方法学会读取另一个模型冻结的权重更新,把参数层面的变化翻译成自然语言描述,判断这次更新到底注入了什么知识或诱导了什么行为习惯。 Q2:MetaEdit能做什么? A:MetaEdit能把Reader对目标行为的梯度信号迁移到原始模型上,不需要任何目标任务的训练数据,只用一句行为描述就能定向调整模型参数,已经在安全性拒绝率调节、数学推理风格转变和工具调用能力提升上验证过效果。 Q3:Imprint Reader读权重更新的准确率高吗? A:目前还不高,知识类目标的Pass@100只有2%,行为类目标是16%,说明自由生成描述的可靠性有限,但相比原始模型完全答不出来的0%基线,已经证明了权重更新的自然语言读取是可行的。
统一多模态模型能看也能画,理论上该会自我修正。本论文提出UMM-Reflection,用整条轨迹的强化学习联合训练诊断和修图能力,GenEval提升12分,且迁移到多个未训练基准,揭示RL提升的是选择可靠性而非新能力。