Language Models Need Sleep
作者: Sangyun Lee, Sean McLeish, Tom Goldstein, Giulia Fanti
分类: cs.CL, cs.AI
发布日期: 2026-05-25
💡 一句话要点
提出睡眠机制,解决Transformer在长序列任务中的计算瓶颈
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 长序列建模 Transformer 状态空间模型 睡眠机制 离线学习
📋 核心要点
- Transformer模型在处理长序列时,注意力机制的计算复杂度成为瓶颈,限制了其在长时程任务中的应用。
- 论文提出一种类似睡眠的机制,模型周期性地将上下文信息转化为快速权重,并在状态空间模型中进行更新。
- 实验表明,该方法在合成任务和数学推理任务上优于传统Transformer和SSM-attention混合模型,且增加睡眠时间能进一步提升性能。
📝 摘要(中文)
基于Transformer的大型语言模型越来越多地用于长时程任务;然而,它们的注意力机制在上下文长度上的扩展性较差。为了解决这个问题,我们研究了一种类似睡眠的巩固机制,在该机制中,模型周期性地将最近的上下文转换为持久的快速权重,然后在清除其键值缓存之前进行此操作。在睡眠期间,模型对累积的上下文执行$N$次离线循环传递,并通过学习到的局部规则更新其状态空间模型(SSM)块中的快速权重。在推理过程中,这会将额外的计算转移到睡眠状态,同时保持唤醒时预测的延迟。我们在受控的合成任务(包括细胞自动机和多跳图检索)以及真实的数学推理任务上测试了我们的方法,在这些任务中,常规的Transformer以及SSM-attention混合模型均失败。然后,我们表明,增加模型的睡眠持续时间$N$可以提高性能,并且在需要更深层次推理的示例中,收益最大。
🔬 方法详解
问题定义:Transformer模型在处理长序列时,标准的注意力机制需要计算所有token之间的关系,导致计算复杂度随序列长度呈平方增长。这限制了Transformer在需要处理长上下文信息的任务中的应用,例如长文本理解、复杂推理等。现有的解决方案,如稀疏注意力或线性注意力,虽然降低了计算复杂度,但可能牺牲模型的表达能力。
核心思路:论文的核心思路是引入一种类似睡眠的机制,将模型处理长序列的过程分为“清醒”和“睡眠”两个阶段。“清醒”阶段模型正常进行推理,处理新输入的token;“睡眠”阶段模型将积累的上下文信息进行压缩和巩固,更新模型的内部状态(快速权重)。通过将大部分计算量转移到“睡眠”阶段,可以在保持推理速度的同时,提高模型处理长序列的能力。
技术框架:整体框架包含两个主要阶段:清醒(Wake)阶段和睡眠(Sleep)阶段。在清醒阶段,模型接收输入序列,并使用标准的Transformer或SSM-attention机制进行推理。在睡眠阶段,模型将清醒阶段积累的上下文信息(例如,键值缓存)进行处理,通过$N$次离线循环传递,使用学习到的局部规则更新状态空间模型(SSM)块中的快速权重。更新后的快速权重将在下一个清醒阶段用于指导推理。
关键创新:该方法最重要的创新点在于引入了“睡眠”机制,将长序列处理过程分解为在线推理和离线巩固两个阶段。这种机制允许模型在不影响推理速度的前提下,利用离线计算资源对长上下文信息进行更深入的处理。与现有方法相比,该方法不需要修改注意力机制本身,而是通过一种外部的“睡眠”机制来增强模型处理长序列的能力。
关键设计:睡眠阶段的关键设计包括:1) 使用状态空间模型(SSM)作为快速权重的存储和更新机制;2) 通过学习到的局部规则来更新SSM块中的快速权重,这种局部规则可以有效地捕捉上下文信息中的重要模式;3) 通过调整睡眠持续时间$N$(即离线循环传递的次数)来控制模型对上下文信息的处理深度。损失函数的设计目标是使模型在睡眠阶段能够有效地学习到上下文信息中的重要模式,并在清醒阶段能够利用这些模式进行更准确的推理。
🖼️ 关键图片
📊 实验亮点
实验结果表明,该方法在细胞自动机、多跳图检索和数学推理等任务上均取得了显著的性能提升。特别是在数学推理任务上,传统的Transformer和SSM-attention混合模型均表现不佳,而该方法通过增加睡眠持续时间$N$,能够显著提高模型的推理准确率,证明了其在处理复杂推理任务上的有效性。
🎯 应用场景
该研究成果可应用于需要处理长序列信息的各种场景,例如长文本摘要、对话系统、代码生成、知识图谱推理等。通过引入睡眠机制,可以显著提高模型在这些任务上的性能,并降低计算成本。未来,该方法有望成为构建更高效、更强大的长序列处理模型的重要技术手段。
📄 摘要(原文)
Transformer-based large language models are increasingly used for long-horizon tasks; however, their attention mechanism scales poorly with context length. To handle this, we study a sleep-like consolidation mechanism in which a model periodically converts recent context into persistent fast weights before clearing its key-value cache. During sleep, the model performs $N$ offline recurrent passes over the accumulated context and updates the fast weights in its state-space model (SSM) blocks through a learned local rule. During inference, this shifts extra computation to sleep while preserving the latency of wake-time prediction. We test our method on controlled synthetic tasks, including cellular automata and multi-hop graph retrieval, as well as a realistic math reasoning task, on which a regular transformer as well as SSM-attention hybrid models fail. We then show that increasing sleep duration $N$ for our models improves performance, with the largest gains on examples that require deeper reasoning.