Locas: I tuoi modelli sono inizializzatori principiati di memorie parametriche localmente supportate.

Abstract

In questo articolo, ci proponiamo di collegare il test-time-training con un nuovo tipo di memoria parametrica che può essere flessibilmente scaricata dai parametri del modello o fusa in essi. Presentiamo Locas, una memoria parametrica a Supporto Locale che condivide la progettazione dei blocchi FFN nei transformer moderni, consentendole di essere permanentizzata in modo flessibile nei parametri del modello supportando al contempo un apprendimento continuo efficiente. Discutiamo due varianti principali di Locas: una con una progettazione MLP a due strati convenzionale che ha una garanzia teorica più chiara; l'altra condivide la stessa struttura GLU-FFN con i modelli linguistici all'avanguardia (SOTA LLM) e può essere facilmente integrata in modelli esistenti per un apprendimento continuo sia efficiente in termini di parametri che di computazione. Fondamentalmente, dimostriamo che una corretta inizializzazione di tali memorie di tipo FFN laterale a basso rango – eseguita in modo metodologico riutilizzando parametri del modello, attivazioni e/o gradienti – è essenziale per una convergenza rapida, una generalizzazione migliorata e la prevenzione della dimenticanza catastrofica. Convalidiamo il meccanismo di memoria proposto sui compiti di modellazione linguistica su interi libri PG-19 e di risposta a domande in dialoghi a contesto lungo LoCoMo. Con solo lo 0,02% di parametri aggiuntivi nel caso più contenuto, Locas-GLU è in grado di memorizzare le informazioni dal contesto passato mantenendo una finestra contestuale molto più piccola. Inoltre, testiamo anche la perdita di capacità generale del modello dopo aver memorizzato l'intero libro con Locas, attraverso una valutazione comparativa MMLU. I risultati mostrano la promettente capacità di Locas di permanentizzare il contesto passato in conoscenza parametrica con una dimenticanza catastrofica minimizzata della conoscenza interna preesistente del modello.

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.