From Weights to Features: SAE-Guided Activation Regularization for LLM Continual Learning

2026-06-25

AI Summary

This work addresses catastrophic forgetting in large language models during continual learning, a challenge exacerbated by the lack of semantic selectivity in weight-space regularization methods such as Elastic Weight Consolidation (EWC). The authors propose the first activation-space regularization approach, leveraging a pretrained sparse autoencoder (SAE) to construct a disentangled semantic feature dictionary. Compact feature masks are generated from current-task data to balance stability and plasticity without storing historical examples. Theoretically, EWC is shown to be a special case of the proposed method—specifically, a one-sided penalty in weight space. Empirical results demonstrate that this approach significantly outperforms conventional weight-regularization techniques on the TRACE and MedCL benchmarks, achieving state-of-the-art performance without task-specific architectures while offering superior memory efficiency.

Abstract

Weight-space regularization methods such as Elastic Weight Consolidation (EWC) are the standard approach to catastrophic forgetting in continual learning. However, those methods tend to underperform when applied to large language models. We argue that such underperformance can be partly explained by the ``polysemantic'' nature of large language models: per-weight importance estimates utilized by EWC-style regularization are too coarse and cannot isolate the knowledge that needs protection. In this paper, we propose regularizing instead in the model's activation space, using pretrained Sparse Autoencoders (SAEs) as a monosemantic feature dictionary. From the perspective of constrained optimization, we derive a new loss function that uses the SAE feature dictionary to explicitly balance stability and plasticity, and show that EWC is a special case in the one-sided weight-space penalty setting. Unlike replay-based methods that store or revisit examples from earlier tasks, our method requires no previous-task data after mask construction: current-task data is used to compute a compact SAE feature mask, and only this mask is retained for later training. Further, since the feature space has significantly lower dimensionality than the parameter space, the proposed method is more memory efficient. On the TRACE and MedCL continual learning benchmarks, the method achieves the strongest result among approaches without introducing task-specific architectural components, also surpassing traditional weight-space regularization methods like EWC. Beyond performance comparisons, we provide empirical evidence for the polysemanticity thesis: task-relevant representations are linearly separable in the SAE feature basis but indistinguishable from chance in the weight basis, and weight-space protection is nearly non-selective at the concept level.

Problem

• continual learning
• catastrophic forgetting
• large language models
• polysemanticity
• activation regularization

Innovation

• Sparse Autoencoders
• Activation Regularization
• Continual Learning
• Polysemanticity
• Feature Masking