语言模型需要"睡眠":离线循环巩固以提升在线推理 Language Models Need Sleep

完整训练流程如算法 1 所述。第一步,将 SSM 块的所有快速权重 $S$ 初始化为 0,保证序列起点无先验。第二步,把输入序列 $x$ 与 loss mask $m$ 按窗口长度 $L$ 切分。第三步,遍历每个 chunk:先嵌入 $h=\mathrm{Embed}(c)$;若 $m_c$ 全零进入 consolidation phase,对 $h$ 和 $S$ 连续做 $N$ 次 $B(h,S)$,期间 KV 缓存累积;若 $m_c$ 含 loss 则进入 prediction phase,只做一次 $B(h,S)$;之后无论哪个 phase,KV 缓存立即清空,$S_t$ 被保留进入下一窗口。第四步,仅对 prediction phase 的 token 计算 MaskedCE,得总损失 $\mathcal{L}$。第五步,对跨窗口、跨循环的计算图反向传播并更新参数。所有 SSM 块共享同一个 $S$,与 KV cache 解耦,这种解耦让 $S$ 可独立迭代 $N$ 次而不与 attention 状态混淆。

技术新颖性主要体现在三个层面。第一,'sleep before eviction'的设计范式与传统循环模型、TTT、context distillation 有本质区别:循环次数 $N$ 成为 consolidation 算力预算的旋钮,而 prediction latency 严格不变。第二,跨窗口梯度反传通过 $S_t$ 而非细化后的特征向量 $h$,避免了递归深度模型常见的激活值爆炸问题,并允许在'同一窗口内'同时堆叠 $N$ 次循环。第三,作者明确指出这是一类 learned local update rule(类 Hebbian),而非梯度下降的一次性更新,从而获得更通用的记忆巩固算子。该方法对预测延迟零影响,仅增加训练 wall-clock(Figure 6b 显示 throughput 与 $N$ 大致反比)。在结构上,与 Ouro(depth-recurrent for prediction)、TTT(gradient step on chunk)、上下文压缩(hidden-state token)三条路线均有清晰边界。

独立分析的方法弱点有四。第一,**'睡眠'触发时机固定为 $L$ 满**:无自适应机制判断当前窗口是否真的需要 $N$ 次循环;简单任务上 $N>1$ 是浪费算力。改进方向是引入轻量控制器预测 reasoning depth 动态选 $N$,或 early-exit。第二,**只更新 SSM 快速权重 $S_t$**:固定权重限制可学习算子族;若允许 LoRA-style 微调部分 MLP 层,可能 $N$ 较小时也能获益。第三,**窗口边界硬对齐**:所有 chunk 等长且 evict 全清空;真实长上下文信息重要性分布不均,固定窗口或丢失跨边界依赖。改进方向是分层窗口或基于内容重要性的 eviction。第四,**超长上下文(>100K) scaling 未验证**:Figure 6a 仅到 12K。改进方向是在 RULER、LongBench、NIAH 等榜单系统评估 $N$、$L$ 的 scaling law。

作者在第 7 节和脚注 2 中提到的未来方向包括:用隐式梯度、截断 BPTT 等改进训练稳定性;探索自适应 sleep duration $N$;将 sleep 范式与 RL、test-time training 等其他离线算力机制结合。基于本文成果可延伸的方向有:第一,把 sleep 推广到**多模态 / 长视频 / 长代码**场景,这些场景天然有窗口化处理的需求;第二,结合**稀疏 attention 与 sleep**,只在 eviction 前做局部 sleep,减少整体训练成本;第三,研究 **sleep 的理论收敛性**——给当前 Hebbian-like 规则加上不动点/收敛条件,证明 $N\to\infty$ 时 SSM 状态有界;第四,把 sleep 与**人类睡眠-觉醒周期**的神经科学证据(如 hippocampal replay timing)做更深入类比,可能启发新的离线算力调度策略;第五,与 **Mamba3 / Jamba 等新一代 SSM** 结合,验证在更现代架构上的可移植性。

复现评估偏中等。作者提供较完整的算法 1 伪代码、超参数(Muon optimizer、AdamW lr=$5\times10^{-5}$、Muon lr=$2\times10^{-3}$ 或 $1\times10^{-3}$、batch size 512/128/256)和 backbone(Jet-Nemotron 2B、Ouro 1.4B、4 层 GDN-Att hybrid)。算力明确:Rule 110 <1 A6000 GPU-day,Depo 与 GSM-Infinite 约 1-2 H100 GPU-days/run,随机种子固定。代码与权重开源文中未明确,作者致谢 Modal。复现难度:(a) 实现'跨窗口反传 + 同窗口内 $N$ 次循环'需 custom autograd 或 activation checkpointing;(b) Jet-Nemotron 与 Ouro 均为较新预训练模型需从官方仓库获取;(c) 数据集公开但 prompt 格式与 mask 规则需对照原 paper。整体对具备 SSM 与深度循环经验的团队可复现,工程门槛不低。