模型也需要“睡觉”？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 继续运行。在训练过程中，模型通过整个过程的反向传播进行端到端优化，以最大化睡眠之后的任务表现。

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