模型也需要“睡觉”?CMU新论文让LLM在梦中“巩固记忆”

很长一段时间,「长上下文」一直是各大模型厂商军备竞赛的焦点,从 128K 到 1M,再到更长的上下文窗口,业界已然形成一个固有认知,只要窗口足够大,模型就能记住更多内容,也就能处理更长、更复杂的任务。

但问题也随之而来:上下文越长,KV Cache 越臃肿,不仅导致显存瞬间被「吃光」,推理速度愈发缓慢,成本也迅速上升。

更关键的是,把更多 token 放进窗口,并不等于模型真的把这些信息转化成了可推理的长期记忆,结果是,榜单分数越刷越高,可在一些需要「深度脑暴」的复杂推理任务中,模型常常因为「记不住细节」,频频翻车……

面对这一两难问题,近日,卡内基梅隆大学(CMU)联合马里兰大学等在一篇新论文中提出了有意思的视角:既然人类连续工作久了会变笨,大模型也一样,既然如此为什么不让 LLM 睡一觉呢?

这篇论文的题目言简意赅,《Language Models Need Sleep》,也就是《语言模型需要睡眠》。

当然,这里的「睡眠」不是真的睡眠,更准确地说,是一种类似睡眠的「记忆巩固机制」。

作者认为,基于 Transformer 的大语言模型正越来越多地被用于长程任务,然而,其注意力机制在面对更长上下文时扩展性较差。为此,他们研究出了这一「记忆巩固机制」:

在睡眠过程中,模型会对累积的上下文执行 N 次离线递归前向传播,并通过一种学习得到的局部规则,更新其状态空间模型(SSM)模块中的快速权重(fast weights)。在推理阶段,这种方法把额外计算转移到「睡眠」阶段,同时保持模型在「醒着」进行预测时的延迟不变。

换句话说,它不是让模型一直把所有内容摊在眼前,而是让模型学会在某些节点「停下来想一想」,把刚刚读过的内容消化成之后还能调用的内部状态。

论文地址:https://arxiv.org/pdf/2605.26099

作者在一系列受控的合成任务上测试了该方法,包括细胞自动机、多跳图检索,以及一个更贴近真实场景的数学推理任务。在这些任务上,普通 Transformer 和 SSM-attention 混合模型都会失败,而增加模型的「睡眠」时长 N ,可以提升性能,其中在需要更深层推理的样本上,提升最为明显。

接下来,我们来详细了解一下。

从动物睡眠中获得启发

这篇论文的灵感,来自动物睡眠中的记忆巩固过程。

神经科学的研究认为,动物从短期记忆到长期记忆的转移,是受到海马体 replay 机制的支持,尤其是在睡眠期间。在这一阶段,短期的海马体记忆会被重新激活,并巩固到皮层突触权重中。睡眠会让动物无法对外部刺激作出反应,这也说明睡眠必须带来足够大的认知收益,才值得付出这一代价。

基于这一认知,作者提出了这种把上下文窗口记忆转移到持久权重中的方法,即当模型的上下文窗口在推理过程中被填满时,模型就会进入「睡眠」状态:对累积的上下文执行多次前向传播,并通过学习得到的局部规则递归地更新 fast weights,在这个阶段,模型不会接收外部输入 token。

巩固完成后,上下文窗口会被清空,模型则带着更新后的 fast weights 继续运行。在训练过程中,模型通过整个过程的反向传播进行端到端优化,以最大化睡眠之后的任务表现。

也就是说,大模型的训练过程被划分为两个阶段:

「醒着」阶段:只负责快速响应,模型就像普通的 Transformer 一样正常工作,它接收长文本输入,快速给出预测和回复,这时候它不需要对信息进行深度内化,只管「读」和「答」。

「睡眠」阶段:每隔一段时间,模型就会进入「离线睡眠状态」,期间模型会利用专门的后台时间,对积累的上下文进行 N 次循环往复的离线处理(Recurrent passes),将近期上下文中的关键细节,转化为持久的 fast weights,并写入其状态空间模型(SSM)模块中。

具体如下。

当上下文窗口被填满、模型即将从注意力层中淘汰 token 之前,模型会先进入一个「巩固阶段」,在这一阶段执行递归计算,通过这种方式扩展计算量来处理深度推理任务,对于较大的 时间步 t,仍然满足预测阶段的延迟约束。

例如,如果在全部 D 个模块上进行循环,其形式如下:

其中,N 表示在整个架构上循环执行 N 次传递。

下图对架构进行了详细描述,从一个 SSM-Attention 混合模型初始化,该模型具有固定的上下文窗口大小 L,其中注意力缓存每 L 个 token 就会被完全淘汰。在每 L 个 token 淘汰 KV Cache 之前,模型会执行 N 次递归传递,根据下面的公式 3 迭代更新 SSM 模块内部的快速权重;当 N = 1 时,它就退化为一个普通的 SSM-Attention 混合模型。模型在迭代更新快速权重的这一阶段就是「睡眠阶段」。

在递归式地细化快速权重之后,KV Cache 会被淘汰,模型随后处理接下来的 L 个 token。

在完整上下文处理完毕后,模型会基于已经细化后的记忆和当前上下文,通过一次前向传播来预测答案。训练时,模型通过对公式 6 所示的整个计算图进行反向传播,最小化预测误差,这一点与其他深度递归模型类似。

不同的是,以往的深度递归模型中,梯度会流经递归细化后的特征向量;而在这里,由于睡眠阶段结束后,细化后的特征会被丢弃,梯度实际上流经的是被细化后的快速权重。

完整的训练流程如下所示:

实验:睡得越久,推理越强?

为了验证:增加睡眠时 N,到底能不能提升模型对「旧」上下文的推理能力?作者进行了系列实验。下面我们来看一个更接近自然语言的数学推理任务 GSM-Infinite。 

GSM-Infinite 可以理解为一个长上下文数学推理基准,它会通过添加干扰 token 拉长题目,同时用所需算术操作数控制难度。题目越复杂,需要的推理步骤越多。

作者在 Jet-Nemotron 2B 和 Ouro 1.4B 两个预训练模型上测试了模型的「睡眠」机制。

结果呈现出一个清晰趋势,题目越难,「睡眠」带来的提升越明显:

对于 Jet-Nemotron 2B,6 次 sleep loop 将 6 步运算题准确率从 0.742 提升到 0.812,将 8 步运算题从 0.351 提升到 0.388;

对于 Ouro 1.4B,4 次 sleep loop 将 6 步运算题准确率从 0.419 提升到 0.615,将 8 步运算题从 0.210 提升到 0.272。

也就是说,「睡眠」机制对简单题的帮助相对没有那么明显限,因为模型本来就能做得不错;但当任务变得复杂,需要更多步推理、更强的上下文组织能力时,「睡眠」阶段的额外计算就开始发挥作用了……

局限性:效果明显,代价同样明显

当然,这篇论文并没有把问题说得过于乐观。

作者坦言,这种方法是通过把额外递归计算转移到巩固阶段,保持了预测阶段的单次前向传播延迟。但可这种收益不是免费的:在训练过程中,需要执行 N 次更深的前向和反向传播,这会让训练变慢,也可能变得不稳定。

而执行 N 次,带来效果明显提升是真,训练成本随其线性增长也是真……

因此,这项工作目前仍主要是方法论探索。

作者表示,这一方法主要贡献是方法论层面的,并且评估主要基于受控合成任务和中等规模预训练模型。目前,它还不是一个已经在超大规模商用模型、真实长程 Agent 系统中充分验证的成熟方案。

本文来自微信公众号 “机器之心”(ID:almosthuman2014),作者:关注AI的,36氪经授权发布。

发布时间:2026-06-05 17:00