PaperlayerView exact revision on arXiv

Source-linked AI summary

Machine Unlearning under Retain-Forget Entanglement

Jingpu Cheng, Ping Liu, Qianxiao Li, Chi Zhang

arXiv:2603.26569v1cs.LG

TL;DR

Machine unlearning can unintentionally damage retained samples that are closely related to the forget set, making average retention inadequate for measuring this entanglement. The paper introduces a two-stage framework combining augmented-Lagrangian forgetting with gradient projection and Wasserstein-2 regularization, and reports effective forgetting with strong retention across diverse tasks. Its main limitation is that conventional projected gradient descent can preserve mean forget-set loss while leaving many forget samples correctly classified.

  • Problem

    Retaining correlated samples during unlearning is difficult because forgetting one group can degrade closely related retained data, while average retain performance can mask this damage.

  • Method

    A two-stage framework uses augmented-Lagrangian optimization to enforce forgetting and gradient projection with Wasserstein-2 distributional regularization to recover correlated-retain performance.

  • Results

    Across subclass-level vision tasks and safety-relevant language benchmarks, the method balances forgetting and retention and outperforms prior approaches in removal fidelity and retained accuracy.

  • Takeaways & Limitations

    The approach supports effective unlearning while preserving performance on retained samples closely related to the forget set.

  • Takeaways & Limitations

    Standard projected gradient descent can preserve mean forget-set loss while redistributing loss so that many forget samples remain correctly predicted.

Abstract

from arXiv · show

Forgetting a subset in machine unlearning is rarely an isolated task. Often, retained samples that are closely related to the forget set can be unintentionally affected, particularly when they share correlated features from pretraining or exhibit strong semantic similarities. To address this challenge, we propose a novel two-phase optimization framework specifically designed to handle such retai-forget entanglements. In the first phase, an augmented Lagrangian method increases the loss on the forget set while preserving accuracy on less-related retained samples. The second phase applies a gradient projection step, regularized by the Wasserstein-2 distance, to mitigate performance degradation on semantically related retained samples without compromising the unlearning objective. We validate our approach through comprehensive experiments on multiple unlearning tasks, standard benchmark datasets, and diverse neural architectures, demonstrating that it achieves effective and reliable unlearning while outperforming existing baselines in both accuracy retention and removal fidelity.

1 INTRODUCTION

Machine unlearning must address retain–forget entanglement because removing one group’s influence can degrade performance on closely related retained samples. The paper proposes a two-stage constrained optimization framework and reports effective forgetting with strong retention across diverse settings.

  • Challenge: Retain–forget entanglement makes unlearning risky because forgetting one group can disrupt predictions on closely related retained examples.Existing evaluations often average across the retain set, potentially masking degradation in sensitive correlated subsets.
  • Method: The framework first increases forget-set loss with an augmented Lagrangian while preserving accuracy on less-correlated retained samples.This stage enforces forgetting while constraining performance on remote retained data.
  • Method: A second gradient-projection stage restores performance on strongly correlated retained samples using Wasserstein-2 loss-distribution regularization.The regularization is intended to stabilize the process and enhance generalization.
  • Evaluation: The method is evaluated across diverse subclass-level tasks, neural architectures, and benchmark datasets.The evaluation covers multiple forgetting scenarios rather than a single model or dataset.
  • Results: The experiments report effective forgetting, high retained-data accuracy, and preserved performance on retained samples closely related to the forget set.The introduction also reports advantages over prior methods in structured selective unlearning settings.

2 RELATED WORKS

Related work spans formal machine-unlearning objectives, approximate post-hoc methods, and growing attention to correlated retained subsets. The paper emphasizes that average retain accuracy can conceal concentrated degradation near the forget set.

  • Machine Unlearning: Machine unlearning seeks model behavior comparable to retraining without deleted data, but full retraining is often infeasible for large models.Prior approaches therefore use approximate updates, pruning, prompt editing, and other alternatives.
  • Retain–Forget Entanglement: Retain–forget entanglement concentrates accuracy drops on retained examples most similar to the forget set.Subclass-level forgetting is one example in which semantically close classes create this concern.
  • Retain–Forget Entanglement: Reporting average retain performance can mask degradation on correlated subsets, motivating separate evaluation of adjacent and remote retained data.The paper distinguishes retained samples near the forget set from less-related samples.

3 PROBLEM FORMULATION

The problem formulation models unlearning as removing a forget subset while preserving the remaining data, with special attention to retained samples correlated with that subset. It separates retained data into adjacent and remote groups to represent different sensitivity levels.

  • Objective: Given a trained model and forget subset Df, unlearning seeks updated parameters that remove Df’s influence while preserving performance on Dr.Dr denotes the remaining data after excluding the forget set.
  • Objective: Classical formulations overlook additional structure in Dr, although forgetting can disproportionately affect a correlated portion of the retain set.Average retain performance does not necessarily reflect this concentrated degradation.
  • Retain-Set Structure: The adjacent retain set Dadj consists of retained examples correlated with Df and more sensitive to forgetting.These samples form the entangled subset requiring explicit protection.
  • Retain-Set Structure: The remote retain set Drem contains the remaining retained examples that are less related to the forget set.The framework treats remote samples as a distinct reference group.
  • Sources of Entanglement: Entanglement can arise in subclass-level vision tasks or when retained language samples are semantically related to the forget set.Examples include other subclasses within a superclass and non-offensive comments referring to the same group as offensive comments.

1. The model retains its performance on Dr, especially on samples belonging to Dadj

The formulation requires retaining performance on the remaining data, especially correlated retained samples, while removing or mitigating the forget set’s influence. The meaning of forgetting depends on the application, ranging from retraining equivalence to reducing forget-set performance.

  • Objectives: The model should retain performance on Dr, with particular attention to retained samples strongly correlated with Df.The objective also requires removing or mitigating the forget set’s influence.
  • Application-Dependent Forgetting: In privacy-focused settings, forgetting may require the updated model to emulate retraining on the retained set.This objective differs from application settings that prioritize reducing performance on undesirable forget-set patterns.

4 METHODS

The framework treats retain–forget entanglement as a constrained optimization problem: Stage 1 enforces forgetting while protecting remote retained samples, and Stage 2 restores adjacent-retain accuracy while controlling forget-set loss distributions. Wasserstein-2 regularization makes the second-stage projection more fine-grained and computationally convenient than mean-loss control alone.

  • Framework overview: The method splits the entangled retain set into less-related samples protected during forgetting and adjacent samples whose accuracy is restored afterward.Stage 1 targets Df and Drem_r, while Stage 2 targets Dadj_r without compromising Df or Drem_r.
  • Stage 1: forgetting via controlled optimization: Stage 1 uses an augmented Lagrangian to increase forget-set loss while constraining degradation on the less-related retained subset.Adaptive multiplier updates adjust the penalty according to constraint violation, avoiding manual tuning of a fixed trade-off coefficient.
  • Stage 1: forgetting via controlled optimization: Stage 1 preserves strong remote-retain performance but typically degrades adjacent-retain performance because semantic or distributional overlap creates conflicting effects.The method therefore avoids explicitly optimizing Dadj_r in Stage 1 and addresses it in the second stage.
  • Stage 2: gradient projection: Classical PGD removes adjacent-retain gradient components aligned with forget and remote-retain gradients, but mean-loss preservation can leave many forget samples correctly predicted.Uneven loss redistribution can produce near-zero losses for some forget samples and increase forget-set accuracy despite stable average loss.
  • Stage 2: Wasserstein-regularized projection: W-PGD addresses this failure by adding a Wasserstein-2 penalty that controls the forget-set loss distribution while projecting adjacent-retain gradients away from forget and remote-retain directions.The empirical W2 distance is computed conveniently by sorting one-dimensional loss samples, avoiding density estimation or strong parametric assumptions generally needed for KL divergence.
  • Stage 2: Wasserstein-regularized projection: The modified projection keeps the mean and distributional loss behavior on the forget set mostly unchanged during adjacent-retain recovery, while sufficiently small updates leave forget and remote-retain losses changed only at second order.The procedure is summarized as augmented-Lagrangian optimization followed by W2-guided gradient projection.
  • Stage 2: Wasserstein-regularized projection: Under the proposition’s large-minimum-loss condition, W-PGD bounds forget-set accuracy by a small constant; empirically, α = 0.5 yields a more uniform loss distribution and zero forget-set accuracy.Although α = 1 minimizes the stated bound for fixed ε, the experiments report better overall performance with α = 0.5.
  • Design rationale: W2 provides a closed-form distributional constraint for one-dimensional empirical losses, avoiding approximations and prior assumptions that can make KL-based constraints costly.The paper identifies this computational convenience as a key advantage of the second-stage design.

5 EXPERIMENTS

Experiments evaluate the method across subclass-level vision and safety-critical language unlearning, showing that it balances forgetting with retention under strong retain–forget entanglement. Results span multiple datasets and architectures, including CIFAR-100, ToxiGen, CelebA, and TinyImageNet.

  • Experimental setup: Experiments cover CIFAR-100, TinyImageNet, ToxiGen, and CelebA across multiple unlearning tasks and neural architectures.The setup includes subclass-level vision tasks and a safety-critical language task.
  • CIFAR-100: On CIFAR-100, effective forgetting often reduces accuracy on the adjacent retain subset, exposing the retain–forget trade-off.Several methods achieve low forget-set performance but noticeably degrade accuracy on closely related retained samples.
  • CIFAR-100: 0.00% training accuracy on the forget class accompanies high retained-data and test performance for the proposed method.The result is reported as complete unlearning without compromising generalization or utility on remaining data.
  • ToxiGen: On ToxiGen, the proposed method achieves the lowest forget-set accuracy while preserving 92.73% accuracy on the adjacent retain set.The task corrects toxic LGBTQ samples mislabeled as normal while preserving non-toxic LGBTQ samples.
  • CelebA: On CelebA, the proposed method achieves 92.38% retained-set accuracy while avoiding the adjacent-set collapse seen in methods that drive forget accuracy near zero.The benchmark uses highly similar semantic attributes to make retain–forget entanglement particularly pronounced.
  • Generalization and ablation: The method generalizes to TinyImageNet with a ViT, and removing W2 regularization raises forget-set accuracy to 18.87% on training data and 14.33% on test data.The TinyImageNet experiment supports generalization beyond one dataset or architecture, while the ablation supports W2 regularization for preserving forgetting.

6 CONCLUSION

The paper addresses retain–forget entanglement with a two-stage optimization framework that combines constrained forgetting with recovery of correlated retained samples. Experiments show effective forgetting alongside preserved retained accuracy across vision and language settings.

  • Conclusion: The framework first enforces forgetting while preserving less-related retained samples, then uses gradient projection with a Wasserstein-2 constraint to recover correlated retained performance.The conclusion describes the two stages as constrained optimization followed by distributionally regularized refinement.
  • Conclusion: Experiments across subclass-level vision tasks and safety-relevant language benchmarks show improved balance between removal fidelity and accuracy preservation over prior approaches.The reported evaluation supports the framework across multiple task types.
  • Conclusion: The results emphasize correlation-aware unlearning as a practical approach for handling retained samples strongly related to the forget set.The conclusion frames retain–forget entanglement as a central practical machine-unlearning challenge.

A PROOF FOR PROPOSITIONS AND THEOREMS

The proofs analyze the projected update through first-order expansions and Wasserstein-based bounds. They establish that the adjacent-retain loss changes only at second order and decreases under the stated conditions, while correcting a minor theorem-statement typo without affecting validity.

  • Proof strategy: First-order expansion of the forget loss is used to analyze the projected update.The proof explicitly begins from a first-order Taylor expansion of Lf.
  • Proof strategy: The projection makes the first-order difference in Lf vanish, leaving a dominant term that is second-order in η.This identifies the local effect of the projected step on the forget loss.
  • Proof result: The change in Ladj is likewise second-order in η, leading to a strict decrease in Ladj under the proof conditions.The proof separately tracks the adjacent-retain objective and concludes strict decrease.
  • Proof note: The appendix notes a minor typo in Theorem 4.2 and states that it does not affect the theorem’s validity or presented proof.The correction is confined to the theorem statement.
  • Proof result: The argument uses E[|X − Y|] ≤ W2(P, Q) to relate expected differences between random variables and the Wasserstein-2 distance.This inequality appears in the derivation of the stated bound.
  • Proof result: For n-class classification, correct predictions imply cross-entropy loss at most log n, which supports the proof’s accuracy estimation.The proof combines this classification-loss bound with a lower bound on the loss.

B.1 EXPERIMENTAL DETAILS

The experiments use vision and language datasets with pretrained neural architectures and explicitly separated forget, adjacent-retain, and remote-retain subsets. In the second stage on CIFAR-100, forgetting remains effective while adjacent-retain accuracy improves during training.

  • Base models: CIFAR-100 uses pretrained ResNet-18, while TinyImageNet uses a pretrained ViT-B-32 fine-tuned on superclass data.The CIFAR-100 model is fine-tuned for 30 epochs; the TinyImageNet setup uses the same broad fine-tuning design.
  • Learning dynamics: Figure 2 separates adjacent and remote retain sets while showing training and test accuracy during the second stage on CIFAR-100 with ResNet18.The adjacent set contains adjacent samples, and the remote set contains remote samples.
  • Implementation: The implementation uses one first-stage epoch and six second-stage epochs for CIFAR-100 and TinyImageNet, versus two second-stage epochs for ToxiGen.The first stage uses augmented-Lagrangian optimization, with clipped cross-entropy applied to the forget set.
  • Implementation: Stage 2 uses SGD with dataset-specific learning rates and different batch sizes for remote retain, adjacent retain, and forget sets.For CIFAR-100 and TinyImageNet, the remote-retain batch size is 512, while adjacent-retain and forget batches use 128.
  • Learning dynamics: The forget-set accuracy remains zero, remote-retain accuracy stays consistently high, and adjacent-retain accuracy steadily improves during second-stage training.These dynamics are reported for CIFAR-100 with ResNet18.

B.3 COMPARISON OF TRAINING TIME AND MEMORY USAGE

The method requires somewhat longer training than the fastest baselines but remains far more efficient than full retraining. Its memory use is slightly above several lightweight baselines and below SCRUB and SalUn, while privacy is also evaluated with membership inference attacks.

  • Training time: 2.5 minutes is the reported running time for SSD, while GA, SCRUB, and SalUn each finish in under 10 minutes.Fine-tuning and the proposed method take slightly longer, but all remain significantly more efficient than full retraining.
  • Memory usage: The proposed method uses slightly more memory than fine-tuning, GA, and ℓ1-sparse, but less than SCRUB and SalUn.The comparison uses batch size 128 for all methods.
  • Memory usage: The reported memory results indicate no significant additional memory overhead compared with other unlearning methods.The method’s two-stage optimization accounts for its modest increase over several baselines.
  • Privacy evaluation: Membership inference attacks assess whether forget-class samples can be distinguished as having been part of training using a confidence-based predictor.The resulting MIA efficacy measures the proportion of forget-set samples correctly identified.

B.5 ABLATION STUDIES

The ablations show that first-stage constraints and Wasserstein regularization choices materially affect forgetting and retained accuracy. Across a reasonable penalty range, the method remains comparatively robust, while the accompanying figure evaluates MIA efficacy.

  • Privacy ablation: Figure 3 evaluates the membership-inference-attack efficacy of different unlearning methods on CIFAR100 with ResNet-18.The figure provides a privacy-oriented comparison across methods.
  • First-stage constraints: Including adjacent-retain samples in the first-stage constraint raises forget-set accuracy to 7.13% on training and 5.00% on test data.The same variant lowers adjacent-retain test accuracy to 72.75%.
  • First-stage constraints: Separating the forget set from adjacent-retain samples in stage one is crucial for effective unlearning in this method.Constraining the entire retain set impairs forgetting and adjacent-retain performance relative to constraining only the remote retain set.
  • Wasserstein-term sensitivity: At α = 0, forget-set accuracy reaches 19.67% on training data and 16.33% on test data, indicating ineffective forgetting.Both α = 0.5 and α = 1 produce low forget accuracy and high retained accuracy, while α = 0.5 offers the more balanced overall performance.
  • Wasserstein-term sensitivity: Changing α from 0.5 to 1 increases training forget accuracy by 0.74 percentage points and out-of-class retain accuracy by 0.75 points, with a 0.17-point adjacent-retain gain.The fully incorporated W2-distance term therefore does not necessarily optimize all reported metrics.
  • Penalty sensitivity: For µ ∈ {5, 10, 20}, test forget accuracy stays at or below 3%, while adjacent- and remote-retain accuracies vary only slightly.The results indicate robustness to the penalty parameter within this range.

B.7 ADDITIONAL RESULTS

Additional experiments show that the method maintains effective forgetting and retained-set accuracy across ViT, noisy adjacency definitions, feature-based adjacency, and retraining comparisons.

  • ViT results: 0% training forget accuracy is achieved while preserving high retained-set accuracy on ViT for CIFAR-100 superclass unlearning.Gradient ascent also reaches zero forget accuracy but substantially degrades retained performance, especially in the adjacent subset; SCRUB retains some forget accuracy.
  • Robustness to imperfect adjacency: Under 20% adjacency-label noise, forget accuracy remains 0% on training data and below 6% on test data, with modest retain-accuracy changes.The study tests both incorrectly assigning remote samples as adjacent and incorrectly assigning adjacent samples as remote.
  • Feature-space adjacency: Feature-space kNN adjacency yields strong forgetting and high accuracy on both adjacent and remote retain subsets, comparable to superclass-defined adjacency.The kNN construction uses output features, k = 20 neighbors per forget sample, and the top 10% of retained samples by adjacency score.
  • Comparison with Retraining: Compared with full retraining, the method produces substantially lower forget accuracy while preserving competitive performance on adjacent and remote retain subsets.Retraining generally maintains high retain accuracy but does not always achieve strong erasure, with forget accuracy often remaining relatively high.
Loading 2603.26569v1…