PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Sidi Lu, Zhenwen Liang, Dongyang Ma, Yan Wang, Haitao Mi, Dong Yu

arXiv:2602.05085v1cs.CL

TL;DR

Test-time adaptation must balance context-length limits against the optimization cost and design constraints of parametric updates. The paper introduces Locas, a principled FFN-style parametric memory with MLP and GLU variants, and reports efficient continual learning with low parameter overhead and minimal catastrophic forgetting.

  • Problem

    Existing adaptation methods are constrained by context length or incur additional optimization cost and require careful objective design for useful adaptivity.

  • Method

    Locas uses backbone-guided initialization for locally supported FFN-style memories, offering theoretical MLP and GLU-compatible variants with nonlinear compression.

  • Results

    Locas-GLU achieves comparable or superior performance to TempLoRA using 15% of the additional parameters and 38% of the computational overhead, while MMLU degradation is 0.1-0.2%.

  • Takeaways & Limitations

    Locas supports permanentizing past context into parametric knowledge through capacity expansion with parameter and compute efficiency and minimized catastrophic forgetting.

  • Takeaways & Limitations

    NL-SVD has substantially higher computational cost and no clear final-perplexity advantage over backpropagation, while future work includes dynamic and hierarchical memory allocation.

Abstract

from arXiv · show

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.

1 Introduction

Locas addresses the efficiency limits of test-time adaptation by introducing a locally supported parametric memory with principled initialization. Its variants and compression method target efficient continual learning while retaining theoretical guarantees and practical compatibility.

  • Motivation: Existing in-context learning is bounded by context length, while test-time training can require multiple gradient-based iterations per token and careful objective design.
  • Approach: Locas uses backbone-model behavior to initialize newly introduced memory modules, aiming to improve parameter and compute efficiency, convergence, and generalization.
  • Variants: Locas-MLP provides theoretical guarantees through a two-layer MLP, whereas Locas-GLU matches modern GLU-FFN structures for flexible attachment to existing models.
  • Compression: The proposed Non-Linear SVD extends SVD to two-layer nonlinear memories for compression, though it has practical limitations compared with standard backpropagation.
  • Evaluation: Locas is evaluated against full-attention, context-truncation, and TempLoRA baselines on PG-19 whole-book language modeling and LoCoMo dialogue question answering.

2 Methodology

Locas treats transformer FFNs as persistent soft look-up-table memories and adds locally supported sideway modules for efficient, expandable continual learning. It provides MLP and GLU variants with principled initialization based on activations, gradients, or cloned backbone parameters.

  • FFNs as soft look-up table memories: Locas reinterprets transformer FFNs as persistent, content-addressable memories composed of key–value slots whose capacity scales with intermediate dimensionality.Keys determine slot activation, values determine retrieved content, and the nonlinearity modulates each slot’s contribution.
  • Locas architecture: Locas adds a sideway FFN memory in parallel with the backbone FFN, scales its output, and adds it to the main pathway for capacity expansion.The sideway design preserves pretrained representations while introducing new test-time memory capacity.
  • Locas-MLP: Locas-MLP uses a two-layer MLP with latent width r and offers clearer theoretical guarantees through step-wise optimal key and value solutions.Its compatibility is limited for models lacking two-layer MLP-style FFN blocks because their MLP-input representations may have poor piecewise-linear separability.
  • Locas-GLU: Locas-GLU matches the GLU-FFN structure of modern LLMs and can attach to GLU-based models for parameter-efficient and computation-efficient continual learning.Its gate, key, and value matrices use the GLU formulation with SiLU activation.
  • Principled initialization: Locas-MLP initializes new memory dimensions from activations and likelihood gradients, while Locas-GLU selects highly activated backbone dimensions and clones their parameters.GLU basis selection ranks intermediate dimensions by average absolute activation before cloning corresponding key and gate rows; proper initialization improves convergence and generalization over random initialization.

3 Experiments

Locas is evaluated for long-context language modeling, memory initialization, width efficiency, preservation of general capabilities, and dialogue question answering. Across these settings, Locas-GLU shows competitive performance with reduced parameter and compute costs, while retaining information with less forgetting.

  • Experimental setup: Locas is evaluated on PG-19 whole-book language modeling and LoCoMo long-context dialogue question answering, emphasizing parameter and compute efficiency.The experiments cover long-context extrapolation and rapid domain adaptation.
  • Whole-book language modeling: Locas achieves comparable or lower perplexity than full attention with fewer parameters and lower computational cost, especially on longer documents.It also improves substantially over context truncation and matches TempLoRA with 25% of its additional parameters and 38% of its computational overhead.
  • Effect of initialization strategy: Top-K activation-guided initialization consistently outperforms alternative Locas-GLU initialization strategies, supporting backbone activation patterns as an effective inductive bias.The comparison includes Bottom-K, random selection, random initialization, and normalized activation initialization.
  • Effect of memory width: At r = 16, Locas-GLU reaches PPL 19.14 with 2.8M parameters, comparable to TempLoRA’s PPL 19.13 at r = 64 with 73.4M parameters.This corresponds to a 26× reduction in parameter count with comparable performance.
  • General capability preservation: Locas-GLU causes 0.2% MMLU degradation after memorizing a complete book, compared with 0.6% for TempLoRA; at r = 512, degradation is 0.1% versus 1.2%.The reported results indicate less catastrophic forgetting for Locas-GLU across memory capacities.
  • Long-context dialogue QA: On LoCoMo, Locas-GLU outperforms full attention and TempLoRA across most question types, including single-hop, multi-hop, temporal, and adversarial questions.For Qwen3-4B-Base temporal questions, it achieves 18.1% F1 versus 17.2% for TempLoRA and 13.9% for full attention; its adversarial F1 is -19.8% versus -25.4% and -24.3%.

4 Related Work

Related work spans test-time training, parameter-efficient fine-tuning, FFN-based memory interpretations, knowledge editing, long-context architectures, and external memory systems. Locas differs by using a parallel FFN memory with principled activation-guided initialization for streaming context memorization.

  • Test-Time Training and Adaptation: Test-time training updates parameters during inference, while TempLoRA uses temporary low-rank adapters trained on generated text chunks.Locas instead emphasizes principled memory initialization and a sideway FFN architecture.
  • Parameter-Efficient Fine-Tuning: Locas shares parameter efficiency goals with LoRA but preserves existing weight matrices through a parallel memory pathway.Its activation-guided initialization leverages the backbone’s internal structure as a warm start.
  • FFN Interpretation and Knowledge Editing: FFN-as-memory interpretations motivate Locas, while WISE provides a related side-memory design for discrete knowledge edits.Locas targets continuous online context memorization and adds activation-guided memory-slot selection.
  • Long-Context Modeling and State Space Models: Locas operates alongside attention or recurrence mechanisms rather than replacing them, allowing potential combination with sparse attention, linear attention, position extrapolation, or state-space methods.This modularity distinguishes Locas from architectural approaches that directly modify long-context computation.
  • Memory-Augmented Neural Networks: Unlike external memory systems such as retrieval-augmented generation and graph-based conversational memory, Locas integrates parametric memory within the model’s forward pass.The paper associates this design with avoiding retrieval latency and enabling rapid memory formation through principled initialization.

5 Conclusion

Locas bridges test-time training with efficient continual learning through sideway FFN-style memories initialized from the backbone model. Its variants combine theoretical guarantees, GLU compatibility, compression, and reduced resource use while limiting catastrophic forgetting.

  • Locas uses backbone-provided initialization for sideway FFN-style memories, enabling parameter-efficient and compute-efficient continual learning.
  • Locas-MLP provides theoretical guarantees and Non-Linear SVD compression, while Locas-GLU integrates with modern GLU-based LLMs through activation-guided parameter cloning.
  • Locas-GLU matches or exceeds TempLoRA on PG-19 and LoCoMo using 15% of the additional parameters and 38% of the computational overhead.
  • MMLU ablations show only 0.1-0.2% degradation, indicating minimal catastrophic forgetting compared with methods that directly modify model weights.
  • Future directions include dynamic allocation, hierarchical memories, and integration with retrieval-augmented generation for hybrid memory systems.

A.1 Non-Linear SVD for Memory Compression of Locas-MLP

Non-Linear SVD offers theoretically principled compression for Locas-MLP, but empirical results favor standard backpropagation because it achieves comparable or better performance at lower computational cost.

  • NL-SVD compresses Locas-MLP while providing theoretical guarantees of functional equivalence within the retained activation subspace.
  • Standard backpropagation achieves comparable or better final performance than NL-SVD with significantly lower computational overhead.

A.1.1 Motivation and Theory

The NL-SVD motivation treats two-layer ReLU perceptrons as having potentially redundant intermediate dimensions and exploits rescaling invariances to identify an effective latent representation.

  • Infinite-dimensional ReLU perceptrons can represent finite-width perceptrons by adding inactive virtual dimensions with zero value vectors.
  • Rescaling a nonlinear dimension’s key and value vectors inversely preserves the represented function for ReLU-like activations.
  • NL-SVD is motivated by compressing the effective latent dimensionality while preserving the FFN’s behavior.
  • The activation pattern depends on the key matrix and the product of the key and value vector norms.

A.1.2 Key Matrix Dimension Reduction

NL-SVD normalizes and rescales Locas-MLP matrices to construct an activation-informed key representation, then retains its dominant singular directions as a compact latent basis.

  • NL-SVD row-normalizes K and column-normalizes V, then forms s_i = α_iβ_i for each intermediate dimension.
  • The coefficients rescale each key dimension so the transformed matrix captures activation direction and effective magnitude.
  • SVD is applied to the rescaled key matrix, and the top n singular directions are retained.
  • Retaining the first n columns of U minimizes rank-n reconstruction error and preserves the original FFN’s dominant activation subspace.
  • The reduced key matrix yields n orthogonal unit-norm probe vectors spanning the FFN’s effective latent subspace.

A.1.3 Value Matrix Reconstruction

The reduced value matrix is reconstructed by querying the original FFN with probe vectors and collecting the resulting output contributions.

  • Each probe vector is fed into the original FFN, and its output contribution before the output projection becomes a reduced value vector.
  • This reconstruction uses the two-layer perceptron’s functional equivalence to preserve the original FFN behavior on retained probes.
  • The resulting vector v_j is assigned as the j-th column of the new value matrix V_n.

A.1.4 Functional Equivalence

The reduced FFN is functionally identical to the original on the span of the retained probe vectors, while inputs outside that subspace incur approximation error governed by discarded singular values.

  • The reduced FFN produces exactly the original output for any probe vector and any input whose activation lies in the retained subspace.
  • For inputs outside the retained subspace, approximation error is governed by the singular values discarded during compression.

A.1.5 Algorithm

The algorithm expands FFN memory from streaming activations and gradients, then compresses it with NL-SVD to a target rank while documenting practical limits in precision, cost, and applicability.

  • Algorithm: NL-SVD takes a target rank n and optionally a retention threshold ϵ, producing reduced key and value matrices.
  • Algorithm: The procedure normalizes key and value factors, weights key rows by composed norms, and selects the top-n singular directions.
  • Algorithm: Each retained reduced key is queried against the original FFN, and the collected outputs are stacked into the reduced value matrix.
  • Expansion-compression cycle: During streaming expansion, each layer creates new key–value entries from normalized hidden activations and globally normalized gradients, increasing FFN dimensionality.
  • Expansion-compression cycle: After N_capacity tokens, NL-SVD reduces the over-complete memory to target rank n, typically N_capacity/2, while preserving dominant activation geometry.
  • Expansion-compression cycle: Compression can occur after every token for continuously refreshed fixed-size memory or after longer accumulated spans.
  • Practical limitations: NL-SVD requires float32 arithmetic, adds substantial overhead, and has theoretical guarantees that do not fully transfer to Locas-GLU.
  • Practical limitations: 17.89 vs. 17.90 final perplexity at 200K context length favors neither NL-SVD nor BP clearly, while NL-SVD requires 12.7× relative time versus 1.7× for BP.
Loading 2602.05085v1…