Unlearning for One-Step Generative Models via Unbalanced Optimal Transport

• URL: http://arxiv.org/abs/2603.16489v1
• Date: Tue, 17 Mar 2026 13:19:04 GMT
• Title: Unlearning for One-Step Generative Models via Unbalanced Optimal Transport
• Abstract summary: UOT-Unlearn is a novel plug-and-play class unlearning framework for one-step generative models.<n>Our method formulates unlearning as a principled trade-off between a forget cost, which suppresses the target class, and an $f$-divergence penalty.
• Score: 11.695528552671371
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Coupling Models for One-Step Discrete Generation [50.37847093567075]
  Coupling Models is a one-step discrete generative model that learns a direct coupling between discrete sequences and Gaussian latents.<n>It avoids complex continuous flows over the simplex and hand-specified data-to-noise couplings.<n>It reduces LM1B text-generation perplexity by 33% at its lowest-perplexity operating point.
  arXiv  Detail & Related papers  (2026-05-08T03:40:39Z)
• Toward Reliable Machine Unlearning: Theory, Algorithms, and Evaluation [1.7767466724342065]
  We introduce Adrial Machine UNlearning (AMUN), which surpasses prior state-of-the-art methods for image classification based on SOTA MIA scores.<n>We show that existing methods fail in replicating a retrained model's behavior by introducing a nearest-neighbor membership inference attack (MIA-NN)<n>We then propose a fine-tuning objective that mitigates this leakage by approximating, for forget-class inputs, the distribution over remaining classes that a model retrained from scratch would produce.
  arXiv  Detail & Related papers  (2025-12-07T20:57:25Z)
• TwinFlow: Realizing One-step Generation on Large Models with Self-adversarial Flows [25.487712175353035]
  We propose TwinFlow, a framework for training 1-step generative models.<n>Our method achieves a GenEval score of 0.83 in 1-NFE on text-to-image tasks.<n>Our approach matches the performance of the original 100-NFE model on GenEval and DPG-Bench benchmarks.
  arXiv  Detail & Related papers  (2025-12-03T07:45:46Z)
• Adversarial Flow Models [26.917627135225118]
  We present adversarial flow models, a class of generative models that unifies adversarial models and flow models.<n>Our method supports native one-step or multi-step generation and is trained using the adversarial objective.<n>We show the possibility of end-to-end training of 56-layer and 112-layer models through depth repetition without any intermediate supervision.
  arXiv  Detail & Related papers  (2025-11-27T14:04:08Z)
• Variational Diffusion Unlearning: A Variational Inference Framework for Unlearning in Diffusion Models under Data Constraints [9.885531514020437]
  We propose a machine unlearning methodology that can prevent the generation of outputs containing undesired features from a pre-trained diffusion model.<n>Our approach is inspired by the variational inference framework with the objective of minimizing a loss function consisting of two terms: plasticity inducer and stability regularizer.<n>We validate the effectiveness of our method through comprehensive experiments for both class unlearning and feature unlearning.
  arXiv  Detail & Related papers  (2025-10-05T06:39:30Z)
• Distillation Robustifies Unlearning [36.27570321651185]
  We show that training a model can drastically modify its input-output behavior while leaving its underlying capabilities intact.<n>We propose Unlearn-Noise-Distill-on-Outputs (UNDO), a scalable method that distills an unlearned model into a noised copy of itself.
  arXiv  Detail & Related papers  (2025-06-06T17:58:54Z)
• One-Step Diffusion Model for Image Motion-Deblurring [85.76149042561507]
  We propose a one-step diffusion model for deblurring (OSDD), a novel framework that reduces the denoising process to a single step.<n>To tackle fidelity loss in diffusion models, we introduce an enhanced variational autoencoder (eVAE), which improves structural restoration.<n>Our method achieves strong performance on both full and no-reference metrics.
  arXiv  Detail & Related papers  (2025-03-09T09:39:57Z)
• Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think [72.48325960659822]
  One main bottleneck in training large-scale diffusion models for generation lies in effectively learning these representations.<n>We study this by introducing a straightforward regularization called REPresentation Alignment (REPA), which aligns the projections of noisy input hidden states in denoising networks with clean image representations obtained from external, pretrained visual encoders.<n>The results are striking: our simple strategy yields significant improvements in both training efficiency and generation quality when applied to popular diffusion and flow-based transformers, such as DiTs and SiTs.
  arXiv  Detail & Related papers  (2024-10-09T14:34:53Z)
• Advancing the Robustness of Large Language Models through Self-Denoised Smoothing [50.54276872204319]
  Large language models (LLMs) have achieved significant success, but their vulnerability to adversarial perturbations has raised considerable concerns. We propose to leverage the multitasking nature of LLMs to first denoise the noisy inputs and then to make predictions based on these denoised versions. Unlike previous denoised smoothing techniques in computer vision, which require training a separate model to enhance the robustness of LLMs, our method offers significantly better efficiency and flexibility.
  arXiv  Detail & Related papers  (2024-04-18T15:47:00Z)
• One-Step Diffusion Distillation via Deep Equilibrium Models [64.11782639697883]
  We introduce a simple yet effective means of distilling diffusion models directly from initial noise to the resulting image. Our method enables fully offline training with just noise/image pairs from the diffusion model. We demonstrate that the DEQ architecture is crucial to this capability, as GET matches a $5times$ larger ViT in terms of FID scores.
  arXiv  Detail & Related papers  (2023-12-12T07:28:40Z)
• RanPAC: Random Projections and Pre-trained Models for Continual Learning [59.07316955610658]
  Continual learning (CL) aims to learn different tasks (such as classification) in a non-stationary data stream without forgetting old ones. We propose a concise and effective approach for CL with pre-trained models.
  arXiv  Detail & Related papers  (2023-07-05T12:49:02Z)
• BOOT: Data-free Distillation of Denoising Diffusion Models with Bootstrapping [64.54271680071373]
  Diffusion models have demonstrated excellent potential for generating diverse images. Knowledge distillation has been recently proposed as a remedy that can reduce the number of inference steps to one or a few. We present a novel technique called BOOT, that overcomes limitations with an efficient data-free distillation algorithm.
  arXiv  Detail & Related papers  (2023-06-08T20:30:55Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.