From Optimization to Sampling Through Gradient Flows

Optimization and sampling algorithms play a central role in science and engineering as they enable finding optimal predictions, policies, and recommendations, as well as expected and equilibrium states of complex systems. The notion of “optimality” is formalized by the choice of an objective function, while the notion of an “expected” state is specified by a probabilistic model for the distribution of states. Optimizing rugged objective functions and sampling multimodal distributions is computationally challenging, especially in high-dimensional problems. For this reason, many optimization and sampling methods have been developed by researchers working in disparate fields such as Bayesian statistics, molecular dynamics, genetics, quantum chemistry, machine learning, weather forecasting, econometrics, and medical imaging.

State-of-the-art algorithms for optimization and sampling often rely on ad-hoc heuristics and empirical tuning, but some unifying principles have emerged that greatly facilitate the understanding of these methods and the communication of algorithmic innovations across scientific communities. This article is concerned with one such principle: the use of gradient flows, and discretizations thereof, to design and analyze optimization and sampling algorithms. The interplay between optimization, sampling, and gradient flows is an active research area and a thorough review of the extant literature is beyond the scope of this article.⁠¹ Our goal is to provide an accessible and lively introduction to some core ideas, emphasizing that gradient flows uncover the conceptual unity behind several existing algorithms and give a rich mathematical framework for their rigorous analysis.

¹

AMS Notices limits to $20$ the references per article; we refer to GTSA20CLGL$^{+}$20GIHLS20 for further pointers to the literature.

✖

We present motivating applications in section 1. Section 2 is focused on the gradient descent approach to optimization and introduces fundamental ideas such as preconditioning, convergence analysis, and time discretization of gradient flows. Sampling is discussed in section 3 in the context of Langevin dynamics viewed as a gradient flow of the Kullback–Leibler (KL) divergence with respect to (w.r.t) the Wasserstein geometry. Some modern applications of gradient flows for sampling are discussed in section 4, followed by concluding remarks in section 5.

1. Motivating Applications

We outline two applications in Bayesian statistics and molecular dynamics that illustrate some important challenges in optimization and sampling.

1.1. Bayesian statistics

In Bayesian statistics GCSR95, an initial belief about an unknown parameter is updated as data becomes available. Let $\theta$ denote the unknown parameter of interest belonging to the parameter space $\Theta$ and let $\pi _0(\theta )$ denote the prior distribution reflecting our initial belief. Furthermore, let $y$ be the observed data also belonging to an appropriate space $\mathcal{Y}$. Then Bayes’s rule identifies the distribution of $\theta$ conditioned on the data $y$:

$$\begin{equation} \pi (\theta | y) \propto \pi (y | \theta ) \pi _0(\theta ), {\tag{1.1}} \end{equation}$$

where $\propto$ indicates that the right-hand side should be normalized to define a probability distribution. Here $\pi (y | \theta )$ is called the likelihood function and $\pi (\theta | y)$ is called the posterior distribution. Bayesian inference on $\theta$ is based on the posterior, which blends the information in the prior and the data.

The choice of prior and likelihood is a modeling task which, perhaps surprisingly, is often not the most challenging aspect of Bayesian inference. The main challenge is to extract information from the posterior since (i) it typically does not belong to a standard family of distributions, unless in the restrictive case of conjugate models GCSR95; (ii) the parameter $\theta$ can be high dimensional; and (iii) the normalizing constant $\int _\Theta \pi (y|\theta ) \pi _0(\theta )\, d\theta$ in 1.1 (known as the marginal likelihood) is rarely available and it can be expensive to compute. These practical hurdles inform the design of optimization and sampling algorithms to find posterior statistics.

Of particular importance is the posterior mode or maximum a posteriori (MAP) estimator

$$\theta _{\text{MAP}} \coloneq \operatorname *{arg\,max}_\theta \pi (\theta | y).$$

Many optimization algorithms for MAP estimation start from an initial guess $\theta _0$ and produce iterates $\{\theta _n\}_{n=1}^N$ by discretizing a gradient flow with the property that $\theta _n \approx \theta _{\text{MAP}}$ for large $n$. Such gradient flows in parameter space will be discussed in section 2.

Computing MAP estimators is closely related to classic regularization techniques such as penalized least squares and Tikhonov regularization SAST18. In order to fully leverage the Bayesian framework, it is often desirable to consider other posterior statistics such as mean, variance, credible intervals, and task-specific functional moments, which can be written in the form

$$\mathbb{E}^{\pi (\cdot | y)} [ \phi (\theta )] \coloneq \int \phi (\theta ) \pi (\theta | y) \, d\theta ,$$

where $\phi$ is a suitable test function. Since $\theta$ is often high dimensional, the standard approach to compute these expectations is to use Monte Carlo Liu01: obtain $N$ samples $\{\theta _n\}_{n=1}^N$ from the posterior $\pi (\theta | y)$, and then approximate

$$\mathbb{E}^{\pi (\cdot | y)} [ \phi (\theta )] \approx \frac{1}{N} \sum _n \phi (\theta _n).$$

While Monte Carlo integration is in principle scalable to high dimensions, the task of generating posterior samples is still highly nontrivial. To that end, one may consider sampling $\theta _n \sim \rho _n,$ where the sequence $\{\rho _n\}_{n=1}^N$ arises from discretizing a gradient flow with the property that $\rho _n \approx \pi (\cdot | y)$ for large $n.$ Such gradient flows in the space of probability distributions will be discussed in sections 3 and 4.

To relate the discussion above to subsequent developments, we note that dropping the data $y$ from the notation, the posterior density can be written as

$$\begin{equation} \pi (\theta ) = \frac{1}{Z} \exp \bigl ( - V(\theta ) \bigr ), {\tag{1.2}} \end{equation}$$

where $Z = \int _\Theta \pi (y|\theta ) \pi _0(\theta )\, d\theta$ is the marginal likelihood and $V: \Theta \times \mathcal{Y}\to \mathbb{R}$ is the negative logarithm of the posterior density.

1.2. Molecular dynamics

Another important source of challenging optimization and sampling problems is statistical mechanics, and in particular the simulation of molecular dynamics (see chapter 9 in Liu01). According to Boltzmann and Gibbs, the positions $q$ and momenta $p$ of the atoms in a molecular system of constant size, occupying a constant volume, and in contact with a heat bath (at constant temperature), are distributed according to

$$\begin{equation} \pi (q,p) = \frac{1}{Z} \exp \Bigl (-\beta \bigl (U(q)+K(p)\bigr )\Bigr ), {\tag{1.3}} \end{equation}$$

where $Z$ is a normalizing constant known as the partition function, $\beta$ represents the inverse temperature, $U$ is a potential energy describing the interaction of the particles in the system, and $K$ represents the kinetic energy of the system. Letting $\theta \coloneq (q,p),$ we can write the Boltzmann–Gibbs distribution 1.3 in the form 1.2, with

$$\begin{equation} V(\theta ) = -\beta \bigl (U(q)+K(p)\bigr ). {\tag{1.4}} \end{equation}$$

As in Bayesian statistics, it is important to determine the most likely configuration of particles (i.e., the mode of $\pi$), along with expectations of certain test functions w.r.t. the Boltzmann distribution. These two tasks motivate the need for optimization and sampling algorithms that acknowledge that the potential $U$ is often a rough function with many local minima, that the dimension of $q$ and $p$ is large, and that finding the normalizing constant $Z$ is challenging.

2. Optimization

In this section we discuss gradient flows for solution of the unconstrained minimization problem

$$\begin{equation} \text{minimize } V(\theta ) \quad \text{s.t. } \theta \in \Theta , {\tag{2.1}} \end{equation}$$

where $V(\theta )$ is a given objective function. Henceforth we take $\Theta \coloneq \mathbb{R}^d$ unless otherwise noted. As guiding examples, consider computing the mode of a posterior or Boltzmann distribution by minimizing $V$ given by 1.2 or 1.4. The methods described in this section are applicable beyond the specific problem of finding the modes, however, this interpretation will be of particular interest in relating the material in this section to our discussion of sampling in section 3.

2.1. Gradient systems

One of the most standard approaches to solve 2.1 is gradient descent, an optimization scheme that is based on the discretization of the gradient system

$$\begin{equation} \dot{\theta }_t = - \nabla V(\theta _t), \quad t >0, {\tag{2.2}} \end{equation}$$

with user-defined initial value $\theta _0$; throughout this article $\nabla V(\theta _t)$ will denote the gradient of the function $V$ at the point $\theta _t$, which will be tacitly assumed to exist wherever needed.

While equation 2.2 is perhaps the most popular formulation of the continuous-time gradient descent dynamics, the equivalent integral form below reveals more transparently some of its properties:

$$\begin{equation} V(\theta _t) = V(\theta _s) - \frac{1}{2}\int _{s}^t |\nabla V (\theta _r)|^2 dr - \frac{1}{2} \int _{s}^t |\dot{\theta }_r|^2 dr, {\tag{2.3}} \end{equation}$$

for all $t\geq s >0$. Indeed, notice that from 2.3 it is apparent that $V(\theta _t) \leq V(\theta _s)$ for $s \leq t$, i.e., the value of the function $V$ decreases in time, and in all but a few trivial situations the decrease is strict. Another advantage of the reformulation 2.3 (or its inequality form 2.5 below) is that it can be adapted to more general settings with less mathematical structure than the one needed to make sense of 2.2. In particular, 2.3 can be used to motivate a notion of gradient flow in arbitrary metric spaces; see AGS08 for an in-depth discussion of this topic.

Proposition 2.1.

Suppose $V$ is a $C^1$ function. Then 2.2 and 2.3 are equivalent and they both imply

$$\begin{equation} |\nabla V ( \theta _t)|^2 = |\dot{\theta }_t|^2. {\tag{2.4}} \end{equation}$$

Proof.

By Cauchy–Schwarz and Young’s inequalities, for any $t>0$ it holds that

$$\begin{align*} - \langle \nabla V(\theta _t), \dot{\theta }_t \rangle &\le | \nabla V(\theta _t) | |\dot{\theta }_t | \\ & \le \frac{1}{2} |\nabla V(\theta _t) |^2 + \frac{1}{2} |\dot{\theta }_t |^2, \end{align*}$$

and both inequalities are equalities iff $-\nabla V(\theta _t) = \dot{\theta }_t.$ Therefore,

$$\begin{align*} V(\theta _t) &= V(\theta _s) + \int _s^t \langle \nabla V(\theta _r), \dot{\theta }_r \rangle dr \\ & \ge V(\theta _s) - \frac{1}{2} \int _s^t |\nabla V(\theta _r)|^2 dr - \frac{1}{2} \int _s^t |\dot{\theta _r}|^2 dr, \end{align*}$$

and equality holds iff $-\nabla V(\theta _t) = \dot{\theta }_t$ for all $t>0.$ The identity $|\nabla V(\theta _t)|^2 = |\dot{\theta }_t|^2$ follows directly from 2.2.

Notice that the relationship $\dot{\theta }_t = - \nabla V(\theta _t)$ is only required in proving the energy dissipation inequality

$$\begin{equation} V(\theta _t) \leq V(\theta _s) - \frac{1}{2}\int _{s}^t |\nabla V (\theta _r)|^2 dr - \frac{1}{2} \int _{s}^t |\dot{\theta }_r|^2 dr, {\tag{2.5}} \end{equation}$$

since the reverse inequality is a consequence of Cauchy–Schwarz. Notice further that 2.2 implies 2.4, but in general the converse statement is not true. For example, the flow $\dot{\theta _t} = \nabla V (\theta _t)$ satisfies 2.4, but in general does not satisfy 2.2. Likewise, 2.2 implies $\frac{d}{dt}V(\theta _t) =- |\nabla V(\theta _t)|^2$ (which follows directly from the chain rule), but not conversely. Indeed, in $\mathbb{R}^2$ we may take $A$ to be any orthogonal matrix and consider $\dot{\theta }_t = -A^2 \nabla V(\theta _t)$ so that $\frac{d}{dt}V(\theta _t) =- |\nabla V(\theta _t)|^2$ but 2.2 is not, in general, satisfied. This digression illustrates that equation 2.3 captures in one single identity of scalar quantities the vectorial identity 2.2, even if it is not as intuitive as other scalar relations.

2.2. A note on convergence

Despite the fact that gradient descent satisfies the energy dissipation property, it is in general not true that as time goes to infinity the dynamics 2.2 converge to a global minimizer of 2.1. This could happen for different reasons. First, the problem 2.1 may not have a minimizer (e.g., take $V(\theta ) = e^{-\theta }$ for $\theta \in \mathbb{R}$). Second, $\nabla V$ may have critical points associated with saddle points or local optima of $V$ as we illustrate in the next example.

Example 2.2.

Consider the double well potential

$$\begin{equation} V(\theta ) = \frac{3}{8}\theta ^4 - \frac{3}{4}\theta ^2, \quad \theta \in \mathbb{R}, {\tag{2.6}} \end{equation}$$

so that $\nabla V(\theta ) = \frac{3}{2} \theta (\theta ^2 -1).$ Notice that $\theta = 0$ is an unstable equilibrium of 2.2, which corresponds to a local maximum of $V.$ For each local minima $\theta = \pm 1$ of $V$ there is an associated subregion in the space of parameters (known as a basin of attraction) such that any initial condition $\theta _0$ chosen in this subregion leads the gradient dynamics toward its corresponding local minimizer, see Figure 1.

Figure 1.

Five trajectories of the one-dimensional gradient system 2.2 with double well potential 2.6. The objective $V$ has critical points at $\theta = 0$ and $\theta = \pm 1$. The intervals $(0,\infty )$ and $(-\infty ,0)$ are basins of attraction for $\theta = 1$ and $\theta = -1,$ respectively.

Suitable assumptions on $V$ prevent the existence of local minimizers that are not global and also imply rates of convergence. One such assumption is the Polyak–Lojasiewic (PL) condition KNS16:

$$\begin{equation} \alpha ( V(\theta ) - V^*) \leq \frac{1}{2}|\nabla V (\theta )|^2 , \quad \forall \theta \in \Theta , {\tag{2.7}} \end{equation}$$

where $V^\ast = \inf _{\theta \in \Theta } V(\theta )$ and $\alpha > 0$ is a constant. Note that the PL condition readily implies that any stationary point $\theta ^\ast$ of $V$ is a global minimizer, since

$$\begin{equation*} \alpha (V (\theta ^\ast ) - V^\ast ) \le \frac{1}{2} | \nabla V(\theta ^\ast ) |^2 =0. \end{equation*}$$

Under the PL condition we can easily obtain a convergence rate for continuous-time gradient descent.

Proposition 2.3.

Suppose $V$ satisfies 2.7. Then,

$$\begin{equation} V(\theta _t) - V^* \leq ( V(\theta _0) - V^* ) \exp (- 2 \alpha t ), \quad \forall t \ge 0. {\tag{2.8}} \end{equation}$$

Proof.

Using condition 2.7 in equation 2.5 and recalling 2.4 we conclude that

$$\begin{equation*} V(\theta _t) - V^* \leq (V(\theta _0)- V^*) - 2 \alpha \int _{0}^t ( V(\theta _r) - V^* )dr. \end{equation*}$$

The result follows by Gronwall’s inequality.

One can verify the PL condition under various assumptions on the function $V$ KNS16. Here we present, as an important example, the case of $\alpha$-strong convexity. For $\alpha >0,$ one says that $V$ is $\alpha$-strongly convex if for any $\theta , \theta ' \in \Theta$ it holds that

$$\begin{equation*} V \bigl ( t \theta +(1-t) \theta ' \bigr ) \leq t V(\theta ) + (1-t) V(\theta ') - \frac{\alpha }{2} t(1-t) |\theta - \theta '|^2, \end{equation*}$$

for all $t \in [0,1]$. This condition can be shown to be equivalent to

$$\begin{equation*} V(\theta ' ) \ge V(\theta ) + \langle \nabla V(\theta ),\theta ' - \theta \rangle + \frac{\alpha }{2} | \theta ' - \theta |^2, \quad \forall \theta , \theta ' \in \Theta , \end{equation*}$$

from which we can see, after minimizing both sides w.r.t. $\theta '$, that

$$\begin{equation} V^\ast \ge V(\theta ) - \frac{1}{2\alpha } | \nabla V(\theta )|^2, {\tag{2.9}} \end{equation}$$

which is equivalent to 2.7. From this we conclude that $\alpha$-strong convexity implies the PL condition with the same constant $\alpha$.

Note that strong convexity is a stronger condition than the PL condition. For example, the function $V(\theta )=\frac{1}{2}\theta _1^2$ (where $\theta = (\theta _1, \theta _2)$) satisfies the PL condition with $\alpha =1$, but it is not strongly convex.

2.3. Choice of the metric

Let us consider a function of the form $V(\theta ) = \alpha _1 \theta _1^2 + \alpha _2 \theta _2^2,$ where $\theta = (\theta _1,\theta _2)\in \mathbb{R}^2.$ Suppose that $0 <\alpha _1 \ll \alpha _2$ and that $\alpha _1$ is very close to zero, as in Figure 2. We can now apply Proposition 2.3 with $\alpha =\alpha _1$, but since we assumed $\alpha _1$ is small we see that the right-hand side of 2.8 decreases very slowly. This suggests that gradient descent may take a long time to reach $V$’s global minimum when initialized arbitrarily.

Figure 2.

Level curves of a two-dimensional potential of the form $V(\theta ) = \alpha _1 \theta _1^2 + \alpha _2 \theta _2^2$ with $0 <\alpha _1 \ll \alpha _2$. The anisotropy of this potential causes the gradient system 2.2 to converge slowly.

The poor behavior of gradient descent described above arises whenever there are regions of points away from the minimizer at which the gradient of $V$ is very small. One approach to remedy this issue is to introduce a more general version of gradient descent that accelerates the dynamics in those regions where the gradient of $V$ is small. This is the goal of preconditioning. Let $H: \Theta \rightarrow \mathcal{S}^d_{++}$ be a continuous field of positive definite matrices, i.e., a function that assigns to every point $\theta \in \Theta$ a $d\times d$ positive definite matrix $H(\theta )$. The preconditioned gradient descent dynamics induced by $H$ is defined as:

$$\begin{equation} \dot{\theta }_t = - H(\theta _t)^{-1} \nabla V(\theta _t), \quad t >0. {\tag{2.10}} \end{equation}$$

Observe that 2.10 coincides with the original gradient descent dynamics 2.2 when $H$ is constant and equal to the $d\times d$ identity matrix. On the other hand, when $V$ is convex and twice differentiable, choosing $H(\theta ) = \nabla ^2 V(\theta )$, the Hessian of $V$, results in a continuous-time analog of Newton’s algorithm. In the example in Figure 2, we can directly compute $\nabla ^2 V=\left(\begin{matrix} 2 \alpha _1 & 0 \\ 0 & 2 \alpha _2 \end{matrix}\right)$, i.e., the Hessian is a fixed matrix since the potential $V$ is quadratic. Substituting this choice of $H$ in 2.10 for that example gives the dynamics $\dot{\theta }_t = -\theta _t$, a scheme that achieves a much faster convergence rate.

The reader may wonder if we could have in fact chosen $H = \frac{1}{r} D^2 V$ for large constant $r$ in order to induce a system that converges to equilibrium at a faster rate. However, as implied by our discussion in section 2.4 and, specifically, Remark 2.6 there is no benefit in doing so, as the cost of discretizing becomes correspondingly higher; rescaling the Hessian may be simply interpreted as a change of units. In general, there is a natural tension between accelerating continuous-time dynamics by changing the metric of the space, and producing accurate time discretizations for the resulting flows; see GTSA20 for a related discussion in the context of sampling algorithms. In a similar vein, we notice that the superior convergence rate and affine invariance of Newton’s algorithm comes at the price of utilizing the Hessian matrix $\nabla ^2 V$, which in many applications can be prohibitively costly to compute or store. To this end, constructing matrix fields $H(\theta )$ that are good proxies for the Hessian and that can be computed efficiently is the goal of preconditioning. Perhaps the most well-known family of such algorithms is the family of Quasi-Newton algorithms NW99 and in particular the Broyden–Fletcher–Goldfarb–Shanno (BFGS) algorithm, which approximates the Hessian using gradients calculated at previous iterates.

2.3.1. Geometric interpretation

As a step toward introducing the material in section 3, here we give a geometric interpretation of equation 2.10. Specifically, we will show that 2.10 can still be understood as a gradient descent equation but w.r.t. a different metric on the parameter space $\Theta$. For this purpose it is convenient to recall that a Riemannian manifold $(\mathcal{M}, g)$ is a manifold $\mathcal{M}$ endowed with a family of inner products $g=\{ g_{\theta } \}_{\theta \in \mathcal{M}}$ (often referred to as the metric), one for each point on the manifold, and which can be used to measure angles between vectors at every tangent plane of $\mathcal{M}$. We will use $\mathcal{T}_\theta \mathcal{M}$ to denote the tangent plane at a given $\theta \in \mathcal{M}$. The standard example of a Riemannian manifold is $\mathcal{M}= \mathbb{R}^d$ with $g_\theta$ the Euclidean inner product at every point. More general examples of Riemannian manifolds with $\mathcal{M}= \mathbb{R}^d$ can be generated from a field of positive definite matrices $H$. Consider the family of inner products:

$$\begin{equation*} g_\theta (u,v) \coloneq \bigl \langle H(\theta ) u, v \bigr \rangle , \quad u, v \in \mathbb{R}^d, \end{equation*}$$

where $\langle \cdot , \cdot \rangle$ denotes the standard Euclidean inner product. Notice that $g_\theta$ is indeed an inner product since $H(\theta )$ is positive definite. In what follows we often suppress the dependence of $g$ on $\theta$ for brevity.

We now proceed to define the notion of the gradient of a function $F$ defined over an arbitrary Riemannian manifold. Let $(\mathcal{M}, g)$ be a Riemannian manifold and let $F: \mathcal{M}\rightarrow \mathbb{R}$ be a smooth enough function. The gradient of $F$ at the point $\theta \in \mathcal{M}$ relative to the metric $g$, denoted $\nabla _g F (\theta ),$ is defined as the vector in $\mathcal{T}_\theta \mathcal{M}$ for which the following identity holds:

$$\begin{equation} \frac{d}{dt} F \bigl (\gamma (t)\bigr ) \big |_{t=0} = g\bigl (\nabla _g F(\theta ), \dot{\gamma }(0)\bigr ) {\tag{2.11}} \end{equation}$$

for any differentiable curve $\gamma :(-\varepsilon , \varepsilon ) \rightarrow \mathcal{M}$ with $\gamma (0)=\theta$; by $\dot{\gamma }(0)$ we mean the velocity of the curve $\gamma$ at time $0$, which is an element in $\mathcal{T}_\theta \mathcal{M}$. In the example of $\mathcal{M}=\mathbb{R}^d$ with inner products induced by a field $H$ (which we denote with $g_H$), we have that

$$\begin{align*} \frac{d}{dt} V \bigl (\gamma (t)\bigr ) \big |_{t=0} & = \bigl \langle \nabla V (\theta ) , \dot{\gamma }(0) \bigr \rangle \\ & = \bigl \langle H(\theta ) H(\theta )^{-1} \nabla V (\theta ) , \dot{\gamma }(0) \bigr \rangle \\ & = g \bigl ( H(\theta )^{-1} \nabla V (\theta ) , \dot{\gamma }(0) \bigr ), \end{align*}$$

for any curve $\gamma : (-\varepsilon , \varepsilon ) \rightarrow \mathcal{M}$ with $\gamma (0)=\theta$, from where it follows that $\nabla _{g_H} V (\theta ) = H(\theta )^{-1} \nabla V (\theta )$, where we recall $\nabla$ denotes the usual gradient in $\mathbb{R}^d$.

In this light, 2.10 can be interpreted as a gradient descent algorithm, only that the gradient is taken w.r.t. a metric that is different from the standard Euclidean one. In section 3, where we discuss sampling, we will return to some of the insights that we have developed in this section. In particular, in order to define gradient descent dynamics of a functional over a manifold we need to specify two ingredients: 1) an energy $V$ to optimize, and 2) a metric $g_\theta$ under which we define the gradient. For the last item, it will be convenient to have a clear understanding of how to represent smooth curves in the manifold of interest and characterize their velocities appropriately. Indeed, equation 2.11 explicitly relates the metric $g_\theta$ of the manifold, the target energy $V$, the rate of change of the energy along arbitrary smooth curves $\gamma$, and the gradient $\nabla _g V$ of the energy relative to the chosen metric.

2.3.2. Geodesic convexity

There are analogous conditions to the PL and strong convexity assumptions discussed in section 2.1 that guarantee the convergence of the flow 2.10 toward global minima of $V$. First, write 2.10 as an energy dissipation equality of the form

$$\begin{equation} V(\theta _t) = V(\theta _s) - \frac{1}{2}\int _{s}^t |\nabla _g V (\theta _r)|_{\theta _r}^2 dr - \frac{1}{2} \int _{s}^t |\dot{\theta }_r|_{\theta _r}^2 dr, {\tag{2.12}} \end{equation}$$

where we have used $|\cdot |_\theta ^2$ to denote $g_\theta (\cdot , \cdot )$. The equivalence between 2.10 and 2.12 follows from an identical argument as in Proposition 2.1 applied to an arbitrary inner product. The analogous PL condition in the preconditioned setting takes the form:

$$\begin{equation} \alpha (V (\theta ) - V^\ast ) \le \frac{1}{2} | \nabla _g V(\theta ) |_\theta ^2, {\tag{2.13}} \end{equation}$$

which generalizes the PL condition in the Euclidean setting to general inner products and gradients.

To introduce an appropriate notion of convexity that allows us to generalize the results of section 2.1 we need to introduce a few more ideas from Riemannian geometry. Given a Riemannian manifold $(\mathcal{M},g)$, we define the geodesic distance $d_g$ induced by the metric $g$ as:

$$\begin{equation} d_g^2(\theta , \theta ') = \inf _{t \in [0,1] \mapsto (\gamma _t, \dot{\gamma }_t) } \int _{0}^1 |\dot{\gamma }_t|_{\gamma _t}^2 dt. {\tag{2.14}} \end{equation}$$

We will say that $\gamma : [0,1] \rightarrow \mathcal{M}$ is a constant speed geodesic between $\theta$ and $\theta '$ if $\gamma$ is a minimizer of the right-hand side of the above expression. Equivalently, a constant speed geodesic between $\theta$ and $\theta '$ is any curve with $\gamma (0)=\theta$ and