Locas: Tus Modelos son Inicializadores Principados de Memorias Paramétricas de Soporte Local

Resumen

En este artículo, pretendemos conectar el entrenamiento en tiempo de prueba con un nuevo tipo de memoria paramétrica que puede descargarse o fusionarse de forma flexible con los parámetros del modelo. Presentamos Locas, una memoria paramétrica de Soporte Local que comparte el diseño de los bloques FFN en los transformadores modernos, lo que le permite ser permanentemente integrada en los parámetros del modelo mientras soporta un aprendizaje continuo eficiente. Discutimos dos variantes principales de Locas: una con un diseño convencional de MLP de dos capas que tiene una garantía teórica más clara; la otra comparte la misma estructura GLU-FFN con los LLMs de última generación y puede adjuntarse fácilmente a modelos existentes para un aprendizaje continuo eficiente en parámetros y en computación. Crucialmente, demostramos que una inicialización adecuada de estas memorias de tipo FFN lateral de bajo rango —realizada de manera fundamentada reutilizando parámetros del modelo, activaciones y/o gradientes— es esencial para una convergencia rápida, una mejora en la generalización y la prevención del olvido catastrófico. Validamos el mecanismo de memoria propuesto en las tareas de modelado de lenguaje de libros completos PG-19 y de respuesta a preguntas en diálogos de contexto largo LoCoMo. Con solo un 0.02% de parámetros adicionales en el caso más bajo, Locas-GLU es capaz de almacenar la información del contexto pasado manteniendo una ventana de contexto mucho más pequeña. Además, también probamos la pérdida de capacidad general del modelo después de memorizar el libro completo con Locas, mediante una evaluación comparativa MMLU. Los resultados muestran la prometedora capacidad de Locas para convertir el contexto pasado en conocimiento paramétrico permanente, minimizando el olvido catastrófico del conocimiento interno existente del modelo.

English

In this paper, we aim to bridge test-time-training with a new type of parametric memory that can be flexibly offloaded from or merged into model parameters. We present Locas, a Locally-Supported parametric memory that shares the design of FFN blocks in modern transformers, allowing it to be flexibly permanentized into the model parameters while supporting efficient continual learning. We discuss two major variants of Locas: one with a conventional two-layer MLP design that has a clearer theoretical guarantee; the other one shares the same GLU-FFN structure with SOTA LLMs, and can be easily attached to existing models for both parameter-efficient and computation-efficient continual learning. Crucially, we show that proper initialization of such low-rank sideway-FFN-style memories -- performed in a principled way by reusing model parameters, activations and/or gradients -- is essential for fast convergence, improved generalization, and catastrophic forgetting prevention. We validate the proposed memory mechanism on the PG-19 whole-book language modeling and LoCoMo long-context dialogue question answering tasks. With only 0.02\% additional parameters in the lowest case, Locas-GLU is capable of storing the information from past context while maintaining a much smaller context window. In addition, we also test the model's general capability loss after memorizing the whole book with Locas, through comparative MMLU evaluation. Results show the promising ability of Locas to permanentize past context into parametric knowledge with minimized catastrophic forgetting of the model's existing internal knowledge.