Предобучение с иерархической памятью: разделение знаний о редких и распространённых явлениях

Pretraining with hierarchical memories: separating long-tail and common knowledge

September 29, 2025

Авторы: Hadi Pouransari, David Grangier, C Thomas, Michael Kirchhof, Oncel Tuzel

cs.AI

Аннотация

Впечатляющие улучшения производительности современных языковых моделей в настоящее время основываются на масштабировании параметров: более крупные модели хранят больше знаний о мире и лучше рассуждают. Однако сжимать все мировые знания в параметры нецелесообразно, так как для каждого запроса используется лишь их часть, и это непрактично для устройств с ограниченной памятью и вычислительными ресурсами на этапе вывода. Мы устраняем этот недостаток с помощью архитектуры, дополненной памятью, и стратегии предварительного обучения, согласованной с существующими аппаратными парадигмами. Мы представляем небольшие языковые модели, которые обращаются к крупным иерархическим параметрическим банкам памяти, кодирующим мировые знания. На этапах предварительного обучения и вывода мы извлекаем небольшой контекстно-зависимый блок памяти и добавляем его к модели. Наше предварительное обучение позволяет сохранять редкие мировые знания в параметрах памяти, в то время как небольшая языковая модель выступает в роли ядра, захватывающего общие знания и способности к рассуждению. В экспериментах с масштабом в триллионы токенов мы демонстрируем значительные улучшения: модель с 160 миллионами параметров, дополненная памятью из 18 миллионов параметров, извлеченной из банка памяти объемом 4,6 миллиарда параметров, показывает сопоставимую производительность с обычной моделью, имеющей более чем в два раза больше параметров. В ходе обширных экспериментов мы исследуем оптимальный тип и размер параметрической памяти в трансформерах, масштабируя их до более чем 21 миллиарда параметров. Мы обнаруживаем, что предложенные нами иерархические прямые памяти эффективно работают в различных архитектурах трансформеров, независимо от того, добавляются ли они на этапе предварительного обучения или постфактум.

English

The impressive performance gains of modern language models currently rely on scaling parameters: larger models store more world knowledge and reason better. Yet compressing all world knowledge into parameters is unnecessary, as only a fraction is used per prompt, and impractical for edge devices with limited inference-time memory and compute. We address this shortcoming by a memory-augmented architecture and a pretraining strategy aligned with existing hardware paradigms. We introduce small language models that access large hierarchical parametric memory banks encoding world knowledge. During pretraining and inference, we fetch a small, context-dependent memory block and add it to the model. Our pretraining learns to store long-tail world knowledge in the memory parameters, while the small language model acts as an anchor capturing common knowledge and general reasoning abilities. Through trillion-token-scale experiments, we show significant gains: a 160M-parameters model augmented with an 18M-parameters memory fetched from a 4.6B memory bank obtains comparable performance to a regular model with more than 2x the parameters. Through extensive experiments, we study the optimal type and size of parametric memories in transformers, scaling them to over 21B parameters. We find that our proposed hierarchical feed-forward memories work robustly across transformer architectures, whether added during pretraining or post-hoc.