Языковым моделям нужен сон

Аннотация

Трансформерные большие языковые модели всё чаще используются для задач с долгосрочным горизонтом; однако их механизм внимания плохо масштабируется с ростом длины контекста. Для решения этой проблемы мы исследуем механизм консолидации, напоминающий сон, в котором модель периодически преобразует недавний контекст в постоянные быстрые веса перед очисткой своего кэша ключ-значение. Во время сна модель выполняет N офлайн-рекуррентных проходов по накопленному контексту и обновляет быстрые веса в своих блоках модели пространства состояний (SSM) с помощью изученного локального правила. Во время вывода это переносит дополнительные вычисления на этап сна, сохраняя задержку предсказания в бодрствующем режиме. Мы тестируем наш метод на контролируемых синтетических задачах, включая клеточные автоматы и многопереходный поиск по графу, а также на реалистичной задаче математического рассуждения, на которых терпят неудачу как обычный трансформер, так и гибридные модели SSM-внимание. Затем мы показываем, что увеличение продолжительности сна N для наших моделей улучшает производительность, причем наибольший прирост наблюдается на примерах, требующих более глубокого рассуждения.

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.