moreover satisfy for real numbers $0\leq\lambda^{s}<\lambda^{u}$,

$\textstyle\sup_{t\geq 0}\lVert e^{\mathbf{A}^{s}t}\rVert
e^{-\lambda^{s}t}<\infty,\quad\sup_{t\geq 0}\lVert e^{-\mathbf{A}^{u}t}\rVert
e^{\lambda^{u}t}<\infty.$

Further assume there exists a spectral gap of $\lambda^{u}>k\lambda^{s}$ and
that $\mathscr{E}^{s}$ has the $C^{k}$ extension property. Let
$\operatorname{\mathcal{B}}_{r},\operatorname{\mathcal{B}}_{r}^{s},\operatorname{\mathcal{B}}_{r}^{u}$
denote the balls of radius $r$ around the origin in
$\mathscr{E},\mathscr{E}^{s},\mathscr{E}^{u}$, respectively. Then for
sufficiently small $r>0$, there exists a $C^{k}$ map
$h:\operatorname{\mathcal{B}}_{r}^{s}\to\operatorname{\mathcal{B}}_{r}^{u}$
with $h(0)=0$, $Dh(0)=0$ whose graph
$\mathscr{V}\subset\operatorname{\mathcal{B}}_{r}$ (the local center-stable
manifold) has the following properties.

  1. cmrmn(i)

(Invariance) For all initial values $\bm{z}_{0}\in\mathscr{V}$ there exists a
$C^{1}$ curve $\bm{z}_{t}:\operatorname{\mathbb{R}}_{\geq 0}\to\mathscr{E}$
such that as long as $\bm{z}_{t}\in\operatorname{\mathcal{B}}_{r}$, then
$\bm{z}_{t}\in\mathscr{V}$ and (11) holds.

  2. cmrmn(ii)

(Uniqueness) If $\bm{z}_{t}$ is any solution of (11) such that
$\bm{z}_{t}\in\operatorname{\mathcal{B}}_{r}$ for all $t\geq 0$, then
$\bm{z}_{t}\in\mathscr{V}$ for all $t\geq 0$.

#### Proof of Theorem 4.3.

Let $\mu^{\dagger}\in\mathscr{G}^{\dagger}$ be a strict saddle point. We apply
the local center-stable manifold theorem to the system (5) on
$L^{2}(\Omega,\mu^{\dagger};\operatorname{\mathbb{R}}^{k+d})$. By the spectral
theorem, the operator $\operatorname{\mathcal{H}}_{\mu^{\dagger}}$ has a
complete set of eigenvalues $\lambda_{j}$ and corresponding eigenfunctions
$\psi_{j}$ for $j\in\operatorname{\mathbb{Z}}$, ordered such that

$\lambda_{1}\geq\lambda_{2}\geq\cdots\geq
0,\quad\lambda_{0}=\cdots=\lambda_{-(p-1)}<\lambda_{-p}\leq\cdots\leq 0.$

Since the spectrum may possess a limit point at 0, we cannot separate
$\operatorname{\mathcal{H}}_{\mu^{\dagger}}$ into absolutely convergent and
divergent components. Instead, we set the cutoff at the largest negative
eigenvalue
$\lambda_{0}=\lambda_{\text{min}}(\operatorname{\mathcal{H}}_{\mu^{\dagger}})$,
taking all possibly multiple eigenvalues, and defining the subspace
$\mathscr{E}^{u}$ as the span of the corresponding
$\psi_{0},\cdots,\psi_{-(p-1)}$. Then we are guaranteed a jump
$\lambda_{-(p-1)}<\lambda_{-p}$ since the spectrum is discrete, and we choose
$\lambda_{s}=-\lambda_{-p},\lambda^{u}=-\lambda_{0}$ and
$k\in(1,|\lambda_{0}/\lambda_{-p}|)$ so that the spectral gap condition is
satisfied – we only need continuity (i.e. $k\geq 0$) for our argument.
Moreover, the $C^{k}$ extension property for $\mathscr{E}^{s}$ holds
automatically as $L^{2}(\Omega,\mu^{\dagger};\operatorname{\mathbb{R}}^{k+d})$
is a Hilbert space. Therefore, any convergent local flow $(\bm{v}_{t})$
defined in an open neighborhood $\operatorname{\mathcal{B}}_{\mu^{\dagger}}$
must be contained in a graph
$\mathscr{V}_{\mu^{\dagger}}\subset\operatorname{\mathcal{B}}_{\mu^{\dagger}}$
containing $\mu^{\dagger}$.

The rest of the proof is similar to Lee et al. (2019). Since the collection
$\\{\operatorname{\mathcal{B}}_{\mu^{\dagger}}:\mu^{\dagger}\in\mathscr{G}^{\dagger}\\}$
forms an open cover of $\mathscr{G}^{\dagger}$ and
$\operatorname{\mathscr{P}}_{2}(\Omega)$ is separable with respect to
2-Wasserstein distance (Ambrosio et al., 2005, Proposition 7.1.5), we can
extract a countable subcover
$\\{\operatorname{\mathcal{B}}_{j}:j\in\operatorname{\mathbb{N}}\\}$
containing $\mathscr{G}^{\dagger}$. If the WGF $(\mu_{t})_{t\geq 0}$ converges
to a strict saddle point, there exists an index $j$ and an integer threshold
$\ell$ such that $\mu_{t}\in\operatorname{\mathcal{B}}_{j}$ for $t\geq\ell$.
In particular, $\mu_{t}$ must be contained in the corresponding center-stable
manifold $\mathscr{V}_{j}$ for $t\geq\ell$.

Let $\omega_{t}^{-}(\nu)$ denote the result of running the reversed gradient
flow $\partial\nu_{-t}=-\nabla\cdot(\nu_{-t}\nabla\frac{\delta
F}{\delta\mu}(\nu_{-t}))$, $\nu_{0}=\nu$ for time $t$ whenever it exists; time
inversion $t\mapsto-t$ shows that $\omega_{t}^{-}(\mu_{t})=\mu_{0}$ for the
forward flow $(\mu_{t})_{t\geq 0}$. Since $\mu_{\ell}\in\mathscr{V}_{j}$ for
some integer time $\ell$ and $\mathscr{V}_{j}$, it holds that

$\displaystyle\mathscr{G}_{0}^{\dagger}\subseteq\bigcup_{j\in\operatorname{\mathbb{N}}}\bigcup_{\ell\in\operatorname{\mathbb{N}}}\omega_{\ell}^{-}(\mathscr{V}_{j}),$

hence $\mathscr{G}_{0}^{\dagger}$ must be contained in the countable union of
images of graphs. ∎

##  Appendix E Proofs for Section 5

###  E.1 First-order Improvement

######  Proposition E.1.

Let $F$ be a functional depending on $\mu$ only through the MLP layer
$h_{\mu}$. Suppose MFD (1) at time $t$ admits a distribution
$\bar{\mu}\in\operatorname{\mathscr{P}}_{2}(\Theta)$ with
$\chi^{2}(\bar{\mu},\mu_{t})\leq\bar{\chi}^{2}$ such that along the linear
homotopy $\bar{\mu}_{s}=(1-s)\mu_{t}+s\bar{\mu}$ we have
$\frac{\mathop{}\\!\mathrm{d}}{\mathop{}\\!\mathrm{d}s}\big{|}_{s=0}F(\bar{\mu}_{s})\leq-\delta\leq
0$. Then
$\frac{\mathop{}\\!\mathrm{d}}{\mathop{}\\!\mathrm{d}t}F(\mu_{t})\leq-\bar{\chi}^{-2}\delta^{2}$.

###### Proof.

We may express $F$ as $F(\mu)=J(h_{\mu})$ for an auxiliary functional
$h\mapsto J(h)$ defined on
$C(\operatorname{\mathscr{X}},\operatorname{\mathbb{R}}^{d})$, which implies
that

$\frac{\delta F}{\delta\mu}(\mu,\theta)=\int\frac{\delta J}{\delta
h}(h_{\mu},\bm{x})^{\top}h_{\theta}(\bm{x})\mathop{}\\!\mathrm{d}\bm{x}.$

In particular, since the dependency on the second layer $\bm{a}$ is linear, it
always holds that $\bm{a}^{\top}\nabla_{\bm{a}}\frac{\delta
F}{\delta\mu}=\frac{\delta F}{\delta\mu}$. We can then directly lower bound
the decrease rate of the objective under (1) by isolating the gradient
provided by $\bm{a}$:

$\displaystyle\frac{\mathop{}\\!\mathrm{d}}{\mathop{}\\!\mathrm{d}t}F(\mu_{t})$
$\displaystyle=\int\frac{\delta
F}{\delta\mu}(\mu_{t},\theta)\partial_{t}\mu_{t}(\mathop{}\\!\mathrm{d}\theta)$
$\displaystyle=-\int\left\lVert\nabla_{\theta}\frac{\delta
F}{\delta\mu}(\mu_{t},\theta)\right\rVert^{2}\mu_{t}(\mathop{}\\!\mathrm{d}\theta)$
$\displaystyle\leq-\int\left\lVert\nabla_{\bm{a}}\frac{\delta
F}{\delta\mu}(\mu_{t},\theta)\right\rVert^{2}\mu_{t}(\mathop{}\\!\mathrm{d}\theta)$
$\displaystyle\leq-\int\left(\bm{a}^{\top}\nabla_{\bm{a}}\frac{\delta
F}{\delta\mu}(\mu_{t},\theta)\right)^{2}\mu_{t}(\mathop{}\\!\mathrm{d}\theta)$
$\displaystyle=-\int\left(\frac{\delta
F}{\delta\mu}(\mu_{t},\theta)\right)^{2}\mu_{t}(\mathop{}\\!\mathrm{d}\theta).$

Starting from the first-order condition, by the Cauchy-Schwarz inequality we
can also bound

$\displaystyle\left(\frac{\mathop{}\\!\mathrm{d}}{\mathop{}\\!\mathrm{d}s}\bigg{|}_{s=0}F(\bar{\mu}_{s})\right)^{2}=\left(\int\frac{\delta
F}{\delta\mu}(\mu_{t},\theta)(\bar{\mu}-\mu_{t})(\mathop{}\\!\mathrm{d}\theta)\right)^{2}\leq\chi^{2}(\bar{\mu},\mu_{t})\int\left(\frac{\delta
F}{\delta\mu}(\mu_{t},\theta)\right)^{2}\mu_{t}(\mathop{}\\!\mathrm{d}\theta).$

Joining the two inequalities gives the desired bound. ∎

#### Proof of Proposition 5.2.

Recall that the functional derivative is computed as

$\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu,\theta)=-\mathbb{E}_{\bm{x}}\left[\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}h_{\theta}(\bm{x})\right],$
(12)

where the additive constant has been implicitly normalized such that the
integral with respect to the current measure $\mu$ is zero, i.e.
$\int\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu)\mathop{}\\!\mathrm{d}\mu=0$
as shown in the proof of Theorem 3.3. Due to the spherical symmetry of $\pi$
in the first component, it is also immediate that

$\int\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu)\mathop{}\\!\mathrm{d}\pi=-\mathbb{E}_{\bm{x}}\left[\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}h_{\pi}(\bm{x})\right]=0.$

The chi-square divergence between $\bar{\mu}=\mathbf{R}\sharp\mu^{\circ}$ and
$\mu_{t}$ can be bounded as

$\int\left(\frac{\mathop{}\\!\mathrm{d}\bar{\mu}}{\mathop{}\\!\mathrm{d}\mu_{t}}-1\right)^{2}\mathop{}\\!\mathrm{d}\mu_{t}\leq\left\lVert\frac{\mathop{}\\!\mathrm{d}\bar{\mu}}{\mathop{}\\!\mathrm{d}\mu_{t}}\right\rVert_{\infty}-1\leq\gamma^{-1}\left\lVert\frac{\mathop{}\\!\mathrm{d}\bar{\mu}}{\mathop{}\\!\mathrm{d}\pi}\right\rVert_{\infty}$

where the birth-death mechanism prevents the density ratio
$\frac{\mathop{}\\!\mathrm{d}\mu_{t}}{\mathop{}\\!\mathrm{d}\pi}$ from falling
below the threshold $\gamma$ at any point. Writing the convex decomposition of
$\mathbf{R}$ in the sense of Lemma C.6 as
$\sum_{j=1}^{m}\alpha_{j}\mathbf{R}_{j}$ with
$\mathbf{R}_{j}\in\mathcal{O}(k)$, the density of $\bar{\mu}$ relative to
$\pi$ is further bounded as

$\left\lVert\frac{\mathop{}\\!\mathrm{d}\bar{\mu}}{\mathop{}\\!\mathrm{d}\pi}\right\rVert_{\infty}=\left\lVert\frac{\mathop{}\\!\mathrm{d}\mathbf{R}\sharp\mu^{\circ}}{\mathop{}\\!\mathrm{d}\pi}\right\rVert_{\infty}\leq\sum_{j=1}^{m}\alpha_{j}\left\lVert\frac{\mathop{}\\!\mathrm{d}\mu^{\circ}}{\mathop{}\\!\mathrm{d}\pi}\right\rVert_{\infty}\leq
R_{4}$

by the spherical symmetry of $\pi$. Hence we may apply Proposition E.1 with
$\bar{\chi}^{2}=\gamma^{-1}R_{4}$, showing that the objective decreases along
MFD by a rate of at least
$\frac{\mathop{}\\!\mathrm{d}}{\mathop{}\\!\mathrm{d}t}\operatorname{\mathcal{L}}(\mu_{t})\leq-
R_{4}^{-1}\gamma\delta^{2}$.

Moreover, whenever the discrete linear update is performed, along the homotopy
$\hat{\mu}_{s}:=(1-s\gamma)\mu_{t}+s\gamma\pi$ we have

$\frac{\mathop{}\\!\mathrm{d}}{\mathop{}\\!\mathrm{d}s}\operatorname{\mathcal{L}}(\hat{\mu}_{s})=\gamma\int\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\hat{\mu}_{s},\theta)(\pi-\mu_{t})(\mathop{}\\!\mathrm{d}\theta)=0.$

Hence $t\mapsto\operatorname{\mathcal{L}}(\mu_{t})$ is unaffected by the
discrete updates, justifying the inequality for all time $t\geq 0$. ∎

We note that if the forcing term is applied continuously as in Remark 5.1,
almost the exact same proof for Proposition E.1 applies by bounding

$\frac{\mathop{}\\!\mathrm{d}}{\mathop{}\\!\mathrm{d}t}\operatorname{\mathcal{L}}(\mu_{t})=-\int\left\lVert\nabla_{\theta}\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu_{t},\theta)\right\rVert^{2}\mu_{t}(\mathop{}\\!\mathrm{d}\theta)+\gamma\int\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu_{t},\theta)(\pi-\mu_{t})(\mathop{}\\!\mathrm{d}\theta)\leq-\int\left(\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu_{t},\theta)\right)^{2}\mu_{t}(\mathop{}\\!\mathrm{d}\theta).$

As we mentioned briefly, the proof can also be easily modified to handle
unbounded second layer $\bm{a}$ by invoking the Cauchy-Schwarz inequality to
lower bound the gradient

$\int\left\lVert\nabla_{\bm{a}}\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu_{t},\theta)\right\rVert^{2}\mu_{t}(\mathop{}\\!\mathrm{d}\theta)\geq\left(\int\lVert\bm{a}\rVert^{2}\mu_{t}(\mathop{}\\!\mathrm{d}\theta)\right)^{-1}\int\left(\bm{a}^{\top}\nabla_{\bm{a}}\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu_{t},\theta)\right)^{2}\mu_{t}(\mathop{}\\!\mathrm{d}\theta)$

and bounding the second moment uniformly in time with the following result,

######  Lemma E.2.

Denote the second moment of $\mu\in\operatorname{\mathscr{P}}_{2}(\Theta)$
along the $\bm{a}$ component as
$m_{\bm{a}}(\mu)=\int\lVert\bm{a}\rVert^{2}\mu(\mathop{}\\!\mathrm{d}\theta)$.
Then the mean-field dynamics $\mu_{t}$ for all time $t\geq 0$ satisfies
$m_{\bm{a}}(\mu_{t})\leq m_{\bm{a}}(\mu_{0})\vee m_{\bm{a}}(\pi)$.

###### Proof.

In fact, $m_{\bm{a}}(\cdot)$ remains unchanged by gradient flow:

$\displaystyle\frac{\mathop{}\\!\mathrm{d}}{\mathop{}\\!\mathrm{d}t}m_{\bm{a}}(\mu_{t})$
$\displaystyle=\int\lVert\bm{a}\rVert^{2}\partial_{t}\mu_{t}(\mathop{}\\!\mathrm{d}\theta)$
$\displaystyle=-2\int(\bm{a}\;\;0_{d})^{\top}\nabla_{\theta}\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu_{t},\theta)\mu_{t}(\mathop{}\\!\mathrm{d}\theta)$
$\displaystyle=-2\int\bm{a}^{\top}\nabla_{\bm{a}}\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu_{t},\theta)\mu_{t}(\mathop{}\\!\mathrm{d}\theta)$
$\displaystyle=-2\int\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu_{t},\theta)\mu_{t}(\mathop{}\\!\mathrm{d}\theta)=0.$

Also if the discrete update is performed, the output satisfies
$m_{\bm{a}}((1-\gamma)\mu_{t}+\gamma\pi)=(1-\gamma)m_{\bm{a}}(\mu_{t})+\gamma
m_{\bm{a}}(\pi)$ by linearity of the moment functional $\mu\mapsto
m_{\bm{a}}(\mu)$. Hence $m_{\bm{a}}(\mu_{t})$ always interpolates between
$m_{\bm{a}}(\mu_{0})$ and $m_{\bm{a}}(\pi)$. ∎

#### Proof of Theorem 5.3.

Suppose $\operatorname{\mathcal{L}}(\mu_{0})\leq 0.49r^{\circ}$. Then
$\operatorname{\mathcal{L}}(\mu_{t})\leq 0.49r^{\circ}$ for all $t\geq 0$ and
by Proposition 3.4 we are guaranteed a direction of improvement
$\bar{\mu}_{s}=(1-s)\mu_{t}+s\bar{\mu}$ with $\bar{\mu}=\mathbf{R}\sharp\mu$
for some $\mathbf{R}\in\operatorname{\mathcal{B}}_{1}(k)$ such that

$\frac{\mathop{}\\!\mathrm{d}}{\mathop{}\\!\mathrm{d}s}\bigg{|}_{s=0}\operatorname{\mathcal{L}}(\bar{\mu}_{s})\leq-\frac{4}{R_{1}^{2}}\operatorname{\mathcal{L}}(\mu_{t})\left(\frac{r^{\circ}}{2}-\operatorname{\mathcal{L}}(\mu_{t})\right).$

Proposition 5.2 then ensures the objective decreases along the Wasserstein
flow as

$\frac{\mathop{}\\!\mathrm{d}}{\mathop{}\\!\mathrm{d}t}\operatorname{\mathcal{L}}(\mu_{t})\leq-\frac{16\gamma}{R_{1}^{4}R_{4}}\operatorname{\mathcal{L}}(\mu_{t})^{2}\left(\frac{r^{\circ}}{2}-\operatorname{\mathcal{L}}(\mu_{t})\right)^{2},\quad
0\leq\operatorname{\mathcal{L}}(\mu_{t})\leq\frac{r^{\circ}}{2}.$

We now divide the band into two halves.

  1. cmrmn(i)

$\frac{r^{\circ}}{4}\leq\operatorname{\mathcal{L}}\leq\frac{r^{\circ}}{2}$
(acceleration band). By substituting
$\operatorname{\mathcal{L}}(\mu_{t})^{2}\geq\frac{r^{\circ 2}}{16}$ above and
solving the differential inequality, we obtain

$\operatorname{\mathcal{L}}(\mu_{t})\leq\frac{r^{\circ}}{2}-\left(\frac{100}{r^{\circ}}-\frac{r^{\circ
2}\gamma t}{R_{1}^{4}R_{4}}\right)^{-1}$

and hence $\operatorname{\mathcal{L}}(\mu_{t})$ decreases below
$\frac{r^{\circ}}{4}$ after time $t_{1}\leq\frac{96R_{1}^{4}R_{4}}{r^{\circ
3}\gamma}$.

  2. cmrmn(ii)

$0\leq\operatorname{\mathcal{L}}\leq\frac{r^{\circ}}{4}$ (deceleration band).
By substituting
$(\frac{r^{\circ}}{2}-\operatorname{\mathcal{L}}(\mu_{t}))^{2}\geq\frac{r^{\circ
2}}{16}$ we likewise obtain

$\operatorname{\mathcal{L}}(\mu_{t})\leq\left(\frac{4}{r^{\circ}}+\frac{r^{\circ
2}\gamma(t-t_{1})}{R_{1}^{4}R_{4}}\right)^{-1}$

and hence $\operatorname{\mathcal{L}}(\mu_{t})$ achieves loss $\leq\epsilon$
after time $t_{1}+\frac{R_{1}^{4}R_{4}}{r^{\circ
2}\gamma}\cdot\frac{1}{\epsilon}$.

Finally, note that the second term dominates the first since
$\epsilon=O(r^{\circ})$. ∎

###  E.2 Second-order Improvement

#### Proof of Lemma 5.4.

It is straightforward to show that

$\displaystyle\partial_{t}\left[\nabla_{\theta}\frac{\delta
F}{\delta\mu}(\mu_{t},\theta)\right]$
$\displaystyle=\nabla_{\theta}\int\frac{\delta^{2}\\!F}{\delta\mu^{2}}(\mu_{t},\theta,\theta^{\prime})(\partial_{t}\mu_{t})(\mathop{}\\!\mathrm{d}\theta^{\prime})$
$\displaystyle=-\nabla_{\theta}\int\nabla_{\theta^{\prime}}\frac{\delta^{2}\\!F}{\delta\mu^{2}}(\mu_{t},\theta,\theta^{\prime})\cdot\nabla_{\theta^{\prime}}\frac{\delta
F}{\delta\mu}(\mu_{t},\theta^{\prime})\mu_{t}(\mathop{}\\!\mathrm{d}\theta^{\prime})$
$\displaystyle=-\int\mathbf{H}_{\mu_{t}}(\theta,\theta^{\prime})\nabla_{\theta^{\prime}}\frac{\delta
F}{\delta\mu}(\mu_{t},\theta^{\prime})\mu_{t}(\mathop{}\\!\mathrm{d}\theta^{\prime}).$

Each term is well-defined as soon as the kernel is assumed to be Hilbert-
Schmidt, or due to Lemma E.3 for the case $F=\operatorname{\mathcal{L}}$. ∎

######  Lemma E.3.

The kernel $\mathbf{H}_{\mu}$ for the functional $\operatorname{\mathcal{L}}$
is Hilbert-Schmidt for all $\mu\in\operatorname{\mathscr{P}}_{2}^{+}(\Theta)$.
Moreover, the corresponding integral operator
$\operatorname{\mathcal{H}}_{\mu}f(\theta)=\int\mathbf{H}_{\mu}(\theta,\theta^{\prime})f(\theta^{\prime})\mu(\mathop{}\\!\mathrm{d}\theta^{\prime})$
is compact self-adjoint, hence there exists an orthonormal basis
$\\{\psi_{j}\\}_{j\in\operatorname{\mathbb{Z}}}$ for
$L^{2}(\Theta,\mu;\operatorname{\mathbb{R}}^{k+d})$ consisting of
eigenfunctions of $\operatorname{\mathcal{H}}_{\mu}$.

###### Proof.

We extend our notation to write for example
$\bm{\Sigma}_{\mu,\theta}=\mathbb{E}_{\bm{x}\sim\operatorname{\mathcal{D}}_{\operatorname{\mathscr{X}}}}[h_{\mu}(\bm{x})h_{\theta}(\bm{x})^{\top}]$.
From (12) the second order functional derivative can be derived as

$\displaystyle\frac{\delta^{2}\\!\\!\operatorname{\mathcal{L}}}{\delta\mu^{2}}(\mu,\theta,\theta^{\prime})$
$\displaystyle=-\frac{\delta}{\delta\mu}\mathbb{E}_{\bm{x}}\left[(h_{\mu^{\circ}}(\bm{x})-\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}h_{\mu}(\bm{x}))^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}h_{\theta}(\bm{x})\right](\theta^{\prime})$
$\displaystyle=-\operatorname{tr}\left(\bm{\Sigma}_{\mu^{\circ},\theta^{\prime}}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\theta,\mu^{\circ}}\right)$
$\displaystyle\qquad+\operatorname{tr}\left(\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}(\bm{\Sigma}_{\theta^{\prime},\mu}+\bm{\Sigma}_{\mu,\theta^{\prime}})\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\theta,\mu^{\circ}}\right)$
$\displaystyle\qquad+\operatorname{tr}\left(\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\mu,\mu^{\circ}}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\theta,\theta^{\prime}}\right)$
$\displaystyle\qquad-\operatorname{tr}\left(\bm{\Sigma}_{\mu,\mu}^{-1}(\bm{\Sigma}_{\theta^{\prime},\mu}+\bm{\Sigma}_{\mu,\theta^{\prime}})\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\mu,\mu^{\circ}}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\theta,\mu}\right)$
$\displaystyle\qquad+\operatorname{tr}\left(\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\theta^{\prime},\mu^{\circ}}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\theta,\mu}\right)$
$\displaystyle\qquad+\operatorname{tr}\left(\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\mu,\mu^{\circ}}\bm{\Sigma}_{\mu^{\circ},\theta^{\prime}}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\theta,\mu}\right)$
$\displaystyle\qquad-\operatorname{tr}\left(\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\mu,\mu^{\circ}}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}(\bm{\Sigma}_{\theta^{\prime},\mu}+\bm{\Sigma}_{\mu,\theta^{\prime}})\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\theta,\mu}\right).$

It is tedious but straightforward to check that this expression is symmetric
in $\theta,\theta^{\prime}$ (which would otherwise follow directly if we had a
priori second order regularity estimates for $\operatorname{\mathcal{L}}$). We
then have that

$[\mathbf{H}_{\mu}(\theta,\theta^{\prime})]_{i,j}=\partial_{\theta_{i}}\partial_{\theta_{j}^{\prime}}\frac{\delta^{2}\\!\\!\operatorname{\mathcal{L}}}{\delta\mu^{2}}(\mu,\theta,\theta^{\prime})=\partial_{\theta_{j}^{\prime}}\partial_{\theta_{i}}\frac{\delta^{2}\\!\\!\operatorname{\mathcal{L}}}{\delta\mu^{2}}(\mu,\theta^{\prime},\theta)=[\mathbf{H}_{\mu}(\theta^{\prime},\theta)]_{j,i}$

which implies $\operatorname{\mathcal{H}}_{\mu}$ is self-adjoint as before.
For the proof of the first claim, we refer to the uniform spectral bound for
$\mathbf{H}_{\mu}$ obtained in Lemma E.5; this also shows that
$\operatorname{\mathcal{H}}_{\mu}$ is compact. ∎

In Lemma E.4 and E.5, we derive various regularity bounds of the ICFL
objective $\operatorname{\mathcal{L}}$. The constants $C_{1},\cdots,C_{5}$,
numbered as to be consistent with Theorem E.7, are explicitly defined during
the proofs and have at most polynomial dependency on all problem constants.

######  Lemma E.4.

The gradients of the functional derivative of $\operatorname{\mathcal{L}}$ at
any $\mu\in\operatorname{\mathscr{P}}_{2}^{+}(\Theta)$ such that
$\lambda_{\textup{min}}(\bm{\Sigma}_{\mu,\mu})\geq\lambda$ uniformly satisfy
$\lVert\nabla_{\bm{a}}\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}\rVert\leq
C_{\bm{a}}$,
$\lVert\nabla_{\bm{w}}\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}\rVert\leq
C_{\bm{w}}$ and
$\lVert\nabla\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}\rVert\leq
C_{1}$. Moreover,
$\nabla\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}$ is
$C_{2}$-Lipschitz on $\Theta$, where
$C_{\bm{a}},C_{2}=O(\frac{1}{(k\lambda)^{1/2}})$ and
$C_{\bm{w}},C_{1}=O(\frac{1}{k\lambda})$.

###### Proof.

The gradient with respect to each component is given by

$\nabla_{\bm{a}}\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu,\theta)=-\mathbb{E}_{\bm{x}}\left[\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\sigma(\bm{w}^{\top}\bm{x})\right]^{\top},\quad\nabla_{\bm{w}}\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu,\theta)=-\mathbb{E}_{\bm{x}}\left[\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{a}\sigma^{\prime}(\bm{w}^{\top}\bm{x})\bm{x}\right].$

Hence we can bound

$\displaystyle\left\lVert\nabla_{\bm{a}}\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu,\theta)\right\rVert^{2}$
$\displaystyle\leq\mathbb{E}_{\bm{x}}\left[\lVert\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\sigma(\bm{w}^{\top}\bm{x})\rVert^{2}\right]$
$\displaystyle\leq
R_{1}^{2}\cdot\mathbb{E}_{\bm{x}}\left[\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-2}\bm{\Sigma}_{\mu,\mu^{\circ}}\zeta_{\mu^{\circ},\mu}(\bm{x})\right]$
$\displaystyle\leq\frac{R_{1}^{2}}{\lambda}\operatorname{tr}\left(\mathbf{L}_{\mu}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\mu,\mu^{\circ}}\right)$
$\displaystyle\leq\frac{r^{\circ}R_{1}^{2}}{\lambda}\operatorname{\mathcal{L}}(\mu)-\frac{2R_{1}^{2}}{\lambda}\operatorname{tr}\mathbf{L}_{\mu}^{2}$
$\displaystyle\leq\frac{kr^{\circ 2}R_{1}^{2}}{2\lambda}=:C_{\bm{a}}^{2},$

and also

$\displaystyle\left\lVert\nabla_{\bm{w}}\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu,\theta)\right\rVert^{2}$
$\displaystyle\leq\mathbb{E}_{\bm{x}}\left[\lVert\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{a}\sigma^{\prime}(\bm{w}^{\top}\bm{x})\bm{x}\rVert^{2}\right]$
$\displaystyle\leq
R_{2}^{2}\cdot\mathbb{E}_{\bm{x}}\left[\lVert\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\rVert^{2}\lVert\bm{x}\rVert^{2}\right]$
$\displaystyle\leq\frac{R_{2}^{2}}{\lambda}\cdot\mathbb{E}_{\bm{x}}\left[\lVert\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1/2}\rVert^{4}\right]^{1/2}\mathbb{E}_{\bm{x}}[\lVert\bm{x}\rVert^{4}]^{1/2}$
$\displaystyle\leq\frac{R_{2}^{2}M_{4}^{1/2}}{\lambda}\left(\operatorname{tr}\mathbb{E}_{\bm{x}}\left[\left(\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\mu,\mu^{\circ}}\zeta_{\mu^{\circ},\mu}(\bm{x})\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\right)^{2}\right]\right)^{1/2}$
$\displaystyle\leq\frac{r^{\circ}R_{2}^{2}M_{4}^{1/2}}{\lambda}\left(\operatorname{tr}\mathbb{E}_{\bm{x}}\left[\zeta_{\mu^{\circ},\mu}(\bm{x})\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\mu,\mu^{\circ}}\zeta_{\mu^{\circ},\mu}(\bm{x})\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\right]\right)^{1/2}$
$\displaystyle\qquad-\frac{2R_{2}^{2}M_{4}^{1/2}}{\lambda}\left(\operatorname{tr}\mathbf{L}_{\mu}\mathbb{E}_{\bm{x}}\left[\zeta_{\mu^{\circ},\mu}(\bm{x})\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\mu,\mu^{\circ}}\zeta_{\mu^{\circ},\mu}(\bm{x})\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\right]\right)^{1/2}$
$\displaystyle\leq\frac{r^{\circ}R_{2}^{2}M_{4}^{1/2}}{\lambda}\left(r^{\circ}\operatorname{tr}\mathbb{E}_{\bm{x}}\left[\left(\zeta_{\mu^{\circ},\mu}(\bm{x})\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\right)^{2}\right]-2\operatorname{tr}\mathbf{L}_{\mu}\mathbb{E}_{\bm{x}}\left[\left(\zeta_{\mu^{\circ},\mu}(\bm{x})\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\right)^{2}\right]\right)^{1/2}$
$\displaystyle\leq\frac{r^{\circ
3/2}R_{2}^{2}M_{4}^{1/2}}{\lambda}\textstyle\sup_{\bm{x}}\lVert\zeta_{\mu^{\circ},\mu}(\bm{x})\rVert(2\operatorname{tr}\mathbf{L}_{\mu})^{1/2}$
$\displaystyle\leq\frac{2k^{1/2}r^{\circ
5/2}R_{1}^{3}R_{2}^{2}M_{4}^{1/2}}{\lambda^{2}}=:C_{\bm{w}}^{2},$

where for the last line we have used the coarser bounds
$\lVert\zeta_{\mu^{\circ},\mu}(\bm{x})\rVert\leq R_{1}+R_{1}^{3}\lambda^{-1}$
and
$\lambda\leq\frac{1}{k}\operatorname{tr}\bm{\Sigma}_{\mu_{0},\mu_{0}}\leq\frac{R_{1}^{2}}{k}$.
Combining the two bounds yields

$\left\lVert\nabla\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu,\theta)\right\rVert\leq\left(\frac{r^{\circ
2}R_{1}^{3}}{2\lambda^{2}}(R_{1}+4k^{1/2}r^{\circ
1/2}R_{2}^{2}M_{4}^{1/2})\right)^{1/2}=:C_{1}.$

Furthermore, for $\theta_{1}=(\bm{a}_{1},\bm{w}_{1})$,
$\theta_{2}=(\bm{a}_{2},\bm{w}_{2})$ we have

$\displaystyle\left\lVert\nabla_{\bm{a}}\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu,\theta_{1})-\nabla_{\bm{a}}\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu,\theta_{2})\right\rVert$
$\displaystyle=\left\lVert\mathbb{E}_{\bm{x}}\left[\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}(\sigma(\bm{w}_{1}^{\top}\bm{x})-\sigma(\bm{w}_{2}^{\top}\bm{x}))\right]\right\rVert$
$\displaystyle\leq
R_{2}\cdot\mathbb{E}_{\bm{x}}\left[\lVert\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\rVert\cdot\lVert\bm{w}_{1}^{\top}\bm{x}-\bm{w}_{2}^{\top}\bm{x}\rVert\right]$
$\displaystyle\leq R_{2}M_{2}^{1/2}\left(\frac{kr^{\circ
2}}{2\lambda}\right)^{1/2}\lVert\bm{w}_{1}-\bm{w}_{2}\rVert,$

and also

$\displaystyle\left\lVert\nabla_{\bm{w}}\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu,\theta_{1})-\nabla_{\bm{w}}\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu,\theta_{2})\right\rVert$
$\displaystyle=\left\lVert\mathbb{E}_{\bm{x}}\left[\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}(\bm{a}_{1}\sigma^{\prime}(\bm{w}_{1}^{\top}\bm{x})-\bm{a}_{2}\sigma^{\prime}(\bm{w}_{2}^{\top}\bm{x}))\bm{x}\right]\right\rVert$
$\displaystyle\leq\left\lVert\mathbb{E}_{\bm{x}}\left[\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}(\bm{a}_{1}-\bm{a}_{2})\sigma^{\prime}(\bm{w}_{1}^{\top}\bm{x})\bm{x}\right]\right\rVert$
$\displaystyle\qquad+\left\lVert\mathbb{E}_{\bm{x}}\left[\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{a}_{2}(\sigma^{\prime}(\bm{w}_{1}^{\top}\bm{x})-\sigma^{\prime}(\bm{w}_{2}^{\top}\bm{x}))\bm{x}\right]\right\rVert$
$\displaystyle\leq R_{2}M_{2}^{1/2}\left(\frac{kr^{\circ
2}}{2\lambda}\right)^{1/2}\lVert\bm{a}_{1}-\bm{a}_{2}\rVert$
$\displaystyle\qquad+R_{3}\cdot\mathbb{E}_{\bm{x}}\left[\lVert\zeta_{\mu^{\circ},\mu}(\bm{x})^{\top}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\rVert\cdot\lVert\bm{w}_{1}-\bm{w}_{2}\rVert\lVert\bm{x}\rVert^{2}\right]$
$\displaystyle\leq R_{2}M_{2}^{1/2}\left(\frac{kr^{\circ
2}}{2\lambda}\right)^{1/2}\lVert\bm{a}_{1}-\bm{a}_{2}\rVert+R_{3}M_{4}^{1/2}\left(\frac{kr^{\circ
2}}{2\lambda}\right)^{1/2}\lVert\bm{w}_{1}-\bm{w}_{2}\rVert.$

Combining the two yields that

$\left\lVert\nabla\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu,\theta_{1})-\nabla\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu,\theta_{2})\right\rVert\leq\left(\frac{kr^{\circ
2}}{2\lambda}(2R_{2}^{2}M_{2}+R_{3}^{2}M_{4})\right)^{1/2}\lVert\theta_{1}-\theta_{2}\rVert=:C_{2}\,\lVert\theta_{1}-\theta_{2}\rVert.$

∎

######  Lemma E.5.

For any $\mu\in\operatorname{\mathscr{P}}_{2}^{+}(\Theta)$ such that
$\lambda_{\textup{min}}(\bm{\Sigma}_{\mu,\mu})\geq\lambda$ it holds that
$\lVert\mathbf{H}_{\mu}(\theta,\theta^{\prime})\rVert\leq C_{3}$,
$\mathbf{H}_{\mu}(\theta,\theta^{\prime})$ is uniformly $C_{4}$-Lipschitz
w.r.t. $\theta$ and $\theta^{\prime}$, and $\mathbf{H}_{\mu}$ is
$C_{5}$-Lipschitz w.r.t. $\mu$ in 1-Wasserstein distance, where
$C_{3},C_{4}=O(\lambda^{-2})$ and $C_{5}=O(d\lambda^{-3})$.

###### Proof.

To derive regularity estimates of $\mathbf{H}_{\mu}$, we start from the
expansion in Lemma E.3 and perform explicit computations for only the first
trace term
$t(\mu,\theta,\theta^{\prime})=\operatorname{tr}\left(\bm{\Sigma}_{\mu^{\circ},\theta^{\prime}}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\theta,\mu^{\circ}}\right)$.
$\nabla_{\theta}\nabla_{\theta^{\prime}}t(\mu,\theta,\theta^{\prime})$
consists of block matrices

$\displaystyle\nabla_{\bm{a}}\nabla_{\bm{a}^{\prime}}t(\mu,\theta,\theta^{\prime})$
$\displaystyle=\mathbb{E}_{\bm{x}}\left[\sigma(\bm{w}^{\top}\bm{x})h_{\mu^{\circ}}(\bm{x})^{\top}\right]\mathbb{E}_{\bm{x}}\left[\sigma(\bm{w}^{\prime\top}\bm{x})h_{\mu^{\circ}}(\bm{x})\right]\bm{\Sigma}_{\mu,\mu}^{-1},$
$\displaystyle\nabla_{\bm{a}}\nabla_{\bm{w}^{\prime}}t(\mu,\theta,\theta^{\prime})$
$\displaystyle=\bm{\Sigma}_{\mu,\mu}^{-1}\bm{a}^{\prime}\mathbb{E}_{\bm{x}}\left[\sigma(\bm{w}^{\top}\bm{x})h_{\mu^{\circ}}(\bm{x})^{\top}\right]\mathbb{E}_{\bm{x}}\left[\sigma^{\prime}(\bm{w}^{\prime\top}\bm{x})h_{\mu^{\circ}}(\bm{x})\bm{x}^{\top}\right],$
$\displaystyle\nabla_{\bm{w}}\nabla_{\bm{a}^{\prime}}t(\mu,\theta,\theta^{\prime})$
$\displaystyle=\mathbb{E}_{\bm{x}}\left[\sigma^{\prime}(\bm{w}^{\top}\bm{x})\bm{x}h_{\mu^{\circ}}(\bm{x})^{\top}\right]\mathbb{E}_{\bm{x}}\left[\sigma(\bm{w}^{\prime\top}\bm{x})h_{\mu^{\circ}}(\bm{x})\right]\bm{a}^{\top}\bm{\Sigma}_{\mu,\mu}^{-1},$
$\displaystyle\nabla_{\bm{w}}\nabla_{\bm{w}^{\prime}}t(\mu,\theta,\theta^{\prime})$
$\displaystyle=\mathbb{E}_{\bm{x}}\left[\sigma^{\prime}(\bm{w}^{\top}\bm{x})\bm{x}h_{\mu^{\circ}}(\bm{x})^{\top}\right]\mathbb{E}_{\bm{x}}\left[\sigma^{\prime}(\bm{w}^{\prime\top}\bm{x})h_{\mu^{\circ}}(\bm{x})\bm{x}^{\top}\right]\bm{a}^{\prime\top}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{a}.$

It follows from Lemma C.1 that
$\lVert\nabla_{\theta}\nabla_{\theta^{\prime}}t(\mu,\theta,\theta^{\prime})\rVert\leq(R_{1}^{4}+2R_{1}^{2}R_{2}k^{1/2}r^{\circ
1/2}M_{2}^{1/2}+R_{2}^{2}kr^{\circ}M_{2})\lambda^{-1}=O(\lambda^{-1})$. Each
term of $\mathbf{H}_{\mu}$ is likewise uniformly bounded so that
$\mathbf{H}_{\mu}$ is a valid kernel.

The Lipschitz constant of
$\nabla_{\theta}\nabla_{\theta^{\prime}}t(\mu,\theta,\theta^{\prime})$ w.r.t.
$\theta$ can also be controlled by separately bounding

$\displaystyle\lVert\nabla_{\bm{a}}\nabla_{\bm{a}^{\prime}}t(\mu,\theta_{1},\theta^{\prime})-\nabla_{\bm{a}}\nabla_{\bm{a}^{\prime}}t(\mu,\theta_{2},\theta^{\prime})\rVert$
$\displaystyle\qquad\leq\mathbb{E}_{\bm{x}}\left[|\sigma(\bm{w}_{1}^{\top}\bm{x})-\sigma(\bm{w}_{2}^{\top}\bm{x})|\cdot\lVert
h_{\mu^{\circ}}(\bm{x})\rVert\right]\mathbb{E}_{\bm{x}}\left[|\sigma(\bm{w}^{\prime\top}\bm{x})|\cdot\lVert
h_{\mu^{\circ}}(\bm{x})\rVert\right]\lVert\bm{\Sigma}_{\mu,\mu}^{-1}\rVert$
$\displaystyle\qquad\leq
R_{1}R_{2}M_{2}^{1/2}kr^{\circ}\lambda^{-1}\lVert\bm{w}_{1}-\bm{w}_{2}\rVert,$
$\displaystyle\lVert\nabla_{\bm{a}}\nabla_{\bm{w}^{\prime}}t(\mu,\theta_{1},\theta^{\prime})-\nabla_{\bm{a}}\nabla_{\bm{w}^{\prime}}t(\mu,\theta_{2},\theta^{\prime})\rVert$
$\displaystyle\qquad\leq\lVert\bm{\Sigma}_{\mu,\mu}^{-1}\bm{a}^{\prime}\rVert\cdot\mathbb{E}_{\bm{x}}\left[|\sigma(\bm{w}_{1}^{\top}\bm{x})-\sigma(\bm{w}_{2}^{\top}\bm{x})|\cdot\lVert
h_{\mu^{\circ}}(\bm{x})\rVert\right]\mathbb{E}_{\bm{x}}\left[|\sigma^{\prime}(\bm{w}^{\prime\top}\bm{x})|\cdot\lVert
h_{\mu^{\circ}}(\bm{x})\bm{x}^{\top}\rVert\right]$ $\displaystyle\qquad\leq
R_{2}^{2}M_{2}kr^{\circ}\lambda^{-1}\lVert\bm{w}_{1}-\bm{w}_{2}\rVert,$
$\displaystyle\lVert\nabla_{\bm{w}}\nabla_{\bm{a}^{\prime}}t(\mu,\theta_{1},\theta^{\prime})-\nabla_{\bm{w}}\nabla_{\bm{a}^{\prime}}t(\mu,\theta_{2},\theta^{\prime})\rVert$
$\displaystyle\qquad\leq\mathbb{E}_{\bm{x}}\left[|\sigma^{\prime}(\bm{w}_{1}^{\top}\bm{x})-\sigma^{\prime}(\bm{w}_{2}^{\top}\bm{x})|\cdot\lVert\bm{x}h_{\mu^{\circ}}(\bm{x})^{\top}\rVert\right]\mathbb{E}_{\bm{x}}\left[|\sigma(\bm{w}^{\prime\top}\bm{x})|\cdot\lVert
h_{\mu^{\circ}}(\bm{x})\rVert\right]\lVert\bm{a}_{1}^{\top}\bm{\Sigma}_{\mu,\mu}^{-1}\rVert$
$\displaystyle\qquad\qquad+\mathbb{E}_{\bm{x}}\left[|\sigma^{\prime}(\bm{w}_{2}^{\top}\bm{x})|\cdot\lVert\bm{x}h_{\mu^{\circ}}(\bm{x})^{\top}\rVert\right]\mathbb{E}_{\bm{x}}\left[|\sigma(\bm{w}^{\prime\top}\bm{x})|\cdot\lVert
h_{\mu^{\circ}}(\bm{x})\rVert\right]\lVert(\bm{a}_{1}-\bm{a}_{2})^{\top}\bm{\Sigma}_{\mu,\mu}^{-1}\rVert$
$\displaystyle\qquad\leq
R_{1}R_{3}M_{4}^{1/2}kr^{\circ}\lambda^{-1}\lVert\bm{w}_{1}-\bm{w}_{2}\rVert+R_{1}R_{2}M_{2}^{1/2}kr^{\circ}\lambda^{-1}\lVert\bm{a}_{1}-\bm{a}_{2}\rVert,$
$\displaystyle\lVert\nabla_{\bm{w}}\nabla_{\bm{w}^{\prime}}t(\mu,\theta_{1},\theta^{\prime})-\nabla_{\bm{w}}\nabla_{\bm{w}^{\prime}}t(\mu,\theta_{2},\theta^{\prime})\rVert$
$\displaystyle\qquad\leq\mathbb{E}_{\bm{x}}\left[|\sigma^{\prime}(\bm{w}_{1}^{\top}\bm{x})-\sigma^{\prime}(\bm{w}_{2}^{\top}\bm{x})|\cdot\lVert\bm{x}h_{\mu^{\circ}}(\bm{x})^{\top}\rVert\right]\mathbb{E}_{\bm{x}}\left[|\sigma^{\prime}(\bm{w}^{\prime\top}\bm{x})|\cdot\lVert
h_{\mu^{\circ}}(\bm{x})\bm{x}^{\top}\rVert\right]\lVert\bm{a}^{\prime\top}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{a}_{1}\rVert$
$\displaystyle\qquad\qquad+\mathbb{E}_{\bm{x}}\left[|\sigma^{\prime}(\bm{w}_{2}^{\top}\bm{x})|\cdot\lVert\bm{x}h_{\mu^{\circ}}(\bm{x})^{\top}\rVert\right]\mathbb{E}_{\bm{x}}\left[|\sigma^{\prime}(\bm{w}^{\prime\top}\bm{x})|\cdot\lVert
h_{\mu^{\circ}}(\bm{x})\bm{x}^{\top}\rVert\right]\lVert\bm{a}^{\prime\top}\bm{\Sigma}_{\mu,\mu}^{-1}(\bm{a}_{1}-\bm{a}_{2})\rVert$
$\displaystyle\qquad\leq
R_{2}R_{3}M_{2}^{1/2}M_{4}^{1/2}kr^{\circ}\lambda^{-1}\lVert\bm{w}_{1}-\bm{w}_{2}\rVert+R_{2}^{2}M_{2}kr^{\circ}\lambda^{-1}\lVert\bm{a}_{1}-\bm{a}_{2}\rVert.$

Therefore,
$\nabla_{\theta}\nabla_{\theta^{\prime}}t(\mu,\theta,\theta^{\prime})$ is
uniformly $O(\lambda^{-1})$-Lipschitz w.r.t. both $\theta$ and
$\theta^{\prime}$ by symmetry. All the remaining terms can also be bounded
with at most an $O(\lambda^{-2})$ Lipschitz constant; in particular, the terms
including three factors of $\bm{\Sigma}_{\mu,\mu}^{-1}$ can be controlled by
removing a factor of $\lambda^{-1/2}$ twice and isolating
$\bm{\Sigma}_{\mu,\mu}^{-1/2}\bm{\Sigma}_{\mu,\mu^{\circ}}$ and
$\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1/2}$ as in the proof of
Lemma E.4.

Finally, the third-order functional derivative
$\nabla_{\tilde{\theta}}\frac{\delta}{\delta\mu}\mathbf{H}_{\mu}(\theta,\theta^{\prime})(\tilde{\theta})$
can be bounded in a similar manner with spectral norm at most
$O(\lambda^{-3})$, yielding via Kantorovich-Rubinstein duality that

$\lVert\mathbf{H}_{\mu_{1}}(\theta,\theta^{\prime})-\mathbf{H}_{\mu_{2}}(\theta,\theta^{\prime})\rVert=\left\lVert\int\frac{\delta}{\delta\mu}\mathbf{H}_{(1-s)\mu_{1}+s\mu_{2}}(\theta,\theta^{\prime})(\tilde{\theta})(\mu_{2}-\mu_{1})(\mathop{}\\!\mathrm{d}\tilde{\theta})\right\rVert\lesssim(k+d)\lambda^{-3}\cdot\operatorname{\mathcal{W}}_{1}(\mu_{1},\mu_{2}).$

The additional $k+d$ factor arises from bounding each entry of
$\mathbf{H}_{\mu_{1}}-\mathbf{H}_{\mu_{2}}$ separately. We omit the details. ∎

######  Proposition E.6.

Let $F$ be a functional depending on $\mu$ only through the MLP layer
$h_{\mu}$. Suppose MFD (1) at time $t$ admits a distribution
$\bar{\mu}\in\operatorname{\mathscr{P}}_{2}(\Theta)$ with
$\chi^{2}(\bar{\mu},\mu_{t})\leq\bar{\chi}^{2}$ such that
$\frac{\mathop{}\\!\mathrm{d}^{2}}{\mathop{}\\!\mathrm{d}s^{2}}\big{|}_{s=0}F(\bar{\mu}_{s})\leq-\Lambda$.
Then the smallest eigenvalue $\lambda_{0}$ of
$\operatorname{\mathcal{H}}_{\mu_{t}}$ satisfies
$\lambda_{0}\leq-\bar{\chi}^{-2}\Lambda$.

###### Proof.

The second derivative along the linear homotopy $\bar{\mu}_{s}$ can be
expanded as

$\displaystyle\frac{\mathop{}\\!\mathrm{d}^{2}}{\mathop{}\\!\mathrm{d}s^{2}}\bigg{|}_{s=0}F(\bar{\mu}_{s})$
$\displaystyle=\frac{\mathop{}\\!\mathrm{d}}{\mathop{}\\!\mathrm{d}s}\bigg{|}_{s=0}\int\frac{\delta
F}{\delta\mu}(\bar{\mu}_{s},\theta)(\bar{\mu}-\mu_{t})(\mathop{}\\!\mathrm{d}\theta)$
$\displaystyle=\iint\frac{\delta^{2}\\!F}{\delta\mu^{2}}(\mu_{t},\theta,\theta^{\prime})(\bar{\mu}-\mu_{t})(\mathop{}\\!\mathrm{d}\theta)(\bar{\mu}-\mu_{t})(\mathop{}\\!\mathrm{d}\theta^{\prime}).$

Now similarly to the proof of Proposition 5.2, denoting
$\theta=(\bm{a},\bm{w}),\theta^{\prime}=(\bm{a}^{\prime},\bm{w}^{\prime})$ we
can exploit the fact that
$\frac{\delta^{2}\\!F}{\delta\mu^{2}}(\mu,\theta,\theta^{\prime})$ is bilinear
in $\bm{a},\bm{a}^{\prime}$ to relate it to the kernel $\mathbf{H}_{\mu}$,

$\displaystyle\frac{\delta^{2}\\!F}{\delta\mu^{2}}(\mu,\theta,\theta^{\prime})$
$\displaystyle=\bm{a}^{\top}\left[\nabla_{\bm{a}}\nabla_{\bm{a}^{\prime}}\frac{\delta^{2}\\!F}{\delta\mu^{2}}(\mu,\theta,\theta^{\prime})\right]\bm{a}^{\prime}=(\bm{a}\;\;0_{d})^{\top}\left[\nabla_{\theta}\nabla_{\theta^{\prime}}\frac{\delta^{2}\\!F}{\delta\mu^{2}}(\mu,\theta,\theta^{\prime})\right](\bm{a}^{\prime}\;\;0_{d})$
$\displaystyle=(\bm{a}\;\;0_{d})^{\top}\mathbf{H}_{\mu}(\theta,\theta^{\prime})(\bm{a}^{\prime}\;\;0_{d}).$

Writing the eigenfunction decomposition of $\mathbf{H}_{\mu_{t}}$ as (omitting
the dependency on $t$ for brevity)

$\mathbf{H}_{\mu_{t}}(\theta,\theta^{\prime})=\sum_{j\in\operatorname{\mathbb{Z}}}\lambda_{j}\psi_{j}(\theta)\psi_{j}(\theta^{\prime})^{\top},\quad\int\lVert\psi_{j}\rVert^{2}\mathop{}\\!\mathrm{d}\mu_{t}=1\quad\forall
j\in\operatorname{\mathbb{Z}},\quad\lambda_{1}\geq\lambda_{2}\geq\cdots\geq
0,\quad\lambda_{0}\leq\lambda_{-1}\leq\cdots\leq 0,$

we may thus bound

$\displaystyle-\Lambda\geq\frac{\mathop{}\\!\mathrm{d}^{2}}{\mathop{}\\!\mathrm{d}s^{2}}\bigg{|}_{s=0}F(\bar{\mu}_{s})$
$\displaystyle=\iint(\bm{a}\;\;0_{d})^{\top}\mathbf{H}_{\mu}(\theta,\theta^{\prime})(\bm{a}^{\prime}\;\;0_{d})(\bar{\mu}-\mu_{t})(\mathop{}\\!\mathrm{d}\theta)(\bar{\mu}-\mu_{t})(\mathop{}\\!\mathrm{d}\theta^{\prime})$
$\displaystyle=\sum_{j\in\operatorname{\mathbb{Z}}}\lambda_{j}\left(\int(\bm{a}\;\;0_{d})^{\top}\psi_{j}(\theta)(\bar{\mu}-\mu_{t})(\mathop{}\\!\mathrm{d}\theta)\right)^{2}$
$\displaystyle\geq-\left\lvert\lambda_{0}\right\rvert\cdot\sum_{j\in\operatorname{\mathbb{Z}}}\left(\int(\bm{a}\;\;0_{d})^{\top}\psi_{j}(\theta)(\bar{\mu}-\mu_{t})(\mathop{}\\!\mathrm{d}\theta)\right)^{2}$
$\displaystyle=-\left\lvert\lambda_{0}\right\rvert\cdot\sum_{j\in\operatorname{\mathbb{Z}}}\left(\int\left(\frac{\mathop{}\\!\mathrm{d}\bar{\mu}}{\mathop{}\\!\mathrm{d}\mu_{t}}-1\right)(\bm{a}\;\;0_{d})^{\top}\psi_{j}(\theta)\mu_{t}(\mathop{}\\!\mathrm{d}\theta)\right)^{2}$
$\displaystyle=-\left\lvert\lambda_{0}\right\rvert\int\left(\frac{\mathop{}\\!\mathrm{d}\bar{\mu}}{\mathop{}\\!\mathrm{d}\mu_{t}}-1\right)^{2}\lVert\bm{a}\rVert^{2}\mu_{t}(\mathop{}\\!\mathrm{d}\theta)$
$\displaystyle\geq-\bar{\chi}^{2}|\lambda_{0}|,$

where we have made use of Parseval’s identity. Hence the largest negative
eigenvalue is bounded as $\lambda_{0}\leq-\bar{\chi}^{-2}\Lambda$. ∎

######  Theorem E.7.

Assume $F:\operatorname{\mathscr{P}}_{2}(\Omega)\to\operatorname{\mathbb{R}}$,
$\Omega\subseteq\operatorname{\mathbb{R}}^{m}$ satisfies
$\lVert\nabla\frac{\delta F}{\delta\mu}\rVert\leq C_{1}$, $\nabla\frac{\delta
F}{\delta\mu}$ is $C_{2}$-Lipschitz, $\mathbf{H}_{\mu}$ is Hilbert-Schmidt,
$\lVert\mathbf{H}_{\mu}\rVert\leq C_{3}$,
$\mathbf{H}_{\mu}(\theta,\theta^{\prime})$ is $C_{4}$-Lipschitz w.r.t.
$\theta,\theta^{\prime}$ and $C_{5}$-Lipschitz w.r.t. $\mu$ in
$\operatorname{\mathcal{W}}_{1}$. Further suppose that
$\lambda_{0}:=\lambda_{\textup{min}}(\operatorname{\mathcal{H}}_{\mu^{\dagger}})<0$
and the corresponding eigenfunction $\psi_{0}$ satisfies
$|\int\psi_{0}^{\top}\nabla\frac{\delta\\!\operatorname{\mathcal{L}}}{\delta\mu}(\mu_{t})\mathop{}\\!\mathrm{d}\mu_{t}|\geq\alpha$
for some $\alpha>0$. Then WGF initialized at $\mu_{0}=\mu^{\dagger}$ decreases
$F$ by at least $F(\mu_{\tau})\leq
F(\mu_{0})-\Omega\left(\frac{|\lambda_{0}|\alpha}{\sqrt{m}\tau}\right)$ in
time
$\tau=O\left(\frac{1}{|\lambda_{0}|}\log\frac{|\lambda_{0}|}{\sqrt{m}\alpha}\right)$.

Unlike before, $F$ can be completely general and does not need to depend on
$\mu$ through an MLP layer.

###### Proof.

First note that the function
$\theta^{\prime}\mapsto\mathbf{H}_{\mu^{\dagger}}(\theta,\theta^{\prime})\nabla\frac{\delta
F}{\delta\mu}(\mu_{t},\theta^{\prime})$ is uniformly Lipschitz: for any
$\theta_{1}^{\prime},\theta_{2}^{\prime}$,

$\displaystyle\left\lVert\mathbf{H}_{\mu^{\dagger}}(\theta,\theta_{1}^{\prime})\nabla\frac{\delta
F}{\delta\mu}(\mu_{t},\theta_{1}^{\prime})-\mathbf{H}_{\mu^{\dagger}}(\theta,\theta_{2}^{\prime})\nabla\frac{\delta
F}{\delta\mu}(\mu_{t},\theta_{2}^{\prime})\right\rVert$
$\displaystyle\leq\lVert\mathbf{H}_{\mu^{\dagger}}(\theta,\theta_{1}^{\prime})-\mathbf{H}_{\mu^{\dagger}}(\theta,\theta_{2}^{\prime})\rVert\cdot\left\lVert\nabla\frac{\delta
F}{\delta\mu}(\mu_{t},\theta_{1}^{\prime})\right\rVert+\lVert\mathbf{H}_{\mu^{\dagger}}(\theta,\theta_{2}^{\prime})\rVert\cdot\left\lVert\nabla\frac{\delta
F}{\delta\mu}(\mu_{t},\theta_{1}^{\prime})-\nabla\frac{\delta
F}{\delta\mu}(\mu_{t},\theta_{2}^{\prime})\right\rVert$ $\displaystyle\leq
C_{1}C_{4}\lVert\theta_{1}^{\prime}-\theta_{2}^{\prime}\rVert+C_{2}C_{3}\lVert\theta_{1}^{\prime}-\theta_{2}^{\prime}\rVert.$

We re-expand the evolution equation (7) for the dynamics $(\mu_{t})_{t\geq 0}$
around $\mu^{\dagger}$ as

$\displaystyle\partial_{t}\left[\nabla_{\theta}\frac{\delta
F}{\delta\mu}(\mu_{t},\theta)\right]$
$\displaystyle=-\int\mathbf{H}_{\mu_{t}}(\theta,\theta^{\prime})\nabla_{\theta^{\prime}}\frac{\delta
F}{\delta\mu}(\mu_{t},\theta^{\prime})\mu_{t}(\mathop{}\\!\mathrm{d}\theta^{\prime})$
$\displaystyle=:-\int\mathbf{H}_{\mu^{\dagger}}(\theta,\theta^{\prime})\nabla_{\theta^{\prime}}\frac{\delta
F}{\delta\mu}(\mu_{t},\theta^{\prime})\mu^{\dagger}(\mathop{}\\!\mathrm{d}\theta^{\prime})+e(t,\theta),$

where the difference or error function $e(t,\theta)$ can be bounded as

$\displaystyle\lVert e(t,\theta)\rVert$
$\displaystyle\leq\left\lVert\int(\mathbf{H}_{\mu_{t}}(\theta,\theta^{\prime})-\mathbf{H}_{\mu^{\dagger}}(\theta,\theta^{\prime}))\nabla_{\theta^{\prime}}\frac{\delta
F}{\delta\mu}(\mu_{t},\theta^{\prime})\mu_{t}(\mathop{}\\!\mathrm{d}\theta^{\prime})\right\rVert+\left\lVert\int\mathbf{H}_{\mu^{\dagger}}(\theta,\theta^{\prime})\nabla_{\theta^{\prime}}\frac{\delta
F}{\delta\mu}(\mu_{t},\theta^{\prime})(\mu_{t}-\mu^{\dagger})(\mathop{}\\!\mathrm{d}\theta^{\prime})\right\rVert$
$\displaystyle\leq\left(C_{1}C_{5}+(C_{1}C_{4}+C_{2}C_{3})m^{1/2}\right)\operatorname{\mathcal{W}}_{1}(\mu_{t},\mu^{\dagger}).$
$\displaystyle=:C_{6}\operatorname{\mathcal{W}}_{1}(\mu_{t},\mu^{\dagger}).$

For the second term, we have used the Lipschitz constant derived above to
bound each entry separately. Then the $\psi_{0}$-component
$\alpha_{0}(t):=\int\psi_{0}^{\top}\nabla\frac{\delta
F}{\delta\mu}(\mu_{t})\mathop{}\\!\mathrm{d}\mu^{\dagger}$ of the gradient
evolves according to

$\displaystyle\frac{\mathop{}\\!\mathrm{d}}{\mathop{}\\!\mathrm{d}t}\alpha_{0}(t)$
$\displaystyle=-\iint\psi_{0}(\theta)^{\top}\mathbf{H}_{\mu^{\dagger}}(\theta,\theta^{\prime})\nabla_{\theta^{\prime}}\frac{\delta
F}{\delta\mu}(\mu_{t},\theta^{\prime})\mu^{\dagger}(\mathop{}\\!\mathrm{d}\theta^{\prime})\mu^{\dagger}(\mathop{}\\!\mathrm{d}\theta)+\int\psi_{0}(\theta)^{\top}e(t,\theta)\mu^{\dagger}(\mathop{}\\!\mathrm{d}\theta)$
$\displaystyle=-\lambda_{0}\int\psi_{0}(\theta)^{\top}\nabla_{\theta}\frac{\delta
F}{\delta\mu}(\mu_{t},\theta)\mu^{\dagger}(\mathop{}\\!\mathrm{d}\theta)+\int\psi_{0}(\theta)^{\top}e(t,\theta)\mu^{\dagger}(\mathop{}\\!\mathrm{d}\theta),$

and hence

$\left\lvert\frac{\mathop{}\\!\mathrm{d}}{\mathop{}\\!\mathrm{d}t}\alpha_{0}(t)+\lambda_{0}\alpha_{0}(t)\right\rvert\leq\left(\int\lVert\psi_{0}\rVert^{2}\mathop{}\\!\mathrm{d}\mu^{\dagger}\right)^{1/2}\sup_{\theta\in\Theta}\,\lVert
e(t,\theta)\rVert\leq
C_{8}\operatorname{\mathcal{W}}_{1}(\mu_{t},\mu^{\dagger}).$

Without loss of generality, assume initially $\alpha_{0}(0)$ is positive so
that $\alpha_{0}(0)\geq\alpha$. We consider a 1-Wasserstein ball centered at
$\mu^{\dagger}$ with radius small enough so that the error term is negligible
compared to the exponential growth,

$\operatorname{\mathcal{B}}_{\operatorname{\mathcal{W}}}(\Delta)=\left\\{\mu\in\operatorname{\mathscr{P}}_{2}(\Theta):\operatorname{\mathcal{W}}_{1}(\mu,\mu^{\dagger})\leq\Delta:=\frac{|\lambda_{0}|\alpha}{2C_{6}}\right\\}.$

Then for a set time interval $\tau>0$ to be determined, either of the
following must happen:

  1. cmrmn(i)

$(\mu_{t})_{t\in[0,\tau]}\subset\operatorname{\mathcal{B}}_{\operatorname{\mathcal{W}}}(\Delta)$.
In this case, $\alpha_{0}(t)$ grows exponentially during the entire interval
$t\in[0,\tau]$ as

$\frac{\mathop{}\\!\mathrm{d}}{\mathop{}\\!\mathrm{d}t}\alpha_{0}(t)\geq|\lambda_{0}|\alpha_{0}(t)-C_{6}\Delta=|\lambda_{0}|\left(\alpha_{0}(t)-\frac{\alpha}{2}\right)>0,$

showing that

$\alpha_{0}(t)\geq
e^{|\lambda_{0}|t}\left(\alpha_{0}(0)-\frac{\alpha}{2}\right)+\frac{\alpha}{2}\geq\frac{\alpha(e^{|\lambda_{0}|t}+1)}{2}.$

Then the decrease of $F$ after time $\tau$ can be bounded below by retrieving
the $\psi_{0}$-component as

$\displaystyle F(\mu_{0})-F(\mu_{\tau})$
$\displaystyle=\int_{0}^{\tau}\int\left\lVert\nabla\frac{\delta
F}{\delta\mu}(\mu_{t},\theta)\right\rVert^{2}\mu_{t}(\mathop{}\\!\mathrm{d}\theta)\mathop{}\\!\mathrm{d}t$
$\displaystyle\geq\int_{0}^{\tau}\bigg{(}\int\left\lVert\nabla\frac{\delta
F}{\delta\mu}(\mu_{t},\theta)\right\rVert^{2}\mu^{\dagger}(\mathop{}\\!\mathrm{d}\theta)\mathop{}\\!\mathrm{d}t-2C_{1}C_{2}\operatorname{\mathcal{W}}_{1}(\mu_{t},\mu^{\dagger})\bigg{)}$
$\displaystyle\geq\int_{0}^{\tau}\left(\int\psi_{0}(\theta)^{\top}\nabla\frac{\delta
F}{\delta\mu}(\mu_{t},\theta)\mu^{\dagger}(\mathop{}\\!\mathrm{d}\theta)\right)^{2}\mathop{}\\!\mathrm{d}t-2C_{1}C_{2}\Delta\tau$
$\displaystyle=\int_{0}^{\tau}\alpha_{0}(t)^{2}\mathop{}\\!\mathrm{d}t-2C_{1}C_{2}\Delta\tau$
$\displaystyle\geq\frac{\alpha^{2}}{4}\left(\frac{1}{2|\lambda_{0}|}(e^{2|\lambda_{0}|\tau}-1)+\frac{2}{|\lambda_{0}|}(e^{|\lambda_{0}|\tau}-1)+\tau\right)-\frac{C_{1}C_{2}}{C_{6}}|\lambda_{0}|\alpha\tau.$

  2. cmrmn(ii)

$\mu_{\tau_{e}}\notin\operatorname{\mathcal{B}}_{\operatorname{\mathcal{W}}}(\Delta)$
for some $\tau_{e}\leq\tau$. If the mean-field flow has managed to escape the
ball $\operatorname{\mathcal{B}}_{\operatorname{\mathcal{W}}}(\Delta)$ in time
$\tau_{e}$, the Benamou-Brenier formula (Proposition A.3) immediately
guarantees that

$F(\mu_{0})-F(\mu_{\tau})\geq
F(\mu_{0})-F(\mu_{\tau_{e}})=\int_{0}^{\tau_{e}}\int\left\lVert\nabla\frac{\delta
F}{\delta\mu}(\mu_{t},\theta)\right\rVert^{2}\mu_{t}(\mathop{}\\!\mathrm{d}\theta)\mathop{}\\!\mathrm{d}t\geq\frac{\operatorname{\mathcal{W}}_{2}(\mu_{\tau_{e}},\mu^{\dagger})}{\tau_{e}}>\frac{\Delta}{\tau}.$

Thus we have proved that:

$F(\mu_{0})-F(\mu_{\tau})\geq\left(\frac{\alpha^{2}}{4}\left(\frac{1}{2|\lambda_{0}|}(e^{2|\lambda_{0}|\tau}-1)+\frac{2}{|\lambda_{0}|}(e^{|\lambda_{0}|\tau}-1)+\tau\right)-\frac{C_{1}C_{2}}{C_{6}}|\lambda_{0}|\alpha\tau\right)\wedge\frac{|\lambda_{0}|\alpha}{2C_{6}\tau}.$
(13)

Due to the exponential terms, we see $\tau\approx\log\frac{1}{\alpha}$ is
enough to ensure that the two terms become roughly equal so that the
guarantees is close to optimal. For the remainder of the proof, we derive the
exact formula. Choose

$\tau=\frac{1}{|\lambda_{0}|}\log\frac{C_{7}}{\alpha}$

for some $C_{7}>\alpha$. The first term in the right-hand side of (13) can be
bounded as

$\displaystyle\frac{\alpha^{2}}{4}\left(\frac{1}{2|\lambda_{0}|}(e^{2|\lambda_{0}|\tau}-1)+\frac{2}{|\lambda_{0}|}(e^{|\lambda_{0}|\tau}-1)+\tau\right)-\frac{C_{1}C_{2}}{C_{6}}|\lambda_{0}|\alpha\tau$
$\displaystyle=\frac{\alpha^{2}}{8|\lambda_{0}|}\left(\frac{C_{7}^{2}}{\alpha^{2}}-1\right)+\frac{\alpha^{2}}{2|\lambda_{0}|}\left(\frac{C_{7}}{\alpha}-1\right)+\frac{\alpha^{2}}{4|\lambda_{0}|}\log\frac{C_{7}}{\alpha}-\frac{C_{1}C_{2}}{C_{6}}\alpha\log\frac{C_{7}}{\alpha}$
$\displaystyle\geq\left(\frac{C_{7}^{2}}{24|\lambda_{0}|}-\frac{5\alpha^{2}}{8|\lambda_{0}|}\right)+\left(\frac{C_{7}^{2}}{24|\lambda_{0}|}-\frac{C_{1}C_{2}}{C_{6}}\alpha\log\frac{C_{7}}{\alpha}\right)+\frac{C_{7}^{2}}{24|\lambda_{0}|}+\frac{C_{7}\alpha}{2|\lambda_{0}|}+\frac{\alpha^{2}}{4|\lambda_{0}|}\log\frac{C_{7}}{\alpha}$
$\displaystyle\geq\left(\frac{C_{7}^{2}}{24|\lambda_{0}|}-\frac{5\alpha^{2}}{8|\lambda_{0}|}\right)+\left(\frac{C_{7}^{2}}{24|\lambda_{0}|}-\frac{C_{1}C_{2}C_{7}}{C_{6}e}\right)+\frac{C_{7}^{2}}{24|\lambda_{0}|},$

where we have used the fact that the function $x\mapsto x\log\frac{c}{x}$ has
maximum $\frac{c}{e}$. Then the first term of (13) will dominate the second as
long as

$\frac{C_{7}^{2}}{24|\lambda_{0}|}\geq\frac{5\alpha^{2}}{8|\lambda_{0}|}\vee\frac{C_{1}C_{2}C_{7}}{C_{6}e}\vee\frac{|\lambda_{0}|^{2}\alpha}{2C_{6}\log\frac{C_{7}}{\alpha}}.$

Manipulating terms shows that

$C_{7}=\sqrt{15}\alpha\vee\frac{24C_{1}C_{2}|\lambda_{0}|}{C_{6}e}\vee\left(\frac{24|\lambda_{0}|^{3}\alpha}{C_{6}\log
15}\right)^{1/2}$

is sufficient. For the purposes of the general statement, we focus on
asymptotic behavior w.r.t. $\alpha$ and hide all regularity constants
$C_{1},\cdots,C_{5}$, yielding $C_{6}=O(m^{1/2})$ and
$C_{7}=O(|\lambda_{0}|m^{-1/2})$. ∎

#### Proof of Theorem 5.6.

Let us fix the lower bound
$\lambda_{\textup{min}}(\bm{\Sigma}_{\mu_{r},\mu_{r}})\geq\lambda=\Theta(\frac{1}{k})$.
(The bound only needs to hold either locally for the
$\operatorname{\mathcal{W}}_{2}$-ball of radius $\Delta$ in the proof of
Theorem E.7, or along the dynamics $\mu_{t}$ until escape.) We first need a
robust version of Theorem 3.3ii since $\mu_{t}$ cannot be exactly on a
critical point. If
$\frac{\mathop{}\\!\mathrm{d}}{\mathop{}\\!\mathrm{d}s}\big{|}_{s=0}\operatorname{\mathcal{L}}(\bar{\mu}_{s})>-\delta$
it must hold that
$\lVert\mathbf{L}_{\mu}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\rVert_{*}<\frac{\delta}{2}$
by (8). Then from (9), again choosing $\mathbf{R}\in\mathcal{O}(k)$ such that
$\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\mathbf{R}$ is
symmetric,

$\displaystyle\frac{\mathop{}\\!\mathrm{d}^{2}}{\mathop{}\\!\mathrm{d}s^{2}}\bigg{|}_{s=0}\operatorname{\mathcal{L}}(\bar{\mu}_{s})$
$\displaystyle=-4\operatorname{tr}\left(\mathbf{L}_{\mu}^{2}\mathbf{R}^{\top}\bm{\Sigma}_{\mu,\mu}^{-1}\mathbf{R}\right)+2\operatorname{tr}\left(\mathbf{L}_{\mu}(2\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\mathbf{R}+\mathbf{R}^{\top}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\mu,\mu^{\circ}}-2\mathbf{I}_{k})\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\mathbf{R}\right)$
$\displaystyle=-4\operatorname{tr}\left(\mathbf{L}_{\mu}^{2}\mathbf{R}^{\top}\bm{\Sigma}_{\mu,\mu}^{-1}\mathbf{R}\right)+2\operatorname{tr}\left((2\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\mathbf{R}+\mathbf{R}^{\top}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\mu,\mu^{\circ}}-2\mathbf{I}_{k})^{\top}\mathbf{L}_{\mu}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\mathbf{R}\right)$
$\displaystyle\leq-\frac{4}{kR_{1}^{2}}\operatorname{\mathcal{L}}(\mu_{t})^{2}+2\,\lVert\mathbf{L}_{\mu}\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\rVert_{*}\lVert
2\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1}\mathbf{R}+\mathbf{R}^{\top}\bm{\Sigma}_{\mu,\mu}^{-1}\bm{\Sigma}_{\mu,\mu^{\circ}}-2\mathbf{I}_{k}\rVert$
$\displaystyle\leq-\frac{4}{kR_{1}^{2}}\operatorname{\mathcal{L}}(\mu_{t})^{2}+\left(3\,\lVert\bm{\Sigma}_{\mu^{\circ},\mu}\bm{\Sigma}_{\mu,\mu}^{-1/2}\rVert\cdot\lVert\bm{\Sigma}_{\mu,\mu}^{-1/2}\mathbf{R}\rVert+2\right)\delta$
$\displaystyle\leq-\frac{4}{kR_{1}^{2}}\operatorname{\mathcal{L}}(\mu_{t})^{2}+(3r^{\circ
1/2}\lambda^{-1/2}+2)\delta.$

Hence if $\delta\leq\frac{2}{kR_{1}^{2}(3r^{\circ
1/2}\lambda^{-1/2}+2)}\operatorname{\mathcal{L}}(\mu_{t})^{2}$ then
$\frac{\mathop{}\\!\mathrm{d}^{2}}{\mathop{}\\!\mathrm{d}s^{2}}\big{|}_{s=0}\operatorname{\mathcal{L}}(\bar{\mu}_{s})\leq-\frac{2}{kR_{1}^{2}}\operatorname{\mathcal{L}}(\mu_{t})^{2}$,
and by Proposition E.6 it holds that

$\displaystyle\lambda_{0}=\lambda_{\textup{min}}(\operatorname{\mathcal{H}}_{\mu_{t}})\leq-\frac{2\gamma}{kR_{1}^{2}R_{4}}\operatorname{\mathcal{L}}(\mu_{t})^{2}.$

Then Theorem E.7 applies to $F=\operatorname{\mathcal{L}}$ by virtue of Lemma
E.3 and the regularity constants derived in Lemma E.4 and E.5. One can check
that

$C_{6}=C_{1}C_{5}+(C_{1}C_{4}+C_{2}C_{3})(k+d)^{1/2}=O\left(\frac{d}{k\lambda^{4}}\right)$

and

$C_{7}=O\left(\alpha\vee\frac{\lambda^{5/2}\gamma}{k^{3/2}d}\operatorname{\mathcal{L}}(\mu_{t})^{2}\vee\frac{\lambda^{2}\gamma^{3/2}\alpha^{1/2}}{kd^{1/2}}\operatorname{\mathcal{L}}(\mu_{t})^{3}\right)=O\left(\alpha+\frac{\gamma}{k^{4}d}\right);$

the third term is dominated by the geometric mean of the first two, and
$\operatorname{\mathcal{L}}(\mu_{t})=O(1)$. Hence the time interval of
interest is

$\tau=O\left(\frac{k}{\gamma\operatorname{\mathcal{L}}(\mu_{t})^{2}}\left(\log\frac{\gamma}{k^{4}d\alpha}\right)\vee
1\right),$

and the guaranteed decrease of the objective is

$\operatorname{\mathcal{L}}(\mu_{t})-\operatorname{\mathcal{L}}(\mu_{t+\tau})\geq\frac{|\lambda_{0}|\alpha}{2C_{6}\tau}\geq\Omega\left(\frac{\gamma^{2}\alpha\operatorname{\mathcal{L}}(\mu_{t})^{4}}{k^{5}d}\left(\log\frac{\gamma}{k^{4}d\alpha}\vee
1\right)^{-1}\right).$

###  E.3 Escaping from Saddle Points

The usual theory of Gaussian processes can be readily extended to
multivariable outputs.

######  Definition E.8 (vector-valued Gaussian process).

The random function $\bm{\xi}:\Omega\to\operatorname{\mathbb{R}}^{m}$ is said
to follow a Gaussian process if any finite collection of variables
$\bm{\xi}(\theta_{1}),\cdots,\bm{\xi}(\theta_{j})$ are jointly normally
distributed. The process is determined by the mean function
$\bm{m}:\Omega\to\operatorname{\mathbb{R}}^{m}$,
$\bm{m}(\theta)=\mathbb{E}[\bm{\xi}(\theta)]$ and matrix-valued covariance
function

$\mathbf{K}:\Omega\times\Omega\to\operatorname{\mathbb{R}}^{m\times
m},\quad\mathbf{K}(\theta,\theta^{\prime})=\mathbb{E}[(\bm{\xi}(\theta)-\bm{m}(\theta))(\bm{\xi}(\theta^{\prime})-\bm{m}(\theta^{\prime}))^{\top}].$

We denote this process as $\bm{\xi}\sim\operatorname{GP}(\bm{m},\mathbf{K})$.
See Álvarez et al. (2012) for further details.

######  Lemma E.9.

For any $\mu\in\operatorname{\mathscr{P}}_{2}(\Omega)$, square-integrable test
function $\psi\in L^{2}(\Omega,\mu;\operatorname{\mathbb{R}}^{m})$ and
covariance function
$\mathbf{K}:\Omega\times\Omega\to\operatorname{\mathbb{R}}^{m\times m}$
satisfying
$\int\lVert\mathbf{K}(\theta,\theta)\rVert\mu(\mathop{}\\!\mathrm{d}\theta)<\infty$
the inner product
$\langle\psi,\bm{\xi}\rangle_{L^{2}(\Omega,\mu;\operatorname{\mathbb{R}}^{m})}$
for $\bm{\xi}\sim\operatorname{GP}(0,\mathbf{K})$ is normally distributed.

###### Proof.

Note that the inner product is defined almost surely since

$\mathbb{E}_{\bm{\xi}}\left[\lVert\bm{\xi}\rVert_{L^{2}(\Omega,\mu;\operatorname{\mathbb{R}}^{m})}^{2}\right]=\int\mathbb{E}_{\bm{\xi}}[\lVert\bm{\xi}(\theta)\rVert^{2}]\mu(\mathop{}\\!\mathrm{d}\theta)=\int\operatorname{tr}\mathbf{K}(\theta,\theta)\mu(\mathop{}\\!\mathrm{d}\theta)<\infty.$

We denote by $\mathscr{E}$ the closed linear span of the set of square-
integrable random variables
$\\{\psi(\theta)^{\top}\bm{\xi}(\theta):\theta\in\Omega\\}$. For any
$\mathbf{Z}\in\mathscr{E}^{\perp}$ it holds that
$\mathbb{E}_{\bm{\xi}}[\mathbf{Z}\psi(\theta)^{\top}\bm{\xi}(\theta)]=0$, so
that by Fubini’s theorem

$\mathbb{E}_{\bm{\xi}}\left[\mathbf{Z}\langle\psi,\bm{\xi}\rangle_{L^{2}(\Omega,\mu;\operatorname{\mathbb{R}}^{m})}\right]=\mathbb{E}_{\bm{\xi}}\left[\int\mathbf{Z}\psi(\theta)^{\top}\bm{\xi}(\theta)\mu(\mathop{}\\!\mathrm{d}\theta)\right]=0.$

Hence
$\langle\psi,\bm{\xi}\rangle_{L^{2}(\Omega,\mu;\operatorname{\mathbb{R}}^{m})}\in(\mathscr{E}^{\perp})^{\perp}=\mathscr{E}$,
and so is normally distributed. ∎

For the proposed perturbation process, the change in the gradient field along
the flow of $\bm{\xi}$ can be quantified as

$\displaystyle\nabla\frac{\delta F}{\delta\mu}(\mu_{\Delta
t},\theta)-\nabla\frac{\delta F}{\delta\mu}(\mu^{\dagger},\theta)$
$\displaystyle=\int_{0}^{\Delta t}\partial_{t}\left[\nabla\frac{\delta
F}{\delta\mu}(\mu_{t},\theta)\right]\mathop{}\\!\mathrm{d}t$
$\displaystyle=-\int_{0}^{\Delta
t}\int\nabla_{\theta}\nabla_{\theta^{\prime}}\frac{\delta^{2}\\!F}{\delta\mu^{2}}(\mu_{t},\theta,\theta^{\prime})\bm{\xi}(\theta^{\prime})\mu_{t}(\mathop{}\\!\mathrm{d}\theta^{\prime})\mathop{}\\!\mathrm{d}t$
$\displaystyle=-\int_{0}^{\Delta
t}\operatorname{\mathcal{H}}_{\mu_{t}}\left[\bm{\xi}\right]\mathop{}\\!\mathrm{d}t.$

The resulting $\psi_{0}$-component is

$\displaystyle\alpha(\bm{\xi})$
$\displaystyle=\int\psi_{0}(\theta)^{\top}\nabla\frac{\delta
F}{\delta\mu}(\mu_{\Delta
t},\theta)\mu^{\dagger}(\mathop{}\\!\mathrm{d}\theta)$
$\displaystyle=\int\psi_{0}(\theta)^{\top}\bigg{(}\nabla\frac{\delta
F}{\delta\mu}(\mu^{\dagger},\theta)-\int_{0}^{\Delta
t}\operatorname{\mathcal{H}}_{\mu^{\dagger}}\left[\bm{\xi}\right]\mathop{}\\!\mathrm{d}t+\int_{0}^{\Delta
t}(\operatorname{\mathcal{H}}_{\mu^{\dagger}}-\operatorname{\mathcal{H}}_{\mu_{t}})\left[\bm{\xi}\right]\mathop{}\\!\mathrm{d}t\bigg{)}\mu^{\dagger}(\mathop{}\\!\mathrm{d}\theta)$
$\displaystyle=-\lambda_{0}\Delta
t\int\psi_{0}(\theta)^{\top}\bm{\xi}(\theta)\mu^{\dagger}(\mathop{}\\!\mathrm{d}\theta)+\alpha+O(\Delta
t^{2}),$

and first term is normally distributed by Lemma E.9.