Table of Contents
Fetching ...

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

TL;DR

The paper tackles catastrophic forgetting during domain-specific fine-tuning of large language models by introducing Hierarchical Importance Regularization (HLoRA). It records general-knowledge parameter importance during pretraining via element-wise path-integral-inspired scores and imposes a layer-aware regularization, balanced by a dual-objective loss that also optimizes for domain performance. This approach avoids expensive Fisher-based computations, enabling faster importance estimation and reduced storage while preserving general capabilities and improving domain adaptation. Empirical results on GPT-J and LLaMA-3 across SciQ, PiQA, and MedMCQA show strong forgetting mitigation, competitive or superior domain accuracy, and substantial runtime and memory efficiency improvements over prior methods.

Abstract

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.

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

TL;DR

The paper tackles catastrophic forgetting during domain-specific fine-tuning of large language models by introducing Hierarchical Importance Regularization (HLoRA). It records general-knowledge parameter importance during pretraining via element-wise path-integral-inspired scores and imposes a layer-aware regularization, balanced by a dual-objective loss that also optimizes for domain performance. This approach avoids expensive Fisher-based computations, enabling faster importance estimation and reduced storage while preserving general capabilities and improving domain adaptation. Empirical results on GPT-J and LLaMA-3 across SciQ, PiQA, and MedMCQA show strong forgetting mitigation, competitive or superior domain accuracy, and substantial runtime and memory efficiency improvements over prior methods.

Abstract

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.
Paper Structure (19 sections, 16 equations, 7 figures, 2 tables)

This paper contains 19 sections, 16 equations, 7 figures, 2 tables.

Figures (7)

  • Figure 1: Illustration of catastrophic forgetting: the fine-tuned LLM fails to answer previously known questions.
  • Figure 2: Adaptive constraint combining element-wise and layer-wise importance to preserve general capabilities from the $\nu$ task while learning domain-specific abilities for the $\mu$ task. RECORD means the general importance recording in \ref{['sec:recording']}. REG means the regularization in \ref{['sec:element_reg']} and \ref{['sec:layer_reg']}.
  • Figure 3: Independent samples t-test of EWCLoRA and our method on LLaMA-3: violin plots of perplexity (PPL) and accuracy (Acc) across datasets
  • Figure 4: Loss curves on three datasets: balancing task learning and generalization. The total loss consists of task loss ($\mathcal{L}_{\text{task}}$) and a scaled version of general loss ($\mathcal{L}_{\text{reg}}$), where task loss controls the model learning on new domain data, and general loss helps maintain the model generalization ability.
  • Figure 5: Comparison of computation time and storage for importance $\mathbf{\Omega}_l^\nu$ between previous method and ours.
  • ...and 2 more figures