Preentrenamiento con memorias jerárquicas: separando el conocimiento de cola larga y común

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

Resumen

Los impresionantes avances en el rendimiento de los modelos de lenguaje modernos dependen actualmente del escalado de parámetros: los modelos más grandes almacenan más conocimiento del mundo y razonan mejor. Sin embargo, comprimir todo el conocimiento del mundo en parámetros es innecesario, ya que solo se utiliza una fracción por cada instrucción, e impráctico para dispositivos de borde con memoria y capacidad de cómputo limitadas durante la inferencia. Abordamos esta limitación mediante una arquitectura aumentada con memoria y una estrategia de preentrenamiento alineada con los paradigmas de hardware existentes. Introducimos modelos de lenguaje pequeños que acceden a grandes bancos de memoria paramétrica jerárquica que codifican conocimiento del mundo. Durante el preentrenamiento y la inferencia, recuperamos un bloque de memoria pequeño y dependiente del contexto, y lo añadimos al modelo. Nuestro preentrenamiento aprende a almacenar conocimiento del mundo de cola larga en los parámetros de la memoria, mientras que el modelo de lenguaje pequeño actúa como un ancla que captura conocimiento común y habilidades de razonamiento general. A través de experimentos a escala de billones de tokens, mostramos ganancias significativas: un modelo de 160 millones de parámetros aumentado con una memoria de 18 millones de parámetros, recuperada de un banco de memoria de 4.6 mil millones, obtiene un rendimiento comparable al de un modelo regular con más del doble de parámetros. Mediante experimentos exhaustivos, estudiamos el tipo y tamaño óptimos de memorias paramétricas en transformadores, escalándolas a más de 21 mil millones de parámetros. Encontramos que nuestras memorias jerárquicas de avance propuestas funcionan de manera robusta en diversas arquitecturas de transformadores, ya sea que se añadan durante el preentrenamiento o de manera posterior.

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.