PaperlayerView exact revision on arXiv

Source-linked AI summary

Memory Bank Compression for Continual Adaptation of Large Language Models

Thomas Katraouras, Dimitrios Rafailidis

arXiv:2601.00756v1cs.LGcs.CL

TL;DR

LLMs become outdated, while continual-learning approaches face computational costs, catastrophic forgetting, and memory banks that grow with streaming data. MBC compresses memory with codebook optimization and online resetting, and uses attention-only KV-LoRA for adaptation. Across QA experiments, it reduces memory-bank size to 0.3% relative to baseline size while maintaining high retention accuracy.

  • Problem

    LLMs need continual updates as new information arrives, but full fine-tuning is costly and vulnerable to catastrophic forgetting, while memory banks grow with incoming documents.

  • Method

    MBC compresses memory using codebook indices, resets the codebook online to prevent collapse, and applies KV-LoRA only in attention layers.

  • Results

    0.3% memory-bank size relative to baseline size, with improved QA accuracy and high retention accuracy during online adaptation.

  • Takeaways & Limitations

    MBC provides memory-efficient continual adaptation for streaming LLM settings without full fine-tuning.

  • Takeaways & Limitations

    CaMeLS results are unavailable for LLaMA-2 (7B) because reproducing that baseline exceeded the capacity of one NVIDIA A100 80GB GPU.

Abstract

from arXiv · show

Large Language Models (LLMs) have become a mainstay for many everyday applications. However, as data evolve their knowledge quickly becomes outdated. Continual learning aims to update LLMs with new information without erasing previously acquired knowledge. Although methods such as full fine-tuning can incorporate new data, they are computationally expensive and prone to catastrophic forgetting, where prior knowledge is overwritten. Memory-augmented approaches address this by equipping LLMs with a memory bank, that is an external memory module which stores information for future use. However, these methods face a critical limitation, in particular, the memory bank constantly grows in the real-world scenario when large-scale data streams arrive. In this paper, we propose MBC, a model that compresses the memory bank through a codebook optimization strategy during online adaptation learning. To ensure stable learning, we also introduce an online resetting mechanism that prevents codebook collapse. In addition, we employ Key-Value Low-Rank Adaptation in the attention layers of the LLM, enabling efficient utilization of the compressed memory representations. Experiments with benchmark question-answering datasets demonstrate that MBC reduces the memory bank size to 0.3% when compared against the most competitive baseline, while maintaining high retention accuracy during online adaptation learning. Our code is publicly available at https://github.com/Thomkat/MBC.

1 Introduction

LLMs become outdated because their parameters are static, motivating continual adaptation without repeatedly optimizing the full model. MBC addresses memory-bank growth by compressing stored representations while retaining QA performance and adaptation stability.

  • LLMs cannot easily incorporate new knowledge after training, creating a knowledge-cutoff problem as information changes.
  • Full fine-tuning is computationally expensive for large LLMs and can cause catastrophic forgetting, while PEFT methods still require deployment-time gradient updates.
  • Memory augmentation avoids repeated gradient updates by storing structured representations, but its memory bank grows as more documents arrive.
  • MBC compresses the memory bank with codebook optimization and online resetting, storing codebook indices instead of full document representations.
  • MBC combines attention-only Key-Value LoRA with compressed memory representations to adapt to new data without full fine-tuning.

2 Online Adaptation of LLMs

The online-adaptation formulation models a pretrained, outdated LLM receiving a stream of documents while evaluation queries remain inaccessible during updates. An auxiliary labeled training set guides query-agnostic parameter updates intended to retain past knowledge and improve future queries.

  • Online adaptation continuously updates a pretrained, outdated language model using a stream of new documents.
  • Evaluation queries and labels associated with incoming documents remain inaccessible during adaptation, so updates must be query-agnostic.
  • An auxiliary training set of documents, queries, and labels guides parameter updates while retaining past knowledge and improving future queries.

3 Proposed Model

MBC compresses document-derived memory with a learned codebook, aggregates the compressed representations for query-specific modulation, and adds lightweight KV-LoRA adaptation. Online adaptation stores only code indices and uses forward passes, while codebook resetting and quantization losses support stable, compact learning.

  • Architecture: MBC encodes documents, stores their representations in a memory bank, and aggregates the bank with query representations to produce modulation for the base LLM.Its three components are an amortization network, memory bank, and aggregation network.
  • Architecture: Hierarchical aggregation partitions the context set into subgroups and recursively combines them, reducing memory complexity to O(MT).The procedure is designed to remain scalable as the number of stored documents increases.
  • Memory Bank Compression: Vector quantization maps each continuous context vector to its nearest codebook entry and stores only its integer index in the compressed memory bank.During inference, indices retrieve quantized vectors that the aggregation network combines with the query representation.
  • Memory Bank Compression: EMA-based resetting reinitializes low-usage codes with sampled encoder outputs and mean usage estimates to maintain codebook diversity and prevent collapse.Dead codes are identified below a usage threshold and reset during training without gradients.
  • Modulation Adaptation: KV-LoRA adds low-rank trainable updates to key and value projections in the final transformer layers while retaining the pretrained base parameters.The updates are scaled and regularized with dropout to balance adaptation capacity and computational efficiency.
  • End-to-End Training: The model jointly optimizes the QA and weighted vector-quantization losses while keeping the base model frozen.Trainable components include the amortization, input, aggregation, KV-LoRA, and codebook parameters.
  • Online Adaptation: After training, online adaptation requires no gradient updates: new documents are quantized into stored codes, and queries retrieve and aggregate their corresponding vectors.The procedure operates entirely through forward passes.

4.1 Datasets

The experiments use benchmark question-answering datasets with document streams and question-answer pairs, while the supplied passages also include end-to-end optimization algorithm descriptions. Dataset passages identify StreamingQA, SQuAD, and ArchivalQA and provide several train, validation, and test splits.

  • Datasets: The evaluation uses three question-answering datasets, as stated in the experimental setup.The supplied dataset passages name StreamingQA, SQuAD, and ArchivalQA.
  • StreamingQA: StreamingQA contains timestamped English WMT news articles from 2007–2020 with questions created by annotators or generated with language models.The described setup uses 21K training, 1.7K validation, and 5K test questions and documents.
  • SQuAD: SQuAD contains crowdsourced questions on Wikipedia whose answers are spans within the corresponding articles.The described setup uses 39.9K training, 5.6K validation, and 10.6K test questions.
  • ArchivalQA: ArchivalQA is built from New York Times Annotated Corpus articles and uses language models to generate questions with answers as text spans.The supplied passage introduces the dataset but does not provide its complete split sizes.
  • Optimization Procedures: The optimization algorithms encode documents, quantize their contexts, update code usage, reset dead codes, aggregate contexts with queries, and combine QA and VQ losses.The supplied algorithm passages describe these operations but do not report dataset-level results.

4.2 Evaluation Protocol

MBC is evaluated after adapting on a sampled test-document stream and then answering questions drawn from those same documents. Performance is reported with Exact Match and token-level F1 after standard answer normalization.

  • Protocol: Each dataset uses 1,665 documents sampled from the test stream for adaptation, followed by evaluation on QA pairs drawn from those documents.The configuration follows prior work for comparison across baselines.
  • Metrics: Exact Match measures the fraction of predictions that exactly match normalized ground-truth answers.Normalization lowercases text, removes punctuation and articles, and collapses repeated spaces.
  • Metrics: Token-level F1 measures the harmonic mean of token-level precision and recall.The metric operates on tokenized answer representations.

4.3 Experimental Setup

The experiments compare MBC and continual-adaptation baselines across four LLM backbones, using T5-based auxiliary networks and attention-focused LoRA configurations. Training uses fixed optimization settings with validation-based evaluation and single-GPU execution.

  • Backbones: MBC is evaluated with DistilGPT2, GPT2-Large, GPT2-XL, and LLaMA-2-7B backbones ranging from 82M to 7B parameters.T5-based amortization and input encoders are scaled across the backbone configurations.
  • Training: Training runs for 50 epochs with Adam, linear warmup over the first 1% of steps, learning rate 10^-5, validation after each epoch, and gradient accumulation.Batch size is 64 for DistilGPT2 and 32 for the other models.
  • LoRA Configuration: For GPT-2 models, the LoRA down-projection matrix is shared across key and value projections, and all experiments use one NVIDIA A100 80GB GPU.The larger backbones use rank 32, scaling factor 64, dropout 0.05, and the final 16 layers.
  • Baselines: Baselines include Uniform Fine-Tuning, Salient Spans, and CaMeLS, which apply increasingly selective or learned weighting to tokens during fine-tuning.The baseline configurations follow prior work, and pretrained LLMs are first fine-tuned on QA pairs before online adaptation.

4.4 Experimental Results

MBC improves QA performance while substantially compressing memory banks across datasets and model scales. Its codebook resetting mechanism maintains diverse code usage during training, supporting stable adaptation.

  • QA performance: MBC consistently improves both EM and F1 across datasets and base LLMs compared with the MAC baseline.Table 1 evaluates StreamingQA, SQuAD, and ArchivalQA across different backbones.
  • QA performance: 11.84% in EM and 12.99% in F1 are MBC’s gains compared with MAC.The reported gains are attributed to KV-LoRA and an efficiently learned codebook.
  • Memory footprint: 98.27% to 99.2% are the average memory-bank reductions achieved by MBC relative to MAC across model scales.The reductions average 98.27% for DistilGPT2, 99.1% for GPT2-Large and GPT2-XL, and 99.2% for LLaMA-2-7B.
  • Parameter overhead: Less than 0.5% is the additional trainable-parameter overhead from MBC’s codebook and KV-LoRA across all base LLMs.The comparison is reported in the offline setting, excluding documents stored during online adaptation.
  • Codebook resetting: Codebook resetting preserves higher and more stable perplexity than no resetting across all four base LLMs on StreamingQA.Without resetting, perplexity collapses toward 12–24, whereas resetting maintains substantially higher ranges, including above 90 for GPT2-XL and above 100 for LLaMA-2-7B.

5 Conclusion

MBC addresses the growing memory-bank problem in memory-augmented LLMs by combining codebook compression, online resetting, and attention-only KV-LoRA for continual adaptation. On QA datasets, it reduces memory use while improving accuracy and retaining performance during online adaptation.

  • 0.3% memory-bank footprint relative to the most competitive baseline accompanies improved EM and F1 scores.MBC also maintains high F1 retention during online adaptation.
  • Codebook-based compression and online resetting prevent codebook collapse and maintain balanced code utilization.The resetting mechanism supports stable use of compressed memory representations.
  • Attention-only KV-LoRA enables MBC to exploit query-memory modulations without full fine-tuning.The modules provide targeted adaptation within the attention mechanism.
  • MBC provides scalable continual adaptation for LLMs in streaming settings where new documents continually expand the memory bank.The conclusion identifies memory efficiency and reduced catastrophic forgetting as outcomes of the design.
Loading 2601.00756v1…