# Unlearning for One-Step Generative Models via Unbalanced Optimal Transport

arXiv: https://arxiv.org/abs/2603.16489 · HTML: https://mlsanity.com/paper/2603.16489/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Vision-language and text-to-image safety, one of 12 directions in LLM safety and security.
- Builds on 23 of its 29 references that are in the map.
- First posted on arXiv 17 Mar 2026.

## Placed in

- LLM safety and security / [Vision-language and text-to-image safety](https://mlsanity.com/safety/vision-language-and-text-to-image-safety/)

## Builds on (most cited first, 23 of 23)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Machine Unlearning](https://mlsanity.com/paper/1912.03817/) (2019)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Mean Flows for One-step Generative Modeling](https://mlsanity.com/paper/2505.13447/) (2025)
- [Erasing Concepts from Diffusion Models](https://mlsanity.com/paper/2303.07345/) (2023)
- [Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion](https://mlsanity.com/paper/2310.02279/) (2023)
- [SalUn: Empowering Machine Unlearning via Gradient-based Weight Saliency in Both Image Classification and Generation](https://mlsanity.com/paper/2310.12508/) (2023)
- [Ablating Concepts in Text-to-Image Diffusion Models](https://mlsanity.com/paper/2303.13516/) (2023)
- [A Survey of Machine Unlearning](https://mlsanity.com/paper/2209.02299/) (2022)
- [Forget-Me-Not: Learning to Forget in Text-to-Image Diffusion Models](https://mlsanity.com/paper/2303.17591/) (2023)
- [Unrolling SGD: Understanding Factors Influencing Machine Unlearning](https://mlsanity.com/paper/2109.13398/) (2021)
- [Selective Amnesia: A Continual Learning Approach to Forgetting in Deep Generative Models](https://mlsanity.com/paper/2305.10120/) (2023)
- [Optimal Entropy-Transport problems and a new Hellinger-Kantorovich distance between positive measures](https://mlsanity.com/paper/1508.07941/) (2015)
- [Robust Optimal Transport with Applications in Generative Modeling and Domain Adaptation](https://mlsanity.com/paper/2010.05862/) (2020)
- [Generative Modeling through the Semi-dual Formulation of Unbalanced Optimal Transport](https://mlsanity.com/paper/2305.14777/) (2023)
- [Adapt then Unlearn: Exploring Parameter Space Semantics for Unlearning in Generative Adversarial Networks](https://mlsanity.com/paper/2309.14054/) (2023)
- [Unlearning in Diffusion models under Data Constraints: A Variational Inference Approach](https://mlsanity.com/paper/2510.04058/) (2025)

## Built on by (most cited first, 2 of 2)

- [One-Step Generative Modeling via Unbalanced Optimal Transport](https://mlsanity.com/paper/2609.32708/) (2026)
- [Data Unlearning via Inverse Distillation](https://mlsanity.com/paper/2609.36099/) (2026)

## Abstract (arXiv)

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.