Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Unlearning for One-Step Generative Models via Unbalanced Optimal Transport

Hyundo Choi, Junhyeong An, Jinseong Park, Jaewoong Choi

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.

Unlearning for One-Step Generative Models via Unbalanced Optimal Transport

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 -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.
Paper Structure (36 sections, 27 equations, 10 figures, 2 tables, 1 algorithm)

This paper contains 36 sections, 27 equations, 10 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. One-Step Generative Models via Probability Flow
  4. Unbalanced Optimal Transport (UOT)
  5. Proposed Method
  6. Problem Formulation
  7. Unlearning via Unbalanced Optimal Transport
  8. Cost Design for Unlearning
  9. Related Works
  10. Experiments
  11. 2D Synthetic Data
  12. Image Unlearning Benchmarks
  13. Experimental Setup.
  14. Evaluation Metrics.
  15. Experimental Results.
  16. ...and 21 more sections

Figures (10)

  • 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.
  • Figure 3: Step-wise unlearning trajectories on CIFAR-10. We analyze the trade-off between concept erasure (PUL $\uparrow$) and generative fidelity (u-FID $\downarrow$) 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 $\lambda$ (left) and the semantic distance margin $m$ (right). The star marker ($\star$) denotes our optimal configuration, which effectively balances the concept erasure efficacy (PUL $\uparrow$) and the generative fidelity (u-FID $\downarrow$).
  • ...and 5 more figures