大模型也要睡觉,而且睡得越久推理越强?

人需要睡觉,很大程度上是为了整理记忆。神经科学的主流理论认为,人在白天形成的短期记忆存储在海马体中,睡眠期间海马体反复重放这些记忆,逐步将其固化到大脑皮层的突触连接里。尽管睡眠会让动物的警觉性大幅下降,提高它们被捕食的风险,但自然选择始终没有淘汰掉它,足见这个过程对生存的价值。
既然生物都是如此,那么大模型呢?近期,卡内基梅隆大学和马里兰大学的一个研究团队从这个过程中获得启发,给语言模型也设计了一套类似的机制。
他们5月27日在arXiv发表一篇题为《语言模型需要睡眠吗?通过离线循环提升在线推理性能》(Do Language Models Need Sleep?Offline Recurrence for Improved Online Inference)的论文提出:当模型的上下文窗口装满时,不急着丢弃注意力缓存,而是先进入一个“睡眠”阶段,对当前窗口内的内容执行N次离线递归前向传播,将信息逐步写入状态空间模型(state-space model,SSM)层的快速权重中。睡醒之后注意力缓存清空,模型带着更新后的权重继续处理新输入。和真正的睡眠一样,模型在这个阶段不接收任何新的外部token,只做内部加工。

图丨在缓存清空边界处,SSM-注意力混合模型在丢弃注意力缓存之前,对当前上下文执行N次离线循环处理。这些循环处理更新了SSM模块中的快速权重,使得后续预测能够使用整合后的上下文,而无需在“唤醒”推理时进行循环。(来源:arXiv)
论文的第一作者是卡内基梅隆大学博士生Sangyun Lee,合作者包括马里兰大学的Sean McLeish和Tom Goldstein,以及卡内基梅隆大学的Giulia Fanti。
语言模型怎么处理超出注意力窗口的上下文,是一个存在已久的问题。当前主流的SSM-注意力混合架构(如NVIDIA的Jet-Nemotron、字节跳动的Ouro等)用固定大小的权重记忆来弥补有限的注意力窗口:窗口内的token靠注意力精确访问,窗口外的信息靠SSM权重压缩存储。这种设计在长序列上节省了显存和计算量,但过去的研究通常把它的局限性归结为存储容量不够。这篇论文的判断则略有不同:瓶颈不只是容量,更是把上下文转化为有用权重记忆所需的计算量。
研究者用一个简单的实验来说明这一点。他们让一个4层的GDN-注意力混合模型处理Rule 110元胞自动机任务:模型先观察初始状态,然后在注意力缓存被清空后预测若干步演化的结果。当演化步数t较小时(比如4步),模型准确率接近100%;t增大到16步时,准确率急剧下降。
整个过程中序列长度始终不变,需要存储的信息量也没有增加,变的只是从已有信息推导出答案所需的推理深度。也就是说,SSM层的快速权重有足够空间存下输入内容,但模型没有足够的计算时间把这些内容“消化”成能支撑后续推理的形式。
论文的方案是在缓存清空之前给模型一段离线计算时间。当上下文窗口填满后,模型先对当前窗口内容做N次完整的前向传播,覆盖所有Transformer块,每次传播都更新SSM块中的快速权重。这N次循环就是所谓的“睡眠”。睡眠结束后缓存清空,模型继续处理后续输入。
在最终预测阶段,模型只做一次标准前向传播,额外的计算已经花在了睡眠期,预测时延迟不受影响。整个过程端到端训练,梯度直接通过N次递归更新的权重回传。当N=1时,它退化为普通的SSM-注意力混合模型。
团队在三个任务上验证了效果。在Rule 110元胞自动机任务上(演化32步),不做循环的基线模型在50亿训练token后准确率约10%,和随机猜测差不多。2次循环提升到约20%,3次和4次循环均超过30%。
在Depo多跳知识检索任务上,这一机制的效果更明显。这个任务要求模型在注意力缓存已清空的情况下,沿着一个有向图做k跳遍历。不循环的模型在4跳及以上的查询上几乎没有学到任何东西,2次循环的模型在8跳以上同样停滞,只有4次循环的模型开始在最难的16跳任务上取得进展。

图丨增加N可提升Depo的性能。4层GDN-attention混合模型在k跳知识检索任务上的测试损失。额外的离线循环加速了学习过程,对于推理密集型、跳数更高的查询效果尤为显著。(来源:arXiv)
在GSM-Infinite数学推理任务上,研究者分别对预训练的Jet-Nemotron 2B和Ouro 1.4B进行了微调验证。对于较简单的2步和4步运算题,不同循环次数的差异不大,准确率都接近饱和。
但在6步和8步运算题上差距拉开了:Jet-Nemotron上,6次循环将6步运算的准确率从0.742提升到0.812(提升9%),8步运算从0.351到0.388(提升11%)。Ouro上的提升更大,4次循环将6步运算的准确率从0.419提升到0.615(提升47%),8步运算从0.210到0.272(提升30%)。
论文还测试了滑动窗口驱逐策略(而非硬清空),在窗口大小L=512的条件下,4次循环将2步运算的准确率从0.596提升到0.905。这表明当注意力窗口远小于序列总长时,更长的“睡眠”不仅对深度推理有帮助,对基本的信息检索和压缩也有帮助。
当然,多循环几次也是有代价的。训练阶段,N次递归前向和反向传播使得训练吞吐量大致与N成反比。不过论文指出,当窗口大小L足够大时,跨窗口的串行依赖不会显著拖慢训练,GPU在处理单个窗口时仍然是满载的。
推理阶段则不同:由于所有额外计算都发生在上下文处理阶段,最终预测时仍然只需一次前向传播,不增加推理延迟。这是它和思维链(chain-of-thought)等推理时扩展方法的一个区别,它把计算开销转移到了记忆整合阶段而非预测阶段。论文也承认,递归深度训练本身存在不稳定性,是当前深度递归模型研究中的活跃问题。
从实验数据看,这笔额外开销放在更难的问题上能带来更大的收益,简单任务上多循环几次和不循环差别不大,只有在需要更深推理的任务上,差距才真正拉开。如果把上下文转化为有用记忆本身是一项计算密集型工作,那么更难的问题自然需要更多的“消化”时间。
参考资料:
1.https://arxiv.org/abs/2605.26099




