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

TL;DR

This paper introduces a hierarchical layer-wise and element-wise regularization method to mitigate catastrophic forgetting in LLM fine-tuning, improving knowledge retention and task adaptation efficiently.

Contribution

It proposes a novel element-wise importance computation and dynamic layer-wise coefficients for effective regularization during fine-tuning.

Findings

Reduces catastrophic forgetting in LLMs by 20 times faster than previous methods.

Requires only 10-15% of storage compared to existing solutions.

Enhances model adaptability across scientific, medical, and physical tasks.

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.