Table of Contents
Fetching ...

No-Regret Generative Modeling via Parabolic Monge-Ampère PDE

Nabarun Deb, Tengyuan Liang

TL;DR

This work introduces a generative framework built on a discretized parabolic Monge-Ampère PDE, treated as a dynamic, no-regret refinement of Brenier transport maps toward the target measure. By establishing a novel Evolution Variational Inequality on the Wasserstein space and a three-point Bregman identity, the authors derive average- and last-iterate convergence guarantees that extend beyond log-concave targets. The framework connects Sinkhorn limits to a continuous-time PDE, and it supports practical neural-PDE implementations via logistic regression density-ratio learning or score matching, while also offering variational-inference perspectives. Overall, the approach unifies sampling efficiency, learning tractability, and theoretical guarantees for non-log-concave targets, with concrete neural architectures and VI interpretations provided. This yields a versatile, theoretically grounded pathway for advanced generative modeling and approximate Bayesian inference.

Abstract

We introduce a novel generative modeling framework based on a discretized parabolic Monge-Ampère PDE, which emerges as a continuous limit of the Sinkhorn algorithm commonly used in optimal transport. Our method performs iterative refinement in the space of Brenier maps using a mirror gradient descent step. We establish theoretical guarantees for generative modeling through the lens of no-regret analysis, demonstrating that the iterates converge to the optimal Brenier map under a variety of step-size schedules. As a technical contribution, we derive a new Evolution Variational Inequality tailored to the parabolic Monge-Ampère PDE, connecting geometry, transportation cost, and regret. Our framework accommodates non-log-concave target distributions, constructs an optimal sampling process via the Brenier map, and integrates favorable learning techniques from generative adversarial networks and score-based diffusion models. As direct applications, we illustrate how our theory paves new pathways for generative modeling and variational inference.

No-Regret Generative Modeling via Parabolic Monge-Ampère PDE

TL;DR

This work introduces a generative framework built on a discretized parabolic Monge-Ampère PDE, treated as a dynamic, no-regret refinement of Brenier transport maps toward the target measure. By establishing a novel Evolution Variational Inequality on the Wasserstein space and a three-point Bregman identity, the authors derive average- and last-iterate convergence guarantees that extend beyond log-concave targets. The framework connects Sinkhorn limits to a continuous-time PDE, and it supports practical neural-PDE implementations via logistic regression density-ratio learning or score matching, while also offering variational-inference perspectives. Overall, the approach unifies sampling efficiency, learning tractability, and theoretical guarantees for non-log-concave targets, with concrete neural architectures and VI interpretations provided. This yields a versatile, theoretically grounded pathway for advanced generative modeling and approximate Bayesian inference.

Abstract

We introduce a novel generative modeling framework based on a discretized parabolic Monge-Ampère PDE, which emerges as a continuous limit of the Sinkhorn algorithm commonly used in optimal transport. Our method performs iterative refinement in the space of Brenier maps using a mirror gradient descent step. We establish theoretical guarantees for generative modeling through the lens of no-regret analysis, demonstrating that the iterates converge to the optimal Brenier map under a variety of step-size schedules. As a technical contribution, we derive a new Evolution Variational Inequality tailored to the parabolic Monge-Ampère PDE, connecting geometry, transportation cost, and regret. Our framework accommodates non-log-concave target distributions, constructs an optimal sampling process via the Brenier map, and integrates favorable learning techniques from generative adversarial networks and score-based diffusion models. As direct applications, we illustrate how our theory paves new pathways for generative modeling and variational inference.

Paper Structure

This paper contains 26 sections, 15 theorems, 107 equations, 1 figure, 3 algorithms.

Key Result

Proposition 2.1

Suppose $\psi:\mathbb{R}^d\to\mathbb{R}$ is a uniformly strongly convex $\mathcal{C}^2$ function. Let $\tilde{\psi}^{\epsilon}(y):=\mathcal{V}^{\epsilon}[\psi](y)$. Then, as $\epsilon\to 0$, we have:

Figures (1)

  • Figure 1: The probability densities $\rho_1,\rho_2,\pi$ live on the target domain $\mathcal{X}$ while the densities $\pi_1,\pi_2,e^{-g}$ live on the reference domain $\mathcal{Y}$. $\nabla\phi_{\rho_i}$ is the Brenier map from $\rho_i$ to $e^{-g}$, $i=1,2$. As $\pi_i=\nabla \phi_{\rho_i}\#\pi$, Brenier's Theorem brenier1991polar implies that $\nabla\phi_{\rho_i}$ is also the Brenier map from $\pi$ to $\pi_i$, for $i=1,2$. The telescoping term in \ref{['lem:mirrid']}, $B_G(\pi|\rho_1)-B_G(\pi|\rho_2)$ involves densities $\rho_1,\rho_2,\pi$ all of which live on the target space. One might naturally expect the transportation cost to be $B_G(\rho_1 | \rho_2)$. In sharp contrast, the corresponding term in \ref{['lem:mirrid']} is $B_{G_{\pi}}(\pi_1 | \pi_2)$ where $\pi_i$ lives on the reference domain and is the image of $\pi$ under the Brenier map $\nabla\phi_{\rho_i}$, $i=1,2$.

Theorems & Definitions (35)

  • Proposition 2.1
  • Definition 1: Bregman divergence over probability measures
  • Example 1: Entropy as mirror
  • Example 2: Wasserstein distance as mirror
  • Lemma 3.1: A new Bregman three-point identity
  • Remark 1
  • Lemma 3.2: Bregman and KL: relative convexity
  • Remark 2
  • Theorem 4.1
  • Remark 3
  • ...and 25 more