🤖 machine learning

UOTIP: Unbalanced Optimal Transport Map for Unpaired Inverse Problems

The paper proposes UOTIP, a novel framework for unpaired image inverse problems that leverages Unbalanced Optimal Transport with a likelihood-based cost function to achieve robust, state-of-the-art reconstruction performance across diverse noise scenarios while theoretically guaranteeing the existence and uniqueness of the transport map.

Original authors: Donggyu Lee, Taekyung Lee, Jaewoong Choi

Published 2026-05-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Donggyu Lee, Taekyung Lee, Jaewoong Choi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Mismatched Puzzle"

Imagine you are trying to restore a beautiful, clear painting (the clean signal) that has been splattered with mud and rain (the noisy measurement).

Usually, to teach a computer how to clean this painting, you would show it thousands of pairs: a muddy version on the left and the clean version on the right. The computer learns, "Oh, when I see this specific mud pattern, I should paint that specific flower."

But what if you don't have pairs?
What if you have a bucket of muddy paintings and a separate bucket of clean paintings, but they don't match up? You don't know which muddy painting came from which clean one. This is the "Unpaired Inverse Problem." It's like trying to match socks from two different laundry piles without knowing which sock belongs to which pair.

The Old Way: The "Rigid Matchmaker"

Previous methods tried to solve this using a concept called Optimal Transport (OT). Think of OT as a rigid matchmaker.

It says: "Every single muddy sock in the left pile must be paired with exactly one clean sock in the right pile. No exceptions."
The Flaw: In the real world, things aren't that neat.
- Different Noise Levels: Some muddy socks are just slightly damp; others are soaked in a swamp. A rigid matchmaker struggles to pair a "lightly damp" sock with a "swamp-soaked" sock if the rules say they must match 1-to-1 perfectly.
- Class Imbalance: Maybe your muddy pile has 100 socks, but your clean pile only has 50. The rigid matchmaker gets stuck trying to force 100 inputs into 50 outputs, creating a mess.

The New Solution: UOTIP (The "Flexible Matchmaker")

The authors propose a new method called UOTIP (Unbalanced Optimal Transport Map for Inverse Problems). Instead of a rigid matchmaker, UOTIP is a flexible, smart matchmaker who understands that life is messy.

Here is how it works, broken down into three simple concepts:

1. The "Likelihood" Rule (The Detective's Logic)

The paper introduces a special rule called a Likelihood Cost.

Analogy: Imagine you are a detective trying to guess what a suspect looked like before they got covered in mud. You don't just guess randomly; you ask, "If this person looked like this, would the mud have landed on them in this pattern?"
In the paper: The model learns to predict the clean image by asking, "Does this clean image, when I 'muddy' it mathematically, look like the noisy image I was given?" This ensures the solution is faithful to the actual data, not just a random guess.

2. The "Unbalanced" Flexibility (The Elastic Band)

This is the core innovation. Standard Optimal Transport is like a rubber band stretched tight between two points. If the points move, the band snaps or breaks.

UOTIP uses a special kind of elastic band (Unbalanced Optimal Transport).
The Benefit:
- Handling Noise: If you have a muddy sock that is very wet and a clean sock that is dry, the elastic band can stretch to connect them without breaking the rules. It allows the model to say, "Okay, this one muddy sample might correspond to two clean samples, or maybe just a fraction of one."
- Handling Imbalance: If you have 100 muddy socks and 50 clean socks, the elastic band can stretch to cover the gap. It doesn't force a perfect 1-to-1 match; it finds the best overall arrangement, even if the numbers don't match perfectly.

3. The "Safety Net" (The Quadratic Cost)

The paper mentions a theoretical problem: sometimes, the math gets "ill-posed," meaning there are infinite wrong answers that look right.

The Fix: The authors add a "safety net" called a Quadratic Cost.
Analogy: Imagine you are trying to find a lost key in a dark room. The "Likelihood" rule tells you where the key might be based on the sound it makes. But that could lead you in circles. The "Quadratic Cost" is like a gentle magnetic pull that keeps you from wandering too far away from the original spot. It ensures the computer finds one unique, stable answer instead of getting confused by infinite possibilities.

What Did They Prove?

The authors tested this "Flexible Matchmaker" on four different types of image restoration tasks:

Removing Blur (Deblurring): Making a blurry photo sharp.
Super-Resolution: Making a small, pixelated image large and clear.
HDR Reconstruction: Fixing images that are too bright or too dark.
Nonlinear Deblurring: Fixing complex, weird distortions.

The Results:

Better Quality: UOTIP produced clearer, sharper images than previous methods (like NOT, OTUR, and RCOT).
Robustness: It handled situations where the "mud" (noise) was different for every single image (multi-level noise) without getting confused.
Adaptability: It worked even when the number of muddy images didn't match the number of clean images (class imbalance).

Summary

Think of UOTIP as a new, super-smart restoration artist.

Old Artists needed a perfect photo of the "before" and "after" to learn. If the conditions changed (more mud, less mud), they failed.
UOTIP learns by understanding the relationship between mud and clean paint. It is flexible enough to handle messy, real-world data where the numbers don't match and the noise varies, all while guaranteeing it finds the one true, correct answer.

The paper claims this is the first time this specific "Unbalanced" approach has been successfully applied to these kinds of image problems, resulting in state-of-the-art performance.

Technical Summary: UOTIP – Unbalanced Optimal Transport Map for Unpaired Inverse Problems

Problem Statement
The paper addresses unpaired image inverse problems, a challenging setting where the goal is to recover an unknown clean signal $x$ from noisy measurements $y$ (modeled as $y = A(x) + n$ ), but the training data consists of independent, non-paired sets of noisy measurements and clean target signals. Traditional supervised methods require paired $(y, x)$ data, which is often unavailable in real-world scenarios. While generative models and Optimal Transport (OT) based methods have been proposed, standard OT formulations enforce strict marginal matching constraints ( $T_\#\mu = \nu$ ), which can be brittle when dealing with outliers, class imbalances between datasets, or varying noise levels. Furthermore, inverse problems are often ill-posed, meaning the mapping $y \to x$ may not have a unique solution, complicating the learning of a deterministic transport map.

Methodology
The authors propose UOTIP (Unbalanced Optimal Transport Map for Inverse Problems), a novel solver that formulates the reconstruction task as learning an Unbalanced Optimal Transport (UOT) map from the distribution of noisy measurements ( $\mu$ ) to the distribution of clean signals ( $\nu$ ).

UOT Formulation: Unlike standard OT, UOT relaxes the exact marginal matching constraint by introducing $f$ -divergence penalties ( $D_{\Psi_1}, D_{\Psi_2}$ ) on the marginals. This allows the transport plan to handle deviations in mass between the source and target distributions, providing robustness to outliers and class imbalances.
Likelihood-Based Cost Function: To align the transport map with the Maximum a Posteriori (MAP) estimate, the authors introduce a likelihood cost function:
$c_l(y, x) = \|A(x) - y\|_2^2$
Under Gaussian noise assumptions, this is proportional to the negative log-likelihood. Minimizing this cost ensures data fidelity (the reconstructed signal produces the observed measurement).
Quadratic Regularization for Well-Posedness: To address the ill-posed nature of inverse problems (where the likelihood cost alone may violate the "twist condition" required for the existence and uniqueness of an OT map), the authors incorporate a quadratic cost term:
$c_q(y, x) = \|y - x\|_2^2$
The total cost is a weighted sum: $c(y, x) = \tau(c_l(y, x) + c_q(y, x))$ . Theoretically, the authors prove that adding this quadratic term ensures the overall cost satisfies the twist condition, guaranteeing the existence and uniqueness of the transport map even for ill-posed problems (e.g., deblurring).
Neural Implementation: The model learns the UOT map $T_\theta$ and a potential function $v_\phi$ using a semi-dual formulation of UOT. The learning objective minimizes a loss derived from the convex conjugates of the penalty functions, effectively acting as a global, unpaired MAP estimator.

Key Contributions

First UOT-based Unpaired Solver: The paper introduces the first model for unpaired inverse problems based on Unbalanced Optimal Transport, specifically integrating a likelihood-based cost function.
Theoretical Guarantees: The authors provide a theoretical proof that incorporating a quadratic cost term ensures the existence and uniqueness of the transport map for ill-posed inverse problems by satisfying the twist condition.
Robustness to Real-World Constraints: By relaxing marginal constraints, the model demonstrates inherent robustness to:
- Multi-level observation noise: Handling scenarios where noise levels vary across samples.
- Class imbalance: Adapting to situations where the source and target datasets have different proportions of classes (e.g., different modes).
- Diverse noise types: Generalizing to non-Gaussian noise distributions (Laplace, Poisson) despite being trained with a Gaussian-based likelihood cost.
State-of-the-Art Performance: The model achieves superior performance on both linear (Gaussian deblurring, Super-resolution) and nonlinear (HDR reconstruction, Nonlinear deblurring) inverse problem benchmarks.

Experimental Results
The authors evaluated UOTIP on the FFHQ and AFHQ datasets across four tasks.

Quantitative Performance: UOTIP consistently achieved the best or second-best scores in PSNR, SSIM, and FID across all tasks and datasets, outperforming existing baselines such as NOT (Neural OT), OTUR (GAN-based OT), and RCOT.
Multi-Level Noise: In experiments with mixed noise levels ( $\sigma_y \in \{0.025, 0.05, 0.1, 0.2\}$ ), UOTIP significantly outperformed standard OT variants and other baselines, demonstrating that the unbalanced formulation effectively handles heterogeneous noise distributions.
Class Imbalance: When source and target distributions had imbalanced class ratios (e.g., cat vs. dog images), UOTIP maintained superior performance compared to baselines, which often degraded as the imbalance increased.
Ablation Studies: Removing the likelihood cost (using only quadratic cost) still yielded competitive results, suggesting utility as a blind solver. However, the full model (likelihood + quadratic) achieved the best overall performance. The study also showed that the model generalizes to Laplace and Poisson noise types even when trained with a Gaussian likelihood cost.

Significance and Claims
The paper claims that UOTIP offers a significant advancement for unpaired inverse problems by bridging the gap between optimal transport theory and practical inverse problem solving. The significance lies in:

Relaxing Strict Constraints: The UOT framework's ability to relax marginal matching makes the solver adaptable to real-world data where strict alignment between noisy and clean datasets is rarely achievable.
Theoretical Soundness: The work provides a rigorous theoretical foundation for solving ill-posed inverse problems via neural transport maps, ensuring solution uniqueness through the twist condition.
Practical Robustness: The model's ability to handle multi-level noise and class imbalance without retraining or complex architectural changes makes it particularly effective for real-world applications where data conditions are variable and imperfect.

The authors note a limitation: the method currently relies on fixed-form cost functions without training capacity for the cost itself, though the Gaussian-based likelihood showed promising generalization to other noise types.