威斯康星大学麦迪逊分校与独立研究者破解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原版H1以及MHARH等于KV头数。实验结果呈现出一个非常清晰的规律。在100M规模上单头注意力残差比基线好0.039个验证损失MHAR比基线好0.049MHAR比单头好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越大越好从H1到H8呈单调下降改善幅度0.028远超单次随机种子的噪声水平。这说明在训练早期更多的路由头总是有帮助的相当于给模型更多的自由度去探索不同的读取方式。然而当模型训练到充分收敛时最优的H并不在最大值处而是饱和在一个特定位置——恰好等于KV头数。在100M的完整训练实验中KV4和KV8的最优H恰好落在HKV这条对角线上。在350M的实验中KV≥4的最优点同样在HKV对角线上。研究团队对此的解释是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检验p0.004GPQA研究生级别科学问题任务提升了3.1个百分点从31.5%到34.6%p0.038。MMLU多任务语言理解、MATH数学竞赛题、HumanEval和MBPP代码生成任务的变化在统计上不显著。这表明MHAR带来的提升主要集中在推理能力方面代码和一般知识任务的收益则不明显——至少在这个规模和训练预算下如此。八、为什么会有效直接探测训练后的路由查询研究团队不满足于只看实验结果还进一步通过直接分析训练好的模型参数来验证被迫妥协机制。他们在训练好的单头H1模型上做了一个探测实验把那个训练好的单一路由查询向量切成若干片段计算每个片段如果单独决策的话会想要哪种历史读取分布然后与实际使用的共享分布做KL散度比较。这个KL散度就是子空间不一致度的量化指标。结果非常清晰训练好的查询向量中各片段之间的不一致度KL散度从100M到1B模型单调增长了2.6倍0.235→0.606而用一个随机权重相同的随机向量做同样的测试得到的不一致度只有训练版本的约5%到15%。这说明这种不一致是模型主动学习出来的不是特征空间几何结构的固有属性而是因为不同子空间确实学到了对历史深度有不同偏好的表示。更重要的是研究团队还做了一个控制实验在保持层数L12和KV头数4不变的情况下只把模型宽度从d512增加到d768观察不一致度的变化。结果是训练查询的不一致度上升了14%从0.273到0.311而随机查询的基线和信息源相似度几乎没有变化。这就把宽度才是驱动因素这一论断从相关性提升到了接近因果的程度。与此同时在训练好的MHAR模型上可以看到8个路由头确实学到了不同的历史读取模式每个头相对于所有头平均分布的偏差最大达到±0.28而随机权重的对照只有±0.067而且这些偏差在不同评估文本上高度一致相关系数r0.77且不同头之间几乎不相关证明各头确实分工明确而非学到了冗余的相同模式。九、工程挑战与高效内核让理论落地MHAR在数学上几乎不增加计算量但在实际训练中深度路由机制有一个严重的工程瓶颈内存带宽。每个子层在计算路由混合时需要读取所有历史输出最多2L1个而这个操作几乎没有数据复用完全受内存带宽限制。一个朴素的参考实现会反复复制和存储大量中间张量内存占用和速度都非常糟糕——在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 GB350M: 20.0 GB vs 19.4 GB1B: 20.1 GB vs 19.0 GB。这使得在真实规模上使用深度路由机制从理论上可行变成了工程上实用。十、结论多头设计让深度路由真正起作用说到底这项研究的核心发现可以用一句话概括如果注意力机制需要多头设计那么对历史深度做注意力的路由机制同样需要多头设计而且随着模型变大这个需求只会越来越迫切。研究团队通过从零训练实验、直接参数探测、宽度控制实验、中途训练实验和专用高效内核这五个维度对这一论断做了非常完整的论证。单头路由在小模型上有帮助但在大模型上会变成障碍多头路由通过零参数、零额外计算的简单改动在各个规模上都能持续改善表现且优势随规模扩大。H等于KV头数这个无需调参的默认规则提供了很强的工程可操作性不需要额外的超参数搜索。身份保持转换方案让已有的大模型也能受益而融合内核则消除了工程上的主要障碍。这项研究留给未来的问题也很清晰为什么最优的H恰好等于KV头数背后是否有更深的理论解释各个路由头具体学到了怎样的历史读取模式是否和注意力头的功能分工有对应关系更激进的中途训练方案能否进一步放大MHAR的收益这些都是值得继续探索的方向。有兴趣深入了解技术细节的读者可以通过arXiv编号2607.27230查阅完整论文。QAQ1多头注意力残差MHAR和普通Transformer的残差连接有什么区别A普通Transformer的残差连接让每层只能看到上一层的输出。注意力残差Kimi 2025年提出允许每层从所有历史层的输出中加权读取但用一个共享的权重分布服务所有特征维度。MHAR则进一步把这个共享权重分配给H个子空间每个子空间有自己独立的历史读取方案解除了所有特征维度被迫共用一套历史参考的限制。Q2MHAR训练速度会不会比普通Transformer慢很多A使用研究团队开发的融合Triton内核后100M模型的训练速度是普通Transformer的88%350M是71%1B是55%。速度有损失但相比于朴素实现只有20%到50%的基线速度改善非常显著。而且MHAR和单头注意力残差的训练开销完全相同多头分割本身不带来任何额外成本。Q3MHAR中路由头数H应该怎么设置A研究团队的实验表明把H设为模型的KV头数键值头数是接近最优的无需调参默认值。在350MKV8和1BKV8规模上最优H都在或接近HKV。对于训练不充分的情况更大的H只会有帮助不会有害所以这个规则是安全的。

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

相关推荐

3步掌握FanControl风扇控制：Windows散热管理的终极解决方案

YOLOv8 多类别检测涨点实战｜全网复现 10559 张 18 类电力金具多尺度识别、不均衡样本调参、无人机巡检全链路落地

GetQzonehistory：三步实现QQ空间历史说说的永久备份方案

更多资讯

H.265全平台Web播放实战：智能自适应渲染与带宽存储优化

PHP服务端集成活体识别：风控第一道防线的完整落地实践

Oracle-NetSuite与用友ERP会计信息系统差异：从循环设计到准则落地的选型指南

仿站工具小飞兔V19.8.1去更新版使用教程与避坑指南

Gromacs NVT与NPT系综原理及平衡实操指南

计算机网络核心知识梳理：分层模型、IP计算与排障实战

今日更新

PHP htmlentities()函数用法讲解

AI Agent工程实战：从七要素到七个决策点的系统设计指南

AI入局芯片设计验证：验证收敛50倍加速的真相与应对

本周精选

CMOS设计三大物理效应：体效应、沟道长度调制与亚阈值导电

微信小程序在线阅读毕设全解析：Java后端+MySQL+避坑实战

Android音频问题五层分析法：从App到Codec的系统化诊断

本月热门

我发现了一个新思路：用 Remotion + Claude Code 像写代码一样自动化生成短视频

Windows下 Codex 中 Chrome 和 Computer Use 插件不可用问题排查及解决参考方式：TaoToken 统一 Key 配置与验证

2026 大模型集体涨价：用 Python 做企业 Token 成本测算与选型避坑（附配置）

读完文章，想聊聊您的网站？