论文

语言模型需要睡眠

语言模型需要睡眠

基于Transformer的大语言模型越来越多地用于长时任务,但其注意力机制随上下文长度扩展而性能下降。为解决此问题,我们研究了一种类似睡眠的巩固机制:模型周期性地将近期上下文转换为持久快速权重,然后清空其键值缓存。在睡眠阶段,模型对累积的上下文执行N次离线循环传递,并通过一个学习到的局部规则更新其状态空间模型块中的快速权重。推理时,该方法将额外计算转移到睡眠阶段,从而保持清醒时预测的延迟。 我们在受控的合成任务(包括元胞自动机和多跳图检索)以及一个现实数学推理任务上测试了该方法。这些任务上,常规Transformer以及SSM-注意力混合模型均失败。我们证明,增加睡眠时长N可以提升模型性能,特别是在需要更深层推理的样本上获得最大收益。

论文精读

TL;DR 提出语言模型“睡眠”机制:推理间歇时离线递归更新快权重,将算力从推理转移至睡眠,在不牺牲响应速度的同时大幅提升长程推理能力。

问题

问题背景

当前大语言模型(LLM)普遍基于 Transformer 架构,其核心 自注意力机制 在长上下文场景下面临显著的扩展性瓶颈:注意力计算复杂度与上下文长度呈二次增长,KV 缓存内存占用则线性增长。这使得处理长文档、多轮对话或需要长期推理的任务时,推理成本和延迟急剧上升。

现有方法局限

为缓解该问题,业界引入了多种 高效序列模型,例如 SSM-attention 混合架构(如 Mamba、Jamba 等),它们通过引入固定大小的 快速权重记忆(Fast Weights)来压缩存储超出注意力窗口的历史信息。然而,现有工作主要聚焦于记忆容量,却忽视了长期推理能力:研究发现,当关键信息被移出活跃的 KV 缓存后,普通的 SSM-attention 混合模型在需要多跳推理(如 Depo 图检索任务)或深层逻辑推导(如 GSM-Infinite 数学推理)时性能大幅下降。这暴露了 快速权重记忆仅能支持简单的信息召回,无法有效执行对已压缩信息的复杂运算

为什么这个问题难且重要

长期推理需要模型不仅“记住”过去,还要能在后续阶段对记忆进行操作、组合与推理。注意力机制天然支持这种动态关联,但代价过高;权重记忆则牺牲了计算灵活性。如何在限定的唤醒计算预算内维持深层推理能力,是制约 LLM 走向无限上下文的关键障碍。该问题受到学界和工程界的高度关注,因为实际部署中,用户对长上下文推理(如代码库理解、长链数学证明)的需求日益增长,而现有优化方案往往顾此失彼。

行业类比

这类似于 自动驾驶系统 在车辆停止时对历史路况数据进行离线学习,以更新路径规划模型,从而在下次行驶时无需实时重放全部历史数据也能做出精准决策。

核心洞察

  • 睡眠式离线记忆巩固将长上下文推理的计算负担从在线预测转移至离线阶段,使模型保持低延迟的同时扩展有效处理范围。与现有混合模型仅依赖固定快权重或滑动窗口注意力不同,该方法周期性地将近期上下文压缩为持久快权重,并通过循环传递进行离线更新,实现了记忆的主动重组,而非简单压缩。
  • 增加睡眠时长 N 可单调提升需要深度推理的样本性能,表明类似人类睡眠的离线回放对复杂计算具有本质增益。这一发现区别于传统观点中“更大模型或更长上下文即更好”的 scaling 路径,为设计自适应推理计算资源分配提供了新维度:让模型根据问题难度动态规划所需的“思考时间”。

方法

该方法在 Transformer-SSM 混合架构 上引入睡眠巩固机制,将长上下文的计算压力从推理时转移至离线睡眠阶段,同时保持推理延迟不变。

输入与唤醒阶段

模型持续处理 token 序列,注意力层维护 KV 缓存 以支持精确的近期上下文访问,而 SSM 块 包含固定大小的快速权重 (fast weights),用于压缩更早的信息。当 KV 缓存达到容量上限或经过固定步数后,触发睡眠事件。

睡眠巩固模块

  • KV 缓存清除:睡眠开始前,近期积累的 KV 缓存被整体清空,释放资源。
  • 循环快速权重更新:模型对积累的上下文序列执行 N 次离线递归传递 (recurrent passes),每次传递时,SSM 块根据学习的局部规则更新快速权重矩阵,将此段上下文的信息压缩进持久记忆。
  • 局部学习规则:更新规则由一个小型神经网络参数化,以当前 SSM 状态、输入 token 和现有快速权重为输入,输出权重增量,实现类似 Hebbian 或塑性规则的元学习。

推理阶段

唤醒后,模型直接使用更新后的快速权重,无需重新访问被清除的 KV 缓存,即可基于压缩记忆进行推理。对于需要深层推理的示例,增加睡眠传递步数 N 可进一步提升表现。

与同类方法的差异:不同于纯 transformer(依赖无限增长的 KV 缓存)或纯 SSM(固定压缩率但缺乏动态上下文关联),该方法将动态记忆巩固与选择性遗忘结合,利用混合架构的优势,通过睡眠时的额外计算换取推理时的高效,解决了"可扩展记忆但不可扩展推理"的问题。

实验

实验设计

实验围绕三个核心任务展开:

  • Cellular Automaton:一维细胞自动机状态预测,要求模型根据较长历史推断下一步状态,测试长程依赖建模。
  • Depo:多跳图检索任务,需在节点间进行多步推理,检验模型对已逐出上下文的信息整合能力。
  • GSM-Infinite:基于 GSM8K 的扩展数学推理,问题经改写以增加推理链长度,评估深层逻辑推理。

对比基线包括普通 TransformerSSM‑attention 混合模型(无睡眠机制)。消融实验考察睡眠时长 N 的影响,并验证滑动窗口 KV 缓存驱逐策略下的表现。

关键发现

  • 普通 Transformer 在上下文超长时因 KV 缓存限制而失效,混合模型虽有压缩记忆,但无法对已不在注意力窗口内的 token 执行深度推理。
  • 引入睡眠式记忆巩固后,模型通过离线循环更新 SSM 快速权重,将近期上下文压缩进持久存储,唤醒时推理能力得以保留
  • 随着睡眠时长 N 增加,性能持续提升,增益在需深层推理的样本上尤为显著;在 Cellular Automaton 和 Depo 上,此前失败的基线模型被超越,在 GSM-Infinite 上也展现出明确改进。
  • 滑动窗口驱逐实验证实,单纯扩大窗口无法解决推理断层,而睡眠机制从架构层面提供了解耦计算方案。

基线对比解读

常规 Transformer 的注意力缓存随序列长度线性增长,计算成本平方级增加,无法扩展至极长上下文。SSM‑attention 混合模型通过固定大小的快速权重缓解了存储压力,但它们本质是前馈式的信息压缩,缺乏对历史信息的反复重处理能力。当关键推理步骤所在的 token 被移出 KV 缓存后,混合模型便丧失对其的访问,推理性能急剧退化。

本工作提出的睡眠机制突破性地将离线循环重计算引入 SSM 更新过程:在睡眠阶段模型可多次扫查积累的上下文,通过学到的局部规则调整快速权重,从而将临界信息固化到 SSM 状态中。这一过程将计算密集的操作转移到离线,保持唤醒预测的低延迟,实现了记忆压缩与深度推理的解耦,为长上下文推理提供了新范式。

行业影响

落地场景

该睡眠式记忆巩固机制直接瞄准 Transformer 模型在长上下文任务中的效率瓶颈,可嵌入各类需要长期依赖处理的产品中:

  • 企业级文档分析:合同审查、财报解读等场景,模型需在海量文本中跨段落推理,定期睡眠巩固可避免反复回扫注意力缓存,降低推理延迟。
  • 对话式 AI 助手:持续性会话中,历史交互被压缩为快速权重,模型在用户思考间隙进行离线循环,既不打断实时响应,又能维持对早期上下文的理解。
  • 自动驾驶与机器人:长序列传感器数据流处理,模型在决策间隙将观测历史固化进权重,确保实时控制回路中历史信息可追溯,而不消耗额外 KV 缓存。
  • 视频理解与生成:长视频的帧间关联推理,通过睡眠窗口将已处理帧转化为紧凑记忆,使后续帧生成或问答只依赖少量上下文,提升吞吐。

商业价值

  • 降本:大幅减少长上下文推理所需 GPU 内存和计算量,可直接缩小缓存占用,降低硬件成本;同时,离线计算可被调度到空闲或低成本设备,提高资源利用率。
  • 增收:能够支持更复杂的长文本服务,如超长法律文书分析、科研文献综述等,开辟高价值付费场景;同时,在同等硬件上支持更大并发,提升 API 服务的吞吐收益。
  • 体验提升:用户感知的延迟明显下降,因为唤醒预测不再等待长序列注意力计算;模型还能通过增加睡眠时长 N 提升复杂推理准确率,对于需要深层逻辑链的任务(如数学题解、多跳问答)效果尤其显著。

现有产品/工作流接口

该机制可作为即插即用模块嵌入主流 Transformer 或混合 SSM-Attention 架构(如 Mamba、Jamba),对现有推理栈改动较小:

  • 模型侧:在 SSM 块中添加快速权重更新规则和触发睡眠的逻辑(如达到 token 阈值或低负载周期);训练阶段可复用原有损失,仅增加睡眠循环的梯度计算。
  • 推理引擎:在服务侧实现一个睡眠调度器,根据请求流量低谷或设定的固定间隔发起 N 轮循环处理,将 KV 缓存转化为快速权重后释放内存;需协调睡眠期间的资源占用,避免与实时请求竞争 GPU。
  • 部署集成:可在模型转换时通过配置开关启用睡眠模式,配合 sliding-window 注意力等现有 eviction 策略,形成分层记忆管理;对 vLLM、TGI 等推理框架,仅需扩展自定义注意力层和调度逻辑。

具体落地 Use Case

  1. 金融研报问答系统:处理上千页的 SEC 文件或财报,在用户阅读提问间隙进行离线巩固,当被问及跨章节细节时,模型直接利用已固化的快速权重回答,而无需重新遍历全文注意力,延迟降至亚秒级且成本远低于全注意力方案。
  2. 医疗病例分析:分析跨越数年的电子健康记录,定期睡眠将旧记录压缩,新就诊时模型快速调用既往病史,辅助诊断,同时满足实时交互的延迟要求,且不因历史过长而溢出缓存。

局限

  • 论文主要局限在实验验证上:仅在可控合成任务(细胞自动机、Depo 图检索)与小规模数学推理(GSM-Infinite)上评估,未在真实世界长文本任务(如长文档问答、多轮对话、代码理解)上测试,泛化能力存疑。睡眠阶段引入了额外的离线计算开销(N 次循环传递),且 N 需人工设定,无法自动权衡任务难度与成本,实际部署时可能导致资源浪费或性能不足。
  • 方法层面,快速权重更新依赖学习到的局部规则,这额外增加了可学习参数与训练复杂度,可能对超参数敏感,不易调优。架构上要求模型中存在 SSM 块才能执行睡眠巩固,无法直接应用于纯注意力 Transformer 或其它非混合结构,限制了其适用范围。与现有记忆压缩方法(如 Compressive Transformer、Memorizing Transformers)相比,该方法在工程复杂性上更高,且未能证明其在多种记忆增强方案中的独特优势。
  • 论文对睡眠机制的消融分析有限,仅展示了睡眠时长 N 的影响,未深入探讨睡眠触发时机、不同 SSM 变体下规则学习的效果差异。此外,虽然声称将计算转移到睡眠阶段以保持唤醒延迟,但未定量对比总推理成本(在线+离线)与纯注意力模型的差异,在吞吐量指标上仅做了训练吞吐量的分析,缺乏端到端推理侧效率对比,难以完整评估其工程可行性。
论文Sangyun Lee2026-05-25原文

相关内容