Pré-entraînement avec des mémoires hiérarchiques : séparation des connaissances rares et communes

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

September 29, 2025

papers.authors: Hadi Pouransari, David Grangier, C Thomas, Michael Kirchhof, Oncel Tuzel

cs.AI

papers.abstract

Les gains impressionnants de performance des modèles de langage modernes reposent actuellement sur la mise à l'échelle des paramètres : les modèles plus volumineux stockent davantage de connaissances du monde et raisonnent mieux. Cependant, compresser toutes les connaissances du monde dans les paramètres est inutile, car seule une fraction est utilisée par prompt, et peu pratique pour les appareils périphériques disposant d'une mémoire et d'une puissance de calcul limitées lors de l'inférence. Nous abordons cette lacune grâce à une architecture augmentée par mémoire et à une stratégie de pré-entraînement alignée avec les paradigmes matériels existants. Nous introduisons de petits modèles de langage qui accèdent à de grandes banques de mémoire paramétrique hiérarchique encodant les connaissances du monde. Pendant le pré-entraînement et l'inférence, nous récupérons un petit bloc de mémoire dépendant du contexte et l'ajoutons au modèle. Notre pré-entraînement apprend à stocker les connaissances du monde à longue traîne dans les paramètres de mémoire, tandis que le petit modèle de langage agit comme une ancre capturant les connaissances communes et les capacités de raisonnement général. Grâce à des expériences à l'échelle de milliers de milliards de tokens, nous montrons des gains significatifs : un modèle de 160 millions de paramètres augmenté d'une mémoire de 18 millions de paramètres extraite d'une banque de mémoire de 4,6 milliards obtient des performances comparables à un modèle standard avec plus de deux fois plus de paramètres. À travers des expériences approfondies, nous étudions le type et la taille optimaux des mémoires paramétriques dans les transformateurs, en les mettant à l'échelle à plus de 21 milliards de paramètres. Nous constatons que nos mémoires hiérarchiques à propagation avant proposées fonctionnent de manière robuste à travers les architectures de transformateurs, qu'elles soient ajoutées pendant le pré-entraînement ou a posteriori.

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.