Do Language Models Need Sleep? Offline Recurrence for Improved Online Inference

LLM Long Context Latent Reasoning ML Mamba/SSM
基于Transformer的大语言模型正日益被用于长程任务,但其注意力机制在上下文长度增加时存在严重的可扩展性问题。为应对这一挑战,我们研究了一种类睡眠的巩固机制:模型周期性地将近期上下文转化为持久化的快速权重,随后清空其键值缓存。在“睡眠”阶段,模型对已累积的上下文执行 $N$ 轮离线循环处理,并通过一个可学习的局部规则更新其状态空间模型(SSM)模块中的快速权重。在推理阶段,该机制将额外计算转移至“睡眠”阶段完成,从而在保持清醒期预测低延迟的同时提升整体性能。我们在若干受控的合成任务上验证了本方法的有效性,包括元胞自动机模拟与多跳图检索任务;此外,还在一项贴近实际应用的数学推理任务上进行了测试——在此任务中,标准Transformer模型以及SSM-注意力混合模型均告失败。进一步实验表明,随着本方法中“睡眠”轮数 $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.
许愿