We present a deep generative model, named Monge-Amp\`ere flow, which builds on continuous-time gradient flow arising from the Monge-Amp\`ere 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\`ere 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\`ere equation, optimal transport, and fluid dynamics into reversible flow-based generative models. Generative modeling is a central topic in modern deep learning research BID20 which finds broad applications in image processing, speech synthesis, reinforcement learning, as well as in solving inverse problems and statistical physics problems. The goal of generative modeling is to capture the full joint probability distribution of high dimensional data and generate new samples according to the learned distribution. There have been significant advances in generative modeling in recent years. Of particular interests are the variational autoencoders (VAEs) BID33 BID47 , generative adversarial networks (GANs) BID19 , and autoregressive models BID18 BID34 BID55 a; BID44 . Besides, there is another class of generative models which has so far gained less attention compared to the aforementioned models. These models invoke a sequence of diffeomorphism to connect between latent variables with a simple base distribution and data which follow a complex distribution. Examples of these flow-based generative models include the NICE and the RealNVP networks BID12 , and the more recently proposed Glow model BID32 . These models enjoy favorable properties such as tractable likelihoods and efficient exact inference due to invertibility of the network.A key concern in the design of flow-based generative models has been the tradeoff between the expressive power of the generative map and the efficiency of training and sampling. One typically needs to impose additional constraints in the network architecture BID12 BID46 BID34 BID44 , which unfortunately weakens the model compared to other generative models. In addition, another challenge to generative modeling is how to impose symmetry conditions such that the model generates symmetry related configurations with equal probability.To further unleash the power of the flow-based generative model, we draw inspirations from the optimal transport theory BID56 BID57 BID45 and dynamical systems BID30 . Optimal transport theory concerns the problem of connecting two probability distributions p(z) and q(x) via transportation z → x at a minimal cost. In this context, the Brenier theorem BID4 states that under the quadratic distance measure, the optimal generative map is the gradient of a convex function. This motivates us to parametrize the vector-valued generative map as the gradient of a scalar potential function, thereby formulating the generation process as a continuous-time gradient flow BID0 . In this regard, a generative map is naturally viewed as a deterministic dynamical system which evolves over time. Numerical integration of the dynamical system gives rise to the neural network representation of the generative model. To this end, E (2017) proposes a dynamical system perspective to machine learning, wherein the training procedure is viewed as a control problem, and the algorithm like back-propagation is naturally derived from the optimal control principle BID36 . Moreover, implemented the generative map as an ODE integration and employed efficient adjoint analysis for its training.In this paper, we devise the Monge-Ampère flow as a new generative model and apply it to two problems: density estimation of the MNIST dataset and variational calculation of the Ising model. In our approach, the probability density is modeled by a compressible fluid, which evolves under the gradient flow of a learnable potential function. The flow has tractable likelihoods and exhibits the same computational complexity for sampling and inference. Moreover, a nice feature of the MongeAmpère flow is that one can construct symmetric generative models more easily by incorporating the symmetries into the scalar potential. The simplicity and generality of this framework allow the principled design of the generative map and gives rise to lightweight yet powerful generative models. Gradient flow of compressible fluids in a learnable potential provides a natural way to set up deep generative models. The Monge-Ampère flow combines ideas and techniques in optimal transport, fluid dynamics, and differential dynamical systems for generative modeling.We have adopted a minimalist implementation of the Monge-Ampère flow with a scalar potential parameterized by a single hidden layer densely connected neural network. There are a number of immediate improvements to further boost its performance. First, one could extend the neural network architecture of the potential function in accordance with the target problem. For example, a convolutional neural network for data with spatial or temporal correlations. Second, one can explore better integration schemes which exactly respect the time-reversal symmetry to ensure reversible sampling and inference. Lastly, by employing the backpropagation scheme of through the ODE integrator one can reduce the memory consumption and achieve guaranteed convergence in the integration.Furthermore, one can employ the Wasserstein distances BID1 instead of the KLdivergence to train the Monge-Ampère flow. With an alternative objective function, one may obtain an even more practically useful generative model with tractable likelihood. One may also consider using batch normalization layers during the integration of the flow BID13 BID44 . However, since the batch normalization breaks the physical interpretation of the continuous gradient flow of a fluid, one still needs to investigate whether it plays either a theoretical or a practical role in the continuous-time flow.Moreover, one can use a time-dependent potential ϕ(x, t) to induce an even richer gradient flow of the probability densities. BID2 has shown that the optimal transport flow (in the sense of minimizing the spatial-time integrated kinetic energy, which upper bounds the squared Wasserstein-2 distance) follows a pressureless flow in a time-dependent potential. The fluid moves with a constant velocity that linearly interpolates between the initial and the final density distributions. Practically, a time-dependent potential corresponds to a deep generative model without sharing parameters in the depth direction as shown in FIG1 . Since handling a large number of independent layers for each integration step may be computationally inefficient, one may simply accept one additional time variable in the potential function, or parametrize ϕ(x, t) as the solution of another differential equation, or partially tie the network parameters using a hyper-network BID22 .Besides applications presented here, the Monge-Ampère flow has wider applications in machine learning and physics problems since it inherits all the advantages of the other flow-based generative models BID12 BID46 BID34 BID13 BID44 . A particular advantage of generative modeling using the Monge-Ampère flow is that it is relatively easy to impose symmetry into the scalar potential. It is thus worth exploiting even larger symmetry groups, such as the permutation for modeling exchangeable probabilities BID35 . Larger scale practical application in statistical and quantum physics is also feasible with the Monge-Ampère flow. For example, one can study the physical properties of realistic molecular systems using Monge-Ampère flow for variational free energy calculation. Lastly, since the mutual information between variables is greatly reduced in the latent space, one can also use the Monge-Ampère flow in conjunction with the latent space hybrid Monte Carlo for efficient sampling BID38 .
A gradient flow based dynamical system for invertible generative modeling
{"src": "scitldr-AIC_train", "paper_id": "rkeUrjCcYQ"}