Accelerated Information Gradient flow
Authors: Yifei Wang
Wuchen Li
Venue:  (2020)
Keywords: Optimal transport
Information geometry
Nesterov accelerated gradient method

Abstract: We present a systematic framework for the Nesterov's accelerated gradient flows in the spaces of probabilities embedded with information metrics. Here two metrics are considered, including both the Fisher-Rao metric and the Wasserstein-$2$ metric. For the Wasserstein-$2$ metric case, we prove the convergence properties of the accelerated gradient flows, and introduce their formulations in Gaussian families. Furthermore, we propose a practical discrete-time algorithm in particle implementations with an adaptive restart technique.  We formulate a novel bandwidth selection method, which learns the Wasserstein-$2$ gradient direction from Brownian-motion samples. Experimental results including Bayesian inference show the strength of the current method compared with the state-of-the-art.

Under review as a conference paper at ICLR 2020
ACCELERATED INFORMATION GRADIENT FLOW
Anonymous authors Paper under double-blind review
ABSTRACT
We present a systematic framework for the Nesterov's accelerated gradient flows in the spaces of probabilities embedded with information metrics. Here two met- rics are considered, including both the Fisher-Rao metric and the Wasserstein-2 metric. For the Wasserstein-2 metric case, we prove the convergence properties of the accelerated gradient flows, and introduce their formulations in Gaussian families. Furthermore, we propose practical discrete-time algorithm in parti- cle implementations with an adaptive restart technique. We formulate a novel bandwidth selection method, which learns the Wasserstein-2 gradient direction from Brownian-motion samples. Experimental results including Bayesian infer- ence show the strength of the current method compared with the state-of-the-ar..
INTRODUCTION
Recently, optimization problems the space of probability and probability models attract increas- ing attentions from machine learning communities. These problems include variational inference (Blei et al. (017)), Bayesian inference (Liu & Wang 2016), Generative Adversary Networks (GAN, Goodfellow et al.] (2014)), and policy optimizations (Zhang et al.] 2018), etc. For instance, vari- ational inference methods approximate a target density by minimizing the Kullback-Leiberr (KL) divergence as the loss (objective) function.
Gradient descent methods with sampling efficient properties play essential roles solve these opti- mization problems. Here the gradient descent direction often relies on the information metric over the probability space. This direction naturally reflects the change of the loss function with respect to the metric. In literature, two important information metrics, such as the Fisher-Rao metric and the Wasserstein-2 (in short, Wasserstein) metric, are of great interests (Amari) 1998 Otto 2001 Laf- ferty] 1988) For the Fisher-Rao gradient, classical results including Adam (Kingma & Bay 2014)
and K-FAC Martens & Grosse 2015) demonstrate its effectiveness probability models. For the Wasserstein gradient, many classical methods such as Markov chain Monte Carlo (MCMC) meth-
ods (Geman & Geman 1987 Neal et al., 2011 Welling & Teh 2011) and particle- variational
inference VI) methods (Liu & Wang 2016 Chen & Zhang
2017: Chen et al.) 2018) are based The strength of using the Wasserstein gradient is also
on this framework in the probability space.
shown in probability models such GANs. (Arjovsky et al.
2017: Lin et al. 2018 Li et al. 2019).
The Nesterov's accelerated method (Nesterov) 1983) is widely applied in accelerating the vanilla gradient descent under the Euclidean metric. corresponds to a damped Hamiltonian flow, known as the accelerated gradient flow (Su et al., (0166). A natural question is whether there exists counterpart of the accelerated gradient flow in the probability space under information metrics. For optimization problems on Riemannian manifold, the accelerated gradient methods are studied by Liu et al. (2017): 7): Zhang & Sra (2018). The probability space embedded with information metric can be viewed as a Riemannian manifold. Several previous works explore accelerated methods in this manifold under the Wasserstein metric. [Lu al. (2018) 2019) propose an acceleration framework of ParVI methods based on manifold optimization. Taghvaci & Mehta (2019) introduce the accelerated flow from an optimal control perspective. On the other hand, Cheng et al. (017); y): Ma et al. (2019) explore and analyze the acceleration MCMC, based the underdamped Langevin dynamics.
In this paper, we present a unified framework of accelerated gradient flows in the probability space embedded with information metrics, named Accelerated Information Gradient (AIG) flows. From an cooomation-differenen gefferentiallgeomettyy perspective, we derive AIG flows by damping Hamiltonian flows, concerning both the Fisher-Rao metric and the Wasserstein metric. Then focus the Wasser-

---

Under review as a conference paper at ICLR 2020
stein metric with the KL divergence loss function. In Gaussian families, verify the existence of the solution to AIG flows. Here we show that the AIG flow corresponds to a well-posed ODE system in the space of symmetric positive definite matrices. We rigorously prove the convergence rate of AIG flows based on the geodesic convexity of the loss function. Here we note that our proof removes the unnecessary technical assumption in (Taghvaei & Mehta 2019
Theorem 1).
Besides, we handle two difficulties in numerical implementations of AIG flows. On the one hand, as pointed out by (Tghvaei & Mehta (2019); [iu et (2019), the logarithm of density term (Wasser- stein gradient of KL divergence) is hard to approximate in partiele formulations. We propose a novel kernel selection method, whose bandwidth learned by sampling from Brownian motions. We call it the BM method. On the other hand, we notice that the AIG flow can be numerically stiff system, especially high-dimensional sample spaces. This is because the solution of AIG flows can be close to the boundary of the probability space. To handle this issue, we propose an adaptive restart technique, which accelerates and stabilizes the AIG algorithm. Numerical results in toy examples, Gaussian measures and Bayesian Logistic regression indicate the validity of the BM method and the acceleration effects the proposed AIG flow.
This paper is organized as follows. Section 2 brieffl reviews the information metrics and their corre- sponding gradient flows and Hamiltonian flows in the probability space. Section we formulate various forms of AIG flows and analyze W-AIG flows in Gaussian measures. We theoretically prove the convergence rate of W-AIG flows in Section 4 Section presents the discrete-time algorithm for W-AIG flows, including the BM method and the adaptive restart technique. Section provides numerical experiments.
2 METRIC AND FLOWS IN THE PROBABILITY SPACE
Suppose that Ω is a region in R". Let F(22 denote the set of smooth functions on Ω. (s.)) and || the Euclidean inner product and norm in R". V, V. and A represent the gradient, divergence and Laplacian operators in R". Denote the set of probability density
P(2)= { {
[ pdl =, p 02
The tangent space at p P(2) follows (2) {ar E F(Q): f adx 0.}. The cotangent space at p, P(R), can be treated the quotient space F(1)//R, which are functions in F(22) defined up to addition constants.
Definition (Metric in the probability space) metric tensor G(p) TTP((2) + P(2) is an invertible mapping from the tangent space at p to the cotangent space at p. This metric tensor defines the metric (inner product) on the tangent space ((I). Namely, for 01, 02 E (S2(),), we define
(0(01,0)==000 01
o1G(p)o2dz
/ Gpp(pp 12dxx
where 4, is the solution to a; - 10 1,2.
We present two important examples of metrics on the probability space P(2)): the Fisher-Rao metric from information geometry and the Wasserstein metric from optimal transport. Example 1 (Fisher-Ras Rao metric) The inverse of the Fisher-Rao metric tensor is defined by
G (p) 1p
2
f @ppdt))
+ 4 TTP(().
The Fisher-Rao metric on the tangent space is given by
of (0),000)) 02) = P2pdr -( (f ( BBspdde): 01,02 01,02E TTP(())
where &; is the solution a; p - (ipdz)) 1,
Example 2 (Wasserstein metrie) The inverse of the Wasserstein metric tensor defined by
G -p - -VVV), o E PT)().
2

---

Under review as a conference paper at ICLR 2020
The Wasserstein metric on the tangent space is given by
W Op (01, 02) / (V V dx, 01 E TTP()))
where P, is the solution to 5j = - (pV ==1,2
2.1 GRADIENT FLOWS
In learning, many problems can be formulated as the optimization problem in the probability space,
min
E(p).
(EP
Here E(p is divergence or metric loss functional between p and target density pt E P(2).. One typical example of Ep) is the KL divergence from p to p",
E(p) DxL (plp"")
log
pda.
Another example the maximum mean discrepancy (MMD, Gretton et al. (012))
(x(y) (() MMD(p, p²) p²)= - (a)
p" (v))
where K(x, y) is a given kernel function. The gradient flow for E(p) in (P(1), go) takes the form
S.pr = -G(:)) (b(p)) -16(()))))
8p:
Here 8E(ps) is the 2 first variation w.r.t. Pt- We formulate the gradient flow under either the 8p,
Fisher-Rao metric or the Wasserstein metric.
Example 3 (Fisher-Rao gradient flow) The Fisher-Rao gradient flow given by
8
SE Ptdx
dep: = =
8p:
öp.
Example 4 (Wasserstein gradient flow) The Wasserstein gradient flow writes
depe= = (p
8E
0pr
2.2 HAMILTONIAN FLOW
In this subsection, briefly review the Hamiltonian flow in the probability space. By using the metric Io in the probability space, we can define a Lagrangian by
C(pr, d,pt) = 90. (Deps, aspt) E(p.).
The LalernLagrange equation for the Lagrangian follows
(1 az
5L
8C
+ C(t),
(I)
8p:
where C(t) is a spatially-constant function.
Proposition 1 we let o. =L//(((pp)) G(p.)d.pr. the equation system of (pt. 4t), i.e.
be formulated
p.p.
1 18
(f0,CLm)
8E
(2)
0. +
d, -
+
8pt = 0.
2ppt
where o is up to spatially-constant function shrift. Here ② is the Hamiltonian flow
0z
Pt
-1
0
1
H(Pr, T.) 0, bgc
(3)
+
0
H(pp. Tb:)
54,
with respect to the Hamiltonian H(pr, Φ.) 1/2 (.G((m) -b + E(pe).
3

---

Under review as a conference paper at ICLR 2020
Similarly, write the Hamiltonian flow under the Fisher-Rao metric the Wasserstein metric. Example 5 (Fisher-Rao Hamiltonian flow) The Fisher-Rao Hamiltonian flow follows
Oept - Deps + Dipidx= = 0:
1 4
(ʃ (((erdz)) Dt
5E
0
0.
2
8pt
The corresponding Hamiltonian is H H D1) = ( - (f E(ps).
Example 6 (Wasserstein Hamiltonian flow) The Wasserstein Hamiltonian flow writes
Oepe V (pe v
=0, 0,
6E
0,0; + 1 1
= 0. 5pt
The corresponding Hamiltonian is H. (pt, 2x)
/ p + E(rr). This identical to
the Wasserstein Hamiltonian flow introduced by Chow et al. (2019).
3 ACCELERATED INFORMATION GRADIENT FLOW
Let 0₁ 2 0 be a scalar function of t. We add a damping term ag 2 to the Hamiltonian flow 3
0z
Pt Dt
+
0
0
J
1
bp:
H(pr, 4t) H(r, Dt)
0.
(4)
-1 0
This renders the Accelerated Information Gradient (AIG) flow
(ppt - G(p:)-((pp)) -.Pt)
=0,
1 1
( (fo,C(a)-BBdd
5F
(AIG)
0. a. +
"B(
2 8pc
8pt
with initial values Po n0 and to = 0. The choice of depends on the geodesic convexity of Ep)) which has an equivalent definition as follows.
Definition 2 For functional E(p) defined the probability space, say that E(p) has 8- positive Hessian (in short, Hcss(5)) w.r.t. the metric 9, if there exists a constant 2 such that for any p P(2)) and any T,2P( have
gp (Hess E(p)o, σ) 2 Sgo(o,
Here Hess is the Hessian operator w.r.t. 9p.
E() is Hess(B) for 8 > 0, then a = 2/3; if E(p) is Hess(O), then a, = 3/tt.
Remark The Nesterov's accelerated method (Nesterov) 1983) a first-order method to optimize f(z) in the Euclidean space. The corresponding accelerated gradient flow by (Su et al. (2016) is equivalent damped Hamilronian system
0
0 - -1
/
V H. (r. 1
1 = 0, H (). = llel1 + f(x).
+
P
0.P
0
with initial values x00) Zo and p(0) 0. The choice of a, depends on the property of f(x). f(x is B-strongly convex, then a₁ = 2/3: f(x) is convex, then a₂ = 3/t. We apply this Hamiltonian flow interpretation to construct (AIG) in the probability space with information metries.
We give examples of AIG flows under either the Fisher- Rao metric or the Wasserstein metric. Example (Fisher-Rao AIG flow) The Fisher-Ran AIG flow writes
(foipndte) Pr. = 00
O.p: - Dept +
(F-AIG)
5E v, -
- - (foevrda))
0 + + + 0
0.
8pt
4

---

Under review as a conference paper at ICLR 2020
Example 8 (Wasserstein AIG flow) The Wasserstein AIG flow writes
Oept + V (pr v = 0,
σE
(W-AIG)
0,0; + R. ||
0.
5pc
For the rest of this paper, we mainly focus on the Wasserstein metric. Here the AIG flow (Eulerian formulation in fluid dynamics) has counterpart in the particle level (Lagrangian formulation).
Proposition 2 Suppose that X, ~ Pt and V₁ = (X((X() are the position and the velocity of particle at time t. Then, the differential equation of the particle system corresponding (W-AIG) writes
dx, = Vidt, dV. = -o.Vdd V
(Xe)dl.
(W-AIG-P)
8p:
If Ep) evaluates the KL divergence and pt x exp(-f(z)) (W-AIG-P) is equivalent to dX; Vidt,
dV; = a4 Vidt (X(e) log pe (Xt)dt.
(W-AIG-P-KL)
Remark The Vlog (X; (Xt)dt term cannot simply replaced by Brownian motion dB, because P: is the marginal distribution on X,. Several previous works have studied the accelerated gradient flow of KL divergence the probability space under the Wasserstein metric. [Tahhaai & Mehta 2019) construct the accelerated gradient flow in the probability space based on Wibisono et al. 2016 ' variational formulation on the Nesterov's 's accelerated method. Their flows coincide with (- with a2 3/t after rescaling. The underdamped Langevin dynamics in (Cheng et al.] 2017; 7 Ma et al. (209)) damps the Hamiltonian flow of the particles, which is different from P- shown in (Taghvaei & Mehta 2019). [iu et al. (2018) 2019) give the discrete-time accelerated algorithm from the perspective of manifold optimization.
3.1 WASSERSTRIN METRIC RESTRICTED TO GAUSSIAN
In this subsection, we demonstrate that (W-AIG) has an ODE formulation in Gaussian family. De- note A to the multivariate Gaussian densities with zero means. Namely, if p', ,p E Noo then we show that (W-AIG) has a solution (p:, .t) and Pt E AVo
Let P and S" represent symmetric positive definite matrix and symmetric matrix with size x respectively. Each E A1 can be uniquely expressed by its covariance matrix Σ E pa by p(x; E) = (2mr n/2
/det(E) exp ---====''>) The Wasserstein metric on P(R") induces a metric on N,, which is
2018 So there exists a Wasserstein metric on P", also known as the Bures metric. For Σ P the tangent space and cotangent space follow TsP²" = TTTP" S
geadee- submanifold in P(R"), see (Takatsu 2008 Modin 2016 Malagò et al.
Definition (Wasserstein metric in Gaussian) For Σ IP, the metric tensor G(E) : S" + S" is defined by G(E) 's 2 + SE). The Wasserstein metric on S" is gr (A, A2) tr(A;G(E) A2) =
where S, E S" is the solution to A, = ES; + S,E, i= 1,2.
Proposition 3 Given an energy function E(E), the Wasserstein gradient flow Gaussian writes
Σ =
+ E(E,)
Here Vs, is the standard matrix derivative. The Hamiltonian flow is system of (Er, i.e..
-E,, + =0
(5)
5 E(E,) 0.
The corresponding Hamiltonian writes H(2,, St) = tr(S,2,,,)) 22(2...
Therefore, by adding the damping term a,S,, we obtain the Wasserstein AIG flow in Gaussian.
S2 +.,.,+,,,,. S,S,) 0,
(W-AIG-G)
Š S? =VVz, = 0,
5

---

Under review as a conference paper at ICLR 2020
with initial values En Σ and So = 0. For now, consider E(E) to be the KL divergence.
E(E) - E((( Σ)) |t<(E(E))=)) log dot det(E(E")))) n²]
(6)
where Ex is the covariance matrix of p². The following theorem proves the well sell-posedness of W-AIG-G and illustrates the connection between W-AIG flows in IP" and P(R)).
E and E". E(E) defined in (6) evaluates the KL divergence from p to g². Let (E,, St) be the solution to (W-AIG-) with
Theorem Suppose that p² No and their covariance matrices
initial values Eo = Σd and So O. Then, for any t z 0, E, is well-defined and stays positive definite. Furthermore, we denote
Pc(x)= (2x)-n/2 det((i)) exp 1-2""
1 1
(,(x) x="
Sex +
where C(t) -t + L0 log det (Et( Then, (pt. 1.) is the solution initial values Po pd and 4 0.
W-AIG with
4 CONVERGENCE RATE ANALYSIS ON W-AIG FLOWS
In this section, we prove the convergence rate of (W-AIG).
Theorem 2 Suppose that E(p) satisfies Hess((3 for B > 0. The solution pz. to (WAAGG with
a₂ = 2// satisfies
E(p.) S Cue VBt VBt
If (p) only satisfies Hess(O), then the solution P: to (W-AIG with = 3/t satisfies
E(pt) S C/t O)
Here the constants Co, C. only depend on Pu-
Remark 3 Here Hess(3) is equivalent to the 8-geodesic convexity in the probability space w.r.t. go. For the Wasserstein metric, it is also known as 3-displacement convexity; see (Villani 2008 Chap 16). Consider the case where E(p) is the KL divergence and the target density takes the form 0² < exp(- (z)).. A sufficient condition for Hess(3) is that (xx is /--ttrongly convex, see (Otto)
& Villani 2000 Bakry & Emery 1985). If E(p) satisfies Hess(B) for B3 > 0, then the classical analysis indicates that the solution the Wasserstein gradient flow has an Ole 23t) convergence
rate. The W-AIG flow improves the convergence rate to ( especially when close to 0.
Here we provide a sketch in the proof of Theorem Given Pt. we can find the optimal transport plan T. from Pt to p". Let T#p denote the posh-forward density from p by the mapping T. The following proposition characterizes the inverse of the exponential map in probability space with the Wasserstein metric.
Proposition 4 Denote the geodesic curve 7(8) that connects pt and p² by 7(8) = (sT, + (1 s) Id) E [0, [0. 1]. Here ld is the identity mapping from 12 to itself. Then. 7()) corresponds to a tangent vector - V (pT())(Ti(z)) z)) E Tp, P(2).
We first consider the case where E(p) satisfies Hess(B) for 3 > 0. Motivated by the Lyapunov function for Nesterov's ODE the Euclidean case, we construct the following Lyapunov function.
E(t) = cVit (2 -))
v=)(2) pı(z)dz + E(p.) - El0').
(7)
Proposition 5 Suppose that E(p) satisfies Hess(3) ffrss( B > 0. Pt is the solution to a, 2/B.. Then, E(t) defined in 7 satisfies E() < 0. As result,
W-AIG with
F(pt) S - < < (() VBL)

---

Under review as a conference paper at ICLR 2020
Note that (()) only depends Po- This proves the first part of Theorem 2 case where E(p) satisfies Hess(O), We construct the following Lyapunov function.
We now consider the
1 E()) - | x))
L Vo pr(x)de + E
+ (Blp))
(8)
Proposition 6 Suppose that Ep) satisfies Hess(O). Pt. is the solution Then, E(t) defined in (8) satisfies E(t) S 0. As a result,
W-AIG with 02 3/L.
E(pt) S {t) v 80) = O(
Because {()) only depends on Po. we complete the proof.
Remark 4 For the Hess(0) case, we obtain the same result in Taghvaci & Mehta 2019 Theorem 1). Their proof comes from the Lagrangian formulation (W-ATG-P) and our the Eulerian formulation (W-ATG). However, their technical assumption (X2 + e - - (X,)) dt (X,)] is only valid in 1-dimensional case. In Appendix C.4 we prove that this quantity is non-negative. This is due to the Hodge decomposition behind
proof is based
the optimal transport, see Lemma in Appendix
5 DISCRETE-TIME ALGORITHM FOR WIAIG FLOWS
In this section, we present the discrete-time time implementation of (W-AIG-P-KL) This implementation simpler and more stable than the in (Tghvaei & Mehta 2019). Suppose that initial positions of a particle system {X) N :=1 are given and V/ = 0. The time parameter is related to the step size
V via t Vik. The update rule follows
V V V=(V)(X)) (x()),,
X X VTV 1
for = 1,2. N. If E(p) is Hess(3), then 1-V3r. if E(p) is Hess(O) or 3 is unknown, then 1-V3;;
k-1 Here & (x) is an approximation of V (z)). K+2
0%
We review two common choices of & follows. If g follows Gaussian distribution, then (x) -x (z mg),
(10)
where my and Dg are the mean and the covariance matrix of (X For the non-Gaussian case, we use the kernel density estimation (KDE, Singh (1977). (x(x)) approximate p=(x.) Here K K(x, y) is kernel function. Then, &A writes
3/4 K
K(x, X to
ΣN ((,(x,X)) ΣN K(x,
(x) V (zo
(11)
A common choice of K(x, y) is Gaussian kernel with the bandwidth h (" h) (2mh)-n/22cc exp (||= / (2h)). There two difficulties in the discretization. For one thing, the bandwidth strongly affects the estimation of V log Pt. so we propose the BM method to adaptively learn the bandwidth from samples. For another, the second equation in (W-AIG) is the Hamilton- Jacobi equation, which usually has strong stiffness. We propose an adaptive restart technique to deal with this problem.
Remark 5 Our numerical implementations of W-AIG flows can be viewed special case of ParVI methods. Compared to traditional MCMC methods, ParVI methods are more sample-efficient because make full use of finite number of particles by taking particle interaction into account.
5.1 LEARN THE BANDWIDTH VIA BROWNIAN MOTION
SVGD uses a median (MED) method to choose the bandwidth, i.e.,
hy = 2 2ogg(N + median (1III) - x2113) Ja,==1
(12)
7

---

Under review as a conference paper at ICLR 2020
Liu et al. (2018) propose Heat Equation (HE) method to adaptively adjust bandwidth. Motivated by the HE method, we introduce the Brownian motion (BM) method to adaptively learn the bandwidth.
Given the bandwidth h. {X*) and step size
compute two particle systems
(h) X - (<x(2;; (2; h), Z X + V2sB', i= i..,...N,
where B is the standard Brownian motion. We want to minimize MMD(jy, the MMD be- tween the empirical distribution (x) >i= - and (z(x) > 8(x ((= Here, the kernel (x. y) for the MMD is the Gaussian kernel with a bandwidth of 1. So we optimize over h to minimize MMD (py, (22), using the bandwidth h from the last iteration the initialization. For simplicity we denote (M ( as the output of the BM method.
Remark 6 Besides KDE, there are other methods that approximate the V log (x(x) (compute {x) via a kernel function, such as the blob method (Carrillo et al. 2019) and the diffusion map (Taghvaei & Mehta 2019). The BM method can also select the kernel bandwidth for these methods.
5.2
ADAPTIVE RESTART
To enhance the practical performance, we introduce an adaptive restart technique, which shares the same idea of gradient restart in (Odonoghue & Candes 2015 Wang et al. 2019) under the Euclidean case. Consider
Pk Σ V (x
(13)
which can be viewed as discrete- time approximation of - (d.pt: GW (p:)- 14E) = (p.)
If 1P& O, then restart the algorithm with initial values X X and Vi 0. This essentially keeps 0, (p:)) negative along the trajectory. The numerical results show that the adaptive restart
accelerates and stabilizes the discrete-time algorithm. The overall algorithm is summarized below.
Algorithm Particle implementation of Wasserstein AIG flow
Require: initial positions [Xi] i-= step size т, number of iteration L. Set 1: 0, V/ 0, 1....A N. Set the bandwidth he by MED T2)
2: for / 1, ... do
3: Compute by based on BM: 11 = (M (X
R (X)
4: Calculate & (Xt) by (10) or by (T with the bandwidth hr.
5: Set Or based on whether (p)) is Hess(B) or Hess(O). For = 1, N, update
V = aa V
VT(Vf((X) + (xg)),
X = X + VTVk+t
6: RESTART then
7:
Compute Pk = = - (Vê VS(X)) + (x())
8: If Pk < 0, set X; = X and Vi = and = O; otherwise set = k +
9: else
10: Set k = k +11. 1.
end if'
12:
end for
6 NUMERICAL EXPERIMENTS
In this section, we present several numerical experiments to demonstrate the validity of the BM method, the acceleration effeet of the Wasserstein AIG flow, and the strength of the adaptive restart technique. Implementation details can be found in Appendix
6.1 TOY EXAMPLE
We first investigate the validity of the BM method in selecting the bandwidth. The target density p" is toy bimodal distribution (Rezende & Mohamed 2015). We compare two types of particle implementations of the Wasserstein gradient flow over the KL divergence:
X X -VS(X)) Bt.B X X +(VS(X)) +x(()))
8

---

Under review as a conference paper at ICLR 2020
Here B ~ -O(0, 1) is the standard Brownian motion and takes the form (TI. The first method known as the MCMC method and the second method is called the ParVI method. For the second method, the bandwidth h is selected by MED/HE/BM respectively. Figure shows the distribution of 200 samples based on different methods. Samples from MCMC match the target distribution in stochastic way; samples from MED collapse: samples from HE align tidily around the contour lines; samples from BM arrange neatly and are closer to samples from MCMC. This indicates that the BM method makes the particle system behave similar MCMC, though in a deterministic way.
MCMC
MED
HE
BM
Figure I:
The effect of the BM method. Samples are ploutd as blue dots. Left to right: MCMC, MED, HE and BM. All methods run for 200 iterations with the same initialization.
6.2 GAUSSIAN MEASURES
Next, explore the effectiveness of (W-AIG) flow compared to the Wasserstein gradient flow and demonstrate the strength of the adaptive restart. The target density p² is a Gaussian distribution with zero
R100 the covariance matrix of p" is Σ" and W = - Let L and 3 be the largest/smallest eigenvalue of W*. E(p) satisfies Hess(3) and the step size is T 1/(4L). The condition number of W' is defined as k L/3. The large L indicates Σ is close be singular.
We first demonstrate the effectiveness of (W-AIG-G) in the ODE level. Detailed discretization is For now, WGF denotes the discretization of the Wasserstein gradient flow; (IG-(r)(s) denotes the discretization of the Wasserstein AIG flow. For letters in the parentheses, 'r' denotes using the adaptive restart and 's' denotes utilizing B3. Figure 2 presents the convergence of the KL. divergence on two target distributions with small/large L. We observe that AIG converges faster compared to WGF, which verifies Theorem The adaptive restart also accelerates the algorithm.
left in Appendix D.2] The initial value is set be Eo
10P
10"4
-Q-BROR
10 -Q-1K0F
10
AAC
AGC-
AD-
104
10
2007
2000
022
700
1000
2007
400
703
800
900
Berittio
Paration
Figure 2: The acceleration effect of W-AIG flow and the strength of adaptive restart (ODE level). The target density is Gaussian distribution with zero mean on R100 Left: L = 1, = 3.8 x 103 Right:/ 1,K 4.0 x 103.
Then, we demonstrate the results the particle level. The setting of pa same as the previous experiment. The initial distribution of samples follows (0, I) and the number of samples is N 600. For particle system { covariance matrix Σ. The left part Figure 3 (small L) is almost identical to Figure 2 which verifies the acceleration effect of AIG flows. It also indicates that the adaptive restart helps to accelerate the convergence. From the right part of Figure 3 (large L). AIG and AIG-s diverge because of the ill target distribution, and the adaptive restart solves this problem.
record the KL divergence E(xx) (6) using the empirical

---

Under review as a conference paper at ICLR 2020
11P
10"
10
-0-RK06
101
-Q-IKIs
10"
A-2-
AD2
A20-12
104
10'13
200
300
800
000
700
1000
500
400
700
800
900
Baration
Baration
Figure 3: The acceleration effect of W-AIG flow and the strength of adaptive restart (particle level). The setting of target densities is identical to the ones in Figure
6.3 BAYESIAN LOGISTIC REGRESSION
We perform the standard Bayesian logistic regression experiment on the Covertype dataset, follow- ing the same settings as Liu & Wang (2016). We compare our methods with MCMC, SVGD (Lu Wang 2016), WNAG (Liu et al. 2018) and WNes (Liu et al.) 2019). We select the bandwidth using either the MED method or the proposed BM method. Figure S iindicates that the BM method accelerates and stabilizes the performance WGF and AIG. The performance of MCMC and WGF are similar and they achieve the best log-likelihood. In test accuracy, AIG-r converges faster than other methods and is more stable. The adaptive restart improves the overall performance AIG.
0.78
0.75
+0.52
20.72
40.50
0.00
WE10
0.58
MONS WNAD
WM40
0.00
WCT
-0.00
WCF
AGA
Neratio
0.77
0.75
40.50
0..72
+0
0.68
MEVC
0.58
MENC
WM46
WN.G
0.00
00.00
WGF
AMr
Figure 4: Comparison of different methods on Bayesian logistic regression, averaged over 10 inde- pendent trials. The shaded areas show the variance 10 trials. Top: BM; Bottom: MED. Left: Test accuracy; Right: Test log-likelihood.
CONCLUSION
In summary, we propose the framework of AIG flows by damping Hamiltonian flows with respect to certain information metrics in the probability space. AIG flows have been carefully studied in Gaussian families. Theoretically, we establish the convergence rate of W-AIG flows. Numerically, we propose the discrete-time algorithm and the adaptive restart technique to overcome the numerical stiffness of W-AIG flows. We introduce novel kernel selection method by learning from Brownian- motion samples. The numerical experiments verify the acceleration effect of AIG flows and the strength of the adaptive restart. In future works, we intend to systematically explain the stiffness of AIG flows and the effect of the adaptive restart. We shall apply results general information metrics, especially for the Fisher-Ran Rao metric. We expect to study the related sampling efficient optimization methods and discrete-time algorithms general probability models.
10

---

Under review as a conference paper at ICLR 2020
REFERENCES
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251- 276, 1998.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pp. 211-2112222 2017. 2017
Dominique Bakry and Michel Emery. Diffusions hypercontractives. In Séminaire de Probabilités XIX 1983/84, pp. 177-206. Springer, 1985.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisti- cians. Journal of the American Statistical Association, 11:(518):55--877 2017.
José Antonio Carrillo, Katy Craig, and Francesco Patacchini. A blob method for diffusion. Cal- culus of Variations and Partial Differential Equations, 58(2):53, 2019.
Changyou Chen and Ruiyi Zhang. Particle optimization in stochastic gradient meme. arXiv preprint 1rXiv: 10.1:7711.1900 2017.
Changyou Chen, Ruiyi Zhang, Wenlin Wang, Bai Li, and Liqun Chen. A unified particle- optimization framework for scalable bayesian sampling. arXiv preprint arXiv: /80tv /8800//6599 1659, 2018.
Xiang Cheng, Niladri Chatterji, Peter L Bartleu, and Michael Jordan. Underdamped langevin meme: A non-asymptotic analysis. arXiv preprint arXiv: 2017.
Shui-Nee Chow, Wuchen Li, and Haomin Zhou. Wasserstein hamiltonian flows. arXiv preprint /903.00038008 2019.
Stuart Geman and Donald Geman. Stochastic relaxation, gibbs distributions, and the bayesian restoration of images. In Readings in computer vision, pp. 564-584. Elsevier, 1987.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in neural infor- mation processing systems, pp. 2672-2680, 2014.
Arthur Gretton, Karsten M Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. Journal of Machine Learning Research, 13(Mar):723-773, 2012.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.69880 6980, 2014.
John Lafferty. The density manifold and configuration space quantization. Transactions of the American Mathematical Society, 305(():699 1988.
Wuchen Li, Alex Tong Lin, and Guido Montúfar. Affine natural proximal learning. Geometric science of information, 2019, 2019.
Alex Tong Lin, Wuchen Li, Stanley Osher, and Guido Montúfar, Wasserstein proximal of gans. CAM report 18-53, 2018.
Chang Liu, Jingwei Zhuo, Pengyu Cheng, Ruiyi Zhang, Jun Zhu, and Lawrence Carin. Ac- celerated first-order methods on the wasserstein space for bayesian inference. arXiv preprint 1007v:1807.01777 2018
Chang Liu, Jingwei Zhuo, Pengyu Cheng, Ruiyi Zhang, and Jun Zhu. Understanding and acceler- ating particle-based variational inference, In International Conference on Machine Learning, pp- 4082-4094 2019.
Qiang Liu and Dilin Wang. Stein variational gradient descent: general purpose bayesian inference algorithm. In Advances in neural information processing systems, pp. 2378-2386. 2016.
Yuanyuan Liu, Fanhua Shang, James Cheng, Hong Cheng, and Licheng Jiao. Accelerated first-order methods for geodesically convex optimization on riemannian manifolds. Advances in Neural Information Processing Systems, pp. 4868 4877, 2017.

---

Under review as a conference paper at ICLR 2020
Yi-An Ma, Niladri Chatterji, Xiang Cheng, Nicolas Flammarion, Peter Bartlett, and Michael I Jor- dan. Is there an analog of nesterov acceleration for meme? arXiv preprint arXiv:1922.00990 2019.
Luigi Malagö, Luigi Montrucchio, and Giovanni Pistone. Wasserstein riemannian geometry of positive definite matrices. arXiv preprint arXiv:/801.0920 2018.
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factoree approximate curvature. In International conference on machine learning, pp. 2408-2417, 2015.
Klas Modin. Geometry of matrix decompositions seen through optimal transport and information geometry. arXiv preprint arXiv:1601.01873 2016.
Radford M Neal al. Meme using hamiltonian dynamics. Handbook of markov chain monte carlo, 2(11):2, 2011.
Yurii Nesterov. A method of solving convex programming problem with convergence rate (11kk)). Soviet Mathematics Doklady, 372))332--76, 1983.
Felix Otto. The geometry of dissipative evolution equations: the porous medium equation. Commu- nications Partial Differential Equations, 26(1-2):101-174, 2001.
Felix Otto and Cédric Villani. Generalization of an inequality by talagrand and links with the loga- rithmic sobolev inequality. Journal of Functional Analysis, 173(2):311-400, 2000.
Brendan Odonoghue and Emmanuel Candes. Adaptive restart for accelerated gradient schemes. Foundations of computational mathematics, 15(3):715-73,, 2015.
Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows. arXiv preprint arXiv: 05770, 2015.
Radhey S Singh. Improvement on some known nonparametrie uniformly consistent estimators of derivatives of density. The Annals of Statistics, pp. 394-399 1977.
Weijie Su, Stephen Boyd., and Emmanuel J. Candés. A differential equation for modeling Nesterov's accelerated gradient method: Theory and insights. Journal of Machine Learning Research, 2016.
Amirhossein Taghvaei and Prashant G Mehta. Accelerated flow for probability distributions. arXiv preprint arXiv:1901.000 00317, 2019.
Asuka Takatsu. On wasserstein geometry of the space of gaussian measures. arXiv preprint #rXivv0801.22500 2008
Cédric Villani. Topies in optimal transportation. American Mathematical Soc., 2003.
Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
Yifei Wang, Zeyu Jia, and Zaiwen Wen. The search direction correction makes first-order methods 2019.
faster. arXiv preprint
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pp. 681-688, 2011.
Andre Wibisono, Ashia C. Wilson, and Michael I. Jordan. A variational perspective accelerated methods in optimization. proceedings of the National Academy of Sciences, 113(47)EE7311- E7358, 2016.
Hongyi Zhang and Suvrit Sra. Towards riemannian accelerated gradient methods. arXiv preprint arXiv: /806.02822. 2018.
Ruiyi Zhang, Changyou Chen, Chunyuan Li, and Lawrence Carin. Policy optimization wasser- stein gradient flows. arXiv preprint 18xiv:1808.00 03030, 2018.
12

---

Under review as a conference paper at ICLR 2020
A THE PROOFS AND DERIVATIONS IN SECTION 2
In this section, we provide the derivation of the Euler-Lagrange equation, proofs in Section 2 and the Euler-Lagrange formulation (AIG).
A.1 DERIVATION OF THE EULER-LAGRANE EQUATION
We derive the Lagranne equation | in this subsection. For a fixed T > and two given densities p", .p"", consider the variational problem
T(p:) = inf Pt {L
C(p, Aepi)dl po , p²)
Let hy E F(2) be the smooth perturbation function that satisfies J heda 0, [0, and ho = hr = O. Denote pi P Ehe. Note that
T(p)) S'" L(pr, depe) + J [
8L 0,h1 drdl + o(e). 5(0.p.)
From di(pi)
0, it follows that
J." (
(8L 6pc
5L + 8(0:pt) ashs dzdt 0.
Note that hg hr 0. Perform integration by parts w.r.t. yields
8C J." 1 8(0(p))
8c o 5pt
hydadl 0.
Because f Shddz 0, ① holds with spatially constant function C(t).
A.2 THE PROOF OF PROPOSITION I
SECTION
In this subsection, derive the Hamiltonian flow in the probability space. First, we give useful identity. Given a metric tensor G(p) TTP( + P(?)), we have
f @ @iiG())odd / G(p)o102de = f Φb G(p) dzdr = / G(p) (lo P2de.
(14)
Here 4 G((p)-lor 01 and 422 G(p)-loz.. We then check that
8pt ( (fona(i))))pdd))
8
(f0,Cln) (0GG(p)(e1dd1
(15)
8p:
Let P. = Pt + eh, where To P()). For all a P, follows
G(p + G + eh)o
The first-order derivative w.r.t of the left hand side shall be 0, i.e..
[OG(p.) h G(p.)o G(p)) C(p)))+((())+
h (
[BG(pt) Op:
0.
Op.
Because Dipi G(p) "r, applying 14 yields
| depi
(G((p) h Oepedx ( (8G(pi)
11) depida
0pt
Op:
(16)
=-f0r Φ
_ h1 G(pr)aapda=-/
(aG(pi) (D(((e))) Op.
D,dr.
Op,
13

---

Under review as a conference paper at ICLR 2020
Based on basic calculations, compute that
@.p / (rppC(pp)))pddx =e dipt
(DG(Pe)..) h)
(Tpp Op.
Oepide + o(e), (17)
- (,G((p))1 - - = (
(oG(p))
Dide + o(e). (18)
Op:
Combining (16),
and (8)) yields
15
Hence, the Fuler-Lagraage equation is equivalent
1 à 0,0, ppp
(f
1 à 2pp
óE Öp:
8pt
This equation combining with S.p. = G(p)-La, recovers the Hamiltonian flow (). In short, the
Fuler-Lagraage equation (1 is from the primal coordinates (pt, Pspt) and the Hamiltonian flow 2
from the dual coordinates (pt. Dt). Similar interpretations can be found (Chow et al. 2019).
A.3
THE BULER-LAGRANGIAN FORMULATION OF AIG FLOWS
We can formulate AIG as a second-order equation of Pt.
D2 DeaPt Pc + a.dip: + G(pt)
SE
8pt
D² Here D2 Dt is the covariant derivative in metric G. We can also explicitly write D Pt as follows.
D Dt2 = Sup: - (8.G(()) + 2 G(p))- (SopAiii))nhpe
8
(ʃ
5pt
B THE PROOFS SECTION
3
In this section, we present proofs of propositions and theorems in Section
B. THE PROOF OF PROPOSITION 2 SECTION
We start with an identity. For twice differentiable ((x), we have
2 -' (V V
(19)
From W-AIG follows that
Eept V (pr VΦ.)
(20)
This is the continuity equation of Pe. Hence, on the particle level, XG shall follows
dx, = (X,)dt..
Let V1 = (X,)..,. Then, by the conservation of momentum and W-AIG). we have dV. = (a, + (X, VX,
( -a. vX., 2
òE
dt + (VΦ. V)
8pt
8E 5pt
8E 5pt
(Xt) 04 V,
(X,)
B.2 THE PROOF OF PROPOSITION
3 SUBSECTION 3.1
In this subsection, we derive the Hamiltonian flow in Gaussian. For A E S", we define the linear operator MA - S" by
MAB AB BA, B
14

---

Under review as a conference paper at ICLR 2020
It is easy to verify that if A P", then MA is well-defined. For flow Σ, E P 0, we define the Lagrangian L(2,, Ex)= 10,, 2xr E(2) The corresponding Euler-Lagrange equation writes
dl dL dt dÊ. dE'
(21)
Let St MM, E, i.., Σ, =, + S,S, Then, it follows
1 2 tr((S,,, + +ES,),, tr(E,,,) = tr, E.). gT, (E,, = ==((,t,,
This leads
M, St. For simplicity, we denote 9 9r, (E,, Σ.). First, show that dL 0b.
dg
S
Because S: -5,1 Σ. Given Eg, St can be viewed continuous function of Σ, For any A E S",
define LA 1r((E,,) + A).
0 dlA
as OLA or, as.
(AE, + E,A) + (AS, + S,A). aL as,
Or, o.,
Here view / linear operator S". Let B AΣ, E,A, then A My. B. as. B + Ms. M5. = holds for all B E S". Therefore, we have us. GE.
83, Mg. -5. Hence,
dg as. 0g 0g
d>. σc, os. Sc, Ms,
+ E,S,) + S = Ms, St + 5 = S-S
As a result, the LagernLagrange equation
21
equivalent to St 2
$ dl
VE(E))
(22)
2 dΣ,
Combining 22 with Σ. = S,E, + E,S, renders the Hamiltonian flow (5)
B.3 THE PROOF OF THEOREM SUBSECTION 3.1
We first show that Σ, stays in Pr, Suppose that E, E P, for 0 S t T. Define H; = H.,, tr(S,E.5.)/22 + E(E.). We observe that (W-AIG-G) is equivalent to
ÖH,
= - -2 HH, "E.
', as,"
(23)
We show that He is decreasing with respect to l.
dH,, dt
ÖH,
dH, E = tr
0H.
(ayst - 0H,
HHe 0H,
= tr
$,
+
+ 2
as,
as
oE, as,
Se 0H,,
at ir(Sr(E + S,E,)) = -0, (r(S, S
= = artr Ir
as,
2
For simplicity, we denote W (E*)-1 Let d be the smallest eigenvalue of Σx. Then, log det(L,WW)))
2 log del (E.) = n log Xo. Therefore.
~2og (log d + +1) S log det( +n] E(E) < H(t) < H()).
which indicates that
X > exp (-2 (-2H()-11 -1)
(24)
This means that long as Σ, E P the smallest eigenvalue of Σ, has positive lower bounded. If there exists T > 0 such that Er & P,, Because Σ, is continuous with respect to t, there exists Ti < T, such that E, E P.. 0 < t < Ti and \T, < <exp (-
(- (0)/n n 1), which violates 24
15

---

Under review as a conference paper at ICLR 2020
We then reveal the relationship between (W-AIG) in P(R") and IP.. We compute that
i
det(Li) (2, tr(2; Ex,
2 -ー -x, - -.
0l. 2
Combining with E = EnS, + S,En, we obtain
1r( 1r(E; = tr(S -
tr( - E x) = tr E1, Sax xx S = (t E xxx))
Therefore, follows
Oep((x)= = 0
(viet((s)
1
v (det(Zi)))(z)) + tr(( tr(a"22 "
01 Vdel(E.))
=-2 2 tr(2i, -
( + tr(S xx t (p(x)= -t tr(S)(1 Mx xx ))((2))
Note that (1.(x)) Sr... Hence,
-.(VV = Σ (()) = Σ (
=-
p((x) (t() (S,z) = t-rp((tttr(())) =-=))((= -,,,-E,'zzz 2x))) 0.p(
The first equation of W-AIG holds. Because 0,0((2)) -"' S,x/2 C(I),
Stx + (z) 4 (x) 2 [+)(1)) x -T 3₁ *T,
2 -= S7x C())
= =-z'vz, ((t)
- W*)
= 2
1
C(t)+ C(t).
Note that p² is the Gaussian density with the covariance matrix M. Because C(t) = log det E W*) 1, can compute
8E
log (z) - log (x) 1= 1
-1==11=-00 - Wz)x log det det(E, +
6p:
=
2
2
1 T ( W*xx (()))=--()=(0)----0))0 + 1. + (z)
1
2
Therefore, the second equation of (W-AIG) holds. Because Eg = E0, So = and C(0) = 0, we have Pu pd and o = 0. This completes the proof.
C THE PROOFS IN SECTION 4
In this section, we briefly review some geometric properties of the probability space as Riemannian manifold and present proofs of propositions in Section 4
C.1 A BRIEF REVIEW ON THE GEOMETRIC PROPERTIES OF THE PROBABILITY SPACE
Suppose that we have metric 9p in the probability space P(Q). Given two probability densities p",p² E P(S)), we define the distance as follows
(int -(nng (10)) ( (r(o)(())) -,000)0))))))00000-00000000)))=) 1/2
d(p'p''= (op')=(in)))
The minimizer (x) of the above problem is defined as the geodesic curve connecting pa and p². An exponential map at pd P(2)) mapping from the tangent space P(2) P(2). It requires that E P(I) is mapped to a point p² E P(?)) such that there exists geodesic curve (s)) satisfying (0) (0) and (1) p².
16

---

Under review as a conference paper at ICLR 2020
C.2 THE PROOF OF PROPOSITION
4 IN SECTION
In this subsection, we characterize the inverse of the exponential map in the probability space with the Wasserstein metric. Let 7 (sT, + (1 - 5) Id)- [o, 1). Then, based on the theory of optimal transport (Villani 2003). we can write the explicit formula of the geodesic curve 7(8) by
7(8) - 7 #p: = det (V; o T,
Through basic calculations, we can compute that
d ds
Ti
(sT, + (1 -)) Id)
= Ild- -T,.
ds
d
d
det
(VT,)
det(I + s(1 DTi) + x(())
= =tr(I -Di).
ds
ds
==0
===0
Therefore,
d
ds (8)
(x) = ir(T VT:)p( (V)(
=V (x ()())) + (v), -x())))))))))) = - -V. (p((() - =))),
which completes the proof.
C.3 THE PROOF OF PROPOSITION
5 AND 6 IN SECTION
4
The main goal of this subsection is to prove the Lyapunov function (t) is non-increasing.
Preparations. We first give a better characterization of the optimal transport plan T, We can write T. = vu,,, where v, is a strictly convex function, see (Villani 2003). This indicates that VT; is symmetric. We then introduce the following proposition.
Proposition? 7 Suppose that Ep) satisfies Hess( for B 2 Let Tx(2) be the optimal transport plan from Pr p", then
E(p') 2 E(pi) / (1,())=
SE pdx + B
f [[TT(x)) =112 pıda.
V 8p.
This is direct result of 3-geodesic convexity E(p) based on Proposition
Next, let us denote = - o Tg. We show that de satisfies
V (pe(1e - V = 0.
(25)
Because #p" = let (T;)) o7, and X; = (T;)-1X,, Xo. where Xo ~ p., This yields d X. ue (X.). The distribution of X, follows Pt. By the Euler's equation, shall follows di
0,p; + V (p:w.) = 0.
Combining this with the continuity equation 20 yields 25)
Then, formulate (x)) with u, By the Taylor expansion,
(1)(()) = (0)T))))++)))))+++++++))))))) +
Let y (T.)-1x it follows
(Te+e) (x) = (Ti) (T:) (x) + sns (()) os)) = y + sue(y) + oss).
=Ti
Therefore,
+ o(s)) 1
0 0= (T) - -
Te(y
= -T( + +0
o(s)
= (V 0(s).
We shall have VT,(y)u.()) + 0. Replacing by x gives
a,T,(x) (z) The following lemma illustrates two important properties of 44 and T,
(26)
17

---

Under review as a conference paper at ICLR 2020
Lemma For u, satisfying 25 we have
/ f - ue VI. pı 2 0,
/ (V - ue VTT - 2)) =0.
PROOF We first notice that u, v divergence- free in term of pr. From - -VT,4 = = = v0,V, we observe that -T,, is the gradient of 0 v, Therefore,
(v Ue VT p: /
(a, v, V (pc uz)) 0.
Based our previous characterization the optimal transport plan T, VT, - is symmetric positive definite. This yields that
( u., VTVV)) pidx = / (V u4 VI.VD,) pidz - (V ue VT.u) Pt
= / (V N1, VT. uz)) p.dx 2 0.
The last inequality utilizes that VT, is positie definite and P is non-negative. Then, we prove the equality in Lemma gradient. Similarly, It follows
Because (T) x) (((T((())) x1|2 + Ti(x) || is
/ ( uz, (T(()) - x)) = 0.
Lemma
D and the relationship 26 gives
- f (0.11, v pidxx / (un, V1.V41) pide S / (V4r, VT.V4.) pida,
(27)
/ (8.Ti, Tx)) - x) pide = (Vo,, (T(()) x)) ptdr.
(28)
Proof of Proposition 5 Based on the definition of the Wasserstein metric, we have
O,E(p:) = =- - àE v (p: Dd)
5pt
Differentiating {()
7 w.r.t. L renders
ε(t)e =B / (arTi, Ti(x) - x) p.dx-
B
[TT((t) =x12V (p.
2
VB / (D,Ti, v pidx vB (Ti(z) - x, V22)) p.dr
+ V / (1)(0)) - I, v V (pr V4, + / (V4 dn v pidz
1 / /22/// (p.V /
5E / (p.V
2
ōp:
VB
p.ds - 3 / (Ti(s)) x, (.)( pede
|V 1
V33
[1T((x) - x1| prd vB(E(p)) E(p")).
(29)
2
For the part (29)
Proposition renders
B
/
pedx + (BE(P)) S -V8 / (i(x)) -
SE
[T(())
Pidx.
(30)
2
8pt
We first compute the terms with the coefficient 3r in ((1)e- VBt, We observe that /|V (pe vdr
/ (V 0.) pida:
ôE
V (p. (D.)
5pr
(31)
= 0. V, + 1 vi
SE Vo, prde = -2/B aad==-2// 1V op
V
18

---

Under review as a conference paper at ICLR 2020
where the last equality uses (W-AIG) with az 2/B. Substituting 330 and (11 into the expression of (t)ee- vBt yields
/ (0.Te, Ti(x) 2) p,dr-
8 -/ / [TT((x) x/12V (d.V
E(t)e vBt 53
2
- 8/ (Te(x)) x, v pıdx - vB / (0,T pydx
(32)
VB f (T((2)) x, 0. v., pıd v3 / TT(x)) - J.
8E
pıdr
óp
/ 1r0,.l'pdde 3v3
VB (T((t) - Vo.) V (pr V )dl
2
Then, we deal with the terms with V (pr v), We have the following two identities
f -, 2,, V (p.V = / 1V (T(x)) - x, Vo.). v pıdr
/ (V.,, V2 (0)) =)+ (VT((z) 1) p.dr
(33)
1 / (Ti(0) x, v 13) prdx pidx- -
f
f (V VI,VD,) prde + / |V prdae
- / / (1T((s) - x12r (pr Vo: ((VT)())- 2)),, v..)) padx
(34)
| (T((=)- x, VT.V pıd - / (T<i(o) Voxe pidx.
Hence, we can proceed to compute the terms with the coefficient /3. 27 and
33
yields
-v8/ / (0.Ti, v.,) prdk /B (T(Co) x, 0 V.. + V dE
(T((a))
prdx
5p.
3/B
2 / /IV0,12 pade + VB (Ti(s)) - x,, Vo..) V. (pr Do
vB / (0.T,+ + VI. Vot)) pidr
vB
/ | 11 pidx
(35)
2
/B / ( (T(()) ( + V
5E
+ v 11,11 pedx
6p 2
vB / |V prx + 23 / (Ti(x) x, v prdx.
< 2
Substituting 34p and 35 into 32 gives
VB
/ || pede
ε(t)e
Vet
2
(0,Ti, x) p.dx= 3 / [[T((x) x1|2 - (,)
<8
3 / (T((T) - x, v pade 23 / (Ti(x) - x, v ptdx
=3 1 (O.T1
V..V VV.V+ T((x) - x) p. = 0,
where the last equality uses
28 In summary, we have
vB / 11 p.d S 0.
ε(f)e VBt
2
19

---

Under review as a conference paper at ICLR 2020
Proof of Proposition 6 Differentiating {t)
8 w.r.t. we compute that
2 / [T+(z) - x1|² (pr V )de
ε(1) = (0,T,, (x( x) prde
/
OcTi, t VΦ,
prdle / /(T.(t)) - J,, 2 VØ, + t a. V prde
(T((x))
2
2
+ / (I((1)) - It (VD))
F.. (/200
1
v + L 20..00 prdx
/
1
V V (pc
1 dF v (p:Vo.) 1 ((((m)) Elpr)).
4 ópc
2
(36)
Because Epp) is Hess(O), Proposition 7 yields
(T1(6)= V OF 8p
E(p() F(p)) E(p') S -
pidx.
(37)
Utilizing the inequality (37) and substituting the expressions of terms involving 0,T, and Vd. in (36) with the expressions in
(28) and (33) (34), we obtain
ε(1) (V (T)()) x)) pide / (Ti(()) x. VT.V0.) p.dx
(T((x) x, Voh) pade + 1 / (V VT:Vo.) pıde
1 / (Ti(s) v pidex
~ / (0. Tx(x) x)) pide
1
2
2
/ / (Ie(s) - x, v/[V2.11>) 33) prd
t
2/2 / (Vel, VT:V pedx
(38)
/ /V + &
4 / /
1² / (V O v.)) pedx
+ 2
4
/ (Vo,, vI[V0.11?) pido 8² #/
4 [
8E
V
prde
5pr
t / (T((z)- - x, V 8E ōpr
Ptdz.
The expression of 388 can be reformulated into
(Te(x) x, V..)) pidx + / [IV pıdr
E(t) <
2
4
/ (Ti(x) x, v.
1
5E
/ 1|2
v
8p: )
pidz
2
t²
/ ( ( a. VΦ
SE V
+ 4
1|
pidr.
8p:
From W-AIG with a; 3/tt, we have the following equalities.
t²
SE öpt
3t
/ 17
pı =
2
4
-
t / -((x))2.0VV0) v
V
6E
pidx
3
ʃ /(Tile) xz V pidz.
óp.
2
As result, E(1) < 0. This completes the proof.
20

---

Under review as a conference paper at ICLR 2020
C.4 COMPARISON WITH THE PROOF IN THE ACCELERATED FLOW
The accelerated flow in TTghvaei & Mehta 2019) is given by
dX. c -
dY: -a
OF
(Xe).
(39)
dt
dt
6p.
Here the target distribution satisies Poo (x) p² (x) & exp(-f(x)) Suppose that we take O2 = logp- p log 3, - p log i log C' and 74 = plog Here we specify p 2 and C' 1/4. Then the accelerated flow (39) is identical to (W-AIG-P-KL) if we replace Y, by 2t VV. The Lyapunov
function in (Taghvaei & Mehta 2019) follows
(tt) X, + 6 YY Th (x.)1 + e (E(()) E(p'))
2
2 EI1x, + f V/- (x.)
8 (B(p)) - E((0))
4
-
1 / (rias - -)+
( (r(x)dx + - (B(pt) - E(p")).
The last equality is based on the fact that Vi = (X,, and T, = T the optimal transport plan from Pr to p0." This indicates that the Lyapunov function in (Tahhaae & Mchta 2019) is identical to ours The technical assumption in (Taghvaei & Mehta
2019) follows
d X1 + e " "Y, Tg. (xi)
0 EE
ae Tg (Xi)
X₂ + 1 Vi T((Xi)
d (
E
2
dt
X, + t Vi - T((Xi) ((0)T)(()) + VTNV)
=E
2
(x)x) + (x), + V..Vo, ped
Based on 0,T, = VT, and Lemma D we have
/ (r - (o-TT(x),, 0.T, VT, pt = / (1- VTi(V,; - ue)) pedr
( T.T, + VT.V p, = / (V VT,(V - -4)) p.dz
= -/ (VB. us VT, ue)) pidx 20.
As result, we have
E (X++e Tp (Xi))
dl To (X:)
t (V Ue, VT. - u.)) ptdx 2 0. 2
In 1-dimensional case, because V (pr(u: - (.))) = 0 indicates that Pe(u: V = 0. For (.))) > O, we have (x) (z) 0. So the technical assumption holds. In general cases, although u, = (T,) o T, satisfies V (p V).)) = O, but this does not necessary indicate that u, = v. Hence, E [X, + e "Y, TP (X+)) 4TR (X.) = 0 does not necessary hold except for 1 1-dimensional case.
D IMPLEMENTATION DETAILS IN THE NUMERICAL EXPERIMENTS
In this section, we claborate on the implementation details in the numerical experiments.
D.1 DETAILS IN SUBSECTION
The initial distribution of the particle system follows the standard Gaussian (0, The objective density function is
('(x) x exp( 1 2([|=|| - (ex))cxp-22(x)))) - 3)²) + exp(-2(xx + 3) x E R.?
All methods run for 200 iterations using the same fixed step size 0.1.
21

---

Under review as a conference paper at ICLR 2020
D.2 DETAILS IN SUBSECTION 6.2
The discretization the Wasserstein gradient flow in Gaussian follows
Ex+1 Σg 2r ( F(EE)) + Vs, F(E))E))
(40)
The discretization of (WAIG-GG is given by
Sk+1 a v (S} E(E,))
Ex+1 + VTSK+1) VFSk+1).
(41)
The choice of aç is same as the one for (9) based on whether 3 is known. To ensure the update of Σk+1 won't 't blow up, we check whether tr(rS²) > n. If this holds, we restart the algorithm. In fact, the update on Σn can be viewed as the exponential map from En with the direction S + Se+1Ex, Nevertheless, this an exponential map and only if + VTSk+1 is positive definite. The update rule of VTSk+1 can rewrite into
VTS =xVTS (VTSx)² 2T VE,
(42)
Because of the existence of - (VTS))) as long as the spectral radius (the eigenvalue with the largest absolute value) of VTS² is greater than 1, then / + VTSk+1 cannot maintain to be positive definite. So at this time we need restart the algorithm. Nevertheless, to compute the spectral radius is computational costly. Instead, we use a weaker condition tr(T5k) > n.
We also introduce the adaptive restart technique for 41
Consider
Pp = tr(Sk+1 V B(*+))
This restarting criterion corresponds to the particle version the algorithm using initial values En Σg and So 0.
13
Similarly, you < O, we restart
For the particle implementation, the update rule of WGF writes
X -(((((() + (X(E))
where is computed based on KDE
D.3 DETAILS IN SUBSECTION 6.3
We follow the same setting as (Liu & Wang 2016), which is also adopted Liu et al. (2018 2019).
The data set split into 80% for training and 20% testing. The mini-batch size is taken as 50. For MCMC, the number of particles is N 1000; for other methods, the number of particles is N 100. The BM method not applied to SVGD in selecting the bandwidth.
The initial step sizes for the compared methods are given in Table I The step size of SVGD is adjusted by Adagrad, which is same as (Liu & Wang 2016). For WNAG, AIG and WRes, the step size is give jet al. 2018) and (Liu et al. 2019). For MCMC, WGF and AIG-r, the step size is multiplied by 0.9 every 100 iterations. We record the cpu-time for each method in Table The computational cost
7 for 1. The parameters for WNAG and Wnes
identical to (Liu
Method
MCMC
SVGD
WNAG
Wnes
WGF
AIG
AIG-r
Step size
le-5
0.05
le-6
le-5
le-5
le-5
le-6
Table 1: Initial step sizes of algorithms in comparison.
of the BM method is much higher than the MED method because we need to evaluate the MMD of two particle systems several times in optimizing the subproblem.
Method
MCMC
SVGD
WNAG
Wnes
WGF
AIG
AIG-r 78.945
BM
13.962
5.539
78.038
78.509
79.006
79.094
MED
13.909
5.581
5.623
5.625
5.395
5.890
5.689
Table Averaged epu time(s) cost for algorithms comparison.
22


Reviews:
[{"review_id": "BklKUV1fcB", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "I acknowledge reading the rebuttal of the authors. Thank you for your clarifications and explanation. My point was this paper would make a good submission to ICLR if it was better motivated presented and explained to a wider audience. Unfortunately in its current form it can only reach a limited audience.\n\n####\nSummary of the paper: \n\nThe paper proposes accelerated information flows under Wasserstein or Fisher rao metric . i.e a method for solving minimization of probability functional , where the probability space is either endowed with the Wasserstein or the Fisher distance.\n\nGradient descent methods in euclidian spaces can be accelerated using a type of momentum , this paper extends this to gradient flows, using similar formalism of Hamiltonian that appeals to the dynamic of a the particles and the velocity (or momentum, $H(x,p)=\\frac{1}{2}||p||^2+ \\mathcal{E}(x)$).  Writing down the lagrangian one obtains two co-evolving PDE one of the dynamic of the density and one for the potential . The PDEs are specified for both Fisher Rao, and Wasserstein distance. The hamiltonian for example for the Wasserstein distance $H(\\rho_t,\\Phi_t)=\\frac{1}{2}\\int||\\nabla_x \\Phi_t||^2+ \\mathcal{E}(\\rho_t)$.  and PDEs amounts a continuity equation for the density  evolving with drift $\\nabla \\Phi_t$, the evolution of the momentum $\\Phi_t$ is also given by a PDE. \n\nProposition 2 of the paper gives the particles differential equation corresponding to the system of PDEs. For the energy being the KL divergence an explicit expression is given , this expression remains difficult in practice since it needs the knowledge of the density $\\rho_t$. Authors propose in the application section to use gaussian approximation , or using a kernel density estimators. The Bandwidth of the kernel is then choosen using a heuristic proposed in the paper.  \n\nThe paper then focuses on deriving expression for flows when the densities are centered gaussians, and this amounts to an ODE on the covariance , the ODE is discretized in Appendix D.2 to lead to  computational method. Then convergence of the flow is analyzed for the wasserstein accelerated flows, under \"\\beta- convextiy\" in the wasserstein sense of the functional. \n\nSome experiments of the particle based method are shown on synthetic experiments and in bayesian logistic regression. \n\nReview: \n\nContribution/ Clarity:\n\nThe main contribution of the paper is in deriving the accelerated gradient flow for the wasserstein distance this was also addressed in a recent paper [Taghevia and Mettha 2019]. \n\nThe technical contribution is interesting but given that this field of flows in probability space is still not very well spread in the ML community, I wonder if ICLR is the best fit for this type of work.  I support good theoretical work, but I think the authors could have done a better job in exposing the ideas how they extend form euclidean space, to manifolds, to probability spaces gradually. Simple derivations of euclidean space Hamiltonian will help the reader that is not exposed to such literature. I think the paper will benefit from a less technical writing in introducing the ideas coming from euclidean space and in conveying the intuitions. \n\nComments: \n\n- In the proof of Proposition 2 you give the expression of evolution of $dV_t$ by conservation of the momentum. Could you please elaborate more how you obtain this expression, and where you proved the conservation of momentum?\n\n- In term of damping if ones uses the Wasserstein Fisher Rao flows  , one obtain also accelerartion , maybe you can comment on that ? since you analyze both flows , would be interesting to discuss the relation to Global convergence of neuron birth-death dynamics, that shows that an acceleration is obtained via WFR flows, since it will introduce a damping as well. \n\n- since MCMC and BM method lead to similar result what is the advantage of the wasserstein accelerated flow? one could also implement also an accelerated langevin dynamic \n\n\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "B1eJdibycS", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper attempts to develop a counterpart of the well-known Nesterov accelerated gradient method for gradient flows on the space of probability measures equipped with an information metric. This is an important problem which is useful for optimization on probability spaces. The accelerated gradient flow is developed by leveraging a damping Hamiltonian flow. The paper focuses mainly on the case with the Wasserstein metric and provides a convergence analysis. Practical considerations such as discretizing the accelerated flow and bandwidth selection are developed for the use of the method in practical problems.\n\nAlthough the paper has some important merit, I find the paper extremely hard to follow, partly because of its writing style. There is not enough motivation and explanation for the ideas presented. Some discussions and sentences either do not make much sense to me or read badly. For example, this sentence \"For the Wasserstein gradient, many classical methods such as Markov Chain Monte Carlo .... are based on this framework...\" doesn't make sense, as the development of MCMC is never based on Wasserstein gradient. Or the sentence right before it \"For the Fisher-Rao gradient, classical results including Adam ... and K-FAC .. demonstrate its effectiveness in ...\": it's not clear what the authors are trying to say here. Adam is not relevant to the Fisher-Rao natural gradient while K-FAC is just an approximation method and isn't a good reference for demonstrating the effectiveness of the natural gradient. Also, there are many English typos and grammar errors.  \n\nI didn't read the proof carefully due to the time constraint, so I cannot judge on the theoretical part of the paper. The numerical experiment is quite limited as it considers very simple problems (a toy example, a single Gaussian distribution and a logistic regression problem). As such, I think there isn't enough evidence to judge the usefulness of the proposed method in practice. \n\nHaving said that, I believe this paper can be an important contribution if the authors invest more time on refining its presentation and if more thorough experimental studies are conducted.\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper makes its contribution by deriving an accelerated gradient flow for the Wasserstein distances. It is technically strong and demonstrates it applicability using examples fo Gaussian distributions and logistic regression.

Reviewer 3 provided a deep technical assessment, pointing out the relevance to our ML community since these ideas are not yet widespread, but had concerns about the clarity of the paper. Reviewer 2 had similar concerns about clarity, and was also positive about its relevance to the ML community. The authors provided details responses to the technical questions posed by the reviewers. The AC believes that such work is a good fit for the conference. The reviewers felts that this paper does not yet achieve the aim of making this work more widespread and needs more focus on communication.

This is a strong paper and the authors are encouraged to address the accessibility questions. We hope the review offers useful points of feedback for their future work.
Decision: Reject