Efficient Utility-Preserving Machine Unlearning with Implicit Gradient Surgery

TL;DR

This work tackles the fundamental utility-unlearning tradeoff in machine unlearning by formulating the Problem as Utility-Preserving Unlearning (UPUP), a constrained optimization that maximizes unlearning while bounding utility loss. It establishes that solving UPUP is equivalent to unilateral gradient surgery on the unlearning gradient, and proposes Efficient Utility-Preserving Machine Unlearning (EUPMU) with an implicit gradient-surgery mechanism to avoid extra backpropagations. Theoretical analysis proves convergence to Pareto-optimal or stationary solutions under standard smoothness and Lipschitz assumptions, and empirical results on classification and diffusion-model generation show that EUPMU achieves better forgetting-utility tradeoffs and faster runtime than baselines. The approach offers a practical, provably sound, and computationally efficient path for responsible MU across diverse models and tasks, including image generation and NSFW content removal.

Abstract

Machine unlearning (MU) aims to efficiently remove sensitive or harmful memory from a pre-trained model. The key challenge is to balance the potential tradeoff between unlearning efficacy and utility preservation, which involves forgetting undesirable information as defined while maintaining the model's original performance. One potential way to tackle this problem is to use multi-objective optimization to jointly optimize both the unlearning and utility preservation objectives. However, existing multi-objective methods only guarantee finding a Pareto-optimal solution without fine-grained control, which causes under-optimization of the unlearning objective. To this end, we first model MU as a constrained optimization problem, that is, optimizing the unlearning objective under the constraint of a bounded increase for utility loss. We then show that solving this optimization problem is equivalent to unilateral gradient surgery on the unlearning objective. To resolve the additional computational cost brought by gradient surgery, we propose an implicit gradient surgery method, which approximates the solution to the aforementioned constrained optimization problem via only one backpropagation, thereby achieving efficient utility-preserving MU. Theoretically, we provide a tight convergence analysis of the algorithm. Empirically, our extensive experiments show that the proposed algorithm achieves better tradeoff results than existing baselines. Codes are available at https://github.com/anseryuer/EUPMU-Efficient-Utility-Preserving-Machine-Unlearning.

Figures (7)

• Figure 1: A showcase where both linear scalarization and MGDA fail to find a good tradeoff solution that balance the utility and unlearning objectives: (a) Linearization often leads to extreme solutions with deterioration on one objective; (b) The pre-trained model fully optimizes the utility objective, leaving little room to improve unlearning objective for MGDA algorithm that fairly optimizes all objectives; (c) Our EUPMU algorithm provides a tolerance for small degradation for utility objective, and returns a better tradeoff solution.
• Figure 2: Illustration of Unilateral Gradient Surgery. The computation of the update direction is divided into two scenarios: (a) If the gradient of the unlearning objective conflicts with the gradient of the retaining objective, that is, the angle between them is greater than 90 degrees, the desired direction is obtained by removing the projection (to the retaining gradient) component from the unlearning gradient; (b) If the gradient of the unlearning objective does not conflict with the gradient of the retaining objective, then the unlearning gradient itself is the desired direction.
• Figure 3: Illustration of the role of $\varepsilon_t$. (a) When $\varepsilon_t = 0$, if the retaining gradient and the unlearning gradient are opposite in direction, the update direction $d = 0$, at which point optimization halts and no further progress can be made in optimizing the unlearning objective. (b) When $\varepsilon_t \neq 0$, the update direction $d \neq 0$, and optimization of the unlearning objective proceeds under the condition that it does not unduly affect the retaining objective.
• Figure 4: Visualization of instance & style unlearning. Concepts that have been unlearned are indicated in red.
• Figure 5: Examples of generated images using SDs w/ and w/o MU. Each column represents generated images with the same prompt (denoted by $P_i$) and seed.

Theorems & Definitions (25)

• Definition 2.1
• Proposition 3.1
• Proposition 3.2
• Remark 3.3
• Theorem 3.4: Approximate $\lambda^*$
• Theorem 3.5: Pareto Optimality
• Remark 3.6
• Theorem 3.7: Pareto Stationary
• proof
• proof