We explore catastrophic forgetting in the context of large pre-trained models. By considering forgetting as a geometric problem in the input space of each weight matrix, we uncover a natural retention objective under which updates produced by gradient-based optimizers are suboptimal. Following this observation, we propose Local Support Learning (LSL), a general-purpose framework that augments gradient-based training for retention of prior capabilities without access to prior data. During a new learning phase, LSL pairs two components with distinct roles: a standard weight adapter, trained as usual to minimize the loss, and a gating function that enables the adapter only on input activations from its own training distribution, making the update local to that distribution. The key challenge is that this gate must route data from all learning phases while training only on data from the current one. We address this with a gate based on a Gaussian Mixture Model (GMM), whose likelihood decays rapidly away from its training data, giving it a natural tendency to stay closed on data from prior phases. We show that this post-training approach can resolve forgetting in LLMs of up to 7 billion parameters, retaining both pretrained and finetuned capabilities across multiple training phases, while being efficient in memory and compute, robust to hyperparameter choice, and showing scaling potential.
Figures & tables
Figure 1: Avoiding catastrophic forgetting with Local Support Learning. ( Left ) We augment gradient-based training with a mechanism that retains prior capabilities without access to prior data. The key component is a gating function that enables the weight adapter only on input activations from its training distribution. Interestingly, conventional MLP classifiers are not well suited for this task. Instead, we propose a GMM-based gate that tends to stay closed on prior data it has never encountered. ( Right ) 1D illustration of the GMM-based gate. Φpos (orange) captures the distribution of the finetuning data, while the wider Φneg (blue) dominates elsewhere - ensuring that data outside the current distribution is routed only to the pretrained weights.
Figure 2: Catastrophic forgetting - toy example . ( Data ) We train over six classes in two phases: “Pretraining” (points in blue), and “Finetuning” (points in red). ( Pretrained ) Classifier trained on the pretraining data. The background colors are the learned decision boundaries - depicting a perfect score on the observed distribution. ( Pretrained + FT ) We continually train the model on two other classes. The model performs well on data from the second phase (red), yet the boundaries learned during the first phase are heavily deformed - a clear case of forgetting. ( Pretrained + LSL ) Same experiment, this time with LSL. We restrict weight updates to the region of the input space that produced them, leaving model behavior unchanged elsewhere - thus achieving a perfect score.
Figure 3: Local Support Learning. We describe the training and inference algorithms. Inference is presented per-module (a weight matrix in some layer) and training is presented for the full model. See Alg. 3 (FitGMMs) and Alg. 4 (FitAdapters) for full details. I{⋅} is the indicator function.
Figure 4: Post-training with LSL . We evaluate forgetting on a Qwen2.5-7B-Instruct model for three diverse downstream tasks. The x -axis measures the new task’s performance and the y -axis measures capability retention - the average score across three pretraining benchmarks. The upper-right corner reflects stronger performance. LSL learns the downstream task while achieving near-optimal retention across all settings. This is in contrast to the baselines, which apply new updates to all inputs and thus are prone to interference on prior task data.
Figure 5: Post-Training with multiple phases . Each plot tracks the evaluation score of one task throughout training; the phase whose training data matches the evaluated task is highlighted in bold. The training order is (1) Igbo translation, (2) Chemistry, (3) Cybersecurity. In the leftmost plot, LSL retains pretrained capabilities throughout, while the baseline drops by 44% already after one phase. On the three finetuning benchmarks, LSL retains its gains after each task’s phase ends, whereas the baseline’s performance declines steadily once the peak is reached.
Figure 6: Robustness to hyperparameters . In each plot, we sweep one of the hyperparameters. While conventional fine-tuning trades off learning against forgetting, our method decouples the two, allowing the model to reach its peak capability.
Figure 7: Scaling behavior. We measure the retention of pretraining tasks, defined as the ratio of performance before and after finetuning. LSL exhibits desirable scaling trends.
Figure 8: Efficiency benchmarks on Nvidia B200 with Qwen2.5-7B-Instruct. LSL achieves strong retention with memory, inference, and training overhead on par with established baselines that do not achieve comparable retention.
Figure 9: Gate ablations. ( Left ) Classification accuracy of the MLP and GMM gates at an intermediate layer. The GMM tends to stay closed on data from previous learning phases, a property that enables stronger retention. ( Right ) Retention gains stem from support locality rather than a specific architecture. Replacing the GMM with a different local gate, Union of Spheres (UoS, green), yields equivalent performance, whereas a non-local MLP gate (red) degrades retention significantly.
Figure 10: Single router ablation. We replace LSL’s per-matrix gates with a single gate after the embedding layer, whose decision controls all adapters. The single router retains far more than LoRA but falls short of LSL, especially on the new task. It also degrades retention on cybersecurity. Per-matrix gates decide from contextual features at every depth, allowing full learning and retention.
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 11: Ablating the Number of GMM Components and Temporal Smoothing . We measure performance via overall model performance (left column) and average hit rate across gates (right column). An optimal gate achieves 100% hit rate on the new task (blue) and 0% on pretrain data (yellow, red). Note that each plot sweeps the number of GMM components. We find that LSL is largely insensitive to the number of GMM components. In addition, the GMM accounts for the vast majority of retention, while temporal smoothing adds the final increment, mainly by further reducing the hit rate on pretraining tasks.
k
512
128
64
16
4
1
ρk
1.70%
0.44%
0.27%
0.12%
0.06%
0.02%
Appendix
Table 1: Evaluating the effectiveness of OP-LoRA. We report the ρk metric from the OP-LoRA paper, which measures how much energy of a standard LoRA adapter falls within the subspace of the top- k singular values of the pretrained weights (trained on the chemistry dataset). We find that forgetting is only weakly related to the preservation of the top- k singular-value subspace, explaining why OP-LoRA performs similarly to the LoRA baseline.
Hyperparameter
Value
Optimizer (AdamW)
Weight decay
0.1
β1,β2
0.9, 0.95
Gradient clip norm
1.0
Schedule
Warmup ratio
0.1
Appendix
Table 2: Shared hyperparameters across all tasks.
Chemistry
Igbo Translation
Cybersecurity
LoRA Adapter
Learning rate
1e−4
1e−4
1e−4
Batch size
32
32
16
Adapter rank r
128
128
512
Epochs
3
3
1
GMM Gate
Appendix
Table 3: Per-task hyperparameters.
Figure 12: Function Behavior Change - In-Distribution Data. For each token of the finetuning (chemistry) data, we measure the KL divergence between the outputs of the finetuned and base models, as a function of the fraction of LSL gates open on that token. Most tokens open nearly all gates, where LSL and LoRA change the model’s behavior similarly, as desired.
Figure 13: Function Behavior Change - Out-of-Distribution Data. Same analysis on pretraining benchmarks. Here most tokens open few gates, and LSL changes the model’s behavior less than LoRA, sometimes by an order of magnitude. The LoRA baseline changes the model’s behavior more on these out-of-distribution tokens, degrading pretraining capabilities.
Figure 14: ROC Curves for UoS Radius Calibration . We show two representative curves, one for the gate projection (left) and one for the down projection (right), both in layer 17. Each curve color indicates different compaction rates (dictated by the maximal value of representative points M). The selected radius is the one with the maximal compaction rate, subject to being within the allowed TPR,FPR region (defined globally for all gates).
Models trained on a new task typically degrade on prior tasks, a phenomenon known as forgetting. Traditionally, mitigating forgetting has required replaying stored exemplars from prior tasks, which is often impractical. By contrast, language models can sample from their own training distribution, and we show that these self-generated samples serve as effective replay data, nearly eliminating forgetting. We find that forgetting nonetheless persists when the model has little remaining capacity: models pretrained close to saturation cannot absorb new information without overwriting prior knowledge. When capacity is not the limiting factor, low learning rates reduce forgetting but require substantially more training steps. Replay breaks this tradeoff, enabling fast, high-learning-rate finetuning without forgetting.
LLM post-training typically propagates task gradients through the full depth of the model. Although this end-to-end structure is simple and general, it couples task adaptation to full-depth activation storage, long-range backward dependencies and direct task-gradient access to pretrained representations. We argue that this full-depth backward coupling can be unnecessarily expensive and intrusive, particularly when post-training supervision is much narrower than pre-training. To this end, we propose \textbf{LoPT}: Local-Learning Post-Training, a simple post-training strategy that makes gradient reach an explicit design choice. LoPT places a single gradient boundary at the transformer midpoint: the second-half block learns from the task objective, while the first-half block is updated by a lightweight feature-reconstruction objective to preserve useful representations and maintain interface compatibility. LoPT shortens the task-induced backward path while limiting direct interference from narrow task gradients on early-layer representations. Extensive experiments demonstrate that LoPT achieves competitive performance with lower memory cost, higher training efficiency and better retention of pretrained capabilities. Our code is available at: https://github.com/HumyuShi/LoPT
Hengyu Shi, Tianyang Han, Peizhe Wang +3
1Independent Researcher · 2D4 Lab · 3Southeast University
Large language models (LLMs) often suffer from catastrophic forgetting in continual learning: after learning new tasks sequentially, they perform worse on earlier tasks. Existing methods mitigate catastrophic forgetting by data replay, parameter freezing, or regularization. However, these methods lack understanding of LLM mechanisms and cannot distinguish which parameters store important knowledge from previous tasks and which parameters can be updated for new tasks. To address this, we propose the attribution-guided continual fine-tuning framework that leverages Layer-wise Relevance Propagation (LRP) to estimate parameter importance based on the internal computational process of LLMs. During continual learning, parameters critical to previous tasks are constrained to receive smaller updates, while less relevant parameters remain available for learning new tasks. Extensive experiments show that, compared with baseline methods, our approach reduces catastrophic forgetting while preserving adaptability to new tasks, highlighting the value of mechanistic attribution for continual fine-tuning of LLMs.
Yazheng Liu, Yuxuan Wan, Rui Xu +3
The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · The Beijing University of Posts and Telecommunications, Beijing, China