Physics Informed Convex Artificial Neural Networks (PICANNs) for Optimal Transport based Density Estimation
Authors: Amanpreet Singh
Martin Bauer
Sarang Joshi
Venue: ICLR 2022 Submitted (2022)
Keywords: Optimal Mass Transport
Density Estimation
Physics Informed Neural Networks
Input Convex Neural Networks
Monge Ampere Equation

Abstract: 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. Since the original formulation of Monge in 1781 significant theoretical progress been made on the existence, uniqueness and properties of the optimal transport maps. The actual numerical computation of the transport maps, particularly in high dimensions, remains a challenging problem. In the past decade several neural network based algorithms have been proposed to tackle this task. In this paper, building on recent developments of input convex neural networks and physics informed neural networks for solving PDE's, we propose a new Deep Learning approach to solve the continuous OMT problem. Our framework is based on Brenier's theorem, which reduces the continuous OMT problem to that of solving a non-linear PDE of Monge-Ampere type whose solution is a convex function. To demonstrate the accuracy of our framework we compare our method to several other deep learning based algorithms. We then focus on applications to the ubiquitous density estimation and generative modeling tasks in statistics and machine learning. Finally as an example we present how our framework can be incorporated with an autoencoder to estimate an effective probabilistic generative model.

Under review as a conference paper at ICLR 2022
PHYSICS INFORMED CONVEX ARTIFICIAL NEURAL NETWORKS (PICANNS) FOR OPTIMAL TRANSPORT BASED DENSITY ESTIMATION
Anonymous authors Paper under double-blind review
ABSTRACT
Optimal Mass Transport (OMT) is well-studied problem with variety of appli- cations in a diverse set of fields, ranging from Physies to Computer Vision and in particular Statistics and Data Science. Since the original formulation of Monge in 1781 significant theoretical progress been made on the existence, uniqueness and properties of the optimal transport maps. The actual numerical computation of the transport maps, particularly in high dimensions, remains a challenging problem. In the past decade several neural network based algorithms have been proposed to tackle this task. In this paper, building on recent developments of input con- neural networks and physics informed neural networks for solving PDE's, we propose a new Deep Learning approach to solve the continuous OMT prob- lem. Our framework is based on Brenier's theorem, which reduces the continuous OMT problem to that of solving nonlinear PDE of Monge-Amin Ampere type whose solution a convex function. To demonstrate the accuracy of our framework we compare our method to several other deep learning based algorithms. We then fo- cus on applications to the ubiquitous density estimation and generative modeling tasks in statistics and machine learning. Finally as example we present how our framework can be incorporated with an autoencoder to estimate an effective probabilistic generative model.
INTRODUCTION
Optimal Mass Transport (OMT) is a well-studied problem with a variety of applications in diverse set of fields, ranging from physics to computer vision and in particular statistics and data science. In this paper propose
framework for the estimation of the solution to the I,2 T.2-optimal transport problem between two densities. Our algorithm, which is based on Brenier's '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 investigations and recent developments in the field.
Density estimation and random sampling: The density estimation problem is to estimate smooth probability density based on a discrete finite set observations. In traditional parametric density estimation techniques, the data is assumed to be drawn from known parametric family of distribu- tions. One of the most ubiquitous parametric techniques is Gaussian Mixture Modeling (McLachlan & Basford 1988).
Nonparametric techniques were first proposed by Fix Hodges (95T) (Silverman & Jones (9899 to move away from rigid distributional assumptions. The most used approach is the kernel density estimation, which dates back to Rosenblatt (1952a) and Parzen (962).. 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 particularly exacerbated in high dimensions with the curse of dimensionality.
Recently, diffeomorphic transformation based algorithms have been proposed tackle this prob- lem (Dinh et al.) 2017) Marzouk et al.] 2016)) Younes (2020; [auer et al., (017). The basic concept of transformation bbasdd algorithms is to find a diffeomorphic mapping between reference proba-

---

Under review as a conference paper at ICLR 2022
bility distribution and the unknown target distribution, from which the data is drawn. Consequently, transformation-base 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-aase approaches is how to estimate and select the trans- formation: 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., (219)) and Glow (Kingma & Dhariwal 2018). For a review of the vast NF literature, we refer to the the overview article (Kobyzev et ata.]] Although these methods have shown to perform well in density estimation applications, the interpretability of the obtained transformation is less clear, c.g. in Real NVP (Dinh et al. 2017). the solution selection is obtained by restricitng the transformations to the class of diffcomorphiimm with triangular Jacobians that are easy to invert, which is closely related to the Knothe-Rosenblatt rearrangement (Knothe 1957 Rosenblatt 1952b).
Optimal mass transport, on the other hand, formulates the transport map selection the minimizer of 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 correspond- ing 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, 201 3a Papadakis] 2015). However, these algorithms do not lead to a continuous transformation map and thus can't '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 finite iifference-bssed methods (Benamou et al. 2010; Benamou & Duval 2019)
Optimal mass transport:
and steepest descent-based energy minimization approaches (Angenent et al.) 20033 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 net- work model. Recent work by Makkuua 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. (2117). 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 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 approx- imate the optimal transport maps. The approach we present relies on Brenier's celebrated theo- rem (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 net- works (Shin et al. 2020), which traditional finite element or finite difference methods for solving
2

---

Under review as a conference paper at ICLR 2022
PDEs do not possess. Brenier's theorem further states that the optimal transport map is given by the gradient of 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 so- lutions are known. Our experiments show that algorithm indeed approximates the true solution well, even in high dimensions. To further quantify the performance of the new framework, we com- pare it two other deep learning-basee algorithms, for which we guided the selection by the results of the recent benchmarking paper by Korotin et al. (202T). in which they evaluate the methods pre- sented in Seguy al. (207); Nhan Dam al. (2019); Taghvaei & Jalali (2019); Makkuva et al. 22020) Liu et (2019); 99; 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 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 demon- strate how our framework can be combined with a traditional autoencoder to obtain a generative
flows: RealNVP (Dinh et al. (2016)) and iResNet (Behrmann et al. (2019). Finally
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, 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 the vast literature on OMT, e.g.. Villani (003 2008). Let Ω be a convex and bounded domain of IR" and let dæ denote the standard measure on R " For simplicity, we restrict our presentation to the set P(22) of all absolutely continuous measures on Z i.e, P(Ω) э = f with f E I.- (2)) such that Sg fdr = 1. From here, on we will identify the measure y with its density function f. We aim to minimize the cost of transporting density to density y using (transport) map T', which leads to the so-called Monge Optimal Transport Problem. We will consider only the special case of quadratic cost function as Brenicr'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 μt, E P(S2). minimize M((T) So ||- ((x) T(=)||ddu(() over all p-maasureable maps T Ω + S2 subject to v Tap. We will call an optimal T an optimal transport map.
Here, the constraint formulated in terms of the push forward action of measurable map T Ω Ω. which is defined via Tzp(B) ppT (B)). for every measurable set A c Ω. By a change of coordinates, the constraint T.μ (f(dd = - gdv can be thus reduced to the equation
f() (I)))) (DTT
(1)
This equation can also be expressed via the pullback action as f1 = :T 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 indeed guaranteed. First, we will introduce more general formulation of the Monge problem, called Kantorovich formulation. Therefore, we define the space of all transport plans II(H, v). i.c., of all measures on the product space Ω x Ω such that the first marginal is μ and the second marginal is Then we have:
3

---

Under review as a conference paper at ICLR 2022
Definition 2.2 (L'-Kantorovich's Optimal Transport Problem) Given J.v E P(D), minimize ((T) = Sax x - d/l"dT(z, y) over all π E [[(44 v)).
Note that the 1,2 Ta-sasssrtteiin metric W2( v) between jr and v 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 μl V E P(2). Then there exists a unique optimal transport plan T (f(f.g), g). which is given by n (x, = (id xT) where T = Vu is the gradient of a convex function u that pushes P forward to v, i.e., (Vu)- = v. The inverse T is also given by the gradient of a convex function that is the Legendre transform of the convex function
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-Ampère equation:
(w)((((w)(z)) - g(Vu(z)) = f(x)
where D is the Hessian, = fdr and = gdz. We obtain equation 2 directly from equation using the constraint that T = Vu as required by Brenier's theorem. We will also refer to this map as the Brenier map. This map is diffeomorphism as it is gradient of a strictly convex function. In general, the Monge-Ampère equation a nonlinear second-orde partial differential equation. If limit ourselves to convex solutions, then the differential equation is elliptic. If the two densities and are smooth (C*s), positive and absolutely continuous with respect to each other then an unique smooth convex solution is guaranteed exist. See e.g. Forzani & Maldonado (2004); [Bakelman (1883); De Philippis & Figalli] (2013) for further details on existence and regularity of solutions this equation. Using methods of classical numerical analysis, Brenier's theorem has been used e.g. in PPeyré 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.
2.2 SOLVING OMT USING PINNS
Physics Informed Neural Networks (PINNs) were proposed by Raissi al, (017)) to solve general nonlinear partial differential equations (PDEs). The basic concept is to use the universal approxi- mation property of deep neural networks to represent the solution of a PDE via a network. Using the automatic differentiation capability of modern machine learning frameworks, a loss function is formulated, such that its minimizer solves the PDE in a weak sense. Such a loss function encodes the structured information, which results in the amplification of the information content of the data the network sees (Raissi et al. 2017). This formulation of the PDE results in good generalization even when only few training examples are available. PINNs have found widespread applications in a short period of time since their introduction. These applications include a wide variety of PDEs. including the Navier-Stokes equation (Jin et al. 2021). nonlinear stochastic PDEs (Zhang et al. 2020) or Allen Cahn PDEs (McClenny & Braga- Neto 2020).
In this work, we propose to use the PINN approach solve the Monge-Ampere equation, as pre- sented in equation g and hence implicitly the OMT problem. This equation has been extensively studied and the properties of its solutions are well established. By Theorem 2.33 we know that the solution is given by convex function u. Recently, Ammss et al. (2017) proposed new architecture of neural networks, Input Convex Neural Networks (ICNNs), that explicitly constrains the function approximated by the network to be convex. Furthermore, ICNNs inherit the universal approximation powers of feedforward networks (Amos et al. 20117 Proposition 2) and thus we can approximate any convex function arbitrarily well by using ICNN architecture of sufficient depth and width. Con- sequently, this architecture naturally lends itself to our proposed application, as directly encodes Brenier's theorem. In the ICNN architecture, the activation function is nondecreasing convex func-
tion and the internal weights (") are constrained to be non-negative; see Figure in Appendix A for schematic description of this class of networks. This architecture is derived from two simple facts: non-negative sums of convex functions are also convex, and the composition of a convex and convex nondecreasing function is again convex.
We are now ready to describe our procedure for solving the continuous OMT problem: given abso- lutely continuous densities μ fdr and v gdx we reduce, using Brenier's theorem, the OMT
4

---

Under review as a conference paper at ICLR 2022
problem the Monge-Ampert equation consider the loss function
2 Using the PINN approach to solve this PDE leads to
|| det det((D(u)) 9(Vu) - fll?..
(3)
which corresponds to a weak formulation of the PDE. To incorporate the convexity of u we construct the solution space using an ICNN of sufficient depth and width. This loss function scales exponen- tially with respect to dimensions of the domain of integration ?.. This stems from the fact that the Monge-Ampere equation contains the determinant of the Jacobian. Using the positive definiteness of the Jacobian (being the Hessian of convex function) one could potentially make use of efficient determinant estimators, such as Stochastic Chebyshev Expansions as developed in Han et al. (2018), which would scales linearly in dimension. Once we have estimated the optimal transport map, the L Wasserstein metric between μ and v is given by
/ / - g(z)dz.
Vu(x)]]
We call this combination of the PINN approach with the ICNN structure, Physics Informed Convex Artificial Neural Networks (PICANNs).
In several applications, we are interested in computing the inverse transformation at the same time. know that this map is also given by the gradient of convex function. Thus, we use a second ICNN (with the same architecture) to compute the inverse optimal transport map (Vv) by solving the minimization problem:
By a duality argument,
Vv(Vu(x)) x1| L2,
(S)
where Vu is the optimal transport map solving (Vu).f = v, c.f. Figure 4 in Appendix
A
2.3 DIFFEOMORPHIC RANDOM SAMPLING AND DENSITY ESTIMATION
In many applications, such as the Bayesian estimation, generating samples from a given density is not trivial. Traditional methods include Markov Chain Monte Carlo methods, c.g., the Metropolis Hastings algorithm (Hastings 1970). An alternative idea is to use diffeomorphic density matching between the given density v and standard density from which samples can be drawn easily. Once have calculated the transport map, standard samples are transformed by the push-forward diffeomorphism to generate samples from the target density v. This approach has been followed in several articles, where the optimal transport map selection was based on both, the Fisher-Raa Rao metric (Bauer et al. 2017) and the -oooehe-ossn rearrangement (Marzouk et al., 2016). The efficient implementation of the present paper directly leads to an efficient random sampling algorithm in high dimensions.
evaluate the density rather easily but
We now recall the density estimation problem using the OMT framework. We are given samples I; drawn from an unknown density 4 E (P(2) that we aim to estimate. The main idea of our algorithm to represent the unknown density as the pullback via (diffeomorphic) Brenier map Vu of a given background density v gdv. i.e., (Vu)* v 1s equivalently using the push forward action as (Vu) +/4 = v, where the pull back and push forward a density are defined in Section 2.1
As we do not have an explicit target density, but only finite number of samples, we need to find replacement for the 2nom used in equation 3 to estimate the transport map Vu. We do this by maximizing the log-likelihood of the data with respect the density (Vu)"v:
N Σ log (dct(D)(u(x)))) - v(Vu(z))))
(6)
Using our PINNS framework, we represent the convex function w again via an ICNN of the same architecture which serves as an implicit regularizer, see Sivaprasad et al. (2021). This equation can be alternatively interpreted as minimizing the empirical Kullback-Leblle divergence between and the pullback of the background density v. To generate new samples from the estimated density, we use the inverse map to transform the samples from the background density v. We calculate the inverse map using second neural network and explicit loss function given by equation 5
5

---

Under review as a conference paper at ICLR 2022
3 EXPERIMENTAL RESULTS
In this section, we will detail our implementation and present several experiments demonstrating both the applicability and accuracy of our framework. In particular, we will compare our results in several experiments to state-of-the-art deep learning- bbased OMT solvers, and we will show that we outperform these methods in terms of accuracy.
3.1 NETWORK DETAILS
As explained in Section 2.2 we use an ICNN architecture for both the forward and the backward map in all of our experiments, c.f. Figure 4 in Appendix [A] As with every deep learning approach, we need to tune the hyperparameters, including width/depth of the network, activation functions and batch size. The width of the network needs to increase with the dimension the ambient space of the data ensure sufficient flexibility. For our experiments in lower dimensions, with three hidden layers with 128 neurons in each layer, whereas for experiments in 30d, we used a first
used network
network with four hidden layers with 128 neurons in each layer. To initialize the network, train the networks to learn the identity transformation, i.e., Vu = 7, which we use as the initial starting point for all our experiments. In all our experiments, 10,000 target samples were used.
To guarantee the convexity of the output function, the activation functions need to be convex and non-decreasing. Since simple ReLUs are not strictly convex and have vanishing second derivative almost everywhere, we experimented with the family of Rectified Power Units (RePUs), the log exponential family and the 'Softplus' function. The Softplus function to the power of a, which defined via Softplus" (z) = (log (1 + exp z))", turned out to be best suited for our applications, where we chose a = 1.1. In particular our experiments suggested that networks with this activation function were able to generalize well to regions where no or only limited training data were available.
3.2 VALIDATION AND COMPARISON TO OTHER METHODS
To demonstrate the accuracy of our implementation, we present two different types of experiments: first will conduct series of experiments in which analytic solutions the OMT problem available, where we compare our results to results obtained with two state-of-the-art deep learning- based OMT solvers: OT-ICNN from Makkuva al. (2020) and W2Gen from Korotin et al. (2019). We choose these two specific algorithms among the available plethora of available OMT solvers based on the recent benchmark paper by Korotin al. (202T). In a second set of experiments we will consider density estimation problems, where we will compare the quality of our results with four other methods: the two OMT solvers from above and two methods from the family of normalizing flows: RealNVP as introduced in Dinh et al. (2016) and iResNet from Behrmann et al. (2019). Approximating the OMT map: Since both OT-ICNN and W2Gen are also based on an ICNN structure, were able to choose the same architecture with same hyperparameters for all three algorithms, thereby ensuring a fair comparison. Nevertheless, we want to emphasize that these parameters could be further fine tuned for all the algorithms and specific experiments to improve the results. Nevertheless, with exception of several experiments for the OT-ICNN solver, we observed good convergence behavior. We present selected convergence graphs in appendix B Furthermore we demonstrate the scalabilty of our algorithm by performing the same experiment in dimensions 2, 5, 8, 15 and 30.
We do not present comparisons of our approach to more traditional OMT algorithms such as the Sinkhorn algorithm (Cuturi) 2013b) or the linear programming approaches (Peyré al. 2019), as these frameworks, although they approximate the OMT distances, do not compute the continuous optimal transport map. which essential for the proposed density estimation. While finite element or finite difference based Monge-Ampert solvers, see e.g. (Benamou & Duval [019 Pacobs & Léger 2020 Benamou & Brenier, 2000). calculate the continuous OMT map, they are not suitable in dimensions greater than two or three.
To quantify the quality of an estimated transport plan T, we calculate the -VV unexplained vari- ance percentage (UVP), which is given by -- (T) = 100 T - T 122= Var(v), where T* denotes the (analytic) optimal transport plan. In appendix D we present in addition the percent- age error between the analytic Wasserstein distance and the approximated distance, which can be viewed as second, albeit coarser, measure of quality. Our first series of experiments is the same as

---

Under review as a conference paper at ICLR 2022
- Unexplained Variance Percentage UV²)UVP)
Dimensions
Experiment
Method
2d
3d
5d
8d
15d
30d
PICANNs
0.004
0.007
0.021
0.034
0.144
0.38
Random Cvx Function
OT-ICNN
0.043
0.052
0.145
0.276
0.746
3.98
W2GEN
0.040
0.043
0.046
0.052
0.150
0.60
PICANNs
0.33
0.15
0.15
0.28
0.30
1.13
Random Gaussian
OT-ICNN
0.28
0.71
0.86
2.38
2.84
2.24
W2GEN
0.17
0.14
0.16
0.23
0.37
0.67
PICANNs
0.29
0.43
0.63
1.61
7.53
21.71
Annulus
OT-ICNN
23.84
9.98
28.21
43.87
44.52
2725.15
W2GEN
1.33
6.86
18.31
20.50
23.28
34.19
We present a comparison between our PICANN approach and OT-ICNN from Makkuva ctt al. (2020) and W2Gen from Korotin et al. (201)). We present the L2-UVP for three experiments. In all these experiments, the source density is the unit Gaussian. The target density in the case of the "Random Cvx Function" experiment is the unit Gaussian deformed by the gradient of a random convex function. In the case "Random Gaussian", the target density is another Gaussian with randomly sampled mean and co-variance matrix. In the third experiment, the target density is the annulus distribution. The results in the first two experiments are averages of 20 realizations.
Table 1:
in the benchmark paper (Koottin et al. 202T): we use the gradient of a random convex function to transport the unit Gaussian to random density. By Brenier's Theorem, the optimal transport map is the gradient of a convex function, this map is the optimal transport map. In each dimen- sion, we repeated this experiment 20 times to compute the error statistics, which are presented in Table T Wheras all three algorithms seem to work well for this experiment, PICANNs consistently outperform the other algorithms. As already observed in Korotin et al. (2021), this experiment is favoring the ICNN architecture, as the true solution was chosen to be of the same nature, which explains the nearly perfect performance of all three algorithms. Next, we turn cases where the analytical solution is known in closed form. The first is the special case where both densities are from family of Gaussians distributions. In that case, the OMT map is simply given by an affine transform and the OMT distance is again given in closed form. We again repeat this experiment 20 times for each dimension, where we generate Gaussian distributions with a random mean and covariances. Here the means are sampled from uniform distribution on [-1,1]. 1]. To construct the random covariance matrices, we recall that we need to enforce the matrix be positive definite and symmetric. Therefore, we generate random matrix A of dimension d x 3d, where d is the dimen- sion of the space and where the entries are i.i.d. chosen from a uniform distribution on [0, 0.75). Then, random covariance matrix can be constructed by letting Σ = AAT (the particular form of almost surely guarantees positive definiteness). The results and comparisons with the other two methods again presented in Table T In general, all three algorithms still lead a good ap- proximation, where one can see that W2Gien and PICANNS are performing significantly better than OT-ICNN. To further validate our algorithm, we choose more challenging problem: an annulus density for which we know the transport map in closed form. The annulus distribution is given by push forward of the Gaussian distribution by a gradient of radially symmetric convex function and is given by f = 9((X' ([XXXX)) X)X) 3(x x)d where g is the unit Gaussian and d is the number of dimensions. The transport map + xxXXX is the gradient of the convex function (X"X) and thus by Brenier's 's theorem the optimal transport map. The results in Tabley] demonstrates that, in particular for this more challenging problem, the proposed PICANNS outperform the other algorithms by orders of magnitudes.
Density estimation comparisons: Next we consider the problem of estimating a continuous density from discrete finite samples. In this part we will compare our method in addition two methods from the normalized flow family, namely RealNVE and iResNet. For both these methods we choose the standard network structure as used in the examples provided by the developers, i.c., for RealNVP we used a fully connected network with 3 hidden layers with 256 neurons in each layer and for iResNets we used fully connected network with 4 hidden layers with 128 neurons in each layer.

---

Under review as a conference paper at ICLR 2022
We want to emphasize that NODE and normalized flows do not solve the OMT problem and con- sequently there is no physically motivated optimality conditions of the transport map generated by these methods. This is also reflected in experiments: in the OMT approach map that is as close as possible to the identity among all transformations that satisfy the required density matching constraint (note, that there is infinite dimensional set of transformations that sat- isty that satisfy this constraints). Consequently the obtained transformation maps for OMT are more regular compared the maps obtained with other density estimation methods such iResNet or RealNVP, cf. Figures 1,2,7 and 8.
search for the
As first example, we consider again the annulus density, that we introduced in the last experiment of the previous section. Figure[[shen shows the comparison of five different methods for estimating the density as well the transport maps. To quantify the ability to approximate the true density we also report the Kl--divergence between the estimated density (i.e., the push-forward of the unit gaussian by the estimated transport map) and the given samples. While W2Gen, RealNVP and iResNet all showed a decent estimation of the annulus density, the quality of their approximations is clearly inferior to our PICANNs approach. Even after extensive parameter tuning we could not get the OT- ICNN algorithm to converge the true density. Finally, note that RealNVP and iResNet do not aim to compute the L optimal transport map. This is exemplified in Figure I where one can see that the estimated maps are very far from the optimal transport map for these
algorithms. A second example in which we consider a nonsymmetric distribution that has been constructed in |Bauer et al.) (2117) can be seen Figure Similar as for the annulus density our PICANNs approach produces the best result. Further examples and comparisons can be found in Appendix C
Ground Truth
PICANNS
W2GEN
OT-ICNN
Real NVP
iResNet
Figure 1: Comparison of different density estimation frameworks. First line: the estimated densities. Second line: the corresponding transport map.
Ground Truth
PICANNs
W2GEN
OT-ICNN
Real NVP
iResNet
Figure 2: Comparison of different density estimation frameworks, First line: the estimated densities.
Second line: the corresponding transport map. Note, that we do not have access to the analytic optimal transport map in this example.
8

---

Under review as a conference paper at ICLR 2022
3.3 GENERATIVE MODELING WITH PICANNS
Finally, we present preliminary results to demonstrate the application of the PICANN approach to develop a generative model. has been well established that for most real world image modeling applications the support of the target probability distribution is lower dimensional manifold Gerber al. (2009) Lee & Verleysen (2007). Autoencoders are very effective in non-linear dimensionality reduction and map high-dimensionas data to a latent space of lower dimensions, which then can be transformed back to the original space using the "decoder" part of the network. One of the major shortcoming of an autoencoder is that the distribution of data mapped in to the latent space we summarizes how the
not know and cannot be assumed to be Gaussian. In Appendix A Fig 5 autoencoder be naturally included in PICANN approach to obtain an efficient generative model for high- sigh-dimensional data.
To demonstrate the generative algorithm in toy example, we consider the MNIST dataset encoded using a simple fully connected autoencoder to a latent space of dimension 2. Figure shows the encoded data points with each class assigned unique color. We train PICANN network of hidden layers with 128 neurons each layer, to learn the forward and the inverse mapping from a Gaussian with the mean and covariance of the encoded data the unknown "latent MNIST distribu- tion". Figure 3 presents comparisons to a similar setup with the four other methods discussed in the previous section. The results show sigmificant differences in the quality of the generated samples, where PICANNs and RealNVP visually lead to the best results. In future work we plan to continue these investigations, by applying the setup to more challenging image sets and carefully comparing the results using state-of-the- -ort quality metrics.
6
4
7
5
2
6
6
2
7
8
9
8
3
6
2
3
7
3
7
/
0
8
72
S
4
9
9
7
5
5
5
5
0
9
3
5
/
7
8
4
/
2
0
3
9
8
4
9
3
5
5
2 O 6
2
8
5
4
6
8
6
2
8
/
7
8
8
5
6
3
4
4
2
6
6
2
8
0
4
8
9
6 7
Ground Truth
PICANNS
W2GEN
OTHICNN
Real NVP
iResNet
Generative model: this figure details the application of our generative framework to the MNIST dataset. Five different algorithms were trained to estimate this distribution and learn the forward and inverse transport map between the encoded 'MNIST Distribution' and a Gaussian. The estimated density can be seen in the first line. The second line shows 25 samples passed through the decoder part of the autoencoder for each of the methods.
Figure 3:
CONCLUSION
In this paper, we the La=sasssrnteiin metric and optimal mass transport (OMT) theory for- mulate a density estimation and generative modeling framework. We develop a new deep learning- based solver for the continuous OMT problem, which rooted in Brenier's celebrated theorem. This theorem allows us to formulate the density estimation problem as solution to a nonlincar PDE -a Monge--Amert equation. Recent developments in deep learning for PDEs, namely PINNS and ICNNs, allow us to develop an efficient solver. We demonstrate the accuracy of our framework by comparing our results to analytic Wasserstein distances. To further quantify the quality of our re- sults we compare them to the results obtained with the two best performing algorithms of the recent benchmark paper for deep learning based OMT solvers (Korotin et al. 2021) )) Our experiments show that approach significantly outperforms these methods in term of accuracy. Finally, present examples of diffeomorphic density estimation within our framework and showcase an example of a generative model.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. In International Confer- ence on Machine Learning, pp. 146-155. PMLR, 2017.
Sigurd Angenent, Steven Haker, and Allen Tannenbaum. Minimizing flows for the monge-
kantorovich problem. SIAM journal on mathematical analysis, 35(1):61-97, 2003.
Ilya J Bakelman. Variational problems and elliptic monge-ampere equations. Journal of differential geometry, 18(4):669 699, 1983.
Martin Bauer, Sarang Joshi, and Klas Modin. Diffeomorphic random sampling using optimal infor- mation transport. In International Conference Geometric Science of Information, pp. 135-142. Springer, 2017.
Jens Behrmann, Will Grathwohl, Ricky TQ Chen, David Duvenaud, and Jörn-Henrik Jacobsen. Invertible residual networks. International Conference on Machine Learning, pp. 573-582. PMLR, 2019.
Jean-David Benamou and Yann Brenier. computational fluid mechanics solution to the monge- kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393, 2000.
Jean-David Benamou and Vincent Duval. Minimal convex extensions and finite difference discreti- sation of the quadratic monge- kantoroiih problem. European Journal of Applied Mathematics, 1016):0041-107,, 2019.
Benamou, Brittany D Froese, and Adam M Oberman. Two numerical methods for the elliptic monge-ampere equation. ESAIM: Mathematical Modelling and Numerical Analysis, 44 (4):737-758, 2010.
Yann Brenier. Polar factorization and monotone rearrangement vector-valued functions. Commu- nications on pure and applied mathematics, 44(4):375-417, 1991.
Guillaume Carlier, Alfred Galichon, and Filippo Santambrogio. From knothe's 's transport brenier's map and a continuation method for optimal transport. SIAM Journal on Mathematical Analysis, 41(6):2554-2576 2010.
Marco Cuturi. Sinkhorn distances: lightspeed computation of optimal transport. In NIPS, volume 2, pp. 4, 2013a.
Marco Cuturi. Sinkhorn distances: Lightspeed computation optimal transport. Advances in neural information processing systems, 26:2292-223 2300, 2013b.
Guido De Philippis and Alessio Figalli. W 2, regularity for solutions of the monge-ampère equa- tion. Inventiones mathematicae, 192(1):55-69, 2013.
Laurent Dinh, Jascha Dickssics,,,,,, and Samy Bengio. Density estimation using real nvp. arXiv preprint arXiv: 100.000005,08803 2016.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. htttttts////onnnni net. n forum? d=#kpbbH91
Evelyn Fix and Joseph Lawson Hodges. Nonparametric discrimination: Consistency properties. Randolph Field, Texas, Project, pp. 21-49, 1951.
Liliana Forzani and Diego Maldonado. Properties of the solutions the monge-ampère equation. Nonlinear Analysis: Theory, Methods & Applications, 57(5-6):815-829, 2004.
Samuel Gerber, Tolga Tasdizen, Sarang Joshi, and Ross Whitaker. On the manifold structure of the space of brain images. In International conference on medical image computing and computer- assisted intervention, pp. 305-312. Springer, 2009.
10

---

Under review as a conference paper at ICLR 2022
Ian J Goodfellow, Jean Pouget-Abadie Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. stat, 1050:10, 2014.
Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. Im- proved training of wasserstein gans. arXiv preprint (xXx::774.00023, 2017.
Insu Han, Haim Avron, and Jinwoo Shin. Stochastic chebyshev gradient descent for spectral opti- mization. In 32nd Conference Neural Information Processing Systems (NIPS). Neural Infor- mation Processing Systems Foundation, 2018.
W Keith Hastings. Monte carlo sampling methods using markov chains and their applications. 1970.
Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbeel. Flow++: Improving flow- based generative models with variational dequantization and architecture design. In International Conference Machine Learning, pp. 7722-2730. PMLR, 2019.
Matt Jacobs and Flavien Léger. A fast approach optimal transport: The back-and-forth method. Numerische Mathematik, 146(3):513-544, 2020.
Xiaowei Jin, Shengze Cai, Hui Li, and George Em Karniadakis. Nsfnets (navier-stokes flow nets): Physics-informed neural networks for the incompressible navier-stokes equations. Journal of Computational Physics, 426: 109951, 2021.
Diederik Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1xl convolutions.
arXiv preprint arXiv:1807.0388 2018.
Herbert Knothe. Contributions to the theory of convex bodies. Michigan Mathematical Journal, 4 (1):39-52, 1957.
Ivan Kobyzev, Simon Prince, and Marcus Brubaker. Normalizing flows: An introduction and review of current methods. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2020.
Alexander Korotin, Vage Egiazarian, Arip Asadulaex, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. arXiv preprint 1uXiv:1999.10082 2019.
Alexander Korotin, Lingxiao Li, Aude Genevay, Justin Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark. arXiv preprint arXiv:2106. 20t.:9944000.01944. 2021.
John Lee and Michel Verleysen. Nonlinear dimensionality reduction. Springer Science Busi- ness Media, 2007.
Huidong Liu, Xianfeng Gu, and Dimitris Samaras. Wasserstein gan with quadratic transport cost. In Proceedings of the IEEE/CVE International Conference on Computer Vision, pp. 4832-4841, 2019.
Clive Loader et Bandwidth selection: classical or plug-in? The Annals of Statistics, 27(2): 415-438, 1999.
Grégoire Loeper and Francesca Rapetti. Numerical solution of the monge-ampèrs equation by newton's algorithm. Comptes Rendus Mathematique, 340(4):319-24,, 2005.
Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pp. 6672- 6681. PMLR, 2020.
Anton Mallasto, Jes Frellsen, Wouter Boomsma, and Aasa Feragen. (q. p)- wassesessnein gans: Com- paring ground metrics for wasserstein gans. arXiv preprint arXiv:1902 03642, 2019.
Youssef Marzouk, Tarek Moselhy, Matthew Parno, and Alessio Spantini. Sampling via measure transport: An introduction. Handbook of uncertainty quantification, pp. 1-41, 2016.
Levi McClenny and Ulisses Braga-Neto. Self-adaptive physics- informee neural networks using a soft attention mechanism. arXiv preprint arXiv:2009.04544 2020.

---

Under review as a conference paper at ICLR 2022
Geoffrey J McLachlan and Kaye E Basford. Mixture models: Inference and applications cluster- ing, volume 38. M. Dekker New York, 1988.
Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Histoire de l'Académie Royale des Sciences de Paris, 1781.
Quan Hoang Nhan Dam, Trung Le, Tu Dinh Nguyen, Hung Bui, and Dinh Phung. Threeplayer wasserstein gan via amortised duality. In Proc. the the 28th Int. Joint Conf. on Artificial Intelligence (IJCAI), 2019.
Nicolas Papadakis. Optimal transport for image processing. PhD thesis, Université de Bordeaux; Habilitation thesis, 2015.
George Papamakarios, Theo Pavlakou, and Jain Murray. Masked autoregressive flow for density estimation. arXiv preprint arXiv:1705.070577 2017.
Emanuel Parzen. On estimation of a probability density function and mode. The annals of mathe- matical statistics, 33()):006--1076, 1962.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning, 11(5-6)3355-077, 2019.
Maziar Raissi, Paris Perdikaris, and George Em Karniadakis. Physics informed deep learn- ing (part i): Data-driven solutions of nonlinear partial differential equations. arXiv preprint arXiv:1711.1056/ 2017.
Murray Rosenblatt. Remarks on multivariate transformation. The annals of mathematical statis- tics, 23(3):470-472, 1952a.
Murray Rosenblatt. Remarks on multivariate transformation. The annals of mathematical statis- tics, 23(3):470-477 1952b.
Vivien Seguy. Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large-scale optimal transport and mapping estimation. arXiv preprint
arXiv:1711.02283, 2017.
Yeonjong Shin, Jerome Darbon, and George Em Kariiadakis. On the convergence and generaliza- tion of physics informed neural networks. arXiv preprint arXiv:2004.01866, 2020.
Bernard W Silverman and M Christopher Jones. E. fix and jl hodges (1951); An important con-
tribution to nonparametric discriminant analysis and density estimation: Commentary on fix and hodges (1951). International Statistical Review/Revue Internationale de Statistique, pp. 233-238, 1989.
Sarath Sivaprasad, Ankur Singh, Naresh Manwani, and Vineet Gandhi. The curious case of convex neural networks. In Nuria Oliver, Fernando Pérez-Cruz, Stefan Kramer, Jesse Read, and Jose A. Lozano (eds.), Machine Learning and Knowledge Discovery in Databases. Research Track, pp- 738-754, Cham, 2021. Springer International Publishing.
Amirhossein Taghvaei and Amin Jalali. 2-wasserstein approximation via restricted convex potentials with application to improved training for gans. arXiv preprint arXiv:1902.07997, 2019.
Cédric Villani. Topics in optimal transportation. Number 58. American Mathematical Soc., 2003.
Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
Laurent Younes. Diffeomorphic learning. Journal of Machine Learning Research, 21(220):1-28, 2020. htRLttt/://3m1 /tt.///ttt////mi.oooooppp/ppesse
dependent stochastic pdes using wttp://rr.oo.ppppp///////// 111111/11111.11111 10151.tmml Computing, (11)))633-A663 2020.
Dongkun Zhang, Ling Guo, and George Em Karniadakis. Learning in modal space: Solving time- neural networks. SIAM Journal on Scientific
12

---

Under review as a conference paper at ICLR 2022
A NETWORK STRUCTURES
In this section we present schematics our PICANNs approach (Figure 4 and how can be in- cluded in generative model (Figure 5
(g
(2
e
[Ides(D((u)) [VVaj a(Va)-ff. -1f-
Wy (9)
3x
ICNN (x 0)
WoP
W2
3x
[VN[Vu(a]] =le
6
W,,00
W (0)
ICNN Bt 01)
Figure 4: PICANN architecture. We present how a combination of two ICNN networks can be used learn the forward and the inverse map between two distributions. Both these networks trained independently with their respective loss functions. The inverse network uses the gradient of the output of the first network as its input.
13

---

Under review as a conference paper at ICLR 2022
81
02
Autoencoder
og
Wy
g
[]des (4) o( -713
w,00 00)
W
8x
ICNN (x; 0)
[Vu[Vu(x)) -
0
Wy
70)
ICNN 8x 01)
Figure 5: How a combination of an autoencoder with our PICANN approach can be used to develop generative model. Note how the latent space the autoencoder becomes the input to the PICANN network. In such a setting, the PICANN estimates the latent space density and samples from the estimated distribution. Using random samples from this distribution, one can pass them through the decoder to generate new samples.
B CONVERGENCE OF EXPERIMENTS
In this appendix (Figure we present the convergence plots for a several experiments of Table should be noted that PICANNs and W2GEN are pure minimisation algorithms while OT-ICNN is min-max algorithm and hence the covergence plots for OT-ICNN differ from the convergence plots for the other two algorithms. These convergence plots show that PICANNs and W2GEN converge for all the experiments presented here while the OT-ICNN, as remarked in Section 3.2 does not converge in the case of the annulus experiments. This is also the case for the annulus experiments in we also present the average timing for 100 Epochs
higher dimensions (not shown here). table 2 for these three algorithms in different dimensions.
Method
2d
3d
5d
8d
15d
30d
PICANNs
90s
95s
130s
165s
233s
310s
W2Gen
110s
110s
112s
115s
118s
120s
OT-ICNN
180s
180s
185s
186s
197s
210s
Table 2: Computational cost for 100 epochs in various dimensions.
14

---

Under review as a conference paper at ICLR 2022
13
24
19
37
(a) PICANNs Annulus in 2d
(b) W2GEN Annulus in 2d
(c) OT-ICNN Annulus in 2d
25
15
(d) PICANNs RG in 2d
(e) W2GEN RG in 2d
( OT-ICNN RG in 2d
Figure 6: In this figure we present the convergence plots for PICANNs, W2GEN and OT-ICNN. In row 1 present convergence plots for the annulus experiment in 2d for PICANNs, W2GEN and OT-ICNN in Panels (a), (b) and (c) respectively. In row we present convergence plots for the gaussian experiments performed using again PICANNs, W2GEN and OT-ICNN in Panels (d), (e) and (f) respectively. As the OT-ICNN algorithm a min-max algorithm these convergence plots have be read slightly differently. For the gaussian experiment, the loss starts with small negative values and climbs up to small positive quantity but around the 50 epoch mark it has converged about 0 and just bounces that neighbourhood thereafter. The OT-ICNN algorithm does not converge for the annulus experiment, as remarked in Section 3.2
C FURTHER DENSITY ESTIMATION EXAMPLES AND COMPARISONS
Ground Truth
PICANNS
W2GEN
OT-ICNN
Real NVP
iResNet
Figure 7: Comparison of different density estimation frameworks for a Gaussian mixture model. First line: the estimated densities. Second line: the corresponding transport map. Note, that do not have access to the analytic optimal transport map in this example.
15

---

Under review as a conference paper at ICLR 2022
Ground Truth
PICANNs
W2GEN
OT-ICNN
Real NVP
iResNet
Figure 8: Comparison of different density estimation frameworks. First line: the estimated densities. Second line: the corresponding transport map. Note, that do not have access the analytic optimal transport map in this example.
D ERROR BETWEEN TRUE AND APPROXIMATED WASSERSTEIN DISTANCE
In this appendix present the percentage error between the theoretical and the approximated Wasserstein metric for three different algorithms: PICANNs, OTICNN from Makkuva et al. (2020)
and W2Gen from Korotin al. (2019). The exact experimental setup and the network details
described in section
Avg error between frue and approximated Wasserstein distance
PICANNs
0.12
0.05
0.03
0.03
0.02
0.04
Random Cvx Function
OT-ICNN
0.10
0.10
0.08
0.07
0.10
0.09
W2GEN
0.09
0.067
0.03
0.04
0.06
0.52
PICANNs
1.56
0.88
0.35
0.21
0.19
0.15
Random Gaussian
OT-ICNN
1.66
1.40
0.93
0.95
0.27
0.31
W2GEN
1.59
0.75
0.41
0.25
0.35
0.19
PICANNs
5.37
1.25
1.81
1.56
7,44
5.07
Annulus
OT-ICNN
6.54
25,89
33.50
20.88
2.03
36.13
W2GEN
12.36
19.39
8.10
3.34
0.96
0.66
Table 3: We present a comparison between our PICANN approach and Makkuva et al. 2020) and Korotin al. (2019). In this table, we present the L2-UVP and the percentage error between the theoretical W2 metric and the approximated W2 metric for three experiments. In all these experi- ments, the source density the unit Gaussian. The target density in the case of the "Random Cvx Function" experiment is the unit Gaussian deformed by the gradient of random convex function. In the case "Random Gaussian". the target density is another Gaussian with randomly sampled mean and co-variance matrix. In the third experiment, the target density is the annulus distribution. The results the first two experiments are averages 20 realizations.
16


Reviews:
[{"review_id": "spqetNouau-", "reviewer": "Reviewer_44BN", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength: + The probabilistic universal approximation framework can be much more efficient than the finite element or difference methods which is known to have exponential scaling with respect to dimensions.\n+ The paper provides sufficient details in terms of network, hyperparameters, evaluation methods in the experimental section.\n\nWeakness: - The main required definitions are provided, however, the main relevant technical concepts are explained heuristically. For example, while relevant background for optimal transport theory is provided in detail, there is little to no background on partial differential equations.\n- The main loss functions (4) and (5) proposed do have exponential scaling with respect to dimensions of the domain of integration \\Omega.\n- Density estimation experiments are preliminary.\n", "summary_of_the_paper": "The paper proposes a new physics informed neural architecture for efficiently learning and optimizing optimal transport maps using techniques from partial differential equations.", "summary_of_the_review": "Justification: The paper makes a solid contribution in terms of technical content. The main high level idea is that solving optimal transport maps under some conditions can be posed as a PDE. While this is not something new, for instance, sinkhorn algorithm does the same -- turn a nonsmooth optimization problem, a linear program, into a smooth version, as given in equations (4) and (5). However, many technical details regarding the algorithm are missing -- (i) what is the overall procedure?; (ii) near equation (7), there is some discussion regarding representing u again using a neural network, is the overall architecture different from what is shown in Figure (1)? It seems like some of these questions regarding how and what PDE is exactly being solved can be answered by reading various other papers cited in the experiments, but it is not clear from the submission. In this sense, I feel like the paper spends too much time on background material (which the reader can pick up from any standard book on Applied PDE, for example, Introduction in Optimal\nTransport for Applied Mathematicians by Filippo Santambrogio, 2015) rather than focusing on the key insights of the paper. \n\nThe network details are provided in abundant details for verification of the main idea. However, the main task density estimation which the paper mentions quite in detail about in the beginning of the paper is only minimally experimented. For example, the paper claims that finite elements solvers are slow in high dimensions. However, the paper contains no comparisons with these methods on the problems considered in the paper, which are also low dimensional.\n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "fvBfUC5EI9V", "reviewer": "Reviewer_U1ba", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The proposed method in this paper has the potential to be a very useful method for solving the continuous optimal transport problem using DNNs.\nIn addition, the manuscript is clear and well constructed.\nOn the other hand, the effectiveness of the method for application problems in machine learning cannot be confirmed because there are few implications about real applications and few experiments.\nIn order for the paper to be accepted at this conference, it might be necessary to apply the method to a real problem such as the image generation problem of natural images, and then conduct experiments to compare the method with existing methods which solved the application problem.\nIn addition to that, the following points were unclear and I would like to ask for your response.\n\nDoes ICNN have the ability to represent all convex functions?\n\nOn page5, it says \"L^2-norm used in equation 5\", but isn't it a mistake of \"equation 4\"?\n\nIn the experiment of method comparison, is there any evidence that the parameters of DNN in the previous study are optimized? (If there is no evidence, it would be better to specify the possibility that it is not optimized?　That would be positively evaluated as a scientific paper, as it would present the limits of verification).", "summary_of_the_paper": "This paper proposes to solve the continuous Optimal Mass Transport problem as a convex mapping function estimation problem based on Brenier's theorem using input convex neural networks.\nThe proposed method is compared with other DNN-based methods in the ubiquitous density estimation and generative modeling tasks in statistics, and good results are obtained.", "summary_of_the_review": "The proposed method in this paper has the potential to be a very useful method for solving the continuous optimal transport problem using DNNs.\nIn addition, the manuscript is clear and well constructed.\nOn the other hand, the effectiveness of the method for application problems in machine learning cannot be confirmed because there are few implications about real applications and few experiments.\nIn order for the paper to be accepted at this conference, it might be necessary to apply the method to a real problem such as the image generation problem of natural images, and then conduct experiments to compare the method with existing methods which solved the application problem.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "Eyi67KNqhv9", "reviewer": "Reviewer_ecZV", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is in general well-written and the idea is clearly presented. Yet, I have a few concerns as listed below:\n\n1. The authors stated that \"We will, however, see, that in our situation, i.e., for absolutely continuous measures, the existence and uniqueness is indeed guaranteed\". This refers to the Brenier's theorem concerning that a unique optimal transport map exists which is given by the gradient of a convex function. This consequently leads to the Monge-Ampere equation that help formulate the loss function in the proposed work. Yet, the physics informed neural networks adopted can only be shown to be a convex function, but there is no guarantee that the optimized network is indeed the solution of the Monge-Ampere equation. Hence, the uniqueness of the map does not seem to be guaranteed. If the alternative is true, a proof is needed here.\n\n2. The presentation of the density estimation in Section 2.3 is a bit confusing. What does it mean to be \"or equivalently $\\bigtriangledown u)_* \\mu = \\nu$\"? A bigger question is whether the proposed work can be used in density estimation, given that the uniqueness of the mapping is not guaranteed as in Comment 1 above.\n\n3. Although literature review covers other density estimation techniques such as normalizing flow, there is no comparison with any of these algorithms in Section 3.3. The authors stated in Section 1 that \"Although these methods (normalizing flow) have shown to perform well in density estimation applications, the interpretability of the obtained transformation is less clear\". But it is not clear to me how the interpretability of the proposed method is clearer, especially given the fact that the computed density may not be exact/correct (due to the lack of the proof of the uniqueness of the obtained approximate mapping), while normalizing flow admits exact density estimation.", "summary_of_the_paper": "The paper presented an input convex neural networks (ICNN)-based methods to approximate the optimal transport maps. The learning objective is based on the Monge-Ampere equation, i.e. zero loss corresponds to the solution to the Monge-Ampere equation. Physics Informed Neural Networks (PINNs) are adopted to minimise the loss. The author(s) also extended the work to density estimation. The problem is interesting as estimation of optimal maps has a wide range of applications. Experiment results look promising when compared with other PINNs in one experiment setup, but only qualitative results are provided for the density estimation experiment with no comparison of other baselines.", "summary_of_the_review": "The paper provides an interesting approach to incorporate ICNN and Brenier's Theorem to approximate optimal transport between distributions. Experiment results are encouraging. But there is no proof to support the claim on the uniqueness of the mapping. Therefore, its application on density estimation is also not fully supported.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "6PCBxHdjiYu", "reviewer": "Reviewer_U76A", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper combines a range of interesting ideas to solve density estimation problems. The authors reformulate the original OMP problem as a special PDE. In turn, PINNs can be used to approximate the solution of this special PDE. In addition, the authors use input convex neural networks (ICNNs) to  incorporate the Brenier potential into the PINN. This approach bypasses some of the challenges of directly minimizing the original OMP problem.\n\n* _Originality:_ The presented ideas in this paper are original and well suited for this conference. I not aware of another paper that uses PINNs for density estimation. \n\n* _Quality:_ The proposed framework seems plausible and is demonstrated on a few experiments. There is no theory that backups some of the claims, and the amount of experiments is on the lower end of the spectrum for an empirical ML conference paper. Further, it would help to better discuss the assumption that have been made and some of the limitations. I assume that the authors have a fair understanding for these details and they should make better use of the unlimited space of the appendix to provide additional discussions and experiments. \n\n* _Clarity:_ The proposed ideas are well organized and clearly presented. The detailed discussion of the background materials helped reading the paper. \n\n* _Significance:_ This paper is of interest for two communities within ML. First, it is of interest for the PINN community and the paper has the potential to motivate future research in the intersection of PINNs and density estimation. On the other hand, density estimation is an important topic in ML and the presented approach performs well on the presented tasks.However, it is not clear whether the presented approach is competitive with other state-of-the-art methods in this area, or whether it just does better on the selected tasks compared to 2 baselines. A better discussion of the advantages and limitations and additional baselines would help to strengthen the paper and improve its significance. \n\n\nHere are a few comments that might help to improve the paper:\n\n* The weakness of this paper is the experimental section. Only two other algorithms are used for comparison. The authors state that these two algorithms performed best on their benchmark problems, still it would be interesting to show results for other algorithms that have been tested. Also, how does your approach compare to state-of-the-art GANs or neural ordinary differential equations (NODEs) for density estimation methods? At least, it would be helpful to discuss some qualitative advantages / disadvantages of your proposed approach in comparison to other methods such as GANs or NODEs. In which situation should one use your approach, and which situation does your approach fail? \n\n* You state that you only consider the simple case of a quadratic cost function. Does Brenier's Theorem apply to other cost functions?\n\n* You state that ICNNs introduce implicit regularization. Can you provide theory to back this claim up, or is this only an educated guess?\n\n* The results in Figure 6 look interesting. Why don't you include these results in the main text? I would rather move the network details to the appendix.   \n\n* Can you comment on the computational costs of your framework as compared to the other methods that you used for comparison.\n\n* On page 2: I assume that there is some typo in:  'the convex function itself should be indenity helps in avoiding'.\n\n* On page 9: word repetition in 'a density estimation estimation'.\n\n* The quality of the figures is poor. It would be nice to better format the figures so that they eat up less space and use vector graphics instead of low-resolution pngs. ", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper presents interesting and novel ideas for density estimation, and to the best of my knowledge this is the first paper that uses PINNs for density estimation. The paper is clearly written and provides a good overview of related work and background to put the proposed ideas into context. However, since this paper provides only empirical results and no theory, I only think that the paper is slightly above the acceptance threshold. Additional results would help to make the story more compelling, and given additional results I am happy to reconsider my rating. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work proposes to define densities via the pushforward of a base density through the gradient field of a convex potential as studied in OT theory and, in particular, inspired by Brenier's theorem.

More concretely, it proposes to use ICNNs to parametrize the convex potentials and considers two mechanisms to match a target density: 1) with a known (normalized) target approximately solve the Monge-Ampere equation via optimization; 2) with only samples available, they propose to use the maximum-likelihood approach.

While the paper is overall well-written, the idea is very close to existing work that was not mentioned or discussed in the paper. The paper would benefit from a substantial revision to incorporate the missing references and emphasize the relative novelty.
Decision: Reject