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

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 LLMs. We argue that such underperformance can be partly explained by the ``polysemantic'' nature of LLMs: 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.

• The paper proposes shifting from weight-space to SAE-guided activation-space regularization to selectively safeguard task-specific features against catastrophic forgetting.
• It employs a two-stage pipeline that constructs task-relevant feature masks via k-NN propagation and applies squared-hinge loss for balancing stability and plasticity.
• Empirical results on TRACE and MedCL benchmarks demonstrate significant performance gains over traditional methods, ensuring efficient continual learning in LLMs.

SAE-Guided Activation Regularization for LLM Continual Learning

Motivation and Problem Statement

Catastrophic forgetting remains a critical challenge in the continual adaptation of LLMs, particularly as they are deployed across evolving domains and task formats. Classical weight-space regularization methods, typified by Elastic Weight Consolidation (EWC), Synaptic Intelligence (SI), and Memory Aware Synapses (MAS), suffer from structural limitations—most notably, the polysemanticity of parameters in modern transformer architectures. Due to superposition, individual weights and neurons are entangled across multiple, often unrelated, semantic concepts, fundamentally weakening parameter-level importance estimates. This entanglement undermines the selectivity required for task-specific knowledge preservation, leading to an inescapable stability–plasticity tradeoff: retention is achieved at the cost of plasticity and vice versa.

Approach: Regularization in Feature Space Using Sparse Autoencoders

This paper proposes regularization in activation space, leveraging pretrained Sparse Autoencoders (SAEs) as monosemantic feature dictionaries. The core insight is that SAE features exhibit substantially higher interpretability and task-alignment than raw parameter units, allowing for concept-level protection via feature masks derived from task data.

The proposed method follows a two-stage pipeline:

• Stage 1: Mask Construction: Task data is encoded through a fixed SAE and the frozen base model. The resulting feature activations yield a relevance profile, which is expanded through k-NN propagation in the decoder vector space to generate a continuous mask $\mathbf{M}$ indicating task-relevant features.

  Figure 1: Overview of the SAE-guided regularization pipeline. Stage 1: offline mask construction using task data and the frozen SAE; Stage 2: per-step drift regularization over masked SAE features during training.

• Stage 2: Optimizer Step Regularization: During training, both the current and reference models are mapped into SAE feature space, and per-feature drift is regularized via a squared-hinge loss. Features with low relevance scores are protected (stability constraint), while those with high relevance are encouraged to adapt (plasticity constraint).

The training objective is derived from a rigorous constrained optimization formulation. Stability and plasticity constraints are enforced through a squared-hinge relaxation, producing two regularization terms: the protect loss (on low-relevance features) and the guide loss (on high-relevance features). EWC appears as a special case within this framework, corresponding to a one-sided weight-space penalty without explicit plasticity.

Empirical Results

TRACE and MedCL Benchmarks

Experiments are conducted on TRACE (cross-domain, 8 tasks) and MedCL (biomedical, 10 tasks), using Gemma-2 9B-it with LoRA adapters. SAE-guided regularization substantially outperforms all regularization-based baselines, including EWC, SI, MAS, and gradient-projection methods (ELLA). It also surpasses experience replay under strict storage constraints (no retained previous task data).

• TRACE: SAE achieves OP $=$ 0.545 (95% CI [0.529, 0.562]), with nearly full retention and high plasticity, representing a $+19.4$ point absolute OP improvement over unprotected fine-tuning.
• MedCL: SAE leads among methods with no replay buffer, OP $=$ 0.510 (95% CI [0.495, 0.524]), while larger replay buffers further improve performance at greater memory cost.

  Figure 2: Stability–plasticity tradeoff on TRACE. SAE-guided regularization uniquely attains both high retention and high plasticity, closely approaching the architectural upper bound.

EWC and SI optimize for retention but suffer from catastrophic intransigence (plasticity <0.46); MAS maintains plasticity but at the cost of severe forgetting. SAE-guided regularization breaks this tradeoff by targetting selective feature-space drift.

Mechanistic Analyses

Superposition and Selective Protection

Empirical evidence supports the polysemanticity thesis. When comparing the separability of task-relevant units:

• Linear Separability: Task-relevant SAE features are visually and quantitatively separable (AUC up to 0.89 in deep layers), while task-relevant neurons are indistinguishable from chance (AUC $\approx$ 0.50), demonstrating the selectivity advantage of the SAE basis.

  Figure 3: t-SNE projections reveal that task-relevant SAE features cluster distinctly by task and layer, emphasizing feature-basis separability.

• Collateral Constraint: Protecting weights for one task constrains features relevant to another task at 91–96% of the rate for its own, indicating non-selective protection. In contrast, SAE feature-space protection reduces collateral constraint to 43–61%.

This evidence correlates with EWC's observed failure mode: its rigidity derives not from hyperparameter miscalibration but from non-selective anchoring inherent to weight-space under superposition.

Efficiency and Scalability

SAE-guided regularization offers a scalability advantage. Traditional weight-space regularization requires per-task anchors and importance vectors proportional to the number of trainable parameters, rapidly exceeding practical device and storage capacity as model scale and LoRA rank rise. SAE-guided regularization requires only a compact feature mask per task (sub-megabyte scale) and a fixed SAE dictionary, shared across all tasks. Training speed is comparable to unprotected fine-tuning, with negligible inference overhead.

Hyperparameter Robustness and Mask Construction Variants

Hyperparameter sweeps confirm that performance is primarily sensitive to protection budget and penalty strength, affirming the mechanistic role of stability constraints. Plasticity guidance parameters exert a secondary effect. Various importance signals (activation magnitude, causal gradient, empirical Fisher) yield similar performance, implying that the principal benefit arises from the SAE feature basis itself.

Implications and Future Directions

Practical implications include:

• Selective Knowledge Retention: Activation-space regularization enables task-specific preservation without storing prior data or allocating task-specific parameters.
• Efficient Deployment: The method is well-suited for adaptation scenarios where memory footprint is critical.
• Theoretical Advancement: This approach clarifies the limits of parameter-space protection and sets a new standard for regularization-based continual learning in LLMs.

Future work will include extending validation to other model families, exploring synergistic integration with replay or architectural isolation, and increasing SAE capture for broader activation variance.

Conclusion

Regularizing LLMs in SAE feature space addresses a structural failure of weight-space methods aggravated by polysemanticity and superposition. Derived via constrained optimization with explicit stability and plasticity terms, the method achieves efficient, selective continual learning, outperforming classical baselines in both aggregate and mechanistic evaluations. The results suggest that concept-aligned feature-space protection is vital for scaling continual learning in modern LLMs.