Continual Learning via Sparse Memory Finetuning

TL;DR

The paper tackles catastrophic forgetting in large language models by introducing sparse memory finetuning. It leverages memory layers that are sparsely updated by design and uses TF-IDF-based ranking to update only the most relevant memory slots, preserving pretraining knowledge. On fact learning and document QA tasks, sparse memory finetuning achieves comparable new knowledge acquisition to full finetuning and LoRA but with substantially reduced forgetting (e.g., 11% drop vs 89% and 71%). This indicates that sparse updates to memory layers can enable continual learning at scale, offering a practical path toward models that continually improve without catastrophic degradation.

Abstract

Modern language models are powerful, but typically static after deployment. A major obstacle to building models that continually learn over time is catastrophic forgetting, where updating on new data erases previously acquired capabilities. Motivated by the intuition that mitigating forgetting is challenging because trainable parameters are shared across all tasks, we investigate whether sparse parameter updates can enable learning without catastrophic forgetting. We introduce sparse memory finetuning, leveraging memory layer models (Berges et al., 2024), which are sparsely updated by design. By updating only the memory slots that are highly activated by a new piece of knowledge relative to usage on pretraining data, we reduce interference between new knowledge and the model's existing capabilities. We evaluate learning and forgetting compared to full finetuning and parameter-efficient finetuning with LoRA on two question answering tasks. We find that sparse memory finetuning learns new knowledge while exhibiting substantially less forgetting: while NaturalQuestions F1 drops by 89% after full finetuning on new facts and 71% with LoRA, sparse memory finetuning yields only an 11% drop with the same level of new knowledge acquisition. Our results suggest sparsity in memory layers offers a promising path toward continual learning in large language models.

Figures (7)

• Figure 1: Memory layer architecture berges2024memorylayersscalehe2024mixturemillionexperts. We replace one FFN in the middle of the transformer with a memory lookup. For each token, we project the output of the previous layer into a query, which is used to identify the top $k=32$ keys. An input-dependent gating is applied to the weighted sum of the top $k$ values, which then becomes the output of the memory layer.
• Figure 2: Sparse memory finetuning. To identify the minimal set of memory slots to update, we aggregate counts for all the memory indices on a particular batch, and then rank the accesses relative to a background corpus (e.g. pretraining data) with TFIDF. We train the top $t$ memory indices that are highly accessed on this batch but accessed relatively infrequently on pretraining, keeping the rest of the memory pool and model frozen.
• Figure 3: Learning single facts in the small-data regime. To simulate a "small data" regime where the model must perform immediate updates on a small amount of data, we train the model to learn a sequence of 1000 facts from TriviaQA. Sparse memory finetuning learns more on the target facts, while forgetting much less on held-out benchmarks (NaturalQuestions and GSM8K). LoRA and full finetuning exhibits catastrophic forgetting on the held-out metrics.
• Figure 4: Learning from documents. We evaluate whether methods can learn from a stream of documents by training on Active Reading-augmented lin2025learning documents for a subset of 100 Wikipedia-grounded SimpleQA questions wei2022emergentabilitieslargelanguage. The best full finetuning and LoRA configurations achieve the same performance on the target task, but suffer from much more forgetting on held-out tasks. Sparse memory finetuning achieves Pareto improvements at both high and low learning rates and generally exhibits much less degradation.
• Figure 5: Tradeoff between learning and forgetting for different methods. To fully characterize the tradeoff for different methods (full finetuning, LoRA, and sparse memory finetuning), we plot performance on the target (TriviaQA) and a held-out task (HellaSwag) across a hyperparameter sweep. For full finetuning, we sweep across learning rate {$2\text{e-}6, 5\text{e-}6, 2\text{e-}5, 5\text{e-}5\}$. For LoRA, we sweep across rank $\{32, 128, 256\}$, alpha $\{1/2, 1, 2, 4\}$x rank, and learning rate $\{2\text{e-}4, 5\text{e-}5, 5\text{e-}6\}$. For sparse memory finetuning, we sweep across the number of trainable indices $t$$\{25, 50, 100, 200, 500, 1000\}$ and learning rate $\{0.1, 2\}$. We omit runs where learning fails (e.g. for full finetuning, learning rates above $5\text{e-}6$ "break" the model and result in lower performance on both learning and forgetting than initialization). The size of each point depicts the number of trainable parameters per batch. We observe a continuous trend in most methods: as the learning capacity increases (e.g. with increased learning rate, or parameter count), forgetting also increases, up to a point where further increasing learning capacity leads to reduced performance. Sparse memory finetuning exhibits high learning capacity with minimal forgetting on the held-out benchmark.