Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks
Authors: David Alvarez-Melis
Yair Schiff
Youssef Mroueh
Venue: Accepted by TMLR ()

Abstract: Gradient flows are a powerful tool for optimizing functionals in general metric spaces, including the space of probabilities endowed with the Wasserstein metric. A typical approach to solving this optimization problem relies on its connection to the dynamic formulation of optimal transport and the celebrated Jordan-Kinderlehrer-Otto (JKO) scheme. However, this formulation involves optimization over convex functions, which is challenging, especially in high dimensions. In this work, we propose an approach that relies on the recently introduced input-convex neural networks (ICNN) to parametrize the space of convex functions in order to approximate the JKO scheme, as well as in designing functionals over measures that enjoy convergence guarantees. We derive a computationally efficient implementation of this JKO-ICNN framework and experimentally demonstrate its feasibility and validity in approximating solutions of low-dimensional partial differential equations with known solutions. We also demonstrate its viability in high-dimensional applications through an experiment in controlled generation for molecular discovery.

Published in Transactions on Machine Learning Research (07/202)
Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks
David Alvarez-Melis Microsoft Research
daalvare@miierossfl.com
Yair Schiff IBM Watson
yaär.schiffükbm..
Youssef Mroueh
mrooch Crrouch@uu.com ibm. com
IBM Research Al
Reviewed on OpenReview:
Abstract
Gradient flows are a powerful tool for optimizing functionals in general metric spaces, including the space of probabilities endowed with the Wasserstein metric. A typical approach to solving this optimization problem relies on its connection the dynamic formulation of optimal transport and the celebrated Jordan- Kintean-KKndrreeeer-Otttttoo (JKO) scheme. However, this formulation involves optimization over convex functions, which is challenging, especially in high dimensions. In this work, we propose an approach that relies on the recently introduced input-convex neural networks (TCNN) to parametrize the space of convex functions in order to approximate the JKO scheme, as well as in designing functionals over measures that enjoy convergence guarantees. We derive a computationally efficient implementation of this JKO-ICNN framework and experimentally demonstrate its feasibility and validity in approximating solutions of low-dimensional partial differential equations with known solutions. We also demonstrate its viability in high-dimensional applications through an experiment in controlled generation for molecular discovery.
1 Introduction
Numerous problems in machine learning and statistics can be formulated as finding a probability distribution that minimizes some objective function of interest. One recent example of this formulation is generative modcling, where one seeks to model a data-generating distribution Pdata by finding, among a parametric family po; the distributtin that minimizes some notion of discrepancy to Pdata: i.e., ming D(pe, Pdata). Different. choices of discrepancies give rise to various training paradigms, such as generative adversarial networks (Goodfellow et al., 2014) (Jensen-Shannon divergence), Wasserstein GAN (Arjovsky al., 2017) (I-Wasserstein distance) and maximum likelihood estimation (KL divergence) (Murphy, 2012; Rezende & Mohamed, 2015; Kingma& & Welling, 2013). In general, such problems can be cast as finding p² argmin, F(p), for functional F' on distributions.
Beyond machine learning and statistics, optimization on the space of probability distributions prominent in applied mathematics, particularly in the study of partial differential equations (PDE). The seminal work of Jordan et al. (1998), and later Otto (2001), Ambrosio et. al. (2005), and several others, showed that many classic PDEs can be understood as minimizing certain functionals defined on distributions. Central to those works is the notion of gradient flows on the probability space endowed with the Wasserstein motric. Jordan, Kinderlehrer, and Otto (1998) set the foundations a theory establishing connections between optimal transport, gradient flows, and differential equations. In addition, they proposed a general iterative method, popularly referred was later extended to more general PDEs and in turn to more general functionals over probability space (Ambrosio et al., 2005). The JKO scheme can be seen generalization of the implicit Euler method on the
the JKO scheme, to solve PDEs of the Fokker-Planck type. This method

---

Published in Transactions on Machine Learning Research (07/2022)
probability space endowed with the Wasserstein metric. This approach has various appealing theoretical convergence properties owing to a notion of convexity of probability functionals, known as geodesic convexity (see Santambrogio (2017) for more details).
Several computational approaches to JKO have been proposed, among them an elegant method introduced in Benamou et al. (2014) that reformulates the JKO variational problem on probability measures as an optimization problem on the space of convex functions. This reformulation is made possible thanks to Brenier's Theorem (Brenier, 1991). However, the appeal of this computational scheme comes at price: computing updates involves solving an optimization over convex functions at each step, which is challenging in general. The practical implementations in Benamou et al. (2014) make use of space discretization to solve this optimization problem, which limits their applicability beyond two dimensions.
In this work, we proposc a computational approach to the JKO scheme that is scalable in high-dimensions. At the core of our approach are Input-Convex Neural Networks (ICNN) (Amos el al., 2017), a recently proposed class of deep models that are convex with respect to their inputs. We use ICNNS to find parametric solutions to the reformulation of the JKO problem as optimization on the space of convex functions by Benamou et al. (2014). This leads to an approximation of the JKO scheme that we call JKO-ICNN. In practice, we implement JKO-ICNN with finite samples from distributions and optimize the parameters of ICNNs with adaptive gradient descent using automatic differentiation.
To evaluate the soundness of our approach, we first conduct experiments well-known PDEs in low dimensions that have exact analytic solutions, allowing us to quantify the approximation quality of the gradient flows evolved with our method. We then use our approach in high-dimensional setting, where we optimize a dataset of molecules to satisfy certain properties such as drug-likeness (QED). The results show that our JKO-ICNN approach is successful at approximating solutions of PDEs and has the unique advantage of scalability in terms of optimizing generic probability functionals on the probability space in high dimensions. When compared to direct optimization methods or particle gradient flows methods, JKO-ICNN has the advantage of computational stability and amortization of computational cost, since the maps found while training JKO-ICNN on one sample from a distribution po generalize at transporting new sample unseen during the training at no additional cost.
2 Background
Notation Let X' be a Polish space equipped with metric d and P(X) be the set of non-negative Borel measures with finite second-order moment on that space. The space P(X)) contains both continuous and discrete measures, the latter represented as an empirical distribution: ΣNV, p:oz,, where 5, is a Dirac at position x E X. For a measure and measurable map T : X' + ,,, we use Tep to denote the fush-forward measure, and Jy the Jacobian of T. For a function u : X' - R, Vu is the gradient and Hx(x) the Hessian. V. denotes the divergence operator. For a matrix A, |A| denotes ils determinant. When clear from the context, we use P interchangeably to denote measure and its density.
Gradient flows in Wasserstein space Consider first a functional F X R and point zo X. A gradient flow is an absolutely continuous curve xxt) that evolves from zo in the direction of steepest descent of F. When X' is Hilbertian and F is sufficiently smooth, its gradient flow can be expressed as the solution of a differential equation z1(1)= = (()()),) with initial condition x00)= 20-
Gradient flows can be defined in probability space too, as long as a suitable notion of distance between probability distributions is chosen. Formally, let us consider P(X') equipped with the p- p-Wasserstein distance, which for measures a, E P(X) is defined as:
()( min xEII(a, / |11 (1
(1)
Here IT(a, () is the set of couplings (Iransportation plans) between a and B, formally: IT (8) & (x E X(XxXX PISTT a, PayT = B]. Endowed with this metric, the Wasserstein space Wp(X)- = (P(X), W,) is a complete and separable metric space. In this case, given a functional in probability space F: P(X) + R, its gradient flow in (X) is curve p(t) R, + P(X) that satisfies dep(t)- (Vt)((p(t). Here Vw, is a
2

---

Published in Transactions on Machine Learning Research (07/2022)
=- natural notion of gradient in (x( given by: Vw, F(p) -V (pV aF the functional F. Therefore, the gradient flow of F solves the following PDE:
where oF 6p
is the first variation of
6p
Oep(t) - (2(V
V.
(2)
also known as continuity equation.
3 Gradient flows via JKO-ICNN
In this section introduce the JKO scheme for solving gradient flows, show how to cast it optimization over convex functions, and propose method to solve the resulting problem via ICNN parametrization.
3.1 JKO scheme measures
Throughout this work we consider problems of the form minpe F()), where F: P(X) + functional over probability measures encoding some objective of interest. Following the gradient flow literature (e.g., Santambrogio (2015; 2017)), focus three quite general families of functionals:
F(p) = / f(o(x)) dz,
1 v(p) = f V(x) dp, w(p) = 2 // d'd dp(x'), 2
(3)
where f: IR -> R is convex and superlinear and V,W: : X -> R are convex and sufficiently smooth. These functionals are appealing for various reasons. First, their gradient flows enjoy desirable convergence properties, as we discuss below. Second, they have a physical interpretation as internal, potential, and interaction energies, respectively. Finally, their corresponding continuity equations (2) turn out recover various classic PDEs (sce Table for oquivalences). Thus, in this work, we focus on objectives that can be writlen as linear combinations of these three types of functionals.
Table 1: Equivalence between gradient flows and PDEs. In each case, the gradient flow of the functional F(p) in Wasserstein space in the rightmost column satisfies the PDE in the middle column.
Class
PDE d.p-
Flow Functional F(p)
Heat Equation
Ap
I p(z) log p(x) de
Advection
V (VVV)
[V()ddp(a)
Fokker-Planck
Ap+ V. (PVV)
p(x)) log p(x) dr ((())
Porous Media
APp")+ + V (VVV)
m- J pl de + (V2) (z) dp(x) + [f(e(x)) ( do(a))
Adv.+Dif..-nnts Inter.
((VF(() +VV+ (VWW)))))))))
For a functional F of this form, it can be shown that the corresponding gradient flow defined in Section 2 con- verges exponentially fast to unique minimizer (Santambrogio, 2017). This suggests solving the optimization problem min, F(p) by following the gradient flow, starting from some initial configuration Po- A convenient method to study this PDE is through the time discretization provided by the Jordan Kinderlehrer Otto (JKO) iterated movement minimization scheme (Jordan et al., 1998):
P argmin F()) + W2(e.p)). (xW2(x))
1
(4)
2T
where T > 0 is a time step parameter. This scheme will form the backbone of our approach.
3.2 From measures to convex functions
The general JKO scheme (4) discretizes the gradient flow (and therefore, the corresponding PDE) in time, but it is still formulated on on-ootentallly difnnteldimensionaal, and therefore intractable- probability space P(X)). Obtaining an implementable algorithm requires recasting this optimization problem in terms of a space that is easier handle than that of probability measures. As first step, do so using convex functions.
A cornerstone of optimal transport theory states that for absolutely continuous measures and suitable cost functions, the solution of the Kantorovich problem concentrates around deterministic map T (the
3

---

Published in Transactions on Machine Learning Research (07/2022
Monge map). Furthermore, for the quadratic cost, Brenier's theorem (Brenier, 1991) states that this map is given by the gradient of a convex function u, i.e., T(x) = Vu(x). Hence given a measure o, the mapping u E cvx(X) + (Vu)ya E P(X) can be seen as a paramctrization, which depends on a, of the space of probabilities (McCann, 1997). We lurthermore have for any u E cvx(X):
(a (Vu)ga) 1 [Vu(()) - x112 do.
Using this expression and the parametrization (Vu)apī in Problem (4), we obtain a reformulation of Wasserstein gradient flows as optimization over convex functions (Benamou al., 2014);
" E argmin F(((U))pF) + $ [Iv(()) - x112 dp..
(6)
uEcvx(a)
which implicitly defines a sequence of measures via pb+ = (Vu(+1) (pf)- For potential and interaction functionals, Lemma 3.1 shows that the first term in this scheme can be written in a form amenable optimization on
Lemma 3.1 (Potential and Interaction Energies). For the pushforward measure p (Vu)sPt, the functionals V and W can be written as:
V(p) / (V (p)(z)
W(p) 2 // W(Vu(s)- - Vu(()) (p(())
Crucially, Pr appears here only as the integrating measure. We will exploit this property for finite-sample computation in the next section. In the case of internal energies F. however, the integrand itself depends on Pt. which poses difficulties for computation. To address this, we start in Lemma 3.2 by tackling the change of density when using strictly convex potential pushforward maps:
Lemma 3.2 (Change of variable). Given strielly conver WECvx(X), Vu invertâble, and (Vu)-i_Vu', where x² is the conver conjugate of u, (v) =uup=uuuedomttu (z, y) u(z). Given measure a with density Par the density P3 of the measure B (Vu), is given by:
pa(w)= = Pa
o (Vu) (V))-t(y)
Pa o Vy)
H
H
In other words log (pa(v)) log (po (Vu" (Vu((u)) log( (Vu' (v))))- Iterating Lemma 3.2 across time in the JKO steps obtain:
Corollary 3.3 (Iterated Change of Variables in JKO). Assume po has a densily. Let Tí, Vuf Vui, where of optimal convez potentials in the JKO sequence that assume are strictly convez. We use the convention Tí (z) x. We have (Ti=))-- (uu))* * .... V(ug)* where (uf)* is the conver conjugate of ut. At time l of the JKO ilerations we have: (x) (T),)()), and therefore:
log (pī) log (po) - Σ log( (H, ()) -T.
From Corollary 3.3, we see that the iterates the JKO scheme imply change of densities that shares similaritics with normalizing flows (Rezende & Mohamed, 2015; Huang al., 2021), where the depth of the flow network corresponds to the time in JKO. Whereas the normalizing flows of Rezende & Mohamed (2015) draw connections to the Fokker-Planck equation in the generative modeling context, JKO is more general and allows for rigorous optimization of generic functionals on the probability space.
Armed with this expression of pf. now write F in terms of the convex potential
4

---

Published in Transactions on Machine Learning Research (07/2022)
Lemma 3.4 (Internal Energy). Let Pt be the measure at time t of the JKO iterations. In the notation of Corollary 3.9, for the measure p = (Vu)aPt = (Vu o Tia);Po, we have:
F(e) ((((=)) IJT (x) dr,
(8)
where E(x) (o(x)(((H,,Ti.a(()) (x)())))) Assuming po > 0, we have: F(p) Exrpat ((x)/((a))
3.3 From convex functions to finite parameters
Solving Problem (6) requires: (i) tractable parametrization of cvx(X), the space of convex functions, (ii) method to evaluate and compute gradients of the Wasserstein distance term, and (ii)) a method to evaluate and compute gradicnts of the functionals as expressed in Lemmas 3.1 and 3.4.
For (i), we rely on the recently proposed Input Convex Neural Networks (Amos et al., 2017). See Appendix B.1 for a background on ICNN. Given po = Σi=1 dz, we solve for = I...7:
0[+1 argmin L(O),
(ICNNN))(
L(0) A F(((2 + 2T / [V - =0112 dp:
1
where ICN is the space of Input Convex Neural Networks, and the 0 denotes parameters of the ICNN. Problem (9) defines JKO sequence of measures via Pī+1 (Vzuor,, ( (p.). We call this iterative process JKO-ICNN, where each optimization problem can be solved with gradient descent the parameter space of the ICNN, using backpropagation and automatic differentiation.
For (ii), we note that this term can be interpreted as an expectation, namely, 27 Ez=pp [o)) - x112, so we can approximate
with finite samples (particles) of Pf obtained via the pushforward map of previous 1 Ei= [z,,) - x:ll2: {x ~ P.. Finally, for (iii) we point clouds in the JKO sequence, i.e., using 2rm first note that the V and W functionals can also be written as expectations over PT:
= E (V)nu)))),
(10)
xmpT
(((V2u0):P))
1
E 2.y=p
W(V vz))(y)).
Thus, as long as we can parametrize the functions V and in a differentiable manner, we can estimate the value and gradients of these two functionals through finite samples too. In many cases V and W will be simple analytic functions, such as in the PDEs considered in Section 5. To model more complex optimization objectives with these functionals we can leverage ICNNs once more to parametrize the functions V and W as neural networks in way that enforces their convexity, This is what we do in the molecular discovery experiments in Section 6.
Particular cases of internal energies Equation (8) simplifies for some choices of f, e.g., for f(t) = t tog (which yields the heat equation) and strictly convex u we get:
= / (H.. (x) HH..
(z) log
(z(
f loo Hin (x) dr
(x) H dx = -
(11)
where we drop t from the notation for simplicity. This expression has a notable interpretation: pushing forward measure Pe by Vug increases its entropy by a log-determinant barrier term on ug's Hessian. Note that only the second term in equation (11) depends on uæ, hence (=F((V -v Ea log HHn (z)]-
Since the latter can be approximated aa before- by empirical expectation, it can be used as surrogate objective for optimization.
Another notable case is given by f(t) -"
t" m > 1, which yields nonlinear diffusion term as in the porous medium equation (Table 1). Tn this case, equation (8) becomes F(p) - F ((x))m= whose gradient with rospect to is:
E log log log Hu, (TTa(<))
(12)
prppo
Table 3 in the Appendix collects all the surrogate objectives described in this section.
5

---

Published in Transactions on Machine Learning Research (07/2022)
Algorithm 1: JKO-ICNN: JKO variational scheme using input convex neural networks
Inputs:
Functional to optimize: F(p)
Outer loop (JKO scheme) parameters: step size > 0 and number of steps T E N Inner loop (ICNN fitting) parameters: learning rate 1 > 0 and number of iterations nu E warmslart boolean: whether to initialize ICNN weights from prior iteration's 's solution Initialize ug ICNN, parameters of ICNN, P N Ei=15,, 8z,- for / = 0 to - do if not warmstort: 0 InitializeWeightss for = to nhe do {JKO inner loop: Updating ICNN }
L(Ø)= = ([(Vuu))#P; + 27 Eppllz [o
- Adam(0, T/, VoLg)
end for
{JKO outer loop: Updating point cloud (measures)
P[+1 =V((u)
end for Output: PT
Implementation and practical considerations We implement Algorithm 1 in PyTorch (Paszke et al., 2019), relying on automatic differentiation to solve the inner optimization loop. The finite-sample approxima- tions of V and W (equation (10)) can be used directly, but the computation of the surrogate objectives for internal energies FF (equations (11) and (12)) require computing Hessian log-determinants- peseemmmnnnnts-pobbbiiiiit in high dimensions. Thus, we use a stochastic log-trace estimator based on the Hutchinson method (Hutchinson, 1989), as used by Huang et al. (2021) (see Appendix B.4 for details on implementation and complexity of this method). To enforce strong convexity needed (e.g., for evaluation), we estimate the true internal energy functional JF by using Corollary 3.3 to compute densities. Appendix B.5 provides full implementation details.
the ICNN, clip its weights away from O after each update. When
Remark 3.5 (Approsimation of Brenier Potential). As pointed in Benamou et al. (2014), the Brenier potential can be non smooth and not strictly conver. Thus, JKO-ICNN can be understood as effectively optimizing not on the full space of conver functions, but rother on smooth subset of it (if we use smooth activation). As a consequence, JKO-ICNN does not seek tthe' Brenier potential, but instead a famely of smooth Brenier potentials, which may be distinct from the former. Note that this argument is similar the one used in (Paty et al., 2020).
4 Related Work
Optimization over Probability Measures and Connections to Convex Optimization The general problem of optimization over the space of measures has been studied extensively, particularly the functional analysis and applied mathematics literature (Jordan et al., 1998; Carlen & Gangbo, 2003; Léonard, 2008). Despite its unique challenges, optimization over measures exhibits various similarities to classic (finite- dimensional) optimization. For example, Chizat (2021) recently studied the convergence of Bregman proximal gradient methods for convex objectives over measures - extensions of their classical counterparts (Nesterov, 2004). Similarly, Wibisono (2018) studies sampling as an optimization over measures, revealing deep connections between these two problems.
Computational gradient flows common general approach for solving optimization problems involving distributions relies on gradient flows, which have been implemented through various computational methods, Benamou et al. (2016) propose an augmented Lagrangian approach for convex functionals based on the dynamical optimal transport implementation of Benamou & Brenier (2000). Another approach relying on the dynamic formulation of JKO and an Eulerian discretization of measures (i.e. via histograms) is the recent primal dual algorithm of Carrillo et al. (2021). Closer to our work is the formulation of Benamou

---

Published in Transactions on Machine Learning Research (07/2022)
et al. (2014) that casts the problem as an optimization over convex functions. This work relies on a Lagrangian discretization of measures, via cloud points, and on a representation of convex functions and their corresponding subgradients via their evaluation at these points. This method does not scale well in high dimensions since it computes Laguerre cells in order to find the subgradients. A different approach by Peyré (2015) defines entropic gradient flows using Eulerian discretization of measures and Sinkhorn- algorithms that leverage an entropic regularization of the Wasserstein distance. Caluya & Halder (2019) propose a method that combines entropic regularization and particle methods to solve gradient flows for the Fokker-Planck equation. Frogner & Poggio (2020) propose kernel approximations to compute gradient flows. On a different line of work, Salim et al. (2020) propose a forward- baakwwdd discretization scheme for Wasserstein gradient flows, and show that it enjoys convergences guarantees akin to those of proximal gradient hoth in usual Euclidean spaces. Finally, blob methods have been considered in Craig & Bertozzi (2014) and Carrillo et al. (2019) for the aggregation and diffusion equations. Blob methods regularize velocity fields with mollifiers (convolution with a kernel) and unlike many other alternative methods allow for the approximation of internal energies.
Forward and particle descent methods In addition, particle descent methods were explored for defining gradient flows for various geometries, e-g- for Sliced Wasserstein distance (Bonet et al., 2021; Liutkus et al., 2019), MMD (Mroueh et al., 2019; Mroueh & Rigotti, 2020; Mroueh & Nguyen, 2021; Arbel et al., 2019), Stein discrepancy (Liu & Wang, 2016; Liu, 2017; Duncan et al., 2019; Korba et al., 2021), and Kullback-Leibler divergence (Glaser et. al., 2021; Feng et al., 2021; di Langosco et al., 2021). Wang et al. (2021) proposes a projected Wasserstein gradient descent that exploits the manifold assumption.
ICNN, optimal transport, and generative modeling ICNN architectures were originally proposed by Amos et al. (2017) to allow for efficient inference in settings like structured prodiction, data imputation, and reinforcement learning. Since their introduction, they have been exploited in various other settings that require parametrizing convex functions, including optimal transport. For example, Makkuva et al. (2020) propose using them to learn an explicit optimal transport map between distributions, which under suitable assumptions, can be shown to be the gradient of a convex function (Brenier, 1991). The ICNN parametrization has been also exploited in order to learn continuous Wasserstein barycenters by Korotin et al. (2021). Using this same characterization, Huang et al. (2021) recently proposed to use ICNNs to parametrize flow-based invertible probabilistic models, an approach they call conver polential flows. These are instances of normalizing flows (not to be confused with gradient flows), and are useful for learning generative models when samples from the target (i.c., optimal) distributions are available and the goal is to learn a parametric generative model. Our use of ICNNs differs from these prior works and other approaches to generative modeling in two important ways. First, unlike those, our approach can be used even if the target distribution is available only up to a constant, or even if it cannot be sampled from and is only implicitly characterized as the minimizer of an optimization problem over distributions. Indoed, any functional defined over distributions
can be used as a plug-in objective in our framework, making it far more general than traditional generative modeling. Additionally, compared to other works that use ICNNs for optimal transport, we leverage these networks not only for solving a single optimal transport problem, but rather for a sequence of JKO step optimization problems that involve various terms in addition to the Wasserstein distance.
Neural networks methods for JKO In concurrent work, Mokrov et al. (2021) and Bunne et al. (2021) adopt a similar approach for approximating JKO with ICNNs. While the former is concerned exclusively with the Fokker-Planck equation, here we consider additional classes of PDEs. Bunne et al. (2021) tackle a different problem: learning dynamics with JKO, i.e., the inverse problem of learning the functional whose JKO flow follows empirical observations, which they also approach in its stochastic version using SDE approach (Bunne et al., 2022). Other subsequent works have built upon these initial set of parallel works. For example, Hwang et al. (2021) provide non-asymptotic approximation guarantecs and an empirical validation for a method similar to the one proposed here Nodozi & Halder (2022), on the other hand, introduce a variant of the JKO-ICNN approach that relies on the alternating direction method of multipliers (ADMM). On a different line of work, a min-max formulation of the JKO scheme was explored by Fan et al. (2021) and link between proximal training methods of GANs with Wasserstein gradient flows was proposed by Lin et al. (2021).

---

Published in Transactions on Machine Learning Research (07/2022)
Density
Density Estimated
Osjective flox-Evv Density
Conivergence of via Wasserstein
Fov
Pddd
twb
081
4-4 (0o-3
(0e-0
Pall
07-
6 (0x-0
2:52-22
C=33 10c 03
000-42
1.202-222
R.dte-31
0.5
Exam
2-
0.3
0.2
0.1
0.0
-1.00
-5.75
0.50
-0.25
0.00
6.25
0.50
0.75
1.00
-1.0
-275
-0.00
-025
0.00
025
0.00
0.55
2.00
0.0
3.0
Time
10-2
Tire
Le-1
(a) Porous medium equation with pure diff'ssio $5.1.
Density Estimated Flow
Dejective value Elso-Svalied Denshty
Density from Floeed Particles
Objective Value Flow-Evolved Density
2.00
1.4
1.75
5.0e00-:
1.4
0.5
1.50
(=1:te-0
te 1-500-5 23
5.00
1:25-
Fo (it-1)
FANTT
1.0
0.3
Py ae
3
P. (itasy
towad)
10-
0.-
0.3
0.4
0.0
0.5
1.0
1.5
2.0
3
0
10
Thmm
(b) Pokker-Planck with nonlinear diffusion $5.2.
(c) Aggregation equation $5.3.
Figure 1: Flows on PDEs with known solution. We use KDE on the flowed particles for the density plots, and the iterated push-forward density method (Corollary 3.3) to evaluate F in (a) and (b).
5 PDEs with known solutions
We first evaluate our method on gradient flows whose corresponding PDEs have known solutions. We focus on three examples from Carrillo et al. (2021) that combine the three types of functionals introduced in Section 3: porous medium, non-linear Fokker-Planck, and aggregation equations. We tackle 1D) versions of these equations in this section, and consider a higher-dimensional where we show that the JKO-ICNN dynamics closely follow Langevin dynamics in recovering the Wasserstein gradient flow corresponding of this PDE (Fig. 3). Throughout this section, we use T == 10-3 in the notation of Algorithm 1, and use the ICNN weights from step f to warm-start those of step t +.1.
(IR2") Fokker-Planck equation in Appendix D,
5.1 Porous medium equation
The porous medium equation is a classic non-linear diffusion PDE. We consider diffusion-only system: uip Ap",, m > 1, corresponding to a gradient flow of the internal energy F(p) - p("((r) dr, which we implement using our JKO-ICNN with objective (12). A known family of exact solutions of this PDE is given by the Barenblatt-Patee profiles (Zel'dovich & Kompaneetz, 1950; Barenblatt., 1952; Pattle, 1959):
p(x, (tt,t)==((( (Ckk]ll( k||x=||2 -
x E t 0,
where C>0 0 is a constant and a = d/(d(m 1) + 2), 3 = a/d, and k = a(m - 12/(m
This exact solution provides a trajectory of densities to compare our JKO-ICNN approach against. Specifically, starting from particles sampled from p(z; 0), we can compare the trajectory pe(x) estimated with our method to the exact density t(a,t). Although this system has stcady-state solution, its asymptotic behavior can be expressed analylically loo. For the case d= 1, m = 2, C (3/16)2/3, Figure 1a shows that our method reproduces the dynamics of the exact solution (here the flow density is estimated from particles via KDE and aggregated over 10 repetitions with random initialization) and that the objective value F(p) has the correct asymptotic behavior.
5.2 Nonlinear Fokker-Planck equation
Next, we consider a Fokker-Planck equation with a non-linear diffusion term as before:
app = V (pVV)
V V.Rd + R. m 1.
(13)
8

---

Published in Transactions on Machine Learning Research (07/2022)
This PDE corresponds to a gradient flow of the objective F(p) = m- p dx + (x)))())) For some V's its solutions approach a unique steady state (Carrillo & Toscani, 2000):
Poo (x) = (C m VV)(()))
(14)
where the constant C depends on the initial mass of the data. For d=1, m= and V(z)=z², we solve this PDE using JKO-ICNN with Objectives (10) and (12), using initial data drawn from a Normal distribution with parameters (H, o²) = 0.2). Unlike the previous example, in this case we do not have a full solution p(a, t) to compare against, but we can instead evaluate convergence of the flow Pro (2). Figure 1b shows that the density (e(x) derived from the JKO-ICNN flow converges to the steady state Peo (2), and so does the value of the objective, i.e., F(A) FP
5.3 Aggregation equation
Next, we consider an aggregation equation: 3,p V (pVW p), W. Rd - R. which corresponds to a gradient flow an interaction functional W(p) = 2 Jf W( x²) dp(=)) We consider the same setting as Carrillo et al. (2021): d=1, Pa (N 1), and the kernel W(x) = log((x]), which enforces repulsion at short length scales and attraction at longer scales. This choice of W has the advantage of yielding a unique steady-state equilibrium (Carrillo et al., 2012), given by (xx == 1222 x²)+ Our JKO-ICNN encodes W using Objective (10). As in the previous section, we investigate the convergence of this flow to this steady state distribution. Figure Ic shows that this case too we observe convergence of densities (x(x) + (r) and objective values F(A.) F(xx).
6 Molecular discovery
To demonstrate the flexibility and efficacy of our approach, we apply it in an important high dimensional setting: controlled generation molccular discovery. In our experiments, the goal is to increase the drug- likeness of given distribution of molecules while staying close to the original distribution, an important task in drug discovery and drug re-purposing. Formally, given an initial distributions of molecules po, convex potential energy function V that models the property of interest, and divergence D, we solve:
min F(p) = V2))) + DzD(p. PEP(X)
(15)
We use our JKO-ICNN scheme to optimize this functional on the space of probability measures, given an initial distribution Po(x) = 인 Ei 6n, (x) where is a molecular embedding.
In what follows, we show how we model each component of this functional via: (i) training a molecular embedding using a Variational Auto-encoder (VAE), (ii) training a surrogate potential to model drug- likeness, (iii) using automatic differentiation via the divergence D.
Choice of divergence term D For the divergence D in Equation (15), we use the 2-sasersn distance with entropic regularization (Cuturi, 2013).
JKO-ICNN for the functional we optimize in (15) is given for t2 0 for p = Pu:
f /V(VV V(V==o(=))ddt + A + 2T /
TT E argmin A
[Vzue(z) - =x12
6: 6:
p[+ = (Vzuo,))
(16)
We use the Sinkhorn algorithm (Cuturi, 2013) to compute D and backpropagate through this objective as proposed by Genevay et al. (2018), using the geomloss toolbox for efficiency (Feydy et al., 2019) (sce Appendix E.2 and E.3 for a discussion of alternative forms of divergence and ablation studies).
Embedding of molecules using VAEs We start by training a VAE on string representation of molecules (Polykovskiy et al., 2020; Chenthamarakshan et al., 2020) known as SMILES (Weininger, 1988) to reconstruct these strings with a regularization term that ensures smoothness of the encoder's latent space (Kingma & Welling, 2013; Higgins et al, 2016). We train the VAE molecular dataset known as MOSES (Polykovskiy
9

---

Published in Transactions on Machine Learning Research (07/2022)
Table 2: Comparison between JKO-ICNN and the direct optimization bascline. For each sctup, we report validity, uniqueness, median QED for the final point cloud of embeddings, and Sinkhorn divergence between the initial and final point clouds (Final SD), Each measurement value cell contains mean values = one standard deviation for 5 repeated runs with different random initialization seeds, and Po corresponds to the initial point cloud.
A2
LR
Validity
Uniqueness
QED Median
Final SD
Po N/A
N/A
100.000
0.000
99.980
0.045
0.630 0.001
N/A
JKO-ICNN
le
le-4
93.940 + 0.336
100.000 + 0.000
0.750 + 0.001
0.620 + 0.010
Baseline SGD
5c-
43.440 + 1.092
100.000 + 0.000
0.772 + 0.004
9792.93 + 76.913
5e
49.440 + 1.128
100.000 + 0.000
0.768
0.006
8881.38 + 69.736
1e3
5e-1
87.240 + 0.777
100.000 + 0.000
0.767 0.002
2515.08
49.870
Baseline ADAM 1e-1
92.080 + 0.973
100.000 + 0.000
0.793 = 0.005
18.261 = 0.134
0
le
93.900 + 0.781
99.979 + 0.048
0.758 + 0.006
1.650 + 0.006
1e-1
91.200 + 0.539
99.978 + 0.049
0.792 = 0.005
17.170 = 0.097
le
le-1
99.980 + 0.045
99.980 + 0.045
0.630 + 0.001
0.077 + 0.003
le*
1e-1
99.900 = 0.122
99.980 + 0.045
0.630 = 0.001
0.240 + 0.019
et al., 2020), which is subset of the ZINC database (Sterling & Irwin, 2015), released under the MIT license. This dataset contains about 1.6M training and 176k test molecules (see Appendix F, for results of this experiment run different dataset, QM9 (Ramakrishnan et al., 2014; Ruddigkeit, et al., 2012)). Given a molecule, we embed it using the VAE encoder to represent it with vector E 12128
Training a convex surrogate for the desired property (high QED) The quantitative estimate of drug-likeness (QED) (Bickerton et al., 2012) can be computed with the RDKit library (Landrum, 2013a) but is not differentiable nor convex. Hence, we propose to learn convex surrogate using Residual TCNNs (Amos ct al., 2017; Huang et al., 2021). This ensures that it can be used as convex potential functional V, as described in Section 3. To do so, we process the MOSES dataset via RDKit and obtain labeled set with QED values. We set QED threshold of 0.85 and give lower value label for all QED values above that threshold and higher value label for all QED values below it so that minimizing the potential functional with this convex surrogate will lead to higher QED values. Given VAE embeddings of the molecules, we train a ICNN classifier on this dataset. See Appendix F.1 for experimental details.
Optimization with JKO With the molecule embeddings coming from the VAE serving as the point cloud to be transported and the potential functional defined by the convex QED classifier, we run Algorithm to move an initial point cloud of molecule embeddings Po with low drug-likeness (QED 0.7) to a region of the latent space that decodes to molecules with distribution of with higher drug-likeness. We the following hyperparameters for JKO-ICNN: N, number of original embeddings, was 1,000. The JKO rate J set to le-4 and the outer loop steps T was set to 100. For the inner loop, the number of iterations n., set to 500, and the inner loop learning rate was set to 1e-3, For JKO-ICNN, we used a fully-comnected ICNN with two hidden layers, each of dimension 100. Finally, we ran the JKO-ICNN flow without warm starts between steps, as we observed empirically that resulted in better performance. The full pipeline for this experiment setting displayed in Figure in Appendix F. All computations were done with CPU and V100 GPU.
Evaluation We set A and X2 10,000 (see Table 4 in Appendix E.3 for details on hyperparameter search). We start JKO with initial cloud point po of embeddings that have QED 0.7 randomly sampled from the MOSES test set. In the second row of Table 2, we that JKO-ICNN is able to optimize the
10

---

Published in Transactions on Machine Learning Research (07/2022)
functional objective and leads to molecules that satisfy low energy potential, i.e., improved drug-likenesse while staying close to the original point cloud, as measured by Sinkhorn divergence (Final SD).
Comparison with direct optimization We also compare the JKO-ICNN flow to a bascline approach that optimizes the same functional objective via direct gradient descent on the molecule embeddings. Formally, this baseline approach is described by the process: dX; = VV (X)))t X2 VD( + Ei= Po), 1.... N, which we discretize using either vanilla gradient descent ADAM updates. For the bascline, we run grid search over various hyperparameters and reproduce a selection of configurations in Table 2 where we report validity, uniqueness, median QED (which corresponds to the first term in Equation (15)) of the final point cloud and Sinkhorn divergence between the initial and final point clouds (Final SD, corresponding to the second term in (15)). For the full grid search, see Tables and 7 in Appendix E.6. We note that the only baseline configurations that are able to meaningfully increase median QED are those where X2 orders of magnitude smaller than in the JKO-ICNN flow. Direct optimization therefore cannot accomplish the joint goals of the objective function. This is consistent with findings by Bunne al. (2021) with regards to the superiority of JKO over forward methods (direct optimization). This is significant because in many setting it is often crucial to stay close to the original set, c.g., drug re-purposing.
Benefit of computational amortization We show that the maps calculated at each step of the JKO- ICNN flow can be re-used to transport a new set of embeddingg with similar gains in QED distribution, without having to retrain the flow (Table 8 in Appendix E.7). We perform this comparison for various sample sizes of initial point clouds and observe linear scaling of the speedup of using the JKO-ICNN map relative direct optimization. This is key advantage of JKO-ICNN relative to direct optimization, which needs to be re-optimized for each new set of embeddings.
7 Discussion
In this paper, proposed JKO-ICNN, scalable method for computing Wasserstein gradient flows. Key to our approach is the parametcrization of the space of convex functions with Input Convex Neural Networks. We showed that JKO-ICNN succeeds at optimizing functionals on the space of probability distributions in low-dimensional settings involving known PDES as well as in large-scale and high-dimensional experiments on molecular discovery via controlled generation. Studying the convergence of solutions of JKO-ICNN an interesting open question that we leave for future work. To mitigate potential risks in biochemical discoveries, generated molecules should be verified in the laboratory, in vilro and in vivo, before being deployed.
References
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savare. Gradient flows in metric spaces and in the Wasserstein space of probability measures. Lectures in Mathematics. ETH Zürich. Birkhäuser Basel, 2005. ISBN 9783764373092. doi: 10.1007/b137000
Brandon Amos, Tei Xu, and J Zico Kolter. Input convex neural networks. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 146-155. PMLR, 2017.
Michael Arbel, Anna Korba, Adil Salim, and Arthur Gretton. Maximum mean discrepancy gradient flow. Advances in Newral Information Processing Systems, 32, 2019.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasscrstein generative adversarial networks. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th Internalional Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 214-223. PMLR, 2017.
G I Barenblaut On some unsteady motions of a liquid and gas in a porous medium. Prikl. Mal. Mekh., 16: 67-78, 1952. ISSN 0032-8235.
Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the Monge- Kantorovih mass transfer problem. Numerische Mathematik, 2000.
Jean-David Benamou, Guillaume Carlier, Quentin Mérigot, and Edouard Oudet. Discretization of functionals involving the monge-ampère operator, 2014.
11

---

Published in Transactions on Machine Learning Research (07/2022)
Jean-David Benamou, Carlier Guillaume, and Maxime Laborde. An augmented lagrangian approach to wasserstein gradient flows and applications. ESAIM: ProcS, 54:1-17, 2016. doi: 10.1051 2r110///rre/21166/66 URL https://doi. org/10. 0tt6//proc22016540000
G Richard Bickerton, Gaia V Paolini, Jérémy Besnard, Sorel Muresan, and Andrew L Hopkins. Quantifying the chemical beauty of drugs. Nature chemistry, 4(2):90 98, 2012.
Clément Bonet, Nicolas Courty, François Septier, and Lucas Drumetz. Sliced-wasserstein gradient flows.
arXiv preprint (rXix:2110010977 2021.
Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on Pure and Applied Mathemalics, 44(4):375-417, 1991.
Charlotte Bunne, Laetitia Meng- Meng--apaxanntos, Andreas Krause, and Marco Cuturi. Jkonet: Proximal optimal transport modeling of population dynamics, 2021.
Charlotte Bunne, Ya-Ping Hsich, Marco Cuturi, and Andreas Krause. Recovering stochastic dynamics via gaussian schrödinger bridges, 2022. URL https://arxx
http://aaxwvv.oogaa//////ae/ orggabs/2222222 05722..
Kenneth F Caluya and Abhishek Halder. Proximal recursion for solving the fokker-planck equation. In 2019 American Control Conference (ACC), pp. 4098-4103. IEEE, 2019.
E A Carlen and W Gangbo. Constrained steepest descent in the wasssrsstiiin metric. Ann. Math., 157(3): 807 846, 2003. ISSN 0003-486X.
A Carrillo and G Toscani. Asymptotic -eecy of solutions of the porous medium equation to self-smmilarity Indiana Univ. Math. J., 49(1):1133 142, 2000. ISSN 0022-2518, 1943-5258.
José A Carrillo, Lucas C F Ferreira, and Juliana C Precioso. A mass-transportation approach to a one dimensional fluid mechanics model with nonlocal velocity. Adv. Math., 231(1):306 327, September 2012. ISSN 0001-8708.
doi: 10.1016/j.aim /0000//..i..0122.3..... 2012.03.036.
José A Carrillo, Katy Craig, Li Wang, and Chaozhen Wei. Primal dual methods for wasserstein gradient flows. Found. Comul. Moth., March 2021. ISSN 1615-3375, 1615-3383. doi: 10.0007/110225--21--0000-3
José Antonio Carrillo, Katy Craig, and Francesco S. Patacchini. A blob method for diffusion, 2019.
Ricky Q Chen, Jens Behrmann, David K Duvenaud, and Joern-Henrik Jacobsen. Residual flows for invertible generative modeling. In H Wallach, H Larochelle, A Beygelzimer, F d' Alcch Buc, E Fox, and R Garnett (eds.); Advances in Newral Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Vijil Chenthamarakshan, Payel Das, Samuel C Hoffman, Hendrik Strobelt, Inkit Padhi, Kar Wai Lim, Ben Hoover, Matteo Manica, Jannis Born, Teodoro Laino, al. Cogmol: Target-specific and selective drug design for covid-19 using deep generative models. arXiv preprint 0xX1:2204.01225,, 2020.
Lénaïc Chizat. Convergence rates gradient methods for convex optimization in the space of measures, 2021.
URL /tttps///ariiv... www.iii..orgabsss.sseeeesosso htgg/bbs/1000 220.00066... 08368.
Katy Craig and Andrea L. Bertozzi. A blob method for the aggregation equation, 2014.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Burges, L Bottou, M Welling, Z Ghahramani, and K Q Weinberger (ods.). Advances in Neural Information Processing Systems 26, pp. 2292 2300. Curran Associates, Inc., 2013.
Lauro Langosco di Langosco, Vincent Fortuin, and Heiko Strathmann. Neural variational gradient descent. arXiv preprint arXtv:2107,10771, 2021.
Andrew Duncan, Nikolas Nüsken, and Lukasz Szpruch. On the geometry of stein variational gradient descent. arXiv preprint ar:iv: 00894, 2019.
12

---

Published in Transactions on Machine Learning Research (07/2022)
William Falcon et al. Pytorch lightning. GitHub. Note: lightning, 3, 2019.
Jiaojiao Fan, Amirhossein Taghvaei, and Yongxin Chen. 1:Xiu2212..22222 2021.
Variational wasserstein gradient flow. arXiv preprint
Xingdong Feng, Yuan Gao, Jian Huang, Yuling Jiao, and Xu Liu. Relative entropy gradient sampler for unnormalized distributions. arXiv preprint 0rxi::1110.02P 02778777 2021.
Jean Feydy, Thibault Séjourné, Maangis--aavier Vialard, Shun-ichi Amari, Alain Trouve, and Gabriel Peyré. Interpolating between optimal transport and mund using sinkhorn divergences. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 2681-268800 26900, 2019.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalic T.H. Gayraud, Hicham Janati, Alain Rakotomamoniy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1 8, 2021. URL http: /mmrr pwtrg//aprr////// httpsr//v22/////// html.
Charlie Frogner and Tomaso Poggio. Approximate inference with wasserstein gradient flows. In Silvia Chiappa and Roberto Calandra (ods.), Proceedings of the Twenty Third International Conference Artificial Intelligence and Stalislics, volume 108 of Proceedings of Machine Learning Research, pp. 2581 2590. PMLR, 8 2020.
Aude Genevay, Gabriel Peyre, and Marco Cuturi. Learning Generative Models with Sinkhorn Divergences. In International Conference on Artificial Intelligence and Statistics, volume 84, pp. 1608 1617. PMLR, 2018. URL http: pttt:e//pocceeinggs.o mlr.
Pierre Glasor, Michael Arbel, and Arthur Gretton. Kale flow: A relaxed kl gradient flow for probabilities with disjoint support. Advances in Neural Information Processing Systems, 34:8018 8031, 2021.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Z Ghahramani, M Welling, C Cortes, N Lawrence, and K Q Weinberger (eds.), Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014.
Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. JMLR, 2012.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. International Conference on Learning Representations (ICLR), 2016.
Chin-Wei Huang, Ricky T Q Chen, Christos Tsirigotis, and Aaron Courville. Convex potential flows: Universal probability distributions with optimal transport and convex optimization. In International Conference on Learning Representations, 2021.
M F Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics Simulation and Computation, 18(3):1099-1076, January 1989. ISSN 0361-0918. doi: (0.1080/03610091800881228
Hyung Ju Hwang, Cheolhyoong Kim, Min Sue Park, and Hwijac Son. The deep minimizing movement scheme, 2021. URL https://arxlv.org/ orwp////21///9/999 14851.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the Fokker Planck equa- tion. SIAM J. Math. Anal., 29(1):1-17, Jamuary 1998. ISSN 0036-14141 doi: 10.1177/000614400663333
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.611 2013.
13

---

Published in Transactions on Machine Learning Research (07/2022)
Anna Korba, Pierre- Cyril Aubin- Buain-Brankowski, Szymon Majewski, and Pierre Ablin. Kernel stein discrepancy descent. In International Conference on Machine Learning, pp- 5719-5730. PMI.R, 2021.
Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev. Continuous wasserstein-2 barycenter estimation without minimax optimization. In International Conference on Learning Representations, 2021.
URL https://openreviss .wt. for
Greg Landrum. Rdkit: A software suite for cheminformatics, computational chemistry, and prodictive modeling, 2013a.
Greg Landrum. Rdkit: Open-source cheminformatics, 2013b. URL https::/www.wddttttt org.
Alex Tong Lin, Wuchen Li, Stanley Osher, and Guido Montúfar. Wasserstein proximal of gans. In International Conference on Geometric Science of Information, pp. 524 533. Springer, 2021.
Qiang Liu. Stein variational gradient descent as gradient flow. Advances in neural information processing systems, 30, 2017.
Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances in neural informalion processing systems, 29, 2016.
Antoine Liutkus, Umut Simsekli, Szymon Majewski, Alain Durmus, and Fabian-Robert Stöter. Sliced- wasserstein flows: Nonparametric generative modeling via optimal transport and diffusions. In International Conference on Machine Learning, pp. 4104 4113. PMLR, 2019.
Christian Léonard. Minimization of entropy functionals. J. Math. Anal. Appl, 346(1):183-22) 204, October 2008. ISSN 0022-247X. doi: 0..0016/.jmaa.200.0411. 20.......8...
Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In Hal Daumé li And Singh (ed.), Proceedings of the 37th International Conference on Machine Learning, volume 119, pp. 1672-6681. PMLR, 2020.
Robert J. McCann. A convexity principle for interacting gases. Advances in Mathematics, 128(1):153 179, 1997.
Raphael A Meyer, Cameron Musco, Christopher Musco, and David P WoodrufT. Hutch++: Optimal stochastic trace estimation. In Sympostum on Simplicity in Algorithms (SOSA), pp. 142- 155. SIAM, 2021.
Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin Solomon, and Evgeny Burnaev. Large- scale wasserstein gradient flows, 2021.
Youssef Mrouch and Truyen Nguyen. On the convergence of gradient descent. in gans: Mmd gan as a gradient. flow. In International Conference on Artificial Intelligence and Statistics, pp. 1720 1728. PMLR, 2021.
Youssef Mroueh and Mattia Rigotti. Unbalanced sobolev descent. Advances in Neural Information Processing Systems, 33:17034 17043, 2020.
Youssef Mrouch, Tom Sercu, and Anant Raj. Sobolev descent. In The 22nd International Conference Artificial Intelligence and Statistics, pp. 2976 2985. PMLR, 2019.
Kevin P Murphy. Machine Learning: A Probabilistic Perspective. Adaptive Computation and Machine Learning series. MIT Press, 2012. ISBN 9780262018029.
doi: 10.1007 1..007/Sppinneereeeeenneey beeee 35834.
Yurii Nesterov. Introductory Lectures on Conver Optimization: Basic Course. Springer, Boston, MA, 2004. ISBN 9781402075537, 9781441988539. doi: 10.1007/978-1-4419-888
Iman Nodozi and Abhishek Halder. A distributed algorithm for measure-valued optimization with additive objective, 2022. URL nttps:://rat org/abs//222. /ttps///2222.2..22
Felix Otto. The geometry of dissipative evolution equations: the porous medium equation. Comm. Partial Differential Equations, 26(1-2):101-17) 114, January 2001. ISSN 0360-5302, doi:
14

---

Published in Transactions on Machine Learning Research (07/2022)
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperalive style, high- performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché Buc, E. Fox, and R. Garnett (ods.), Advances in Neural Information Processing Systems 32, pp. 8024 8035. Curran Associates, Inc., 2019. URL http://papers neur pes.ipsscc.. cc/paper/ 9015-pytorch-an-imeer oow. setrraiii--ttlll-igg--eerrr
.ddff
R. Pattle. Diffusion from an instantaneous point source with a concentratioon-deendeeet coefficient. Quart. 10tmmamm12.4407. 12.44.077
J. Mech. Appl. Math., 12(4):407 409, January 1959. ISSN 0033-5614. doi: /a//9//jmmm/////
Frangois-Pierre Paty, Alexandre d' A'Apremont,, and Marco Cuturi. Regularity as regularization: Smooth and strongly convex brenier potentials in optimal transport. In Silvia Chiappa and Roberto Calandra (eds.), Proceedings of the Twenly Third International Conference on Artificial Intelligence and Statislics, volume 108 of Proceedings of Machine Learning Research, pp. 1222 1232. PMLR, 2020.
Gabriel Peyré. Entropic approximation of wasserstein gradient flows SIAM J. Imaging Sci., 8(4):2323 2351, January 2015. doi: 10.1137/16M10100ST.
Daniil Polykovskiy, Alexander Zhebrak, Benjamin Sanchez-Lengeling, Sergey Golovanov, Oktai Tatanov, Stanislav Belyaev, Rauf Kurbanov, Aleksey Arlamonov, Vladimir Aladinskiy, Mark Veselov, Artur Kadurin, Simon Johansson, Hongming Chen, Sergey Nikolenko, Alan Aspuru-Guzik, and Alex Zhavoronkov. Molecular Sets (MOSES): A Benchmarking Platform for Molecular Generation Models. Frontiers in Pharmacology, 2020.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 1530-1538, Lille, France, 2015. PMLR.
R. Tyrrell Rockafellar. Convez analysis. Princeton Mathematical Series. Princeton University Press, Princeton,
N. J. 1970.
Lars Ruddigkeit, Rund Van Deursen, Lorenz C Blum, and Jean-Louis Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17. Journal of chemical information and modeling, 52((1)) 2012.
Adil Salim, Anna Korba, and Giulia Luise. The wasserstein proximal gradient algorithm. In H Larochelle, M Ranzato, R Hadsell, M F Balcan, and H Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 12356 12366. Curran Associates, Inc., 2020.
Filippo Santambrogio. Optimal Transport for Applied Mathematicions: Calculus of Variations, PDEs, and Modeling. Birkhäuser, Cham, 2015. ISBN 9783319208275. doi: 10.1007/998-3331111 208-2022-22.....
Filippo Santambrogio. (Euclidean, metric, and Wasserstein] gradient flows: an overview. Bull. Moth. Sci., (1):87-154, April 2017. ISSN 1664-3607, 1664-3615. doi: 10.1007/413333---7-000001
Teague Sterling and John J Irwin. Zine 15-ligand discovery for everyone. Journal of chemical information and modeling, 55(11):2324 2337, 2015.
Shashanka Ubaru, Jie Chen, and Yousef Saad. Fast estimation of Str(f(a))8 via stochastic lanczos quadrature. SIAM Journal on Matria Analysis and Applications, 38(4):1075-1099, 2017.
Yifei Wang, Peng Chen, and Wuchen Li. Projected wasserstein gradient descent for high-dimensional bayesian inference. arXiv preprint XrXtx:2102.00:: 2021.
15

---

Published in Transactions on Machine Learning Research (07/2022)
David Weininger. Smiles, chemical language and information system. 1. introduction to methodology and encoding rules. Journal of chemical information and computer sciences, 28(1):31-33 -30, 1988.
Andre Wibisono. Sampling as optimization in the space of measures: The langevin dynamics as composite optimization problem. In Sébastien Bubeck, Vianney Perchet, and Philippe Rigollet (eds.), Proceedings of the 31st Conference On Learning Theory, volume 75 of Proceedings of Machine Learning Research, pp. 2093 3027. PMLR, 06 09 Jul 2018. URL https: /ttpss////ooeedii/o mlr. htt....ees/////wiiiiiisssss
8a. html.
Yákov B Zel' dovich and A S Kompaneelz. Towards theory of heat conduction with chermal conductivity depending the temperature. Collection of papers dedicated to 70th birthday of Academician AF Ioffe, Izd. Akad. Nauk SSSR, Moscow, pp. 61-71, 1950.
16

---

Published in Transactions on Machine Learning Research (07/202)
A Proofs
A.1 Proof of Lemma 3.1
Starting from the original form of the potential energy functional V in Equation (3), and using the expression (Vu)pPt we have:
v(p) - (VV)))((V = / V(z) [(Ve))p)) (v
(17)
On the other hand, for an interaction functional W we first note that it can be written as
1
/f W(2 2')
1 / (W ()))))(()))
W(p)
(18)
In addition, we will need the fact that
W(x Vu(v)) W * [(Vu)ar] = W(x / (VVV)((())
(19)
Hence, combining the two oquations above we have:
2 / (w * (Vab)))))))d((((w))c))
1
W((V))=pL)
1 / (/ (V(==-v()))) (/W(=-Vuo))a())) d[(Vu))p(==)
1 2 // /W(Vu(x)- vuly)) dpr(z),
as stated. L
A.2 Proof of Lemma 3.2
Following Santambrogio (2017), we note that whenever u is convex and v is absolutely continuous, then Tyv is absolutely continuous too, with a density given by
V
T-
(20)
Jrl
where Jr is the Jacobian matrix of T. In our case (Vu)aPt. so that
Pt ((Vu)- (v) H. ((Vu)-1(v)
p(v) = Pr H
(Vu)
(21)
where H the Hessian u. When strictly convex it is known that invertible and that (Vu) Vu", where u² the convex conjugate of u (see e.g. Rockafellar (1970)). O
A.3 Proof of Corollary 3.3
In this proof, we drop the index As before, we use the change of variables pt = (Vu+)pt-1- Thus, by induction,
Pt (Vu o Vu, ... Vuz)aPo
(22)
Let T = (Vu, o Vu.... Vuz), so that P: (Tia)aPo (Vu Tra-114PP- The Jacobian of this map is given by the chain rule as:
Jnn (z) = H (T) (x)
(23)
Hence by induction we have:
JTn (x) 1I H (T (x))
(24)
17

---

Published in Transactions on Machine Learning Research (07/2022)
On the other hand, as long as the inverses exist, we have:
T
(Vuz)-'o.... (Vui)-1= =Vu, e .... ou
Hence,
pc(y)
Po
Ti:t (y) =
PO
o TTi (3)
[T
IT: 6=1
Hu,, (T::0-1)
Po 111 (v)
11 - [H, (Zi:s-1 (1))
and finally taking the log we obtain:
log (p:)
log (Po) - -
log Hu, 1))) (T 7
(25)
8=1
A.4 Proof of Lemma 3.4 As before, let Pt+11 (Vuc+1) Thus, by induction,
Pt+1 (Vu4+1 Vu, ... Vui)aPo
(26)
Let Tia (Vue o Vug... Vui), so that Pc+1 (Tr++1)4Po (Vu++1 Trt)sPo- The Jacobian of this map is given by the chain rule as:
Jre (z) Hump (TT.(E))JTT..... (z)
(27)
On the other hand, using the fact that (whenever the inverses exist) (f o 1 f-1, in our case we have
Th4 Tise Vut+1
(28)
Using Equations (20) and (26), we can write the density of p++1 as
(J Pn
Pc+1(y)
Tist1
(v)
(29)
Pol (T (T 11441 (y))
Po Tial (urti())
(30)
JTan (1t41(1)) (1))
H H (Vu (y)) Tin (Ti.l V- (y))
Finally, using the change of variables y Vutt TTa()), Tuil Vuzti in the integral in the definition of F, we get
F(pet1)= = =f$ Hu (Txe(x)))
Po(x)
Hunt (T):( dr
po(x) Exaspo ( [mn (Tia(x)) (JT (
H H (Tis(2))) (x) Po(x)
(x)]
B Practical Considerations
B.1 Input Convex Neural Networks
Input Convex Neural Networks were introduced by Amos et al. (2017). A k-layer fully input convex neural network (FICNN) is one in which each layer has the form:
2++ (W( (( y wibb) i == O,..,,
(31)
18

---

Published in Transactions on Machine Learning Research (07/202)
where 9: are activation functions. Amos et al. (2017) showed that the function f + is convex with respoct x all the Wi-1 are non-negative, and all the activation functions 9i are convex and non-decreasing. Residual skip connections from the input with linear weights are also allowed and preserve convexity (Amos et al., 2017; Huang al., 2021).
In experiments, parametrize the Brenier potential ug FICNN with two hidden layers, with (100, 20) hidden units for the simple PDE experiments in Section 5 and (100, 100) for the molecule generation experiments Section In order to preserve the convexity of the network, we clip the weights of W after every gradient update using Wiy - max Way, 10 10-*). Alternatively, one can add small term +X||a enforce strong convexity. In all our simple PDE experiments in Section 5, 10-3 initial learning rate, and a JKO step-size 10-3. Optimization details for the molecular experiments are provided in that section (Section 6).
the ADAM optimizer with
B.2 Surrogate loss for entropy
For the choice f(t) = t logt in the internal energy functional F, we not use Lemma 3.4 but rather derive the expression from first principles:
/ Pe(a) log
(c(x)
F((Vau))=n)
Hus (x) da
Hug (x)
H ug (x)
/ log
pe(
Hu,, (x:) Pr
dr
/ Pe log (x() dx - dz-
log (z)) = F(p.) - E
[log H.g (x)]]
arups
As mentioned carlier, this expression has an intercsting interpretation reducing negative entropy (incrcasing entropy) of Pt by amount given by a log-determinant barrier term w's Hessian. We see that the only term depending on 0 is - Ez=-p [log [Hu, (x)]]. We discuss how to estimate this quantity and backpropagate through Hu, in Appendix B.4.
B.3
Surrogate losses for internal energies
Let
(uo)) = log(E(x)) = log(po(z)) - log Hno (Ti+())) - log((/i<zzz))
From Lemma 3.4, our point-wise loss is
L(wo) f o exp(r=(us))
exp(r=(uo))
Computing gradient w.r.t 0, parameters of us:
0 L(u) ((exp(rr(uu))))exx(rrccuu 20, (u((u) exp(r=(ue)) 20, rx (uo)f exp(r=(u6))
00;
(xx)(==(uo)]
[((exp(ra(uo)) (xp(r ( (exp((x exp(r=(uo))
(uo)
00,
Also,
à log Ha, (Tia()))
00; Ta (v)
00;
Hence the Surrogate loss that has same gradient can be evaluated as follows:
L(uo) = (((exp(ra(wo)) exppprrr (mp(r((u)))
- (ooomi)))(((xxxe (exp(ra(uu))) log Hu,, (Tie()))
grad
19

---

Published in Transactions on Machine Learning Research (07/2022
For the particular case of porous medium internal energy, let a = exp(r). For f(a) m- a" exp(m log(a) m-T exp(mr)), we have (a( m mm m- exp((m 1)log(a)) (a m
exp((m 1)r).
(a) f(a) a m exp((m - 1)r) m
1
cxp((m 1)r) exp((m- 1)r)
Hence we have finally the surrogate loss:
((us) = ((u))==<>p(mm- 1)rz (uo)) log Hap (Tia()))),
no grad
for which discuss estimation in Appendix B.4. Table summarizes surrogates losses for common energies in gradient flows.
Table 3: Surrogate optimization objectives used for computation. Here È denotes an empirical expectation, E(z) is defined in Lemma 3.4, and SG denotes the STOPGRAD operator.
Functional Type
Exact Form F(P)
Surrogate Objective F(uo)
Potential energy
J ((x) dp(z)
Bzwpi (oV
Interaction energy
JJ ((x - z²) dp(x)) dp(x')
E
Eyypp W(V, 2 Ug (x) (o))(y))
Meg-Entropy
p(x)) log p(x) da
Rapp. log H (z)
Nonlinear diffusion
S ((="" dx
e e ( E(m)) log H (T))
B.4 Stochastic log determinant estimators
For numerical reasons, we use different methods to evaluate and compute gradients of Hessian log- determinants.
Evaluating log-determinants Following Huang al. (2021), we use Stochastic Lanczos Quadrature (SLQ) (Ubaru et al., 2017) to estimate the log determinants. We refer to this step as LOGDETESTIMATOR.
Estimating gradients of log-determinants The SLQ procedure involves an eigendecomposition, which is unstable to back-propagate through. Thus, to compute gradients, Huang et al. (2021), inspired by Chen et al. (2019), instead use the following expression of the Hessian log- log-determinant:
log H
H | (H)
tr (H- 0H E
9H 80
(32)
Here e a random Rademacher vector, This last step is the Hutchinson trace estimator (Hutchinson, 1989).
Algorithm Density estimation for JKO-ICNN
1: Input:
Query point x, sequence of Brenier potentials (u,) obtained with JKO-ICNN, initial density evaluation function po(-).
As in Huang et al. (2021), we avoid constructing and inverting the Hessian in this expression by instead solving a problem that requires computing
2: Initialize y: +
3: for t T to do
only Hessian-vector products:
4: 11-1 argmax, (Ve, ue (v)
argmin 1/2 H: 1---7- T
(33)
5: {u--1 satisfies (Vu))(n y/
6: end for
Since H is symmetric positive definite, this strictly convex problem has unique minimizer, 2", that satisfies =* = H- This problem can be solved using the conjugate gradient method with fixed number of iterations or a error stopping condi- tion. Thus, computing the last expression in Equa- tion (32) can be done with automatic differentia-
7: Do + yo
8: for l = to - d
9: I++1 + (2,
end for
11: { [Compute 5 & log JJn (xo)1}
12: 8 -LOGGDETTSSTIAAATRRR(((
13: p t log (o -
tion by: (i) sampling Rademacher vector v, (ii)
14: Output: p satisfying p = log (x)
running conjugate gradient for m iterations
Problem (33) to obtain z (iii) computing & (>>*) Huj with automatic differentiation.
20

---

Published in Transactions on Machine Learning Research (07/2022)
Complexity of the Hutchinson method The Hutchinson estimator needs O(1/c²) random projections to give an 1 + guarantee for trace approximation. newer version, Hutch+++ proposed by Meyer et al. (2021), reduces the number of random projections needed to (11cc).
B.5 Implementation Details
Apart from the stochastic log-determinant estimation (Appendix B.4) needed for computing internal energy functionals, the other main procedure that requires discussion is the density estimation. This is needed, for example, to obtain exact evaluation of the internal energy functionals and requires having access to the exact density (or an estimate thereof, e.g., via KDE) from which the initial set of particles were sampled. For this, we rely on Lemma 3.2 and Corollary 3.3, which combined provide way estimate the density (r()) using pa(a;, the sequence of Brenier potentials {ui)== and their combined Hessian log-determinant. This procedure is summarized in Algorithm 2.
21

---

Published in Transactions on Machine Learning Research (07 (07/2022)
C
Additional qualitative results on 2D datasets
Tietel
Time o 0,800 2 68
Timet 2200. Pote-1.3
Timets to 2.0005, Plole
(a) Heat equation corresponding to internal energy functional F(p) p(z)) log p(t) da (see Table 1). Tmmt= 21001
ToateD
Tmma 92000 Reletis
Timet= R.000, 0Z
(b) Potential-only flow using functional V(p) = f || - 2013 dp(x).
Tmete0
Tmet= 0,820 37
Timt 3el=0
Time 0.000 25
(c) Fokker-Planck with linear diffusion: same potential as 2b, plus F(p) p(x) log p(z) dz
D
Ime 0,000 Alain 0.32.
Timete 02000 0.3
Time 04000 Apgle 022.
(d) Fokker-Planck with noulinear diffusion: same potential as 2b, plus F(p) = f(p)=f{(())d (x)
Lmm 00r0 rinlenas
Lime te uobon Holesa
limgte 21000 Aielmna
(e) Aggregation equation (same functional as $5.3).
Figure JKO-ICNN flows of a 2D point cloud with density estimated via KDF.
22

---

Published in Transactions on Machine Learning Research /0022)
C.1 Experimental details: PDEs with known solutions
In all the experiments in Section 6, we use the same JKO step-size 10-3 for the outer loop and ADAM optimizer with initial learning rate 77 10-3 for the inner loop. We run the inner optimization loop for 400 iterations. In the plots in Figure 1, we show snap-shots at different intervals to facilitate visualization. For these experiments, we parametrize Wg as an 2-hidden-layer ICNN with layer width: (100, 20). For the non-linear diffusion term needed for the PDEs in Sections 5.1 and 5.2, we use the surrogate Objective (12). all cases, we impose strict positivity on the weight matrices W(=) (Equation (31)) with minimum value 8 10 10-18
to enforce strong convexity.
D Comparing high dimensional PDEs to Wasserstein gradient flow
To confirm that JKO-ICNN recovers the Wasserstein gradient flow in high dimension at all time steps, we considered the Fokker-Planck equation and ran the following experiment on the Langevin dynamics high dimension considering a convex potential: V(z) = (x - n) A(x - μ)), where æ, μ E IR is positive-defiiite matrix, F()) = (<)p(()) + H(), and H(p) = f log(p)dz is the negative entropy.
We chose this functional since Langevin dynamics can be implemented using particles thanks to the Unadjusted Langevin Algorithm (ULA). ULA with learning rate is given by:
Xk-1 Xx VX v 2m18-1an
where & ~ (O, Ia), 3 is temperature term, and Xo is sampled form (OO,, This initial distribution is used for ULA and JKO-ICNN.
Flow Comparison: Distance True (ULA) Wass. GF
0.014
Flow Pair
0.012
JKO-ICNN ULA
GD ULA
0.010
0.008
MMD(X
0.006
0.004
0.002
0.000
3
Time (adj. step size # steps)
le-4
Figure 3: Wasserstein Gradient Flow Recovery in high dimension: comparison between the JKO-ICNN flow and direct gradient descent the objective (GD), evaluated against the 'truc' flow obtained by following dynamics using the Unadjusted Langevin Algorithm (ULA) to simulate the Pokker-Planck equation in 20D.
It known that ULA implements the minimization of F(p) using particles, and at the limit of an infinite number particles and as 17 + 0, the intermediate distribution of Xx ULA corresponds to the Wasserstein Gradient Flow (WGF) dynamics Pt. Hence, we compare the distance between the JKO-ICNN intermediate cloud point ULA's 's all times, showing that the JKO-ICNN recovers the Wasserstein gradient flows all steps (i.e., the MMD of JKO-ICNN's intermediate clouds to ULA's remains small at all time steps, see Figure 3, for d = 20).
23

---

Published in Transactions on Machine Learning Research (07/2022)
In order to put these results in context, we also provide the MMD distance of intermediate clouds of ULA versus direct optimization (gradient descent, GD). We see that JKO-ICNN faithfully tracks ULA, relative to GD.
When trying to go beyond d = 20, we run into issues with ULA, which is known to become unstable in high dimensions. Although this can be addressed via anncaling rates or temperatures, this would make it deviate from the true WGF, defeating the purpose the comparison, and out of the scope for this work.
E Experimental details: molecular discovery with JKO-ICNN (MOSES dataset)
In what follows, we present the experimental details of the experiments in Section 6, on the MOSES dataset. The MOSES dataset (Polykovskiy et al., 2020) is is subset of the ZINC database (Sterling & Irwin, 2015). MOSES dataset is available for download at https:/ www.th/.c.o/oooleeaaaaiieeo /moses,, released under the MIT license.
All Molecular discovery JKO-ICNN experiments run in a compute environment with CPU and 1 V100 GPU submitted as resource-rest jobs to cluster. This applies to both convex QED surrogate classifier training and evaluation runs and to the JKO-ICNN flows for each configuration of hyperparameters and random seed initialization. The full pipeline for this experiment is detailed in Figure
Molecule SMILES Strings
VAE decoder
Molecule embeddings
JKO-ICNN
(Convex) QED classifier
D (Distribution distance)
VAE encoder
Nc1-cc2-c c-c02
QED annotation
Figure 4: JKO-ICNN for molecular discovery. We apply our method to generate molecules whose Brug-likeness and closeness to Po are maximized in latent VAE space.
E.1 Convex QED surrogate classifier
We first describe the hyperparameters and results of the convex surrogate that was trained predict high (> 0.85) and low (< 0.85) QED values from molecule embeddings coming from a pre-trained VAE. Molecules with high QED were given lower value labels compared to the low QED molecules so that when this model would be used as potential, minimizing this functional would lead to higher QED values. For this convex surrogate, we trained a Residual ICNN (Amos et al., 2017: Huang et al., 2021) model with four hidden layers, each with dimension 128, which was the dimensionality of the input molocule embeddings as well. We trainod the model with binary cross-ent ropy loss. To maintain convexily in the potential functional however, we used the last layer before the sigmoid activation for V. The model was trained with an initial learning rate of 0.01, batch sizes of 1,024, ADAM optimizer, and learning rate scheduler that decreased learning rate on validation set loss plateau. The model was trained for 100 epochs, and the weights from the final epoch were used to initialize the convex surrogate in the potential functional. For this epoch, the model achieved 85% accuracy on the test set. In Figure 5, we display the test set confusion matrix for this final epoch.
E.2 Automatic differentiation via D
For the divergence D in Equation (15), we use either the 2- Wassersttnn distance with entropic regularization (Cuturi, 2013) or the Maximum Mean Discrepancy (Gretton et al., 2012) with a Gaussian kernel (MMD). When is the entropx-regullarized Wasserstein distance, we use the Sinkhorn algorithm (Cuturi, 2013) compute (henceforth denoted Sinkhorn). We backpropagate through this objective as proposed by Genevay et al. (2018), using the geomloss toolbox for efficiency (Feydy et al., 2019). We also geomloss for evaluation and backpropagation when D is chosen to be MMD.
24

---

Published in Transactions on Machine Learning Research (07 (07/2022)
0.8
0
0.8
0.2
0.7
Ground Truth
0.6
0.5
-0.4
0.11
0.89
0.3
0.2
Predicted
Figure 5: Confusion matrix for the convex surrogate trained to predict QED labels from molecule embeddings for MOSES.
E.3 D ablation study and X hyperparameter search
We fixed Àι 1 for all experiments and used either Sinkhorn MMD for D. The weight X2 on D set to cither 1,000 or 10,000. We start JKO with an initial cloud point of embedding that have QED < 0.7 randomly sampled from the MOSES test set. In the Table 4, we report several measurements. First is validity, which the proportion of the decoded embeddings that have valid SMILES strings according to RDKit. Of the valid strings, we calculate the percent that are unique. Finally, we use RDKit to get the QED annotation of the decoded embeddings and report median values for the point cloud. We re-ran experiments five times with different random seed initializations and report means and standard deviations. In the first of Table 4, report the initial values for the point cloud at time t 0. In the last four rows of the table, report the measurements for different hyperparameter configurations al the end of the JKO scheme for t = 7 100. The results are stable across different runs as can be seen from the reported standard dcviations from the different random initializations. We notice that Sinkhorn divergence with X2 =10,000 prevents mode collapse and preserves uniqueness of the transported embedding via JKO-ICNN. While MMD yields higher drug-likeness, it leads to a deterioration between the transformed point cloud and the original one, which preserves better uniqueness than MMD, which merely matches mean embeddings of the distribution.. The JKO-ICNN experiment reported in Table 2 in Section 6 and in Table 8 uses Sinkhorn as the divergence term.
uniqueness. Using Sinkhorn allows for a matching
E.4 T hyperparameter search
Our search of optimal F was done under the following setup: learning rate for training the ICNN was fixed at = 0.001. Åg set to 10,000, the JKO outer loop was run for 100 steps, and inner loop optimization was set 500 steps. The results are presented Table 5. We notice that unlike direct optimization (sce Tables and 7), JKO-ICNN is robust across learning rates
E.5 JKO-ICNN QED histograms
In Figure 6, we present several time steps of the histograms of the QED values for the decoded SMILES strings corresponding to the point clouds for the experiment that used Sinkhorn as the distribution distance with weight 10,000.
E.6 Direct optimization baseline hyperparameter grid search
The full grid search over hyperparameters for the direct optimization baseline discussed in Section is available in Tables and 7. In order to ensure 'fair shot' competing with our JKO-ICNN approach, we performed a grid search over the following hyperparameters:
25

---

Published in Transactions on Machine Learning Research (07/2022
Table 4: Molccular discovery with JKO-ICNN experiment results Measures of validity, uniqueness, and median, average, and standard deviation QED are reported in each row for the corresponding point cloud of embeddings. For each point cloud, we decode the embeddings to get SMILES strings and use RDKit to determine whether the corresponding string is valid and get the associated QED value. In the first row, we display the point cloud at time step zero of JKO-ICNN. In the subsequent rows, we display the values for each measurement at the final time step T of JKO-ICNN for different hyperparameter configurations of distribution distance D (either Sinkhorn or MMD) and weight on this distance X2 (either 1,000 or 10,000). Each mcasurement value cell contains mean values + one standard deviation for five repeated runs of the experiment with different random initialization seeds. We find that the setup that uses Sinkhorn with weight, 10,000 yields the best results in terms of moving the point cloud towards regions with higher QED without sacrificing validity and uniqueness of the decoded SMILES strings. This table contains results for the MOSES dataset.
Mrasure
Ag
Validity
Uniqueness
QED Median
QEI Avg-
QED Std.
Po
N/A
N/A
100,000 + 0.000
99.980 + 0.045
0.630 + 0.001
0.621 0.000
0.063 +: 0.002
P
Sinkhorn
1,000
92.460
2.096
69.919 + 4.906
0.746 + 0.016
0.735
0.009
0.110 + 0,003
Pr
Sinkhorn
10,000
93.020
1.001
99.245
1 0.439
0.769
1 0.002
0.754
0.003
0.112 1: 0.002
PT
MMD
1,000
94.560 + 1.372
51.668 + 2.205
0.780
: 0.009
0.767 +: 0.013
0.107 + 0.012
PT
MMD
10,000
92.020
3.535
53.774 1 3.013
0.776 1 0.014
0.767
0.009
0.102 1
Table 5: Molecular discovery with JKO-ICNN experiment hyperparameter search. Measures of validity, uniqueness, and median QED of final point cloud of embeddings are presented for each т. Fach measurement value cell contains mean values + one standard deviation for five repeated runs of the experiment with different random initialization sceds. For this search, we fix the other hyperparameters: 17 = 0.001, λ2 = 10,000, the JKO outer loop was run for 100 steps, and inner loop optimization was set to 500 steps. This table contains results for the MOSES datasct.
T
Validity
Uniqueness
QED Median
0.01
94.600 +'2 0.620
99.979 + 0.047
0.708 + 0.007
0.001
94.620 + 0.907
99.979 + 0.047
0.716 + 0.005
0.0001
93.320 + 0.687
99.957 + 0.059
0.751 + 0.007
Optimizer: (SGD, ADAM]
Learning rate (LR): {5e-1 1e-1, 1e-2, le 1e-3, 1e-4}
A2: {0, 1, 10, 100, 1, 1,000, 10,000}
and report validity, uniqueness, median QED of the final point cloud and Sinkhorn divergence between the initial and final point clouds (Final SD).
E.7 Computation amortization results
In Table 8, we compare the results and computation time for the the best direct optimization hyperparameter configuration found during the grid search (Optimizer: ADAM, LR: 0.1, Ag: 1) vs. re-using the maps found during the JKO-ICNN flow applied to new set of embeddings. We observe linear scaling speed-up when reusing the JKO-ICNN maps vs. re-optimizing the baseline on a new set of points.
F Molecule JKO-ICNN experimental setup and results for QM9 dataset
In this section, we repeat the results and analysis presentation of Appendix E but for the experiments that used the QM9 dataset (Ramakrishnan et al., 2014; Ruddigkeit et al., 2012). The QM9 dataset contains on average smaller molecules than MOSES. The MOSES dataset is larger in size than QM9 (Training: 1.6M molecules in MOSES vs. 121k in QM9; Test: 176k in MOSES molecules vs. 13k in QM9). For these experiments separate VAE model and convex surrogate classifier were trained using the QM9 dataset.
26

---

Published in Transactions on Machine Learning Research
QED Oitribution Step #g
QED Distrbbtion Stop #10
QED Distibution Stop #20
Distibutiin delance Sintherm, Weight 10,000 Mediam 82%, buy 821;
Dstribbtion distance Sinchom Weight 10,000
Dstriibttin dislance: Siithmm Weigh: 10,000
Mediam 6.67, Auy 0 714: Std 0.136
Median Aug B 763, 117
8.50
0.55
6.60
0.55
0.70
0.5
06 Q2D
0.
0.4
07
00
OC
QTD
GED Distribution Stop #33
QED Distribution Stop #440
OED Distribution Stop #50
Distribution detance: Sinthorn, Weight 10,000
Distribution distanca:: Senchom Weight: 10 000
Ditttibutinn distance: Sinthom, Weight: 10.000
Mediam 780, Awg FR3,
Meciam: 755, Jogg Sidi 0.134
Mediam: Aay: 788, 116
0
0.4
0.5
0.7
0.
0.7
08.
120
060
GED Distribution Step #50
QED Distribution Step TO
QED Distribution Step #BD
Distribution detance. Sinkhorn, Weight. 10,000
Distribution distance: Sincham; Weight 10.000
Distribution distance: Sinshom,, Weight 10.000
Medion. Avg.
Aug. 8.758: 0.111
Modian, 0.767: Aug.
30
30
0.4
4.5
05
07
0.6
09
0.4
0.
06
0.
09
0.5
06
6.7
0.6
69
003
030
QED Distributionn Stop #90
QED Distribution Stop #1000
Distributien Sinkhorn, Weight: 10,000
Distribution distance: Smithom Weight
Mediam Auy
40-
20-
18-
044 05
0.6 0. 0.8 0.5
03 0.4
0.5
06
07 0.5 0.9
Figure 6: Histograms of QED values for the decoded SMILES strings corresponding to the point clouds at time step L = 0, 10, 100 for the JKO-ICNN experiment that uses Sinkhorn as the distribution distance with weight 10,000. We observe a clear shift to the right in the distribution, which corresponds to increased drug-likeness of the decoded molccule strings. This figure displays results for the MOSES dataset experiment.
F.1 Convex QED surrogate classifier (QM9)
For QM9, the convex surrogate was trained to predict high (> 0.5) and low 0.5) QED values from molecule embeddings coming from a pre-trained VAE. The threshold for QM9 molecules was set to a lower value compared to that for the MOSES datasct because the underlying QED distribution of train and test data for QM9 molecules has significantly lower values compared to MOSES. As above, molecules with high QED were given lower value labels compared to the low QED molecules so that when this model would be used as potential, minimizing this functional would lead to higher QED values. Similar to the MOSES dataset pipeline, for this convex surrogate, trained Residual ICNN (Amos et al., 2017; Huang et al., 2021) model with four hidden layers, each with dimension 128, which was the dimensionality of the input molecule
27

---

Published in Transactions on Machine Learning Research (07/2022
Table Molecular discovery with direct optimization of functional objective using the SGD optimizer: Measures of validity, uniqueness, and median QED are reported in each row for the corresponding final point cloud of embeddings in each configuration. We also report the final Sinkhorn divergence between the initial cloud point and that from the final time step (Final SD). Each mcasurement value cell contains mean values + one standard deviation for five repeated runs of the experiment with different random initialization seeds. This table contains results for the MOSES dataset. t For the SGD optimizer, the only configurations that meaningfully increase the QED distribution of the given point cloud are those with relatively large learning rate, 0.5. However, this leads the point cloud to regions that decode to invalid molocule strings and Final SD being several orders of magnitude higher compared to the JKO-ICNN approach.
Grid search for SGD optimizer I.R
Validity
Uniqueness
QED Median
Final SD
0.5
43.440 1: 1.092
100.000 1: 0.000
0.772 1 0.004
9792.929 1: 76.913
90.960 + 0.055
99.980 + 0.045
0.630 + 0.002
4.584 +:
0
0.01
99.960 + 0.055
99.950 = 0.045
0.630 = 0.001
0.022 = 0.000
0.001
99.980 1: 0.045
99.980 1 0.045
0.630 1 0.001
0.017 1: 0.000
0.0001
100.000 0.000
99.980 0.045
0.630 0.001
0.017 0.000
0.5
19.140 + 1.128
100.00 + 0.000
0.768 = 0.006
8881,378 + 796
0.1
99 99.9960 + 0.055
99.980 = 0.045
0.630 = 0.002
4.496 + 1.156
0.01
99 960 + 0.055
99,980 0.045
0,630
0.022 0.000
0.001
99.980 + 0.015
99.980 = 0.045
0.630 = 0.001
0.017 + 0.000
0.0001
100.000 0.000
99.980 + 0.045
0.630 = 0.001
0,017 0.000
0.5
83.660 1 0.918
100.000 0.000
0.771 0.007
3696.65 650 1: 28.212
10
0.1
99.960 t 0.055
99.980 = 0.U45
0.630 = 0.002
3.832 + 0.917
0.01
99.960 1 0.055
99,980 1
0,630 1 0.001
0,021
10
0.00
90.980 + 0.045
99.980 = 0. 0.045
0.630 = 0.000
0.017 + 0.000
10
0.0001
100.000 0.000
99. 980 = 0.045
0.630 = 0.001
0,017 + 0.000
100
0.5
96.920 +: 0.576
98.885 + 0.308
0.635 + 0.002
547.310 36.557
100
0.1
09.960 + 0.055
99.980 = 0. 0.045
0.630 = 0.002
0.822 + 0.525
100
0.01
99 960 + 0.055
99.980 = 0.045
0,630 = 0.001
0.021 + 0.000
100
0,0U1
99.980 + 0.045
99.980 = 0.045
0,630 0,001
0,017 + 0,000
100
0.0001
100.000 0.000
99.980 1 0.045
0.630 T 0.001
0.017 1 0.000
1,000
0.5
87 240 + 0.777
100.000 0.000
0.767 = 0.002
2515. 075 + 19.870
1,000
0.1
99.960 1 0.055
99.980 = 0.045
0.630 = 0.002
1.833 1 0.611
1,000
0.01
99 00. 1 0.055
99.980 + 0.045
0.630 7F 0.001
0.019 1 0.000
1,000
0.001
99. 99.980 + 0.045
99.980 0.045
0.630 0.001
0.017 0.000
1,000
0,0001
100.000 + 0.000
99.980 = 0.045
0.630 = 0.001
0.017 + 0.000
10,000
nan uan
nan nan
nan uan
man nan
10,000
0.1
98.50 98.500 + 0.235
99,950 = 0.045
0,641 = 0.003
296.809 + 9.311
10,000
0.01
99.980 + 0.045
99.980 = 0.045
0,630 = 0.001
U.017 + 0.000
10,000
0.001
99.980
1 0.045
99.980 0.045
0.630 0.001
0.017 0.000
10,000
0.0001
100.000
+ 0.000
99.980 - 0.045
0.630 - 0.001
0.017 + 0.000
embeddings as well. We trained the model with binary cross-entropy loss. To maintain convexity in the potential functional, we used the last layer before the sigmoid activation for V. The model was trained with an initial learning rate of 0.01, batch sizes of 1,024, ADAM optimizer, and a learning rate scheduler that decreased learning rate on validation set loss plateau. The model was trained for 100 epochs, and the weights from the final epoch were used to initialize the convex surrogate in the potential functional. For this epoch, the model achieved 88.6% accuracy on the test set. In Figure 7, we display the test set confusion matrix for this final epoch.
F.2 JKO-ICNN QED histograms (QM9)
In Table 9, we present the same results as in Appendix F.5. For the JKO-ICNN flow on the QM9 dataset, we started with an initial distribution of embeddings Po that had corresponding QED value of < 0.35. This initial point cloud was taken from the QM9 train set since the test set is quite small and does not contain enough data points below the starting QED threshold.
As seen with the experiment on MOSES, all four combinations are able to increase QED values. However, the sctups that use MMD or X 1,000 lead to mode collapse, see the discussion in Appendix E.3.
In Figure 8, we present several time steps of the histograms of the QED values for the decoded SMILES strings corresponding to the point clouds for the experiment that used Sinkhorn as the distribution distance with weight 10,000.
28

---

Published in Transactions on Machine Learning Research (07/202)
Table 7: Molecular discovery with direct optimization of functional objective using the ADAM optimizer: Measures of validity, uniqueness, and median QED are reported in each row for the corresponding final point cloud of embeddings in each configuration. We also report the final Sinkhorn divergence between the initial cloud point and that from the final time step (Final SD). Each measurement value cell contains mcan values + one standard deviation for five repeated runs of the experiment with different random initialization seeds. This table contains results for the MOSES dataset. f For the ADAM optimizer, although there are some hyperparameter configurations that are able to meaningfully increase QED of the point cloud without sacrificing the validity of the decoded molecule strings, these configurations require X2 to either be 0 or several orders of magnitude lower compared to in the JKO-ICNN approach. This direct optimization approach is therefore unable to successfully meet both goals of the objectíve, higher QED with small distance from the original point cloud, as seen by the larger Final SD values.
Grid search for ADAM optimiser Validity
LIt
Uniqueness
QED Median
Pinal SD
7.420
0.729
98.479
1.595
0.654 0.016
441.779
1.168
0.1
92.080 + 0.973
100.000
0.000
0.793 0.005
18:261
0.134
0,01
93.900 + 0.781
99.979 + 0.018
0.758 + 0.006
1.650 = 0.006
0.001
99.320 + 0.164
99.980 0.045
0.632 0.001
0.073 = 0.000
0,0001
99.900 + 0.100
99.980 0.045
0.630 0.002
0.018 = 0.000
0.5
16.240 + 1.553
100.000 = 0.000
0.740 + 0.005
394.993 = 1.093
0.1
91.200 +: 0.539
99.978 0.049
0.792 + 0.005
17.170 + 0.097
0.01
95. + 0.505
99 99.979 0.047
0.702 0.004
1.551 = 0.009
0.001
99.340 + 0.167
99.980 0.015
0.632 0.001
0.072 = 0,000
0.0001
99.900 1 0.100
99.980 0.045
0.630 1. 0.002
0.018 1 0.000
10
0.5
96.940 1: 0.182
92.635
0.496
0.656 0.003
241.505 1.616
10
0.1
95.380 + 0.618
99.931 + 0.057
0.701 + 0.005
13.360 + 0.030
0.01
98.900 1: 0.464
99.980 0.045
0.637 0.002
1.253 + 0.005
10
0.001
99.680 + 0.164
99.980 + 0.045
0.631 + 0.001
0.068 = 0.000
10
0.0001
99.900 + 0.100
99.980 0.045
0.630 + 0.001
0.018 = 0.000
100
0.5
99.680 + 0.084
90.610 0.783
0.631 0.002
18,000 + 0.605
100
0.1
99.840 1 0.182
99.960 0.055
0.630 1 0.001
4.528 1 0.068
100
0.01
99.920 +: 0.084
99.980 0.045
0.630 0.001
0.246 = 0.000
100
0.001
99.980 + 0.015
99.980 + 0.015
0.630 + 0.001
0.054 = 0.000
100
0.0001
99.980 + 0.045
99.980 0.045
0.630 0.002
0.018 + 0.000
1,000
0.5
96.140 + 0.279
99.105 + 0.262
0.662 + 0.006
12.146 + 0.347
1,000
0.1
99.980 + 0.045
99.980 0.045
0.630 + 0.001
0.077 = 0.003
1,000
0.01
99.94 1 0.055
99.980
0.045
0.630 0.001
0.023 1 0.000
1,000
0.001
99.980 1 0.045
99.980 + 0.045
0.630 + 0.001
0.021 + 0.000
1,000
0.0001
99.980 + 0.045
99.980 0.045
0.630 0.001
0.018 = 0,000
10,000
0.5
99.020 + 0.295
97 97.697 + 0.221
0.635 + 0.002
2.616 = 0.062
10,000
99.900 + 0.122
99.980 0.045
0.630 0.001
0.240 = 0.019
10,000
0.01
99.980 - 0.015
99.980 0.015
0.630 0.001
0.018 = 0.000
10,000
0.001
99.980
4 0.045
99.980 0.045
0.630 0.001
0.017 = 0.000
10,000
0.0001
99.980
1 0.045
99.980
0.045
0.001 0.630
0.017 1 0.000
G Assets
Software Our implementation of JKO-ICNN relies various open-source libraries, including pytorch (Paszke et al., 2019) (license: BSD), pytorch--ightning (Falcon et al., 2019) (Apache 2.0), POT (Flamary et al., 2021) (MIT), geomloss (Feydy et al., 2019) (MIT), rdkit (Landrum, 2013a;b) (BSD 3- 3-Clause).
Data All data used in Section 5 is synthetic. The MOSES dataset is released under the MIT license. The QM9 dataset does not explicitly provide a license in their website nor data files.
29

---

Published in Transactions on Machine Learning Research (07/2022)
Table Computational amortization gains. Comparison between re-use of maps calculated from JKO-ICNN flow and the direct optimization baseline applied to various sizes of embedding point clouds. For each sctup, report median QED for the final point cloud of embeddings from the respective flows and the Sinkhorn divergence between the initial and final point clouds (Final SD). We observe a linear trend in speed-up up when re- re-using JKO-ICNN maps. This table contains results for the MOSES dataset experiment.
Size
QED Median
Final SD
Time (s)
Speedup
Baseline
1,000
0.789
17.134
1.416
2,000
0.795
17.333
1.550
3,000
0.791
17.255
2.300
4,000
0.790
17.238
3.433
5,000
0.792
17.212
4.696
JKO-ICNN
maps
1,000
0.778
1.739
0.086
1.330
2,000
0.773
1.440
0.086
1.464
3,000
0.771
1.363
0.086
2.213
4,000
0.767
1.373
0.086
3.347
5,000
0.764
1.384
0.088
4.608
0.8
0,83
0.17
Ground Truth
0.6
0.4
0.087
0.91
-0.2
Predicted
Figure 7: Confusion matrix for the convex surrogate trained to predict QED labels from molecule embeddings for QM9.
30

---

Published in Transactions on Machine Learning Research (07 (07/2022)
Table Molecular discovery with JKO-ICNN experiment results (QM9) Measures of validity, uniqueness, and median, average, and standard deviation QED reported in each row for the corresponding point cloud of embeddings. For each point cloud, decode the embeddings to get SMILES strings and use RDKit to determine whether the corresponding string is valid and get. the associated QED value. In the first. row, we display the point cloud at time step zero of JKO-ICNN. In the subsequent rows, we display the values for each measurement al the final time step T of JKO-ICNN for different hyperparameter configurations of distribution distance (either Sinkhorn or MMD) and weight on this distance X2 (either 1,000 or 10,000). Each mcasurement value cell contains mean values + one standard deviation for five repeated runs of the experiment with different random initialization seeds. This table contains results for the QM9 dataset. experiment.
Measure
D
Aa
Validity
Uniqueness
QED Median
QED Avg.
QED Std.
N/A
N/A
100.000 + 0.000
99.840 + 0.134
0.315 + 0.001
0.308 = 0.001
0.041 = 0.001
Sinkhorn
1,000
90.840 + 1.457
33.367 + 2.491
0.381 + 0.024
0.373 = 0.016
0.096 = 0.005
PT
Sinkborn
10,000
92.700
+ 0.828
81.925
= 1.982
0.419
0.005
0.404
0.005
0.096 = 0.004
MMD
1,000
91.680
+ 4.463
22.424
+ 1.185
0.452 0.024
0.434
0.019
0.094 = 0.005
MMD
10,000.
88.800
1 4.661
28.664
1 1.500
0.448
0.013
0.432 1 0.010
0.093 1 0.005
31

---

Published in Transactions on Machine Learning Research (07/2022)
Distribttin Weight: 10,000
QED Distniuttin Stop-
Datribution distance Sincham Weight 10,000
QED Distribtiinn Step #10
Detribution distance Snichom Weight
QED Distrbutiin Stop #20
Median 0.315: Avg. 0.302: Std. 0.042
fug:
Aug
0.35
0.00
0.25
0.32
0.59
62
03
04
05
06
0.2
0.3
05
06
QTD
QED Distributtinn Step #30
QED Distribution Step #:00
QED Distribution Step #50
Distribution Sinkhorn: Weight 30,000 Median. 0.400:
Distibutionn distance: Weight Median. 0:403: 0.094
Distribution Sinthom: Weight 10.000 Medion. 0:422:
0.057
Coart
30
0.1
02
0.5
0.4
0.5
06
0.2
0.3
0.4
05
06
01
02
0.5
0.4
0.5
2.6
GED Distribution Step #50
DED Distribution Step #00
DED Distribution Step #80
Distribution Sinkhorn, Weight: Medion. 0.427:
Distribution Weight: Medion. 0.407: 0.094
Distribution Mesionn 0.422:
0,065
Chart
20
20
02
0.5
0.6
0.2
0.5
06
0.5
Distribution Stop 10,000 Distribution distance:
Distribution Distribution Step #1000
10.000
Mediaam 0/422. 0.609, Stdd 2,094
Nediam: 425, D.S,, Sidi D.096
Figure 8: Histograms of QED values for the decoded SMILES strings corresponding to the point clouds at time step t 0, 10. 100 for the JKO-ICNN experiment that uses Sinkhorn as the distribution distance with weight 10,000. We observe a clear shift to the right in the distribution, which corresponds to increased drug-likeness of the decoded molecule strings. This figure displays results for the QM9 dataset experiment.
32


Reviews:
[{"review_id": "sY4lpvYblw", "reviewer": "Reviewer_jczu", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "- Strengths \n  - Using ICNNs to approximate JKO schemes appears to be a useful methodological contribution and it appears to perform well empirically and scales to moderately high dimension\n  - Experiments are interesting, thorough, and validate the method both for closed-form PDEs and for the molecular experiment at dimension 128 (latent of VAE)\n\n- Weaknesses\n  - My understanding is that ICNNS require training at each time-step or fine-tuning from a warm-start initialization. I feel this would be quite slow for a large number of JKO steps if training to completion. \n  - On page 9 it is mentioned that the authors do not use warm-start, use only 500 inner training steps for the ICNN, and a relatively small only 2 layer ICNN, does this converge in so few steps? Does warm-start not work given it is not used? I imagine pretraining without meta-learning may result in the network being initialized from a local mode and not adapting.\n\n  - Although TMLR does not review for novelty or originality and I agree that it may be concurrent, there is little methodological contribution beyond previously published [1].  I understand that unlike [1], the authors demonstrate the proposed method in high dimensional settings, for a larger number of JKO steps, however the underlying methodological contribution is very similar. It may be worth expanding on the differences.\n\n\n[1] Proximal Optimal Transport Modeling of Population Dynamics, Bunne 2021", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "- The authors propose to use ICNNs to approximate JKO schemes with a number of fixed energy functionals\n- The authors validate the proposed method in moderately high dimension (128 dimension VAE latent) and compare to PDEs with known solutions\n\n", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "I would like the authors' comments on the following points. These are not requests but more request for discussion.\n\n- The authors mention the direct optimization approach multiple times, including in discussion, and use it as the main baseline but do not explain what it is until the appendices. It may be worth commenting what it is briefly in the main text. This is perhaps more taste than needed.\n\n- It is commented briefly on page 10 that direct optimization does not work very well, is there any detailed reason why? The direct optimization approach mentioned in the appendices is interesting as it seems to not require training an ICNN at each iteration. A a side note, this approach could also be used to learn the functional similar to [1] without needing to backprop through the training of an ICNN and hence perhaps more scalable than [1] which appears to use only 4 JKO steps in low dimension.\n\n- It would be useful to give an indication of how long it takes to run a single forward pass of the JKO-ICNN scheme, training ICNN at each step etc., compared to the other approaches and detail the effect of warm-start as it is not clear if it is used.\n\n- It appears the ICNNs used are relatively small, just two layers with width 100. Is this standard in ICNN literature or do practitioners typically use larger networks? Does this impact the training time and feasibility of the proposed method for a large number of JKO steps?\n\n- The experiment runs at dimension 128 in the latent space of a VAE, does the method scale to higher dimensions or will this require larger ICNNs and hence longer training per time-step and would this be prohibitively slow?\n\n[1] Proximal Optimal Transport Modeling of Population Dynamics, Bunne 2021", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "y7xtEtcR4k", "reviewer": "Reviewer_83vF", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\nOverall, this is a very nice work, well written and with an interesting experiment.\n\nStrengths:\n- The paper is well written, and simple to follow\n- The molecular discovery is an interesting and original experiment. This shows that gradient flows in high dimension could be used on \"real\" experiments.\n\nWeaknesses:\n- Incremental work (mainly plug ICNN into the JKO scheme, the change of variable formulas are well known for normalizing flows, notably for convex potential flows [1] which are basically the same transformations used in this work, but trained for different objectives).\n- High dimensional experiment \"only\" in R^128\n- I believe some points in the experiments could be clarified or improved.\n\nI have some question about the experiments. In Section 5.1, it is stated \"here the flow density is estimated from particles via KDE\". However, the ICNN learned is supposed to be strictly convex if I understand correctly, and hence we should be able to evaluate directly the density of the flow. Is there a reason why the KDE is used?\n\nIn Section 5.2, for the experiment performed, there is no closed-form solution to compare with. But it would be possible for the particular case m=1, a quadratic V and an initial Gaussian. I noticed that it is the setting tested in Appendix D. However, Figure 3 plots the MMD between the JKOICNN trajectory and the Unadjusted Langevin Algorithm (ULA/SGLD). However, in this case, there is a closed-form (see e.g. [2]). Hence, I believe this experiment could be improved. Furthermore, ULA might actually perform worse than JKO-ICNN (see e.g. [2]) so I'm not convinced by this Figure.\n\nFor the molecular discovery experiment, I do not believe that the used divergence D used is described in the main part of the paper. I found it in appendix E.2, but I believe it should be in Section 6.\n\n[1] Huang, C. W., Chen, R. T., Tsirigotis, C., & Courville, A. (2020). Convex potential flows: Universal probability distributions with optimal transport and convex optimization. arXiv preprint arXiv:2012.05942.\n\n[2] Mokrov, P., Korotin, A., Li, L., Genevay, A., Solomon, J. M., & Burnaev, E. (2021). Large-scale wasserstein gradient flows. Advances in Neural Information Processing Systems, 34.", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this work, authors build on the work of [1] in which it was proposed to solve the JKO scheme by reparametrizing it as a minimization problem over convex functions. In [1], the convex functions were approximated using a space discretization which does not scale in dimension higher than 2. In this paper, it is proposed to use Input Convex Neural Networks (ICNN) to parametrize the set of convex functions. The method is called JKO-ICNN. This has the advantage to scale better with the dimension. Then, experiments are conducted to show that JKO-ICNN are able to learn well Wasserstein gradient flows for simple PDEs. Furthermore, they conduct an experiment in higher dimension in controlled generation for molecular discovery.\n\n[1] Benamou, J. D., Carlier, G., Mérigot, Q., & Oudet, E. (2016). Discretization of functionals involving the Monge–Ampère operator. Numerische mathematik, 134(3), 611-636.", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "\nI have few remarks on the Related Works section. First, in the paragraph ICNN, optimal transport, and generative modeling, it is written \"we consider the setting where the target distribution cannot be sampled from, and is only implicitly characterized as the minimizer of an optimization problem over distributions\". I find this sentence confusing since JKO-ICNN can also cover the case where one would want to minimize some loss w.r.t. a distribution from which we have only access to samples. Maybe it should be clarified by saying that the setting here is more general than the other generative modeling works since any loss can basically be plugged in and minimized.\n\nAlso, about convex potential flows, it is written that they \"are useful for learning generative models when samples from the target (i.e., optimal) distributions are available and the goal is to learn a parametric generative model\". While it is true, it can also be used when we only have access to the target up to a constant (using the forward KL instead of the reverse KL). \n\nAnother remark is on the related works cited. For example, the JKO scheme (or variations of it) with neural networks is also used in [1, 2, 3]. A method to solve the JKO scheme with the entropic regularization is proposed in [4]. Some works which solve Wasserstein gradient flows using particle schemes (via the forward scheme) might also be cited, e.g. [5,6,7,8,9] for different functionals or [10, 11] for the KL with neural networks.\n\nTypos:\n- In Section 4, after quotation of Salim et al, \"show that in enjoys\"\n\n\n[1] Lin, A. T., Li, W., Osher, S., & Montúfar, G. (2021, July). Wasserstein proximal of GANs. In International Conference on Geometric Science of Information (pp. 524-533). Springer, Cham.\n\n[2] Fan, J., Taghvaei, A., & Chen, Y. (2021). Variational Wasserstein gradient flow. arXiv preprint arXiv:2112.02424.\n\n[3] Bonet, C., Courty, N., Septier, F., & Drumetz, L. (2021). Sliced-Wasserstein Gradient Flows. arXiv preprint arXiv:2110.10972.\n\n[4] Caluya, K. F., & Halder, A. (2019, July). Proximal recursion for solving the Fokker-Planck equation. In 2019 American Control Conference (ACC) (pp. 4098-4103). IEEE.\n\n[5] Liutkus, A., Simsekli, U., Majewski, S., Durmus, A., & Stöter, F. R. (2019, May). Sliced-Wasserstein flows: Nonparametric generative modeling via optimal transport and diffusions. In International Conference on Machine Learning (pp. 4104-4113). PMLR.\n\n[6] Arbel, M., Korba, A., Salim, A., & Gretton, A. (2019). Maximum mean discrepancy gradient flow. Advances in Neural Information Processing Systems, 32.\n\n[7] Korba, A., Aubin-Frankowski, P. C., Majewski, S., & Ablin, P. (2021, July). Kernel Stein Discrepancy Descent. In International Conference on Machine Learning (pp. 5719-5730). PMLR.\n\n[8] Glaser, P., Arbel, M., & Gretton, A. (2021). KALE Flow: A Relaxed KL Gradient Flow for Probabilities with Disjoint Support. Advances in Neural Information Processing Systems, 34.\n\n[9] Wang, Y., Chen, P., & Li, W. (2021). Projected Wasserstein gradient descent for high-dimensional Bayesian inference. arXiv preprint arXiv:2102.06350.\n\n[10] Feng, X., Gao, Y., Huang, J., Jiao, Y., & Liu, X. (2021). Relative Entropy Gradient Sampler for Unnormalized Distributions. arXiv preprint arXiv:2110.02787.\n\n[11] di Langosco, L. L., Fortuin, V., & Strathmann, H. (2021). Neural Variational Gradient Descent. arXiv preprint arXiv:2107.10731.", "broader_impact_concerns": "No comments for this part.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "aHojTfVImi", "reviewer": "Reviewer_x751", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "STRENGTH:\n\n* The paper is well written with few typos. It is well motivated and the proofs are correct.\n\n* The extension of the parametric JKO setting to internal energies is interesting. \n\n* I also enjoyed the applications to non-linear PDEs and to molecular discovery. In particular, I found interesting the thorough comparison between JKO + ICNN and the \"direct approach\" (which consists into discretizing the ODE associated with the Wasserstein gradient flow). The JKO approach seems to improve on the direct approach by managing to propose configurations with good QED which are close to the original point cloud. \n\nWEAKNESSES:\n\n* It is hard to understand what are the true contributions of the paper, especially when compared to [1]. In particular, Algorithm 1 was already proposed in [1]. Not taking into the account the other contribution of [1] which considers the learning of the energy, the only difference I can see between the two approaches is that the current paper consider more general energies. From a methodological point of view, this seems incremental. I think the authors should do a better job at motivating the application of ICNN + JKO to these extended energies.\n\n* The molecular discovery application is interesting but not enough details in the main document are given about the functional and its optimization. In particular, how does the divergence term D fits in the general picture presented by the authors? In the supplementary material it is stated that D is either the (entropy regularized) Sinkhorn divergence, or the Maximum Mean Discrepancy. It is not clear to me how such divergence can be understood in the theoretical framework of the authors (even though I agree that the computations can be done by unrolling the minimization in the computation of the divergence).\n\n[1] -- Bunne, Meng-Papaxanthos, Krause, Cuturi -- JKOnet: Proximal optimal transort modeling.", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper deals with the optimization of functionals on the space of probability measures. Following a gradient flow approach, the authors propose to investigate the JKO scheme to approximate the continuous-time Wasserstein gradient flow. Following earlier works on the solution of JKO, they parametrize the solutions of the JKO problems using convex pushforward maps. Using the recently introduced Input Convex Neural Networks (ICCN) and restricting further to the class of pushforward defined by ICNN, the authors obtain an approximate solution of the JKO problem by minimizing the JKO loss w.r.t the parameters of the ICNN. They investigate different energy function (potential, interaction, internal). They then evaluate their methologies on different non-linear PDEs whose solutions or equilibirum are known. Finally, they propose to apply their algorithm for molecular discovery. To do so, they define the JKO problem in the latent space of a VAE.\n\nMain contributions are :\n* Extending ICNN + JKO framework to more general energies.\n* Application to PDEs + Molecular discovery (comparison with direct optimization).", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "MAJOR :\n\n* More details are needed in the main document, regarding the molecular experiment and in particular how the authors deal with the Sinkhorn or MMD divergence term.\n\n* It should be made more clear that the ICNN + JKO scheme was already introduced in the literature. If I understood correctly the main contribution of the authors in fact resides into the extension of the ICNN + JKO method to other energies.\n\nOTHER COMMENTS :\n\n* Computing and backpropagating through the logarithm of the determinant of the Hessian can be extremely costly as the dimension increase. The authors solve this problem using the Hutchinson method to obtain an unbiased estimator of the trace (which is needed when differenciating). However, more and more samples are needed in order to control the variance of the estimator as the dimension grows. Hence, I am wondering how the method proposed by the authors would really scale to high dimension. Could the authors discuss the limitations of the Hutchinson estimator (if any)?\n\n* In Appendix D, the authors refer to the Stochastic Gradient Langevin Dynamics (SGLD) (there is a typo in the recursion X_K should be X_k). However, the dynamics they use is a very special case of SGLD where there is no noise on the gradient term and therefore SGLD is simply the discretization of the Langevin dynamics (sometimes referred to as ULA - Unadjusted Langevin Algorithm).\n\n* Table 6 is quite hard to read. Highlighting the important and relevant lines would help the reader, trying to check the behavior of the direct optimization baseline. \n\n* Do the authors have any explanation as of why the direct optimization does not work as well as JKO+ICNN? This is an interesting and potentially important findings but it is not obvious to me why JKO would performed better.\n\n* In the molecular experiment the ADAM baselin with \\lambda_2 = 0 and 1e-2 give reasonable result except for the final SD (which makes sense since there is no regularization term to fit the distribution to \\rho_0). However, would it be possible to compare the point cloud generated with this configuration and the one obtained with JKO? \n\n* In Table 7, the authors compare the speed-up when using the JKO-ICNN to other configurations while the direct approach needs to be retrained entirely. I totally agree that in this case, it will benefit JKO. However (except if I missed something), JKO+ICNN will likely take longer to train in the first place than just the optimization using the direct approach. Could the authors comment on that as well? This will of course depend on the dimensionality + the number of parameters of the ICNN.", "broader_impact_concerns": "No concern.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper makes an interesting contribution by developing a novel scheme to implement gradient flows optimizing functionals on the space of probability distributions. It proposes an approximation of the celebrated JKO scheme using ICNN. The methodology is demonstrated on a variety of applications including molecular discovery.

The paper has now been reviewed by three reviewers: two recommend accept and one leans towards acceptance (however one of the main concerns of this reviewer is novelty but, as pointed out by the authors, the first version of this work appeared at the same time as other JKO-neural net schemes). Thus I recommend acceptance of the paper.

A minor point has been raised by reviewer x751: Appendix D appears to use erroneously the terminology SGLD. SGLD is an approximation of Unadjusted Langevin algorithm (ULA) using a Monte Carlo approximation of the gradient (mini batches in the original paper of Welling & Teh). It seems you are using here ULA and not SGLD. Please correct it.
Decision: Accept (poster)