Technical Summary: To forget is to preserve: Machine Unlearning for 3D medical image segmentation

1. Problem Statement

The rapid adoption of deep learning in medical imaging, particularly for 3D volumetric data like MRI and CT, raises critical data privacy concerns. Regulations such as the GDPR (Article 17), HIPAA, and India's DPDPA 2023 grant individuals the "right to erasure," requiring that personal data be removed not only from storage but also from trained machine learning models.

In the context of 3D medical image segmentation, this presents a unique challenge. Unlike 2D classification tasks where "forgetting" might involve removing entire classes, medical unlearning requires removing data at the subject level (i.e., specific patients). Furthermore, the output is a dense, voxel-wise prediction rather than a single label. The standard approach to compliance—retraining a model from scratch on the remaining data—is computationally prohibitive for 3D volumetric datasets, often requiring hundreds of GPU-hours. Consequently, there is a need for approximate machine unlearning strategies that can modify an existing pre-trained model to "forget" specific subjects while maintaining high performance on the retained data, without the cost of full retraining.

2. Methodology

2.1 Backbone Architecture

The study utilizes a 3D ResNet-50 architecture pre-trained with the Med3D framework. This backbone is chosen for its ability to learn robust volumetric representations from diverse 3D medical datasets (lung, liver, brain) across CT and MRI modalities. The model is decomposed into:

• Encoder ($E_\phi$): A 3D ResNet-50 that extracts spatiotemporal features.
• Decoder/Segmentation Head ($G_\psi$): A multi-layer decoder involving 3D transposed convolutions and refinement layers that map features to voxel-wise class logits.

The model was fine-tuned on the MRBrainS18 dataset, which contains expert annotations for brain structures in elderly subjects.

2.2 Experimental Setup

• Dataset: MRBrainS18 (10 subjects).
• Forget Set ($D_f$): 2 subjects (20% of the data).
• Retain Set ($D_r$): The remaining 8 subjects.
• Task: Binary segmentation (brain tissue vs. non-brain tissue).
• Training Horizons: Unlearning strategies were evaluated over 20 and 50 epochs.
• Metrics:
  • Dice Similarity Coefficient (DSC): Measures spatial overlap.
  • Mean Absolute Error (MAE): Measures prediction probability deviation.
  • Gap Metrics: $\Delta$Dice ($D_r - D_f$) and $\Delta$MAE ($D_f - D_r$) to quantify the trade-off between forgetting and retention.

2.3 Unlearning Strategies

The study evaluates several approximate unlearning strategies adapted from the I2I framework and other literature, categorized into four mechanics:

1. Gradient-Based Approaches (GA):

   • Retain Label / Max Loss: Modifies the loss function by negating the gradient for the forget set (maximizing error) while minimizing error for the retain set.
   • Risk: Unbounded gradient ascent can lead to catastrophic degradation of the retain set.

2. Label and Data Manipulation:

   • Noisy Label (NL): Replaces the ground-truth voxel labels of the forget subjects with uniformly random labels during training. The model is optimized on this manipulated distribution, effectively overwriting memorized representations for the forget subjects while reinforcing the retain subjects.

3. Feature-Space Interventions:

   • Random Encoder (RE): Overwrites the latent feature maps of forget subjects with Gaussian noise before they reach the decoder.
   • Learn Others (LO) & Learn Noise (LN): Uses a teacher-student distillation paradigm. The student encoder is forced to map forget subjects to a generic target (either the average features of the current batch or dynamic Gaussian noise) while mirroring the teacher for retain subjects.

4. Fine-tuning and Selective Retraining:

   • Fix Decoder (FD): Freezes the segmentation head (decoder) and only updates the encoder parameters via distillation. This aims to confine forgetting to the latent feature space to prevent decoder drift.

3. Key Results

The experiments revealed significant differences in the trade-offs between forgetting effectiveness and utility retention:

• Noisy Label (NL) Strategy: This approach demonstrated the best overall trade-off.

  • After 50 epochs, it achieved a Dice Gap of +0.6963 and an MAE Gap of +0.091.
  • It reduced the Dice score of the forget set by 93% (from 0.9226 to 0.0613).
  • Crucially, it maintained 84% accuracy on the retain set (Dice score of 0.7576).
  • Unlike other methods, its performance improved with longer training horizons, simultaneously strengthening forgetting and retention.

• Gradient-Based Methods (Retain Label/Max Loss): While effective at removing forget-set performance, these methods suffered from catastrophic degradation of the retain set, especially at 50 epochs. The unconstrained loss negation drove the model into poor regions of parameter space.

• Feature-Space Interventions (RE, LO, LN):

  • Random Encoder preserved retain performance but achieved very limited forgetting (Dice gap ~0.01–0.04), suggesting the decoder could compensate for the corrupted encoder features.
  • Learn Others/Noise caused a near-complete collapse of retain performance, indicating that aggressive feature manipulation disrupted learned representations too severely.

• Fix Decoder (FD): Showed moderate unlearning at 20 epochs but became unstable at 50 epochs, eventually reversing the intended behavior by performing better on the forget set than the retain set.

4. Contributions and Significance

The paper claims the following contributions and significance:

1. First Systematic Benchmark: To the best of the authors' knowledge, this is the first study to systematically evaluate multiple approximate unlearning strategies specifically for 3D medical image segmentation using subject-level forget sets. Previous benchmarks were largely limited to 2D classification tasks (e.g., CIFAR-10, ImageNet).
2. Empirical Evaluation of 3D Unlearning: The study adapts and evaluates these strategies on a Med3D ResNet-50 framework, providing the first comprehensive empirical data on removing patient-specific knowledge from 3D volumetric models.
3. Reproducible Evaluation Protocol: The authors establish a protocol using segmentation-specific metrics (Dice and MAE) to quantify the trade-off between forgetting effectiveness and retained clinical utility, addressing the difficulty of applying 2D metrics to 3D dense prediction tasks.
4. Identification of Optimal Strategy: The study identifies the Noisy Label strategy as the most robust method for this domain, capable of achieving strong forgetting while preserving the majority of the model's utility on retained data.

5. Conclusion and Future Work

The study concludes that while many unlearning strategies exist, they often fail in the high-dimensional, complex context of 3D medical imaging, either failing to forget effectively or destroying the model's utility. The Noisy Label strategy emerges as the most viable solution for subject-level unlearning in this domain.

The authors note that future work should focus on:

• Formal adversarial privacy evaluations using membership inference or model inversion attacks to rigorously prove the removal of residual information.
• Investigating adaptive optimization algorithms (e.g., dynamic alpha scheduling) and parameter-efficient fine-tuning (e.g., LoRA) to further stabilize unlearning in feature space.

The paper emphasizes that developing effective unlearning methods is a necessary step toward building trustworthy clinical AI systems that respect privacy principles in real-world healthcare settings.