微信扫一扫,关注公众号

  • 科技行者

  • 算力行者

见证连接与计算的「力量」

首页 高通AI研究院让AI"深度思考"不再耗尽内存:一种让推理更聪明又更省资源的新架构

高通AI研究院让AI"深度思考"不再耗尽内存:一种让推理更聪明又更省资源的新架构

2026-05-18 09:03
分享至:
----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.-
2026-05-18 09:03 科技行者

这项由高通AI研究院(Qualcomm AI Research)主导的研究成果发表于2026年5月,论文编号为arXiv:2605.07721,有兴趣深入了解的读者可以通过该编号查询完整论文。

**研究背景:当AI越想越费内存,我们该怎么办**

手机导航会在出发前规划好整条路线,而一个真正聪明的向导则会边走边思考,遇到路障时随机应变,反复斟酌哪条弯路最省时。现在的AI大模型正在经历类似的转变——从"一次性给出答案"走向"反复推敲、深度思考"。这种反复推敲的能力,被研究者称为"推理能力",而它带来的代价,就是越来越高的内存消耗。

高通AI研究院的研究团队注意到,当AI模型开始"循环思考"时,有一个隐藏的问题正在悄悄膨胀。他们提出了一种全新的架构——**记忆高效循环Transformer**(Memory-Efficient Looped Transformer,简称MELT),专门解决这个问题。在不牺牲推理能力的前提下,MELT让AI的内存消耗保持恒定,无论它思考多少轮。

**一、循环思考的代价:为什么AI越聪明越耗内存**

要理解这项研究解决的问题,先得了解"循环Transformer"(Looped Transformer)是什么。普通的AI模型就像流水线工人,输入进来,经过固定层数的处理,直接输出答案。而循环Transformer则像一个反复审稿的编辑——把同一篇文章(也就是当前的思维状态)反复交给同一批人审读,每轮审读之后再传给下一个人,直到质量满意为止。这种"循环多次处理"的方式,能够让模型在不增加参数规模的情况下,大幅提升处理复杂问题的能力。

这一思路的代表性工作是一个叫做Ouro的模型(也被称为LoopLM)。研究者发现,Ouro这类循环模型能够媲美甚至超越两倍参数量的普通模型——换句话说,通过"多想几遍",一个小模型可以打败一个大模型。

然而,这种"多想几遍"的策略有一个明显的缺陷。AI在处理文字时,需要把之前看过的内容存储在一种叫做"KV缓存"(Key-Value Cache,键值缓存)的结构里,方便后续回头参考。你可以把它理解为AI的"便条纸",每次阅读一段文字,就在便条纸上记一条笔记。

在普通模型里,每个词语只需要记一条笔记。但在循环模型里,同一个词语每循环一次就要记一条新笔记,循环10次就有10条笔记,循环20次就有20条笔记。这导致内存消耗随着思考轮数线性增长。生成32000个词的内容时,Ouro需要约28GB的内存,而类似规模的普通模型只需要约7GB。这个差距,在实际部署时往往是致命的。

**二、MELT的核心思路:便条纸不用越记越多,更新就够了**

高通团队提出的解决方案,核心思想可以用一个生活场景来理解。假设你是一名侦探,正在审查一份证词。普通做法是每次重读后都拿一张新便条纸,把最新的理解写下来,结果桌上堆满了纸。而MELT的做法是:只用一张便条纸,每次重新理解之后,把上面的内容擦掉一部分,写上更新的认识。这张便条纸始终只有一张,无论你重新思考了多少遍。

具体来说,MELT为每一层神经网络维护一个"潜在状态"(latent state),可以理解为那张会被不断更新的便条纸。每次循环时,模型不是新添一条笔记,而是通过一个可学习的"门控机制"(gating mechanism)来决定:旧的理解保留多少,新产生的认识写入多少。这个门控机制就像一个滑块——完全向左,代表完全保留旧认识;完全向右,代表用新认识完全替换旧认识;停在中间,则按比例混合。

这种设计带来的内存效益是根本性的。Ouro的KV缓存大小与"词数×循环次数"成正比,而MELT的KV缓存大小只与词数成正比,与循环次数无关。无论模型循环思考4次还是40次,内存占用都一样。

从数学角度说,这个更新规则写成公式是:门控值乘以上一时刻的潜在状态,加上(1减门控值)乘以当前输入状态。门控值本身由当前输入和上一个潜在状态共同计算得出,每个维度都有自己独立的门控值(元素级门控),而不是所有维度共用一个值。这种精细化设计让模型能对不同类型的信息采取不同的保留策略。

更新后的潜在状态,会通过两个学习得到的投影矩阵(分别叫做WK和WV)转化为注意力机制所需的"键"和"值",加入到当前这个词的缓存位置——替换旧的,而不是追加新的。

**三、为什么不直接让模型"只记最后一次"的想法行不通**

高通的研究者在设计MELT之前,其实也仔细评估了更简单的替代方案:比如直接只保留最后一轮循环产生的KV缓存,或者取所有轮次的平均值,或者用指数移动平均(EMA)。他们把Ouro模型直接套用这些策略,测试了四种变体:只保留最后一轮的缓存(无论是否保留提示词缓存)、只保留第一轮的缓存(同样分有无保留提示词缓存两种情况)。

结果令人咋舌——这四种策略在多个推理基准测试上得分全部为零。不是差一点,是完全失效。研究团队分析,这种失效不是随机的,而是一种"累积漂移":在靠近提示词的位置,缓存替换带来的误差还不明显,但随着生成的文字越来越长、越来越远离原始提示,错误会不断叠加。研究者给出了一个典型的失败案例:模型开始还在认真推导数学题,后来思路越来越混乱,最后输出的文字完全是无意义的重复。这就好比侦探走出案发现场太远之后,完全忘记了最初的线索,开始胡乱猜测。

这说明,简单地共享或复用缓存并不可行,必须通过训练让模型学会如何在单张便条纸上有效地整合信息,这正是MELT门控机制存在的价值。

**四、让老模型穿上新衣:两阶段训练如何实现平稳过渡**

MELT的架构改变是相当剧烈的——原本往缓存里"追加"内容,现在变成"覆盖更新",而且还引入了全新的门控参数。如果直接用这套新架构从零开始训练,代价极大。高通团队的思路是:从已经训练好的Ouro模型出发,通过一套精心设计的两阶段过渡流程,让MELT"继承"Ouro的知识,同时适应新的架构。

**第一阶段:分块训练与插值过渡**

MELT有一个让训练变复杂的特性:因为每个词的KV缓存依赖于前一个词处理完成后的潜在状态,所以无法像普通模型那样对整个序列并行计算。这就像一条流水线,后面的工人必须等前面的工人完工才能开始。

为了既保证训练质量又不让速度变得难以接受,团队提出了"分块训练"(chunk-wise training)策略:把一个长序列切成若干段(每段500个词),同一段内部并行计算,不同段之间按顺序传递状态。块越小,训练越接近真实推理,但越慢;块越大,训练越快,但与推理时的行为差异越大。500个词的块大小是实验中找到的较好平衡点。

但即便有了分块训练,另一个问题仍然存在:MELT的架构改变太剧烈,从Ouro的权重出发直接训练会导致模型一开始表现得像一个完全没有训练过的网络。为了让过渡更平滑,团队引入了"插值过渡"(interpolated transition)机制:在训练早期,同时计算两套KV缓存——一套按照Ouro的原始方式,另一套按照MELT的新方式。实际使用的是两套缓存的加权混合,混合系数α从0线性增长到1,耗时500步。在最开始,完全使用Ouro的缓存(模型行为等同于Ouro);随着训练推进,MELT缓存的比重越来越大,最终完全切换到MELT的行为。

在第一阶段,还额外加入了知识蒸馏(Knowledge Distillation):以原始Ouro为"教师",让MELT的每一个循环步骤的输出都去学习教师的输出。这种密集的监督信号帮助模型收敛得更快、更稳定。

**第二阶段:注意力对齐蒸馏**

当α达到1之后,MELT已经完全在自己的架构下运行了。但实验发现,如果就这样不加约束地继续训练,模型会逐渐"忘记"Ouro的推理风格,性能开始下滑。这就好比一个学生临近毕业时突然失去了老师的指导,开始走偏。

为了防止这种漂移,第二阶段引入了"注意力对齐蒸馏"(attention-aligned distillation):把Ouro模型完全冻结,作为固定的教师。MELT在每一层、每一个循环步骤的注意力机制之后产生的"中间表示",必须尽量贴近教师在同等位置产生的中间表示。这个额外的对齐损失项用一个超参数β(设为0.1)控制强度,与知识蒸馏损失一同优化。

第二阶段训练了300步,消耗了9600万个词的数据。整个两阶段训练共处理约2.56亿词,在8块H100 GPU(每块80GB显存)上运行了130小时,总计1040 GPU小时。

**五、实验结果:省了多少内存,又损失了多少性能**

研究团队将MELT-1.6B与多个竞争对手进行了系统性对比。对手包括来自Ouro的1.4B参数版本(也就是MELT的"前身"),以及同等规模的普通非循环模型:Qwen3-1.7B、Gemma4-E2B、Qwen3.5-2B和DeepSeek-R1-1.5B。评测覆盖了10个基准测试,包括AIME24、AIME25、AIME26、AMC23、MATH-500、OlympiadBench等数学推理测试,以及GPQA、HLE、MMLU-Red、HumanEval等通用推理与代码测试。所有评测使用最多32000个词的完成长度,温度参数1.0,top-p为0.7。

**内存方面**,数据非常直观。从vLLM工具中提取的精确数字显示,Ouro每生成一个词需要0.786MB的KV缓存空间,而MELT只需要0.197MB,减少了整整四倍。生成32000词的完整序列时,Ouro的KV缓存占用约25GB,加上模型本身的约2.9GB,总共接近28GB。MELT的KV缓存只有约6.3GB,总计约9.5GB,减少约2.95倍。而Qwen3-1.7B总计约7.1GB,MELT比Qwen3多出约2.5GB,这个差距主要来自Qwen使用了多查询注意力(MQA,一种通过共享键值来进一步压缩缓存的技术),而MELT目前尚未采用这一技术。

**性能方面**,MELT与Ouro相比有一定差距,但差距并不悬殊。以AIME24为例,Ouro的pass@1(提交一次答对的概率)为50.2%,MELT为46.7%;AIME25分别为36.7%和33.3%;AIME26分别为44.0%和41.0%。在数学类测试的综合平均pass@1上,Ouro得62.3分,MELT得59.9分,相差不到3分。在通用推理类测试上,MELT反而以50.1分超过了Ouro的48.6分,其中HumanEval代码测试MELT以81.7%明显领先Ouro的76.8%。

与普通非循环模型相比,MELT的优势更加突出。在数学综合平均pass@1上,Qwen3-1.7B得56.9分,Gemma4-E2B得56.0分,Qwen3.5-2B只有40.7分,DeepSeek-R1-1.5B为46.9分,而MELT以59.9分全面领先。通用推理方面,MELT的50.1分同样超过Qwen3的45.9分和Gemma4的45.5分。这些结果说明,MELT在与普通模型内存占用相近的条件下,显著超出了普通模型的推理能力。

**六、门控设计的重要性:为什么不能用简单规则替代**

研究团队专门做了一组消融实验,验证元素级门控机制是否真的必要,还是说更简单的替代方案也能达到类似效果。他们设计了四种替代方案,均在第一阶段训练结束后进行比较。

第一种替代是"均值融合":把所有循环轮次产生的KV表示取平均,作为缓存内容。第二种是"指数移动平均"(EMA,衰减因子0.2,与训练后观察到的平均门控值相近):每次循环按固定比例混合新旧缓存,等同于门控值固定为0.2的特殊情况。第三种是"只用最后一轮":完全抛弃之前轮次的信息,只保留最后一轮的缓存,在前面提到的缓存共享实验中这是性能最好的简单策略。第四种是"标量门控":只用一个数字(而不是每个维度一个数字)来控制整个状态的保留比例。

实验结果以第一阶段训练后MELT-1.6B为基准:AIME24的pass@1为44.8,AIME25为32.9,AMC23为77.7,MATH-500准确率为92.8。均值融合的对应数字分别为29.0、23.3、68.8、83.2,下降幅度明显。EMA方案略好,分别为30.2、21.5、68.6、84.6。只用最后一轮的方案表现最佳:33.7、24.0、69.7、84.0。标量门控则为34.4、23.1、66.9、85.6,与只用最后一轮相近。

所有简单替代方案都显著落后于完整的元素级门控机制,差距在10到16个百分点之间。这表明让每个维度独立学习保留比例是关键——不同类型的信息需要以不同的方式随时间演化,一刀切的规则无法满足这种需求。

**七、去掉任何一块,效果就会崩塌:组件必要性验证**

除了门控机制的变体实验,研究团队还对整个训练流程做了逐步拆解的消融实验,从完整的MELT开始,依次移除各个组件并重新训练,验证每个部分的贡献。

完整MELT-1.6B在AIME24的pass@1为46.7,pass@10(提交10次至少一次对的概率)为79.9;AIME25为33.3和61.9;AMC23为80.2和97.8;MATH-500为93.4。

移除第二阶段的注意力对齐蒸馏,退回只用第一阶段的版本,性能明显下降:AIME24降至44.8/78.1,AIME25降至32.9/66.1,AMC23降至77.7/99.3,MATH-500降至92.8。进一步移除第一阶段的插值过渡机制(直接从Ouro切换到MELT,不做平滑过渡),性能再次大幅下滑:AIME24降至35.4/63.7,AIME25降至26.9/57.9,AMC23降至73.0/93.1,MATH-500降至86.6。再移除对所有循环步骤的知识蒸馏,改为纯粹的SFT(有监督微调),数字进一步下降:AIME24为35.8/63.9,AIME25为24.4/48.6,AMC23为67.2/95.9,MATH-500为85.2。最后,如果连分块训练也取消,改用完全并行的SFT,结果是灾难性的——所有测试得分全部归零,模型完全失效。

这个逐步拆解的实验清楚地显示,这几个组件不是锦上添花,而是缺一不可:分块训练是运行的基础,知识蒸馏是收敛的保障,插值过渡是稳定的关键,注意力对齐蒸馏是性能的最后一公里。

**八、Ouro遗留的两个局限与MELT的未来方向**

研究团队在论文中坦诚地指出了MELT目前尚存的局限,其中一部分继承自Ouro,另一部分是MELT自身的新问题。

第一个局限是固定的循环次数。目前MELT(和Ouro一样)在推理时使用固定的4次循环,不管问题是"1加1等于几"还是"证明黎曼猜想",消耗的计算量是一样的。理想情况下,简单的问题应该少思考几轮,复杂的问题多思考几轮。研究者指出,MELT的常数内存设计实际上为动态调整循环深度提供了更好的基础——因为内存不会随循环数增长,增加循环次数不会带来额外的内存代价,这为"按需深度思考"的未来版本留下了空间。

第二个局限是缺少MQA支持。MQA(多查询注意力)是一种让不同注意力头共享键值数据的技术,可以进一步压缩内存。Qwen3等模型已经使用了这项技术,这也是为什么Qwen3的KV缓存比MELT还小的主要原因。研究团队认为,把MQA引入MELT是一个有潜力的方向,有望进一步缩小MELT与普通模型之间的内存差距,甚至超越普通模型的内存效率。

第三个局限是训练并行性受限。因为MELT的KV缓存依赖于前一个词的处理结果,无法像普通Transformer那样对整个序列完全并行处理,这让训练速度慢于普通模型。分块训练是当前的折中方案,但开发更高效的并行化策略仍是未来需要攻克的工程难题。

此外,研究团队在复现Ouro原论文的性能时遇到了困难,发现论文中的部分实现细节描述不够具体,导致实验结果与原论文有出入。他们也注意到,Ouro声称的"早退出"机制(让模型在认为不需要更多思考时提前结束循环)在实际代码中并未真正工作——默认配置实际上是禁用早退出的,即使触发了早退出,后续循环的计算依然全部执行。研究者认为,这是由于Ouro的KV缓存结构使得真正的早退出在不破坏自回归一致性的前提下难以实现,而MELT的常数内存设计恰好不受这个限制。

---

说到底,MELT解决的是一个实实在在的工程瓶颈:让AI模型"想更多"而不用"花更多内存"。通过把"每想一次就记一张新便条纸"改成"不断更新同一张便条纸",MELT把Ouro那套强大的循环推理能力移植到了一个与普通模型相当的内存预算之内。研究者用不到3亿参数量的新增门控组件、约2.56亿词的训练数据、1040个GPU小时的主训练,就完成了这次架构转换,并且在10个测试基准上全面超越了同等规模的普通模型。

这项研究对普通用户的意义在于:当未来这类技术落地到手机、平板或边缘计算设备时,有限的内存不再是AI深度推理的硬性门槛。更聪明的AI助手,有可能在不升级硬件的前提下出现在你的日常设备中。对研究者而言,MELT展示了一条从现有循环模型出发、以轻量级后处理训练实现架构升级的可行路径,避免了从零训练的巨大代价,这套过渡方法本身就有独立的参考价值。

感兴趣的读者可以通过arXiv编号2605.07721查阅完整论文,文中包含所有数学推导、完整的超参数设置以及更多实验细节。

---

Q&A

Q1:MELT和普通大语言模型相比,推理能力强在哪里?

A:MELT通过循环Transformer架构让模型对同一输入反复处理多轮,相当于"多想几遍"。实验显示,MELT-1.6B在数学推理综合得分上以59.9分超过了Qwen3-1.7B的56.9分、Gemma4-E2B的56.0分和DeepSeek-R1-1.5B的46.9分,同时总内存占用约9.5GB,与这些普通模型的7~9GB处于同一量级,比循环模型Ouro的28GB少了近三倍。

Q2:MELT的门控机制具体是怎么工作的?

A:门控机制本质上是一个可学习的"混合比例",对每个维度独立控制。每次循环时,模型计算一个0到1之间的门控值,决定旧状态保留多少、新输入写入多少,公式为:新状态 = 门控值×旧状态 + (1-门控值)×新输入。这个门控值本身由当前输入和上一个状态共同决定,通过训练学会在不同情况下灵活调整,比固定比例的指数移动平均或直接取平均效果显著更好。

Q3:MELT训练为什么必须用分块训练,直接并行不行吗?

A:不行。MELT的KV缓存更新存在顺序依赖:第N+1个词的缓存需要等第N个词处理完成后才能计算。普通Transformer的KV缓存只依赖当层激活值,因此可以完全并行。如果强行跳过这个依赖、用完全并行的方式训练MELT,实验结果显示所有测试得分归零,模型完全失效。分块训练通过在块内并行、块间顺序传递状态的方式,在训练速度和推理一致性之间找到了可行的平衡。

分享至
0赞

好文章,需要你的鼓励

推荐文章
----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.- ----..---.-...-/--...-.-......./-...-....-..--../-............-.-