Multi-Objective Reference-Aligned Machine Unlearning

Technical Summary: Multi-Objective Reference-Aligned Machine Unlearning (RAUL)

Problem Statement
Machine unlearning (MU) aims to remove the influence of specific training samples (the "forget set," $F$) from a deployed model while preserving performance on the remaining data (the "retain set," $R$). Existing approximate methods typically rely on single-objective optimization, such as gradient ascent (maximizing loss on $F$) or random relabeling. The authors argue that these approaches induce catastrophic forgetting because the optimization dynamics of "forgetting" (repelling correct predictions) conflict with "retention" (maintaining correct predictions on $R$). This conflict causes the model to drift from its pre-trained knowledge, resulting in significant utility loss on the retain set.

Methodology: Reference-Aligned UnLearning (RAUL)
The paper proposes RAUL, a multi-objective optimization (MOO) framework that reframes unlearning as a joint optimization of forgetting and retention, rather than a single adversarial task.

1. Objective Formulation:

   • Retention Objective ($L_{retain}$): Standard cross-entropy minimization on the retain set $R$ to preserve model utility.
   • Reference-Based Forgetting Objective ($L_{forget}$): Instead of maximizing loss on $F$, RAUL minimizes the Kullback-Leibler (KL) divergence between the model's predictions on $F$ and a reference distribution ($p_{ref}$) representing "unseen" data. This aligns the model's confidence on forgotten samples with how it would behave on data it has never encountered.
     $$\min_{\theta} L_{forget}(f_\theta, F) = \mathbb{E}_{x_u \sim F} [KL (p_\theta(y | x_u) \parallel p_{ref}(y))]$$

2. Reference Distribution Strategies:
   The paper introduces two strategies to instantiate $p_{ref}$:

   • Uniform Distribution (UniDist): Sets $p_{ref}(y) = 1/C$ (where $C$ is the number of classes). This acts as a principled baseline requiring no additional data, encouraging maximum uncertainty on forgotten samples. Minimizing KL to a uniform distribution is equivalent to maximizing prediction entropy.
   • Held-out Distribution (HeldDist): Computes an empirical reference distribution from a disjoint, unseen dataset ($D_{ref}$) that was not part of the original training, validation, or test sets. This provides a stable anchor that mimics the behavior of truly unseen data, potentially offering better retention stability.

3. Optimization via Jacobian Descent:
   The problem is treated as a bi-objective optimization: $\theta_u = \arg \min_{\theta} (L_{retain}, L_{forget})$. To solve this without manually weighting the objectives (which can lead to conflicting gradients), the authors employ Jacobian descent using the UPGrad (Unconflicting Projection of Gradients) aggregator.

   • The gradients of the two objectives ($g_R$ and $g_F$) are stacked into a Jacobian matrix.
   • UPGrad projects each gradient onto the dual cone of all objective gradients and averages them.
   • This aggregation finds a Pareto-stationary point where the update direction does not ascend one objective while descending the other, effectively resolving the gradient conflict.

Key Contributions

• Reformulation of Unlearning: Shifts the paradigm from adversarial loss maximization to reference-aligned prediction, constraining the forgetting objective to prevent model drift.
• Multi-Objective Framework: Explicitly models the tension between forgetting and retention as a bi-objective problem, solved via Pareto-stationary optimization (Jacobian descent) rather than linear scalarization.
• Reference Strategies: Proposes two distinct reference distributions (Uniform and Held-out) to control the forgetting behavior, offering a trade-off between data-free implementation and retention stability.
• Stable Optimization: Demonstrates that replacing unbounded loss maximization with bounded KL alignment reduces gradient conflict, leading to more stable updates.

Experimental Results
The authors evaluated RAUL on image classification tasks (CIFAR-10 with ResNet-18) under 10% and 50% random forgetting ratios, comparing against baselines including Gradient Ascent (GA), Random Relabeling (RL), Influence Functions (IU), and SalUn.

• Performance Gap: RAUL variants achieved the smallest average performance gaps compared to full retraining.
  • At 10% forgetting, HeldDist-RAUL with UPGrad achieved the best average gap of 2.23, outperforming the strongest baseline (IU at 6.06) and significantly beating GA (16.90).
  • At 50% forgetting, where baselines like GA collapsed (gap of 49.19), UniDist-RAUL with UPGrad achieved the best gap of 6.75, while HeldDist-RAUL variants maintained gaps around 9.22.
• Retention vs. Forgetting Trade-off:
  • HeldDist-RAUL prioritized retention, achieving high Retain Accuracy (RA ~98.89%) and Test Accuracy (TA ~92.12%) with conservative forgetting.
  • UniDist-RAUL offered a more aggressive forgetting profile (higher Unlearning Accuracy) while still maintaining the lowest overall gap to retraining.
• Privacy Metrics: RAUL methods maintained Membership Inference Attack (MIA) scores close to the retraining baseline (e.g., ~15.31 vs. 13.13 for retraining), whereas GA dropped to near-random levels (1.93), indicating RAUL removes influence without the instability of adversarial attacks.
• Stability: RAUL variants exhibited markedly lower standard deviations in performance compared to GA and $\ell_1$-sparse methods, confirming the stability of reference-aligned updates.

Significance and Claims
The paper claims that RAUL addresses the fundamental bottleneck of explicit conflict between forgetting and retention in machine unlearning. By replacing unbounded loss maximization with a bounded KL alignment to a reference distribution, the method:

1. Reduces Gradient Conflict: The MOO formulation with UPGrad aggregation ensures updates do not degrade the retain set while forgetting the target set.
2. Prevents Catastrophic Forgetting: The bounded nature of the forgetting objective prevents the model from drifting too far from its pre-trained knowledge.
3. Offers Flexibility: The framework provides a choice between a data-free approach (UniDist) and a data-dependent, higher-stability approach (HeldDist).

The authors conclude that while UniDist-RAUL serves as a robust fallback requiring no extra data, the availability of a held-out set allows HeldDist-RAUL to achieve tighter retention. They note that future work will explore extending this to generative models, LLMs, and utilizing Multi-Objective Evolutionary Algorithms (MOEAs) for more diverse Pareto-optimal solutions.