Unlearning for One-Step Generative Models via Unbalanced Optimal Transport

Hyundo Choi; Junhyeong An; Jinseong Park; Jaewoong Choi

cs.CVMar 17, 2026

Unlearning for One-Step Generative Models via Unbalanced Optimal Transport

Authors: Hyundo Choi, Junhyeong An, Jinseong Park, Jaewoong Choi

Organizations: Sungkyunkwan University, Seoul, Republic of Korea · Korea Institute for Advanced Study, Seoul, Republic of Korea

Abstract

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.

Figures & tables

Method	Optimization-phase Data
Method	Real Forget	Real Retain	Generated data
Gradient Ascent (GA)	Required	None	None
VDU	Required	None	None
Selective Amnesia	None	None	Required
SalUn	Required	Required	Required ∗
UOT-Unlearn (Ours)	None †	None	Required

Table 1: Comparison of data requirements during the unlearning optimization phase . Our UOT-based framework is the only method that operates with zero real data , relying solely on synthetic samples and a pre-computed centroid.

Panel of Figure 1 — Figure 1: Unlearning results on a 2D toy dataset where the forget mode is located at $(0,1)$ . (a) Pretrained one-step generator. (b) VDU leads to overall distribution distortion. (c) Our method redistributes the forget mode to the remaining modes.

Panel of Figure 2 — Figure 2: Qualitative unlearning results on CIFAR-10 . Unconditional samples of target classes (1, 6, and 8) generated by CTM (top) and MF (bottom) architectures. To clearly illustrate the semantic erasure of targeted concepts, the Unlearned outputs are generated using the exact same initial noise seeds as their Pretrained counterparts.

Unlearned	Original	GA		SA		SalUn		VDU		UOT-Unlearn (Ours)
Class	FID $\downarrow$	PUL (%) $\uparrow$	u-FID $\downarrow$	PUL (%) $\uparrow$	u-FID $\downarrow$	PUL (%) $\uparrow$	u-FID $\downarrow$	PUL (%) $\uparrow$	u-FID $\downarrow$	PUL (%) $\uparrow$	u-FID $\downarrow$
Target Model: Consistency Trajectory Models (CTM) [ 13 ]
Class 1 (Auto)	4.53	88.07	160.03	91.16	49.80	90.64	41.36	53.86	71.98	80.32	9.90
Class 6 (Frog)	5.02	39.16	72.17	61.85	42.54	43.16	39.00	51.25	47.11	90.98	5.11
Class 8 (Ship)	4.36	95.40	208.80	64.50	62.48	48.65	45.70	48.43	52.07	85.23	5.88
Average	4.64	74.21	147.00	72.50	51.61	60.82	42.02	51.18	57.05	85.51	6.96

Table 2: Unlearning performance on CTM and Meanflow. Best results are highlighted in bold , and second-best results are underlined . UOT-Unlearn consistently achieves the highest Percentage of Unlearning (PUL) while preserving the original FID.

Figure 3: Step-wise unlearning trajectories on CIFAR-10. We analyze the trade-off between concept erasure (PUL

\uparrow

) and generative fidelity (u-FID

\downarrow

) across unconditional CTM (top row) and MF (bottom row) models. Our framework achieves a superior trade-off, erasing target class while preserving distributional consistency. Metrics are computed sequentially using fixed initial noise seeds for fair variance analysis.

Panel of Figure 4 — Figure 4: Unlearning performance on ImageNet-256 (Goldfish). Samples for the Forget and Retain classes. GA suffers from severe structural corruption (u-FID: 79.89) to achieve concept suppression. In contrast, UOT-Unlearn achieves robust erasure (85.08% PUL) while preserving generative fidelity (u-FID: 20.16) relative to the baseline (FID: 11.57). u-FID is computed over 36 aquatic classes.

Panel of Figure 5 — Figure 5: Ablation study of key hyperparameters for our UOT framework, conducted on CTM model for CIFAR-10 with target class 8. We independently vary the forget loss weight $\lambda$ (left) and the semantic distance margin $m$ (right). The star marker ( $\star$ ) denotes our optimal configuration, which effectively balances the concept erasure efficacy (PUL $\uparrow$ ) and the generative fidelity (u-FID $\downarrow$ ).

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 6 — Figure 6: Qualitative comparison of unlearning methods on the Ship class (Class 8) using the Consistency Trajectory Model (CTM). Generated under fixed seed setting, this figure illustrates how the target class is erased across different approaches. Notably, unlike other baselines generate corrupted images, UOT-Unlearn demonstrates a distinct tendency to transition the forgotten concept into features of other classes while preserving the overall spatial layout.

Panel of Figure 7 — Figure 7: Visual comparison of images generated from the pretrained CTM versus our UOT-Unlearn-applied CTM trained to forget the Ship class (Class 8). The results demonstrate that UOT-Unlearn successfully erases the forget class, seamlessly transitioning it into features of other classes, while strictly preserving the visual quality and structural layout of the remaining retain classes.

Panel of Figure 8 — Figure 8: Qualitative comparison of unlearning methods on the Ship class (Class 8) using the Meanflow (MF). Generated under fixed seed setting, this figure illustrates how the target class is erased across different approaches.

Panel of Figure 9 — Figure 9: Visual comparison of randomly generated images from the Pretrained model and our UOT-Unlearn (MF) trained to forget the Ship class (Class 8), using the same initial noises.

Panel of Figure 10 — Figure 10: Unlearning results on the ImageNet 256 $\times$ 256 dataset using the MeanFlow (MF) model. The top row of each generated batch displays the Goldfish class (forget class), while the subsequent bottom rows show 9 other aquatic animal classes (retain classes) using the exact same initial noises.

Machine Unlearning

Machine unlearning removes the influence of selected data or concepts from trained models while preserving useful behavior and avoiding full retraining.

See all 343 papers

UNO: Unlearning via Orthogonalization in Generative models
Jun 5, 2025Pinak Mandal, Georg A. GottwaldMachine Unlearning Diffusion Model Unlearning
University of Sydney, Australia
Data Unlearning via Inverse Distillation
Sep 28, 2026Aleksei Leonov, Nikita Kornilov, Zhenhe Zhang +3Machine Unlearning Diffusion Model Unlearning
AI Foundation lab, Moscow, Russia · MIRAI, Moscow, Russia · Applied AI Institute, Moscow, Russia +2
One-Step Generative Modeling via Unbalanced Optimal Transport
Sep 26, 2026Yirong Shen, Mengfei Xia, Junpeng Jing +2One-Step Generative Modeling Optimal Transport
Imperial College London · Ant Group · Brunel University of London

Jun 5, 2025J/K move · Enter open · S save

UNO: Unlearning via Orthogonalization in Generative models

Pinak Mandal, Georg A. Gottwald

University of Sydney, Australia

Machine Unlearning Diffusion Model Unlearning

Open paper AskPDF arXiv

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.

Abstract

Machine Unlearning

Explore similar work

UNO: Unlearning via Orthogonalization in Generative models

Data Unlearning via Inverse Distillation

One-Step Generative Modeling via Unbalanced Optimal Transport