LLM Sleep 机制发布：通过类睡眠离线递归提升长上下文推理

卡内基梅隆大学与马里兰大学的研究人员提出一种名为 LLM Sleep 的离线递归机制, 旨在解决 Transformer 架构在处理长上下文时计算量随长度平方级增长的问题。该方法通过模拟生物睡眠中的记忆巩固过程, 将短期上下文转换为持久的 fast weights, 从而在清空 KV cache 后仍能支持深层推理。

LLM Sleep 机制：从 KV cache 到持久权重

传统的 Transformer 模型依赖 attention 缓存来存储上下文, 但这种机制在长序列任务中扩展性较差。据 arXiv 预印本论文 Language Models Need Sleep 披露, 研究团队设计了一种类睡眠的巩固机制。当模型在推理过程中上下文窗口满载时, 会进入 “睡眠”(Sleep) 状态。

在睡眠阶段, 模型不再接收外部输入, 而是对累积的上下文执行 N 次离线递归前向传递。通过学习到的局部规则, 模型将近期上下文信息递归地更新到状态空间模型 (SSM) 模块的持久 fast weights 中。完成巩固后, 模型清空 KV cache 并恢复运行, 此时更新后的权重已承载了此前被剔除的上下文信息。

解决长上下文中的深层推理瓶颈

研究团队指出, 目前的 SSM-attention 混合架构虽然能通过固定大小的权重内存支持长距离召回, 但在处理需要深层计算的任务时表现不佳。他们发现, 即使存储信息量保持不变, 随着推理深度增加, 常规混合模型的性能会显著下降。

LLM Sleep 的核心洞察在于, 将观察到的 token 转换为有用的权重内存本身就是一项复杂的计算任务, 单次前向传递难以完成。通过增加递归深度 (即延长睡眠时长 N), 模型拥有更多步骤将上下文转化为支持后续预测的内部表示。这种设计将额外的计算量转移到了睡眠阶段, 从而在保持预测时低延迟的同时, 提升了推理能力。

合成任务与数学推理验证

为了验证该机制, 研究者在受控的合成任务 (如细胞自动机和多跳图检索) 以及自然语言数学推理数据集 GSM-Infinite 上进行了测试。实验结果显示:

• 在推理深度增加时, 传统的 SSM(如 Gated Delta Nets) 即使拥有足够的权重容量也会失效;
• 增加递归次数 N 能显著提升 LLM Sleep 模型的表现, 且在需要最深层推理的样本中提升最明显;
• 该架构在预训练 LLM 初始化的基础上, 成功在数学推理任务中证明了其有效性。

这种离线递归机制证明了通过在非预测阶段投入计算资源, 可以将瞬时上下文组织成结构化权重, 从而在单次前向传递中完成复杂的长程推理任务。