Large-Scale Wasserstein Gradient Flows

Abstract:
Wasserstein gradient flows provide a powerful means of understanding and solving many diffusion equations. Specifically, Fokker-Planck equations, which model the diffusion of probability measures, can be understood as gradient descent over entropy functionals in Wasserstein space. This equivalence, introduced by Jordan, Kinderlehrer and Otto, inspired the so-called JKO scheme to approximate these diffusion processes via an implicit discretization of the gradient flow in Wasserstein space. Solving the optimization problem associated to each JKO step, however, presents serious computational challenges. We introduce a scalable method to approximate Wasserstein gradient flows, targeted to machine learning applications. Our approach relies on input-convex neural networks (ICNNs) to discretize the JKO steps, which can be optimized by stochastic gradient descent. Unlike previous work, our method does not require domain discretization or particle simulation. As a result, we can sample from the measure at each time step of the diffusion and compute its probability density. We demonstrate our algorithm’s performance by computing diffusions following the Fokker-Planck equation and apply it to unnormalized density sampling as well as nonlinear filtering.

heading: 1 Introduction
text: Stochastic differential equations (SDEs) are used to model the evolution of random diffusion processes across time, with applications in physics [63], finance [22, 52], and population dynamics [35]. In machine learning, diffusion processes also arise in applications filtering [34, 21] and unnormalized posterior sampling via a discretization of the Langevin diffusion [70].
The time-evolving probability density ρt of these diffusion processes is governed by the FokkerPlanck equation. Jordan, Kinderlehrer, and Otto [32] showed that the Fokker-Planck equation is
∗Equal contribution.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
equivalent to following the gradient flow of an entropy functional in Wasserstein space, i.e., the space of probability measures with finite second order moment endowed with the Wasserstein distance. This inspired a simple minimization scheme called JKO scheme, which consists an implicit Euler discretization of the Wasserstein gradient flow. However, each step of the JKO scheme is costly as it requires solving a minimization problem involving the Wasserstein distance.
One way to compute the diffusion is to use a fixed discretization of the domain and apply standard numerical integration methods [18, 49, 15, 17, 40] to get ρt. For example, [50] proposes a method to approximate the diffusion based on JKO stepping and entropy-regularized optimal transport. However, these methods are limited to small dimensions since the discretization of space grows exponentially.
An alternative to domain discretization is particle simulation. It involves drawing random samples (particles) from the initial distribution and simulating their evolution via standard methods such as Euler-Maruyama scheme [36, M9.2]. After convergence, the particles are approximately distributed according to the stationary distribution, but no density estimate is readily available.
Another way to avoid discretization is to parameterize the density of ρt. Most methods approximate only the first and second moments ρt, e.g., via Gaussian approximation. Kalman filtering approaches can then compute the dynamics [34, 39, 33, 61]. More advanced Gaussian mixture approximations [65, 1] or more general parametric families have also been studied [64, 69]. In [48], variational methods are used to minimize the divergence between the predictive and the true density.
Recently, [24] introduced a parametric method to compute JKO steps via entropy-regularized optimal transport. The authors regularize the Wasserstein distance in the JKO step to ensure strict convexity and solve the unconstrained dual problem via stochastic program on a finite linear subset of basis functions. The method yields unnormalized probability density without direct sample access.
Recent works propose scalable continuous optimal transport solvers, parametrizing the solutions by reproducing kernels [10], fully-connected neural networks [62], or Input Convex Neural Networks (ICNNs) [37, 44, 38]. In particular, ICNNs gained attention for Wasserstein-2 transport since their gradients ∇ψθ : RD → RD can represent OT maps for the quadratic cost. These continuous solvers scale better to high dimension without discretizing the input measures, but they are too computationally expensive to be applied directly to JKO steps.
Contributions. We propose a scalable parametric method to approximate Wasserstein gradient flows via JKO stepping using input-convex neural networks (ICNNs) [6]. Specifically, we leverage Brenier’s theorem to bypass the costly computation of the Wasserstein distance, and parametrize the optimal transport map as the gradient of an ICNN. Given sample access to the initial measure ρ0, we use stochastic gradient descent (SGD) to sequentially learn time-discretized JKO dynamics of ρt. The trained model can sample from a continuous approximation of ρt and compute its density dρtdx (x). We compute gradient flows for the Fokker-Planck free energy functional FFP given by (5), but our method generalizes to other cases. We demonstrate performance by computing diffusion following the Fokker-Planck equation and applying it to unnormalized density sampling as well as nonlinear filtering.
Notation. P2(RD) denotes the set of Borel probability measures on RD with finite second moment. P2,ac(RD) denotes its subset of probability measures absolutely continuous with respect to Lebesgue measure. For ρ ∈ P2,ac(RD), we denote by dρdx (x) its density with respect to the Lebesgue measure. Π(µ, ν) denotes the set of probability measures on RD×RD with marginals µ and ν. For measurable T : RD → RD, we denote by T] the associated push-forward operator between measures.

heading: 2 Background on Wasserstein Gradient Flows
text: We consider gradient flows in Wasserstein space (P2(RD),W2), the space of probability measures with finite second moment on RD endowed with the Wasserstein-2 metricW2. Wasserstein-2 distance. The (squared) Wasserstein-2 metricW2 between µ, ν ∈ P2(RD) is
W22 (µ, ν) def = min
π∈Π(µ,ν) ∫ RD×RD ‖x− y‖22 dπ(x, y), (1)
where the minimum is over measures π on RD × RD with marginals µ and ν respectively [68].
For µ ∈ P2,ac(RD), there exists a µ-unique map ∇ψ∗ : RD → RD that is the gradient of a convex function ψ∗ : RD → R t {∞} satisfying ∇ψ∗]µ = ν [46]. From Brenier’s theorem [13], it follows that π∗ = [idRD ,∇ψ∗]]µ is the unique minimizer of (1), i.e.,
W22 (µ, ν) = ∫ RD ‖x−∇ψ∗(x)‖22 dµ(x).
Wasserstein Gradient Flows. In the Euclidean case, gradient flows along a function f : R → R follow the steepest descent direction and are defined through the ODE dxtdt = −∇f(xt). Discretization of this flow leads to the gradient descent minimization algorithm. When functionals are defined over the space of measures equipped with the Wasserstein-2 metric, the equivalent flow is called the Wasserstein gradient flow. The idea is similar: the flow follows the steepest descent direction, but this time the notion of gradient is more complex. We refer the reader to [4] for exposition of gradient flows in metric spaces, or [59, Chapter 8] for an accessible introduction.
A curve of measures {ρt}t∈R+ following the Wasserstein gradient flow of a functional F solves the continuity equation
∂ρt ∂t = div(ρt∇xF ′(ρt)), s.t. ρ0 = ρ0, (2)
where F ′(·) is the first variation of F [4, Theorem 8.3.1]. The term on the right can be understood as the gradient of F in Wasserstein space, a vector field perturbatively rearranging the mass in ρt to yield the steepest possible local change of F . Wasserstein gradient flows are used in various applied tasks. For example, gradient flows are applied in training [8, 43, 25] or refinement [7] of implicit generative models. In reinforcement learning, gradient flows facilitate policy optimization [55, 72]. Other tasks include crowd motion modelling [45, 58, 50], dataset optimization [2], and in-between animation [26].
Many applications come from the connection between Wasserstein gradient flows and SDEs. Consider an RD-valued stochastic process {Xt}t∈R+ governed by the following Itô SDE:
dXt = −∇Φ(Xt)dt+ √
2β−1dWt, s.t. X0 ∼ ρ0 (3) where Φ : RD → R is the potential function, Wt is the standard Wiener process, and β > 0 is the magnitude. The solution of (3) is called an advection-diffusion process. The marginal measure ρt of Xt at each time satisfies the Fokker-Planck equation with fixed diffusion coefficient:
∂ρt ∂t = div(∇Φ(x)ρt) + β−1∆ρt, s.t. ρ0 = ρ0. (4)
Equation (4) is the Wasserstein gradient flow (2) for F given by the Fokker-Planck free energy functional [32] FFP(ρ) = U(ρ)− β−1E(ρ), (5) where U(ρ) = ∫ RD Φ(x)dρ(x) is the potential energy and E(ρ) = − ∫ RD log dρ dx (x)dρ(x) is the entropy. As the result, to solve the SDE (3), one may compute the Wasserstein gradient flow of the Fokker-Planck equation with the free-energy functional FFP given by (5). JKO Scheme. Computing Wasserstein gradient flows is challenging. The closed form solution is typically unknown, necessitating numerical approximation techniques. Jordan, Kinderlehrer, and Otto proposed a method—later abbreviated as JKO integration—to approximate the dynamics of ρt in (2) [32]. It consists of a time-discretization update of the continuous flow given by:
ρ(k) ← arg min ρ∈P2(Rn)
[ F(ρ) + 1
2h W22 (ρ(k−1), ρ)
] (6)
where ρ(0) = ρ0 is the initial condition and h > 0 is the time-discretization step size. The discrete time gradient flow converges to the continuous one as h → 0, i.e., ρ(k) ≈ ρkh. The method was further developed in [4, 60], but performing JKO iterations remains challenging thanks to the minimization with respect toW2. A common approach to perform JKO steps is to discretize the spatial domain. For support size / 106, (6) can be solved by standard optimal transport algorithms [51]. In dimensions D ≥ 3, discrete supports can hardly approximate continuous distributions and hence the dynamics of gradient flows. To tackle this issue, [24] propose a stochastic parametric method to approximate the density of ρt. Their method uses entropy-regularized optimal transport (OT), which is biased.

heading: 3 Computing Wasserstein Gradient Flows with ICNNs
text: We now describe our approach to compute Wasserstein gradient flows via JKO stepping with ICNNs.

heading: 3.1 JKO Reformulation via Optimal Push-forwards Maps
text: Our key idea is to replace the optimization (6) over probability measures by an optimization over convex functions, an idea inspired by [11]. Thanks to Brenier’s theorem, for any ρ ∈ P2,ac there exists a unique ρ(k−1)-measurable gradient ∇ψ : RD → RD of a convex function ψ satisfying ρ = ∇ψ]ρ(k−1). We set ρ = ∇ψ]ρ(k−1) and rewrite (6) as an optimization over convex ψ:
ψ(k) ← arg min Convex ψ
[ F(∇ψ]ρ(k−1)) + 1
2h W22 (ρ(k−1),∇ψ]ρ(k−1))
] . (7)
To proceed to the next step of JKO scheme, we define ρ(k) def= ∇ψ(k)]ρ(k−1).
Since ρ is the pushforward of ρ(k−1) by the gradient of a convex function ∇ψ, theW22 term in (7) can be evaluated explicitly, simplifying the Wasserstein-2 distance term in (7):
ψ(k) ← arg min Convex ψ
[ F(∇ψ]ρ(k−1)) + 1
2h ∫ RD ‖x−∇ψ(x)‖22dρ(k−1)(x) ] . (8)
This formulation avoids the difficulty of computing Wasserstein-2 distances. An additional advantage is that we can sample from ρ(k). Since ρ(k) = [∇ψ(k) ◦ · · · ◦ ∇ψ(1)]]ρ0, one may sample x0 ∼ ρ(0), and then∇ψ(k) ◦ · · · ◦ ∇ψ(1)(x0) gives a sample from ρ(k). Moreover, if functions ψ(·) are strictly convex, then gradients ∇ψ(·) are invertible. In this case, the density dρ (k)
dx of ρ (k) = ∇ψ(k) ◦ · · · ◦
∇ψ(1)]ρ0 is computable by the change of variables formula (assuming ψ(·) are twice differentiable) dρ(k)
dx (xk) = [det∇2ψ(k)(xk−1)]−1 · · · [det∇2ψ(1)(x0)]−1 ·
dρ(0)
dx (x0), (9)
where xi = ∇ψ(i)(xi−1) for i = 1, . . . , k and dρ (0) dx is the density of ρ (0).

heading: 3.2 Stochastic Optimization for JKO via ICNNs
text: In general, the solution ψ(k) of (8) is intractable since it requires optimization over all convex functions. To tackle this issue, [11] discretizes the space of convex function. The approach also requires discretization of measures ρ(k) limiting this method to small dimensions.
We propose to parametrize the search space using input convex neural networks (ICNNs) [6] satisfying a universal approximation property among convex functions [20]. ICNNs are parametric models of the form ψθ : RD → R with ψθ convex w.r.t. the input. ICNNs are constructed from neural network layers, with restrictions on the weights and activation functions to preserve the input-convexity, see [6, M3.1] or [37, MB.2]. The parameters are optimized via deep learning optimization techniques such as SGD.
The JKO step then becomes finding the optimal parameters θ∗ for ψθ:
θ∗ ← arg min θ
[ F(∇ψθ]ρ(k−1)) + 1
2h ∫ RD ‖x−∇ψθ(x)‖22dρ(k−1)(x) ] . (10)
If the functional F can be estimated stochastically using random batches from ρ(k−1), then SGD can be used to optimize θ. FFP given by (5) is an example of such a functional: Theorem 1 (Estimator of FFP). Let ρ ∈ P2,ac(RD) and T : RD → RD be a diffeomorphism. For a random batch x1, . . . , xN ∼ ρ, the expression [ÛT (x1, . . . , xN )− β−1∆̂ET (x1, . . . , xN )], where
ÛT (x1, . . . , xN ) def =
1
N N∑ n=1 Φ ( T (xn) ) and
∆̂ET (x1, . . . , xN ) def =
1
N N∑ n=1 log |det∇T (xn)|,
is an estimator of FFP(T]ρ) up to constant (w.r.t. T ) shift given by β−1E(ρ).
Proof. ÛT is a straightforward unbiased estimator for U(T]ρ). Let p and pT be the densities of ρ and T]ρ. Since T is a diffeomorphism, we have pT (y) = p(x) · | det∇T (x)|−1 where x = T−1(y). Using the change of variables formula, we write
E(T]ρ) = − ∫ RD pT (y) log pT (y)dy
= − ∫ RD p(x) · | det∇T (x)|−1 log [ p(x) · | det∇T (x)|−1 ] · | det∇T (x)|dx
= − ∫ RD p(x) log p(x)dx+ ∫ RD p(x) log |det∇T (x)|dx
= E(ρ) + ∫ RD p(x) log |det∇T (x)|dx,
=⇒ ∆ET (ρ) def = E(T]ρ)− E(ρ) = ∫ RD log |det∇T (x)|dρ(x)
which explains that ∆̂ET is an unbiased estimator of ∆ET (ρ). As the result, ÛT − β−1∆̂ET is an estimator for FFP(T]ρ) = U(T]ρ)− β−1E(T]ρ) up to a shift of β−1E(ρ).
To apply Theorem 1 to our case, we take T ← ∇ψθ and ρ← ρ(k−1) to obtain a stochastic estimator for FFP(∇ψθ]ρ(k−1)) in (10). Here, β−1E(ρ(k−1)) is θ-independent and constant since ρ(k−1) is fixed, so the offset of the estimator plays no role in the optimization w.r.t. θ.
Algorithm 1 details our stochastic JKO method for FFP. The training is done solely based on random samples from the initial measure ρ0: its density is not needed.
Algorithm 1: Fokker-Planck JKO via ICNNs Input :Initial measure ρ0 accessible by samples;
JKO discretization step h > 0, number of JKO steps K > 0; target potential Φ(x), diffusion process temperature β−1; batch size N ;
Output : trained ICNN models {ψ(k)}Kk=1 representing JKO steps for k = 1, 2, . . . ,K do
ψθ ← basic ICNN model; for i = 1, 2, . . . do
Sample batch Z ∼ ρ0 of size N; X ←∇ψ(k−1) ◦ · · · ◦∇ψ(1)(Z); Ŵ22 ← 1N ∑ x∈X ‖∇ψθ(x)− x‖22;
Û ← 1N ∑ x∈X Φ ( ∇ψθ(x) ) ;
∆̂E ← 1N ∑ x∈X log det∇2ψθ(x); L̂ ← 12hŴ 2 2 + Û − β−1∆̂E ; Perform a gradient step over θ by using ∂L̂∂θ ;
ψ(k) ← ψθ
This algorithm assumes F is the Fokker-Planck diffusion energy functional. However, our method admits straightforward generalization to any F that can be stochastically estimated; studying such functionals is a promising avenue for future work.

heading: 3.3 Computing the Density of the Diffusion Process
text: Our algorithm provides a computable density for ρ(k). As discussed in M3.1, it is possible to sample from ρ(k) while simultaneously computing the density of the samples. However, this approach does not provide a direct way to evaluate dρ (k)
dx (xk) for arbitrary xk ∈ R D. We resolve this issue below.
If a convex function is strongly convex, then its gradient is bijective on RD. By the change of variables formula for xk ∈ RD, it holds dρ (k) dx (xk) = dρ(k−1) dx (xk−1) · [det∇ 2ψ(k)(xk−1)]
−1 where xk = ∇ψ(k)(xk−1). To compute xk−1, one needs to solve the convex optimization problem:
xk = ∇ψ(k)(xk−1) ⇐⇒ xk−1 = arg max x∈RD
[ 〈x, xk〉 − ψ(k)(x) ] . (11)
If we know the density of ρ0, to compute the density of ρ(k) at xk we solve k convex problems
xk−1 = arg max x∈RD
[ 〈x, xk〉 − ψ(k)(x) ] . . . x0 = arg max
x∈RD
[ 〈x, x1〉 − ψ(1)(x) ] to obtain xk−1, . . . , x0 and then evaluate the density as
dρk dx (xk) = dρ0 dx (x0) · [ k∏ i=1 det∇2ψ(i)(xi−1) ]−1 .
Note the steps above provide a general method for tracing back the position of a particle along the flow, and density computation is simply a byproduct.

heading: 4 Experiments
text: In this section, we evaluate our method on toy and real-world applications. Our code is written in PyTorch and is publicly available at
https://github.com/PetrMokrov/Large-Scale-Wasserstein-Gradient-Flows
The experiments are conducted on a GTX 1080Ti. In most cases, we performed several random restarts to obtain mean and variation of the considered metric. As the result, experiments require about 100-150 hours of computation. The details are given in Appendix A.
Neural network architectures. In all experiments, we use the DenseICNN [37, Appendix B.2] architecture for ψθ in Algorithm 1 with SoftPlus activations. The network ψθ is twice differentiable w.r.t. the input x and has bijective gradient ∇ψθ : RD → RD with positive semi-definite Hessian ∇2ψθ(x) 0 at each x. We use automatic differentiation to compute∇ψθ and ∇2ψθ. Metric. To qualitatively compare measures, we use the symmetric Kullback-Leibler divergence
SymKL(ρ1, ρ2) def = KL(ρ1‖ρ2) + KL(ρ2‖ρ1), (12) where KL(ρ1‖ρ2) def = ∫ RD log dρ1 dρ2
(x)dρ1(x) is the Kullback-Leibler divergence. For particle-based methods, we obtain an approximation of the distribution by kernel density estimation.

heading: 4.1 Convergence to Stationary Solution
text: Starting from an arbitrary initial measure ρ0, an advection-diffusion process (4) converges to the unique stationary solution ρ∗ [56] with density
dρ∗
dx (x) = Z−1 exp(−βΦ(x)), (13) where Z = ∫ RD exp(−βΦ(x))dx is the normalization constant. This property makes it possible to compute the symmetric KL between the distribution to which our method converges and the ground truth, provided Z is known.
We use N (0, 16ID) as the initial measure ρ0 and a random Gaussian mixture as the stationary measure ρ∗. In our method, we perform K = 40 JKO steps with step size h = 0.1. We compare with a particle simulation method (with 103, 104, 105 particles) based on the Euler-Maruyama bEMe approximation [36, M9.2]. We repeat the experiment 5 times and report the averaged results in Figure 1.
In Figure 2, we present qualitative results of our method converging to the ground truth in D = 13, 32.

heading: 4.2 Modeling Ornstein-Uhlenbeck Processes
text: Ornstein-Uhlenbeck processes are advection-diffusion processes (4) with Φ(x) = 12 (x− b) TA(x− b) for symmetric positive definite A ∈ RD×D and b ∈ RD. They are among the few examples where we know ρt for any t ∈ R+ in closed form, when the initial measure ρ0 is Gaussian [67]. This allows to quantitatively evaluate the computed dynamics of the process, not just the stationary measure.
We chooseA, b at random and set ρ0 to be the standard Gaussian measureN (0, ID). We approximate the dynamics of the process by our method with JKO step h = 0.05 and compute SymKL between the true ρt and the approximate one at time t = 0.5 and t = 0.9. We repeat the experiment 15 times in dimensions D = 1, 2 . . . , 12 and report the performance at in Figure 3. The baselines are bEMe with 103, 104, 5× 104 particles, EM particle simulation endowed with the Proximal Recursion operator bEM PRe with 104 particles [16], and the parametric dual inference method [24] for JKO steps bDual JKOe. The detailed comparison for times t = 0.1, 0.2, . . . 1 is given in Appendix C.

heading: 4.3 Unnormalized Posterior Sampling in Bayesian Logistic Regression
text: An important task in Bayesian machine learning to which our algorithm can be applied is sampling from an unnormalized posterior distribution. Given the model parameters x ∈ RD with the prior distribution p0(x) as well as the conditional density p(S|x) = ∏M m=1 p(sm|x) of the data S = {s1, . . . , sM}, the posterior distribution is given by
p(x|S) = p(S|x)p0(x) p(S) ∝ p(S|x)p0(x) = p0(x) · M∏ m=1 p(sm|x).
Computing the normalization constant p(S) is in general intractable, underscoring the need for estimation methods that sample from p(S|x) given the density only up to a normalizing constant.
In our context, sampling from p(x|S) can be solved similarly to the task in M4.1. From (13), it follows that the advection-diffusion process with temperature β > 0 and Φ(x) = − 1β log [ p0(x) · p(S|x) ] has dρ ∗
dx (x) = p(x|S) as the stationary distribution. Thus, we can use our method to approximate the diffusion process and obtain a sampler for p(x|S) as a result. The potential energy U(ρ) = ∫ RD Φ(x)dρ(x) can be estimated efficiently by using a trick similar to the ones in stochastic gradient Langevin dynamics [70], which consists in resampling samples in S uniformly. For evaluation, we con-
sider the Bayesian linear regression setup of [42]. We use the 8 datasets from [47]. The number of features ranges from 2 to 60 and the dataset size from 700 to 7400 data points. We also use the Covertype dataset2 with 500K data points and 54 features. The prior on regression weights w is given by p0(w|α) = N (w|0, α−1) with p0(α) = Gamma(α|1, 0.01), so the prior on parameters x = [w,α] of the model is given by p0(x) = p0(w,α) = p0(w|α) · p0(α). We randomly split each dataset into train Strain and test Stest ones with ratio 4:1 and apply the inference on the posterior p(x|Strain). In Table 1, we report accuracy and log-likelihood of the predictive distribution on Stest. As the baseline, we use particle-based Stein Variational Gradient Descent [42]. We use the author’s implementation with the default hyper-parameters.

heading: 4.4 Nonlinear Filtering
text: We demonstrate the application of our method to filtering a nonlinear diffusion. In this task, we consider a diffusion process Xt governed by the Fokker-Planck equation (4). At times t1 < t2 < · · · < tK we obtain noisy observations of the process Yk = Xtk + vk with vk ∼ N (0, σ). The goal is to compute the predictive distribution pt,X(x|Y1:K) for t ≥ tK given observations Y1:K . For each k and t ≥ tk predictive distribution pt,X(x|Y1:k) follows the diffusion process on time interval [tk, t] with initial distribution ptk,X(x|Y1:k). If tk = t then
ptk,X(x|Y1:k) ∝ p(Yk|Xtk = x) · ptk,X(x|Y1:k−1). (14)
For k = 1, . . . ,K, we sequentially obtain the predictive distribution ptk,X(x|Y1:k) by using the previous predictive distribution ptk−1,X(x|Y1:k−1). First, given access to ptk−1,X(x|Y1:k−1), we approximate the diffusion on interval [tk−1, tk] with initial distribution ptk−1,X(x|Y1:k−1) by our Algorithm 1 to get access to ptk,X(x|Y1:k−1). Next, we use (14) to get unnormalized density and Metropolis-Hastings algorithm [57] to sample from ptk,X(x|Y1:k). We give details in Appendix B. For evaluation, we consider the experimental setup of [24, M6.3]. We assume that the 1-dimensional diffusion process Xt has potential function Φ(x) = 1π sin(2πx) + 1 4x
2 which makes the process highly nonlinear. We simulate nonlinear filtering on the time interval tstart = 0 sec., tfin = 5 sec. and take the noise observations each 0.5 sec. The noise variance is σ2 = 1 and p(X0) = N (X0|0, 1). We predict the conditional density ptfinal,X(x|Y1:9) and compare the prediction with ground truth obtained with numerical integration method by Chang and Cooper [19], who use a fine discrete grid. As the baselines, we use bDual JKOe [24] as well as the Bayesian Bootstrap filter bBBFe [27], which combines particle simulation with bootstrap resampling at observation times.
We repeat the experiment 15 times. In Figure 4a, we report the SymKL between predicted density and true p(Xtfin |Y1:9). We visually compare the fitted and true conditional distributions in Figure 4b.

heading: 5 Discussion
text: Complexity of training and sampling. Let T be the number of operations required to evaluate ICNN ψθ(x), and assume that the evaluation of Φ(x) in the potential energy U takes O(1) time.
2https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/binary.html
Recall that computing the gradient is a small constant factor harder than computing the function itself [41]. Thus, evaluation of ∇ψθ(x) : RD → RD requires O(T ) operations and evaluating the Hessian ∇2ψθ(x) : RD → RD×D takes O(DT ) time. To compute log det∇2ψθ(x), we need O(D3) extra operations. Sampling from ρ(k−1) = ∇ψ(k−1) ◦ · · · ◦ ∇ψ(1)]ρ0 involves pushing x0 ∼ ρ0 forward by a sequence of ICNNs ψ(·) of length k − 1, requiring O ( (k − 1)T
) operations. The forward pass to evaluate the JKO step objective L̂ in Algorithm 1 requires
O(DT +D3) operations, as does the backward pass to compute the gradient ∂L̂∂θ w.r.t. θ.
The memory complexity is more difficult to characterize, since it depends on the autodiff implementation. It does not exceed the time complexity and is linear in the number of JKO steps k.
Wall-clock times. All particle-based methods considered in M4 and bDual JKOe require from several seconds to several minutes CPU computation time. Our method requires from several minutes to few hours on GPU, the time is explained by the necessity to train a new network at each step.
Advantages. Due to using continuous approximation, our method scales well to high dimensions, as we show in M4.1 and M4.2. After training, we can produce infinitely many samples xk ∼ ρ(k), together with their trajectories xk−1, xk−2, . . . , x0 along the gradient flow. Moreover, the densities of samples in the flow dρ (k)
dx (xk), dρ(k−1) dx (xk−1), . . . , dρ(0) dx (x0) can be evaluated immediately.
In contrast, particle-based and domain discretization methods do not scale well with the dimension (Figure 3) and provide no density. Interestingly, despite its parametric approximation, bDual JKOe performs comparably to particle simulation and worse than ours (see additionally [24, Figure 3]).
Limitations. To train k JKO steps, our method requires time proportional to k2 due to the increased complexity of sampling x ∼ ρ(k). This may be disadvantageous for training long diffusions. In addition, for very high dimensions D, exact evaluation of log det∇2ψθ(x) is time-consuming.
Future work. To reduce the computational complexity of sampling from ρ(k), at step k one may regress an invertible network H : RD → RD [9, 31] to satisfy H(x0) ≈ ∇ψ(k) ◦ · · · ◦ ∇ψ(1)(x0) and use H]ρ0 → ρ(k) to simplify sampling. An alternative is to use variational inference [12, 54, 71] to approximate ρ(k). To mitigate the computational complexity of computing log det∇ψθ(x), fast approximation can be used [66, 28]. More broadly, developing ICNNs with easily-computable exact Hessians is a critical avenue for further research as ICNNs continue to gain attention in machine learning [44, 37, 38, 30, 23, 5].
Potential impact. Diffusion processes appear in numerous scientific and industrial applications, including machine learning, finances, physics, and population dynamics. Our method will improve models in these areas, providing better scalability. Performance, however, might depend on the expressiveness of the ICNNs, pointing to theoretical convergence analysis as a key topic for future study to reinforce confidence in our model.
In summary, we develop an efficient method to model diffusion processes arising in many practical tasks. We apply our method to common Bayesian tasks such as unnormalized posterior sampling (M4.3) and nonlinear filtering (M4.4). Below we mention several other potential applications:
• Population dynamics. In this task, one needs to recover the potential energy Φ(x) included in the Fokker-Planck free energy functional FFP based on samples from the diffusion obtained at timesteps t1, . . . , tn, see [29]. This setting can be found in computational biology, see M6.3 of [29]. A recent paper [14] utilizes ICNN-powered JKO to model population dynamics.
• Reinforcement learning. Wasserstein gradient flows provide a theoretically-grounded way to optimize an agent policy in reinforcement learning, see [55, 72]. The idea of the method is to maximize the expected total reward (see (10) in [72]) using the gradient flow associated with the Fokker-Planck functional (see (12) in [72]). The authors of the original paper proposed discrete particle approximation method to solve the underlying JKO scheme. Substituting their approach with our ICNN-based JKO can potentially improve the results.
• Refining Generative Adversarial Networks. In the GAN setting, given trained generator G and discriminator D, one can improve the samples from G by D via considering a gradient flow w.r.t. entropy-regularized f -divergence between real and generated data distribution (see [7], in particular, formula (4) for reference). Using KL-divergence makes the gradient flow consistent with our method: the functional F defining the flow has only entropic and potential energy terms. The usage of our method instead of particle simulation may improve the generator model.
• Molecular Discovery. In [3], in parallel to our work the JKO-ICNN scheme is proposed. The authors consider the molecular discovery as an application. The task is to increase the drug-likeness of a given distribution ρ of molecules while staying close to the original distribution ρ0. The task reduces to optimizing the functional F(ρ) = Ex∼ρΦ(x)+D(ρ, ρ0) for a certain potential Φ (V - in the notation of [3]) and a discrepancy D. The authors applied the JKO-ICNN method to minimize F on MOSES [53] molecular dataset and obtained promising results.
ACKNOWLEDGEMENTS. Skoltech acknowledges the support of the Ministry of Science and Higher Education grant No. 075-10-2021-068. The MIT Geometric Data Processing group acknowledges the generous support of Army Research Office grants W911NF2010168 and W911NF2110293, of Air Force Office of Scientific Research award FA9550-19-1-031, of National Science Foundation grants IIS-1838071 and CHS-1955697, from the CSAIL Systems that Learn program, from the MIT–IBM Watson AI Laboratory, from the Toyota–CSAIL Joint Research Center, from a gift from Adobe Systems, from an MIT.nano Immersion Lab/NCSOFT Gaming Program seed grant, and from the Skoltech–MIT Next Generation Program.

Review summaries:
In this work the authors introduce an innovative discretisation of the Jordan Kinderleher Otto scheme (which describes the evolution of the gradient flow of relative entropy in the Wasserstein metric).  To do so the authors iteratively construct a transformation map as a composition of gradients of input-convex neural networks, which satisfy a universal approximation theory among convex functions.  Using Brenier's theorem, an iterative variational formulation for the density at discrete time-steps is obtained, from which one can generate samples and or evaluate an unnormalised density.  A variety of applications are presented, including finding the stationary distribution, as well as nonlinear filtering.   

The Fokker-Planck equation can be treated as gradient descent over entropy functionals in Wasserstein space. To solve this problem, Jordan Kinderlehrer and Otto proposed the JKO scheme. Later, through the connection between the Brenier theorem and this scheme, [10] treat it as an optimization problem which aims at finding the convex Brenier functional $\psi$ in equation (7). However, the solution given by [10] is usually intractable in high dimensional space. Thus, this paper proposes to use the input convex neural networks (ICNNs) proposed by [5] to replace $\psi$. And then use SGD to optimize the model parameter. This forms the core contribution of this work. Additionally, the experiments are interesting.

This paper is on the computation of Wasserstein gradient flow, a type of gradient flow based on the optimal transport theory. The proposed algorithm is built on the famous JKO scheme. To implement the JKO scheme, input-convex neural networks (ICNNs) are used to model the optimal transport map between the distributions at consecutive steps. 

This paper proposes an input-convex neural networks (ICNNs) based method to approximate the JKO scheme for discretizing Wasserstein gradient flows. Following a previous work [1], the key idea is to reformulate the distribution optimization problem in each JKO iteration as optimization over convex functions due to the well-known Brenner’s polar factorization theorem [2]. ICNNs are used to parameterize these convex functions from computational considerations. The proposed method is applied to simulate Wasserstein gradient flows for sampling from unnormalized density and nonlinear filtering. [1] Discretization of functionals involving the Monge-Ampere operator. [2] Polar factorization and monotone rearrangement of vector-valued functions.