PaperlayerView exact revision on arXiv

Source-linked AI summary

How to Alleviate Catastrophic Forgetting in LLMs Finetuning? Hierarchical Layer-Wise and Element-Wise Regularization

Shezheng Song, Hao Xu, Jun Ma, Shasha Li, Long Peng, Qian Wan, Xiaodong Liu, Jie Yu

arXiv:2501.13669v2cs.CLcs.AI

TL;DR

Domain-specific fine-tuning can cause LLMs to lose pretrained knowledge, limiting their broader applicability. The paper combines parameter-importance regularization with cross-entropy learning and layer-wise weighting. Across scientific, medical, and physical tasks, it reports reduced forgetting and improved adaptability, alongside faster, lower-storage importance estimation.

  • Problem

    Domain-specific fine-tuning can overwrite pretrained knowledge, limiting LLMs’ broader applicability.

    Source: §1 · p. 1

  • Method

    The method combines element-wise importance-based regularization, cross-entropy for domain learning, and layer-wise coefficients to balance the losses.

    Source: §1 · p. 2, §4.3 · p. 5

  • Results

    Across scientific, medical, and physical tasks, the method mitigates forgetting while enhancing adaptability; importance estimation is 20x faster and requires 10%∼15% of prior methods’ storage.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The reported speed and storage reductions support the practicality of parameter-importance estimation for fine-tuning LLMs.

    Source: §1 · p. 2

    The parameter-update analysis assumes a gradient exists that completes the LoRA update process.

    Source: §4.1 · p. 4

Abstract

from arXiv · show

Large Language Models (LLMs) exhibit strong general language capabilities. However, fine-tuning these models on domain-specific tasks often leads to catastrophic forgetting, where the model overwrites or loses essential knowledge acquired during pretraining. This phenomenon significantly limits the broader applicability of LLMs. To address this challenge, we propose a novel approach to compute the element-wise importance of model parameters crucial for preserving general knowledge during fine-tuning. Our method utilizes a dual-objective optimization strategy: (1) regularization loss based on element-wise parameter importance, which constrains the updates to parameters crucial for general knowledge; (2) cross-entropy loss to adapt to domain-specific tasks. Additionally, we introduce layer-wise coefficients to account for the varying contributions of different layers, dynamically balancing the dual-objective optimization. Extensive experiments on scientific, medical, and physical tasks using GPT-J and LLaMA-3 demonstrate that our approach mitigates catastrophic forgetting while enhancing model adaptability. Compared to previous methods, our solution is approximately 20 times faster and requires only 10-15% of the storage, highlighting the practical efficiency. The code will be released.

1. Introduction

Domain-specific fine-tuning can overwrite pretrained knowledge, threatening both general capabilities and performance on specialized tasks. The paper proposes element-wise and layer-wise importance to balance preserving general knowledge with learning domain-specific expertise.

1. Introduction

  • Catastrophic forgetting occurs when fine-tuning loses or overwrites pretrained knowledge, potentially impairing reasoning, domain-task performance, and answers to previously known general questions.The paper identifies domain data’s fixed composition and format as a possible contributor.Source: §1 · p. 1
  • Element-wise importance constrains updates to parameters important for general capabilities, combining a regularization loss with cross-entropy loss for domain-specific learning.The importance is calculated from parameter-update path integrals, avoiding Fisher-matrix computation and storage.Source: §1 · p. 2
  • Layer-wise coefficients dynamically adjust regularization because layers contribute differently to generalization and domain-specific learning.The coefficients use the L2 norm of each layer’s computed element-wise importance.Source: §1 · p. 2
  • The framework achieves state-of-the-art performance across scientific, physical, and medical tasks while mitigating forgetting; importance estimation is 20x faster and uses 10%∼15% of prior methods’ storage.Experiments use GPT-J and LLaMA-3.Source: §1 · p. 2

2. Related Work

The related work reviews continual-learning strategies and methods for reducing forgetting during LLM fine-tuning, including LoRA-based approaches with distinct computational and storage trade-offs.

2.1. Continual Learning

  • Regularization-based continual learning includes SI (Zenke et al., 2017) and EWC (Kirkpatrick et al., 2016), alongside architecture-, replay-, and distillation-based approaches.SI estimates parameter importance online; EWC uses the Fisher Information Matrix to constrain updates.Source: §2.1 · pp. 2–3

2.2. Catastrophic Forgetting in LLM and LoRA

  • Full parameter fine-tuning can cause forgetting, while LoRA updates a small parameter subset and mitigates—but does not eliminate—it.EWCLoRA (Xiang et al., 2024) combines EWC with LoRA but incurs Fisher-matrix overhead; I-LoRA (Ren et al.) requires an additional set of LoRA parameters.Source: §2.2 · p. 3

3. Preliminary

LoRA adapts pretrained models through low-rank updates, optimizing only added matrices while keeping the pretrained weights frozen. Despite its efficiency, task-specific fine-tuning can impair general ability.

3. Preliminary

  • LoRA fine-tuning updates only low-rank matrices B and A while keeping pretrained weights θ0 frozen.The adapted parameter is θt = θ0 + δt, where δt = BtAt and the rank r is much smaller than the weight dimension.Source: §3 · p. 3
  • LoRA’s optimization objective uses a task-specific loss function.The task-specific loss is generally computed using cross-entropy.Source: §3 · p. 3
  • Task-specific fine-tuning with LoRA can cause catastrophic forgetting that harms general ability.Source: §3 · p. 3

4. Hierarchical Importance Regularization

The framework records parameter importance during a general task and uses it to constrain later domain-task updates. It combines element-wise importance with layer-wise weighting to balance preserving general capabilities and adapting to the new task.

4.1. General Element-Wise Importance Recording

  • During ν-task training, LoRA parameter updates are treated as a trajectory, with a hypothetical gradient assumed to represent the low-rank updates.The importance calculation adapts parameter and gradient updates to LoRA’s low-rank matrices.Source: §4.1 · pp. 3–4
  • Element-wise importance scores measure each parameter’s contribution to general-task loss reduction, using the negative gradient path integral and accounting for update magnitude.Higher scores identify parameters considered more important for maintaining ν-task performance.Source: §4.1 · pp. 3–5

4.2. Element-Wise Regularization in Domain Tuning

  • During µ-task tuning, the regularization loss keeps high-importance parameters near their ν-task reference values while allowing less-important parameters more flexibility.The µ-task objective combines task-specific and regularization losses; cumulative importance is updated at the end of ν and the running contribution values are reset.Source: §4.2 · p. 5

4.3. Layer-Wise Coefficient Regularization

  • Layer-wise importance adjusts regularization across layers, while φ controls the balance between domain-task and general-ability losses.The layer importance uses the L2 norm of the parameter-importance matrix for each layer.Source: §4.3 · p. 5

5. Experiments

Experiments use GPT-J and LLaMA-3, comparing the proposed approach with tuning and regularization baselines. General ability is evaluated on Pile, while domain adaptation is tested on medical, scientific, and physics datasets.

5.1. Backbone LLMs and Baseline Methods

  • GPT-J and LLaMA-3 serve as evaluation backbones, with their differing architectures and parameter counts intended to provide a robust evaluation.The compared approaches include an untuned base, LoRA on domain data or general-then-domain data, EWCLoRA, and rsLoRA.Source: §5.1 · p. 5

5.2. Tasks, Metrics, and Hyperparameters

  • General ability is evaluated on Pile, where LoRA training on the ν task also records parameter importance for Synaptic Intelligence.Source: §5.2 · p. 5
  • Domain adaptation and forgetting are evaluated on MedMCQA, SciQ, and PiQA for medical, scientific, and physics tasks.Source: §5.2 · p. 5
  • Table 1 reports domain ability as accuracy (Acc↑) and general ability as perplexity (PPL↓).Source: Table 1 · p. 6
  • The experiments use GPT-J-6B and LLaMA 3.2-3B, training both ν and µ tasks for five epochs.Batch size is 20, learning rate is 8e-4, LoRA rank is 8, and LoRA alpha is 32.Source: §5.2 · p. 6

6. Results and Analysis

HLoRA preserves general ability while maintaining domain accuracy, and the analyses report lower computation and storage costs than EWCLoRA. The coefficient, parameter-importance, and ablation analyses examine how regularization balances these capabilities and where importance lies.

6.1. Comparison of General and Domain Capabilities

  • Table 1 reports the lowest PPL for general ability while domain-specific accuracy is comparable to, or better than, previous methods.Source: §6.1 · p. 6
  • Perplexity differences on SciQ, PiQA, and MedMCQA, and accuracy differences on SciQ and PiQA, were statistically significant versus EWCLoRA (p < 0.05).Source: §6.1 · p. 6
  • Across three datasets, task loss decreases while general loss first rises and then falls as parameters move away from their prior-task values.The total loss combines task loss and general loss to constrain learning against both objectives.Source: §6.1 · p. 6

6.2. Complexity Comparision

  • GPT-J-6B: EWCLoRA takes 27.17 hours versus HLoRA's 1.15 hours; LLaMA-3-3B: EWCLoRA takes 25.97 hours versus HLoRA's 1.19 hours.Source: Figure 5 · p. 7, §6.2 · p. 7
  • GPT-J-6B: EWCLoRA's Fisher matrix uses 22.65 GB versus HLoRA's 3.5 GB; LLaMA-3-3B: 11.97 GB versus 1.3 GB.Source: Figure 5 · p. 7, §6.2 · p. 7

6.3. Regularization Coefficient Analysis

  • As regularization coefficient φ increases, PPL decreases but average PiQA accuracy is negatively affected; e−3 is selected to balance them.Source: §6.3 · p. 7, Figure 6 · p. 8

6.4. Parameters Importance Visualization

  • The log-scaled importance heatmap shows v_proj layers, especially the first four and last layer, as more important for general knowledge than q_proj layers.Source: Figure 7 · p. 8, §6.4 · p. 8

6.5. Ablation Study

  • Removing layer-wise and element-wise importance lowers performance across three datasets for both backbone LLMs.Source: Table 2 · p. 8, §6.5 · p. 8

7. Conclusion

The conclusion presents a fine-tuning framework intended to mitigate catastrophic forgetting while retaining general capabilities and adapting to new domains. It reports experiments across scientific, physical, and medical tasks that maintain domain-task performance.

  • The framework targets catastrophic forgetting during domain-specific fine-tuning while preserving general capabilities and adapting to new domains.Source: §7 · p. 8
  • A layer-wise coefficient dynamically adjusts the balance between regularization loss and cross-entropy loss to reflect layers' differing contributions.Source: §7 · p. 8
  • Experiments in scientific, physical, and medical tasks report mitigation of catastrophic forgetting while maintaining domain-specific task performance.Source: §7 · p. 8
Loading 2501.13669v2…