Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Energy-guided Entropic Neural Optimal Transport

Petr Mokrov, Alexander Korotin, Alexander Kolesov, Nikita Gushchin, Evgeny Burnaev

TL;DR

The paper addresses learning truly conditional optimal transport plans under entropy regularization by uniting Energy-Based Models with Entropy-regularized OT. It develops an energy-based reformulation of the weak dual EOT problem, yielding a closed-form conditional minimizer and a practical gradient-based optimization using Langevin sampling. The approach comes with generalization guarantees via Rademacher complexities and is demonstrated on toy benchmarks, Gaussian-to-Gaussian comparisons, and high-resolution unpaired I2I translation in the StyleGAN latent space, achieving competitive results and scalable performance. This framework enables principled, conditional data-to-data translation with a theoretically grounded single-energy function backbone, offering a pathway to large-scale OT maps in complex domains.

Abstract

Energy-based models (EBMs) are known in the Machine Learning community for decades. Since the seminal works devoted to EBMs dating back to the noughties, there have been a lot of efficient methods which solve the generative modelling problem by means of energy potentials (unnormalized likelihood functions). In contrast, the realm of Optimal Transport (OT) and, in particular, neural OT solvers is much less explored and limited by few recent works (excluding WGAN-based approaches which utilize OT as a loss function and do not model OT maps themselves). In our work, we bridge the gap between EBMs and Entropy-regularized OT. We present a novel methodology which allows utilizing the recent developments and technical improvements of the former in order to enrich the latter. From the theoretical perspective, we prove generalization bounds for our technique. In practice, we validate its applicability in toy 2D and image domains. To showcase the scalability, we empower our method with a pre-trained StyleGAN and apply it to high-res AFHQ $512\times 512$ unpaired I2I translation. For simplicity, we choose simple short- and long-run EBMs as a backbone of our Energy-guided Entropic OT approach, leaving the application of more sophisticated EBMs for future research. Our code is available at: https://github.com/PetrMokrov/Energy-guided-Entropic-OT

Energy-guided Entropic Neural Optimal Transport

TL;DR

The paper addresses learning truly conditional optimal transport plans under entropy regularization by uniting Energy-Based Models with Entropy-regularized OT. It develops an energy-based reformulation of the weak dual EOT problem, yielding a closed-form conditional minimizer and a practical gradient-based optimization using Langevin sampling. The approach comes with generalization guarantees via Rademacher complexities and is demonstrated on toy benchmarks, Gaussian-to-Gaussian comparisons, and high-resolution unpaired I2I translation in the StyleGAN latent space, achieving competitive results and scalable performance. This framework enables principled, conditional data-to-data translation with a theoretically grounded single-energy function backbone, offering a pathway to large-scale OT maps in complex domains.

Abstract

Energy-based models (EBMs) are known in the Machine Learning community for decades. Since the seminal works devoted to EBMs dating back to the noughties, there have been a lot of efficient methods which solve the generative modelling problem by means of energy potentials (unnormalized likelihood functions). In contrast, the realm of Optimal Transport (OT) and, in particular, neural OT solvers is much less explored and limited by few recent works (excluding WGAN-based approaches which utilize OT as a loss function and do not model OT maps themselves). In our work, we bridge the gap between EBMs and Entropy-regularized OT. We present a novel methodology which allows utilizing the recent developments and technical improvements of the former in order to enrich the latter. From the theoretical perspective, we prove generalization bounds for our technique. In practice, we validate its applicability in toy 2D and image domains. To showcase the scalability, we empower our method with a pre-trained StyleGAN and apply it to high-res AFHQ unpaired I2I translation. For simplicity, we choose simple short- and long-run EBMs as a backbone of our Energy-guided Entropic OT approach, leaving the application of more sophisticated EBMs for future research. Our code is available at: https://github.com/PetrMokrov/Energy-guided-Entropic-OT
Paper Structure (35 sections, 5 theorems, 49 equations, 6 figures, 6 tables, 1 algorithm)

This paper contains 35 sections, 5 theorems, 49 equations, 6 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Optimal Transport
  4. Energy-based models
  5. Taking up EOT problem with EBMs
  6. Energy-guided reformulation of weak dual EOT
  7. Optimization procedure
  8. Generalization bounds for learned entropic plans
  9. Related works
  10. Energy-Based Models for unpaired data-to-data translation
  11. Entropy-regularized OT
  12. Experimental illustrations
  13. Toy 2D
  14. Gaussian-to-Gaussian
  15. High-Dimensional Unpaired Image-to-image Translation
  16. ...and 20 more sections

Key Result

Theorem 1

Let $f \in \mathcal{C}(\mathcal{Y})$ and $x \in \mathcal{X}$. Then inner weak dual objective $\min_{\mu \in \mathcal{P}(\mathcal{Y})} \mathcal{G}_{x, f}(\mu)$weak_c_EOT_transform permits the unique minimizer $\mu_x^f$ which is given by where $Z(f, x) \stackrel{\text{def}}{=} \int_{\mathcal{Y}} \exp\left(\frac{f(y) - c(x, y)}{\varepsilon}\right) \dd y$.

Figures (6)

  • Figure 1: AFHQ $512\times512$Cat$\rightarrow$Dog unpaired translation by our Energy-guided EOT solver applied in the latent space of StyleGAN2-ADA. Our approach does not need data2latent encoding.Left: source samples; right: translated samples.
  • Figure 2: Performance of Energy-guided EOT on Gaussian$\rightarrow$Swissroll 2D setup.
  • Figure 3: AFHQ $512\times512$Wild$\rightarrow$Dog unpaired I2I by our method in the latent space of StyleGAN2-ADA. Left: source; right: translated.
  • Figure 4: Quantitative performance of Energy-guided EOT on Colored MNIST.
  • Figure 5: Uncurated Image-to-Image translation by our method in the latent space of StyleGAN.
  • ...and 1 more figures

Theorems & Definitions (11)

  • Theorem 1: Optimizer of weak $C_{\text{EOT}}$-transform
  • Theorem 2: Bound on the quality of the plan recovered from the dual variable
  • Theorem 3: Gradient of the weak dual loss $L(\theta)$
  • Theorem 4: Finite sample learning guarantees
  • proof
  • proof
  • proof
  • Definition 1: Rademacher complexity $\mathcal{R}_N( \mathcal{F}, \mu)$
  • proof
  • Lemma 1
  • ...and 1 more