Redistributive Unlearning for One-Step Generators via Unbalanced Optimal Transport

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 these powerful generators remains entirely unexplored, as existing diffusion unlearning methods rely on multi-step denoising and are inherently incompatible with one-step models. Moreover, prior methods focus only on suppressing the target concept, often leaving the unlearned probability mass to collapse into noise-like samples. In this work, we propose *UOT-Unlearn*, a redistributive 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 confines the generated distribution to the data support via relaxed marginal constraints. As a result, the probability mass of the unlearned class is smoothly redistributed to the remaining classes rather than discarded, preserving overall generation fidelity. Experiments on CIFAR-10 and ImageNet-256 show that UOT-Unlearn achieves superior unlearning performance while maintaining generation quality, and it further extends to one-step text-to-image generation.