语言模型需要睡眠

# 语言模型需要睡眠
来源：https://arxiv.org/html/2605.26099
Sangyun Lee 卡内基梅隆大学 & Sean McLeish 马里兰大学 & Tom Goldstein 马里兰大学 & Giulia Fanti 卡内基梅隆大学

###### 摘要

基于 Transformer 的大型语言模型越来越多地用于长时程任务，然而它们的注意力机制随上下文长度扩展性不佳。为了解决这个问题，我们研究了一种类似睡眠的整合机制，其中模型定期将最近的上下文转换为持久的快速权重，然后清空其键值缓存。在睡眠期间，模型对累积的上下文执行 `N` 次离线循环传递，并通过学习到的局部规则更新其状态空间模型（SSM）块中的快速权重。在推理时，这将额外的计算转移到睡眠阶段，同时保持清醒时预测的延迟。我们在受控的合成任务上测试了我们的方法，包括元胞自动机和多跳图检索，以及一个实际的数学推理任务，而常规 Transformer 和 SSM-注意力混合模型在这些任务上均失败。然后我们表明，增加我们的模型的睡眠持续时间 `N` 会提升性能，在需要更深推理的示例中增益最大。

## 1 引言

大型语言模型（LLMs）通常基于 Transformer 架构[51 (https://arxiv.org/html/2605.26099#bib.bib4)]，该架构将上下文存储在注意力缓存中，并根据需要检索过去的标记。这种记忆机制是其性能的核心，但其扩展性不佳：总的注意力计算量随上下文长度呈二次方增长，而缓存内存呈线性增长。

近期的高效序列模型[42 (https://arxiv.org/html/2605.26099#bib.bib49), 18 (https://arxiv.org/html/2605.26099#bib.bib50), 16 (https://arxiv.org/html/2605.26099#bib.bib51), 2 (https://arxiv.org/html/2605.26099#bib.bib31)] 通过引入与完全自注意力交错的固定大小快速权重记忆[53 (https://arxiv.org/html/2605.26099#bib.bib1), 14 (https://arxiv.org/html/2605.26099#bib.bib52), 43 (https://arxiv.org/html/2605.26099#bib.bib41)] 来缓解这一成本。这种混合设计将两种互补的记忆形式结合：注意力用于高保真地访问最近的标记，而基于权重的记忆用于压缩超出活跃上下文窗口的信息。混合模型如今在大型前沿模型中很常见[49 (https://arxiv.org/html/2605.26099#bib.bib61)]。

然而，可扩展的记忆并不等同于可扩展的推理。快速权重记忆可能支持长距离回忆[42 (https://arxiv.org/html/2605.26099#bib.bib49)]，但尚不清楚它是否支持对不再存在于 KV 缓存中的标记进行深度计算。我们发现，在相同的标记预算下，随着所需推理深度的增加，即使是标准的 SSM-注意力混合模型的性能也会下降，即使要存储的信息量保持不变。这表明瓶颈不仅仅是先前工作[27 (https://arxiv.org/html/2605.26099#bib.bib32), 2 (https://arxiv.org/html/2605.26099#bib.bib31)] 所暗示的记忆容量，而是可用于将被驱逐的上下文转换为有用内部状态的计算量。

**睡眠。** 在动物中，短时记忆向长时记忆的转移被认为由海马回放[33 (https://arxiv.org/html/2605.26099#bib.bib15)] 支持，尤其是在睡眠期间[41 (https://arxiv.org/html/2605.26099#bib.bib16)]；在这一阶段，海马的短时记忆被重新激活并整合到皮层突触权重中。睡眠使动物无法对外部刺激做出反应，这表明它必须提供足够的认知益处来证明这种代价是合理的[41 (https://arxiv.org/html/2605.26099#bib.bib16)]。受这些生物过程的启发，我们提出了一种将上下文窗口记忆转移到持久权重的方法。当模型在推理过程中上下文窗口变满时，模型进入“睡眠”状态，在睡眠期间它对累积的上下文执行多次前向传递，并通过学习到的局部规则递归地更新其快速权重。与动物睡眠一样，模型在此阶段不接收外部输入标记。整合完成后，上下文窗口被清空，模型以更新后的快速权重恢复运行。在训练期间，通过整个过程的端到端反向传播来优化模型，以最大化睡眠后的任务性能。

我们的架构也受到深度循环或循环神经网络结果[23 (https://arxiv.org/html/2605.26099#bib.bib22), 17 (https://arxiv.org/html/2605.26099#bib.bib24), 4 (https://arxiv.org/html/2605.26099#bib.bib25)] 的启发。先前的研究表明，在顺序推理任务上，动态深度模型可以优于固定深度模型，并通过在预测时扩展计算量来解决固定深度模型无法解决的困难问题实例。我们的关键见解是，循环不仅可用于预测，还可用于记忆整合。将观察到的标记转换为有用的权重记忆本身就是一项不平凡的计算，并且不能在一次传递中完成。事实上，许多学习算法（例如梯度下降）通过迭代更新权重来改进。因此，在快速权重形成期间分配更多的循环计算，可为模型提供更多步骤来将上下文转换为支持后续预测的表示。我们发现，增加循环深度（即睡眠持续时间）可以改善睡眠后的推理。与先前的循环模型不同，我们的模型在预测时不需要循环：额外的计算已用于形成支持后续单次预测的快速权重。

我们引入并评估了 LLM 睡眠在精心设计的合成任务上的表现，在这些任务中，模型必须仅使用一次前向传递来回答关于已被驱逐的上下文的问题。这些合成任务使我们能够在保持记忆负载不变的情况下变化推理深度，为睡眠期间的计算是否可以将临时上下文转换为支持后续推理的快速权重提供了一个清晰的压力测试。我们将贡献总结如下：

- • 在受控设置中，我们表明，随着问题推理深度的增加，即使有足够的快速权重容量，标准状态空间模型（SSM），例如门控 Delta 网络（GDN），也会失败。
- • 我们提出了一种将循环计算与快速权重记忆块相结合的架构，并表明增加我们架构的递归次数可以提高优于 GDN 的性能。我们在需要最深推理的问题实例上观察到最大的增益。
- • 我们通过使用预训练 LLM 初始化，在自然语言数学推理数据集 GSM-Infinite 上进一步验证了我们架构的有效性。

总的来说，这些结果支持核心观点：类似睡眠的离线循环可以将被驱逐的上下文组织成权重，以支持后续推理。

## 2 相关工作

**快速权重和线性递归神经网络。** 线性递归神经网络或 SSM 可以被视为维护在线快速权重记忆，而不是随序列长度呈二次方增长的 KV 缓存。在这种观点下，线性注意力对应于对固定大小的矩阵值状态的递归更新，其中键值映射被写入和查询[29 (https://arxiv.org/html/2605.26099#bib.bib42), 43 (https://arxiv.org/html/2605.26099#bib.bib41)]。最近的变体通过 Delta 规则更新和门控改进了这种记忆，实现了更具选择性的写入、覆盖和遗忘[54 (https://arxiv.org/html/2605.26099#bib.bib46), 53 (https://arxiv.org/html/2605.26099#bib.bib1), 55 (https://arxiv.org/html/2605.26099#bib.bib45), 14 (https://arxiv.org/html/2605.26099#bib.bib52)]。这些机制支撑了最近的高效混合语言模型[24 (https://arxiv.org/html/2605.26099#bib.bib10), 39 (https://arxiv.org/html/2605.26099#bib.bib9)]，并有助于解释为什么线性网络可以提供有利的召回、吞吐量和内存权衡。如先前工作[2 (https://arxiv.org/html/2605.26099#bib.bib31), 27 (https://arxiv.org/html/2605.26099#bib.bib32)] 所指出的，由于固定记忆大小，它们在某些情况下相对于完全注意力仍然难以实现精确复制和检索。与这些工作相反，我们表明，当解决任务所需的推理深度增加时，即使要存储的信息量保持不变，此类模型也可能失败。

**上下文压缩。** 有几种方法通过在测试时处理长上下文来压缩上下文信息。Ge 等人[21 (https://arxiv.org/html/2605.26099#bib.bib11)] 提出使用语言模型将长上下文压缩成更短的隐藏状态序列，然后将这些隐藏状态传递给语言模型以替代原始的长上下文。Eyuboglu 等人[20 (https://arxiv.org/html/2605.26099#bib.bib47)] 使用离线自学来学习一个小型 KV 缓存，该缓存可以替代全上下文缓存。这一系列工作与我们目标相同：花费离线计算一次性将长上下文转换为紧凑状态，以便以后重用。这些方法缩短了注意力上下文中保留的内容，而我们的方法将驱逐的上下文转移到基于权重的记忆中。

**上下文蒸馏。** 上下文蒸馏[46 (https://arxiv.org/html/2605.26099#bib.bib66), 3 (https://arxiv.org/html/2605.26099#bib.bib65)] 旨在通过训练没有上下文的模型来模仿有上下文的教师模型[46 (https://arxiv.org/html/2605.26099#bib.bib66), 3 (https://arxiv.org/html/2605.26099#bib.bib65), 8 (https://arxiv.org/html/2605.26099#bib.bib68)]，重构它[11 (https://arxiv.org/html/2605.26099#bib.bib67)]，预测其续文[8 (https://arxiv.org/html/2605.26099#bib.bib68), 11 (https://arxiv.org/html/2605.26099#bib.bib67)]，或回答关于它的问题[47 (https://arxiv.org/html/2605.26099#bib.bib70), 9 (https://arxiv.org/html/2605.26099#bib.bib69), 8 (https://arxiv.org/html/2605.26099#bib.bib68)]，从而将活跃上下文蒸馏到模型权重中。与在预定义损失上进行梯度下降不同，我们的方法使用学习到的循环前向传递将上下文转移到权重中。

**测试时训练。** Tandon 等人[48 (https://arxiv.org/html/2605.26099#bib.bib48)] 用滑动窗口注意力替换完全注意力，并在 MLP 层的一个子集上执行测试时梯度更新。在推理时，他们的方法在观察到的上下文上优化标准交叉熵损失，将长距离信息存储在临时参数更新中，而不是存储在完整的 KV 缓存中。他们只为蒸馏每个上下文块执行一个梯度步。相比之下，我们的方法使用学习到的循环前向传递作为记忆更新规则，允许更灵活的整合形式，而不必对应于固定标量目标上的单步梯度下降。他们主要在通用网络文本数据上评估困惑度，其中检索和推理需求交织在一起；而我们使用合成任务独立控制推理深度和问题长度，表明当推理深度增加时，额外的睡眠时间计算最为有益。Zhang 等人[56 (https://arxiv.org/html/2605.26099#bib.bib72)] 附加了一个 LoRA 适配器，该适配器根据当前上下文块更新模型权重，并在强化学习设置中评估这种方法。与我们的不同，他们的方法每个块只更新一次权重。

**深度循环模型。** 增加语言模型的深度已知可以提高其表达能力[35 (https://arxiv.org/html/2605.26099#bib.bib23)]。深度循环是增加 Transformer 模型深度的一种方法，也是使其成为图灵完备的一种方法[17 (https://arxiv.org/html/2605.26099#bib.bib24)]。此外，这些模型的深度可以是自适应的[23 (https://arxiv.org/html/2605.26099#bib.bib22), 19 (https://arxiv.org/html/2605.26099#bib.bib38), 44 (https://arxiv.org/html/2605.26099#bib.bib39), 5 (https://arxiv.org/html/2605.26099#bib.bib40)]。最近的工作已将这些深度自适应语言模型扩展到大规模，既包括从头开始训练[22 (https://arxiv.org/html/2605.26099#bib.bib20), 58 (https://arxiv.org/html/2605.26099#bib.bib34)]，也包括作为训练后目标[34 (https://arxiv.org/html/2605.26099#bib.bib37)]。关于如何最好地训练深度循环模型的详细分析表明，循环深度应与训练计算量成比例缩放[40 (https://arxiv.org/html/2605.26099#bib.bib35), 45 (https://arxiv.org/html/2605.26099#bib.bib36)]。

**离线规划。** 在结构化环境中成功的规划通常需要将新观察到的信息与早期状态的记忆相结合。一个长期的观点认为，动物在决策时刻在线进行这种整合[50 (https://arxiv.org/html/2605.26099#bib.bib57), 36 (https://arxiv.org/html/2605.26099#bib.bib56)]。然而，在决策时刻整合遥远的记忆可能耗费时间，而在任务休息期间进行离线规划可以分摊这种成本[36 (https://arxiv.org/html/2605.26099#bib.bib56)]。与此观点一致，Momennejad 等人[36 (https://arxiv.org/html/2605.26099#bib.bib56)] 表明，休息期间离线回放的神经证据可以预测人类受试者规划表现的改善。机器学习领域的近期工作使用人工神经网络研究了相关机制。Lin 等人[30 (https://arxiv.org/html/2605.26099#bib.bib54)] 提出通过让 LLM 生成用户可能提出的预期问题并预计算解决它们所需的量来扩展离线计算。Chalvidal 等人[10 (https://arxiv.org/html/2605.26099#bib.bib55)] 在强化学习环境中训练单层网络，并表明递归的类 Hebbian 权重更新支持快速适应。在本文中，我们表明在类似睡眠的离线阶段递归更新快速权重可以改善对被驱逐上下文的推理，同时保持严格的预测阶段延迟约束。

## 3 预备知识

### 3.1 序列混合器

**注意力。** Softmax 注意力[51 (https://arxiv.org/html/2605.26099#bib.bib4)] 是一种序列混合操作，其中每个标记根据查询-键相似度从先前的标记中检索信息。对于时间步 \(t\) 的标记表示 \(\boldsymbol{x}_{t}\)，定义

\[
\begin{aligned}
\boldsymbol{q}_{t} &= \mathbf{W}_{Q}\boldsymbol{x}_{t}, \\
\boldsymbol{k}_{t} &= \mathbf{W}_{K}\boldsymbol{x}_{t}, \\
\boldsymbol{v}_{t} &= \mathbf{W}_{V}\boldsymbol{x}_{t},
\end{aligned}
\]

其中 \(\boldsymbol{q}_{t}, \boldsymbol{k}_{t}, \boldsymbol{v}_{t} \in \mathbb{R}^{d}\) 是列向量，\(\mathbf{W}_{Q}, \mathbf{W}_{K}, \mathbf{W}_{V}\) 是形状兼容的学习投影矩阵。自注意力将所有先前的键 \(\boldsymbol{k}_{t}\) 和值 \(\boldsymbol{v}_{t}\) 存储在 \(\mathbf{K}_{t} = [\boldsymbol{k}_{1}, \ldots, \boldsymbol{k}_{t}]^{\top} \in \mathbb{R}^{t \times d}\) 和 \(\mathbf{V}_{t} = [\boldsymbol{v}_{1}, \ldots, \boldsymbol{v}_{t}]^{\top} \in \mathbb{R}^{t \times d}\) 中，然后计算

\[
\boldsymbol{o}_{t} = \mathbf{V}_{t}^{\top} \operatorname{softmax}\!\left( \frac{\mathbf{K}_{t}\boldsymbol{q}_{t}}{\sqrt{d}} \right).
\]

这使得 \(\boldsymbol{x}_{t}\) 可以关注任何先前的标记，但需要存储 \(\mathbf{K}_{t}\) 和 \(\mathbf{V}_{t}\)，即 KV 缓存，其大小呈线性增长。