Unified Gradient-Based Machine Unlearning with Remain Geometry Enhancement

Shanghai Jiao Tong University
NeurIPS 2024 Spotlight
MY ALT TEXT

Overview of our Saliency Forgetting in the Remain-preserving manifold online (SFR-on).

Abstract

Machine unlearning (MU) has emerged to enhance the privacy and trustworthiness of deep neural networks. Approximate MU is a practical method for large-scale models. Our investigation into approximate MU starts with identifying the steepest descent direction, minimizing the output Kullback-Leibler divergence to exact MU inside a parameters' neighborhood. This probed direction decomposes into three components: weighted forgetting gradient ascent, fine-tuning retaining gradient descent, and a weight saliency matrix. Such decomposition derived from Euclidean metric encompasses most existing gradient-based MU methods. Nevertheless, adhering to Euclidean space may result in sub-optimal iterative trajectories due to the overlooked geometric structure of the output probability space. We suggest embedding the unlearning update into a manifold rendered by the remaining geometry, incorporating second-order Hessian from the remaining data. It helps prevent effective unlearning from interfering with the retained performance. However, computing the second-order Hessian for large-scale models is intractable. To efficiently leverage the benefits of Hessian modulation, we propose a fast-slow parameter update strategy to implicitly approximate the up-to-date salient unlearning direction. Free from specific modal constraints, our approach is adaptable across computer vision unlearning tasks, including classification and generation. Extensive experiments validate our efficacy and efficiency. Notably, our method successfully performs class-forgetting on ImageNet using DiT and forgets a class on CIFAR-10 using DDPM in just 50 steps, compared to thousands of steps required by previous methods.

Contribution

  • We provide a novel perspective to unify previous approaches by decomposing the vanilla gradient descent direction of approximate MU into three components: weighted forgetting gradient ascent,remaining gradient descent, and a weight saliency matrix.
  • We derive the steepest descent direction for approximate MU on the remain-preserved manifold.
  • We propose a fast-slow weight method to implicitly approximate online Hessian-modulated salient forgetting updates.
  • We conduct experiments on a wide range of CV unlearning tasks across multiple datasets and models of different architectures, verifying the effectiveness and efficiency of our method.

Results on random forgetting

Random subset forgetting

Performance summary of MU methods for image classification (including RT, six baselines, our proposed SFR-on, and ablations on our designed components), assessing unlearning 10% random subset of CIFAR-10 using ResNet-18 and TinyImageNet using Swin-T.

Results on class-forgetting in image generation tasks

BibTeX

@article{huang2024unified,
  title={Unified Gradient-Based Machine Unlearning with Remain Geometry Enhancement},
  author={Huang, Zhehao and Cheng, Xinwen and Zheng, JingHao and Wang, Haoran and He, Zhengbao and Li, Tao and Huang, Xiaolin},
  journal={arXiv preprint arXiv:2409.19732},
  year={2024}
}