Unlearning as multi-task optimization: A normalized gradient difference approach with an adaptive learning rate (2025)

Abstract: Machine unlearning has been used to remove unwanted knowledge acquired by large language models (LLMs). In this paper, we examine machine unlearning from an optimization perspective, framing it as a regularized multi-task optimization problem, where one task optimizes a forgetting objective and another optimizes the model performance. In particular, we introduce a normalized gradient difference algorithm, enabling us to have better control over the trade-off between the objectives, while integrating a new, automatic learning rate scheduler. We provide a theoretical analysis and empirically demonstrate the superior performance of among state-of-the-art unlearning methods on the TOFU and MUSE datasets while exhibiting stable training.

## introduction large language models (llms) consume a large amount of data during pre-training. after the model is built, we may have to unlearn certain data points that contain potentially sensitive, harmful, or copyrighted content. as re-training from scratch in such a case is not feasible due to the associated costs, researchers have developed a number of machine unlearning methods applied after training. existing machine unlearning methods are formulated primarily as minimizing memorization through the language model loss @xcite @xcite . in particular, the gradient ascent (ga) method maximizes the language model (lm) loss (i.e., minimizes the negative lm loss) on the target forget set (f). however, this approach can also negatively affect the utility of the model. to mitigate the utility loss, the gradient difference (gdiff) method selects a subset of the training data as the retain set (r), minimizing the sum of the negative lm loss on the forgetting set and the standard lm loss on the retaining set. this approach has been empirically shown to effectively preserve the model's performance @xcite . similarly, negative preference optimization (npo) @xcite assigns a lower likelihood of forgetting data, thereby balancing the unlearning performance with model utility. despite these successes, there are still two key issues preventing the methods from reaching their full potential. first, balancing retaining and forgetting losses is difficult (figure 1 , details are in appendix a) given the disproportionate sizes of the forget and retain datasets. second, the optimization methods for unlearning are usually sensitive to the learning rate (cf., appendix a, figure 7 ). for instance, various learning rates can lead to substantial changes in the rouge scores and loss values even for the same algorithm, making the unlearning methods unstable and difficult to use in practice. in this paper, we carefully examine unlearning from an optimization perspective and formulate it as a multi-task optimization (mto) problem @xcite : we aim to minimize the lm loss (i.e., maximize the utility) on the retaining set and maximize the lm loss on the forgetting set (i.e., minimize memorization), simultaneously. 1 to solve this two-task problem, we study the rich literature of multi-task methods that seeks the pareto optimality of two tasks (e.g., imtl @xcite , gradnorm @xcite , rlw @xcite , pcgrad @xcite , and scalarization @xcite ), and design an approach specifically for the llm unlearning problem. inspired by the simplicity and strong empirical performance of linear scalarization methods, 2 which minimize a linearly weighted average of task losses, we propose an llm unlearning method, ngdiff, based on dynamic scalarization, and analyze its theoretical properties. building on the analysis, we introduce an automatic learning rate adaptation method tailored for llm unlearning. we showcase the effectiveness of our method through extensive experiments on multiple datasets, different llms and vision models. for example, on tofu @xcite , ngdiff achieves 40% higher model utility while maintaining comparable unlearning performance with llama2-7b. figure 1 highlights the effectiveness of ngdiff. contributions are summarized as follows: ## unlearning with normalized gradient difference while theorem 3 shows the pareto optimality of θ * as @xmath0, it does not shed insight on the convergence through intermediate steps θ t . put differently, although many mto and unlearning methods are all pareto optimal upon convergence, they may converge to different pareto points at different convergence speeds. therefore, it is important to understand and control the algorithm dynamics to maintain high performance for r throughout the training. specifically, the dynamics are determined by the choices of g @xmath1. ( 4 ). in this section, we propose to use gradient normalization for g un and automatic learning rate for η t , so as to achieve stable convergence, effective unlearning, high retaining utility, without manually tuning the learning rate. ## conclusion and discussion we formulated the machine unlearning problem as a two-task optimization problem and proposed a novel unlearning method ngdiff based on normalized gradient difference and automatic learning rate adaption. by leveraging insights from multi-task optimization, ngdiff empirically improves forgetting quality while maintaining utility. we hope this paper helps establish a connection between llm unlearning and multi-task optimization, and inspires further advancements in this field.