Technical Summary: SAE-FD for Continual Learning of Large Language Models

Problem Statement

Large Language Models (LLMs) face the challenge of catastrophic forgetting when sequentially adapting to multiple tasks. Existing continual learning (CL) methods typically operate in dense representation spaces (weight, gradient, or output spaces). A fundamental limitation of these approaches is feature superposition, where multiple concepts are encoded in overlapping dimensions within the dense activation space. This entanglement makes it difficult to selectively protect previously learned knowledge without inadvertently constraining the learning of new tasks, leading to a suboptimal plasticity-stability tradeoff.

Methodology: SAE-FD

The authors propose SAE-FD (Sparse Autoencoder Feature Distillation), a method that anchors model representations in the sparse feature space of a pre-trained Sparse Autoencoder (SAE). By decomposing dense activations into a sparse, over-complete basis, SAE-FD aims to disentangle superimposed representations, allowing for more targeted regularization.

The method operates as a three-stage pipeline:

1. SAE Training (Stage 1)

Before the continual learning sequence begins, a Gated SAE is trained on diverse text activations from the base model.

• Input: Last-layer MLP activations ($h \in \mathbb{R}^d$) from diverse sources (classification, QA, knowledge benchmarks).
• Architecture: The Gated SAE separates encoding into a gating pathway ($g$) and a magnitude pathway ($m$), producing pre-ReLU features ($f_{pre}$) and post-ReLU features ($f$).
• Objective: Minimize reconstruction loss ($\|\hat{h} - h\|_2^2$) and an L1 sparsity penalty ($\lambda_{\ell1}\|f\|_1$).
• Outcome: A frozen SAE that decomposes dense activations into sparse features ($D \gg d$), effectively disentangling concepts.

2. Task Training and Anchor Capture (Stage 2)

For each task $t$ in the sequence:

• Fine-tuning: The model is fine-tuned using LoRA with the standard language modeling objective.
• Anchor Capture: After training on a task, the model performs a forward pass on a small set of anchor samples ($N_a$). The per-token MLP outputs ($H_i$) are extracted and stored.
• Buffer: These anchors accumulate across tasks. During subsequent training, a batch is sampled uniformly from the buffer to compute distillation loss, amortizing the cost across all previous tasks.

3. Sparse-Aware Feature Distillation (Stage 3)

During the training of a new task ($t+1$), the model is regularized by comparing its current representations against the stored anchors, both encoded through the frozen SAE. The distillation loss ($L_{FD}$) consists of two complementary terms:

• Cosine Direction Loss ($L_{cos}$): Penalizes angular deviation between the current pre-ReLU feature vector and the anchor feature vector. This preserves the global direction of the representation.
• Active Magnitude Loss ($L_{mag}$): Applies Mean Squared Error (MSE) only on features that were active (post-ReLU $> 0$) in the anchor. This constrains the magnitude of features contributing to previous tasks while leaving inactive features free to be recruited for new tasks.

Adaptive $\lambda$ Mechanism

To balance task learning and regularization, the authors introduce an adaptive mechanism for the weighting coefficient $\lambda$ in the total loss $L = L_{task} + \lambda \cdot L_{FD}$.

• Instead of a fixed $\lambda$, the system targets a fixed contribution ratio $\rho$ of the distillation loss to the total loss.
• $\lambda$ is updated via exponential moving average (EMA) to satisfy $\rho = \frac{\lambda L_{FD}}{\lambda L_{FD} + L_{task}}$.
• Behavior: $\lambda$ starts high when a new task begins (when $L_{task}$ is large and forgetting risk is high) and relaxes as the model converges, preventing excessive constraint during fine-tuning.

Key Contributions

1. SAE-FD Framework: The first continual learning method to perform regularization in the sparse, decoupled feature space of SAEs, addressing the feature superposition issue inherent in dense spaces.
2. Sparse-Aware Distillation Loss: A novel loss function combining cosine direction preservation with selective magnitude matching on active features, leveraging the sparsity of SAE representations to minimize interference with new tasks.
3. Adaptive Regularization: An adaptive $\lambda$ mechanism that dynamically adjusts regularization strength based on the relative magnitude of task and distillation losses, ensuring strong protection during high-risk periods without hindering convergence.

Experimental Results

The authors evaluated SAE-FD on two benchmarks across three model architectures (LLaMA-2-7B, Vicuna-7B, Mistral-7B).

• TRACE Benchmark (8 Heterogeneous Tasks):
  • On LLaMA-2-7B, SAE-FD achieved 52.70% Average Accuracy (AA) with only −0.46% Backward Transfer (BWT).
  • This outperforms the strongest published regularization-based baseline (GORP) by +2.3% AA and reduces forgetting by 34% (−0.46 vs −0.70 BWT).
  • It significantly outperforms other LoRA-based methods (EWC, O-LoRA, TreeLoRA), which showed BWT worse than −3%.
• Generalization:
  • SAE-FD consistently achieved the best AA and BWT across all three model families.
  • On Vicuna-7B, it surpassed O-LoRA by +10.1% AA and achieved the best BWT (−0.69%) among all methods, including full fine-tuning baselines.
• Ablation Studies:
  • Loss Components: Both cosine direction and active magnitude losses are necessary; using either alone resulted in lower AA or higher forgetting.
  • Adaptive $\lambda$: Replacing the adaptive mechanism with a fixed $\lambda=1.0$ degraded BWT from −0.46% to −2.38%.
  • SAE vs. Raw MSE: Distilling in the SAE feature space outperformed raw MSE distillation on dense activations, confirming the benefit of the decoupled space.
• Efficiency: The computational overhead is approximately 1.85× that of SeqLoRA, with a constant cost per step regardless of the number of previous tasks. Storage requires ~0.8 GB per task for anchor activations.

Significance and Claims

The paper claims that SAE-FD successfully connects mechanistic interpretability (via SAEs) with continual learning. By utilizing the sparse, disentangled representations originally developed for understanding neural networks, the method provides an effective inductive bias for controlling them.

The authors assert that operating in a sparse feature space allows for more targeted preservation of knowledge, solving the "feature superposition" problem that plagues dense-space regularization. The method achieves a superior plasticity-stability tradeoff, enabling LLMs to adapt to evolving tasks without retraining from scratch while maintaining high performance on previously learned tasks.

Limitations noted by the authors:

• Requires a one-time training cost for a model-specific SAE.
• Storage requirements for anchor activations grow linearly with the number of tasks (though compressible).
• Experiments were conducted on a fixed task order for the TRACE benchmark, though multi-model consistency suggests robustness.
• Comparisons with some baselines must be interpreted with caution due to differences in experimental settings (e.g., full fine-tuning vs. LoRA, epoch counts).