Pretraining met hiërarchische geheugens: het scheiden van lange-staart en algemene kennis

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

September 29, 2025

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

cs.AI

Samenvatting

De indrukwekkende prestatieverbeteringen van moderne taalmodelle berusten momenteel op het schalen van parameters: grotere modellen slaan meer wereldkennis op en redeneren beter. Het comprimeren van alle wereldkennis in parameters is echter onnodig, aangezien slechts een fractie ervan per prompt wordt gebruikt, en onpraktisch voor edge-apparaten met beperkt geheugen en rekenkracht tijdens inferentie. Wij pakken dit tekort aan met een geheugen-augmented architectuur en een pretrainingsstrategie die aansluit bij bestaande hardwareparadigma's. We introduceren kleine taalmodelle die toegang hebben tot grote hiërarchische parametrische geheugenbanken die wereldkennis coderen. Tijdens pretraining en inferentie halen we een klein, contextafhankelijk geheugenblok op en voegen dit toe aan het model. Onze pretraining leert om langstaart-wereldkennis op te slaan in de geheugenparameters, terwijl het kleine taalmodel fungeert als een anker dat algemene kennis en redeneervaardigheden vastlegt. Door experimenten op schaal van biljoenen tokens tonen we aanzienlijke verbeteringen aan: een model met 160M parameters, versterkt met een 18M-parameters geheugen dat wordt opgehaald uit een 4,6B geheugenbank, behaalt vergelijkbare prestaties als een regulier model met meer dan twee keer zoveel parameters. Door uitgebreide experimenten bestuderen we het optimale type en de grootte van parametrische geheugens in transformers, waarbij we deze opschalen tot meer dan 21B parameters. We ontdekken dat onze voorgestelde hiërarchische feed-forward geheugens robuust werken in verschillende transformer-architecturen, ongeacht of ze tijdens pretraining of achteraf worden toegevoegd.

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.