UOTIP: Unbalanced Optimal Transport Map for Unpaired Inverse Problems

Abstract: We investigate unpaired image inverse problems, a challenging setting where only independent, non-paired sets of noisy measurements and clean target signals are available for training. We propose a novel inverse problem solver based on Unbalanced Optimal Transport, called Unbalanced Optimal Transport Map for Inverse Problems (UOTIP). Our method formulates the reconstruction task, predicting clean target signals from noisy measurements, as learning a UOT Map from noisy measurement distribution to clean signal distribution by incorporating a likelihood-based cost function. By relaxing the exact marginal constraint, the UOT framework provides key advantages to our model: robustness to multi-level observation noise, adaptability to class imbalance between noisy and clean datasets, and generalizability to diverse noise-type scenarios. Furthermore, we theoretically demonstrate that incorporating a quadratic cost term ensures the existence and uniqueness of the transport map by satisfying the twist condition, even for ill-posed inverse problems. Our experiments demonstrate that UOTIP achieves state-of-the-art performance on unpaired image inverse problem benchmarks, across linear and nonlinear inverse problems.

• The paper introduces UOTIP, a likelihood-based unbalanced optimal transport map that performs MAP-style image reconstruction without paired measurements and targets.
• The paper proves that adding a quadratic cost to the likelihood cost restores the transport-map twist condition when the forward operator is Lipschitz, supporting existence and uniqueness even for ill-posed problems.
• The paper reports stronger PSNR, SSIM, and FID than transport and diffusion baselines across deblurring, super-resolution, HDR, and nonlinear reconstruction, with added robustness to mixed noise and class imbalance.

Problem setting and motivation

The paper addresses unpaired image inverse problems, where the training data consist of two independently sampled sets: noisy measurements $Y = \{y_i\} \sim \mu$ and clean target signals $X = \{x_j\} \sim \nu$, with no correspondence between samples. The forward model is $\mathbf{y} = \mathcal{A}(\mathbf{x}) + \mathbf{n}$, where $\mathcal{A}$ may be nonlinear and the noise level $\sigma_y > 0$ is assumed unknown. Existing unpaired solvers such as OTUR (Rütte et al., 2022), NOT (Dadarlat et al., 2022), and RCOT (Erba et al., 2023) typically assume a single fixed noise level and enforce exact marginal matching. The authors propose UOTIP, which formulates reconstruction as learning an Unbalanced Optimal Transport (UOT) map from $\mu$ to $\nu$ under a likelihood-based cost function.

OT formulation of inverse problems as a global MAP estimator

The core observation is that a MAP estimate requires both prior fidelity (outputs consistent with $p(x)$) and data fidelity (high likelihood $\log p(y_0 \mid \cdot)$). The Kantorovich OT framework naturally encodes prior fidelity through the marginal constraint $T_\#\mu = \nu$. The paper introduces a likelihood cost

$X = \{x_j\} \sim \nu$0

which is proportional to the negative log-likelihood under Gaussian noise. Minimizing expected transport cost under this cost is shown to be equivalent to minimizing the negative log-posterior over couplings in $X = \{x_j\} \sim \nu$1, since the marginal constraints fix the $X = \{x_j\} \sim \nu$2 term. The Neural OT framework thus acts as a "global, unpaired MAP estimator."

Well-posedness via the quadratic cost

Because inverse problems are ill-posed, the likelihood cost alone violates the left twist condition required for existence and uniqueness of an optimal transport map (via Theorem 1 of Fathi–Figalli). Specifically, with only $X = \{x_j\} \sim \nu$3, the map $X = \{x_j\} \sim \nu$4 is not injective when $X = \{x_j\} \sim \nu$5 is non-invertible. The paper proves that adding a quadratic cost $X = \{x_j\} \sim \nu$6 restores injectivity:

Proposition: If $X = \{x_j\} \sim \nu$7 is $X = \{x_j\} \sim \nu$8-Lipschitz continuous, then $X = \{x_j\} \sim \nu$9 satisfies the twist condition whenever $\mathbf{y} = \mathcal{A}(\mathbf{x}) + \mathbf{n}$0, guaranteeing existence and uniqueness ($\mathbf{y} = \mathcal{A}(\mathbf{x}) + \mathbf{n}$1-a.e.) of the transport map even for ill-posed problems such as Gaussian deblurring or HDR reconstruction.

This is a notable theoretical contribution: it establishes well-posedness of the learned solver precisely where the underlying inverse problem is ill-posed. A practical caveat is that the Lipschitz constant $\mathbf{y} = \mathcal{A}(\mathbf{x}) + \mathbf{n}$2 is generally unknown; in implementation the role of $\mathbf{y} = \mathcal{A}(\mathbf{x}) + \mathbf{n}$3 is absorbed into a cost intensity hyperparameter $\mathbf{y} = \mathcal{A}(\mathbf{x}) + \mathbf{n}$4 rather than explicitly enforcing $\mathbf{y} = \mathcal{A}(\mathbf{x}) + \mathbf{n}$5, so the guarantee holds only implicitly in practice.

Learning objective

UOTIP builds on UOTM's semi-dual formulation of UOT (Duc, 2023), jointly parameterizing a potential network $\mathbf{y} = \mathcal{A}(\mathbf{x}) + \mathbf{n}$6 and transport map $\mathbf{y} = \mathcal{A}(\mathbf{x}) + \mathbf{n}$7 satisfying the optimality condition $\mathbf{y} = \mathcal{A}(\mathbf{x}) + \mathbf{n}$8. The objective uses KL-divergence marginal penalties ($\mathbf{y} = \mathcal{A}(\mathbf{x}) + \mathbf{n}$9), reducing to standard Neural OT when $\mathcal{A}$0 is the indicator at $\mathcal{A}$1. The overall cost is $\mathcal{A}$2.

Three advantages of the unbalanced formulation are argued: (a) relaxed marginals allow higher-likelihood regions under both prior and data fidelity; (b) sample reweighting via $\mathcal{A}$3 handles class imbalance—for instance, matching a source mode covering 20% of mass to a target mode covering 30%; and (c) improved training dynamics via gradient-norm bounding of the potential.

Empirical results

Experiments cover four tasks on FFHQ and AFHQ-dog at $\mathcal{A}$4: Gaussian deblurring, $\mathcal{A}$5 super-resolution, HDR reconstruction, and nonlinear deblurring (a pretrained neural operator), with AWGN at $\mathcal{A}$6. Representative results on FFHQ:

Task	Method	PSNR	SSIM	FID
Gaussian deblurring	OTUR	23.82	0.7106	24.337
	UOTIP	24.06	0.7139	21.210
Super-resolution 4×	OTUR	24.09	0.7243	22.751
	UOTIP	24.35	0.7371	19.475
HDR reconstruction	OTUR	25.32	0.8545	16.458
	UOTIP	26.02	0.8642	20.840
Nonlinear deblurring	OTUR	26.94	0.8594	12.538
	UOTIP	28.52	0.8841	11.370

UOTIP attains best scores on nearly all metrics across both datasets, with particularly large FID gains on AFHQ (e.g., 12.566 vs. 30.773 for OTUR on deblurring). Additional comparisons show UOTIP outperforming DPS across all four tasks by large margins (e.g., 28.52 vs. 19.01 PSNR on nonlinear deblurring); notably, DPS exhibits a large gap between test-set FID and full-dataset FID, which the authors attribute to memorization. Against KIDOT on linear tasks, UOTIP again dominates all metrics.

Stress tests enabled by unbalancedness

Under multi-level noise (mixture of $\mathcal{A}$7 in proportions 4:3:2:1), UOTIP beats both OTUR and its own balanced OT variant (OTIP) on all metrics—e.g., FID 65.664 vs. 91.309 (OTIP) on Gaussian deblurring—directly supporting the claim that robustness stems from the unbalanced formulation rather than other design choices. Under class imbalance (AFHQ cat:dog ratios up to 4:1 between target and source), UOTIP remains stable while OTIP's FID degrades sharply (e.g., 78.869 at ratio 4 vs. 39.894 for UOTIP). Under noise-type mismatch, the Gaussian likelihood cost generalizes to Laplace and Poisson noise, outperforming baselines; interestingly, a Poisson-specific likelihood variant performs worse than the Gaussian one due to numerical instability of the scaled quadratic approximation, indicating the framework is not tied to—but also does not automatically benefit from—the exact likelihood.

Cost ablation

Removing either cost term degrades performance modestly, but the quadratic-only variant (blind to $\mathcal{A}$8) still surpasses OTUR on linear tasks, suggesting potential as a blind solver when the operator preserves signal structure. For super-resolution, the quadratic cost must be modified via bicubic interpolation $\mathcal{A}$9, under which the twist condition no longer holds; the authors concede that existence is then not guaranteed and attribute empirical success to the generator's local smoothing inductive bias—an explicit gap between theory and practice.

Limitations and open questions

The paper acknowledges that the method is evaluated only with fixed-form cost functions without learned capacity, and that the Gaussian likelihood's generalization to other noise types, while demonstrated empirically, could be improved by broader cost designs. The condition $\sigma_y > 0$0 is never enforced explicitly; developing a principled adaptive strategy for the intensity parameter $\sigma_y > 0$1 is left open. The super-resolution setting lacks any theoretical existence guarantee. Finally, the blind (quadratic-only) variant's applicability beyond structure-preserving operators remains untested.

Conclusion

UOTIP casts unpaired inverse problems as neural UOT map estimation with a likelihood cost, yielding a globally defined MAP-style solver that requires neither paired data nor large-scale pretraining. Its main contributions are the likelihood-cost formulation, a twist-condition argument restoring map existence and uniqueness despite ill-posedness, and state-of-the-art results among direct transport methods—with distinct advantages under multi-level noise, class imbalance, and noise-type mismatch that are attributable specifically to the unbalanced relaxation.