Unlearning as Multi-Task Optimization: a normalized gradient difference approach with adaptive learning rate

Authors: Xiaomeng Jin, Zhiqi Bu, Bhanukiran Vinzamuri, Anil Ramakrishna, Kai-Wei Chang, Mingyi Hong
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: machine unlearning, multi-task optimization, learning rate scheduler, large language models

Abstract:
Unlearning techniques have been proposed as a cost-effective post-training way to remove undesired knowledge learned by large language models (LLMs). However, existing methods often fail to effectively unlearn the targeted information or cause a significant drop in model performance. In this paper, we frame machine unlearning as a multi-task optimization problem to balance this tradeoff -- one task maximizes forgetting loss, while the other minimizes retaining loss. We introduce a novel unlearning method, Normalized Gradient Difference (NGDiff), which guarantees Pareto optimality upon convergence. Specifically, NGDiff dynamically normalizes task gradients, enabling the model to unlearn targeted forgetting data while preserving utility on the retaining set. We also identified that unlearning methods are sensitive to learning rate and integrate an automatic learning rate scheduler that selects the locally optimal learning rate to stabilize and accelerate the convergence. Experiments with various LLMs demonstrate that NGDiff outperforms state-of-the-art unlearning methods on the TOFU and MUSE datasets.

Decision: reject

Reviews and discussion:

Metareview: