Pre-addestramento con memorie gerarchiche: separazione della conoscenza a coda lunga e comune

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

September 29, 2025

Autori: Hadi Pouransari, David Grangier, C Thomas, Michael Kirchhof, Oncel Tuzel

cs.AI

Abstract

I notevoli miglioramenti nelle prestazioni dei moderni modelli linguistici attualmente si basano sulla scalabilità dei parametri: modelli più grandi memorizzano più conoscenza del mondo e ragionano meglio. Tuttavia, comprimere tutta la conoscenza del mondo nei parametri è inutile, poiché solo una frazione viene utilizzata per ogni prompt, e impraticabile per dispositivi edge con memoria e capacità di calcolo limitate durante l'inferenza. Affrontiamo questa limitazione con un'architettura aumentata dalla memoria e una strategia di pre-addestramento allineata con i paradigmi hardware esistenti. Introduciamo piccoli modelli linguistici che accedono a grandi banche di memoria parametrica gerarchica che codificano la conoscenza del mondo. Durante il pre-addestramento e l'inferenza, recuperiamo un piccolo blocco di memoria dipendente dal contesto e lo aggiungiamo al modello. Il nostro pre-addestramento impara a memorizzare la conoscenza del mondo a coda lunga nei parametri della memoria, mentre il piccolo modello linguistico funge da ancoraggio che cattura la conoscenza comune e le capacità di ragionamento generale. Attraverso esperimenti su scala trilioni di token, dimostriamo significativi miglioramenti: un modello da 160M di parametri aumentato con una memoria da 18M di parametri recuperata da una banca di memoria da 4.6B ottiene prestazioni comparabili a un modello regolare con più del doppio dei parametri. Attraverso esperimenti estensivi, studiamo il tipo e la dimensione ottimali delle memorie parametriche nei transformer, scalandole fino a oltre 21B di parametri. Troviamo che le nostre memorie gerarchiche feed-forward proposte funzionano in modo robusto tra le architetture transformer, sia aggiunte durante il pre-addestramento che in seguito.

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.