Unlearning for One-Step Generative Models via Unbalanced Optimal Transport

This paper introduces UOT-Unlearn, a novel plug-and-play framework that leverages Unbalanced Optimal Transport to enable effective class unlearning in one-step generative models by smoothly redistributing the probability mass of forgotten classes while preserving overall generation fidelity.

The Gist

🎨 The Big Picture: The "Instant Artist" Problem

Imagine you have a magical artist who can paint a masterpiece in one single brushstroke.

• Old Artists (Diffusion Models): These artists take 50 to 100 tiny, slow strokes to turn a blank canvas into a picture. They are great, but they are slow.
• New Artists (One-Step Models): These are the new "Instant Artists." They look at a blank canvas and poof!—a perfect image appears instantly. They are incredibly fast and powerful.

The Problem:
Sometimes, these Instant Artists learn to paint things they shouldn't. Maybe they were trained on copyrighted art, or they accidentally learned to draw "NSFW" (Not Safe For Work) content.

• The Old Solution: To fix an old artist, you could tell them, "Don't paint that specific thing," and they would slowly adjust their technique over many strokes.
• The New Problem: You can't do that with an Instant Artist. They don't have "middle steps" to tweak. They just go from "Blank" to "Done" in one jump. If you try to tell them to change their mind halfway through, they can't, because there is no "halfway."

If you try to force them to forget using old methods, the whole painting usually turns into a blurry, ugly mess.

🧠 The Solution: "UOT-Unlearn" (The Smart Redistributor)

The authors of this paper created a new method called UOT-Unlearn. Think of it as a traffic controller for the artist's imagination.

Instead of trying to "erase" a specific type of image (which is like trying to delete a file from a hard drive), they use a concept called Unbalanced Optimal Transport (UOT).

The Analogy: The Moving Company

Imagine the artist's brain is a warehouse full of different types of furniture (images):

• Furniture A: Cars (The "Forget" class we want to remove).
• Furniture B: Dogs, Trees, and Houses (The "Keep" classes).

The Old Way (Gradient Ascent):
You try to throw the "Cars" out the back door. But because you are throwing them out so hard, you accidentally knock over the "Dogs" and "Trees," and the whole warehouse collapses into a pile of junk.

The New Way (UOT-Unlearn):
Instead of throwing things out, you use a Smart Moving Truck (the UOT framework).

1. The Goal: You want to get rid of the "Cars."
2. The Trick: You tell the truck, "It costs a huge penalty (a lot of money) to park a Car in the warehouse."
3. The Result: The truck looks at the "Cars" and says, "Okay, I can't park these here because it's too expensive."
4. The Magic: Instead of just deleting the cars, the truck smoothly moves the "Car" space into the empty spots next to the "Dogs" and "Trees."

Because the truck is "unbalanced" (it allows for some flexibility), it doesn't force the cars to vanish into thin air. Instead, it gently reshapes the warehouse so the "Car" area becomes part of the "Dog" or "Tree" area. The result? The warehouse is still full of beautiful, organized furniture, but there are no cars left.

🔑 How It Works (The 3 Simple Steps)

1. Identify the "Bad" Stuff:
   The system takes a few examples of the thing we want to forget (e.g., pictures of "Goldfish") and creates a mental anchor (a target point in the brain).

2. Set the "Expensive" Rule:
   The system tells the model: "If you try to make a Goldfish, it will cost you a lot of 'energy' (penalty)."

3. The Smooth Shift:
   The model realizes, "Okay, making Goldfish is too expensive." So, it takes the energy it would have used to make Goldfish and redistributes it to make better Dogs or Trees instead.

   • Crucially: It does this without needing to see the original "Good" pictures again. It just uses the pictures it generates itself to figure out where to move the energy.

🏆 Why Is This Better?

The paper tested this on famous datasets like CIFAR-10 (small toy images) and ImageNet (real photos).

• Other Methods: When they tried to remove "Goldfish," the other methods made the "Dogs" look like melted blobs. The quality dropped drastically.
• UOT-Unlearn: It removed the "Goldfish" almost completely (90%+ success), but the "Dogs" and "Trees" still looked crisp and high-quality.

The Metaphor:

• Other methods are like using a sledgehammer to remove a weed; you break the flowerbed.
• UOT-Unlearn is like a gardener who gently pulls the weed out and fills the hole with soil so the surrounding flowers grow even better.

🚀 The Takeaway

This paper solves a major safety problem for the next generation of super-fast AI art generators. It proves that you can teach an AI to "forget" dangerous or unwanted concepts without breaking its ability to create beautiful art, and it does it all in a single, lightning-fast step.

In short: They figured out how to tell a super-fast AI, "Don't draw that," without making the AI forget how to draw anything else.

Technical Summary

1. Problem Statement

Context: One-step generative models (e.g., Consistency Models, Flow Matching) have revolutionized image generation by mapping noise to data in a single forward pass, offering near-diffusion quality with significantly faster inference speeds.
The Gap: As these models become faster and more powerful, the risk of generating harmful content (NSFW, copyrighted material) increases. While "machine unlearning" (removing specific concepts from a model without retraining) is well-studied for multi-step diffusion models, it remains unexplored for one-step models.
The Challenge: Existing unlearning methods for diffusion models rely on iterative denoising processes (modifying gradients or predictions at intermediate timesteps). These methods are incompatible with one-step architectures because they lack intermediate steps; the generation happens in a single mapping $G_\theta(x_0)$. Furthermore, standard unlearning often leads to "catastrophic forgetting" or severe degradation of the remaining data distribution.

2. Methodology: UOT-Unlearn

The authors propose UOT-Unlearn, a plug-and-play framework that formulates class unlearning as an Unbalanced Optimal Transport (UOT) problem.

Core Concept: Unbalanced Optimal Transport (UOT)

Unlike standard Optimal Transport (OT), which enforces strict marginal constraints (mass must be conserved exactly), UOT relaxes these constraints. It minimizes a trade-off between:

1. Transport Cost: The cost of moving probability mass from the source to the target.
2. Divergence Penalty: An $f$-divergence penalty (e.g., KL-divergence) that allows for marginal mismatch.

In the context of unlearning, this flexibility allows the model to "remove" the probability mass of the forget class by redistributing it to other valid classes, rather than forcing it into a rigid, potentially low-quality state.

The Unlearning Objective

The framework treats the unlearning process as finding a transport map $\Delta T_\theta$ that transforms the pre-trained distribution ($p_{pre}$) into an unlearned distribution ($\pi_1$).

• Source: Pre-trained distribution ($p_{pre}$).
• Target: The full data distribution ($p_{data}$), approximated by $p_{pre}$ during optimization to avoid needing real retain data.
• Cost Function ($c_{ul}$): A custom cost function is designed to penalize the generation of the "forget" concept:
  • Forget Cost: If a generated sample falls within the "forget region" (defined by a feature-space anchor $\mu_f$ and margin $m$), a heavy hinge-like penalty is applied to push the features away from the forget anchor.
  • Retain Cost: For samples outside the forget region, a squared $L_2$ distance is imposed between the new output and the pre-trained output to preserve fidelity.

Optimization

The method uses a semi-dual formulation of UOT involving a neural transport map ($G_\theta$) and a dual potential ($v_\phi$).

• Data-Free: Crucially, the optimization requires zero real retain data. It relies solely on generated samples from the pre-trained model and a pre-computed forget centroid.
• Algorithm: The process iteratively updates the dual potential to maximize the lower bound and updates the generator to minimize the transport cost, effectively shifting probability mass away from the forget class and toward the retain classes.

3. Key Contributions

1. First Framework for One-Step Unlearning: Introduces the first plug-and-play unlearning framework specifically designed for one-step generative models, addressing a critical gap in the literature.
2. UOT-Based Formulation: Proposes a novel objective based on Unbalanced Optimal Transport that smoothly redistributes the probability mass of the forgotten class to remaining classes via $f$-divergence penalties, preventing the collapse into noise.
3. Data-Free Optimization: The method operates without access to real retain datasets during the unlearning phase, relying only on synthetic samples and a pre-computed anchor, making it highly efficient and privacy-preserving.
4. Superior Trade-off: Demonstrates a significantly better balance between Unlearning Success (erasing the target) and Retention Quality (preserving the rest of the distribution) compared to adapted baselines.

4. Experimental Results

The method was evaluated on CIFAR-10 and ImageNet-256 using Consistency Trajectory Models (CTM) and MeanFlow architectures.

• Metrics:
  • PUL (Percentage of Unlearning): Measures the reduction in frequency of the forget class.
  • u-FID (Unlearned FID): Measures the quality of the remaining classes compared to real data.
• Performance:
  • CIFAR-10: UOT-Unlearn achieved the highest PUL scores (e.g., ~90% for Class 1) while maintaining the lowest u-FID (e.g., ~5.11 for Class 6), significantly outperforming baselines like Gradient Ascent (GA), Selective Amnesia (SA), and Variational Diffusion Unlearning (VDU).
  • ImageNet-256: In high-resolution tasks, baselines like GA caused severe structural corruption (u-FID ~79.89) when erasing a class. UOT-Unlearn achieved robust erasure (85.08% PUL) while keeping u-FID low (20.16), preserving the structural integrity of the retain classes.
• Qualitative Analysis: Visualizations on 2D synthetic data and real images show that UOT-Unlearn smoothly redistributes the "forget" mode to valid "retain" modes, whereas baselines often distort the distribution or leave residual artifacts.

5. Significance and Impact

• Safety for Fast Generators: As one-step models become the standard for real-time AI applications, this work provides an essential safety mechanism to remove harmful concepts without the prohibitive cost of full retraining.
• Theoretical Insight: It reframes machine unlearning not as a parameter-space update or a denoising step modification, but as a constrained probability transport problem. This offers a principled way to balance removal strength with distributional consistency.
• Efficiency: By eliminating the need for real retain data and leveraging the single-pass nature of the models, the method is computationally efficient and scalable to large datasets.

In conclusion, UOT-Unlearn successfully bridges the gap between the speed of one-step generative models and the safety requirements of machine unlearning, offering a robust, data-free, and theoretically grounded solution.