Monge-Ampère Flow for Generative Modeling

TL;DR

The paper introduces the Monge-Ampère flow, a continuous-time, gradient-flow-based generative model derived from the Monge-Ampère equation in optimal transport. By linearizing the Monge-Ampère equation, it yields ODEs that describe a gradient-flow of a learnable potential, enabling tractable likelihoods, efficient sampling, and symmetry-aware generation. The approach is demonstrated on MNIST density estimation, achieving competitive test log-likelihoods with far fewer parameters, and on the Ising model at criticality via variational learning with symmetry constraints. This work bridges optimal transport, fluid dynamics, and reversible flow-based models, offering a flexible framework for principled generative modeling with potential impact across physics-informed machine learning.

Abstract

We present a deep generative model, named Monge-Ampère flow, which builds on continuous-time gradient flow arising from the Monge-Ampère equation in optimal transport theory. The generative map from the latent space to the data space follows a dynamical system, where a learnable potential function guides a compressible fluid to flow towards the target density distribution. Training of the model amounts to solving an optimal control problem. The Monge-Ampère flow has tractable likelihoods and supports efficient sampling and inference. One can easily impose symmetry constraints in the generative model by designing suitable scalar potential functions. We apply the approach to unsupervised density estimation of the MNIST dataset and variational calculation of the two-dimensional Ising model at the critical point. This approach brings insights and techniques from Monge-Ampère equation, optimal transport, and fluid dynamics into reversible flow-based generative models.

Figures (4)

• Figure 1: (a) Schematic illustration of the gradient flow of compressible fluid in one-dimension. The instantaneous velocities of fluid parcels are determined by the gradient of the potential function. The fluid density evolves to the final one after a finite flow time. (b) Numerical integration of the Monge-Ampère flow is equivalent to forward passing through a deep residual neural network. Each integration step advances these fluid parcels by one step according to the instantaneous velocity. The density distribution changes accordingly. $d$ is the total number of integration steps.
• Figure 2: Schematic illustration of two applications (a) unsupervised density estimation and (b) variational free energy calculation for a statistical mechanics problem. In both cases, we integrate equations \ref{['eq:dxdt']} and \ref{['eq:dlnpdt']} under a parametrized potential function $\varphi({\bm{x}})$, and optimize $\varphi({\bm{x}})$ such that the density at the other end matches to the desired one.
• Figure 3: (a) The NLL of the training (blue) and the test (orange) MNIST dataset. The horizontal lines indicate results obtained with previous flow-based models reported in Papamakarios. (b) From top to bottom, Monge-Ampère flow of test MNIST images to the base Gaussian distribution.
• Figure 4: The training objectivity equation \ref{['eq:variational']} is the variational free energy of the Ising model, which is lower bounded by the exact solution indicated by the horizontal red line Onsager1944Li2018. Inset shows the directed generated Ising configurations at epochs $0, 500, 1000, 1500$ respectively. Each inset contains a $5\times 5$ tile of a $16^2$ Ising model.