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

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

Problem Statement
Continual learning (CL) in Large Language Models (LLMs) faces the challenge of catastrophic forgetting, where sequential fine-tuning overwrites representations needed for prior tasks. While regularization-based methods like Elastic Weight Consolidation (EWC) are standard for smaller models, they underperform significantly in LLMs. The authors argue this failure is structural, stemming from the "polysemantic" nature of LLMs. Due to the superposition hypothesis, individual neurons and their associated weights encode multiple unrelated concepts simultaneously. Consequently, weight-space importance estimates (e.g., Fisher information) are too coarse; protecting a weight to preserve Concept A inadvertently constrains Concept B, creating a non-selective penalty that stifles plasticity. Furthermore, existing methods often fail to explicitly enforce the plasticity required to learn new tasks, leading to "intransigence" where the model retains old knowledge but cannot adapt.

Methodology
The paper proposes SAE-guided activation regularization, a method that shifts the regularization target from the entangled parameter space to a more selective activation space using pre-trained Sparse Autoencoders (SAEs).

1. Coordinate System Shift: Instead of penalizing weight drift, the method uses a frozen, pre-trained SAE as a fixed, monosemantic feature dictionary. SAEs decompose dense model activations into sparse feature directions that are often more interpretable and less entangled than individual neurons.
2. Constrained Optimization Formulation: The training objective is derived from a constrained optimization problem rather than ad-hoc penalty design. The model must minimize the current task loss subject to two constraints:
   • Stability: Task-irrelevant features (low relevance) must not drift beyond a specific budget ($\varepsilon_s$).
   • Plasticity: Task-relevant features (high relevance) must drift sufficiently ($\ge \varepsilon_p$) to avoid degenerate solutions where the model learns nothing.
3. Loss Function: Using Lagrangian relaxation and a squared-hinge penalty, the authors derive a training loss comprising a protect loss (penalizing excessive drift in protected features) and a guide loss (penalizing insufficient drift in relevant features). EWC is shown to be a special case of this framework where the representation space is weights and the plasticity constraint is absent.
4. Mask Construction: For each new task, a continuous relevance mask ($M_t$) is constructed using current-task data passed through the frozen base model and SAE. This mask identifies relevant features via activation magnitude and propagates relevance to neighboring features using k-NN on decoder vectors. Crucially, once the mask is constructed, all task data is discarded; only the compact mask is retained.
5. Training: During training, the current model (with LoRA) and the frozen base model are both encoded by the SAE. The feature drift is weighted by the mask to compute the protect and guide losses.

Key Contributions

• SAE-Guided Regularization: A practical CL method that regularizes in activation space using SAE features. It requires no previous-task data, no stored activations, no task-specific parameters, and no inference-time routing after mask construction.
• Principled Derivation: The stability-plasticity tradeoff is formulated as a constrained optimization problem in SAE feature space, yielding a squared-hinge training loss via Lagrangian relaxation.
• Empirical Evidence for Polysemanticity: The paper provides mechanistic evidence that weight-space protection is non-selective. Experiments show task-relevant representations are linearly separable in SAE feature space (AUC ~0.88) but indistinguishable from chance in weight/neuron space (AUC ~0.50). Additionally, protecting weights for one task constrains features for a subsequent task at 91–96% of the rate of the original task, whereas feature-space protection reduces this collateral constraint to 43–61%.
• Scalability: Due to the lower dimensionality of feature space compared to parameter space, per-task storage is reduced from gigabytes (weight anchors) to sub-megabytes (feature masks).

Results
The method was evaluated on TRACE-5000 (8 cross-domain tasks) and MedCL (10 biomedical tasks) using Gemma-2 9B-it with LoRA.

• Performance: On TRACE, SAE-guided regularization achieved the strongest performance among non-architectural approaches (OP = 0.545), outperforming weight-space regularizers (EWC: 0.447, SI: 0.513), gradient-projection methods, and replay baselines (even with 10% data). On MedCL, it led all methods that retain no previous-task examples (OP = 0.510).
• Stability-Plasticity Tradeoff: Unlike EWC, which achieved high retention (ARA = 0.987) by severely sacrificing plasticity (0.453), the proposed method maintained high plasticity (0.567, comparable to unprotected fine-tuning) while retaining 95% of learned performance.
• Efficiency: The method adds minimal computational overhead (~1.88s/step vs 1.70s baseline) and drastically reduces storage requirements (412 KB per task vs 6.5 GB for EWC anchors).

Significance and Claims
The paper claims that the underperformance of weight-space regularization in LLMs is not a hyperparameter tuning issue but a structural consequence of polysemanticity. By moving the regularization constraint to a monosemantic feature basis, the method achieves a more selective protection of knowledge. The authors assert that the effectiveness of regularization depends critically on the representation space in which constraints are imposed. The method demonstrates that it is possible to achieve strong continual learning performance without storing past data or modifying the model architecture at inference time, provided the regularization operates on a coordinate system aligned with the model's internal concept structure. The paper concludes that SAE-based regularization offers a more selective and scalable alternative to traditional weight-space methods for LLMs.