Los modelos de lenguaje necesitan sueño

Resumen

Los modelos de lenguaje grandes basados en transformadores se utilizan cada vez más para tareas de horizonte largo; sin embargo, su mecanismo de atención escala deficientemente con la longitud del contexto. Para abordar esto, estudiamos un mecanismo de consolidación similar al sueño en el cual un modelo convierte periódicamente el contexto reciente en pesos rápidos persistentes antes de limpiar su caché de clave-valor. Durante el sueño, el modelo realiza N pasadas recurrentes fuera de línea sobre el contexto acumulado y actualiza los pesos rápidos en sus bloques de modelo de espacio de estados (SSM) mediante una regla local aprendida. Durante la inferencia, esto desplaza cómputo adicional al sueño mientras preserva la latencia de la predicción en estado de vigilia. Probamos nuestro método en tareas sintéticas controladas, incluyendo autómatas celulares y recuperación de grafos con múltiples saltos, así como en una tarea realista de razonamiento matemático, en la que fallan un transformador regular y modelos híbridos de atención y SSM. Luego mostramos que aumentar la duración del sueño N en nuestros modelos mejora el rendimiento, con las mayores ganancias en ejemplos que requieren un razonamiento más profundo.

English

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.