Locas: Your Models are Principled Initializers of Locally-Supported Parametric Memories

Abstract: 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.

• The paper introduces Locas, a parametric memory mechanism that employs principled initialization to enhance test-time training and continual learning in transformers.
• It integrates Locas as a parallel module alongside Feed-Forward Networks, enabling capacity expansion without modifying pretrained model representations.
• Experimental results on language modeling and dialogue QA demonstrate that Locas achieves superior memory efficiency and performance compared to existing methods.

Locas: Principled Initializers of Locally-Supported Parametric Memories

Introduction

The paper "Locas: Your Models are Principled Initializers of Locally-Supported Parametric Memories" (arXiv ID: (2602.05085)) addresses the challenge of enhancing test-time training and continual learning in transformers. It introduces Locas, a locally-supported parametric memory mechanism designed to be integrated easily into transformer models like those with Feed Forward Network (FFN) structures. Locas leverages principled initialization techniques to improve memory efficiency and effectiveness.

Design and Variants of Locas

Locas is conceptualized as a parametric memory that can be smoothly integrated into existing transformer architectures. Transformers typically utilize dense layers like FFNs that can be seen as soft look-up table memories (Figure 1).

Figure 1: Illustration of a typical dense transformer layer with FFN interpreted as a soft look-up table memory, in comparison with the attention mechanism, which is a contextual soft look-up table mechanism. The GLU variant follows a similar structure but with an additional gating mechanism.

The memory module operates alongside the main Feed-Forward Network (FFN) to expand model capacity during testing without altering the pretrained model representations (Figure 2).

Figure 2: Architecture of the proposed Locas parametric memory integrated as a sideway FFN module in transformer layers. The memory module operates in parallel with the backbone FFN, with its output scaled and added to the main pathway. This design enables genuine model capacity expansion at test time while preserving the backbone model's pretrained representations.

The paper outlines two Locas variants:

• Locas-MLP: Employs a traditional two-layer MLP with ReLU activation, allowing theoretical guarantees in its initialization strategy.
• Locas-GLU: Utilizes the GLU-FFN structure and efficiently integrates with modern architectures like that of LLMs, ensuring seamless continual learning.

Memory Initialization

Effective initialization of the parametric memory is crucial to its performance. The paper proposes methods utilizing activation and gradient reuse for initialization in Locas-MLP, presenting a step-wise optimal approach. For Locas-GLU, an activation-guided parameter cloning strategy is suggested, selecting critical model dimensions by analyzing activation patterns during model processing.

Experimental Results

Whole-Book Language Modeling: The paper evaluates Locas on the PG-19 text corpus, demonstrating substantial efficiency in parameter use and computational resources. The Locas-GLU variant showed noteworthy performance advantages compared to TempLoRA, with a significantly lower parameter count.

Long-Context Dialogue QA: Locas outperforms baselines on the LoCoMo benchmark, especially on multi-hop and temporal reasoning questions, underscoring its ability to store and recall long-context information effectively.

Ablation Studies

The research conducted ablation studies to explore the impact of initialization strategies and memory width on model performance. Locas demonstrates consistent superiority with top-activated dimension selection for its memory initialization strategy.

Related Work

The work connects with existing research in test-time training, parameter-efficient fine-tuning, and memory-augmented neural networks. Locas contributes to the landscape by providing a method for capacity expansion without substantial modification of pretrained model parameters.

Conclusion

Locas provides an innovative solution for integrating efficient memory into transformer models, emphasizing parameter and compute efficiency. The activation-guided initialization ensures rapid convergence and minimizes catastrophic forgetting. This architecture holds promise for future refinement and integration in diverse AI applications, with potential developments focusing on dynamic allocation strategies and hybrid memory systems, combining parametric and non-parametric approaches.