Unlearning-Aware Minimization
Hoki Kim, Keonwoo Kim, Sungwon Chae, Sangwon Yoon
Machine unlearning aims to remove the influence of specific training data (i.e., forget data) from a trained model while maintaining performance on the remaining data (i.e., retain data). Existing approximate unlearning methods typically rely on fine-tuning or negative gradient updates. However, these methods often degrade performance on the retain data or fail to eliminate the influence of forget data.In this paper, we propose Unlearning-Aware Minimization (UAM), a novel min-max optimization framework that identifies weights with high loss on the forget data and uses their gradients to minimize the loss on the retain data. We derive an efficient optimization method for this min-max problem, enabling the successful removal of the forget data and uncovering optimal solutions that existing methods struggle to find.Our empirical results demonstrate that the proposed framework consistently outperforms existing methods in machine unlearning across various benchmark datasets, including image classification datasets (CIFAR-10, CIFAR-100, and TinyImageNet) and multiple-choice question-answering datasets using large language models (WMDP-Bio and WMDP-Cyber).