Locas: Uw Modellen zijn Principiële Initialisatoren van Lokaal Ondersteunde Parametrische Geheugens

Samenvatting

In dit artikel streven we ernaar om test-time-training te verbinden met een nieuw type parametrisch geheugen dat flexibel kan worden losgekoppeld van of geïntegreerd in modelparameters. Wij presenteren Locas, een lokaal ondersteund parametrisch geheugen dat het ontwerp deelt met FFN-blokken in moderne transformers. Hierdoor kan het flexibel worden gepermanentiseerd in de modelparameters, terwijl het efficiënt continu leren ondersteunt. Wij bespreken twee belangrijke varianten van Locas: één met een conventioneel twee-laags MLP-ontwerp dat een duidelijker theoretische garantie biedt; de andere deelt dezelfde GLU-FFN-structuur met state-of-the-art LLM's en kan eenvoudig aan bestaande modellen worden toegevoegd voor zowel parameter- als reken-efficiënt continu leren. Cruciaal is dat we aantonen dat een juiste initialisatie van dergelijke low-rank zijwaartse-FFN-geheugens – uitgevoerd op een principiële manier door hergebruik van modelparameters, activaties en/of gradiënten – essentieel is voor snelle convergentie, verbeterde generalisatie en het voorkomen van catastrofaal vergeten. Wij valideren het voorgestelde geheugenmechanisme op de PG-19 whole-book taalmodellering en LoCoMo lange-context dialoogvraag-antwoordtaken. Met slechts 0,02% extra parameters in het gunstigste geval, is Locas-GLU in staat om informatie uit eerdere context op te slaan terwijl een veel kleiner contextvenster wordt aangehouden. Daarnaast testen we ook het verlies van algemene capaciteit van het model na het memoriseren van een volledig boek met Locas, via vergelijkende MMLU-evaluatie. De resultaten tonen de veelbelovende mogelijkheid van Locas om eerdere context te permanentiseren in parametrische kennis met geminimaliseerd catastrofaal vergeten van de bestaande interne kennis van het model.

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.