威斯康星大学麦迪逊分校与独立研究者破解AI大脑"记忆读取"难题

这项由威斯康星大学麦迪逊分校与独立研究者共同完成的研究，以预印本形式发布于2026年7月22日，论文编号为arXiv:2607.27230，感兴趣的读者可通过该编号检索完整原文。

神经网络，尤其是近年来风头无两的Transformer模型（也就是驱动ChatGPT、Gemini等大语言模型的核心架构），正在变得越来越庞大、越来越复杂。然而，在这种越做越大的趋势背后，有一个被长期忽视的设计缺陷正在悄悄拖累模型的表现。研究团队在深入审视这个问题后，提出了一个优雅而高效的解决方案——多头注意力残差（Multi-Head Attention Residuals，简称MHAR）。这个方案的核心改动只有寥寥几行代码，却能让模型在各种规模下都表现得更好，而且几乎不增加任何额外的计算成本。

一、神经网络的"集体记忆"：为什么模型需要回头看？

要理解这个研究解决了什么问题，先得弄清楚神经网络是怎么工作的。可以把一个深度神经网络比作一条精密的流水线：原材料（也就是输入的文字或数据）从第一道工序进入，经过几十道甚至上百道加工工序，最终输出成品。每道工序的工人只会接过上一道工序传来的零件，加工后传给下一道，不会去翻找仓库里更早期的半成品。

这种"只看前一步"的机制在工程上叫做残差连接（residual connection），是2016年由何恺明等人提出的经典设计。它的好处是简单稳定，缺点是如果某个早期工序产出了一个非常有价值的中间产品，到了后期工序时，那个中间产品已经被后续无数次加工覆盖在一堆累加的结果里，想单独取出来用几乎是不可能的。

2025年，Kimi团队提出了一个改进方案，叫做"注意力残差"（Attention Residuals）。他们的想法是：与其让每道工序只看上一步的结果，不如给每道工序一个"记忆索引"，让它能够在所有历史记录（包括原始输入、每一道工序的输出）中自由检索，找到自己最需要的那一份，然后把那份内容取来作为当前加工的起点。这种机制用的是一种叫做"softmax加权"的数学操作，本质上就是给历史上每个工序的输出打一个重要性分数，然后按分数加权平均，得到一个融合了历史精华的输入。

这个想法听起来非常聪明，但研究团队发现，Kimi的方案里藏着一个微妙但严重的缺陷——随着模型变大，这个缺陷会从可以忽略的小问题，变成拖累整体性能的大障碍。

二、"被迫妥协"：单一路由查询的致命短板

Kimi的注意力残差方案中，每道工序用一个"查询向量"（query vector）来决定从历史记录里读取什么。这个查询向量的维度和整个网络的宽度（即特征维度d）相同，它会对所有历史输出打分，算出一个统一的权重分布，然后所有特征维度都按照这同一个分布去读取历史信息。

用流水线的比喻来说：每道工序里有d个不同的"工人"，他们负责加工零件的不同部分（不同的特征子空间）。Kimi的方案相当于让这d个工人共用一份"历史档案调阅单"——无论是负责加工零件外形的工人，还是负责打磨内部结构的工人，他们调阅的都是同一套历史档案，按同一个比例混合。

问题在于，负责外形的工人可能最需要第3道工序的成果，负责内部结构的工人可能最需要第17道工序的成果，而这两类工人被迫共用一张调阅单，这张单子只能是某种折中——每类工人都没有得到自己最想要的。研究团队把这种现象称为"被迫妥协"（forced compromise）。

更关键的是，这种妥协的代价会随着模型变宽而急剧增加。模型越宽，意味着有越多种类的"工人"，他们对历史档案的需求差异就越大，单一调阅单带来的妥协成本就越高。在100M参数的小模型上，这个问题还不明显；但到了1B参数的大模型上，这种强制妥协已经严重到让整个注意力残差机制的表现不如不用——比普通的残差连接还差0.105个验证损失单位。

三、一个零成本的解决方案：把调阅单一分为H份

研究团队的解决思路非常简洁：既然自注意力机制（self-attention）早就因为"单个查询无法服务所有子空间"而改用了多头设计，那么用于历史深度检索的查询向量，凭什么还要坚持用单头？

MHAR的核心改动就是：把那个形状为(d,)的单一查询向量，重新排列成(H, d/H)的形状——也就是H个更小的查询向量，每个负责d/H个特征维度。每个小查询向量各自独立地对历史输出打分，算出自己独立的权重分布，然后各自读取自己负责的那一部分特征维度的历史信息。最后，H个子结果拼接在一起，还原为完整的d维特征。

这个改动的妙处在于：参数量完全没有增加。原来是一个d维的向量，现在是H个d/H维的向量，总参数量还是d个数字，只是排列方式变了。计算量也几乎没有变化，因为每个特征元素还是只被处理一次，只是额外多算了H个小型的softmax（对历史长度轴做归一化），而这个额外计算相对于整个注意力和MLP模块来说可以忽略不计。将H设为1时，MHAR完全退化回Kimi的原始单头版本，因此MHAR是对注意力残差的严格推广。

用流水线的比喻：现在每道工序里，d个工人被分成了H个小组，每个小组有d/H个工人，每个小组有自己专属的调阅单，可以独立决定从哪些历史档案中获取灵感，不再被迫共享。负责外形的小组可以重点参考第3道工序，负责内部结构的小组可以重点参考第17道工序，两组互不干扰。

四、实验数据：小改动，大收益

研究团队在FineWeb-Edu数据集上，从零开始训练了100M、350M和1B三个规模的语言模型，每个规模训练2万步，并与四种方法做了严格对比：标准Transformer（基线）、超连接（hyper-connections，另一种改进残差结构的方案）、单头注意力残差（即Kimi原版，H=1）以及MHAR（H等于KV头数）。

实验结果呈现出一个非常清晰的规律。在100M规模上，单头注意力残差比基线好0.039个验证损失，MHAR比基线好0.049，MHAR比单头好0.010。看起来差距不大，两者都有帮助。

到了350M规模，局势开始分化：单头注意力残差比基线差了0.055（即变差了），而MHAR比基线好了0.080。这意味着MHAR比单头版本好了整整0.135。

到了1B规模，分化更加剧烈：单头注意力残差在其自身最优学习率下仍比基线差0.105，而MHAR比基线好0.063。MHAR对单头的优势扩大到了0.168。

换句话说，随着模型规模从100M增长到1B，单头路由从"有一点帮助"变成了"明显有害"，而MHAR始终保持改进，且优势持续扩大。这正好印证了"被迫妥协"代价随模型宽度增加而增长的理论预测。

为了把损失改善转换成更直觉的计算量等价值，研究团队还计算了"计算等价增益"（CEGFLOPs）：要让基线模型通过单纯增加训练计算量来达到MHAR的性能，需要多花1.27倍（100M）、1.49倍（350M）和1.38倍（1B）的计算。而MHAR本身只增加了0.5%到1.2%的计算量，代价和收益的比例相当悬殊。

与超连接方案相比，MHAR在三个规模上的表现都更好（分别好0.035、0.050、0.061），而且超连接的收益在1B规模几乎消失（仅比基线好0.002，几乎在误差范围内），MHAR的优势则随规模持续扩大。

五、H到底设为多少最合适？

既然H越大越好，是不是应该把H设为尽可能大的值？答案并非如此简单。研究团队做了一个4×4的网格搜索实验，同时变化H（路由头数）和KV（键值头数），测量不同组合下的验证损失。

在模型训练不充分的早期阶段（100M模型只训练5000步），有一个清晰的规律：当KV头数固定为1时，H越大越好，从H=1到H=8呈单调下降，改善幅度0.028，远超单次随机种子的噪声水平。这说明在训练早期，更多的路由头总是有帮助的，相当于给模型更多的自由度去探索不同的读取方式。

然而，当模型训练到充分收敛时，最优的H并不在最大值处，而是"饱和"在一个特定位置——恰好等于KV头数。在100M的完整训练实验中，KV=4和KV=8的最优H恰好落在H=KV这条对角线上。在350M的实验中，KV≥4的最优点同样在H=KV对角线上。

研究团队对此的解释是：KV头数代表了模型在特征空间上的"消费粒度"，也就是说模型在进行注意力计算时，本来就是以KV头为单位来处理不同的特征子空间的。把路由头数设为和KV头数相同，正好对齐了这个粒度，既充分释放了不同子空间的读取自由度，又不会因为过度细分而引入不必要的随机性。

这个发现给了实践者一个非常方便的默认规则：不用为H调参，直接把H设为模型的KV头数就是接近最优的选择，而且这个规则是"宁可多不可少"的安全方向——如果模型训练不够充分，更多的H只会有帮助，不会有害。

六、不仅仅是训练指标：在真实任务上的表现

验证损失的改善是否真的意味着模型在实际任务中更有用？研究团队对100M、350M、1B三个规模的基线和MHAR模型做了零样本评测，测试了WikiText-2困惑度（一种衡量语言模型对文本理解能力的指标）、LAMBADA（需要理解长文上下文才能预测最后一个词的任务）和HellaSwag（常识推理补全任务）。

MHAR在每个规模上都改善了WikiText-2困惑度和LAMBADA准确率，而且随着规模增大，困惑度的相对改善幅度也在扩大：100M改善了约10%，350M改善了约15%，1B改善了约19%。HellaSwag方面，350M有提升，1B持平，100M略低于基线但在误差范围内。

这说明验证损失的改善不是针对训练数据的过拟合，而是真实泛化能力的提升，能够迁移到完全不同的评测数据集上。

七、8B大模型的中途改造：不重新训练也能受益

从零训练小模型固然可以证明方法的有效性，但现实中，大多数高性能模型都是花费了海量资源预训练出来的，不可能轻易抛弃重来。研究团队因此设计了一套"身份保持转换"方案，让已有的预训练模型可以在继续训练（mid-training）过程中逐渐接入MHAR机制，而不会产生任何突然的性能跳变。

他们采用的是"delta注意力残差"（delta attention residuals）形式：不是用路由混合结果完全替换残差流，而是把路由结果作为一个附加项叠加到原有的残差流上（h = 原有残差流输出 + α × 路由混合结果）。关键在于，这个附加项的输出门控α被初始化为零，这意味着在转换的第一步，模型的计算结果和原始模型完全相同，不会有任何数值上的突变。随着继续训练的推进，α逐渐学习到非零值，MHAR机制才慢慢发挥作用。

研究团队在Marin-8B这个开源的8B参数模型上验证了这套方案。他们用一个约1.9万亿token的高质量混合语料库（anneal_pt_v3）进行约100亿token的继续训练，同时设置了一个完全相同训练配置（相同学习率、相同数据顺序、相同训练步数）的对照组，唯一的区别是是否加入MHAR。

从训练损失曲线来看，两条曲线在最初几乎完全重合——第一步的训练损失差值只有约10??，远小于单批次间的随机波动（约0.028）。这正是零初始化门控的效果：转换后的模型和原模型在数值上完全等价，不会引起任何优化上的震荡。

在下游任务评估上，MHAR相对于对照组带来了统计显著的改善。GSM8K数学推理任务提升了3.2个百分点（从47.0%到50.2%，配对McNemar检验p=0.004），GPQA研究生级别科学问题任务提升了3.1个百分点（从31.5%到34.6%，p=0.038）。MMLU多任务语言理解、MATH数学竞赛题、HumanEval和MBPP代码生成任务的变化在统计上不显著。这表明MHAR带来的提升主要集中在推理能力方面，代码和一般知识任务的收益则不明显——至少在这个规模和训练预算下如此。

八、为什么会有效：直接探测训练后的路由查询

研究团队不满足于只看实验结果，还进一步通过直接分析训练好的模型参数来验证"被迫妥协"机制。他们在训练好的单头（H=1）模型上做了一个探测实验：把那个训练好的单一路由查询向量切成若干片段，计算每个片段"如果单独决策的话会想要哪种历史读取分布"，然后与实际使用的共享分布做KL散度比较。这个KL散度就是"子空间不一致度"的量化指标。

结果非常清晰：训练好的查询向量中，各片段之间的不一致度（KL散度）从100M到1B模型单调增长了2.6倍（0.235→0.606），而用一个随机权重相同的随机向量做同样的测试，得到的不一致度只有训练版本的约5%到15%。这说明这种不一致是模型主动学习出来的，不是特征空间几何结构的固有属性，而是因为不同子空间确实学到了对历史深度有不同偏好的表示。

更重要的是，研究团队还做了一个控制实验：在保持层数（L=12）和KV头数（4）不变的情况下，只把模型宽度从d=512增加到d=768，观察不一致度的变化。结果是训练查询的不一致度上升了14%（从0.273到0.311），而随机查询的基线和信息源相似度几乎没有变化。这就把"宽度才是驱动因素"这一论断从相关性提升到了接近因果的程度。

与此同时，在训练好的MHAR模型上，可以看到8个路由头确实学到了不同的历史读取模式：每个头相对于所有头平均分布的偏差最大达到±0.28，而随机权重的对照只有±0.067，而且这些偏差在不同评估文本上高度一致（相关系数r=0.77），且不同头之间几乎不相关，证明各头确实分工明确，而非学到了冗余的相同模式。

九、工程挑战与高效内核：让理论落地

MHAR在数学上几乎不增加计算量，但在实际训练中，深度路由机制有一个严重的工程瓶颈：内存带宽。每个子层在计算路由混合时，需要读取所有历史输出（最多2L+1个），而这个操作几乎没有数据复用，完全受内存带宽限制。一个朴素的参考实现会反复复制和存储大量中间张量，内存占用和速度都非常糟糕——在1B模型上，朴素实现甚至会超出80GB的GPU显存限制。

研究团队为此专门开发了融合的Triton自定义GPU内核（fused Triton routing kernels）。这套内核的核心思想类似于FlashAttention的在线softmax技巧：在正向传播中，每个token位置只需一遍扫描所有历史源，用寄存器里的累加器保持运行中的softmax统计量，完全不需要把归一化后的键张量材料化（保存到显存中）；反向传播则通过两遍扫描完成所有梯度计算，并把所有路由调用对源张量的梯度贡献直接累加到一个共享的fp32缓冲区中，避免了大量小张量加法。唯一需要为反向传播保存的是路由权重张量（形状为N×B×T×H），比朴素实现保存的堆叠源张量小d/H倍。

实际测试结果显示，融合内核把路由操作本身的速度提升到了torch.compile版本的2到5.3倍，端到端训练吞吐量从朴素实现的20%到50%基线速度，提升到55%到88%基线速度，内存占用也回落到接近基线水平（100M: 42.0 GB vs 基线41.5 GB，350M: 20.0 GB vs 19.4 GB，1B: 20.1 GB vs 19.0 GB）。这使得在真实规模上使用深度路由机制从理论上可行变成了工程上实用。

十、结论：多头设计让深度路由真正起作用

说到底，这项研究的核心发现可以用一句话概括：如果注意力机制需要多头设计，那么对历史深度做注意力的路由机制同样需要多头设计，而且随着模型变大，这个需求只会越来越迫切。

研究团队通过从零训练实验、直接参数探测、宽度控制实验、中途训练实验和专用高效内核这五个维度，对这一论断做了非常完整的论证。单头路由在小模型上有帮助，但在大模型上会变成障碍；多头路由通过零参数、零额外计算的简单改动，在各个规模上都能持续改善表现，且优势随规模扩大。

H等于KV头数这个无需调参的默认规则提供了很强的工程可操作性，不需要额外的超参数搜索。身份保持转换方案让已有的大模型也能受益，而融合内核则消除了工程上的主要障碍。

这项研究留给未来的问题也很清晰：为什么最优的H恰好等于KV头数，背后是否有更深的理论解释？各个路由头具体学到了怎样的历史读取模式，是否和注意力头的功能分工有对应关系？更激进的中途训练方案能否进一步放大MHAR的收益？这些都是值得继续探索的方向。有兴趣深入了解技术细节的读者，可以通过arXiv编号2607.27230查阅完整论文。

Q&A

Q1：多头注意力残差（MHAR）和普通Transformer的残差连接有什么区别？

A：普通Transformer的残差连接让每层只能看到上一层的输出。注意力残差（Kimi 2025年提出）允许每层从所有历史层的输出中加权读取，但用一个共享的权重分布服务所有特征维度。MHAR则进一步把这个共享权重分配给H个子空间，每个子空间有自己独立的历史读取方案，解除了所有特征维度"被迫共用一套历史参考"的限制。

Q2：MHAR训练速度会不会比普通Transformer慢很多？

A：使用研究团队开发的融合Triton内核后，100M模型的训练速度是普通Transformer的88%，350M是71%，1B是55%。速度有损失，但相比于朴素实现（只有20%到50%的基线速度），改善非常显著。而且MHAR和单头注意力残差的训练开销完全相同，多头分割本身不带来任何额外成本。

Q3：MHAR中路由头数H应该怎么设置？

A：研究团队的实验表明，把H设为模型的KV头数（键值头数）是接近最优的无需调参默认值。在350M（KV=8）和1B（KV=8）规模上，最优H都在或接近H=KV。对于训练不充分的情况，更大的H只会有帮助不会有害，所以这个规则是安全的。