Physics Informed Convex Artificial Neural Networks (PICANNs) for Optimal Transport based Density Estimation

1 INTRODUCTION . Optimal Mass Transport ( OMT ) is a well-studied problem with a variety of applications in a diverse set of fields , ranging from physics to computer vision and in particular statistics and data science . In this paper we propose a new framework for the estimation of the solution to the L2-optimal transport problem between two densities . Our algorithm , which is based on Brenier ’ s theorem , and builds on recent developments of input convex neural networks and physics-informed neural networks for solving PDE ’ s . Before we describe the contributions of the article in more detail , we will briefly summarize the motivation of our investigations and recent developments in the field . Density estimation and random sampling : The density estimation problem is to estimate a smooth probability density based on a discrete finite set of observations . In traditional parametric density estimation techniques , the data is assumed to be drawn from a known parametric family of distributions . One of the most ubiquitous parametric techniques is Gaussian Mixture Modeling ( McLachlan & Basford , 1988 ) . Nonparametric techniques were first proposed by Fix & Hodges ( 1951 ) ( Silverman & Jones ( 1989 ) ) to move away from rigid distributional assumptions . The most used approach is the kernel density estimation , which dates back to Rosenblatt ( 1952a ) and Parzen ( 1962 ) . Many challenges remain regarding the implementation and practical performance of kernel density estimators , including in particular , the bandwidth selection and the lack of local adaptivity resulting in a large sensitivity to outliers ( Loader et al. , 1999 ) . These problems are particularly exacerbated in high dimensions with the curse of dimensionality . Recently , diffeomorphic transformation-based algorithms have been proposed to tackle this problem ( Dinh et al. , 2017 ; Marzouk et al. , 2016 ; Younes , 2020 ; Bauer et al. , 2017 ) . The basic concept of transformation-based algorithms is to find a diffeomorphic mapping between a reference proba- bility distribution and the unknown target distribution , from which the data is drawn . Consequently , transformation-based density estimation leads at the same time to an efficient generative model , as new samples from the estimated density can be generated at a low cost by sampling from the reference density and transforming the samples by the estimated transformation . The fundamental problem in diffeomorphic transformation-based approaches is how to estimate and select the transformation : from a theoretical point of view there exists an infinite set of transformations that map two given probability densities onto each other . Recently , several deep learning methods have been devised for this task , where Normalizing Flows ( NF ) stand out among these methods . Examples of such models include Real NVP ( Dinh et al. , 2017 ) , Masked Autoregressive Flows ( Papamakarios et al. , 2017 ) , iResNets ( Behrmann et al. , 2019 ) , Flow++ ( Ho et al. , 2019 ) and Glow ( Kingma & Dhariwal , 2018 ) . For a review of the vast NF literature , we refer to the the overview article ( Kobyzev et al. , 2020 ) . Although these methods have shown to perform well in density estimation applications , the interpretability of the obtained transformation is less clear , e.g . in Real NVP ( Dinh et al. , 2017 ) , the solution selection is obtained by restricitng the transformations to the class of diffeomorphisms with triangular Jacobians that are easy to invert , which is closely related to the Knothe-Rosenblatt rearrangement ( Knothe , 1957 ; Rosenblatt , 1952b ) . Optimal mass transport : Optimal mass transport , on the other hand , formulates the transport map selection as the minimizer of a cost function ( Villani , 2008 ; 2003 ) . The optimal transportation cost induces a metric structure , the Wasserstein metric , on the space of probability densities and is sometimes referred to as the Earth Mover ’ s Distance . This theory , which dates back to 1781 , was originally formulated by the French mathematician Gaspard Monge ( 1781 ) . The difficulty in applying this framework to the proposed density estimation problem lies in solving the corresponding optimization problem , which in a dimension greater than one is highly non trivial . The fully discrete OMT problem ( optimal assignment problem ) can be solved using linear programming and can be approximated by the Sinkhorn algorithm ( Cuturi , 2013a ; Papadakis , 2015 ) . However , these algorithms do not lead to a continuous transformation map and thus can ’ t be used for the proposed diffeomorphic density estimation and generative modelling . Previous algorithmic solutions for the continuous OMT problem include fluid mechanics-based approaches ( Benamou & Brenier , 2000 ) , finite element or finite difference-based methods ( Benamou et al. , 2010 ; Benamou & Duval , 2019 ) and steepest descent-based energy minimization approaches ( Angenent et al. , 2003 ; Carlier et al. , 2010 ; Loeper & Rapetti , 2005 ) . In recent years , several deep learning methods have been deployed for solving the OMT problem . In these methods , the OMT problem is typically embedded in the loss function for the neural network model . Recent work by Makkuva et al . ( 2020 ) proposed to approximate the OMT map as the solution of min-max optimization using input convex neural networks ( ICNN ) , see Amos et al . ( 2017 ) . The min-max nature of this algorithm arises from the need to train an ICNN to represent a convex function and the conjugate of the convex function . Building upon this approach , Korotin et al . ( 2019 ) imposed a cyclic regularisation that converts the min-max optimization problem to a standard minimization problem . This change results in a faster converging algorithm that scales well to higher dimensions and also prevents convergence to local saddle points and instabilities during training , as is the case in the min-max algorithm . Another class of neural networks which have been proposed to solve OMT problems are Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) . GANs are defined through a min-max game of two neural networks where one of the networks tries to generate new samples from a data distribution , while the other network judges whether these generated samples originate from the data population or not . Later , Gulrajani et al . ( 2017 ) proposed using the Wasserstein-1 distance in GANs instead of the Jensen-Shannon divergence between the generated distribution and the data distribution as in the original formulation . They demonstrated that this new loss functions leads to better stability of the training of networks attributed to the Wasserstein metric being well defined even when the two distributions do not share the same support . Contributions : In this paper , we propose a different deep learning-based framework to approximate the optimal transport maps . The approach we present relies on Brenier ’ s celebrated theorem ( Brenier , 1991 ) , thereby reducing the optimal transport problem to that of solving a partial differential equation : a Monge-Ampere type equation . We frame this PDE in the recently developed paradigm of Physics Informed Neural Networks ( PINNs ) ( Raissi et al. , 2017 ) . Similar to other deep learning-based algorithms , our framework directly inherits the dimensional scalability of neural networks ( Shin et al. , 2020 ) , which traditional finite element or finite difference methods for solving PDEs do not possess . Brenier ’ s theorem further states that the optimal transport map is given by the gradient of a convex function- the Brenier potential . To incorporate this information in our PINN approach , we parameterize the Brenier potential using an ICNN , thereby guaranteeing its convexity . We test the accuracy of our OMT solver on numerous synthetic examples for which analytical solutions are known . Our experiments show that our algorithm indeed approximates the true solution well , even in high dimensions . To further quantify the performance of the new framework , we compare it to two other deep learning-based algorithms , for which we guided the selection by the results of the recent benchmarking paper by Korotin et al . ( 2021 ) , in which they evaluate the methods presented in Seguy et al . ( 2017 ) ; Nhan Dam et al . ( 2019 ) ; Taghvaei & Jalali ( 2019 ) ; Makkuva et al . ( 2020 ) ; Liu et al . ( 2019 ) ; Mallasto et al . ( 2019 ) ; Korotin et al . ( 2019 ) . We restricted our comparision to the algorithms of Makkuva et al . ( 2020 ) and Korotin et al . ( 2019 ) , as these two showed the best performance in this benchmark . Our results showed that the newly proposed method significantly outperforms these methods in terms of accuracy . As an explicit application of our solution of OMT , we focus on the density estimation problem . In synthetic examples , we show that we can estimate the true density based on a limited amount of samples . We compare the results of the proposed method to four other density estimation algorithm : the two OMT based algorithms mentioned above , and two methods from the family of normalizing flows : RealNVP ( Dinh et al . ( 2016 ) ) and iResNet ( Behrmann et al . ( 2019 ) ) . Finally we demonstrate how our framework can be combined with a traditional autoencoder to obtain a generative framework . In accordance with the best practices for reproducible research , we are providing an open-source version of the code , which is publicly available on github . 2 OMT USING DEEP LEARNING . In this section , we will present our framework for solving the Optimal Mass Transport ( OMT ) problem . Our approach will combine methods of deep learning with the celebrated theorem of Brenier , which reduces the solution of the OMT problem to solving a Monge-Ampere type equation . To be more precise , we will tackle this problem by embedding the Monge-Ampere equation into the broadly applicable concept of Physics Informed Neural Networks . 2.1 MATHEMATICAL BACKGROUND OF OMT . We start by summarizing the mathematical background of OMT , including a description of Brenier ’ s theorem . For more information we refer to the vast literature on OMT , see e.g. , Villani ( 2003 ; 2008 ) . Let Ω be a convex and bounded domain of Rn and let dx denote the standard measure on Rn . For simplicity , we restrict our presentation to the set P ( Ω ) of all absolutely continuous measures on Ω , i.e. , P ( Ω ) 3 µ = fdx with f ∈ L1 ( Ω ) , such that ∫ Ω fdx = 1 . From here , on we will identify the measure µ with its density function f . We aim to minimize the cost of transporting a density µ to a density ν using a ( transport ) map T , which leads to the so-called Monge Optimal Transport Problem . We will consider only the special case of a quadratic cost function as Brenier ’ s theorem , which forms the basis of our algorithm , is only true for this type of cost function . Definition 2.1 ( L2-Monge Optimal Transport Problem ) Given µ , ν ∈ P ( Ω ) , minimize M ( T ) =∫ Ω ‖x − T ( x ) ‖2dµ ( x ) over all µ-measureable maps T : Ω → Ω subject to ν = T∗µ . We will call an optimal T an optimal transport map . Here , the constraint is formulated in terms of the push forward action of a measurable map T : Ω→ Ω , which is defined via T∗µ ( B ) = µ ( T−1 ( B ) ) , for every measurable set A ⊂ Ω . By a change of coordinates , the constraint T∗µ = T∗ ( fdx ) = ν = gdx can be thus reduced to the equation f ( x ) = g ( T ( x ) ) |det ( DT ( x ) ) | . ( 1 ) This equation can also be expressed via the pullback action as µ = T ∗ν : = ( T−1 ) ∗ν . The existence and uniqueness of an optimal transport map is not guaranteed . We will see , that in our situation , i.e. , for absolutely continuous measures , the existence and uniqueness is indeed guaranteed . First , we will introduce a more general formulation of the Monge problem , called Kantorovich formulation . Therefore , we define the space of all transport plans Π ( µ , ν ) , i.e. , of all measures on the product space Ω× Ω , such that the first marginal is µ and the second marginal is ν . Then we have : Definition 2.2 ( L2-Kantorovich ’ s Optimal Transport Problem ) Given µ , ν ∈ P ( Ω ) , minimize K ( π ) = ∫ Ω×Ω ‖x− y‖ 2dπ ( x , y ) over all π ∈ Π ( µ , ν ) . Note that the L2-Wasserstein metric W2 ( µ , ν ) between µ and ν is defined as the infimum of K. We will now formulate Brenier ’ s theorem , which guarantees the existence of an optimal transport map and will be the central building block of our algorithm : Theorem 2.3 ( Brenier ( 1991 ) ) Let µ , ν ∈ P ( Ω ) . Then there exists a unique optimal transport plan π∗ ∈ Π ( f , g ) , which is given by π∗ ( x , y ) = ( id×T ) where T = ∇u is the gradient of a convex function u that pushes µ forward to ν , i.e. , ( ∇u ) ∗µ = ν . The inverse T−1 is also given by the gradient of a convex function that is the Legendre transform of the convex function u . Thus , Brenier ’ s Theorem guarantees the existence and the uniqueness of the optimal transport map of the OMT problem . Consequently , we can determine this optimal transport map by solving for the function u in the form of a Monge-Ampère equation : det ( D2 ( u ) ( x ) ) · g ( ∇u ( x ) ) = f ( x ) ( 2 ) where D2 is the Hessian , µ = fdx and ν = gdx . We obtain equation 2 directly from equation 1 using the constraint that T = ∇u as required by Brenier ’ s theorem . We will also refer to this map as the Brenier map . This map is a diffeomorphism as it is a gradient of a strictly convex function . In general , the Monge–Ampère equation is a nonlinear second-order partial differential equation . If we limit ourselves to convex solutions , then the differential equation is elliptic . If the two densities f and g are smooth ( C∞ ) , positive and absolutely continuous with respect to each other then an unique smooth convex solution is guaranteed to exist . See e.g . Forzani & Maldonado ( 2004 ) ; Bakelman ( 1983 ) ; De Philippis & Figalli ( 2013 ) for further details on existence and regularity of solutions to this equation . Using methods of classical numerical analysis , Brenier ’ s theorem has been used e.g . in Peyré et al . ( 2019 ) to obtain a numerical framework for the continous OMT problem . In the following section we will propose a new discretization to this problem , which will make use of recent advances in deep learning .

Summary: This paper introduces a new method for density estimation that is based on the idea of optimal mass transport (OMP), i.e., finding some function that transports a probability density into another probability density while minimizing the transportation cost. To do so, the authors leverage Brenier's theorem which guarantees the existence and uniqueness of the optimal transport map. As a consequence of this theorem,  the optimal transport map of the OMT problem can be obtained as the solution of a special nonlinear partial differential equation (PDE). The authors then show that the approximate solution can be computed with help of physics informed neural networks (PINNs). The approach is demonstrated for several canonical examples.