Natural gradient via optimal transport

Abstract

We study a natural Wasserstein gradient flow on manifolds of probability distributions with discrete sample spaces. We derive the Riemannian structure for the probability simplex from the dynamical formulation of the Wasserstein distance on a weighted graph. We pull back the geometric structure to the parameter space of any given probability model, which allows us to define a natural gradient flow there. In contrast to the natural Fisher–Rao gradient, the natural Wasserstein gradient incorporates a ground metric on sample space. We illustrate the analysis of elementary exponential family examples and demonstrate an application of the Wasserstein natural gradient to maximum likelihood estimation.

Appendices

Appendix

In this appendix we review the equivalence of static and dynamical formulations of the $L^2$-Wasserstein metric formally. For more details see [38].

Consider the duality of linear programming.

$$\begin{aligned} \begin{aligned}&\frac{1}{2}W(\rho ^0,\rho ^1)^2\\&\quad =\inf _{\pi \ge 0}\Big \{\int _{\Omega }\int _{\Omega }\frac{1}{2}d_{\Omega }(x,y)^2\pi (x,y)dxdy:\int _\Omega \pi dy=\rho ^0(x),~\int _\Omega \pi dx=\rho ^1(y)\Big \}\\&\quad =\sup _{\Phi ^1, \Phi ^0}\Big \{\int _{\Omega }\Phi ^1(y)\rho ^1(y)dy-\int _\Omega \Phi ^0(x)\rho ^0(x)dx:\Phi ^1(y)-\Phi ^1(x)\le \frac{1}{2}d_\Omega (x,y)^2\Big \}. \end{aligned} \end{aligned}$$

(21)

By standard considerations, the supremum in the last formula is attained when

$$\begin{aligned} \Phi ^1(y)=\sup _{x\in \Omega }~\Phi ^0(x)+\frac{1}{2}d_\Omega (x,y)^2. \end{aligned}$$

(22)

This means that $\Phi ^1$, $\Phi ^0$ are related to the viscosity solution of the Hamilton-Jacobi equation on $\Omega$:

$$\begin{aligned} \frac{\partial \Phi (t,x)}{\partial t}+\frac{1}{2}g_x^\Omega (\nabla \Phi (t,x), \nabla \Phi (t,x))=0, \end{aligned}$$

(23)

with $\Phi ^0(x)=\Phi (0,x)$, $\Phi ^1(x)=\Phi (1,x)$. Hence (21) becomes

$$\begin{aligned}&\frac{1}{2}W(\rho ^0,\rho ^1)^2\\&\quad =\sup _{\Phi }\Big \{\int _{\Omega }\Phi ^1(x)\rho ^1(x)-\Phi ^0(x)\rho ^0(x)dx:\frac{\partial \Phi (t,x)}{\partial t}+\frac{1}{2}g_x^\Omega (\nabla \Phi (t,x), \nabla \Phi (t,x))=0 \Big \}. \end{aligned}$$

By the duality of above formulas, we can obtain variational problem (1). In other words, consider the dual variable of $\Phi _t=\Phi (t,x)$ by the density path $\rho _t=\rho (t,x)$, then

$$\begin{aligned} \begin{aligned}&\frac{1}{2}W(\rho ^0,\rho ^1)^2\\&\quad =\sup _{\Phi _t}\inf _{\rho _t}~\int _{\Omega }\Phi ^1\rho ^1-\Phi ^0\rho ^0dx-\int _0^1\int _{\Omega }\rho _t\big [ \partial _t\Phi _t+\frac{1}{2}g_x^\Omega (\nabla \Phi _t, \nabla \Phi _t)dx\big ] dt\\&\quad =\sup _{\Phi _t}\inf _{\rho _t}~\int _{\Omega }\Phi ^1\rho ^1-\Phi ^0\rho ^0dx-\int _0^1\int _{\Omega }\rho _t \partial _t\Phi _tdxdt\\&\qquad - \int _0^1\int _{\Omega }\frac{1}{2}g_x^\Omega (\nabla \Phi _t, \nabla \Phi _t)\rho _tdx dt\\&\quad =\sup _{\Phi _t}\inf _{\rho _t}~\int _0^1\int _{\Omega }\partial _t\rho _t \Phi _t-g_x^\Omega (\nabla \Phi _t, \nabla \Phi _t)\rho _tdx dt+\int _0^1\int _{\Omega }\frac{1}{2}g_x^\Omega (\nabla \Phi _t, \nabla \Phi _t)\rho _tdx dt \\&\quad =\inf _{\rho _t}\sup _{\Phi _t}~\int _0^1\int _{\Omega }\Phi _t(\partial _t\rho _t+\text {div}(\rho \nabla \Phi _t)) dt+\int _0^1\int _{\Omega }\frac{1}{2}g_x^\Omega (\nabla \Phi _t, \nabla \Phi _t)\rho _tdx dt \\&\quad =\inf _{\rho _t}~\Big \{\int _0^1\int _{\Omega }\frac{1}{2}g_x^\Omega (\nabla \Phi _t, \nabla \Phi _t)\rho _tdx dt:\partial _t\rho _t\\&\qquad +\text {div}(\rho \nabla \Phi _t)=0,~\rho _0=\rho ^0, ~\rho _1=\rho ^1\Big \}. \end{aligned} \end{aligned}$$

The third equality is derived by integration by parts w.r.t. t and the fourth equality is by switching infimum and supremum relations and integration by parts w.r.t. x.

In the above derivations, the relation of Hopf–Lax formula (22) and Hamilton–Jacobi equation (23) plays a key role for the equivalence of static and dynamic formulations of the Wasserstein metric. This is also a consequence of the fact that the sample space $\Omega$ is a length space, i.e.,

$$\begin{aligned} d_\Omega (x,y)^2=\inf _{\gamma (t)}\Big \{\int _0^1g_{\gamma (t)}^\Omega (\dot{\gamma }, \dot{\gamma })dt:\gamma (0)=x,~\gamma (1)=y\Big \}. \end{aligned}$$

However, in a discrete sample space I, there is no path $\gamma (t)\in I$ connecting two discrete points. Thus the relation between (22) and (23) does not hold on I. This indicates that in discrete sample spaces, the Wasserstein metric in Definition 1 can be different from the one defined by linear programming (5). See many related discussions in [12, 26].

Notations

We use the following notations.

Continuous/discrete sample space	$\Omega$	I
Inner product	$g^\Omega$	$g^I$
Gradient	$\nabla$	$\nabla _G$
divergence	$\text {div}$	$\text {div}_G$
Hessian in $\Omega$	Hess
Potential function set	$\mathcal {F}(\Omega )$	$\mathcal {F}(I)$
Weighted Laplacian operator	$-\nabla \cdot (\rho \nabla )$	L(p)

Continuous/discrete probability space	$\mathcal {P}_+(\Omega )$	$\mathcal {P}_+(I)$
Probability distribution	$\rho$	p
Tangent space	$T_\rho \mathcal {P}_+(\Omega )$	$T_p\mathcal {P}_+(I)$
Wasserstein metric tensor	$g^W$	$g^W$
Dual coordinates	$\Phi (x)$	$(\Phi _i)_{i=1}^n$
Primal coordinates	$\sigma (x)$	$(\sigma _i)_{i=1}^n$
First differential operator	$\delta _\rho$	$\nabla _p$
Second differential operator	$\delta ^2_{\rho \rho }$
Gradient operator		$\nabla _W$
Hessian operator		$\text {Hess}_W$
Levi–Civita connection		$\nabla ^W_{\cdot }\cdot$

Parameter space/Probability model	$\Theta$	$p(\Theta )$
Inner product	$g_\theta$	$g_{p(\theta )}$
Tangent space	$T_\theta \Theta$	$T_{p(\theta )}p(\Theta )$
$L^2$-Wasserstein matrix	$G(\theta )$
$L^2$-Wasserstein distance	$\text {Dist}$	$\text {Dist}$
Second fundamental form		$B(\cdot , \cdot )$
Projection operator		H
Levi–Civita connection		$(\nabla ^W_\cdot \cdot )^{||}$
Jacobi operator	$J_\theta$
First differential operator	$\nabla _\theta$
Gradient operator	$\nabla _g$
Hessian operator	$\text {Hess}_g$