Locas enables efficient test-time adaptation by principally initializing memory modules using model activations and gradients.
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.
Locas uses model activations and gradients to initialize memory modules, reducing parameters by 85% compared to existing test-time training methods.
FFN layers can be interpreted as soft lookup tables with fixed memory slots, enabling efficient test-time adaptation without catastrophic forgetting.
Locas-GLU variant is compatible with modern LLMs like LLaMA and Mistral, requiring no architectural changes for deployment.
Locas-MLP has compatibility issues with GLU-style FFN blocks used in state-of-the-art models, limiting its practical applicability.
The paper doesn't fully address scalability challenges for extremely long documents beyond whole-book language modeling tasks.
Shaobo Wang, Xuan Ouyang +10
As high-quality public text approaches exhaustion, a phenomenon known as the Data Wall, pre-training is shifting from more tokens to better tokens. However, existing methods either rely on heuristic s...
Zhongzhi Li, Xuansheng Wu +3
The diversity of post-training data is critical for effective downstream performance in large language models (LLMs). Many existing approaches to constructing post-training data quantify diversity usi...
Zixuan Huang, Xin Xia +12
Recent advancements in large reasoning models (LRMs) have greatly improved their capabilities on complex reasoning tasks through Long Chains of Thought (CoTs). However, this approach often results in ...
Tengyue Xu, Zhuoyang Qian +17
Autonomous scientific discovery with large language model (LLM)-based agents has recently made substantial progress, demonstrating the ability to automate end-to-end research workflows. However, exist...
Zhiyuan Hu, Yucheng Wang +8
Reinforcement learning (RL) has become a central paradigm for post-training large language models (LLMs), particularly for complex reasoning tasks, yet it often suffers from exploration collapse: poli...
