Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Generative Modeling with Optimal Transport Maps

Litu Rout, Alexander Korotin, Evgeny Burnaev

TL;DR

The paper demonstrates that an optimal transport map can serve as a powerful generative model in high-dimensional ambient spaces, not just as a loss function. It introduces a min–max framework for learning Wasserstein-2 OT maps and extends the approach to unequal input/output dimensions via a Q-embedding, with accompanying error bounds. Empirically, the method yields competitive image generation and effective unpaired image restoration (denoising, colorization, inpainting) while maintaining training regimes comparable to GAN-based methods. This work advances practical OT-based mapping by enabling end-to-end learning directly in ambient spaces and provides rigorous guarantees linking duality gaps to map accuracy and distributional closeness.

Abstract

With the discovery of Wasserstein GANs, Optimal Transport (OT) has become a powerful tool for large-scale generative modeling tasks. In these tasks, OT cost is typically used as the loss for training GANs. In contrast to this approach, we show that the OT map itself can be used as a generative model, providing comparable performance. Previous analogous approaches consider OT maps as generative models only in the latent spaces due to their poor performance in the original high-dimensional ambient space. In contrast, we apply OT maps directly in the ambient space, e.g., a space of high-dimensional images. First, we derive a min-max optimization algorithm to efficiently compute OT maps for the quadratic cost (Wasserstein-2 distance). Next, we extend the approach to the case when the input and output distributions are located in the spaces of different dimensions and derive error bounds for the computed OT map. We evaluate the algorithm on image generation and unpaired image restoration tasks. In particular, we consider denoising, colorization, and inpainting, where the optimality of the restoration map is a desired attribute, since the output (restored) image is expected to be close to the input (degraded) one.

Generative Modeling with Optimal Transport Maps

TL;DR

The paper demonstrates that an optimal transport map can serve as a powerful generative model in high-dimensional ambient spaces, not just as a loss function. It introduces a min–max framework for learning Wasserstein-2 OT maps and extends the approach to unequal input/output dimensions via a Q-embedding, with accompanying error bounds. Empirically, the method yields competitive image generation and effective unpaired image restoration (denoising, colorization, inpainting) while maintaining training regimes comparable to GAN-based methods. This work advances practical OT-based mapping by enabling end-to-end learning directly in ambient spaces and provides rigorous guarantees linking duality gaps to map accuracy and distributional closeness.

Abstract

With the discovery of Wasserstein GANs, Optimal Transport (OT) has become a powerful tool for large-scale generative modeling tasks. In these tasks, OT cost is typically used as the loss for training GANs. In contrast to this approach, we show that the OT map itself can be used as a generative model, providing comparable performance. Previous analogous approaches consider OT maps as generative models only in the latent spaces due to their poor performance in the original high-dimensional ambient space. In contrast, we apply OT maps directly in the ambient space, e.g., a space of high-dimensional images. First, we derive a min-max optimization algorithm to efficiently compute OT maps for the quadratic cost (Wasserstein-2 distance). Next, we extend the approach to the case when the input and output distributions are located in the spaces of different dimensions and derive error bounds for the computed OT map. We evaluate the algorithm on image generation and unpaired image restoration tasks. In particular, we consider denoising, colorization, and inpainting, where the optimality of the restoration map is a desired attribute, since the output (restored) image is expected to be close to the input (degraded) one.

Paper Structure

This paper contains 24 sections, 3 theorems, 33 equations, 17 figures, 8 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background on Optimal Transport
  3. Optimal Transport in Generative Models
  4. End-to-end Solution to Learn Optimal Maps
  5. Equal Dimensions of Input and Output Distributions
  6. Unequal Dimensions of Input and Output Distributions
  7. Practical Aspects and Optimization Procedure
  8. Error Analysis
  9. Experiments
  10. Modeling Data distribution from Noise Distribution
  11. Unpaired Image Restoration
  12. Conclusion
  13. Acknowledgment
  14. Proofs
  15. Proof of Equivalence: Equation (\ref{['opt-x']}) and (\ref{['opt-t']})
  16. ...and 9 more sections

Key Result

Lemma 4.1

Let $T^{*}$ be the OT map from $\mu$ to $\nu$. Then, for every optimal potential $\psi^{*}$,

Figures (17)

  • Figure 1: Two existing approaches to use optimal transport in generative models.
  • Figure 2: Monge's OT.
  • Figure 3: The existing most prevalent approach to use OT maps in generative models.
  • Figure 4: The scheme of our approach for learning OT maps between unequal dimensions. In the figure, the setup of \ref{['sec-noise-to-data']} is shown: $\mu$ is a noise, $Q$ is the bicubic upscaling, $\nu$ is a distribution of images.
  • Figure 5: Randomly generated MNIST, CIFAR10, and CelebA samples by our method (OTM).
  • ...and 12 more figures

Theorems & Definitions (7)

  • Lemma 4.1
  • Lemma 4.2
  • Theorem 4.3
  • proof
  • proof
  • proof
  • proof