Pré-treinamento com memórias hierárquicas: separando conhecimento de cauda longa e comum

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

September 29, 2025

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

cs.AI

Resumo

Os impressionantes ganhos de desempenho dos modelos de linguagem modernos atualmente dependem da escalagem de parâmetros: modelos maiores armazenam mais conhecimento do mundo e raciocinam melhor. No entanto, comprimir todo o conhecimento do mundo em parâmetros é desnecessário, pois apenas uma fração é usada por prompt, e impraticável para dispositivos de borda com memória e capacidade de computação limitadas no momento da inferência. Abordamos essa limitação com uma arquitetura aumentada por memória e uma estratégia de pré-treinamento alinhada com os paradigmas de hardware existentes. Introduzimos pequenos modelos de linguagem que acessam grandes bancos de memória paramétricos hierárquicos que codificam conhecimento do mundo. Durante o pré-treinamento e a inferência, buscamos um pequeno bloco de memória dependente do contexto e o adicionamos ao modelo. Nosso pré-treinamento aprende a armazenar conhecimento de cauda longa nos parâmetros da memória, enquanto o pequeno modelo de linguagem atua como uma âncora capturando conhecimento comum e habilidades gerais de raciocínio. Através de experimentos em escala de trilhões de tokens, mostramos ganhos significativos: um modelo de 160M parâmetros aumentado com uma memória de 18M parâmetros buscada de um banco de memória de 4,6B obtém desempenho comparável a um modelo regular com mais que o dobro dos parâmetros. Através de extensos experimentos, estudamos o tipo e tamanho ótimos de memórias paramétricas em transformadores, escalando-as para mais de 21B parâmetros. Descobrimos que as memórias hierárquicas de feed-forward propostas funcionam de forma robusta em diferentes arquiteturas de transformadores, sejam adicionadas durante o pré-treinamento ou posteriormente.

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.