阿里让AI大脑"术业专攻":给不同神经元分配不同任务
(来源:科技行者)
这项由阿里巴巴集团研究团队完成的研究成果以预印本形式发布于2026年6月18日,论文编号为arXiv:2606.20097v1,有兴趣深入了解的读者可以通过该编号查询完整论文。
当你把一本厚达几百页的侦探小说递给AI,要它找出第237页提到的那把钥匙藏在哪里时,它很可能两眼一抹黑。这不是AI不够聪明,而是它的"大脑工作方式"遇到了一道物理上的天花板。现有的AI语言模型处理文本时,需要让每一个词都和其他每一个词"打招呼"——这种全员互动的机制叫做全注意力(Full Attention,简称FA)。当文本只有几百个词时,这还算轻松;但当文本拉长到几十万个词时,计算量会像雪球一样以平方级别暴增,最终把计算机的资源压垮。
研究团队从另一个角度切入了这个难题。他们没有简单地换掉全注意力,而是先问了一个更有趣的问题:AI大脑里每一个"注意力头"(可以理解为大脑里专门负责不同工作的小组)真的都需要做同样高精度的"全员打招呼"吗?答案让人惊讶——绝大多数小组根本不需要。
就这样,一个名叫HydraHead的新架构诞生了。它的核心思路是:把AI大脑里真正负责精准检索的少数关键小组保留"全员打招呼"的权力,其余的大多数小组则改用一种更高效的"流水线记忆"方式工作。这样一来,整个系统既保住了精准检索的能力,又大幅降低了处理长文本时的计算压力。在实验中,仅用150亿个词的训练数据,HydraHead在处理512K长度文本(相当于一部厚厚的百科全书)的检索任务上,比原始模型提升了超过69%,接近了专门为长文本设计的旗舰模型Qwen3.5的水准。
一、AI大脑的"分工之谜":为什么不同小组需要不同待遇
要理解HydraHead为什么有效,先得弄清楚AI大脑是怎么工作的。现代语言模型的核心机制叫做"多头注意力",可以把它想象成一个大型图书馆的管理团队。这个团队有很多小组,每个小组负责在书库里找不同类型的信息——有的小组专门找人名,有的小组负责理清时间顺序,有的小组追踪逻辑关系。每一层楼(对应模型的每一层)都有这样一套小组配置。
过去,主流的"混合架构"方案是按楼层来分配任务:某几层楼的全体小组用高精度的全注意力,其他楼层的全体小组改用高效的线性注意力(Linear Attention,简称LA)。线性注意力就像一个不断更新的备忘录,只记住当前认为最重要的信息,而不是全部保存,所以它处理很长的文本时计算量是线性增长的,比全注意力省力得多。
然而,这种按楼层划分的方式有一个根本问题:同一层楼里不同小组的"功能需求"差异极大。阿里巴巴的研究团队通过一系列精密的因果实验(后面会详细解释)发现,在同一层楼里,往往只有极少数几个小组对精准检索至关重要,绝大多数小组其实可以换用省力的工作方式而不影响大局。换句话说,按楼层一刀切的分法,要么把宝贵的高精度资源浪费给了不需要的小组,要么在转换整层楼时把少数真正关键的小组也给换掉了,导致检索能力下跌。
研究团队用两张图直观地呈现了这一发现。第一张图显示了在一个长文本检索任务中,模型每一层每一个小组对最终答案的贡献——同一层楼里,不同小组的贡献差异悬殊,有的小组贡献巨大,旁边的小组却几乎为零。第二张图则展示了相邻楼层之间的输出相似度,结果是平滑渐变的,没有明显的"断层",这意味着按楼层划分的边界本身就缺乏可靠的功能依据。
这两个发现合在一起,指向了一个清晰的结论:**注意力头**(每个楼层里的每个小组),而不是整层楼,才是分配不同工作机制的最合理粒度。
二、如何找出真正的"关键小组":一套精准的因果诊断工具
发现了头部层面的功能异质性之后,下一个问题随之而来:怎么知道哪个小组才是真正关键的?光看它"看起来很忙"是不够的——一个小组可能表面上在认真工作,但它的输出结果其实被下游的其他小组纠正了,真正关键的信号其实来自别处。
研究团队借用了认知神经科学里的一套经典方法论,叫做"因果干预",通俗地说就是"做实验,不猜测"。具体操作分三步走。
第一步叫做"激活替换"。研究人员构造了一对近乎相同的题目:一道是原题(正确答案是数字"4321"),另一道把答案偷偷换成了另一个数字"8765",其余内容保持一模一样。然后,他们把模型处理原题时某个特定小组的输出,替换成处理那道"偷换答案"的题目时的输出,同时让所有其他小组保持处理原题时的状态。如果这一个小组被替换后,模型原本应该输出"4321"的答案变差了,那这个小组就是关键的——它的输出对最终结果有直接的因果影响。重要性分数用一个归一化的公式来计算:干预前后准确度的下降幅度除以替换整个模型时的最大下降幅度,结果落在0到1之间,越接近1代表越关键。
第二步叫做"路径追踪"。有些小组的影响是间接的——它不直接写结果,而是把信号传给另一个关键小组,再由那个小组传递给最终输出。为了捕捉这类"幕后推手",研究团队使用了更精细的路径追踪技术:记录某个小组在处理偷换答案的题目时,传给下游关键小组的那部分信号,然后在处理原题时把这段特定信号也替换掉,看结果怎么变。这样就能追踪一个小组通过特定路径产生的间接影响。对于长文本检索任务,这个追踪过程只需要大约两轮就能收敛,说明这是一个"浅层电路"——关键信号传递的链条并不长。
第三步是"多任务融合"。研究团队考虑了两类目标能力:长文本检索能力和通用推理能力。对每个小组,先分别计算它对这两类能力的重要性分数,再按任务稳定性加权——一个小组如果只在个别测试中表现重要,而在大多数测试中无足轻重,就会被降权。最后,两类任务的分数按等权重合并成一个总排名,按照总分从高到低筛选出需要保留全注意力的小组。
整个筛选过程非常轻量:只需要对几十个样本做前向推理,不需要任何反向传播或梯度计算,而且只需要大约6个样本就能让排名趋于稳定。在Qwen3-1.7B这个模型上(共28层,每层16个小组,合计448个小组),最终只有约29个小组(大约6.5%)被认定为长文本检索的关键小组,其余90%以上的小组都可以安全地换用线性注意力。更重要的是,这29个关键小组分散在各个楼层,而不是集中在几个特定楼层——这正是按楼层一刀切的方案无法避免误伤的根本原因。
三、混搭方案的工程细节:如何让两种截然不同的工作方式和平共处
找到了关键小组之后,接下来的挑战是如何让全注意力和线性注意力在同一层楼里"和平共处"。这两种机制的工作方式有根本性的差异,就像让一个习惯精读每一行字的人和一个习惯速读只记要点的人坐在同一张桌子上合作——他们产出的"笔记风格"会很不一样,直接把两份笔记混在一起可能反而让人更糊涂。
具体来说,全注意力的输出特征和线性注意力的输出特征在"音量大小"上差异显著。全注意力通过softmax函数处理,会让输出集中在少数几个高权重词上,整体特征的"音量"受到自然约束;线性注意力则没有这种约束,在深层网络里,其输出的均方根值(可以理解为"音量大小")可以比全注意力高出6.2倍。如果直接把这两种"音量"差异巨大的输出拼接在一起送给后续处理层,会导致训练不稳定,模型会一头雾水。
研究团队的解决方案分两步。第一步是独立归一化:对每个小组的输出,不管它是用全注意力还是线性注意力算的,都先各自做一次RMSNorm(一种标准化处理,效果类似于把不同乐器的音量都调到同一个基准)。第二步是可学习的头部缩放:归一化之后,给每个小组分配一个独立的可训练缩放系数,让模型在训练过程中自动学习该给每个小组的输出赋予多大的权重。这样一来,既消除了两种机制之间的"音量差异",又保留了每个小组独特的功能贡献。
在两种机制的内部设计上,研究团队也做了针对性的调整。对于全注意力小组,他们去掉了旋转位置编码(RoPE,一种给词语标注位置信息的机制),改用一个随序列长度对数缩放的系数来处理超长文本时的数值稳定性问题,同时加入了一个额外的"门控"机制来提升表达能力并缓解"注意力汇聚"现象(即模型过于关注某几个固定位置的词)。对于线性注意力小组,研究团队反而给它加上了RoPE,以弥补线性注意力在位置感知上的先天不足,并把键值头的数量从GQA模式(多个查询头共享少数几个键值头)扩展到了MHA模式(每个查询头都有自己的键值头),以提升表达能力。
四、从旧模型到新架构:三阶段接力训练流水线
有了架构设计,还需要一套高效的训练策略把这个混搭架构"训练好"。从零开始训练代价极高,研究团队选择了从已有的预训练模型(Qwen3-1.7B)出发,通过三个阶段的迁移训练完成转换。
第一阶段的任务是"移植和对齐"。研究团队先按照因果筛选结果,把每一层里非关键小组的全注意力替换成线性注意力(具体使用的是Gated DeltaNet,一种带有遗忘门机制的改进版线性注意力)。替换时,新加入的线性注意力小组直接继承原来全注意力小组的查询、键、值投影矩阵权重,而不是随机初始化,这让新小组一开始就能模仿原小组的基本行为。保留下来的全注意力小组则加入一个新的门控分支,这个门控分支的初始权重被设置为接近零,偏置被设置为接近1,效果是让门控一开始几乎是透明的,不改变全注意力小组的原有行为。完成替换后,冻结模型的其他参数,只训练新替换的部分,目标是让每一层的混合注意力输出尽量接近原始全注意力层的输出——用数学语言说就是最小化两者之间的均方误差。
第二阶段的任务是"全局对齐"。第一阶段只保证了每一层的局部行为接近,但层层累积的误差可能让最终输出的概率分布偏移。第二阶段解冻整个模型,用原始Qwen3-1.7B作为"老师",让混合模型学习模仿老师在每个位置的词汇概率分布,同时也用真实的下一个词预测损失来训练。这一阶段的目标是把全局的语义连贯性和知识保留好。
第三阶段的任务是"长文本适应"。前两阶段使用的序列长度都比较短(最多2048个词),这个阶段把序列长度拉长到16384个词,用标准的下一个词预测损失继续训练,帮助模型真正学会处理长上下文的规律。优化后的训练配置把第一阶段的数据量扩大到约8亿词,第二阶段扩大到40亿词,这比最初的配置提升了数倍,实验证明这对最终性能有显著帮助。
五、实验结果:数字背后的真实含义
研究团队在统一的训练配置下,把HydraHead和多种竞争方案做了系统对比,所有方案都从同一个Qwen3-1.7B基础模型出发,使用相同的训练数据和步数。参与对比的方案覆盖了三大类混合架构:按楼层划分的、按词语位置划分的、以及按注意力头划分的。
按楼层划分的方案中,表现最好的是用闪电注意力(一种高效的近似注意力)替换非关键楼层的版本,在长文本扩展性能上平均达到约85%,但通用推理能力较弱。使用Gated DeltaNet的楼层混合方案则表现更差,扩展后长文本性能急剧下降,接近零分。
按词语位置划分的方案(即滑动窗口注意力+线性注意力的组合)在通用推理上表现亮眼,大幅超过按楼层划分的方案,但长文本扩展能力同样非常有限。
HydraHead在这次对比中实现了双赢:长文本扩展性能平均达到约87%(远超所有其他方案),同时通用推理能力也比按楼层划分的最优方案高出超过10个百分点(在困难推理任务上提升了约11%)。
在更激进的压缩比例测试中,HydraHead也展现出了明显优势。当线性注意力与全注意力的比例从3:1提升到7:1时,使用因果筛选+至少保留每层一个全注意力头的约束策略,长文本性能(在16K到256K范围内的平均分)仅下降约1.3个百分点,达到了约54%,而通用推理能力的保留程度也远好于同等压缩比例下的按楼层方案——事实上,7:1压缩比的HydraHead,其综合表现已经和3:1压缩比的最佳楼层混合方案基本持平,但通用推理能力高出约9.7%(困难任务)和约3%(简单任务)。这意味着HydraHead用更少的全注意力头,做到了更好的事情。
当压缩比进一步推到9:1时,简单长文本任务的性能还能保留80%以上,但多关键词检索和复杂推理能力出现了更明显的下滑,说明在极端稀疏配置下仍有优化空间。
在开源模型大比武中,扩大训练数据到150亿词后的HydraHead更是表现抢眼。在处理256K长度文本的针查任务(在海量文本中找到特定信息,如大海捞针)上,单关键词任务得分94.53%,多关键词任务得分52.70%,而绝大多数竞争模型在这个长度上已经趋近于零。与此同时,在通用推理基准测试中,HydraHead平均得分约50.6%,比大多数混合架构模型高出10个百分点以上,虽然略低于专门为推理强化的全注意力模型Qwen3-1.7B(约54%),但差距仅有约3.4个百分点——而这个差距对应的代价是仅150亿词的迁移训练,而非完整的预训练。
六、为什么因果筛选比随机分配强这么多:一次"拆桥实验"
研究团队专门设计了一个实验来验证因果筛选的价值。他们对比了五种不同的头部分配策略:固定比例(每层都保留25%的全注意力头)、层内随机(每层随机选25%的头)、全局随机(从全局随机选总数25%的头)、层内因果筛选(在每层内按因果重要性排名选)、以及全局因果筛选(跨层按因果重要性排名选)。
固定比例和层内随机的表现非常接近,说明在层内打乱顺序对模型影响不大,关键在于每层是否都有全注意力头覆盖。全局随机的表现则大幅下滑,某些层可能完全没有全注意力头,导致关键信号传递链断裂,长文本性能接近崩溃——这恰恰验证了研究团队的核心判断:关键头部是分散的,不是集中的,随机的全局分配很容易把关键层"漏掉"。
层内因果筛选比层内随机有明显提升,说明在每层内优先保留最重要的头确实有价值。全局因果筛选则取得了最佳成绩:通过跨层的全局优先级排名,把宝贵的全注意力配额精确分配给最需要它的那些头,既不浪费,也不遗漏,综合性能全面领先。
研究团队还对因果筛选得到的重要性分数做了"拆桥实验"验证:按重要性从高到低逐步把头部的输出替换成错误答案时对应的激活值(相当于逐步切断关键桥梁),模型的检索准确率从接近100%迅速崩溃,只移除前1%的关键头就能让准确率大幅下滑;而以同等数量随机移除头部时,准确率几乎没有变化。这个实验有力地证明了:因果筛选确实找到了真正负责检索的那些关键小组,而不只是相关性较高的冗余小组。
说到底,HydraHead的核心贡献是把一个直觉上显而易见、但工程上难以实现的想法变成了现实:AI大脑里的不同小组确实各司其职,我们完全可以根据每个小组的真实功能来分配工作方式,而不是用一把尺子量所有人。这个思路之所以奏效,是因为研究团队没有停留在观察层面,而是用严格的因果实验证明了功能分工的存在,再用精心设计的归一化融合机制解决了混搭带来的工程问题,最后用三阶段迁移训练把整套方案以极低的成本落地。
对普通用户而言,这项研究意味着未来的AI助手在处理超长文档时——无论是几十万字的法律合同、完整的学术文献库还是冗长的会议记录——将不再那么容易"迷路",而且这一能力的获得不需要从零开始训练一个庞大的全新模型,只需要在已有模型上进行相对轻量的改造。
一个值得继续思考的问题是:这套因果诊断工具能否进一步自动化,甚至扩展到更大规模的模型上?当模型参数达到千亿级别时,逐一对所有头部做因果实验的成本会急剧上升,如何在更大的"图书馆"里快速找到那几个真正的关键小组,或许是这条研究路线下一步最值得探索的方向。有兴趣深入了解技术细节的读者,可以通过论文编号arXiv:2606.20097v1查阅完整论文。
Q&A
Q1:HydraHead架构是如何决定哪些注意力头保留全注意力、哪些改用线性注意力的?
A:HydraHead通过因果干预实验来判断每个注意力头的重要性。具体方法是构造一对内容相同但答案不同的题目,把某个头的输出替换成错误答案时的激活值,观察模型准确率的下降幅度。下降越大,说明这个头越关键,就保留全注意力;下降不明显的头则替换为线性注意力。整个筛选过程不需要训练,只需要几十个样本的前向推理就能完成。
Q2:线性注意力和全注意力混合使用时为什么需要专门的归一化处理?
A:两种注意力机制输出的特征在"数值大小"上差异很大。全注意力通过softmax函数约束了输出幅度,而线性注意力没有这种约束,在深层网络里其输出的均方根值可以比全注意力高出6倍多。如果直接把两者的输出拼接在一起,数值大的那类会主导后续处理,导致训练不稳定。独立归一化把两类输出都先调整到统一的尺度,再通过可学习的缩放系数让模型自己决定各自的权重,从而消除这种干扰。
Q3:HydraHead和现有的按层混合架构(比如HypeNet)相比,主要优势体现在哪里?
A:HydraHead最大的优势是在长文本能力和通用推理能力上同时优于按层混合方案。按层混合方案要么整层用全注意力(浪费计算),要么整层换成线性注意力(可能切断关键信号链)。HydraHead在头部粒度上精准分配,避免了两种浪费。实验数据显示,HydraHead在困难推理任务上比最优按层混合方案高出超过10%,同时长文本扩展性能也更强。即使把全注意力比例压缩到7:1,其综合表现仍可媲美3:1比例的按层混合方案。