Monge-Amp\`ere Flow for Generative Modeling

1 INTRODUCTION . Generative modeling is a central topic in modern deep learning research ( Goodfellow et al. , 2016 ) 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 ) ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) , generative adversarial networks ( GANs ) ( Goodfellow et al. , 2014 ) , and autoregressive models ( Germain et al. , 2015 ; Kingma et al. , 2016 ; van den Oord et al. , 2016b ; a ; Papamakarios et al. , 2017 ) . 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 ( Dinh et al. , 2014 ; 2016 ) , and the more recently proposed Glow model ( Kingma & Dhariwal , 2018 ) . 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 ( Dinh et al. , 2014 ; Rezende & Mohamed , 2015 ; Kingma et al. , 2016 ; Papamakarios et al. , 2017 ) , 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 ( Villani , 2003 ; 2008 ; Peyré & Cuturi , 2018 ) and dynamical systems ( Katok & Hasselblatt , 1995 ) . Optimal transport theory concerns the problem of connecting two probability distributions p ( z ) and q ( x ) via transportation z 7→ x at a minimal cost . In this context , the Brenier theorem ( Brenier , 1991 ) 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 ( Ambrosio et al. , 2008 ) . 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 ( LeCun et al. , 1988 ) . Moreover , Chen et al . ( 2018 ) implemented the generative map as an ODE integration and employed efficient adjoint analysis for its training . In this paper , we devise the Monge-Ampè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 MongeAmpè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 . 2 THEORETICAL BACKGROUND . Consider latent variables z and physical variables x both living in RN . Given a diffeomorphism between them , x = x ( z ) , the probability densities in the latent and physical spaces are related by p ( z ) = q ( x ) ∣∣det ( ∂x∂z ) ∣∣ . The Brenier theorem ( Brenier , 1991 ) implies that instead of dealing with a multi-variable generative map , one can consider a scalar valued generating function x = ∇u ( z ) , where the convex Brenier potential u satisfies the Monge-Ampère equation ( Caffarelli et al. , 1998 ) p ( z ) q ( ∇u ( z ) ) = det ( ∂2u ∂zi∂zj ) . ( 1 ) Given the densities p and q , solving the Monge-Ampère equation for u turns out to be challenging due to the nonlinearity in the determinant . Moreover , for typical machine learning and statistical physics problems , an additional challenge is that one does not even have direct access to both probability densities p and q . Instead , one only has independent and identically distributed samples from one of them , or one only knows one of the distributions up to a normalization constant . Therefore , solving the Brenier potential in these contexts is a control problem instead of a boundary value problem . An additional computational challenge is that even for a given Brenier potential , the righthand side of equation 1 involves the determinant of the Hessian matrix , which scales unfavorably as O ( N3 ) with the dimensionality of the problem . To address these problems , we consider the Monge-Ampère equation in its linearized form , where the transformation is infinitesimal ( Villani , 2003 ) . We write the convex Brenier potential as u ( z ) = ||z||2/2 + ϕ ( z ) , thus x − z = ∇ϕ ( z ) , and correspondingly ln q ( x ) − ln p ( z ) = −Tr ln ( I + ∂ 2ϕ ∂zi∂zj ) = − ∇2ϕ ( z ) + O ( 2 ) . In the last equation we expand the logarithmic function and write the trace of a Hessian matrix as the Laplacian operator . Finally , taking the continuous-time limit → 0 , we obtain dx dt = ∇ϕ ( x ) , ( 2 ) d ln p ( x , t ) dt = −∇2ϕ ( x ) , ( 3 ) such that x ( 0 ) = z , p ( x , 0 ) = p ( z ) , and p ( x , T ) = q ( x ) , where t denotes continuous-time and T is a fixed time horizon . For simplicity here , we still keep the notion of x , which now depends on time . The evolution of x from time t = 0 to T then defines our generative map . We notice that Chen et al . ( 2018 ) used a more general form of these two equations as a continuous-time normalizing flow ( Rezende & Mohamed , 2015 ) . The two ordinary differential equations ( ODEs ) compose a dynamical system , which describes the flow of continuous random variables and the associated probability densities under iterative changeof-variable transformation . To match p ( x , T ) and the target density q ( x ) , one can optimize a functional I [ p ( x , T ) , q ( x ) ] that measures the difference between p ( x , T ) and q ( x ) . Thus , the training process amounts to solving an optimal control problem for the potential function : min ϕ I [ p ( x , T ) , q ( x ) ] . ( 4 ) p ( x , t ) p ( xt+ ϵ , t + ϵ ) As specified later , in our examples I [ p ( x , T ) , q ( x ) ] is defined as the Kullback-Leibler ( KL ) divergence or reverse KL divergence , which give rise to equations 6 and 7 , respectively . One can interpret equations 2 and 3 as fluid mechanics equations in the Lagrangian formalism . Equation 2 describes the trajectory of fluid parcels under the velocity field given by the gradient of the potential function ϕ ( x ) . While the time derivative in equation 3 is understood as the “ material derivative ” ( Thorne & Blandford , 2017 ) , which describes the change of the local fluid density p ( x , t ) experienced by someone traveling with the fluid . Figure 1 ( a ) illustrates the probability flow in a one dimension example ( Appendix A ) . The velocity field of the given potential pushes the fluid parcel outward , and the fluid density expands accordingly . The fluid mechanics interpretation is even more apparent if we write out the material derivative in equation 3 as d/dt = ∂/∂t+ dx/dt · ∇ , and use equation 2 to obtain ∂p ( x , t ) ∂t +∇ · [ p ( x , t ) v ] = 0 , ( 5 ) which is the Liouville equation , namely , the continuity equation of a compressible fluid with density p ( x , t ) and velocity field v = ∇ϕ ( x ) . Obeying the continuity equation ensures that the flow conserves the total probability mass . Moreover , the velocity field is curl free ∇ × v ≡ 0 and the fluid follows a form of gradient flow ( Ambrosio et al. , 2008 ) . The irrotational flow matches one ’ s heuristic expectation that the flow-based generative model transports probability masses . It should be stressed that although we use the optimal transport theory to motivate model architecture design , i.e. , the gradient flow for generative modeling , we do not have to employ the optimal transport objective functions . The difference is that in generative modeling one typically fixes only one end of the probability density and aims at learning a suitable transformation to reach the other end . While for optimal transport one has both ends fixed and aims at minimizing the transportation cost . Arjovsky et al . ( 2017 ) ; Bousquet et al . ( 2017 ) ; Genevay et al . ( 2017 ) adapted the Wasserstein distances in the optimal transport theory as an objective function for generative modeling . 3 PRACTICAL METHODOLOGY . We parametrize the potential function ϕ ( x ) using a feedforward neural network and integrate the ODEs ( 2 ) and ( 3 ) to transform the data and their log-likelihoods . Then , by applying the backpropagation algorithm through the ODE integration , we tune the parametrized potential function so that the probability density flows to the desired distribution 1 . 1See https : //github.com/ ... /MongeAmpereFlow for code implementation and pre-trained models . In general , one can represent the potential function in various ways as long as the gradient and Laplacian are well defined . Here we adopt a densely connected neural network with only one hidden layer in our minimalist implementation . We use the softplus activation function in the hidden layers so that the potential function is differentiable to higher orders . The gradient and the Laplacian of the potential appearing in the right-hand side of equations 2 and 3 can be computed via automatic differentiation . We integrate equations 2 and 3 using a numerical ODE integrator . At each integration step one advances the ODE system by a small amount , as shown in figure 1 . Therefore , the numerical integration of ODE is equivalent to the forward evaluation of a deep residual neural network ( He et al. , 2016 ; E , 2017 ; Li & Shi , 2017 ; Chang et al. , 2017 ; Lu et al. , 2017 ; Rousseau & Fablet , 2018 ; Ruthotto & Haber , 2018 ; Haber & Ruthotto , 2018 ; Chen et al. , 2018 ) . Alternatively , the integration can also be viewed as a recurrent neural network in which one feeds the output back into the block iteratively . In both pictures , the network heavily shares parameters in the depth direction since the gradient and Laplacian information of the same potential function is used in each integration step . As a result , the network is very parameter efficient . We employ the fourth order Runge-Kutta scheme with a fixed time step , which is set such that the numerical integrator is accurate enough . Thus , the ODE integrator block shown in figure 1 ( b ) contains four layers of neural networks . With a hundred integration steps , the whole ODE integration procedure corresponds to four hundreds layers of a neural network . At the end of the integration , we obtain samples x and their likelihoods ln p ( x , T ) , which depends on the parametrized potential function ϕ ( x ) . Differentiable optimization of such deep neural network is feasible since the integrator only applies small changes to the input signal of each layer . The deep generative model resulting from the discretization of an ODE system shows some nice features compared to the conventional neural networks . At training time , there is then a tradeoff between the total number of integration steps and the expressibility of the potential function . Longer integration time corresponds to a deeper network , and hence a simpler transformation at each step . While at testing time , one can construct a variable depth neural network for the potential function . For example , one can use a larger time step and a smaller number of integration steps for efficient inference . Moreover , by employing a reversible ODE integrator , one can also integrate the equations 2 and 3 backward in time with the same computational complexity , and return to the starting point deterministically .

Summary: This paper introduces a continuous-time flow, which is motivated from continuous-time gradient flow in the Monge-Ampere equation in optimal transport theory. They relate the resulting generative model to a dynamical system with a learnable potential function that controls a compressible fluid (representing the probability density), towards the target density. The resulting set of differential equations for the evolution of the samples and the density is solved through a fourth-order Runge-Kutta ODE solver. By imposing symmetries on the scalar potential function, symmetries of the resulting target distribution can also be enforced, a property that is desirable in many applications.