As generative models become increasingly powerful and pervasive, the ability to unlearn specific data due to privacy concerns, legal requirements, or harmful content has become increasingly important. Unlearning aims to selectively remove the influence of particular data without costly retraining from scratch. Effective unlearning requires (i) forgetting undesired data, (ii) preserving generation quality, (iii) retaining the influence of desired training data on model parameters, and (iv) computational efficiency. We propose fast unlearning algorithms based on loss gradient orthogonalization for unconditional and conditional generative models. We consider the loss function of the original model restricted to the retain data and augment it by a term promoting the orthogonality between the two gradients of the loss function evaluated at the retain dataset and the forget dataset. We coin this method unlearning via orthogonalization (UNO) and show that the additional weak orthogonality constraint accelerates unlearning. The inclusion of the weak orthogonality constraint implies that each gradient descent step utilizes local information about the geometry via the Hessian of the loss function. Our geometry aware gradient descent step significantly accelerates the learning of the retain data loss function compared to a gradient descent which uses only gradient information. By alternating between an UNO and a gradient surgery step, we obtain UNO-S, which further reduces execution time. We demonstrate that our algorithms effectively forget data while preserving the fidelity of the original model. On standard image benchmarks, they achieve orders-of-magnitude faster unlearning than existing methods such as gradient surgery. We validate our approach on datasets of increasing complexity (MNIST, CelebA, and ImageNet-1K) and generative architectures ranging from VAEs to diffusion transformers.
Figures & tables
Figure 1: MNIST samples generated by the original model (left) and after unlearning digit " 1 " with UNO (right), using identical noise inputs for the decoder.
Figure 2: Distribution of generated digits before (left) and after unlearning (right), for a single run of UNO. Each histogram shows data for 500 generated samples. A bar in the right panel is colored green if the fraction of the corresponding digit increases after unlearning, and red if it decreases.
Figure 3: Distribution of digits in the latent space of the original model for MNIST. The colors are obtained by mapping 50,000 images to the two-dimensional latent space, labelling them according to the color associated with the image. To obtain smooth regions we performed averaging over nearby points in latent space.
Figure 4: CelebA samples generated by the original model (top) and after unlearning "male" faces with UNO (bottom), using identical noise inputs for the decoder.
Figure 5: ImageNet-1K samples generated by the original model (top) and after replacement unlearning with UNO-S (bottom). We replace class 207 (golden retriever) with class 208 (labrador retriever), using identical noise inputs for the diffusion transformer. The labels classifying each image were provided by the pretrained classifier.
Figure 6: ImageNet-1K samples generated by the original model (top) and after replacement unlearning with UNO-S (bottom). We replace class 207 (golden retriever) with class 665 (moped), using identical noise inputs for the diffusion transformer. The labels classifying each image were provided by the pretrained classifier.
Figure 7: ImageNet-1K samples generated by the original model (top) and after replacement unlearning with UNO-S for different ratios of the cardinalities of the retain and forget datasets. We show replacement unlearning of class 207 (golden retriever) with class 208 (labrador retriever) (left) and replacement unlearning of class 207 (golden retriever) with class 665 (moped) (right). We use identical noise samples for generating the before and after images.
Dataset
Algorithm
Time to unlearn (s) ↓
FID ↓
Δt (ms)
5 -step f-ratio ↓
MNIST
Gradient ascent (A)
0.024 (0.004)
612.3 (4.9) ✗
5.1 (0.47)
0.00 (0.01)
Ascent descent (A-D)
0.025 (0.006)
266.9 (19.3) ✗
5.1 (0.09)
0.01 (0.01)
Gradient surgery (S)
1.016 (0.907)
23.0 (0.2) ✓
7.1 (0.06)
0.05 (0.03)
UNO
0.055 (0.009)
21.8 (0.2) ✓
20.5 (0.27)
0.00 (0.00)
UNO-S
0.041 (0.010)
21.8 (0.4) ✓
15.1 (0.35)
0.00 (0.00)
Moon et al. (2024)
0.065 (0.052)
25.1 (0.4) ✓
5.3 (0.04)
0.03 (0.02)
Table 1: Performance of various algorithms for class/feature unlearning with VAEs on MNIST and CelebA and diffusion transformers on ImageNet-1K. Each experiment is repeated 10 times, and the standard deviations are shown in parentheses. FID values are calculated using 25,000 samples for MNIST and CelebA and 22,000 samples for ImageNet-1K. ✗ indicates that the generated samples after unlearning are unrecognizably different from the original model. Bold indicates the best score among instances with acceptable FID. ✓ indicates the generated samples after unlearning are perceptually indistinguishable from the original model in terms of visual fidelity. An asterisk (*) denotes cases where the algorithm failed to reach the target fraction of forget samples in the generated images within the allotted number of training steps. We only report on (A) and (A-D) for MNIST since they demonstrate catastrophic forgetting in all cases; for CelebA (A) and (A-D) result in NaN for FID, generating white images. The original models for MNIST, CelebA and ImageNet-1K have FID 20.7 , 166.3 and 12.7 , respectively, and we are unlearning the digit 1 , faces identified as "male" and images of class 207 (golden retriever), respectively.
Retain
Forget
Dataset
Model
ROC-AUC [95% CI]
Attack Adv.(%) [95% CI]
TPR@1% FPR (%) [95% CI]
ROC-AUC [95% CI]
Attack Adv.(%) [95% CI]
TPR@1% FPR (%) [95% CI]
Original
0.534 [0.526, 0.543]
5.36 [4.52, 6.83]
1.06 [0.80, 1.35]
0.484 [0.459, 0.506]
0.70 [0.53, 4.23]
1.32 [0.53, 2.11]
Gradient ascent
0.498 [0.489, 0.506]
1.68 [0.59, 3.12]
0.96 [0.69, 1.41]
0.507 [0.483, 0.531]
2.56 [1.23, 7.22]
0.62 [0.18, 1.68]
Ascent descent
0.493 [0.485, 0.502]
1.15 [0.38, 2.37]
1.13 [0.81, 1.48]
0.490 [0.467, 0.515]
0.53 [0.44, 5.46]
1.15 [0.26, 2.38]
MNIST
Moon et al.
0.530 [0.522, 0.539]
4.99 [4.35, 6.44]
0.91 [0.62, 1.18]
0.492 [0.468, 0.515]
1.23 [0.70, 5.29]
1.41 [0.44, 2.82]
UNO
0.504 [0.496, 0.513]
1.40 [0.96, 2.70]
1.14 [0.81, 1.56]
0.480 [0.457, 0.505]
0.88 [0.44, 4.41]
1.67 [0.53, 2.64]
Table 2: Class-separated loss-based membership inference, for which lower loss predicts membership. Each entry reports the point estimate followed by a stratified 95% percentile bootstrap confidence interval in brackets. Members and non-members were resampled independently at their original group sizes for 1,000 image-level bootstrap replicates. Attack advantage is maxτ(TPRτ−FPRτ) . TPR@1% FPR is the maximum TPR attained at an empirical FPR no greater than 1%.
Model
ROC-AUC [95% CI]
Attack adv. (%) [95% CI]
TPR@1% FPR (%) [95% CI]
ESD-u-1
0.528 [0.448, 0.557]
6.00 [0.46, 10.62]
1.08 [0.15, 2.31]
ESD-3
0.530 [0.453, 0.559]
6.31 [0.31, 10.92]
2.46 [0.46, 4.62]
UNO
0.526 [0.451, 0.554]
5.85 [0.46, 10.15]
1.08 [0.15, 3.08]
UNO-S
0.528 [0.451, 0.557]
5.69 [0.31, 10.77]
3.54 [0.31, 5.69]
Surgery
0.526 [0.454, 0.552]
5.85 [0.31, 10.31]
1.69 [0.31, 2.92]
Table 3: Loss-based membership inference for the unlearned DiT models, treating retained class 208 training images as members and forgotten class 207 training images as non-members. Attack advantage is maxτ(TPRτ−FPRτ) . TPR@1% FPR is the maximum TPR attained at an empirical FPR no greater than 1%. Brackets give 95% bootstrap confidence intervals from 1000 replicates.
Dataset
Model
ΔpF(50)
Gradient ascent
0.000
Ascent–descent
0.315
Moon et al.
0.261
MNIST
UNO
0.270
UNO-S
0.288
Surgery
0.288
Table 4: Normalized change in forgotten-class generation during a single-run of 20-shot relearning attack with 50 update steps.
Figure 8: The f-ratio versus the cosine similarity between loss gradients for various algorithms until a threshold f-ratio of τ=0.18 is reached. The data here correspond to replacement unlearning of golden retrievers with labrador retrievers, discussed in Section 4 . We show the average cosine similarity and f-ratio for 10 independent realizations, together with confidence intervals indicating variability in the cosine similarity. The numbers on each data point refer to the respective update steps.
Figure 9: Computational complexity of unlearning methods across model scales as a function of the cost of a single model forward pass for a single sample, measured in GFLOPs. Left: memory usage per sample, defined as the difference between the peak GPU memory occupied during an update step and the memory occupied immediately before that step, divided by the training batch size. Right: Δt per sample, defined as the duration of a single update step divided by the training batch size. Both panels use logarithmic axes. The dashed lines show power-law fits to the UNO results.
Appendix figures & tables10 assets
Supplementary material from the paper’s appendix.
Appendix
Dataset
Algorithm
η
K
Bβo
Bβh
α
B
NFID
MNIST (Class: 1)
Gradient ascent (A)
10−3
530
-
-
-
128
25,000
Ascent descent (A-D)
10−3
530
-
-
-
128
25,000
Gradient surgery (S)
10−3
530
-
-
-
128
25,000
UNO
10−3
530
103
-
-
128
25,000
UNO-S
10−3
530
103
-
-
128
25,000
H
10−3
530
-
103
10−8
128
25,000
Appendix
Table 5: Experiment hyperparameters
Figure 10: Catastrophic forgetting induced by gradient ascent (left) and ascent-descent (right) for MNIST at a learning rate of 10−3 .
Figure 11: Generated samples after unlearning digit 1 via gradient surgery with ascent in the forget direction (SA), described in Algorithm 3 , for two different learning rates: 10−3 (left), 10−5 (right). SA was run for K=53 training steps on the left and K=530 training steps on the right.
Figure 12: Generated samples after unlearning digit 1 via gradient surgery with descent in the retain direction (S), described in Algorithm 4 , for two different learning rates: 10−3 (left), 10−5 (right). S was run for K=530 training steps for both learning rates.
Figure 13: Results for unlearning on CelebA with UNO, illustrated using 18 pairs of generated images. The images labeled "Before" were generated using the original model. Each image labeled "After" was generated after unlearning using the same noise sample for the decoder as the corresponding "Before" image.
Figure 14: Results for eyeglasses removal on CelebA with UNO-S, illustrated using 18 pairs of generated images. The images labeled "Before" were generated using the original model. Each image labeled "After" was generated after unlearning using the same noise sample for the decoder as the corresponding "Before" image.
Figure 15: Results for multiclass unlearning on CelebA with UNO-S, where the composite class "male with dark hair" i.e. ((Male = 1) & ((Black_Hair = 1) | (Brown_Hair = 1))) is unlearned. The same noise samples were decoded by the original model (top) and the unlearned model (bottom) to generate the faces.
Figure 16: Results for unlearning on ImageNet with UNO-S, illustrated using 18 pairs of generated images. The images labeled "Before" were generated using the original model. Each image labeled "After" was generated after unlearning using the same noise sample for the diffusion transformer as the corresponding "Before" image.
Figure 17: ImageNet-1K samples generated by the original model (top) and after replacement unlearning with UNO for different ratios of the cardinalities of the retain and forget datasets. We show replacement unlearning of class 207 (golden retriever) with class 208 (labrador retriever) (left) and replacement unlearning of class 207 (golden retriever) with class 665 (moped) (right). We use identical noise samples for generating the before and after images.
Dataset
Algorithm
Time to unlearn (s) ↓
FID ↓
Time per step (s)
MNIST (Class: 1) Original FID: 20.7
H
2.181 (0.853)
23.4 (0.5) ✓
0.006 (0.0003)
S^
0.021 (0.003)
24.1 (0.5) ✓
0.008 (0.0006)
UN O^
0.061 (0.009)
22.1 (0.5) ✓
0.022 (0.0002)
UN O^ - S^
0.041 (0.021)
22.1 (0.4) ✓
0.017 (0.0067)
CelebA (Feature: Male) Original FID: 166.3
H
4.345 (1.849)
174.8 (1.9) ✓
0.016 (0.0001)
S^
0.181 (0.209)
176.4 (3.5) ✓
0.023 (0.0002) ∗
Appendix
Table 6: Performance of various algorithms for class/feature unlearning with VAEs on MNIST and CelebA when a classifier able to distinguish between the retain and forget data is available. Each experiment is repeated 10 times, and the standard deviations are shown in parentheses. Bold indicates the best score. ✗ indicates that the generated samples after unlearning are unrecognizably different from the original model. ✓ indicates the generated samples after unlearning are perceptually indistinguishable from the original model in terms of visual fidelity. An asterisk (*) indicates that, without classifier assistance, the algorithm failed to reach the target fraction of forget samples in the generated images within the allotted training steps.
Machine unlearning removes the influence of selected data or concepts from trained models while preserving useful behavior and avoiding full retraining.
Recent advances in one-step generative frameworks, such as flow map models, have significantly improved the efficiency of image generation by learning direct noise-to-data mappings in a single forward pass. However, machine unlearning for ensuring the safety of these powerful generators remains entirely unexplored. Existing diffusion unlearning methods are inherently incompatible with these one-step models, as they rely on a multi-step iterative denoising process. In this work, we propose UOT-Unlearn, a novel plug-and-play class unlearning framework for one-step generative models based on the Unbalanced Optimal Transport (UOT). Our method formulates unlearning as a principled trade-off between a forget cost, which suppresses the target class, and an f-divergence penalty, which preserves overall generation fidelity via relaxed marginal constraints. By leveraging UOT, our method enables the probability mass of the forgotten class to be smoothly redistributed to the remaining classes, rather than collapsing into low-quality or noise-like samples. Experimental results on CIFAR-10 and ImageNet-256 demonstrate that our framework achieves superior unlearning success (PUL) and retention quality (u-FID), significantly outperforming baselines.
Hyundo Choi, Junhyeong An, Jinseong Park +1
Sungkyunkwan University, Seoul, Republic of Korea · Korea Institute for Advanced Study, Seoul, Republic of Korea
Data unlearning aims to remove the influence of specific training samples from a trained model. In fine-tuning methods, data unlearning relies primarily on loss maximization over forget samples, which often leads to quality degradation or incomplete forgetting. Existing methods perform unlearning uniformly across diffusion stages, ignoring diffusion dynamics from noise to data. Our systematic study of diffusion phases shows that forgetting in diffusion models is uneven across time and frequency, with theoretical justification of distributive distortion and forgetting-utility trade-off. By selectively forgetting time and frequency in diffusion models, we achieve both higher unlearning success rates and improved generation quality across diverse settings, including both conditional and unconditional scenarios. We also introduce an improved SSCD metric that measures dissimilarity using a normalized perturbation distance. Together, we provide practical insights for understanding and improving data unlearning in diffusion models.
Jinseong Park, Mijung Park
Korea Institute for Advanced Study · University of British Columbia
Unlearning in diffusion models aims to remove undesirable data or concepts while preserving the utility of pretrained models -- two fundamentally conflicting objectives. We propose a principled constrained optimization framework that formulates unlearning as minimizing the deviation from a pretrained model, subject to explicit separation constraints from the unlearning distributions. Specifically, we formulate three constrained optimization problems based on reverse and forward KL divergences, and likelihood constraints. The first two generalize existing approaches for concept and data unlearning, while the third offers a novel and natural formulation for unlearning. Despite the nonconvexity of the KL constraints, we establish strong duality for all three problems, enabling us to explicitly characterize their optimal solutions as unlearning targets and develop primal-dual algorithms for each formulation. Experimental results demonstrate that our KL-constrained approach achieves superior retention-unlearning tradeoffs compared to weight-based baselines for concept and data unlearning, and that our likelihood-based approach matches unlearning effectiveness while better preserving retained concepts compared to baselines.