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

How can we mitigate catastrophic forgetting in Large Language Models (LLMs) during continual learning without relying on memory-intensive replay buffers or task-specific architectural components? The authors investigate whether the failure of traditional weight-space regularization methods, like Elastic Weight Consolidation (EWC), stems from the polysemantic nature of LLM parameters.

The authors propose a novel regularization framework that operates in the activation space rather than the parameter space. By utilizing pretrained Sparse Autoencoders (SAEs), they decompose dense model activations into a set of monosemantic features. For each new task, the model computes a task-specific relevance mask that identifies which SAE features are critical for the current task. They then derive a training objective from a constrained optimization problem, penalizing drift in protected (low-relevance) features while ensuring sufficient adaptation in task-relevant (high-relevance) features. This approach requires no storage of previous-task data, as only the compact feature mask is retained.

The study demonstrates that task-relevant representations are significantly more separable in SAE feature space than in weight space, confirming that weight-space protection is inherently non-selective. On benchmarks like TRACE and MedCL, the SAE-guided method outperforms traditional weight-space regularizers and other non-architectural approaches. Furthermore, the method is highly memory-efficient, as it replaces gigabyte-scale anchor storage with sub-megabyte feature masks.

This research provides a mechanistic explanation for why standard continual learning techniques struggle with LLMs. By shifting the focus from entangled weights to interpretable, monosemantic SAE features, the authors provide a more principled and scalable way to manage the stability-plasticity tradeoff, enabling LLMs to learn new tasks sequentially without overwriting previously acquired knowledge.