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

• LLM unlearning removes sensitive, harmful, or copyrighted data without the cost of retraining models from scratch.Source: §1 · p. 1
• Balancing forgetting and retaining objectives is difficult, and learning-rate sensitivity makes existing unlearning methods unstable in practice.Source: §1 · p. 2
• The paper formulates unlearning as two-task optimization: maximize forgetting-set loss while minimizing retaining-set loss.Source: §1 · p. 2
• NGDiff dynamically scalarizes forgetting and retaining objectives using gradient norms and adds an automatic learning-rate adaptation method.Source: §1 · p. 3
• 40% higher model utility is reported for NGDiff on TOFU with Llama2-7B while maintaining comparable unlearning performance.Source: §1 · p. 3

• LLM training data raises concerns involving personal information, outdated knowledge, copyright-protected materials, and adversarially embedded harmful information.Source: §2 · p. 3
• Machine unlearning methods aim to remove unwanted information without retraining entire models, but many classical techniques are difficult to scale to LLMs.Source: §2 · p. 3
• Recent LLM unlearning work mainly designs objectives or hides unwanted information, while the paper targets the underlying optimization problem.Source: §2 · p. 3
• Knowledge editing surgically updates specific knowledge, whereas unlearning removes the influence of particular documents.Source: §2 · p. 3
• Multi-task optimization jointly optimizes distinct objectives by modifying per-task gradients to manage trade-offs during optimization.Source: §2 · p. 4

3 Unlearning as multi-task optimization

• Unlearning generally trades off retaining-set performance against forgetting-set performance because forgetting can remove general knowledge such as grammar rules.Source: §3 · p. 4
• Pareto optimality describes models for which improving either retaining or forgetting loss necessarily worsens the other objective.Source: §3 · p. 4

3.1 Static linear scalarization

• Static scalarization combines task losses with a fixed coefficient c, and its gradient lies in the linear span of the per-task gradients.Source: §3.1 · p. 5
• Extended GDiff unifies GD on retaining data at c = 1, GA on forgetting data at c = 0, and vanilla GDiff at c = 0.5.Source: §3.1 · p. 5
• Pareto optimality alone does not guarantee useful unlearning, since perfect memorization can itself be Pareto optimal.Source: §3.1 · p. 6
• For any convergent coefficient sequence with 0 ≤ c_t ≤ 1, the resulting model is Pareto optimal, including NGDiff.Source: §3.2 · p. 6

4.3 Automatic learning rate adaption

• The adaptive schedule is intended to provide stable convergence, effective unlearning, and high retaining utility without manual learning-rate tuning.The curvature condition used by the theory may fail in deep learning, so the method stabilizes training by leaving η unchanged when the condition fails.Source: §4 · p. 6, §4.2 · p. 8

4.1 Loss landscape of unlearning

• The loss landscapes of retaining and forgetting objectives are approximated locally by quadratic functions, with the approximation fitting most iterations on TOFU.Figure 3 visualizes the retaining and forgetting losses for different learning rates using Phi-1.5 at step 10.Source: §4.1 · p. 7, Figure 3 · p. 7

4.2 Normalized gradient difference

• NGDiff normalizes the retaining and forgetting gradients, ensuring positive correlation with the retain gradient and negative correlation with the forget gradient.This construction satisfies the desired directional condition for balancing the two objectives and is always valid for any g_R and g_F.Source: §4.2 · p. 7
• Theorem 5 shows that, under an appropriate learning rate, NGDiff can improve retaining and forgetting losses simultaneously, while other methods may harm retaining utility.The result is interpreted locally for generally small learning rates, such as η ∼ 10^-4 in the experiments.Source: §4.2 · p. 8

4.3 Automatic learning rate adaption

• NGDiff uses an automatically selected learning rate based on the local minimizer of the retaining loss while monotonically increasing the forgetting loss.GeN estimates the needed scalar quantities from loss differences without instantiating the high-dimensional Hessian, and the rate is updated every 10 iterations to amortize overhead.Source: §4.3 · pp. 8–9

Setup and Main Results

• On TOFU, NGDiff achieves success in most cases under criteria requiring low Verbatim memorization or at least 70% retained Utility and TruthRatio.TOFU uses forget10 as the forgetting set and retain90 as the retaining set.Source: §5.1 · p. 9, Table 1 · p. 10, §5.3 · p. 12

Main Results

• Across TOFU and MUSE-NEWS, NGDiff consistently outperforms baseline methods across models and evaluation metrics.On MUSE-NEWS, it achieves lower Verbmem and higher Utility on Llama2-7B and Mistral-7B, while also removing associated forgetting-set knowledge.Source: §5.3 · pp. 11–13, Table 2 · p. 11
• NGDiff avoids the main baseline trade-offs: GDiff and NPO can reduce memorization while damaging Utility and TruthRatio, whereas LossNorm and high-weight GDiff can fail to unlearn on some models.Baseline effectiveness is highly model- and dataset-dependent, unlike NGDiff.Source: §5.3 · p. 12
• In TOFU training curves, NGDiff is the only compared method maintaining Utility above 75% ROUGE while showing the fastest and most stable convergence and retaining loss ≤0.1.The comparison uses Phi-1.5 with GDiff, LossNorm, and NGDiff.Source: Figure 4 · p. 12, §5.3 · p. 13

Ablation Study

• AutoLR improves NGDiff’s training behavior: ROUGE on forgetting data approaches 0 while retaining ROUGE reaches 0.747, 19.5% above the best static-learning-rate result.With static rates, η = 10^-5 fails to unlearn, while η = 10^-4 reduces retain ROUGE from 100% to around 50%.Source: Figure 5 · p. 13, §5.4 · p. 14

• NGDiff combines normalized gradient difference with automatic learning-rate adaptation for two-task machine unlearning.Source: §6 · p. 14
• The method uses multi-task optimization insights to improve forgetting quality while maintaining model utility.Source: §6 · p. 14
• The paper aims to connect LLM unlearning with multi-task optimization and motivate further research.Source: §6 · p. 14

• Complete removal of specific documents’ influence cannot always be guaranteed, so unlearned LLMs may still generate harmful or undesired outputs.Source: Limitations · p. 14
• Several learning-rate-free and multi-task optimization alternatives were not explored, and scaling them to LLMs could be challenging.Source: Limitations · p. 14
• The study mainly evaluates NGDiff on the TOFU and MUSE benchmark datasets, with only an additional computer-vision example.Source: Limitations · p. 14
• Unlearning performance is highly sensitive to the learning rate, producing substantial changes in ROUGE scores and loss values across settings.Source: Figure 6 · p. 20

NGDiff versus GradNorm

• NGDiff constructs its update as the normalized retaining gradient minus the normalized forgetting gradient, then applies a learning rate.Source: B Comparing NGDiff with GradNorm · pp. 20–21
• Unlike GradNorm, NGDiff is model-agnostic, computes full per-task gradients, and uses fixed gradient-norm coefficients rather than learning task weights.Source: B Comparing NGDiff with GradNorm · p. 21
• NGDiff requires two back-propagations per iteration, whereas GradNorm requires three, potentially increasing GradNorm’s training time for large models.Source: B Comparing NGDiff with GradNorm · p. 21
• With GeN, NGDiff is described as simpler and more suitable for unlearning than GradNorm, with stable performance and theoretical support.Source: B Comparing NGDiff with GradNorm · pp. 21–22

GeN learning-rate adaptation

• GeN selects a learning rate from gradient and Hessian directional quantities, estimated using additional forward passes, and can update it periodically to reduce overhead.Source: §C.1 · p. 22
• Because unlearning maximizes forgetting loss while minimizing retaining loss, the learning rate must be designed around the retaining objective rather than the whole dataset.Source: §C.2 · pp. 22–23

• On CIFAR-10, NGDiff has the best model utility among the compared unlearning methods.Source: Table 4 · p. 23
• The study evaluates NGDiff on image classification using ResNet-50 models trained from scratch on CIFAR-10 and CIFAR-100.Source: D Computer Vision Experiments · p. 23
• For CIFAR-10, 500 dog images form the forgetting set, while images from the other nine classes form the retaining set.Source: D Computer Vision Experiments · p. 23

Machine unlearning

• The paper formulates machine unlearning as a two-task problem solvable with established multi-task optimization methods.This framing separates forgetting from retaining model performance.Source: E Other Related Works · p. 24
• Existing unlearning approaches such as SISA, influence functions [47], and differential privacy [17] can be difficult to scale to large models and datasets.The passage attributes this difficulty to algorithmic complexity.Source: E Other Related Works · p. 24

Multi-task optimization

• Multi-task optimization balances task trade-offs by modifying per-task gradients, including through normalization.Without balancing, a multi-task model can underperform models trained separately for each task.Source: E Other Related Works · p. 24

Learning-rate-free methods

• Learning-rate-free methods aim to automate scheduling because tuning learning-rate hyperparameters is computationally infeasible for LLMs.The paper contrasts this motivation with the requirements of unlearning optimization.Source: E Other Related Works · p. 24
• Parameter-free methods based on convexity or G-Lipschitz assumptions face unlearning-specific problems because the minimizer is unclear in multiobjective optimization.Their loss-based optimal learning rate also does not directly match multi-task methods operating on gradients.Source: E Other Related Works · p. 24

Datasets

• TOFU contains fictitious-author question-answer pairs and uses Full for training, with Forget10 and Retain90 for unlearning experiments.The dataset is generated from biographies of 200 authors, with 20 question-answer pairs per author.Source: §F.1 · pp. 24–25, Table 5 · p. 25
• MUSE-NEWS uses Train for target-model fine-tuning, Raw for unlearning, and VerbmemF, KnowmemF, and KnowmemR for evaluation.The dataset is based on BBC news articles from August 2023.Source: §F.1 · p. 25, Table 5 · p. 25

Evaluation metrics

• Evaluation measures no verbatim memorization, no knowledge memorization, and maintained utility using ROUGE-L recall, with TruthRatio additionally reported for TOFU’s Retain10-perturbed split.TruthRatio compares the likelihood of correct versus incorrect answers, and higher values indicate better retained utility.Source: §F.2 · pp. 25–26

MUSE-NEWS evaluation

• MUSE-NEWS results compare unlearning methods on verbatim memorization, model utility, and TruthRatio for Phi-1.5 and Falcon-1B.Lower Verbmem and higher Utility and TruthRatio indicate better unlearning.Source: Table 6 · p. 26, §F.4 · p. 26

NGDiff training behavior

• NGDiff’s automatic learning rate rises from 5e−5 to around 2e−4 while assigning a smaller rate to the forgetting gradient.The reported effect is reduced perturbation of the model and high utility on the retaining set.Source: §F.5 · p. 27, Figure 7 · p. 27

Pareto optimality

• The scalarized-loss solution cannot be dominated, because a dominating point would contradict its minimality.This establishes Pareto optimality for the solution of the scalarized objective.Source: G Proofs · p. 27
• For convergent weighting sequences, the limiting model produced by the update is Pareto optimal.The proof connects the update’s stationary condition to the scalarized loss solution, which is Pareto optimal.Source: G Proofs · p. 27

NGDiff guarantees

• The NGDiff gradient satisfies the paper’s stated gradient inequality for arbitrary retain and forget gradients.The proof invokes the Cauchy-Schwarz inequality and separately establishes the corresponding forget-gradient relation.Source: G Proofs · pp. 27–28
• For sufficiently small learning rates, the analysis gives conditions under which NGDiff’s retain and forget objectives exhibit the desired local behavior.The result excludes the exactly parallel-gradient case and uses error terms that are o(η) or o(η^2).Source: G Proofs · p. 28

Computational overhead

• GeN incurs computational overhead from additional forward passes, while updating its learning rate every 10 iterations amortizes that overhead.The paper provides a FLOPs-based extension of its complexity analysis.Source: G Proofs · p. 29