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

Source: [Unlearning for One-Step Generative Models via Unbalanced Optimal Transport](<https://www.sota2.com/research/paper/69ba318c0636e5bf4afc0820>)

Collection index: [llms.txt](<https://www.sota2.com/research/llms.txt>)

Data updated: 2026-03-19T05:00:05.030Z

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

Published: 2026-03-17T13:19:04.000Z

## 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.

## Recorded benchmark results

A research result applies to its stated dataset, split, metric and evaluation configuration. Compare like-for-like settings, preserve the metric's unit and direction, and distinguish paper-reported results from independent reproduction. The best result recorded here is not a claim of exhaustive, current global coverage.

- [Class Unlearning on CIFAR-10 (Fidelity Metrics)](<https://www.sota2.com/research/sota/class-unlearning-on-cifar-10-fidelity-metrics.md>): Class Unlearning; CIFAR-10; PUL: 96.25.

## Original sources and implementation

[Original paper](<https://arxiv.org/abs/2603.16489>)

[PDF](<https://arxiv.org/pdf/2603.16489.pdf>)