# Error bounds for particle gradient descent,
and extensions of the log-Sobolev and Talagrand inequalities

Rocco Caprio222Department of Statistics, University of Warwick. Email:
{rocco.caprio<EMAIL_ADDRESS>Juan Kuntz333Polygeist. Email:
<EMAIL_ADDRESS>Samuel Power444School of Mathematics, University of
Bristol. Email<EMAIL_ADDRESS>Adam M. Johansen222Department of
Statistics, University of Warwick. Email: {rocco.caprio,
<EMAIL_ADDRESS>
###### Abstract

We prove non-asymptotic error bounds for particle gradient descent (PGD) [23],
a recently introduced algorithm for maximum likelihood estimation of large
latent variable models obtained by discretizing a gradient flow of the free
energy. We begin by showing that, for models satisfying a condition
generalizing both the log-Sobolev and the Polyak–Łojasiewicz inequalities (LSI
and PŁI, respectively), the flow converges exponentially fast to the set of
minimizers of the free energy. We achieve this by extending a result well-
known in the optimal transport literature (that the LSI implies the Talagrand
inequality) and its counterpart in the optimization literature (that the PŁI
implies the so-called quadratic growth condition), and applying it to our new
setting. We also generalize the Bakry–Émery Theorem and show that the LSI/PŁI
generalization holds for models with strongly concave log-likelihoods. For
such models, we further control PGD’s discretization error, obtaining non-
asymptotic error bounds. While we are motivated by the study of PGD, we
believe that the inequalities and results we extend may be of independent
interest.

Keywords: latent variable models, maximum marginal likelihood, gradient flows,
log-Sobolev inequality, Polyak–Łojasiewicz inequality, Talagrand inequality,
quadratic growth condition.

##  1 Introduction

Many tasks in machine learning and statistics require the fitting of a
probabilistic model, with Lebesgue density $p_{\theta}(x,y)$, featuring latent
variables, $x\in\mathbb{R}^{d_{x}}$, to data, $y$, which we have observed.
Often, we achieve this by finding model parameters,
$\theta\in\mathbb{R}^{d_{\theta}}$, that maximize the probability,
$p_{\theta}(y)$, of observing the data we observed (the _marginal
likelihood_): that is, $\theta_{*}$ belonging to

$\mathcal{O}_{*}\mathrel{\mathop{\mathchar
58\relax}}=\operatorname*{arg\,max}_{\theta\in\mathbb{R}^{d_{\theta}}}p_{\theta}(y)=\operatorname*{arg\,max}_{\theta\in\mathbb{R}^{d_{\theta}}}\int
p_{\theta}(x,y)\operatorname{d\\!}x.$ (1)

(We assume throughout that the latent variables and parameters respectively
take values in $\mathbb{R}^{d_{\theta}}$ and $\mathbb{R}^{d_{x}}$.) In many
cases, the latent variables are meaningful or interesting in some way, and we
would like to infer them. Following the empirical Bayes paradigm [33], we do
this using the posterior distribution of the latent variables given parameter
vector $\theta_{*}$:

$p_{\theta_{*}}(x|y)\mathrel{\mathop{\mathchar
58\relax}}=\frac{p_{\theta_{*}}(x,y)}{p_{\theta_{*}}(y)}.$

For most models of practical interest, the integral in (1) is intractable, we
have no closed-form expressions for $p_{\theta}(y)$ or its derivatives, and we
are unable to directly optimize $p_{\theta}(y)$. Often, this hurdle is
overcome by noting that $(\theta_{*},q_{*})$ minimizes the _free energy_
functional,

$F(\theta,q)\mathrel{\mathop{\mathchar
58\relax}}=\left\\{\begin{array}[]{cl}\displaystyle\int\log\left(\frac{q(x)}{p_{\theta}(x,y)}\right)q(\operatorname{d\\!}x)&\text{if
}q\ll\operatorname{d\\!}x\\\
+\infty&\text{otherwise}\end{array}\right.\quad\forall(\theta,q)\in\mathcal{M},$
(2)

if and only if $\theta_{*}$ maximizes the marginal likelihood and
$q_{*}=p_{\theta_{*}}(\cdot\mid y)$; where $\operatorname{d\\!}x$ in (2)
denotes the Lebesgue measure and $\mathcal{M}$ the product of the parameter
space $\mathbb{R}^{d_{\theta}}$ and the space
$\mathcal{P}(\mathbb{R}^{d_{x}})$ probability distributions over the latent
space $\mathbb{R}^{d_{x}}$ with its Euclidean Borel $\sigma$-algebra. For
instance, the well-known expectation-maximization (EM) algorithm [13] can be
viewed as minimizing $F$ using coordinate descent (CD) [27]. Similarly, many
methods in variational inference (approximately) minimize $F$ by (i)
restricting $\mathcal{P}(\mathbb{R}^{d_{x}})$ to a parametric family of
distributions $(q_{\phi})_{\phi\in\Phi}$ such that the surrogate objective
$(\theta,\phi)\mapsto F(\theta,q_{\phi})$ is tractable, (ii) computing a
minimizer $(\theta_{*},\phi_{*})$ thereof using an appropriate optimization
algorithm, and (iii) employing $(\theta_{*},q_{\phi_{*}})$ as a proxy for a
genuine minimizer of $F$; e.g., see [22].

Inspired by the interpretation of the EM algorithm as CD applied to $F$, [23]
examines whether it is possible to minimize $F$ over $\mathcal{M}$, rather
than any restriction thereof, using analogues of optimization algorithms other
than CD. In particular, they identify analogues of gradient descent (GD)
applicable to $F$. To do so, they recall that the application of GD to
minimizing a function $f\mathrel{\mathop{\mathchar
58\relax}}\mathbb{R}^{d_{\theta}}\to\mathbb{R}$ may be viewed as the forward
Euler discretization of the following ordinary differential equation (ODE)
known as the _gradient flow_ :

$\dot{\theta}_{t}=-\nabla_{\theta}f(\theta_{t}),$ (3)

where $\nabla_{\theta}$ denotes the usual Euclidean gradient on
$\mathbb{R}^{d_{\theta}}$. Borrowing ideas from optimal transport, they
identify $\nabla F=(\nabla_{\theta}F,\nabla_{q}F)$, with

$\nabla_{\theta}F(\theta,q)=-\int\nabla_{\theta}\ell(\theta,x)q(\operatorname{d\\!}x),\quad\nabla_{q}F(\theta,q)=\nabla_{x}\cdot\left[q\nabla_{x}\log\left(\frac{p_{\theta}(\cdot,y)}{q}\right)\right],$
(4)

and $\ell(\theta,x)\mathrel{\mathop{\mathchar
58\relax}}=\log(p_{\theta}(x,y))$, as an analogue to $\nabla_{\theta}f$ and
obtain the following gradient flow dynamics for the parameter estimate
$\theta_{t}$ and the latent posterior approximation $q_{t}$:

$\displaystyle\dot{\theta}_{t}=\int\nabla_{\theta}\ell(\theta_{t},x)q_{t}(\operatorname{d\\!}x),\qquad\dot{q}_{t}=\nabla_{x}\cdot\bigg{[}q_{t}\nabla_{x}\log\bigg{(}\frac{q_{t}}{p_{\theta_{t}}(\cdot|y)}\bigg{)}\bigg{]}.$
(5)

They go on to (i) interpret (5) as the Fokker-Planck equation for the
following McKean-Vlasov stochastic differential equation (SDE; e.g., see [8])

$\displaystyle\operatorname{d\\!}\theta_{t}=\int\nabla_{\theta}\ell(\theta_{t},x)q_{t}(\operatorname{d\\!}x)\operatorname{d\\!}t,\quad\operatorname{d\\!}X_{t}=\nabla_{x}\ell(\theta_{t},X_{t})\operatorname{d\\!}t+\sqrt{2}\operatorname{d\\!}W_{t},\hskip
3.00003pt\text{where}\enskip q_{t}\mathrel{\mathop{\mathchar
58\relax}}=\textrm{Law}(X_{t});$ (6)

and (ii) to approximate the measure $q_{t}$ with an empirical measure
$N^{-1}\sum_{n=1}^{N}\delta_{X^{n}}$ of _particles_
$\smash{X^{1},\ldots,X^{N}}$; and (iii) discretize the resulting SDE in time.
The result is a practical algorithm for fitting latent variable models to data
they refer to as particle gradient descent (PGD; Algorithm 1).

1: Inputs: step size $h$, step number $K$, particle number $N$, and initial
particles $X^{1,h}_{0},\dots,$ $X_{0}^{N,h}$ and parameter $\Theta_{0}^{N,h}$.

2: for $k=0,\dots,K-1$ do

3: Update the parameter estimate:
$\Theta_{k+1}^{N,h}=\Theta_{k}^{N,h}+\frac{h}{N}\sum_{n=1}^{N}\nabla_{\theta}\ell(\Theta_{k}^{N,h},X_{k}^{n,h}).$
(7)

4: Update the particles: with $W_{k}^{1},\dots,W_{k}^{N}$ denoting i.i.d.
$\mathcal{N}(0,I_{d_{x}})$ r.v.s,
$X_{k+1}^{n,h}=X_{k}^{n,h}+h\nabla_{x}\ell(\Theta_{k}^{N,h},X_{k}^{n,h})+\sqrt{2h}W_{k}^{n}\quad\forall
n=1,\dots,N.$ (8)

5: end for

6: return $\Theta_{K}^{N,h}$ and $Q^{N,h}_{K}\mathrel{\mathop{\mathchar
58\relax}}=N^{-1}\sum_{n=1}^{N}\delta_{X_{K}^{n,h}}$.

Algorithm 1 Particle gradient descent (PGD).

PGD performs well in experiments [23] and is well-suited to modern computing
environments. Its updates (7,8) require only evaluating
$\nabla\ell\mathrel{\mathop{\mathchar
58\relax}}=(\nabla_{\theta}\ell,\nabla_{x}\ell)$ and adding noise, thus
facilitating its implementation in modern autodiff frameworks such as
TensorFlow, PyTorch, and JAX. Moreover, each particle’s update (8) is
independent of all other particles, offering scope for simple parallelisation
and distribution of these steps. A full run of the algorithm can be naively
implemented at a computational cost scaling as $\mathcal{O}(NK[$evaluation
cost of $\nabla\ell])$, where $N$ denotes the number of particles employed and
$K$ the total number of steps taken; the dependence of running time on $N$ can
be further attenuated in practice through the aforementioned parallelism
strategies.

Here, we validate theoretically PGD and show that, for models whose log-
likelihoods $\ell$ are $\lambda$-strongly concave with $L$-Lipschitz
gradients, its output, $(\Theta_{K}^{N,h},Q^{N,h}_{K})$, satisfies

$\mathsf{d}((\Theta_{K}^{N,h},Q^{N,h}_{K}),(\theta_{*},Q_{*}^{N}))=\mathcal{O}(h^{\frac{1}{2}}+N^{-\frac{1}{2}}+e^{-h\lambda
K}),$ (9)

assuming that the algorithm’s step size $h$ is no greater than
$1/(\lambda+L)$. In the above, $Q_{*}^{N}$ denotes the empirical measure of
$N$ i.i.d. particles drawn from $p_{\theta_{*}}(\cdot|y)$ and $\mathsf{d}$ the
following metric acting on $\mathcal{M}$-valued random variables:

$\mathsf{d}((\Theta,Q),(\Theta^{\prime},Q^{\prime}))\mathrel{\mathop{\mathchar
58\relax}}=\sqrt{\mathbb{E}\left[\mathsf{d}_{E}(\Theta,\Theta^{\prime})^{2}\right]+\mathbb{E}\left[\mathsf{d}_{W_{2}}(Q,Q^{\prime})^{2}\right]}$
(10)

where $\mathsf{d}_{E}$ and $\mathsf{d}_{W_{2}}$ denote the Euclidean and
Wasserstein-2 [38, Chapter 6] metrics, resp.

We begin by studying PGD’s continuous-time infinite-particle limit: the
gradient flow (5). First, we show that the free energy
$(F(\theta_{t},q_{t}))_{t\geq 0}$ evaluated along the flow converges at an
exponential rate to its infimum,

$F_{*}\mathrel{\mathop{\mathchar
58\relax}}=\inf_{(\theta,q)\in\mathcal{M}}F(\theta,q),$

provided that the model satisfies a condition which generalizes the log-
Sobolev inequality (LSI) popular in optimal transport (see e.g. [38, Chapter
21]) and the Polyak–Łojasiewicz inequality (PŁI) used in optimization (see
e.g. [21]). In particular, we assume that there exists a constant $\lambda>0$
such that

$2\lambda[F(\theta,q)-F_{*}]\leq I(\theta,q)$ (11)

for all parameters $\theta$ and distributions $q$, where the functional $I$ is
defined by

$\displaystyle I(\theta,q)\mathrel{\mathop{\mathchar 58\relax}}=$
$\displaystyle\mathinner{\\!\left\lVert\int\nabla_{\theta}\ell(\theta,x)q(\operatorname{d\\!}x)\right\rVert}^{2}+\int\mathinner{\\!\left\lVert\nabla_{x}\log\left(\frac{q(x)}{p_{\theta}(x,y)}\right)\right\rVert}^{2}q(\operatorname{d\\!}x),$
(12)

with $\mathinner{\\!\left\lVert\cdot\right\rVert}$ here denoting the Euclidean
norm. We then extend a result of [31] and show that models which satisfy (11)
also satisfy the following generalization of both the Talagrand inequality
[36] and the _quadratic growth_ condition used in optimization [3, 21]: for
all parameters $\theta$ and distributions $q$, it holds that

$2[F(\theta,q)-F_{*}]\geq\lambda\mathsf{d}((\theta,q),\mathcal{M}_{*})^{2},$
(13)

where $\mathcal{M}_{*}$ denotes the $F$’s optimal set and
$\mathsf{d}((\theta,q),\mathcal{M}_{*})$ the distance from $(\theta,q)$ to
$\mathcal{M}_{*}$:

$\mathcal{M}_{*}\mathrel{\mathop{\mathchar
58\relax}}=\operatorname*{arg\,min}_{(\theta,q)\in\mathcal{M}}F(\theta,q),\quad\mathsf{d}((\theta,q),\mathcal{M}_{*})\mathrel{\mathop{\mathchar
58\relax}}=\inf_{(\theta_{*},q_{*})\in\mathcal{M_{*}}}\mathsf{d}((\theta,q),(\theta_{*},q_{*}).$
(14)

Using (13), we then show that the flow itself $(\theta_{t},q_{t})_{t\geq 0}$
converges to $\mathcal{M}_{*}$ exponentially fast in $\mathsf{d}$. Next, we
generalize the Bakry–Émery Theorem [4] and show that models with strongly
concave likelihoods $\ell$ satisfy (11). For such models, we obtain a slightly
stronger convergence result. Under the further assumption that $\ell$’s
gradients are Lipschitz, we are able to bound the errors introduced by both
approximating distributions with empirical measures and discretizing (5) in
time, and the bound (9) follows.

###  1.1 Related literature

The analysis provided in [23] was limited: as explained in the ‘Our setting,
notation, assumptions, rigour, and lack thereof’ paragraph of the paper’s
introduction, it focused on intuition and practical application rather than
theoretical validation. While the authors obtained no error bounds; they did
argue with their Theorem 3 that the gradient flow (5) converges exponentially
fast to $F$’s minimizer whenever $\ell$ is strongly concave. However, they
imposed an unnecessary technical condition that substantially restricts the
applicability of their result (that the gradient of $\ell$ be bounded
uniformly).

Over the last year’s course, four other works have proposed particle-based
algorithms for maximum marginal likelihood estimation: [1, 24, 9, 35]. [1]
consider an alternative approximation to the flow (5) which they term the
‘Interacting Particle Langevin Algorithm’ or IPLA. IPLA differs from PGD by
the inclusion of an extra additive Gaussian noise term $\sqrt{2h/N}W_{k}$ in
the parameter’s update (7). This difference enables an analysis of the
algorithm using tools similar to those used to study the unadjusted Langevin
algorithm in [17, 16]. In particular, the authors show that the
$\mathcal{L}^{2}$ error of the parameter estimates produced by IPLA satisfies
(9, RHS) for models satisfying the same Lipschitz gradient and strong
concavity assumptions considered in the present work, which we find pleasantly
consistent with some relevant results in Section 3. [24, 9] consider variants
of PGD which incorporate momentum in the updates, which can outperform PGD in
practice. These variants can no longer be viewed as approximations to the
gradient flow (5); and consequently, require a different type of analysis.
Some was conducted in [24], which established exponentially fast convergence
of their algorithm’s limiting (non-gradient) flow under (11). [35] introduced
two new particle-based methods for minimizing $F$. The first, ‘Stein
Variational Gradient Descent EM’ or SVGD EM, is also constructed as an
approximation to a gradient flow of $F$, but with respect to a different
underlying geometry—the so-called ‘Stein Geometry’, as opposed to the
Euclidean/Wasserstein-2 geometry underpinning (4) (whose induced distance
metric is given by (10) omitting the expectations; cf. [23, Appendix A] for
more on this). The second, ‘Coin EM’, involves an entirely different approach
for minimizing $F$ that builds on coin betting techniques from convex
optimization and obviates the need for tuning discretization step sizes. Only
SVGD-EM provides a theoretical analysis, which was conducted under a natural
‘Stein’ analogue of (11): Assumption 7 in [35, Appendix A]. Let us also
mention the recent work of [18], who considers, within an similar Empirical
Bayes framework, a non-parametric maximum likelihood estimation problem in
which they treat the prior as an unknown measure, arriving at a gradient flow
that is somehow related in spirit to (5) but qualitatively different. Our
results in Section 2 relate to earlier results well-known in the optimal
transport and optimization literatures. We discuss these connections as we go
along.

###  1.2 Paper structure

In Section 2, we study the convergence of the gradient flow (5) and the
inequalities (11,13). In Section 3, we analyze the flow’s approximation (PGD;
Algorithm 1) and obtain non-asymptotic bounds on its error. We conclude in
Section 4 by discussing our results beyond the context of PGD.

##  2 Convergence of the gradient flow

In this section, we study the gradient flow (5) and several pertinent
inequalities. As we explain in what follows, these inequalities generalize
others well-known in the literature, and we prefix the former with ‘extended’
to differentiate them from the latter. We start by showing that for models
satisfying the extended LSI (xLSI) (11), the values of $F$ along the flow
converge exponentially fast to the infimum of $F$. We next argue that any
model satisfying the xLSI also satisfies the extended Talagrand-type
inequality (13), and we use this result to show that the trajectory of the
flow itself converges exponentially fast to
$(\theta_{*},p_{\theta_{*}}(\cdot|y))$ for some maximizer $\theta_{*}$ of the
marginal likelihood (Section 2.2). Finally, we show that the xLSI holds for
models with $\lambda$-strongly concave log-likelihoods $\ell$, and we obtain a
slightly stronger convergence result for this case (Section 2.3).

#### Notation and assumptions.

For the remainder of the paper, we shorten our notation by writing
$\rho_{\theta}(\cdot)$ for the unnormalized density $p_{\theta}(\cdot,y)$,
$Z_{\theta}$ for its normalizing constant, $p_{\theta}(y)$, and
$\pi_{\theta}(\cdot)$ for the normalised density $p_{\theta}(\cdot\mid y)$. To
ensure that $\mathsf{d}$ in (10) is a metric, we restrict
$\mathcal{P}(\mathbb{R}^{d_{x}})$ to its subset of measures with finite second
moments and similarly for $\mathcal{M}$:

$\mathcal{P}_{2}(\mathbb{R}^{d_{x}})\mathrel{\mathop{\mathchar
58\relax}}=\bigg{\\{}q\in\mathcal{P}(\mathbb{R}^{d_{x}})\mathrel{\mathop{\mathchar
58\relax}}\int\mathinner{\\!\left\lVert
x\right\rVert}^{2}q(\operatorname{d\\!}x)<\infty\bigg{\\}},\quad\mathcal{M}_{2}\mathrel{\mathop{\mathchar
58\relax}}=\mathbb{R}^{d_{\theta}}\times\mathcal{P}_{2}(\mathbb{R}^{d_{x}}).$

To ensure that the functional $I$ in (12) is well-defined, we further restrict
restrict $\mathcal{P}_{2}(\mathbb{R}^{d_{x}})$ its subset of measures with
a.e. differentiable densities w.r.t. to the Lebesgue measure
$\operatorname{d\\!}x$ and similarly for $\mathcal{M}_{2}$:

$\displaystyle\mathcal{P}_{2}^{1}(\mathbb{R}^{d_{x}})$
$\displaystyle\mathrel{\mathop{\mathchar
58\relax}}=\bigg{\\{}q\in\mathcal{P}_{2}(\mathbb{R}^{d_{x}})\mathrel{\mathop{\mathchar
58\relax}}q\ll\operatorname{d\\!}x,\,\nabla_{x}\frac{\operatorname{d\\!}q}{\operatorname{d\\!}x}(x)\text{
exists a.e.}\bigg{\\}},\quad\mathcal{M}_{2}^{1}\mathrel{\mathop{\mathchar
58\relax}}=\mathbb{R}^{d_{\theta}}\times\mathcal{P}_{2}^{1}(\mathbb{R}^{d_{x}}).$

For the same reasons, we assume that the model has finite second moments and a
positive a.e. differentiable density:

######  Assumption 1 (Model regularity).

(i) For all $x$ in $\mathbb{R}^{d_{x}}$, $\theta\mapsto\pi_{\theta}(x)$ is
differentiable; and $\theta\mapsto Z_{\theta}$ is differentiable; (ii) for all
$\theta$ in $\mathbb{R}^{d_{\theta}}$, $\pi_{\theta}$ belongs to
$\mathcal{P}_{2}^{1}(\mathbb{R}^{d_{x}})$; (iii) for all $\theta$ in
$\mathbb{R}^{d_{\theta}}$ and $x$ in $\mathbb{R}^{d_{x}}$,
$\rho_{\theta}(x)>0$;

We also assume that the gradient flow (5) has solutions in $\mathcal{M}_{2}$:

######  Assumption 2 (Existence of solutions).

For any initial conditions $(\theta,q)$ in $\mathcal{M}_{2}$, (5) has a weak
solution $(\theta_{t},q_{t})_{t\geq 0}$ with $(\theta_{0},q_{0})=(\theta,q)$.
By a weak solution, we mean a function mapping from $[0,\infty)$ to
$\mathcal{M}_{2}$ such that for all $t\geq 0$ and twice continuously
differentiable and compactly supported functions
$\phi\mathrel{\mathop{\mathchar 58\relax}}\mathbb{R}^{d_{x}}\to\mathbb{R}$,

$\displaystyle\frac{\operatorname{d\\!}\theta_{t}}{\operatorname{d\\!}t}$
$\displaystyle=\int\nabla_{\theta}\ell(\theta_{t},x)q_{t}(\operatorname{d\\!}x),$
$\displaystyle\frac{\operatorname{d\\!}}{\operatorname{d\\!}t}\int\phi(x)q_{t}(\operatorname{d\\!}x)$
$\displaystyle=\int\Delta_{x}\phi(x)q_{t}(\operatorname{d\\!}x)-\int\left\langle{\nabla_{x}\ell(\theta_{t},x)},{\nabla_{x}\phi(x)}\right\rangle
q_{t}(\operatorname{d\\!}x);$

where $\left\langle{\cdot},{\cdot}\right\rangle$ denotes the Euclidean inner
product and $\Delta_{x}$ the Laplacian.

We verify in Section 3 that the above holds whenever the model’s log-
likelihood has a Lipschitz gradient. Lastly, we assume that the solutions are
sufficiently regular:

######  Assumption 3 (Regularity of solutions).

Let $(\theta_{t},q_{t})_{t\geq 0}$ be any weak solution of (5). Then, for all
$t>0$, $(\theta_{t},q_{t})\in\mathcal{M}_{2}^{1}$ and

$\frac{\operatorname{d\\!}F}{\operatorname{d\\!}t}(\theta_{t},q_{t})=-I(\theta_{t},q_{t})\quad\forall
t>0.$ (15)

Formally arguing (15) is straightforward (cf. Appendix A.1). In fact, we
believe that the above assumption holds under regularity conditions on the
model of the sort in Assumption 1. However, rigorously arguing this is
technically involved and lies beyond the scope of the paper. In the special
case that $\pi_{\theta}$ is a single distribution $\pi$ independent of
$\theta$, the flow (5) essentially reduces to the following well-known Fokker-
Planck equation,

$\dot{q}_{t}=\nabla_{x}\cdot\bigg{[}q_{t}\nabla_{x}\log\bigg{(}\frac{q_{t}}{\pi}\bigg{)}\bigg{]},$
(16)

whose solutions’ regularity properties are well-understood (e.g., [20, Theorem
5.1] shows that the solutions are smooth whenever $\log(\pi)$ is smooth). In
this case, (15) reduces to de Bruijn’s identity which is known to hold
whenever $\pi$ is sufficiently regular; e.g. see [5, Proposition 5.2.2] or
[31, Lemma 1].

###  2.1 The extended log-Sobolev inequality and $(F(\theta_{t},q_{t}))_{t\geq
0}$’s convergence

We prove the flow’s convergence for models which satisfy the ‘xLSI’ in the
following sense:

######  Definition 1 (Extended log-Sobolev inequality; xLSI).

The measures
$(\rho_{\theta}(\operatorname{d\\!}x))_{\theta\in\mathbb{R}^{d_{\theta}}}$
satisfy the extended log-Sobolev inequality with constant $\lambda>0$ if (11)
holds for all $(\theta,q)$ in $\mathcal{M}_{2}^{1}$.

Under this assumption, the values of $F$ converge along the gradient flow
exponentially fast to $F$’s infimum over $\mathcal{M}_{2}$:

######  Theorem 1 (xLSI $\Rightarrow$ exp. conv. of
$(F(\theta_{t},q_{t}))_{t\geq 0}$).

If Assumptions 1–3 hold, $(\theta_{0},q_{0})$ belongs to $\mathcal{M}_{2}$,
and the measures
$(\rho_{\theta}(\operatorname{d\\!}x))_{\theta\in\mathbb{R}^{d_{\theta}}}$
satisfy the xLSI with constant $\lambda$ (Definition 1), then

$0\leq F(\theta_{t},q_{t})-F_{*}\leq[F(\theta_{0},q_{0})-F_{*}]e^{-2\lambda
t}\quad\forall t\geq 0.$ (17)

###### Proof.

The leftmost inequality in (17) follows from the definition of $F_{*}$. The
rightmost equality follows by combining the xLSI (11) with de Bruijn’s
identity (15) and applying Grönwall’s inequality. ∎

As is well-known,

$F(\theta,q)=\mathrm{KL}(q||\pi_{\theta})-\log(Z_{\theta})\quad\forall(\theta,q)\in\mathcal{M},$
(18)

where KL denotes the Kullback-Leibler divergence:

$\mathrm{KL}(q||\pi)\mathrel{\mathop{\mathchar
58\relax}}=\left\\{\begin{array}[]{cl}\displaystyle\int\log\left(\frac{q(x)}{\pi(x)}\right)q(\operatorname{d\\!}x)&\textup{if
}q\ll\operatorname{d\\!}x\\\
+\infty&\textup{otherwise}\end{array}\right.\quad\forall
q\in\mathcal{P}(\mathbb{R}^{d_{x}}),$

Because $\mathrm{KL}(q||q^{\prime})$ is non-negative and zero iff
$q=q^{\prime}$ a.e., (18) implies that

$\displaystyle
F_{*}=\inf_{(\theta,q)\in\mathcal{M}}F(\theta,q)=\inf_{\theta\in\mathbb{R}^{d_{\theta}}}F(\theta,\pi_{\theta})=-\log\left(\sup_{\theta\in\mathbb{R}^{d_{\theta}}}Z_{\theta}\right)=-\log
Z_{*},$ (19)

where $Z_{*}\mathrel{\mathop{\mathchar
58\relax}}=\sup_{\theta\in\mathbb{R}^{d_{\theta}}}Z_{\theta}$ denotes the
marginal likelihood’s supremum. Putting (17–19) together we find that

$\log(Z_{*})-\log(Z_{\theta_{t}})+\mathrm{KL}(q_{t}||\pi_{\theta_{t}})=F(\theta_{t},q_{t})+\log(Z_{*})=\mathcal{O}(e^{-2\lambda
t}).$

In other words, Theorem 1 shows that as $t$ increases: the free energy
$F(\theta_{t},q_{t})$ converges exponentially fast to $-\log Z_{*}$, the log-
marginal likelihood $\log(Z_{\theta_{t}})$ converges exponentially fast to its
supremum $\log(Z_{*})$, and $q_{t}$ tracks the corresponding posterior
distributions $\pi_{\theta_{t}}$ in the sense that the $\mathrm{KL}$
divergence between the two decays exponentially fast to zero.

#### Connections with the log-Sobolev inequality.

We refer to the inequality in Definition 1 as the ‘xLSI’ because it extends
the classical log-Sobolev inequality (LSI); e.g., see [38, Definition 21.1].
While the xLSI is a statement about a parametrized family of measures, the LSI
is a statement about a single probability measure. In particular, if
$\pi_{\theta}$ is a distribution $\pi$ independent of $\theta$, then (11)
reads

$2\lambda\mathrm{KL}(q||\pi)\leq I(q||\pi)\enskip\forall
q\in\mathcal{P}_{2}^{1}(\mathbb{R}^{d_{x}});\enskip\text{where}\enskip
I(q||\pi)\mathrel{\mathop{\mathchar
58\relax}}=\int\mathinner{\\!\left\lVert\nabla_{x}\log\left(\frac{q(x)}{\pi(x)}\right)\right\rVert}^{2}q(\operatorname{d\\!}x).$
(20)

For this reason, Theorem 1 extends a well-known result (e.g., [5, Section
5.2]) stating that, if $\pi$ satisfies the LSI (20) and $(q_{t})_{t\geq 0}$
solves the Fokker-Planck equation (16), then the $\mathrm{KL}$ divergence from
$q_{t}$ to $\pi$ decays exponentially fast.

#### Connections with the Polyak–Łojasiewicz inequality.

Definition 1 also extends an inequality due to Polyak [32] and Łojasiewicz
[25] commonly used to argue linear convergence of gradient descent algorithms.
A differentiable real-valued function $f$ on $\mathbb{R}^{d_{\theta}}$ with
infimum $f_{*}$ is said to satisfy the PŁI with constant $\lambda>0$ if

$2\lambda[f(\theta)-f_{*}]\leq\mathinner{\\!\left\lVert\nabla_{\theta}f(\theta)\right\rVert}^{2}\quad\theta\in\mathbb{R}^{d_{\theta}}.$
(21)

In particular, because

$\int\nabla_{\theta}\ell(\theta,x)\pi_{\theta}(\operatorname{d\\!}x)=\int\frac{\nabla_{\theta}\rho_{\theta}(x)}{\rho_{\theta}(x)}\pi_{\theta}(\operatorname{d\\!}x)=\frac{\nabla_{\theta}\int\rho_{\theta}(x)\operatorname{d\\!}x}{Z_{\theta}}=\frac{\nabla_{\theta}Z_{\theta}}{Z_{\theta}}=\nabla_{\theta}\log(Z_{\theta}),$
(22)

setting $q\mathrel{\mathop{\mathchar 58\relax}}=\pi_{\theta}$ in (11), we
recover (21) with $f(\theta)\mathrel{\mathop{\mathchar
58\relax}}=-\log(Z_{\theta})$. (The exchange of limits in (22) is valid
whenever Fisher’s identity holds, which is a common assumption in the study of
latent variable models and the EM algorithm; it is satisfied whenever e.g.
$\pi_{\theta}$ belongs to an exponential family [15, Appendix D].) In this
case, Theorem 1 reduces to a well-known result showing that $f$ converges
exponentially fast along the Euclidean gradient flow (3) to its infimum
whenever $f$ satisfies the PŁI; e.g., see [37, Proposition 2.3].

Among other reasons, the PŁI (21) is popular in optimization because—without
requiring the objective to be convex—it implies that the objective’s
stationary points are global minimizers. In our case, the xLSI implies that
the marginal likelihood’s stationary points are global maximizers.

###  2.2 The Talagrand inequality and flow’s exponential convergence

We now show that for models satisfying the xLSI, the gradient flow
$(\theta_{t},q_{t})_{t\geq 0}$ itself converges exponentially fast to the set
of (global) minimizers of the free energy. We measure convergence in terms of
the following metric on $\mathcal{M}_{2}$:

$\displaystyle\mathsf{d}((\theta,q),(\theta^{\prime},q^{\prime}))\mathrel{\mathop{\mathchar
58\relax}}=\sqrt{\mathsf{d}_{E}(\theta,\theta^{\prime})^{2}+\mathsf{d}_{W_{2}}(q,q^{\prime})^{2}}.$
(23)

Since $F$ is a lower semicontinuous (l.s.c.) function on
$(\mathcal{M}_{2},\mathsf{d})$ (see Lemma 15 in Appendix A.2), the set
$\mathcal{M}_{*}$ in (14) of the free energy’s minima is closed (note that
Assumption 1 and (19) imply that $\mathcal{M}_{*}$ is contained in
$\mathcal{M}_{2}$). Using (18,19), we can characterize $\mathcal{M}_{*}$ in
terms of the marginal likelihood’s set of maximizers ($\mathcal{O}_{*}$ in
(1)):

$\mathcal{M}_{*}=\\{(\theta_{*},\pi_{\theta_{*}})\mathrel{\mathop{\mathchar
58\relax}}\theta_{*}\in\mathcal{O}_{*}\\}.$ (24)

To argue the flow’s convergence, we require the following extension of the
Talagrand inequality:

######  Definition 2 (Extension of the Talagrand inequality; xT2I).

The measures
$(\rho_{\theta}(\operatorname{d\\!}x))_{\theta\in\mathbb{R}^{d_{\theta}}}$
satisfy an extension of the Talagrand inequality with constant $\lambda>0$ if
(13) holds for all $(\theta,q)$ in $\mathcal{M}_{2}$.

The inequality holds for all models satisfying the xLSI:

######  Theorem 2 (xLSI $\Rightarrow$ xT2I).

If Assumptions 1–3 hold, and the measures
$(\rho_{\theta}(\operatorname{d\\!}x))_{\theta\in\mathbb{R}^{d_{\theta}}}$
satisfy the xLSI with constant $\lambda$, then they also satisfy the xT2I with
the same constant.

See Appendix A.2 for the proof. The convergence of $(\theta_{t},q_{t})_{t\geq
0}$ then follows from Theorem 1:

######  Corollary 3 (xLSI $\Rightarrow$ exp. conv. of
$(\theta_{t},q_{t})_{t\geq 0}$).

If Assumptions 1–3 hold $(\theta_{0},q_{0})$ belongs to $\mathcal{M}_{2}$, and
the measures
$(\rho_{\theta}(\operatorname{d\\!}x))_{\theta\in\mathbb{R}^{d_{\theta}}}$
satisfy the xLSI with constant $\lambda$, then

$\lambda\mathsf{d}((\theta_{t},q_{t}),\mathcal{M}_{*})^{2}\leq
2[F(\theta_{0},q_{0})-F_{*}]e^{-2\lambda t}\quad\forall t\geq 0.$

#### Connections with the Talagrand inequality.

By (24), if $\pi_{\theta}$ is a distribution $\pi$ independent of $\theta$,
then $\mathcal{M}_{*}=\mathbb{R}^{d_{\theta}}\times\\{\pi\\}$ and (13) reduces
to the Talagrand inequality (‘T2I’) [36, Theorem 1.2]; e.g. see [38]):

$2\mathrm{KL}(q||\pi)\geq\lambda\mathsf{d}_{W_{2}}(q,\pi)^{2}\quad\forall
q\in\mathcal{P}_{2}(\mathbb{R}^{d_{x}}).$ (25)

For this reason, Theorem 2 extends the Otto-Villani Theorem ([31, Theorem 1])
showing that the LSI (20) implies T2I (25).

#### Connections with the quadratic growth condition.

Setting $(q,f(\theta),f_{*})$ to $(\pi_{\theta},$ $-\log(Z_{\theta}),F_{*})$,
(13) reduces to the ‘quadratic growth’ (QG) condition used in optimization to
establish linear convergence rates to local minima of gradient descent
algorithms:

$2[f(\theta)-f_{*}]\geq\lambda\mathsf{d}_{E}(\theta,\mathcal{O}_{*})^{2}\quad\forall\theta\in\mathbb{R}^{d_{\theta}}.$
(26)

In this case, Theorem 2 reduces to the well-known result stating that the QG
condition holds whenever the PŁI (21) holds; see [21, Theorem 2].

###  2.3 Strongly log-concave models

In Section 3, we study PGD’s convergence for models satisfying the following
assumption.

######  Assumption 4 (Strong log-concavity).

There exists a $\lambda>0$ such that

$\ell((1-t)\theta+t\theta^{\prime},(1-t)x+tx^{\prime})\geq(1-t)\ell(\theta,x)+t\ell(\theta^{\prime},x^{\prime})+\frac{\lambda
t(1-t)}{2}\mathinner{\\!\left\lVert(\theta,x)-(\theta^{\prime},x^{\prime})\right\rVert}^{2},$

for all $(\theta,x),(\theta^{\prime},x^{\prime})$ in
$\mathbb{R}^{d_{\theta}}\times\mathbb{R}^{d_{x}}$ and $0\leq t\leq 1$.

Models that satisfy the above also satisfy the xLSI:

######  Theorem 4 (Strong log-concavity $\Rightarrow$ xLSI).

Any model satisfying Assumptions 1(iii) and 4 satisfies the xLSI.

###### Proof.

The result follows from the fact that the free energy is strongly geodesically
convex if log-likelihood is strongly concave; see Appendix A.3 for details. ∎

#### Connection with the Bakry–Émery Theorem.

Given that (11) reduces to the standard LSI (20) whenever $\pi_{\theta}$ is
independent of $\theta$, Theorem 4 extends Bakry–Émery Theorem [4] showing
that distributions with strongly log-concave densities satisfy (20).

#### Connection with the optimization literature.

Given that (11) with $q=\pi_{\theta}$ reduces to the PŁI (21) with
$f(\theta)=-\log(Z_{\theta})$, Theorem 4 extends the well-known fact that the
PŁI holds for strongly convex $f$; e.g., see [21, Theorem 2].

#### The flow’s convergence.

Under Assumption 4, the marginal likelihood has a unique maximizer
$\theta_{*}$; e.g., see [23, Theorem 4 in Appendix B.3]. Hence, Corollary 3
and (24) tell us that the flow converges exponentially fast to
$(\theta_{*},\pi_{\theta_{*}})$:

$\displaystyle\lambda\mathsf{d}((\theta_{t},q_{t}),(\theta_{*},\pi_{\theta_{*}}))^{2}\leq
2[F(\theta_{0},q_{0})-F_{*}]e^{-2\lambda t}\quad\forall t\geq 0.$

In this special case, we can sharpen the above slightly (see Appendix A.4 for
the proof):

######  Theorem 5 (Strong log-concavity $\Rightarrow$ exp. conv. of
$(\theta_{t},q_{t})_{t\geq 0}$).

If Assumptions 1(iii), 2–4 hold and $(\theta_{0},q_{0})$ belongs to
$\mathcal{M}_{2}$, then

$\displaystyle\mathsf{d}((\theta_{t},q_{t}),(\theta_{*},\pi_{\theta_{*}}))\leq\mathsf{d}((\theta_{0},q_{0}),(\theta_{*},\pi_{\theta_{*}}))e^{-\lambda
t}\quad\forall t\geq 0.$

##  3 Error bounds for particle gradient descent

Here, we capitalize on the results of Section 2 to obtain the error bound in
(9) for PGD (Algorithm 1) applied to models with strongly concave log-
likelihoods. In particular, Theorem 5 in Section 2.3 bounds the distance
between the gradient flow and $(\theta_{*},\pi_{\theta_{*}})$, where
$\theta_{*}$ denotes the marginal likelihood’s unique maximizer and the
distance is measured in terms of the metric $\mathsf{d}$ in (23). To exploit
this bound and obtain (9), we first need to connect the gradient flow to PGD,
which we do via the McKean-Vlasov SDE (6). To ensure that the SDE has
globally-defined solutions, we assume that $\ell$’s gradient is Lipschitz:

######  Assumption 5.

The log-likelihood $\ell$ is differentiable and its gradient
$\nabla\ell=(\nabla_{\theta}\ell,\nabla_{x}\ell)$ is $L$-Lipschitz for some
$L>0$:

$\mathinner{\\!\left\lVert\nabla\ell(\theta,x)-\nabla\ell(\theta^{\prime},x^{\prime})\right\rVert}\leq
L\mathinner{\\!\left\lVert(\theta,x)-(\theta^{\prime},x^{\prime})\right\rVert}\quad\forall(\theta,x),(\theta^{\prime},x^{\prime})\in\mathbb{R}^{d_{\theta}}\times\mathbb{R}^{d_{x}}.$

Let $\mathcal{L}^{2}(\mathbb{R}^{d_{x}})$ denote the space of square-
integrable random variables taking values in $\mathbb{R}^{d_{x}}$. The SDE’s
relation to the gradient flow is as follows (see Appendix B.1 for the proof):

######  Proposition 6.

If Assumptions 4–5 hold and $(\theta,X)$ belongs to
$\mathbb{R}^{d_{\theta}}\times\mathcal{L}^{2}(\mathbb{R}^{d_{x}})$, then the
SDE (6) has an unique strong solution $(\theta_{t},X_{t})_{t\geq 0}$ such that
$(\theta_{0},X_{0})=(\theta,X)$. Moreover,
$(\theta_{t},\textup{Law}(X_{t}))_{t\geq 0}$ is a weak solution of the
gradient flow (5) in the sense of Assumption 2.

PGD, on the other hand, is obtained by approximating (6) as we detail in
Section 3.1 below. Bounding the errors introduced by these approximations as
we do there, we obtain (9). In particular, recall that
$(\Theta_{K}^{N,h},Q_{K}^{N,h})$ denotes the PGD’s output after $K$ iterations
and $Q^{N}_{*}$ the empirical distribution of $N$ i.i.d. particles drawn from
$\pi_{\theta_{*}}$. Moreover, note that, the log-likelihood $\ell$ has a
unique maximizer $(\theta_{\dagger},x_{\dagger})$ which, without loss of
generality, we assume lies at the origin $(0,0)$. Our main result for PGD is
then as follows:

######  Theorem 7 (PGD error bound).

Under Assumptions 1(iii) and 3–5, if $X_{0}^{1,h},\dots,X_{0}^{N,h}$ are drawn
independently from a distribution $q_{0}$ in
$\mathcal{P}_{2}(\mathbb{R}^{d_{x}})$ and $h\leq 1/(\lambda+L)$, then

$\displaystyle\mathsf{d}((\Theta_{K}^{N,h},Q_{K}^{N,h}),(\theta_{*},Q_{*}^{N}))\leq\sqrt{h}A_{0,h}+\frac{L\sqrt{2}}{\lambda\sqrt{N}}\sqrt{B_{0}+\frac{d_{x}}{\lambda}}+\mathsf{d}((\theta_{0},q_{0}),(\theta_{*},\pi_{\theta_{*}}))e^{-h\lambda
K}$ (27)

for all $K$ in $\mathbb{N}$; where $B_{0}\mathrel{\mathop{\mathchar
58\relax}}=\mathinner{\\!\left\lVert\theta_{0}\right\rVert}^{2}+\mathbb{E}\left[\mathinner{\\!\left\lVert
X_{0}\right\rVert}^{2}\right]$ and

$\displaystyle A_{0,h}\mathrel{\mathop{\mathchar 58\relax}}=$
$\displaystyle\sqrt{\frac{4h+4/\iota}{\iota}220L^{2}\left(L^{2}h\left[B_{0}+\frac{d_{x}}{\lambda}\right]+d_{x}\right)}\quad\text{with}\quad\iota\mathrel{\mathop{\mathchar
58\relax}}=\frac{2L\lambda}{L+\lambda}.$

We postpone the proof until Section 3.1 below. In short, the Theorem 7 gives
us the bound (9) and an explicit expression for the proportionality constant
as a function of the algorithm’s parameters, the model’s Lipschitz and
concavity constants, the initial condition, and the dimension of the latent
space. It follows that in order to achieve an error of
$\mathcal{O}(\epsilon)$, we need only choose the step size $h$ proportional to
$\epsilon^{2}$, the particle number $N$ proportional to $\epsilon^{-2}$, and
the step number $K$ proportional to $\epsilon^{-2}\log(\epsilon^{-1})$. The
corresponding computational cost is
$\mathcal{O}(\epsilon^{-4}\log(\epsilon^{-1})[\text{eval. cost of }\ell])$;
and we can lower the running time to
$\mathcal{O}(\epsilon^{-2}\log(\epsilon^{-1})[\text{eval. cost of }\ell])$ by
parallelizing/distributing computations across the particles (recall that each
particle’s update is independent of the other particles’).

Given the definition of the metric $\mathsf{d}$ in (10), Theorem 7 bounds the
$\mathcal{L}^{2}$ error of the parameter estimates $\Theta_{K}^{N,h}$ and,
also, the expected $W_{2}$ distance between the particles’ empirical
distribution $Q_{K}^{N,h}$ and that, $Q_{*}^{N}$, of $N$ i.i.d. samples drawn
from the optimal posterior $\pi_{\theta_{*}}$. We find the latter insightful
because, in many ways, $Q_{*}^{N}$ is the best one could hope to achieve using
an $\mathcal{O}(N)$-cost Monte Carlo algorithm that returns $N$-sample
approximations $\pi_{\theta_{*}}$. The expected $W_{2}$ distance between
$Q^{N}_{*}$ itself and the optimal posterior $\pi_{\theta_{*}}$ is essentially
a property of $\pi_{\theta_{*}}$ and has more to do with the fundamental
limitations of approximating $\pi_{\theta_{*}}$ using $N$-sample ensembles
than the method used to obtain the ensemble. To bound this distance (and, by
combining this bound with Theorem 7, to obtain bounds on the distance between
$Q_{K}^{N,h}$ and $\pi_{\theta_{*}}$), we refer the reader to [19, 14] and
references therein.

#### The error bound under warm-starts.

If we ‘warm start’ PGD (see, e.g., [12]) by initializing the parameter
estimates and the particles at the log-likelihood $\ell$ maxima (i.e., with
$\theta_{0}=0$ and $q_{0}=\delta_{0}$), $B_{0}$ in (27) vanishes and we obtain
the following sharper bound.

######  Corollary 8 (Error bound under warm starts).

Under the conditions in Theorem 7’s premise and with $\iota$ as defined
therein, if $(\theta_{0},q_{0})\mathrel{\mathop{\mathchar
58\relax}}=(0,\delta_{0})$, then, for all $K$ in $\mathbb{N}$,

$\displaystyle\mathsf{d}((\Theta_{K}^{N,h},Q_{K}^{N,h}),(\theta_{*},Q_{*}^{N}))\leq\sqrt{h\frac{4h+4/\iota}{\iota}220L^{2}\left(\frac{L^{2}hd_{x}}{\lambda}+d_{x}\right)}+\sqrt{\frac{d_{x}}{\lambda}}\left(\frac{L\sqrt{2}}{\lambda\sqrt{N}}+e^{-h\lambda
K}\right).$

###### Proof.

Combining Theorem 5 together with Proposition 19 in Appendix B.1, we find that

$\mathsf{d}((\theta_{0},q_{0}),(\theta_{*},\pi_{\theta_{*}}))^{2}=\lim_{t\rightarrow\infty}\mathsf{d}((\theta_{0},q_{0}),(\theta_{t},q_{t}))^{2}\leq\mathinner{\\!\left\lVert\theta_{t}\right\rVert}^{2}+\mathbb{E}\left[\mathinner{\\!\left\lVert
X_{t}\right\rVert}^{2}\right]\leq\frac{2d_{x}}{\lambda}.$

∎

This result shows that PGD, with a warm start, achieves an
$\mathcal{O}(\epsilon)$-error as long as step size $h$ is
$\mathcal{O}(\epsilon^{2}d_{x}^{-1})$, the number of particles is
$\mathcal{O}(\epsilon^{-2}d_{x})$ and the number of steps is
$\mathcal{O}(\epsilon^{-2}d_{x}\log(d_{x}^{-1/2}\epsilon^{-1}))$. The
corresponding cost is
$\mathcal{O}(\epsilon^{-4}d_{x}^{2}\log(d_{x}^{-1/2}\epsilon^{-1})[\text{eval.
cost of }\ell])$ and the running time can be brought down to
$\mathcal{O}(\epsilon^{-2}d_{x}\log(d_{x}^{-1/2}\epsilon^{-1})[\text{eval.
cost of }\ell])$ by parallelizing/distributing over particles. In this
context, starting PGD warm is reasonably practical because the cost of running
PGD will typically significantly exceed that of maximizing $\ell$ using a
first-order method. These computations rely on the implicit assumption
widespread in the literature that $L$ and $\lambda$ are themselves independent
of dimension. This is often not the case leading to a further, hidden
dependence upon dimension. However, as we now explore, it is possible to
obtain dimension-free bounds for a class of models in which the observations
are conditionally independent given the parameters.

#### Dimension-free bounds for models with independent latent variables and
conditionally Gaussian observations.

We consider here a simple setting in order to demonstrate that conditional
independence of observations given latent variables can significantly improve
the bounds obtained above. Often, in machine learning applications, the data
$y$ decomposes as a sequence of datapoints $(\tilde{y}_{m})_{m=1}^{M}$, the
model features a latent variable $\tilde{x}_{m}$ per datapoint
$\tilde{y}_{m}$, the pairs $(\tilde{x}_{m},\tilde{y}_{m})_{m=1}^{M}$ are
i.i.d. and are generated through a mechanism of the sort

$\tilde{y}_{m}=f_{\theta}(\tilde{x}_{m})+\eta_{m},\quad\tilde{x}_{m}\sim
p_{\theta}^{x}(\operatorname{d\\!}\tilde{x}),\quad\eta_{m}\sim
p^{\eta}(\operatorname{d\\!}\eta),\quad\forall m=1,\dots,M;$

where $f_{\theta}$ denotes a parametrized function mapping from the latent
space $\mathbb{R}^{\tilde{d_{x}}}$ to the parameter space, $p_{\theta}^{x}$
denotes a distribution with a closed-form density (for example, Gaussian
distributions), and the noise has a Gaussian law, $p^{\eta}$, which is
independent of the model parameters $\theta$. In such cases, the log-
likelihood breaks down into a sum of per-datapoint log-likelihoods:

$\ell(\theta,x)=\sum_{m=1}^{M}\tilde{\ell}(\theta,\tilde{x}_{m};\tilde{y}_{m}),$
(28)

where

$\tilde{\ell}(\theta,\tilde{x}_{m};\tilde{y}_{m})\mathrel{\mathop{\mathchar
58\relax}}=\log p^{\eta}(\tilde{y}_{m}-f_{\theta}(\tilde{x}_{m}))+\log
p_{\theta}^{x}(\tilde{x}_{m})\quad\forall m=1,\dots,M.$ (29)

Because $p^{\eta}$ is Gaussian,
$(\tilde{x},\theta)\mapsto\tilde{\ell}(\tilde{x},\theta;\tilde{y})$’s Hessian
$\nabla^{2}_{(\theta,\tilde{x})}\tilde{\ell}$ does not depend on $\tilde{y}$.
Consequently, if $\tilde{\ell}$ is $\tilde{\lambda}$-strongly concave and its
gradient is $\tilde{L}$-Lispchitz,

$\tilde{\lambda}I\preceq-\nabla^{2}_{(\theta,\tilde{x})}\tilde{\ell}\preceq\tilde{L}I$
(30)

for some positive $\tilde{\lambda}$ and $\tilde{L}$, then $\ell$ is
$(M\tilde{\lambda})$-strongly concave and has a $(M\tilde{L})$-Lipschitz
gradient. Combining this fact with Corollary 8, we obtain the following:

######  Corollary 9.

Let the conditions in Theorem 7’s premise hold. Suppose that
$(\theta_{0},q_{0})=(0,\delta_{0})$, $\ell$ is as in (28) with $\tilde{l}$
twice-differentiable and satisfying (30), and $\epsilon>0$. If

$h\leq\frac{C_{1}\epsilon^{2}}{\tilde{d}_{x}M},\qquad
N\geq\frac{C_{2}\tilde{d}_{x}}{\epsilon^{2}},\qquad
K\geq\frac{C_{3}\tilde{d}_{x}}{\epsilon^{2}}\log\left(\frac{1}{\tilde{d}_{x}^{1/2}\epsilon}\right),$

for some constants $C_{1},C_{2},C_{3}>0$ independent of $M$, then

$\mathsf{d}((\Theta_{K}^{N,h},Q_{K}^{N,h}),(\theta_{*},Q_{*}^{N}))\leq
C\epsilon,$

for another constant $C>0$ independent of $M$.

The corresponding cost is
$\mathcal{O}(\epsilon^{-4}\log(M\epsilon^{-1})[\text{eval. cost of
}\tilde{\ell}])$, where we are ignoring any dependence on $\tilde{d}_{x}$
which is usually small for these models. In particular, the cost depends on
$M$ only through the evaluation cost of $\ell$, which cannot be avoided. We
can, however, mitigate this by parellizing/distributing over both particles
and datapoints, so lowering the running time to
$\mathcal{O}(\epsilon^{-2}\log(\epsilon^{-1})[\text{eval. cost of
}\tilde{\ell}])$.

The corollary provides some theoretical support for the idea that the PGD
algorithm and its derivatives can work well for problems in which $d_{x}$ is
large, even when $N$ is relatively modest, provided that the structure of the
model allows the log-likelihood to be decomposed as a sum of functions of low-
dimensional contributions arising from individual datapoints. This appeared to
be the case in several examples explored in [23, Section 3], including a
‘generator network’ model for which $d_{x}\approx 2.5\times 10^{6}$ but good
performance was obtained with only $10$ particles. The model, however, lies
beyond the scope of Corollary 9 because, in that case, $f_{\theta}$ in (29) is
a non-convex neural network and $\theta$ encompasses its weights. Moreover,
[23]’s authors adapted PGD’s step sizes to account for this and also
subsampled $\ell$’s gradient to facilitate efficient implementation on a
single GPU. Regardless, we view the corollary as a tentative first step in
explaining the behaviour observed in [23].

###  3.1 PGD’s derivation and Theorem 7’s proof

We obtain PGD by approximating the gradient flow (5) via the McKean-Vlasov SDE
(6). The first step is to approximate the law, $q_{t}$, of $X_{t}$ in (6) with
the empirical distribution $Q_{t}^{N}\mathrel{\mathop{\mathchar
58\relax}}=N^{-1}\sum_{n=1}^{N}\delta_{X_{t}^{n}}$ of $N$ i.i.d. copies
$X_{t}^{1},\dots,X_{t}^{N}$ of $X_{t}$. To obtain these copies, consider

$\displaystyle\operatorname{d\\!}\theta_{t}=\frac{1}{N}\sum_{n=1}^{N}\int\nabla_{\theta}\ell(\theta_{t},x)q_{t}^{n}(\operatorname{d\\!}x)\operatorname{d\\!}t,\quad\operatorname{d\\!}X^{n}_{t}=\nabla_{x}\ell(\theta_{t},X^{n}_{t})\operatorname{d\\!}t+\sqrt{2}\operatorname{d\\!}W_{t}^{n},\quad\forall
n\in[N],$ (31)

where $[N]\mathrel{\mathop{\mathchar 58\relax}}=\\{1,\dots,N\\}$,
$W^{1},\dots,W^{N}$ denote $N$ independent Brownian motions,
$q_{t}^{n}\mathrel{\mathop{\mathchar 58\relax}}=\text{Law}(X_{t}^{n})$ for
each $n$, and we assume that the initial condition $X_{0}^{n}$ of each
particle $(X_{t}^{n})_{t\geq 0}$ is drawn independently from the law $q_{0}$
of $(X_{t})_{t\geq 0}$’s initial condition $X_{0}$. Because the particles all
satisfy the same differential equation and have the same initial distribution,
they share the same law (in particular, $q_{t}^{n}$ is independent of $n$).
Hence, for any given $n$, we can re-write the leftmost equation in (31) as
$\operatorname{d\\!}\theta_{t}=\int\nabla_{\theta}\ell(\theta_{t},x)q_{t}^{n}(\operatorname{d\\!}x)\operatorname{d\\!}t$.
Putting this equation together with that for $X_{t}^{n}$ in (31), we recover
(6) with $X^{n}_{t}$ replacing $X_{t}$ therein. In other words, if
$(\theta_{t},X^{1}_{t},\dots,X^{N}_{t})_{t\geq 0}$ solves (31), then
$(\theta_{t},X_{t}^{n})_{t\geq 0}$ solves (6) for any $n$. Consequently,
Proposition 6 tells us that, for any $n$, $(\theta_{t},q_{t}^{n})_{t\geq 0}$
is a weak solution to the gradient flow (5). Exploiting this link and Theorem
5, we can bound the distance between $(\theta_{t},Q_{t}^{N})$ and
$(\theta_{*},Q_{*}^{N})$, where $Q_{*}^{N}\mathrel{\mathop{\mathchar
58\relax}}=N^{-1}\sum_{n=1}^{N}\delta_{X_{*}^{n}}$ and
$X_{*}^{1},\dots,X_{*}^{N}$ denotes $N$ i.i.d. particles with law
$\pi_{\theta*}$:

######  Lemma 10.

Under Assumptions 1(iii) and 3–5, if $q_{0}$ belongs to
$\mathcal{P}_{2}(\mathbb{R}^{d_{x}})$, then (31) has a strong solution
$(\theta_{t},X_{t}^{1},\dots,X_{t}^{N})_{t\geq 0}$ and the solution satisfies

$\mathsf{d}((\theta_{t},Q_{t}^{N}),(\theta_{*},Q_{*}^{N}))\leq\mathsf{d}((\theta_{0},q_{0}),(\theta_{*},\pi_{\theta_{*}}))e^{-\lambda
t}\quad\forall t\geq 0.$

###### Proof.

The argument for the existence of solutions for (31) is similar to that for
(6) in Proposition 6’s proof, and we skip it. The error bound follows by
optimally coupling the two sets of particles; see Appendix B.2 for details. ∎

The next step entails approximating
$N^{-1}\sum_{n=1}^{N}q_{t}^{n}(\operatorname{d\\!}x)$ in (31) with the
particle’s empirical distribution $Q_{t}^{N}$; which leads us to the following
SDE:

$\displaystyle\operatorname{d\\!}\Theta_{t}^{N}=$
$\displaystyle\frac{1}{N}\sum_{n=1}^{N}\nabla_{\theta}\ell(\Theta_{t}^{N},\bar{X}_{t}^{n})\operatorname{d\\!}t,$
$\displaystyle\operatorname{d\\!}\bar{X}_{t}^{n}=$
$\displaystyle\nabla_{x}\ell(\Theta_{t}^{N},X_{t}^{n})\operatorname{d\\!}t+\sqrt{2}\operatorname{d\\!}W_{t}^{n},\quad\quad\forall
n\in[N].$ (32)

where $\bar{X}_{0}^{n}=X_{0}^{n}$ for each $n$ in $[N]$. We can bound the
approximation error between (31,32) as follows:

######  Lemma 11 (Bound on the space discretization error).

Under Assumptions 4–5, if $q_{0}$ belongs to
$\mathcal{P}_{2}(\mathbb{R}^{d_{x}})$, then (32) has a strong solution
$(\Theta_{t}^{N},\bar{X}_{t}^{1},\dots,\bar{X}_{t}^{N})_{t\geq 0}$ and the
solution satisfies

$\mathsf{d}((\Theta^{N}_{t},\bar{Q}^{N}_{t}),(\theta_{t},Q^{N}_{t}))\leq\frac{L\sqrt{2}}{\lambda\sqrt{N}}\sqrt{\mathinner{\\!\left\lVert\theta_{0}\right\rVert}^{2}+\mathbb{E}\left[\mathinner{\\!\left\lVert
X_{0}\right\rVert}^{2}\right]+\frac{d_{x}}{\lambda}}\enskip\forall t\geq
0,\,\text{with}\enskip\bar{Q}_{t}^{N}\mathrel{\mathop{\mathchar
58\relax}}=\frac{1}{N}\sum_{n=1}^{N}\delta_{\bar{X}_{t}^{n}}.$

###### Proof.

(32) is a standard SDE with a Lipschitz drift and diffusion coefficients, so
the existence of its solutions is routine (e.g. [30, Theorem 5.2.1]). The
error bound follows from a synchronous coupling argument; see Appendix B.3 for
details. ∎

Lastly, we obtain PGD by discretizing (32) in time using the Euler-Maruyama
scheme with step size $h$. Given Lemmas 10 and 11, Theorem 7 follows from the
triangle inequality for $\mathsf{d}$ and the following bound on the
discretization error:

######  Lemma 12 (Bound on the time discretization error).

Under Assumptions 4–5. If $q_{0}$ belongs to
$\mathcal{P}_{2}(\mathbb{R}^{d_{x}})$, $h\leq 1/(\lambda+L)$, and $A_{0,h}$ is
as in Theorem 7, then:

$\displaystyle\mathsf{d}((\Theta^{N,h}_{K},Q^{N,h}_{K}),(\Theta^{N}_{Kh},\bar{Q}^{N}_{Kh}))\leq\sqrt{h}A_{0,h}.$

###### Proof.

The bound follows from a synchronous coupling argument; see Appendix B.4. ∎

##  4 Discussion

In this paper, we theoretically validated PGD (Algorithm 1) by bounding its
error for models with strongly concave log-likelihoods. In order to do this we
analyzed the theoretical properties of both PGD and its continuous-time
infinite-particle limit: the gradient flow (5). To study the latter, we
extended certain inequalities well-known in the optimal transport and
optimization literatures and several results interlinking them. In doing so,
we obtain conditions weaker than strong concavity of the likelihood under
which the flow converges at an exponential rate to the set $\mathcal{M}_{*}$
of pairs $(\theta_{*},\pi_{\theta_{*}})$ of marginal likelihood maximizers
$\theta_{*}$ and matching posterior distributions $\pi_{\theta_{*}}$. Under
the additional assumptions that the model log-likelihood is strongly concave
and its gradient is Lipschitz continuous, we then proved explicit non-
asymptotic error bounds for PGD.

Aside from evidencing PGD’s well-foundedness, our bounds serve as a starting
point for theoretical comparisons of PGD with other maximum marginal
likelihood algorithms. For instance, [1] analyze the IPLA algorithm which can
be viewed as a modification of PGD more amenable to standard analysis and
obtain very similar scalings in $\epsilon,d_{x},d_{\theta}$ for $h,N,$ and $K$
necessary to achieve an error of $\epsilon^{-1}$ under warm-start conditions
(see Section 3.6 therein) to those obtained here. The main difference between
these bounds is that $N$ must be proportional to $d_{x}\epsilon^{-2}$ in our
case and proportional to $d_{\theta}\epsilon^{-2}$ in the case of IPLA. Their
dependence on $d_{\theta}$ stems from the extra noise term featuring IPLA’s
parameter updates; see Algorithm 1 and the proof of Proposition 3 in [1]. We
conjecture that the dependence on $d_{x}$ present in our results but not those
of [1] simply reflects the fact that our bounds controlling the error of both
the parameter estimates and the particle approximation, while theirs only
control the former.

Similarly, our bounds could help settle the question of whether it is possible
to ‘accelerate’ PGD and achieve faster convergence rates via the use of
momentum [24] as its the case for both gradient descent [28] and the
unadjusted Langevin algorithm [26]. In particular, [24, Theorem 4.1] shows
that, under the xLSI, the free energy $F$ along the limiting flow of such a
‘momentum-enriched’ version of PGD converges exponentially fast to $F$’s
infimum. Given Theorem 2, applying our extension of the Talagrand inequality,
we can immediately prove convergence in $\mathsf{d}$-distance for the
momentum-enriched PGD studied in [24]. By following arguments similar to those
in Section 3.1, it might be possible to then also obtain similar results for
the momentum-enriched PGD and compare the algorithms’ performances in various
settings.

While our motivation for this work was understanding the theoretical
properties of PGD, we find the results in Section 2 interesting in their own
right and we believe they might find applications beyond the study of
particle-based algorithms like PGD. For instance, for models with
$\lambda$-strongly concave log-likelihoods, Theorem 4 gives us a (to the best
of our knowledge) novel upper bound on the optimal marginal log-likelihood:

$\log(Z_{*})\leq\frac{I(\theta,q)}{2\lambda}-F(\theta,q)\quad\forall(\theta,q)\in\mathcal{M}_{2}^{1}.$
(33)

Similarly, for models satisfying the xLSI, Theorem 2 yields a bound on the
$\mathsf{d}$-distance between $(\theta,q)$ and the optimal set
$\mathcal{M}_{*}$:

$\lambda\mathsf{d}((\theta,q),\mathcal{M}_{*})\leq\sqrt{I(\theta,q)}\quad\forall(\theta,q)\in\mathcal{M}_{2}^{1}.$
(34)

It would be interesting to explore the extent to which the right-hand side of
(33), or of (34) prove informative for models of practical interest. Moreover,
for such models, (34) hints at potential (also to the best of our knowledge)
new variational inference methods that would minimize $(\theta,\phi)\mapsto
I(\theta,q_{\phi})$, where $(q_{\phi})_{\phi\in\Phi}$ denotes a parametrized
family, rather than the conventional approach of minimizing
$(\theta,\phi)\mapsto F(\theta,q_{\phi})$.

## References

  * [1] Ö. Akylidìz, Francesca Romana Crucinio, Mark Girolami, Tim Johnston and Sotirios Sabanis  “Interacting particle Langevin algorithm for maximum marginal likelihood estimation”, 2023
  * [2] Luigi Ambrosio, Nicola Gigli and Giuseppe Savaré  “Gradient Flows: In Metric Spaces and in the Space of Probability Measures”  Springer Science & Business Media, 2005
  * [3] Mihai Anitescu  “Degenerate nonlinear programming with a quadratic growth condition”  In _SIAM Journal on Optimization_ 10, 2000, pp. 1116–1135
  * [4] Dominique Bakry and Michel Émery  “Diffusions hypercontractives”  In _Séminaire de Probabilités XIX 1983/84_ Lecture Notes in Mathematics Springer, 1985
  * [5] Dominique Bakry, Ivan Gentil and Michel Ledoux  “Analysis and Geometry of Markov Diffusion Operators”  Springer, 2014
  * [6] Dmitri Burago, Yuri Burago and Sergei Ivanov  “A Course in Metric Geometry”  American Mathematical Society, 2001
  * [7] René Carmona  “Lectures on BSDEs, Stochastic Control, and Stochastic Differential Games with Financial Applications”  SIAM, 2016
  * [8] Louis-Pierre Chaintron and Antoine Diez  “Propagation of chaos: A review of models, methods and applications. I. Models and methods”  In _Kinetic and Related Models_ 15.6, 2022
  * [9] Rujian Chen  “Approximate Bayesian Modeling with Embedded Gaussian Processes”, 2023
  * [10] Xiang Cheng and Peter Bartlett  “Convergence of Langevin MCMC in KL-divergence”  In _Proceedings of Algorithmic Learning Theory_ 83, 2018, pp. 186–211
  * [11] Sinho Chewi  “Log-concave Sampling” Book draft, 2014  URL: https://chewisinho.github.io
  * [12] Arnak S. Dalalyan  “Theoretical Guarantees for Approximate Sampling from Smooth and Log-Concave Densities”  In _Journal of the Royal Statistical Society Series B: Statistical Methodology_ 79.3, 2016, pp. 651–676
  * [13] Arthur P. Dempster, Nan M. Laird and Donald B. Rubin  “Maximum likelihood from incomplete data via the EM Algorithm”  In _Journal of the Royal Statistical Society, Series B_ 39, 1977, pp. 2–38
  * [14] Steffen Dereich, Michael Scheutzow and Reik Schottstedt  “Constructive quantization: Approximation by empirical measures”  In _Annales de l’Institut Henri Poincaré: Probabilités et statistiques_ 49, 2013, pp. 1183–1203
  * [15] Randal Douc, Éric Moulines and David Stoffer  “Nonlinear Time Series: Theory, Methods and Applications with R Examples”  CRC press, 2014
  * [16] Alain Durmus and Éric Moulines  “High-dimensional Bayesian inference via the unadjusted Langevin algorithm”  In _Bernoulli_ 25.4A, 2019, pp. 2854–2882
  * [17] Alain Durmus and Éric Moulines  “Nonasymptotic convergence analysis for the unadjusted Langevin algorithm”  In _Annals of Applied Probability_ 27, 2017, pp. 1551–1587
  * [18] Zhou Fan, Leying Guan, Yandi Shen and Yinhong Wu  “Gradient flows for empirical Bayes in high-dimensional linear models”, 2023
  * [19] Nicolas Fournier  “Convergence of the empirical measure in expected Wasserstein distance: non-asymptotic explicit bounds in $\mathbb{R}^{d}$”  In _ESAIM: Probability and Statistics_ 27, 2023, pp. 749–775
  * [20] Richard Jordan, David Kinderlehrer and Felix Otto  “The variational formulation of the Fokker–Planck equation”  In _SIAM Journal on Mathematical Analysis_ 29, 1998, pp. 1–17
  * [21] Hamed Karimi, Julie Nutini and Mark Schmidt  “Linear convergence of gradient and proximal-gradient methods under the Polyak-Łojasiewicz condition”  In _Machine Learning and Knowledge Discovery in Databases_ , 2016, pp. 795–811
  * [22] Diederik P. Kingma and Max Welling  “An Introduction to Variational Autoencoders”  In _Foundations and Trends® in Machine Learning_ 12, 2019, pp. 307–392
  * [23] Juan Kuntz, Jen Ning Lim and Adam M. Johansen  “Particle algorithms for maximum likelihood training of latent variable models”  In _Proceedings of The 26th International Conference on Artificial Intelligence and Statistics_ 206, 2023, pp. 5134–5180
  * [24] Jen Ning Lim, Juan Kuntz, Samuel Power and Adam M Johansen  “Momentum particle maximum likelihood”, 2023
  * [25] Stanislaw Łojasiewicz  “Une propriété topologique des sous-ensembles analytiques réels”  In _Les équations aux dérivées partielles_ 117, 1963, pp. 87–89
  * [26] Yi-An Ma, Niladri S. Chatterji, Xiang Cheng, Nicolas Flammarion, Peter L. Bartlett and Michael I. Jordan  “Is there an analog of Nesterov acceleration for gradient-based MCMC?”  In _Bernoulli_ 27.3 Bernoulli Society for Mathematical StatisticsProbability, 2021, pp. 1942–1992
  * [27] Radford M. Neal and Geoffrey E. Hinton  “A View of the EM Algorithm that Justifies Incremental, Sparse, and other Variants”  In _Learning in Graphical Models_ Springer Netherlands, 1998, pp. 355–368
  * [28] Yurii Nesterov  “A method of solving a convex programming problem with convergence rate $\mathcal{O}\bigl{(}\frac{1}{k^{2}}\bigr{)}$”  In _Doklady Akademii Nauk_ 269.3, 1983, pp. 543–547  Russian Academy of Sciences
  * [29] Yurii Nesterov  “Introductory Lectures on Convex Optimization: A Basic Course”  Springer Science & Business Media, 2003
  * [30] Bernt Øksendal  “Stochastic Differential Equations: An Introduction with Applications”  Springer Science & Business Media, 2013
  * [31] Félix Otto and Cédric Villani  “Generalization of an Inequality by Talagrand and Links with the Logarithmic Sobolev Inequality”  In _Journal of Functional Analysis_ 173, 2000, pp. 361–400
  * [32] Boris T. Polyak  “Gradient methods for the minimisation of functionals (in Russian)”  In _Zhurnal Vychislitel’noi Matematiki i Matematicheskoi Fiziki_ 3, 1963, pp. 643–653
  * [33] Herbert Robbins  “An empirical Bayes approach to statistics”  In _Proceedings of the Third Berkeley Symposium on Mathematical Statistics and Probability_ 3.1, 1956, pp. 157–164
  * [34] Filippo Santambrogio  “Optimal Transport for Applied Mathematicians: Calculus of Variations, PDEs, and Modeling”  Birkhäuser/Springer, 2015
  * [35] Louis Sharrock, Daniel Dodd and Christopher Nemeth  “CoinEM: Tuning-free particle-based variational inference for latent variable models”, 2023
  * [36] Michel Talagrand  “Transportation cost for Gaussian and other product measures”  In _Geometric & Functional Analysis_ 6 Springer, 1996, pp. 587–600
  * [37] N. Trillos, Bamdad Hosseini and Daniel Sanz-Alonso  “From Optimization to Sampling Through Gradient Flows”  In _Notices of the American Mathematical Society_ 70.6, 2023
  * [38] Cédric Villani  “Optimal Transport: Old and New”  Springer Science & Business Media, 2009

##  Appendix A Proofs for Section 2

###  A.1 Formal derivation of de Bruijin’s identity

Using the defining equation of the gradient flow (5) and integrating by parts
we compute

$\displaystyle\frac{\operatorname{d\\!}}{\operatorname{d\\!}t}\int\log(q_{t}(x))q_{t}(\operatorname{d\\!}x)$
$\displaystyle=\int(\log(q_{t}(x))+1)\nabla_{x}\cdot\bigg{[}q_{t}(x)\nabla_{x}\log\bigg{(}\frac{q_{t}(x)}{\rho_{\theta_{t}}(x)}\bigg{)}\bigg{]}\operatorname{d\\!}x$
$\displaystyle=-\int\left\langle{\nabla_{x}\log(q_{t}(x))},{\nabla_{x}\log\bigg{(}\frac{q_{t}(x)}{\rho_{\theta_{t}}(x)}\bigg{)}}\right\rangle
q(\operatorname{d\\!}x)$

and

$\displaystyle\frac{\operatorname{d\\!}}{\operatorname{d\\!}t}\int\ell(\theta_{t},x)q_{t}(\operatorname{d\\!}x)$
$\displaystyle=\int\left\langle{\nabla_{\theta}\ell(\theta_{t},x)},{\dot{\theta_{t}}}\right\rangle
q_{t}(\operatorname{d\\!}x)-\int\ell(\theta_{t},x)\nabla_{x}\cdot\bigg{[}q_{t}(x)\nabla_{x}\log\bigg{(}\frac{q_{t}(x)}{\rho_{\theta_{t}}(x)}\bigg{)}\bigg{]}\operatorname{d\\!}x$
$\displaystyle=\mathinner{\\!\left\lVert\int\nabla_{\theta}\ell(\theta_{t},x)q_{t}(\operatorname{d\\!}x)\right\rVert}^{2}+\int\left\langle{\nabla_{x}\log(\rho_{\theta_{t}}(x))},{\nabla_{x}\log\bigg{(}\frac{q_{t}(x)}{\rho_{\theta_{t}}(x)}\bigg{)}}\right\rangle
q_{t}(\operatorname{d\\!}x).$

which upon re-arranging shows that formally
$(\operatorname{d\\!}/\operatorname{d\\!}t)F(\theta_{t},q_{t})=-I(\theta_{t},q_{t})$.

###  A.2 Proof of Theorem 2

Here we follow some arguments from [31]. Throughout this appendix, we assume
that Assumptions 1– 3 hold. Let $(\theta,q)$ be any point in
$\mathcal{M}_{2}$. As we show in Lemma 13 below,

$\frac{\operatorname{d\\!}}{\operatorname{d\\!}t}\mathsf{d}((\theta_{t},q_{t}),(\theta,q))\leq\sqrt{I(\theta_{t},q_{t})}\quad\forall
t>0.$ (35)

By the theorem’s premise, the measures
$(\rho_{\theta}(\operatorname{d\\!}x))_{\theta\in\mathbb{R}^{d_{\theta}}}$
satisfy the xLSI or, in other words,

$\sqrt{I(\theta,q)}\leq\frac{I(\theta,q)}{\sqrt{2\lambda(F(\theta,q)-F_{*})}}\quad\forall(\theta,q)\in\mathcal{M}_{2}^{1}.$

By the de Bruijn’s identity (15), we find that

$\frac{I(\theta_{t},q_{t})}{2\sqrt{{F(\theta_{t},q_{t})-F_{*}}}}=-\frac{\operatorname{d\\!}}{\operatorname{d\\!}t}\sqrt{F(\theta_{t},q_{t})-F_{*}}\quad\forall
t>0.$

Combining the above three, we obtain

$\displaystyle\frac{\operatorname{d\\!}}{\operatorname{d\\!}t}\mathsf{d}((\theta_{t},q_{t}),(\theta,q))$
$\displaystyle\leq\sqrt{I(\theta_{t},q_{t})}\leq\frac{I(\theta_{t},q_{t})}{\sqrt{2\lambda[F(\theta_{t},q_{t})-F_{*}]}}=-\frac{\operatorname{d\\!}}{\operatorname{d\\!}t}\sqrt{\frac{2[F(\theta_{t},q_{t})-F_{*}]}{\lambda}}\quad\forall
t>0.$

Integrating over $t$ in $(t,t^{\prime})$ for $t^{\prime}\geq t\geq 0$, yields

$\mathsf{d}((\theta_{t^{\prime}},q_{t^{\prime}}),(\theta,q))-\mathsf{d}((\theta_{t},q_{t}),(\theta,q))\leq\sqrt{\frac{2[F(\theta_{t},q_{t})-F_{*}]}{\lambda}}-\sqrt{\frac{2[F(\theta_{t^{\prime}},q_{t^{\prime}})-F_{*}]}{\lambda}}.$

If we set $(\theta,q)\mathrel{\mathop{\mathchar
58\relax}}=(\theta_{t},q_{t})$, then

$\mathsf{d}((\theta_{t},q_{t}),(\theta_{t^{\prime}},q_{t^{\prime}}))\leq\sqrt{\frac{2[F(\theta_{t},q_{t})-F_{*}]}{\lambda}}-\sqrt{\frac{2[F(\theta_{t^{\prime}},q_{t^{\prime}})-F_{*}]}{\lambda}}\quad\forall
t^{\prime}\geq t\geq 0.$ (36)

As we show in Lemma 14 below, there exists an increasing sequence
$t_{1}<t_{2}<\dots$ approaching $\infty$ such that

$(\theta_{t_{n}},q_{t_{n}})\to(\theta_{*},q_{*})\quad\text{as}\quad
n\to\infty,$ (37)

for some $(\theta_{*},q_{*})$ belonging to the optimal set $\mathcal{M}_{*}$
in the topology induced by $\mathsf{d}$ on $\mathcal{M}_{2}$. Because
$F(\theta_{t_{n}},q_{t_{n}})$ converges to $F_{*}$ as $n\to\infty$ (Theorem 1)
and

$(\theta^{\prime},q^{\prime})\mapsto\mathsf{d}((\theta,q),(\theta^{\prime},q^{\prime}))$

is a continuous function on $(\mathcal{M}_{2},\mathsf{d})$ for any given
$(\theta,q)$ in $\mathcal{M}_{2}$, setting
$(t,t^{\prime})\mathrel{\mathop{\mathchar 58\relax}}=(0,t_{n})$ in (36) and
taking the limit $n\to\infty$, then implies that

$\mathsf{d}((\theta_{0},q_{0}),\mathcal{M}_{*})\leq\sqrt{\frac{2[F(\theta_{0},q_{0})-F_{*}]}{\lambda}};$

since $(\theta_{0},q_{0})\in\mathcal{M}_{2}$ is arbitrary, the claim follows.

######  Lemma 13.

(35) holds.

###### Proof.

Let $(\theta,q)\in\mathcal{M}_{2}$ and let $(\theta_{t},q_{t})$ satisfy (5)
and define $v_{t}(x)\mathrel{\mathop{\mathchar
58\relax}}=\nabla_{x}\log(q_{t}(x)/\rho_{\theta_{t}}(x))$. [2, Theorem 8.4.7]
shows that

$\frac{\operatorname{d\\!}}{\operatorname{d\\!}t}W_{2}(q_{t},q)^{2}=2\int\left\langle{v_{t}(x_{1})},{x_{1}-x_{2}}\right\rangle\operatorname{d\\!}\varrho(x_{1},x_{2})$
(38)

where $\varrho$ it an optimal transport plan for $(q_{t},q)$. By an
application of the Cauchy-Schwartz inequality and the definition of the
Wasserstein-2 distance we get

$\frac{\operatorname{d\\!}}{\operatorname{d\\!}t}W_{2}(q_{t},q)^{2}\leq
2\sqrt{\int\mathinner{\\!\left\lVert
x_{1}-x_{2}\right\rVert}^{2}\operatorname{d\\!}\varrho(x_{1},x_{2})\int\mathinner{\\!\left\lVert
v_{t}(x)\right\rVert}^{2}q(\operatorname{d\\!}x)}=2W_{2}(q_{t},q)\sqrt{\int\mathinner{\\!\left\lVert
v_{t}(x)\right\rVert}^{2}q(\operatorname{d\\!}x)}.$

On the other hand,
$(\operatorname{d\\!}/\operatorname{d\\!}t)\mathinner{\\!\left\lVert\theta_{t}-\theta\right\rVert}^{2}=2\left\langle{\dot{\theta_{t}}},{\theta_{t}-\theta}\right\rangle\leq
2\lVert\dot{\theta_{t}}\rVert\mathinner{\\!\left\lVert\theta_{t}-\theta\right\rVert}$,
and by these estimates and then again by Cauchy–Schwartz,

$\displaystyle\frac{\operatorname{d\\!}}{\operatorname{d\\!}t}\mathsf{d}((\theta_{t},q_{t}),(\theta,q))=\frac{(\operatorname{d\\!}/\operatorname{d\\!}t)\mathsf{d}((\theta_{t},q_{t}),(\theta,q))^{2}}{2\mathsf{d}((\theta_{t},q_{t}),(\theta,q))}\leq\sqrt{\lVert\dot{\theta_{t}}\rVert^{2}+\int\mathinner{\\!\left\lVert
v_{t}(x)\right\rVert}^{2}q(\operatorname{d\\!}x)}=\sqrt{I(\theta_{t},q_{t})}.$

∎

######  Lemma 14.

For any increasing sequence $t_{1}<t_{2}<\dots$ approaching $\infty$, (37)
holds for some $(\theta_{*},q_{*})$ in $\mathcal{M}_{*}$.

###### Proof.

Inequality (36) tells us that the sequence is Cauchy in
$(\mathcal{M}_{2}^{1},\mathsf{d})$ and, consequently, in
$(\mathcal{M}_{2}\mathrel{\mathop{\mathchar
58\relax}}=\mathbb{R}^{d_{\theta}}\times\mathcal{P}_{2}(\mathbb{R}^{d_{x}}),\mathsf{d})$.
Because both $(\mathbb{R}^{d_{\theta}},\mathsf{d}_{E})$ and
$(\mathcal{P}_{2}(\mathbb{R}^{d_{x}}),\mathsf{d}_{W_{2}})$ are complete metric
spaces [38, Theorem 6.18], so is $(\mathcal{M}_{2},\mathsf{d})$. Hence,
$(\theta_{t_{n}},q_{t_{n}})_{n\in\mathbb{N}}$ converges to a limit
$(\theta_{\infty},q_{\infty})$ in $\mathcal{M}_{2}$. As we show in Lemma 15
below, the free energy $F$ is lower semicontinuous on
$(\mathcal{M}_{2},\mathsf{d})$. Consequently, Theorem 1 implies that

$F(\theta_{\infty},q_{\infty})\leq\liminf_{n\to\infty}F(\theta_{t_{n}},q_{t_{n}})=F_{*},$

where the equality holds because
$(\theta_{t_{n}},q_{t_{n}})_{n\in\mathbb{N}}\subseteq\mathcal{M}_{2}^{1}$.
Since it also holds $F(\theta_{\infty},q_{\infty})\geq F_{*}$, the claim
follows. ∎

######  Lemma 15.

$F$ is lower semicontinuous on $(\mathcal{M}_{2},\mathsf{d})$.

###### Proof.

Given any $\pi$ in $\mathcal{P}_{2}(\mathbb{R}^{d_{x}})$, the duality formula
for the $\mathrm{KL}$ divergence (e.g., [2], Lemma 9.4.4) reads

$\mathrm{KL}(q||\pi)=\sup\left\\{\int\phi(x)q(\operatorname{d\\!}x)-\int(e^{\phi(x)}-1)\pi(\operatorname{d\\!}x)\mathrel{\mathop{\mathchar
58\relax}}\phi\in\mathcal{C}_{b}(\mathbb{R}^{d_{x}})\right\\},$

where $\mathcal{C}_{b}(\mathbb{R}^{d_{x}})$ denotes the set of continuous
bounded functions $\mathbb{R}^{d_{x}}$. Consequently,

$\displaystyle
F(\theta,q)=\mathrm{KL}(q||\pi_{\theta})-\log(Z_{\theta})=\sup\\{G(\theta,q,\phi)\mathrel{\mathop{\mathchar
58\relax}}\phi\in\mathcal{C}_{b}(\mathbb{R}^{d_{x}})\\}-\log(Z_{\theta}),$

where $G(\theta,q,\phi)\mathrel{\mathop{\mathchar
58\relax}}=\int\phi(x)q(\operatorname{d\\!}x)-\int(e^{\phi(x)}-1)\pi_{\theta}(\operatorname{d\\!}x)$.
Because Assumption 1(i) implies that $\theta\mapsto\log(Z_{\theta})$ is
continuous and the pointwise supremum of a family of lower semicontinuous
(l.s.c.) functions is l.s.c., we need only show that $(\theta,q)\mapsto
G(\theta,q,\phi)$ is l.s.c. for each $\phi$ in
$\mathcal{C}_{b}(\mathbb{R}^{d_{x}})$. Since the Wasserstein-2 topology is
finer than the weak topology, it suffices to show that $(\theta,q)\mapsto
G(\theta,q,\phi)$ is continuous in the weak topology for any given $\phi$ in
$\mathcal{C}_{b}(\mathbb{R}^{d_{x}})$. Consider a sequence
$(\theta_{n},q_{n})_{n\in\mathbb{N}}$ such that $\theta_{n}\rightarrow\theta$
and $q_{n}\rightarrow q$ weakly. By definition of weak convergence,
$\int\phi(x)q_{n}(\operatorname{d\\!}x)\rightarrow\int\phi(x)q(\operatorname{d\\!}x)$.
Furthermore, since
$x\mapsto(e^{\phi(x)}-1)\in\mathcal{C}_{b}(\mathbb{R}^{d_{x}})$, we have
$|\int(e^{\phi(x)}-1)\pi_{\theta_{n}}(\operatorname{d\\!}x)-\int(e^{\phi(x)}-1)\pi_{\theta}(\operatorname{d\\!}x)|\leq
c_{\phi}\int|\pi_{\theta_{n}}(x)-\pi_{\theta}(x)|\operatorname{d\\!}x$ for
some $c_{\phi}<\infty$. Since $\theta\mapsto\pi_{\theta}$ is continuous by
Assumption 1(i), an application of Scheffé’s lemma concludes the proof. ∎

###  A.3 Proof of Theorem 4

We preface the proof of Theorem 4 with three auxiliary results which describe
variations of the free energy along geodesics in
$(\mathcal{M}_{2},\mathsf{d})$.

######  Definition 3.

A curve $\gamma(t)\mathrel{\mathop{\mathchar
58\relax}}t\in[0,1]\mapsto\mathcal{M}_{2}$ is a (constant speed) geodesic if

$\mathsf{d}(\gamma(s),\gamma(t))=(t-s)\mathsf{d}(\gamma(0),\gamma(1)),\quad\forall
0\leq s\leq t\leq 1.$

For a probability measure $\mu\in\mathcal{P}(\mathbb{R}^{d_{x}})$ and a
measurable map $T$ with domain $\mathbb{R}^{d_{x}}$, $T_{\\#}\mu=\mu\circ
T^{-1}$ indicates the pushforward of $\mu$ by $T$.

######  Lemma 16.

A curve $\gamma(t)\mathrel{\mathop{\mathchar
58\relax}}t\in[0,1]\mapsto\mathcal{M}_{2}$ is a geodesic if and only if
$\gamma(t)=(\gamma_{\theta}(t),\gamma_{q}(t))$, where $\gamma_{\theta}$ is a
geodesic in
$(\mathbb{R}^{d_{\theta}},\mathinner{\\!\left\lVert\cdot\right\rVert})$ and
$\gamma_{q}$ is a geodesic in $(\mathcal{P}_{2}(\mathbb{R}^{d_{x}}),W_{2})$.
In particular, if $\gamma(t)$ is a geodesic in $(\mathcal{M}_{2},\mathsf{d})$
connecting $(\theta,q)$ and $(\theta^{\prime},q^{\prime})$ then
$\gamma(t)=(\gamma_{\theta}(t)=(1-t)\theta+t\theta^{\prime},\gamma_{q}(t)=(h_{t})_{\\#}\varrho)$
where $\varrho$ is a $W_{2}$ optimal transport plan for $q,q^{\prime}$ and
$h_{t}(x,x^{\prime})=(1-t)x+tx^{\prime}$. Furthermore, if $q$ has a density
w.r.t. the Lebesgue measure, we can also write
$\gamma_{q}(t)=((1-t)\textrm{id}+t\nabla_{x}\Phi)_{\\#}q$ for some convex
function $\Phi$.

###### Proof.

The first claim is almost immediate from the definition, also see [6, Lemma
3.6.4]. The second claim follows from the first using the characterization of
geodesics in $(\mathcal{P}_{2}(\mathbb{R}^{d_{x}}),W_{2})$—see e.g. [34,
Theorem 5.27]. The last claim is Brenier’s Theorem [38, Theorem 9.4]. ∎

The following lemma uses similar arguments to [10, Lemma 9], and it connects
the strong concavity of the log-likelihood function with strong geodesic
convexity of $F$.

######  Lemma 17.

If Assumption 4 holds, $F$ is $\lambda$-strongly geodesically convex i.e. for
every $(\theta,q),(\theta^{\prime},q^{\prime})$ there is constant speed
geodesic $\gamma\mathrel{\mathop{\mathchar 58\relax}}[0,1]\to\mathcal{M}_{2}$
connecting $\gamma(0)=(\theta,q)$ and $\gamma(1)=(\theta^{\prime},q^{\prime})$
such that $t\mapsto F(\gamma(t))$ is $\lambda$-strongly convex:

$F(\gamma(t))\leq(1-t)F(\theta,q)+tF(\theta^{\prime},q^{\prime})-\frac{\lambda
t(1-t)}{2}\mathsf{d}((\theta,q),(\theta^{\prime},q^{\prime}))^{2}.$ (39)

###### Proof.

Let $\gamma=(\gamma_{\theta},\gamma_{q})$ be a geodesic in $\mathcal{M}_{2}$
connecting $(\theta,q)$ and $(\theta^{\prime},q^{\prime})$, and recall that by
Lemma 16 $\gamma_{\theta}$ and $\gamma_{q}$ are geodesics in
$(\mathbb{R}^{d_{\theta}},\mathinner{\\!\left\lVert\cdot\right\rVert})$ and
$(\mathcal{P}_{2}(\mathbb{R}^{d_{x}}),W_{2})$, respectively. Since
$F(\theta,q)=\int\log(q(x))q(\operatorname{d\\!}x)-\int\ell(\theta,x)q(\operatorname{d\\!}x)$,
and since $q\mapsto\int\log(q(x))q(\operatorname{d\\!}x)$ is convex along
$\gamma_{q}$ by [34, Theorem 7.28], we just need to show that
$V(\theta,q)\mathrel{\mathop{\mathchar
58\relax}}=-\int\ell(\theta,x)q(\operatorname{d\\!}x)$ is $\lambda$-strongly
convex along $\gamma$. Using the representation in Lemma 16 above and then
applying Assumption 4 we obtain

$\displaystyle V(\theta(t))=$
$\displaystyle-\int\ell(\gamma_{\theta}(t),x)\gamma_{q}(t)(x)\operatorname{d\\!}x$
$\displaystyle=$
$\displaystyle-\int\ell((1-t)\theta+t\theta^{\prime},(1-t)x+tx^{\prime})\varrho(x,x^{\prime})\operatorname{d\\!}x\operatorname{d\\!}x^{\prime}$
$\displaystyle\leq$
$\displaystyle-\int\int((1-t)\ell(\theta,x)+t\ell(\theta^{\prime},x^{\prime}))\varrho(x,x^{\prime})\operatorname{d\\!}x\operatorname{d\\!}x^{\prime}$
$\displaystyle-\frac{\lambda
t(1-t)}{2}\int\left(\mathinner{\\!\left\lVert\theta-\theta^{\prime}\right\rVert}^{2}+\mathinner{\\!\left\lVert
x-x^{\prime}\right\rVert}^{2}\right)\varrho(x,x^{\prime})\operatorname{d\\!}x\operatorname{d\\!}x^{\prime}$
$\displaystyle=$
$\displaystyle(1-t)V(\theta,q)+tV(\theta^{\prime},q^{\prime})-\frac{\lambda
t(1-t)}{2}\mathsf{d}((\theta,q),(\theta^{\prime},q^{\prime}))^{2}.$

∎

The next result supplies an estimate for the (right) lower derivative of the
free energy along geodesics in $\mathcal{M}_{2}$ (c.f. [38, Theorem 20.1]).

######  Lemma 18.

Let $\gamma$ be a geodesic connecting
$(\theta,q),(\theta^{\prime},q^{\prime})\in\mathcal{M}_{2}^{1}$. If $\ell$ is
differentiable a.e.,

$\liminf_{t\rightarrow
0^{+}}\frac{F(\gamma(t))-F(\gamma(0))}{t}\geq\left\langle{\dot{\gamma}(0)},{\nabla
F(\theta,q)}\right\rangle_{q}$ (40)

where

$\left\langle{\dot{\gamma}(0)},{\nabla
F(\theta,q)}\right\rangle_{q}\mathrel{\mathop{\mathchar
58\relax}}=\left\langle{\gamma_{q}(0)},{\nabla_{q}F(\theta,q)}\right\rangle_{q}+\left\langle{\gamma_{\theta}(0)},{\nabla_{\theta}F(\theta,q)}\right\rangle$

and

$\nabla_{q}F(\theta,q)(x)\mathrel{\mathop{\mathchar
58\relax}}=\nabla_{x}\log\bigg{(}\frac{q(x)}{\rho_{\theta}(x)}\bigg{)},\quad\nabla_{\theta}F(\theta,q)\mathrel{\mathop{\mathchar
58\relax}}=-\int\nabla_{\theta}\ell(\theta,x)q(\operatorname{d\\!}x).$

###### Proof.

Let us consider the geodesic $\gamma(t)$ given by
$\gamma(t)\mathrel{\mathop{\mathchar 58\relax}}=(\theta_{t},q_{t})$ with
$q_{t}(\operatorname{d\\!}x)\mathrel{\mathop{\mathchar
58\relax}}=((1-t)x+t\nabla_{x}\Phi(x))\\#q(\operatorname{d\\!}x)$ and
$\theta_{t}=(1-t)\theta+t\theta^{\prime}$ — see Lemma 16. Since $\Phi$ is a
convex function, by Alexandrov’s theorem (Theorem 14.24 in [38]) $\Phi$ has a
second derivative almost everywhere. Let $\nabla^{2}_{x}$ denote the Hessian
with respect to the variable $x$. We use the change of variables
$x\rightarrow(1-t)x+t\nabla_{x}\Phi(x)$ and then the facts
$\theta_{t}=(1-t)\theta+t\theta^{\prime}$ and
$q_{t}(\operatorname{d\\!}x)\mathrel{\mathop{\mathchar
58\relax}}=((1-t)x+t\nabla_{x}\Phi(x))\\#q(\operatorname{d\\!}x)\Rightarrow
q_{t}((1-t)x+t\nabla_{x}\Phi(x))=q(x)/\det((1-t)I_{d_{x}}+t\nabla^{2}_{x}\Phi(x))$
to write

$\displaystyle F(\gamma(t))=F(\theta_{t},q_{t})=\int
q_{t}(x)(\log(q_{t}(x))-\ell(\theta_{t},x))\operatorname{d\\!}x$
$\displaystyle=\int
q((1-t)x+t\nabla_{x}\Phi(x))\log(q(x+t\nabla_{x}\Phi(x))-q((1-t)x+t\nabla_{x}\Phi(x))$
$\displaystyle\cdot\ell(\theta_{t},(1-t)x+t\nabla_{x}\Phi(x))\det((1-t)I_{d_{x}}+t\nabla^{2}_{x}\Phi(x))\operatorname{d\\!}x$
$\displaystyle=\int
q(x)\log\left(\frac{q(x)}{\det((1-t)I_{d_{x}}+t\nabla^{2}_{x}\Phi(x))}\right)\operatorname{d\\!}x-q(x)\ell(\theta_{t},(1-t)x+t\nabla_{x}\Phi(x))\operatorname{d\\!}x$
$\displaystyle=\mathrel{\mathop{\mathchar 58\relax}}\int
q(x)\log(q(x))\operatorname{d\\!}x+F_{2}(\gamma(t))+F_{3}(\gamma(t)).$

Now, by Fatou’s lemma,

$\displaystyle\liminf_{t\rightarrow
0^{+}}\frac{F_{2}(\theta_{t},q_{t})-F_{2}(\theta_{0},q_{0})}{t}\geq\int\liminf_{t\rightarrow
0^{+}}\frac{\log(\det(I_{d_{x}}))-\log(\det((1-t)I_{d_{x}}+t\nabla^{2}_{x}\Phi(x)))}{t}q(\operatorname{d\\!}x)$

and by positive semi-definiteness of $(1-t)I_{d_{x}}+t\nabla^{2}_{x}\Phi(x)$,
the arithmetic mean–geometric mean inequality yields that

$\displaystyle\det((1-t)I_{d_{x}}+t\nabla^{2}_{x}\Phi(x))\leq\bigg{(}\frac{\textrm{Tr}((1-t)I_{d_{x}}+t\nabla^{2}_{x}\Phi(x))}{d_{x}}\bigg{)}^{d_{x}}=\bigg{(}(1-t)+\frac{t}{d_{x}}\Delta_{x}\Phi(x)\bigg{)}^{d_{x}}$

it follows that

$\displaystyle\liminf_{t\rightarrow
0^{+}}\frac{F_{2}(\theta_{t},q_{t})-F_{2}(\theta_{0},q_{0})}{t}$
$\displaystyle\geq\int\liminf_{t\rightarrow
0^{+}}\bigg{[}-\frac{d_{x}}{t}\log\bigg{(}(1-t)+\frac{t}{d_{x}}\Delta_{x}\Phi(x)\bigg{)}\bigg{]}q(\operatorname{d\\!}x)$
$\displaystyle=-\int(\Delta_{x}\Phi(x)-d_{x})q(\operatorname{d\\!}x).$

Next, we compute $F_{3}(\gamma(t))$’s lower derivative as

$\displaystyle\liminf_{t\rightarrow
0^{+}}\frac{F_{3}(\theta_{t},q_{t})-F_{3}(\theta_{0},q_{0})}{t}\geq\int
q(\operatorname{d\\!}x)\liminf_{t\rightarrow
0^{+}}\bigg{[}\frac{\ell(\theta,x)-\ell(\theta_{t},(1-t)x+t\nabla_{x}\Phi(x))}{t}\bigg{]}$
$\displaystyle=\int-\left\langle{\nabla\ell(\theta,x)},{(\theta^{\prime}-\theta,\nabla_{x}\Phi(x)-x)}\right\rangle
q(\operatorname{d\\!}x)$

where the limit exchange is again justified by Fatou’s lemma. Putting these
results together, integrating by parts and using the divergence theorem then
gives

$\displaystyle\liminf_{t\rightarrow
0^{+}}\frac{F(\theta_{t},q_{t})-F(\theta_{0},q_{0})}{t}$
$\displaystyle\geq-\int q(x)(\Delta_{x}\Phi(x)-d_{x})\operatorname{d\\!}x-\int
q(x)\left\langle{\nabla\ell(\theta,x)},{(\theta^{\prime}-\theta,\nabla_{x}\Phi(x)-x)}\right\rangle\operatorname{d\\!}x$
$\displaystyle=\int\left\langle{\nabla_{x}q(x)},{\nabla_{x}\Phi(x)-x}\right\rangle\operatorname{d\\!}x-\int
q(x)\left\langle{\nabla\ell(\theta,x)},{(\theta^{\prime}-\theta,\nabla_{x}\Phi(x)-x)}\right\rangle\operatorname{d\\!}x$
$\displaystyle=\int\left\langle{\nabla_{x}q(x)},{\nabla_{x}\Phi(x)-x}\right\rangle-q(x)\left\langle{\nabla_{x}\ell(\theta,x)},{\nabla_{x}\Phi(x)-x}\right\rangle\operatorname{d\\!}x-\int\left\langle{\theta^{\prime}-\theta},{\nabla_{\theta}\ell(\theta,x)}\right\rangle
q(x)\operatorname{d\\!}x$
$\displaystyle=\int\left\langle{\nabla_{x}\log\left(\frac{q(x)}{\rho_{\theta}(x)}\right)},{\nabla_{x}\Phi(x)-x}\right\rangle
q(x)\operatorname{d\\!}x-\int\left\langle{\theta^{\prime}-\theta},{\nabla_{\theta}\ell(\theta,x)}\right\rangle
q(x)\operatorname{d\\!}x$ $\displaystyle=\mathrel{\mathop{\mathchar
58\relax}}\left\langle{\dot{\gamma}(0)},{\nabla
F(\theta,q)}\right\rangle_{q}.$

∎

###### Proof of Theorem 4.

Take any two points $(\theta,q)$ and $(\theta^{\prime},q^{\prime})$ in
$\mathcal{M}_{2}^{1}$ and consider a geodesic
$\gamma=(\gamma_{\theta},\gamma_{q})$ connecting those. Since under Assumption
4, $\ell$ is differentiable a.e., taking the right lower derivative w.r.t. $t$
at 0 in (39) and using Lemma 18 we obtain

$\left\langle{\dot{\gamma}(0)},{\nabla
F(\theta,q)}\right\rangle_{q}\leq\liminf_{t\rightarrow
0^{+}}\frac{F(\gamma(t))-F(\gamma(0))}{t}\leq
F(\theta^{\prime},q^{\prime})-F(\theta,q)-\frac{\lambda}{2}\mathsf{d}((\theta,q),(\theta^{\prime},q^{\prime}))^{2}.$
(41)

Because $\gamma_{q}$ is constant-speed w.r.t. $W_{2}$, it follows that
$\mathinner{\\!\left\lVert\dot{\gamma}_{q}(0)\right\rVert}_{q}^{2}=W_{2}(q,q^{\prime})^{2}$,
whence

$\displaystyle\mathsf{d}((\theta,q),(\theta^{\prime},q^{\prime}))^{2}=\mathinner{\\!\left\lVert\theta-\theta^{\prime}\right\rVert}^{2}+W_{2}(q,q^{\prime})^{2}=\mathinner{\\!\left\lVert\dot{\gamma}_{\theta}(0)\right\rVert}^{2}+\mathinner{\\!\left\lVert\dot{\gamma}_{q}(0)\right\rVert}_{q}^{2}=\mathrel{\mathop{\mathchar
58\relax}}\mathinner{\\!\left\lVert\dot{\gamma}(0)\right\rVert}_{q}^{2}$

Plugging the above into (41) and setting
$(\theta^{\prime},q^{\prime})\mathrel{\mathop{\mathchar
58\relax}}=(\theta_{*},\pi_{\theta_{*}})$, we find that

$\displaystyle F(\theta,q)-F_{*}\leq-{\left\langle{\dot{\gamma}(0)},{\nabla
F(\theta,q)}\right\rangle}_{q}-\frac{\lambda}{2}{\mathinner{\\!\left\lVert\dot{\gamma}(0)\right\rVert}}_{q}^{2}.$
(42)

Now, the function

$\displaystyle\tau\mapsto\left\langle{\tau},{\nabla
F(\theta,q)}\right\rangle_{q}+\frac{\lambda}{2}\mathinner{\\!\left\lVert\tau\right\rVert}_{q}^{2}$

is minimized by $-c\nabla F(\theta,q)$ for some constant $c$, and

$\left\langle{-c\nabla F(\theta,q)},{\nabla
F(\theta,q)}\right\rangle_{q}+\frac{\lambda}{2}\mathinner{\\!\left\lVert-c\nabla
F(\theta,q)\right\rVert}_{q}^{2}=\left(\frac{\lambda}{2}c^{2}-c\right)\mathinner{\\!\left\lVert\nabla
F(\theta,q)\right\rVert}^{2}_{q}$

is minimized at $c=1/\lambda$. For these reasons,

$\left\langle{-c\nabla F(\theta,q)},{\nabla
F(\theta,q)}\right\rangle_{q}+\frac{\lambda}{2}\mathinner{\\!\left\lVert-c\nabla
F(\theta,q)\right\rVert}_{q}^{2}\geq-\frac{1}{2\lambda}\mathinner{\\!\left\lVert\nabla
F(\theta,q)\right\rVert}^{2}_{q}=-\frac{1}{2\lambda}I(\theta,q),$

and plugging this into (42) yields the claim. ∎

###  A.4 Proof of Theorem 5

Let $t>0$. Consider a geodesic $\gamma=(\gamma_{\theta},\gamma_{q})$
connecting the points $(\theta_{t},q_{t})$ and $(\theta_{*},\pi_{\theta_{*}})$
in $\mathcal{M}_{2}^{1}$. Setting $(\theta,q)=(\theta_{t},q_{t})$ and
$(\theta^{\prime},q^{\prime})=(\theta_{*},\pi_{\theta_{*}})$ in (41), and
since under Assumptions 1(iii) and 4 the xLSI holds (Theorem 4), we can use
the extended Talagrand-type inequality (Theorem 2) to obtain

$\left\langle{\dot{\gamma}(0)},{\nabla
F(\theta_{t},q_{t})}\right\rangle_{q_{t}}\leq
F_{*}-F(\theta_{t},q_{t})-\frac{\lambda}{2}\mathsf{d}((\theta_{t},q_{t}),(\theta_{*},\pi_{\theta_{*}}))^{2}\leq-\lambda\mathsf{d}((\theta_{t},q_{t}),(\theta_{*},\pi_{\theta_{*}}))^{2}.$

Let $\varrho$ be the optimal transport plan for $(q_{t},\pi_{\theta_{*}})$ and
let $v_{t}(x)=\nabla_{x}\log(q_{t}(x)/\rho_{\theta_{t}}(x))$. Combining the
above inequality with (38) we write

$\displaystyle\frac{\operatorname{d\\!}}{\operatorname{d\\!}t}\mathsf{d}((\theta_{t},q_{t}),(\theta_{*},\pi_{\theta_{*}}))^{2}$
$\displaystyle=2\int\left\langle{v_{t}(x_{1})},{x_{1}-x_{2}}\right\rangle\operatorname{d\\!}\varrho(x_{1},x_{2})+2\left\langle{\dot{\theta_{t}}},{\theta_{t}-\theta_{*}}\right\rangle$
$\displaystyle=2\left\langle{\dot{\gamma}(0)},{\nabla
F(\theta_{t},q_{t})}\right\rangle_{q_{t}}\leq-2\lambda\mathsf{d}((\theta_{t},q_{t}),(\theta_{*},\pi_{\theta_{*}}))^{2},$

upon which the result follows via Grönwall’s inequality. ∎

##  Appendix B Proofs for Section 3

###  B.1 Proof of Proposition 6

######  Proposition 19.

Assume Assumptions 4–5 and let $(\theta_{\dagger},x_{\dagger})$ denote
$\ell$’s unique maximizers. Then, we have the following uniform-in-time second
moment bound

$\mathinner{\\!\left\lVert\theta_{t}\right\rVert}^{2}+\mathbb{E}\left[\mathinner{\\!\left\lVert
X_{t}\right\rVert}^{2}\right]\leq
2\left(\mathinner{\\!\left\lVert\theta_{\dagger}\right\rVert}^{2}+\mathinner{\\!\left\lVert
x_{\dagger}\right\rVert}^{2}+\mathinner{\\!\left\lVert\theta_{0}-\theta_{\dagger}\right\rVert}^{2}+\mathbb{E}\left[\mathinner{\\!\left\lVert
X_{0}-x_{\dagger}\right\rVert}^{2}\right]+\frac{d_{x}}{\lambda}\right)\quad\forall
t\geq 0.$

###### Proof.

Because,

$f(z)\mathrel{\mathop{\mathchar 58\relax}}=\mathinner{\\!\left\lVert
z\right\rVert}^{2}\Rightarrow\nabla f(z)=2z,\quad\nabla^{2}f(z)=2I,$

by setting

$\xi_{t}\mathrel{\mathop{\mathchar 58\relax}}=\mathinner{\\!\left\lVert
X_{t}-x_{\dagger}\right\rVert}^{2}+\mathinner{\\!\left\lVert\theta_{t}-\theta_{\dagger}\right\rVert}^{2}=\mathinner{\\!\left\lVert(\theta_{t},X_{t})-(\theta_{\dagger},x_{\dagger})\right\rVert}^{2},$

applying Itô’s formula [30, Theorem 4.2.1] and noting that $X_{t}$ has law
$q_{t}$, we find

$\displaystyle\operatorname{d\\!}\xi_{t}=\left[2\left\langle{(\theta_{t}-\theta_{\dagger},X_{t}-x_{\dagger})},{(\mathbb{E}\left[\nabla_{\theta}\ell(\theta_{t},X_{t})\right],\nabla_{x}\ell(\theta_{t},X_{t}))}\right\rangle+d_{x}\right]\operatorname{d\\!}t+\sqrt{2}\left\langle{X_{t}-x_{\dagger}},{\operatorname{d\\!}W_{t}}\right\rangle.$
(43)

However,

$\displaystyle\left\langle{(\theta_{t}-\theta_{\dagger},X_{t}-x_{\dagger})},{(\mathbb{E}\left[\nabla_{\theta}\ell(\theta_{t},X_{t})\right],\nabla_{x}\ell(\theta_{t},X_{t}))}\right\rangle$
(44)
$\displaystyle\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad=\left\langle{\theta_{t}-\theta_{\dagger}},{\mathbb{E}\left[\nabla_{\theta}\ell(\theta_{t},X_{t})\right]}\right\rangle+\left\langle{X_{t}-x_{\dagger}},{\nabla_{x}\ell(\theta_{t},X_{t})}\right\rangle$
$\displaystyle\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad=\mathbb{E}\left[\left\langle{\theta_{t}-\theta_{\dagger}},{\nabla_{\theta}\ell(\theta_{t},X_{t})}\right\rangle\right]+\left\langle{X_{t}-x_{\dagger}},{\nabla_{x}\ell(\theta_{t},X_{t})}\right\rangle,$

and, because $\ell$ is strongly log-concave by Assumption 4 and
$(\theta_{\dagger},x_{\dagger})$ maximizes it,

$\displaystyle\left\langle{\nabla\ell(\theta_{t},X_{t})},{(\theta_{t},X_{t})-(\theta_{\dagger},x_{\dagger})}\right\rangle\leq\ell(\theta_{t},X_{t})-\ell(\theta_{\dagger},x_{\dagger})-\lambda\xi_{t}/2\leq-\lambda\xi_{t}/2.$
(45)

(e.g. see pp. 63-64 in [29]). Taking expectations of (44) and applying the
above,

$\displaystyle\mathbb{E}\left[\left\langle{(\theta_{t}-\theta_{\dagger},X_{t}-x_{\dagger})},{(\mathbb{E}\left[\nabla_{\theta}\ell(\theta_{t},X_{t})\right],\nabla_{x}\ell(\theta_{t},X_{t}))}\right\rangle\right]$
$\displaystyle=\mathbb{E}\left[\left\langle{(\theta_{t}-\theta_{\dagger},X_{t}-x_{\dagger})},{\nabla\ell(\theta_{t},X_{t})}\right\rangle\right]\leq-\lambda\mathbb{E}\left[\xi_{t}\right]/2.$

In turn, taking expectations of (43) yields
$\operatorname{d\\!}\mathbb{E}\left[\xi_{t}\right]\leq[-\lambda\mathbb{E}\left[\xi_{t}\right]+d_{x}]\operatorname{d\\!}t\enskip\forall
t\geq 0$ and

$\mathbb{E}\left[\xi_{t}\right]\leq e^{-\lambda
t}\left(\mathbb{E}\left[\xi_{0}\right]-\frac{d_{x}}{\lambda}\right)+\frac{d_{x}}{\lambda}\leq\max\bigg{(}\mathbb{E}\left[\xi_{0}\right],\frac{d_{x}}{\lambda}\bigg{)}\leq\mathbb{E}\left[\xi_{0}\right]+\frac{d_{x}}{\lambda}\quad\forall
t\geq 0,$

where we used $e^{-\lambda t}\leq 1$. Now the $C_{p}$ inequality,
$\mathinner{\\!\left\lVert\theta_{t}\right\rVert}^{2}\leq
2(\mathinner{\\!\left\lVert\theta_{\dagger}\right\rVert}^{2}+\mathinner{\\!\left\lVert\theta_{t}-\theta_{\dagger}\right\rVert}^{2})$
and similarly for $\mathbb{E}\left[\mathinner{\\!\left\lVert
X_{t}\right\rVert}^{2}\right]$, gives the desired bound. ∎

###### Proof of Proposition 6.

This is an adaptation of the arguments in [7, Theorem 1.7], [8, Proposition
1], and [24, Proposition 3.1]. Fix a $T>0$ and define the map

$\Psi\mathrel{\mathop{\mathchar
58\relax}}\mathcal{C}([0,T],\mathbb{R}^{d_{\theta}})\times\mathcal{P}_{2}(\mathcal{C}([0,T],\mathbb{R}^{d_{x}}))\mapsto\mathcal{C}([0,T],\mathbb{R}^{d_{\theta}})\times\mathcal{P}_{2}(\mathcal{C}([0,T],\mathbb{R}^{d_{x}})),$

$(\theta,\nu)\mapsto\Psi(\theta,\nu)$, associating to each pair $(\theta,\nu)$
the law on path space of the solution $(\theta_{t},X_{t})_{t\leq T}$ of the
SDE (6), which is convenient to rewrite as

$\operatorname{d\\!}(\theta_{t},X_{t})=b(\theta_{t},X_{t},\nu_{t})\operatorname{d\\!}t+\sigma\operatorname{d\\!}W_{t}$

where $b(\theta_{t},X_{t},\nu_{t})\mathrel{\mathop{\mathchar
58\relax}}=[\int\nabla_{\theta}\ell(\theta_{t},x)\nu_{t}(\operatorname{d\\!}x),\nabla_{x}\ell(\theta_{t},X_{t})]^{\intercal}$,
$\sigma\mathrel{\mathop{\mathchar
58\relax}}=\sqrt{2}[0_{d_{\theta}},1_{d_{x}}]^{\intercal}$ and $\nu_{t}$ is
the time marginal of $\nu$ at $t$. Notice that $\Psi$ is well defined by
Proposition 19. We aim to show that $\Psi$ has a unique fixed point, which is
both necessary and sufficient for $(\theta_{t},X_{t})_{t\leq T}$ to solve the
SDE above in the strong sense. To do that, we first define the straightforward
extension of the metric $\mathsf{d}$ to the path space
$\mathcal{C}([0,T],\mathbb{R}^{d_{\theta}})\times\mathcal{P}_{2}(\mathcal{C}([0,T],\mathbb{R}^{d_{x}}))$:

$\mathsf{d}_{T}((\theta_{[0,T]},\nu_{[0,T]}),(\theta_{[0,T]}^{\prime},\nu_{[0,T]}^{\prime}))\mathrel{\mathop{\mathchar
58\relax}}=\sup_{t\leq
T}\sqrt{\mathinner{\\!\left\lVert\theta_{t}-\theta_{t}^{\prime}\right\rVert}^{2}+W_{2}(\nu_{t},\nu_{t}^{\prime})}.$

Let $(\theta_{t},X_{t})_{t\leq T},(\theta_{t}^{\prime},X_{t}^{\prime})_{t\leq
T}$ be two solutions of (6) satisfying $\theta_{0}=\theta_{0}^{\prime}$ and
$X_{0}=X_{0}^{\prime}$. For some rolling constant $c$ we compute by Jensen’s
inequality and then Lipschitz continuity and the Kantorovich–Rubinstein
duality formula,

$\displaystyle\mathbb{E}\left[\mathinner{\\!\left\lVert\theta_{t}-\theta_{t}^{\prime}\right\rVert}^{2}+\mathinner{\\!\left\lVert
X_{t}-X_{t}^{\prime}\right\rVert}^{2}\right]$
$\displaystyle\leq\mathbb{E}\left[\mathinner{\\!\left\lVert\int_{0}^{t}(b(\theta_{s},X_{s},\nu_{s})-b(\theta_{s}^{\prime},X_{s}^{\prime},\nu_{s}^{\prime}))\operatorname{d\\!}s\right\rVert}^{2}\right]$
$\displaystyle\leq T\int_{0}^{t}\mathbb{E}\left[\mathinner{\\!\left\lVert
b(\theta_{s},X_{s},\nu_{s})-b(\theta_{s}^{\prime},X_{s}^{\prime},\nu_{s}^{\prime})\right\rVert}^{2}\operatorname{d\\!}s\right]$
$\displaystyle\leq
cT\int_{0}^{t}\mathbb{E}\left[\mathinner{\\!\left\lVert\theta_{s}-\theta_{s}^{\prime}\right\rVert}^{2}+\mathinner{\\!\left\lVert
X_{s}-X_{s}^{\prime}\right\rVert}^{2}+W_{1}(\nu_{s},\nu_{s}^{\prime})^{2}\right]\operatorname{d\\!}s$

Next, using the fact $W_{1}\leq W_{2}$,

$\displaystyle\int_{0}^{t}W_{1}(\nu_{r},\nu_{r}^{\prime})^{2}\operatorname{d\\!}s\leq\int_{0}^{t}\sup_{r\leq
s}W_{1}(\nu_{r},\nu_{r}^{\prime})^{2}\operatorname{d\\!}s\leq\int_{0}^{T}\mathsf{d}_{s}((\theta_{[0,s]},\nu_{[0,s]}),(\theta_{[0,s]}^{\prime},\nu_{[0,s]}^{\prime}))^{2}\operatorname{d\\!}s.$

Combining these two estimates, Grönwall’s inequality then shows that for all
$t\leq T$,

$\displaystyle\mathbb{E}\left[\mathinner{\\!\left\lVert\theta_{t}-\theta_{t}^{\prime}\right\rVert}^{2}+\mathinner{\\!\left\lVert
X_{t}-X_{t}^{\prime}\right\rVert}^{2}\right]\leq
c_{T}\int_{0}^{T}\mathsf{d}_{s}((\theta_{[0,s]},\nu_{[0,s]}),(\theta_{[0,s]}^{\prime},\nu_{[0,s]}^{\prime}))^{2}\operatorname{d\\!}s.$

with $c_{T}\mathrel{\mathop{\mathchar 58\relax}}=cTe^{T}$. Using the
definition of $\Psi$ and $\mathsf{d}_{T}$, taking the supremum across $t\leq
T$,

$\displaystyle\mathsf{d}_{T}(\Psi(\theta_{[0,T]},\nu_{[0,T]}),\Psi(\theta_{[0,T]}^{\prime},\nu_{[0,T]}^{\prime}))^{2}\leq
c_{T}\int_{0}^{T}\mathsf{d}_{s}((\theta_{[0,s]},\nu_{[0,s]}),(\theta_{[0,s]}^{\prime},\nu_{[0,s]}^{\prime}))^{2}\operatorname{d\\!}s.$

Let $\Psi^{k}$ denote the $k$-fold composition of $\Psi$. Iterating the
inequality above $k$-times yields

$\displaystyle\mathsf{d}_{T}(\Psi^{k}(\theta_{[0,T]},\nu_{[0,T]}),\Psi^{k}(\theta_{[0,T]}^{\prime},\nu_{[0,T]}^{\prime}))^{2}$
$\displaystyle\leq(c_{T})^{k}\int_{0}^{T}\frac{s^{k}}{(k-1)!}\mathsf{d}_{s}((\theta_{[0,s]},\nu_{[0,s]}),(\theta_{[0,s]}^{\prime},\nu_{[0,s]}^{\prime}))^{2}\operatorname{d\\!}s$
$\displaystyle\leq\frac{(c_{T})^{k}}{k!}\mathsf{d}_{T}((\theta_{[0,T]},\nu_{[0,T]}),(\theta_{[0,T]}^{\prime},\nu_{[0,T]}^{\prime}))^{2}.$

Letting $k\rightarrow\infty$ shows that $\Psi^{k}$ is a contraction for some
$k\in\mathbb{N}$ and the first claim follows by the fixed point theorem. Now
we want to show that the law of the solution to the SDE (6) solves the
gradient flow (5) weakly, in the sense of Assumption 2. We adapt an argument
in [30, Section 8.3]. Let $\phi$ in $\mathcal{C}_{c}^{2}(\mathbb{R}^{d_{x}})$,
where $\mathcal{C}_{c}^{2}(\mathbb{R}^{d_{x}})$ is the space of twice
continuously differentiable compactly supported functions on
$\mathbb{R}^{d_{x}}$. Formula (7.3.1) in [30] shows

$\displaystyle\phi(X_{t})=\phi(X_{0})+\int_{0}^{t}\bigg{(}\left\langle{\nabla_{x}\ell(\theta_{s},X_{s})},{\nabla_{x}\phi(X_{s})}\right\rangle+\Delta_{x}\phi(X_{s})\bigg{)}\operatorname{d\\!}s+\int_{0}^{t}\left\langle{\nabla_{x}\phi(X_{s})},{\operatorname{d\\!}W_{s}}\right\rangle.$

Now consider the system

$\displaystyle\dot{\theta}_{t}=\int\nabla_{\theta}\ell(\theta_{t},x)\nu_{t}(\operatorname{d\\!}x),\quad
M_{t}\mathrel{\mathop{\mathchar
58\relax}}=\phi(X_{t})-\int_{0}^{t}\bigg{(}\left\langle{\nabla_{x}\ell(\theta_{s},X_{s})},{\nabla_{x}\phi(X_{s})}\right\rangle+\Delta_{x}\phi(X_{s})\bigg{)}\operatorname{d\\!}s$

with $\nu_{t}=\text{Law}(X_{t})$. Comparing with the expression for
$\phi(X_{t})$ above we notice that $M_{t}$ is a martingale with respect to the
natural filtration generated by $(\theta_{t},X_{t})_{t\geq 0}$, as it
corresponds to the Itô integral of an adapted, square integrable process
against Brownian motion. Taking expectations and the time derivative in
$M_{t}$ above and then integrating by parts shows that

$\displaystyle\frac{\operatorname{d\\!}}{\operatorname{d\\!}t}\int\phi(x)\nu_{t}(\operatorname{d\\!}x)$
$\displaystyle=-\int\left\langle{\nabla_{x}\ell(\theta_{t},x)},{\nabla_{x}\phi(x)}\right\rangle\nu_{t}(\operatorname{d\\!}x)+\int\Delta_{x}\phi(x)\nu_{t}(\operatorname{d\\!}x)\quad\forall\phi\in\mathcal{C}_{c}^{2}(\mathbb{R}^{d_{x}}).$

∎

###  B.2 Proof of Lemma 10

Recall that the particles $(X_{t}^{1},\dots,X^{N}_{t})$ are i.i.d. with
distribution $q_{t}$, and that the particles $(X_{*}^{1},\dots,X^{N}_{*})$ are
i.i.d. with distribution $\pi_{\theta_{*}}$. Consider the coupling of these
random vectors such that each $X_{*}^{i}$ is optimally coupled (in the sense
of the Wasserstein-2 distance) with $X_{t}^{i}$. In this case, we write

$\displaystyle\mathsf{d}((\theta_{t},Q^{N}_{t})),(\theta_{*},Q_{*}^{N}))^{2}$
$\displaystyle\leq\mathinner{\\!\left\lVert\theta_{t}-\theta_{*}\right\rVert}^{2}+\frac{1}{N}\sum_{n=1}^{N}\mathbb{E}\left[\mathinner{\\!\left\lVert
X_{t}^{n}-X_{*}^{n}\right\rVert}^{2}\right]$
$\displaystyle=\mathinner{\\!\left\lVert\theta_{t}-\theta_{*}\right\rVert}^{2}+\mathbb{E}\left[\mathinner{\\!\left\lVert
X_{t}^{1}-X_{*}^{1}\right\rVert}^{2}\right]=\mathsf{d}((\theta_{t},q_{t})),(\theta_{*},\pi_{\theta_{*}}))^{2}$

The result follows via Theorem 5 and Proposition 6. ∎

###  B.3 Proof of Lemma 11

To prove the desired bound, we start with Itô’s formula to write

$\displaystyle\operatorname{d\\!}\mathinner{\\!\left\lVert\Theta^{N}_{t}-\theta_{t}\right\rVert}^{2}$
$\displaystyle=2\left\langle{\Theta^{N}_{t}-\theta_{t}},{\frac{1}{N}\sum_{n=1}^{N}\nabla_{\theta}\ell(\Theta_{t}^{N},\bar{X}_{t}^{n})-\int\nabla_{\theta}\ell(\theta_{t},x)q_{t}(\operatorname{d\\!}x)}\right\rangle\operatorname{d\\!}t,$
$\displaystyle\operatorname{d\\!}\mathinner{\\!\left\lVert\bar{X}_{t}^{n}-X^{n}_{t}\right\rVert}^{2}$
$\displaystyle=2\left\langle{\bar{X}_{t}^{n}-X_{t}^{n}},{\nabla_{x}\ell(\Theta_{t}^{N},\bar{X}_{t}^{n})-\nabla_{x}\ell(\theta_{t},X_{t}^{n})}\right\rangle\operatorname{d\\!}t.$

Setting $\bar{\xi}_{t}^{N}\mathrel{\mathop{\mathchar
58\relax}}=N^{-1}\sum_{n=1}^{N}\mathinner{\\!\left\lVert\bar{X}_{t}^{n}-X^{n}_{t}\right\rVert}^{2}+\mathinner{\\!\left\lVert\Theta_{t}^{N}-\theta_{t}\right\rVert}^{2}$,
averaging the second equation over $n$, and adding to the first, we find that

$\displaystyle\operatorname{d\\!}\bar{\xi}_{t}^{N}=$ $\displaystyle
2\left\langle{\Theta^{N}_{t}-\theta_{t}},{\frac{1}{N}\sum_{n=1}^{N}\nabla_{\theta}\ell(\Theta_{t}^{N},\bar{X}_{t}^{n})-\int\nabla_{\theta}\ell(\theta_{t},x)q_{t}(\operatorname{d\\!}x)}\right\rangle\operatorname{d\\!}t$
$\displaystyle+\frac{2}{N}\sum_{n=1}^{N}\left\langle{\bar{X}_{t}^{n}-X_{t}^{n}},{\nabla_{x}\ell(\Theta_{t}^{N},\bar{X}_{t}^{n})-\nabla_{x}\ell(\theta_{t},X_{t}^{n})}\right\rangle\operatorname{d\\!}t.$

Because $X_{t}^{1},\dots,X_{t}^{N}$ all have law $q_{t}$,

$\displaystyle\operatorname{d\\!}\bar{\xi}_{t}^{N}=\frac{2}{N}\sum_{n=1}^{N}\Big{[}$
$\displaystyle\left\langle{\Theta^{N}_{t}-\theta_{t}},{\nabla_{\theta}\ell(\Theta_{t}^{N},\bar{X}_{t}^{n})-\nabla_{\theta}\ell(\theta_{t},X_{t}^{n})}\right\rangle$
$\displaystyle+\left\langle{\bar{X}_{t}^{n}-X_{t}^{n}},{\nabla_{x}\ell(\Theta_{t}^{N},\bar{X}_{t}^{n})-\nabla_{x}\ell(\theta_{t},X_{t}^{n})}\right\rangle\Big{]}\operatorname{d\\!}t+2G_{t}^{N}\operatorname{d\\!}t$

where

$\displaystyle G_{t}^{N}\mathrel{\mathop{\mathchar 58\relax}}$
$\displaystyle=\left\langle{\Theta^{N}_{t}-\theta_{t}},{\frac{1}{N}\sum_{n=1}^{N}\nabla_{\theta}\ell(\theta_{t},X_{t}^{n})-\int\nabla_{\theta}\ell(\theta_{t},x)q_{t}(\operatorname{d\\!}x)}\right\rangle$
$\displaystyle=\frac{1}{N}\left\langle{\Theta^{N}_{t}-\theta_{t}},{\sum_{n=1}^{N}\left[\nabla_{\theta}\ell(\theta_{t},X_{t}^{n})-\mathbb{E}\left[\nabla_{\theta}\ell(\theta_{t},X_{t}^{n})\right]\right]}\right\rangle.$

It then follows from Assumption 4 that

$\displaystyle\operatorname{d\\!}\bar{\xi}_{t}^{N}\leq[-2\lambda\bar{\xi}_{t}^{N}+2G_{t}^{N}]\operatorname{d\\!}t\enskip\forall
t\geq
0,\quad\Rightarrow\quad\frac{\operatorname{d\\!}\mathbb{E}\left[\bar{\xi}_{t}^{N}\right]}{\operatorname{d\\!}t}\leq-2\lambda\mathbb{E}\left[\bar{\xi}_{t}^{N}\right]+2\mathbb{E}\left[G_{t}^{N}\right]\enskip\forall
t\geq 0.$ (46)

As we show below,

$\left|\mathbb{E}\left[G_{t}^{N}\right]\right|\leq
L\sqrt{\frac{2\mathbb{E}\left[\bar{\xi}_{t}^{N}\right]}{N}\left(B_{0}+\frac{d_{x}}{\lambda}\right)}.$
(47)

Because
$(\operatorname{d\\!}/\operatorname{d\\!}t)\mathbb{E}\left[\bar{\xi}_{t}^{N}\right]^{1/2}=2^{-1}\mathbb{E}\left[\bar{\xi}_{t}^{N}\right]^{-1/2}(\operatorname{d\\!}/\operatorname{d\\!}t)\mathbb{E}\left[\bar{\xi}_{t}^{N}\right]$,
(46,47) imply that

$\frac{\operatorname{d\\!}}{\operatorname{d\\!}t}\mathbb{E}\left[\bar{\xi}_{t}^{N}\right]^{1/2}\leq-\lambda\mathbb{E}\left[\bar{\xi}_{t}^{N}\right]^{1/2}+L\sqrt{\frac{2}{N}\left(B_{0}+\frac{d_{x}}{\lambda}\right)}.$

Applying Grönwall’s inequality, we obtain that

$\displaystyle\mathbb{E}\left[\bar{\xi}_{t}^{N}\right]^{1/2}$
$\displaystyle\leq e^{-\lambda t}\sqrt{\bar{\xi}_{0}^{N}}+\frac{(1-e^{-\lambda
t})L}{\lambda}\sqrt{\frac{2}{N}\left(\mu_{0}+\frac{d_{x}}{\lambda}\right)},$

and the result follows because $\bar{\xi}_{0}^{N}=0$ by construction. ∎

###### Proof of (47).

Applying the Cauchy–Schwarz inequality, we obtain

$\displaystyle\left|\mathbb{E}\left[G_{t}^{N}\right]\right|^{2}$
$\displaystyle\leq\frac{1}{N^{2}}\mathbb{E}\left[\mathinner{\\!\left\lVert\Theta^{N}_{t}-\theta_{t}\right\rVert}^{2}\right]\mathbb{E}\left[\mathinner{\\!\left\lVert\sum_{n=1}^{N}\left[\nabla_{\theta}\ell(\theta_{t},X_{t}^{n})-\mathbb{E}\left[\nabla_{\theta}\ell(\theta_{t},X_{t}^{n})\right]\right]\right\rVert}^{2}\right]$
$\displaystyle\leq\frac{\mathbb{E}\left[\bar{\xi}_{t}^{N}\right]}{N^{2}}\sum_{n=1}^{N}\mathbb{E}\left[\mathinner{\\!\left\lVert\nabla_{\theta}\ell(\theta_{t},X_{t}^{n})-\mathbb{E}\left[\nabla_{\theta}\ell(\theta_{t},X_{t}^{n})\right]\right\rVert}^{2}\right]=\mathrel{\mathop{\mathchar
58\relax}}\frac{\mathbb{E}\left[\bar{\xi}_{t}^{N}\right]}{N^{2}}\sum_{n=1}^{N}c_{t}^{n},$

where the last inequality follows from the independence of
$X_{t}^{1},\dots,X_{t}^{N}$. Let $X_{t}^{\prime}$ denote a random variable
independent of $(X_{t}^{1},\dots,X_{t}^{n})$ with law $q_{t}$. Jensen’s
inequality and Lipschitz continuity of $\nabla_{\theta}\ell$ (Assumption 5)
imply

$\displaystyle c_{t}^{n}$
$\displaystyle=\mathbb{E}\left[\mathinner{\\!\left\lVert\nabla_{\theta}\ell(\theta_{t},X_{t}^{n})-\mathbb{E}\left[\nabla_{\theta}\ell(\theta_{t},X_{t}^{n})\right]\right\rVert}^{2}\right]\leq\mathbb{E}\left[\mathinner{\\!\left\lVert\nabla_{\theta}\ell(\theta_{t},X_{t}^{n})-\nabla_{\theta}\ell(\theta_{t},X_{t}^{\prime})\right\rVert}^{2}\right]$
$\displaystyle\leq L^{2}\mathbb{E}\left[\mathinner{\\!\left\lVert
X_{t}^{n}-X_{t}^{\prime}\right\rVert}^{2}\right]\leq
2L^{2}\mathbb{E}\left[\mathinner{\\!\left\lVert
X_{t}^{\prime}\right\rVert}^{2}\right].$

Combining the above two and Proposition 19, we then obtain (47). ∎

###  B.4 Proof of Lemma 12

The following proof extends the arguments in [16, Lemma S2] or [11, Chapter
4.1]. We introduce the linear interpolation of the Euler-Maruyama
discretization of (32):

$\displaystyle\tilde{X}_{t}^{n}=$
$\displaystyle\tilde{X}_{kh}^{n}+(t-kh)\nabla_{x}\ell(\tilde{X}_{kh}^{n},\tilde{\Theta}_{kh}^{N})+\sqrt{2}(W^{n}_{t}-W^{n}_{kh}),\quad\forall
n\in[N];$ (48) $\displaystyle\tilde{\Theta}_{t}^{N}=$
$\displaystyle\tilde{\Theta}_{kh}^{N}+(t-kh)\frac{1}{N}\sum_{n=1}^{N}\nabla_{\theta}\ell(\tilde{X}_{kh}^{n},\tilde{\Theta}_{kh}^{N});$

for all $t\in[kh,(k+1)h)$ and $k\in\mathbb{N}$. Let
$\tilde{Q}_{t}^{h}\mathrel{\mathop{\mathchar
58\relax}}=N^{-1}\sum_{i=1}\delta_{\tilde{X}_{t}^{n}}$. We notice that
$(\tilde{\Theta}_{Kh}^{N},\tilde{Q}_{Kh}^{h})$ coincides in distribution with
$(\Theta^{N,h}_{K},Q^{N,h}_{K})$, hence we just need to derive the bound

$\mathsf{d}((\tilde{\Theta}_{Kh}^{N},\tilde{Q}_{Kh}^{h}),(\Theta^{N}_{Kh},\bar{Q}^{N}_{Kh}))\leq\sqrt{h}A_{0,h}\quad\forall
K\in\mathbb{N}.$

For all $n\in[N]$, $k\in\mathbb{N}$ and $h>0$ we compute directly from the
defining equations,

$\displaystyle\mathinner{\\!\left\lVert\bar{X}_{(k+1)h}^{n}-\tilde{X}_{(k+1)h}^{n}\right\rVert}^{2}=\mathinner{\\!\left\lVert\int_{kh}^{(k+1)h}[\nabla_{x}\ell(\bar{X}_{s}^{n},\Theta_{s}^{N})-\nabla_{x}\ell(\tilde{X}_{kh}^{n},\tilde{\Theta}_{kh}^{N})]\operatorname{d\\!}s\right\rVert}^{2}$
$\displaystyle-2h\left\langle{\bar{X}_{kh}^{n}-\tilde{X}_{kh}^{n}},{\nabla_{x}\ell(\bar{X}_{kh}^{n},\Theta_{kh}^{N})-\nabla_{x}\ell(\tilde{X}_{kh}^{n},\tilde{\Theta}_{kh}^{N})}\right\rangle$
$\displaystyle-2\int_{kh}^{(k+1)h}\left\langle{\bar{X}_{kh}^{n}-\tilde{X}_{kh}^{n}},{\nabla_{x}\ell(\bar{X}_{s}^{n},\Theta_{s}^{N})-\nabla_{x}\ell(\bar{X}_{kh}^{n},\Theta_{kh}^{N})}\right\rangle\operatorname{d\\!}s+\mathinner{\\!\left\lVert\bar{X}_{kh}^{n}-\tilde{X}_{kh}^{n}\right\rVert}^{2},$
$\displaystyle\mathinner{\\!\left\lVert\Theta_{(k+1)h}^{N}-\tilde{\Theta}_{(k+1)h}^{N}\right\rVert}^{2}=\frac{1}{N}\sum_{n=1}^{N}\mathinner{\\!\left\lVert
2\int_{kh}^{(k+1)h}[\nabla_{\theta}\ell(\bar{X}_{s}^{n},\Theta_{s}^{N})-\nabla_{\theta}\ell(\tilde{X}_{kh}^{n},\tilde{\Theta}_{kh}^{N})]\operatorname{d\\!}s\right\rVert}^{2}$
$\displaystyle-\frac{2h}{N}\sum_{n=1}^{N}\left\langle{\Theta_{kh}^{N}-\tilde{\Theta}_{kh}^{N}},{\nabla_{\theta}\ell(\bar{X}_{kh}^{n},\Theta_{kh}^{N})-\nabla_{\theta}\ell(\tilde{X}_{kh}^{n},\tilde{\Theta}_{kh}^{N})}\right\rangle$
$\displaystyle-\frac{2}{N}\sum_{n=1}^{N}\int_{kh}^{(k+1)h}\left\langle{\Theta_{kh}^{N}-\tilde{\Theta}_{kh}^{N}},{[\nabla_{\theta}\ell(\bar{X}_{s}^{n},\Theta_{s}^{N})-\nabla_{\theta}\ell(\bar{X}_{kh}^{n},\Theta_{kh}^{N})]}\right\rangle\operatorname{d\\!}s+\mathinner{\\!\left\lVert\Theta_{kh}^{N}-\tilde{\Theta}_{kh}^{N}\right\rVert}^{2}.$

Averaging the $N$ equations for
$\mathinner{\\!\left\lVert\bar{X}_{(k+1)h}^{n}-\tilde{X}_{(k+1)h}^{n}\right\rVert}^{2}$,
adding the one for
$\mathinner{\\!\left\lVert\Theta_{(k+1)h}^{N}-\tilde{\Theta}_{(k+1)h}^{N}\right\rVert}^{2}$
we obtain, with $\bar{Y}_{t}^{n}\mathrel{\mathop{\mathchar
58\relax}}=(\bar{X}_{t}^{n},\Theta_{t}^{N})$,
$\tilde{Y}^{n}_{t}\mathrel{\mathop{\mathchar
58\relax}}=(\tilde{X}_{t}^{n},\tilde{\Theta}_{t}^{N})$ and
$\xi_{t}\mathrel{\mathop{\mathchar
58\relax}}=N^{-1}\sum_{n=1}^{N}\mathinner{\\!\left\lVert\tilde{Y}_{t}^{n}-\bar{Y}_{t}^{n}\right\rVert}^{2}$,

$\displaystyle\xi_{(k+1)h}=$
$\displaystyle\xi_{kh}+\frac{1}{N}\sum_{n=1}^{N}\mathinner{\\!\left\lVert\int_{kh}^{(k+1)h}[\nabla\ell(\bar{Y}_{s}^{n})-\nabla\ell(\tilde{Y}^{n}_{kh})]\operatorname{d\\!}s\right\rVert}^{2}$
$\displaystyle-\frac{2h}{N}\sum_{n=1}^{N}\left\langle{\bar{Y}_{kh}^{n}-\tilde{Y}_{kh}^{n}},{\nabla\ell(\bar{Y}_{kh}^{n})-\nabla\ell(\tilde{Y}_{kh}^{n})}\right\rangle$
$\displaystyle-\frac{2}{N}\sum_{n=1}^{N}\int_{kh}^{(k+1)h}\left\langle{\bar{Y}_{kh}^{n}-\tilde{Y}_{kh}^{n}},{\nabla\ell(\bar{Y}_{s}^{n})-\nabla\ell(\bar{Y}_{kh}^{n})}\right\rangle\operatorname{d\\!}s.$
(49)

Now, adding and subtracting $\nabla\ell(\bar{Y}_{kh}^{n})$, applying Jensen’s
inequality, expanding the resulting integrand and further applying Young’s
inequalities,

$\displaystyle\mathinner{\\!\left\lVert\int_{kh}^{(k+1)h}[\nabla\ell(\bar{Y}_{s}^{n})-\nabla\ell(\tilde{Y}^{n}_{kh})]\operatorname{d\\!}s\right\rVert}^{2}$
$\displaystyle\leq
2h^{2}\mathinner{\\!\left\lVert\nabla\ell(\bar{Y}_{kh}^{n})-\nabla\ell(\tilde{Y}_{kh}^{n})\right\rVert}^{2}$
$\displaystyle+2h\int_{kh}^{(k+1)h}\mathinner{\\!\left\lVert\nabla\ell(\bar{Y}_{s}^{n})-\nabla\ell(\bar{Y}_{kh}^{n})\right\rVert}^{2}\operatorname{d\\!}s.$
(50)

Furthermore, under Assumptions 4 and 5 we have, by [29, Theorem 2.1.12], there
holds the co-coercivity property

$\left\langle{\nabla\ell(y^{\prime})-\nabla\ell(y)},{y^{\prime}-y}\right\rangle\geq\frac{\iota}{2}\mathinner{\\!\left\lVert
y-y^{\prime}\right\rVert}^{2}+\frac{\mathinner{\\!\left\lVert\nabla\ell(y^{\prime})-\nabla\ell(y)\right\rVert}^{2}}{\lambda+L},$

where $\iota\mathrel{\mathop{\mathchar 58\relax}}=2\lambda L/(\lambda+L)$. It
follows that, whenever $h<1/(\lambda+L)$, we have

$\displaystyle\frac{1}{N}\sum_{n=1}^{N}\left[2h^{2}\mathinner{\\!\left\lVert\nabla\ell(\bar{Y}_{kh}^{n})-\nabla\ell(\tilde{Y}_{kh}^{n})\right\rVert}^{2}\right.$
$\displaystyle\left.-2h\left\langle{\bar{Y}_{kh}^{n}-\tilde{Y}_{kh}^{n}},{\nabla\ell(\bar{Y}_{kh}^{n})-\nabla\ell(\tilde{Y}_{kh}^{n})}\right\rangle\right]$
$\displaystyle\leq-\frac{h\iota}{N}\sum_{n=1}^{N}\mathinner{\\!\left\lVert\bar{Y}_{kh}^{n}-\tilde{Y}_{kh}^{n}\right\rVert}^{2}=-h\iota\xi_{kh}$
(51)

Combining (49,50,51), we obtain

$\displaystyle\xi_{(k+1)h}\leq$ $\displaystyle(1-\iota
h)\xi_{kh}-\frac{2}{N}\sum_{n=1}^{N}\left[\int_{kh}^{(k+1)h}\left\langle{\bar{Y}_{kh}^{n}-\tilde{Y}_{kh}^{n}},{\nabla\ell(\bar{Y}_{s}^{n})-\nabla\ell(\bar{Y}_{kh}^{n})}\right\rangle\operatorname{d\\!}s\right.$
$\displaystyle\left.+2h\int_{kh}^{(k+1)h}\mathinner{\\!\left\lVert\nabla\ell(\bar{Y}_{s}^{n})-\nabla\ell(\bar{Y}_{kh}^{n})\right\rVert}^{2}\operatorname{d\\!}s\right].$

To deal with the remaining terms, fix any $\epsilon>0$. Applying Young’s
inequality, we obtain

$\displaystyle\frac{1}{N}\sum_{n=1}^{N}$
$\displaystyle\left|\left\langle{\bar{Y}_{kh}^{n}-\tilde{Y}_{kh}^{n}},{\nabla\ell(\bar{Y}_{s}^{n})-\nabla\ell(\bar{Y}_{kh}^{n})}\right\rangle\right|$
$\displaystyle\leq\frac{1}{N}\sum_{n=1}^{N}\left[\frac{\epsilon}{2}\mathinner{\\!\left\lVert\bar{Y}_{kh}^{n}-\tilde{Y}_{kh}^{n}\right\rVert}^{2}+\frac{1}{2\epsilon}\mathinner{\\!\left\lVert\nabla\ell(\bar{Y}_{s}^{n})-\nabla\ell(\bar{Y}_{kh}^{n})\right\rVert}^{2}\right]$
$\displaystyle=\frac{\epsilon}{2}\xi_{kh}+\frac{1}{2\epsilon
N}\sum_{n=1}^{N}\mathinner{\\!\left\lVert\nabla\ell(\bar{Y}_{s}^{n})-\nabla\ell(\bar{Y}_{kh}^{n})\right\rVert}^{2}.$

Putting the above two together, we obtain for any $\epsilon>0$ and
$k\in\mathbb{N}$,

$\displaystyle\xi_{(k+1)h}\leq(1-\iota
h+h\epsilon/2)\xi_{kh}+\left(2h+\frac{1}{2\epsilon}\right)\int_{kh}^{(k+1)h}\frac{1}{N}\sum_{n=1}^{N}\mathinner{\\!\left\lVert\nabla\ell(\bar{Y}_{s}^{n})-\nabla\ell(\bar{Y}_{kh}^{n})\right\rVert}^{2}\operatorname{d\\!}s.$
(52)

We will show below that if $h<1/(\lambda+L)$ then, for all $s\in[kh,(k+1)h)$,

$\displaystyle\frac{1}{N}\sum_{n=1}^{N}\mathbb{E}\left[\mathinner{\\!\left\lVert\nabla\ell(\bar{Y}_{s}^{n})-\nabla\ell(\bar{Y}_{kh}^{n})\right\rVert}^{2}\right]\leq
220hL^{2}(L^{2}h(B_{0}+d_{x}/\lambda+d_{x}).$ (53)

Hence, taking expectations in (52), choosing $\epsilon=\iota$ using the bound
above and the fact that $\xi_{0}=0$ we obtain

$\displaystyle\mathbb{E}\left[\xi_{(k+1)h}\right]$
$\displaystyle\leq(2h+2/\iota)220h^{2}(L^{2}h(B_{0}+d_{x}/\lambda+d_{x})\sum_{j=1}^{k}(1-\iota
h/2)^{j}$
$\displaystyle\leq\frac{4h+4/\iota}{\iota}220hL^{2}(L^{2}h(B_{0}+d_{x}/\lambda+d_{x}).$

∎

###### Proof of (53).

Let $\\{\mathcal{F}_{t};t\geq 0\\}$ be the filtration generated by
$(\theta_{t},\bar{X}_{t}^{n})_{n=1}^{N}$ and denote $\mathbb{E}_{kh}[\cdot]$
expectation conditional on $\mathcal{F}_{kh}$. We first show that whenever
$h<1/(\lambda+L)$

$\frac{1}{N}\sum_{n=1}^{N}\mathbb{E}_{kh}\left[\mathinner{\\!\left\lVert\nabla\ell(\bar{Y}_{s}^{n})-\nabla\ell(\bar{Y}_{kh}^{n})\right\rVert}^{2}\right]\leq\frac{1}{N}\sum_{n=1}^{N}220hL^{2}(L^{2}h\mathinner{\\!\left\lVert\bar{Y}^{n}_{kh}\right\rVert}^{2}+d_{x}),$
(54)

after which the result follows from taking expectations and applying
Proposition 20. Since

$\displaystyle\mathbb{E}_{kh}$
$\displaystyle\left[\mathinner{\\!\left\lVert\bar{X}_{s}^{n}-\bar{X}_{kh}^{n}\right\rVert}^{2}\right]=\mathbb{E}_{kh}\left[\mathinner{\\!\left\lVert\int_{kh}^{s}\nabla_{x}\ell(\Theta_{r}^{N},\bar{X}^{n}_{r})\operatorname{d\\!}r+\sqrt{2}(W^{n}_{s}-W^{n}_{kh})\right\rVert}^{2}\right]$
$\displaystyle\leq
2(s-kh)\int_{kh}^{s}\mathbb{E}_{kh}\left[\mathinner{\\!\left\lVert\nabla_{x}\ell(\Theta_{r}^{N},\bar{X}^{n}_{r})\right\rVert}^{2}\right]\operatorname{d\\!}r+4d_{x}(s-kh)$

for all $n\in[N]$, and

$\displaystyle\mathinner{\\!\left\lVert\Theta_{s}^{N}-\Theta_{kh}^{N}\right\rVert}^{2}=\mathinner{\\!\left\lVert\frac{1}{N}\sum_{n=1}^{N}\int_{kh}^{s}\nabla_{\theta}\ell(\Theta_{r}^{N},\bar{X}_{r}^{n})\operatorname{d\\!}r\right\rVert}^{2}\leq
2(s-kh)\frac{1}{N}\sum_{n=1}^{N}\int_{kh}^{s}\mathinner{\\!\left\lVert\nabla_{\theta}\ell(\Theta_{r}^{N},\bar{X}_{r}^{n})\right\rVert}^{2}\operatorname{d\\!}r,$

adding up and since Assumption 5 implies
$\mathinner{\\!\left\lVert\nabla\ell(\bar{Y}_{r}^{n})\right\rVert}^{2}\leq
L^{2}\mathinner{\\!\left\lVert\bar{Y}_{r}^{n}\right\rVert}^{2}$, we obtain

$\displaystyle\frac{1}{N}$
$\displaystyle\sum_{n=1}^{N}\mathbb{E}_{kh}\left[\mathinner{\\!\left\lVert\bar{Y}^{n}_{s}-\bar{Y}^{n}_{kh}\right\rVert}^{2}\right]\leq
2hL^{2}\int_{kh}^{s}\frac{1}{N}\sum_{n=1}^{N}\mathbb{E}_{kh}\left[\mathinner{\\!\left\lVert\bar{Y}^{n}_{r}\right\rVert}^{2}\right]\operatorname{d\\!}r+4d_{x}h$
$\displaystyle\leq
4hL^{2}\int_{kh}^{s}\frac{1}{N}\sum_{n=1}^{N}\mathbb{E}_{kh}\left[\mathinner{\\!\left\lVert\bar{Y}^{n}_{r}-\bar{Y}^{n}_{kh}\right\rVert}^{2}\right]\operatorname{d\\!}r+4h^{2}L^{2}\mathbb{E}_{kh}\left[\mathinner{\\!\left\lVert\bar{Y}^{n}_{kh}\right\rVert}^{2}\right]+4d_{x}h$

and by Grönwall lemma

$\displaystyle\frac{1}{N}\sum_{n=1}^{N}\mathbb{E}_{kh}$
$\displaystyle\left[\mathinner{\\!\left\lVert\bar{Y}^{n}_{s}-\bar{Y}^{n}_{kh}\right\rVert}^{2}\right]\leq\frac{1}{N}\sum_{n=1}^{N}\exp(4L^{2}h^{2})(4L^{2}h^{2}\mathinner{\\!\left\lVert\bar{Y}^{n}_{kh}\right\rVert}^{2}+4d_{x}h).$

If $h<1/(\lambda+L)$, $\exp(4L^{2}h^{2})\leq 55$, and (54) then follows by
Assumption 5. ∎

######  Proposition 20.

Assume Assumptions 4–5. Then, for all $t\geq 0$,

$\frac{1}{N}\sum_{n=1}^{N}\mathbb{E}\left[\mathinner{\\!\left\lVert
X_{t}^{n}\right\rVert}^{2}+\mathinner{\\!\left\lVert\Theta_{t}^{N}\right\rVert}^{2}\right]\leq
2\left(\mathinner{\\!\left\lVert\theta_{0}\right\rVert}^{2}+\mathbb{E}\left[\mathinner{\\!\left\lVert\bar{X}_{0}\right\rVert}^{2}\right]+\frac{d_{x}}{\lambda}\right)$

###### Proof.

Recalling that here $(\theta_{\dagger},x_{\dagger})=(0,0)$, this argument is
very similar to that supporting Proposition 19. For each $n\in[N]$,

$\displaystyle\operatorname{d\\!}\mathinner{\\!\left\lVert
X_{t}^{n}\right\rVert}^{2}$
$\displaystyle=-2\left\langle{X_{t}^{n}},{\nabla_{x}\ell(\Theta_{t}^{N},X_{t}^{n})}\right\rangle\operatorname{d\\!}t+\sqrt{2}\operatorname{d\\!}W_{t}^{n}+2d_{x}\operatorname{d\\!}t$
$\displaystyle\operatorname{d\\!}\mathinner{\\!\left\lVert\Theta_{t}^{N}\right\rVert}^{2}$
$\displaystyle=-\frac{2}{N}\sum_{n=1}^{N}\left\langle{\Theta_{t}^{N}},{\nabla_{\theta}\ell(\Theta_{t}^{N},X_{t}^{n})}\right\rangle\operatorname{d\\!}t$

now averaging the $N$ equations for $\mathinner{\\!\left\lVert
X_{t}^{n}\right\rVert}^{2}$ and then adding
$\mathinner{\\!\left\lVert\Theta_{t}^{N}\right\rVert}^{2}$, taking
expectations and time derivatives, with
$\xi_{t}^{N}=N^{-1}\sum_{n=1}^{N}\mathinner{\\!\left\lVert
X_{t}^{n}\right\rVert}^{2}+\mathinner{\\!\left\lVert\Theta_{t}^{N}\right\rVert}^{2}$,

$\displaystyle\frac{\operatorname{d\\!}}{\operatorname{d\\!}t}\mathbb{E}\left[\xi_{t}^{N}\right]=-2\mathbb{E}\left[\frac{1}{N}\sum_{n=1}^{N}\left\langle{(X^{n}_{t},\Theta_{t}^{N})},{\nabla\ell(\Theta_{t}^{N},X^{n}_{t})}\right\rangle\right]+2d_{x}\leq-2\lambda\mathbb{E}\left[\xi_{t}^{N}\right]+2d_{x},$

and then Grönwall concludes. ∎