Language Models Need Sleep

Sangyun Lee, Sean Michael McLeish, Tom Goldstein, Giulia Fanti

Keywords: LLMs, memory, state space models, sleep, reasoning

TL;DR: We develop a sleep-like consolidation mechanism for LLMs where their attention cache is transferred into model weights via recursive forward passes.

Abstract: Abstract 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 model weights before clearing its key-value cache. During the sleep phase, 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. This shifts extra computation to the sleep phase 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 more 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.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 109