模型也需要「睡觉」？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 系统中充分验证的成熟方案。

更多详情，可查看论文了解！

https://x.com/iScienceLuvr/status/2059221770075562113

https://arxiv.org/pdf/2605.26099