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
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.
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.
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.
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 ↑ ) and generative fidelity (u-FID ↓ ) 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.
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.
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 λ (left) and the semantic distance margin m (right). The star marker ( ⋆ ) denotes our optimal configuration, which effectively balances the concept erasure efficacy (PUL ↑ ) and the generative fidelity (u-FID ↓ ).
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
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.
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.
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.
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.
Figure 10: Unlearning results on the ImageNet 256 × 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 removes the influence of selected data or concepts from trained models while preserving useful behavior and avoiding full retraining.
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.
Multi-step matching models, including flow and diffusion models, produce high-quality outputs but incur substantial inference costs and may reproduce unwanted components of their training datasets. We introduce Inverse Distillation Unlearning (IDU), a unified framework that simultaneously distills a teacher multi-step matching model into an efficient one-step student generator and suppresses outputs corresponding to a designated training subset. We first formulate distillation as a min-max objective over a data distribution and then represent this distribution as a mixture of the forget-set and the generated distributions. This allows us to compare this mixture with the teacher's training distribution and recover only the retained data at the optimum. Our method requires only a pretrained full-data teacher and data from the forget set, without access to retained training examples, extra feature extractors or classifiers. Extensive experiments on MNIST and CIFAR-10 datasets under flow-matching and score-based diffusion settings demonstrate that IDU substantially reduces the generation frequency of forgotten classes while preserving generation quality on the retained classes. To the best of our knowledge, IDU is the first unified framework for simultaneous unlearning and distillation in unconditional flow-matching and score-based models.
Aleksei Leonov, Nikita Kornilov, Zhenhe Zhang +3
AI Foundation lab, Moscow, Russia · MIRAI, Moscow, Russia · Applied AI Institute, Moscow, Russia +2
Drifting models enable one-step generation by amortizing distribution transport into training, but this efficiency places greater demands on the transport field estimated at each update. In large-scale training, the field is computed from finite mini-batches of generated and real samples, which provide only imperfect approximations to the underlying distributions. Balanced optimal transport enforces exact mass matching within every mini-batch, making the estimated field sensitive to the particular composition of the real-data batch. We find that generated and real samples should be treated asymmetrically: letting the mass assigned to real samples adapt while keeping every generated sample fully transported improves generation across six feature-space metrics in controlled ablations, and is more robust to the relaxation strength than relaxing both marginals simultaneously, which falls below balanced transport under stronger relaxation. Motivated by this observation, we propose Unbalanced Optimal Transport Gradient Flow (UOT-GF), which keeps the generated-sample marginal fixed and relaxes only the real-data marginal. Under identical settings at DiT-B/2 on ImageNet-256, UOT-GF improves Fréchet Inception Distance (FID) from 1.53 to 1.46 over the balanced W-Flow baseline; scaling the same recipe yields 1.34 and 1.22 FID at L/2 and XL/2, the best FID among the one-step models we compare. We further derive the induced UOT transport force, establish a kinetic Vlasov--Fokker--Planck formulation whose overdamped zero-temperature limit recovers the drifting dynamics, and characterize non-target stationary states together with sufficient conditions for convergence.
Yirong Shen, Mengfei Xia, Junpeng Jing +2
Imperial College London · Ant Group · Brunel University of London