Multi-Objective Large Language Model Unlearning

TL;DR

Large language models raise safety and copyright concerns, making full retraining impractical. The paper reframes LLM unlearning as a multi-objective optimization over $L_{fgt}$, $L_{KL}$, and $L_{rt}$, replacing the unlearning loss with Unlearning Cross-Entropy $L_{UCE}$ to prevent gradient explosion. It then derives a common descent direction using the Dual Space MGDA (DS-MGDA) to achieve Pareto stationarity, balancing forgetting of targeted data with retention of downstream utility. Empirical results on LLama3-8B with PKU-SafeRLHF show MOLLM outperforms state-of-the-art GA-based methods in unlearning effectiveness while better preserving model utility, with the code publicly available.

Abstract

Machine unlearning in the domain of large language models (LLMs) has attracted great attention recently, which aims to effectively eliminate undesirable behaviors from LLMs without full retraining from scratch. In this paper, we explore the Gradient Ascent (GA) approach in LLM unlearning, which is a proactive way to decrease the prediction probability of the model on the target data in order to remove their influence. We analyze two challenges that render the process impractical: gradient explosion and catastrophic forgetting. To address these issues, we propose Multi-Objective Large Language Model Unlearning (MOLLM) algorithm. We first formulate LLM unlearning as a multi-objective optimization problem, in which the cross-entropy loss is modified to the unlearning version to overcome the gradient explosion issue. A common descent update direction is then calculated, which enables the model to forget the target data while preserving the utility of the LLM. Our empirical results verify that MoLLM outperforms the SOTA GA-based LLM unlearning methods in terms of unlearning effect and model utility preservation. The source code is available at https://github.com/zibinpan/MOLLM.

Figures (2)

• Figure 1: A demo of the gradient conflict during LLM unlearning, where $g_{fgt}$, $g_{KL}$, $g_{rt}$ denote the gradients of $\mathcal{L}_{fgt}$, $\mathcal{L}_{KL}$, and $\mathcal{L}_{rt}$, respectively. $\mathcal{L}_{fgt}^*$, $\mathcal{L}_{KL}^*$, and $\mathcal{L}_{rt}^*$ are the corresponding dual vectors. $\bar{d}^t$ represents the update direction of the weighted-sum method, which takes the weighted-sum of $-g_{fgt}$, $-g_{KL}$, and $-g_{rt}$. The hull in the pink color denotes the dual space of span($-g_{fgt}$, $-g_{KL}$, $-g_{rt}$), where all the vectors inside are the common descent directions. The blue direction $d^t_{MOLLM}$ is the common descent direction obtained by the proposed MOLLM.
• Figure 2: When using CE loss and GA for unlearning, $p_{i,c}$ is driven to 0, which leads to gradient explosion (blue line). As for UCE, it adopts the gradient descent to drive $p_{i,c}$ to 0 to unlearn, preventing gradient explosion.