# Momentum Particle Maximum Likelihood

Jen Ning Lim University of Warwick Juan Kuntz Polygeist  Samuel Power
University of Bristol  Adam M. Johansen University of Warwick

###### Abstract

Maximum likelihood estimation (MLE) of latent variable models is often recast
as an optimization problem over the extended space of parameters and
probability distributions. For example, the Expectation Maximization (EM)
algorithm can be interpreted as coordinate descent applied to a suitable free
energy functional over this space. Recently, this perspective has been
combined with insights from optimal transport and Wasserstein gradient flows
to develop particle-based algorithms applicable to wider classes of models
than standard EM.

Drawing inspiration from prior works which interpret ‘momentum-enriched’
optimisation algorithms as discretizations of ordinary differential equations,
we propose an analogous dynamical systems-inspired approach to minimizing the
free energy functional over the extended space of parameters and probability
distributions. The result is a dynamic system that blends elements of
Nesterov’s Accelerated Gradient method, the underdamped Langevin diffusion,
and particle methods.

Under suitable assumptions, we establish quantitative convergence of the
proposed system to the unique minimiser of the functional in continuous time.
We then propose a numerical discretization of this system which enables its
application to parameter estimation in latent variable models. Through
numerical experiments, we demonstrate that the resulting algorithm converges
faster than existing methods and compares favourably with other (approximate)
MLE algorithms.

##  1 Introduction

In this work, we are interested in parameter estimation for (probabilistic)
latent variable models of the form $p_{\theta}\left(y,x\right)$, where
$\theta\in\mathbb{R}^{d_{\theta}}$ denotes the model parameters,
$x\in\mathbb{R}^{d_{x}}$ is an unobserved (or latent) variable and
$y\in\mathbb{R}^{d_{y}}$ is an observed variable. Throughout this work, we
treat $y$ as fixed. The type II maximum likelihood approach (Good, , 1983)
aims to estimate $\theta$ by maximization of the _marginal_ likelihood
$p_{\theta}\left(y\right):=\int p_{\theta}\left(y,x\right)\,\mathrm{d}x$.
However, in various models of practical interest, this integral is not
available in closed form, so direct optimisation of this objective is
infeasible.

One approach to circumventing this intractability is to construct a tractable
objective which is defined over an extended space, and whose optimal values
coincide with that of the original maximum likelihood estimation (MLE)
problem. For latent variable models, this can be carried out as follows
(Dempster et al., , 1977): writing
$\mathcal{P}\left(\mathbb{R}^{d_{x}}\right)$ for the space of probability
distributions over the latent space, define a ‘free energy’ functional
$\mathcal{E}:\mathbb{R}^{d_{\theta}}\times\mathcal{P}\left(\mathbb{R}^{d_{x}}\right)\to\mathbb{R}$
by

$\displaystyle\mathcal{E}\left(\theta,q\right):=\int
q\left(x\right)\log\left(\frac{q\left(x\right)}{p_{\theta}\left(y,x\right)}\right)\,\mathrm{d}x.$

One can then compute directly that $\mathcal{E}\left(\theta,q\right)=-\log
p_{\theta}\left(y\right)+\mathsf{KL}\left(q,p_{\theta}\left(\cdot\mid
y\right)\right)$, where $\mathsf{KL}$ is the Kullback-Leibler divergence.
Since it holds that

$\min\left\\{\mathcal{E}\left(\theta,q\right):q\in\mathcal{P}\left(\mathbb{R}^{d_{x}}\right)\right\\}=-\log
p_{\theta}\left(y\right),$

minimizing the free energy functional over the extended space of parameters
and probability distributions (over the latent space), one will minimize the
function $\theta\mapsto-\log p_{\theta}\left(y\right)$, hence solving the
original maximum likelihood estimation problem.

This representation suggests exploring practical procedures for minimising the
free energy functional $\mathcal{E}$. Neal and Hinton,  (1998) describe that
the classical EM algorithm is precisely coordinate descent applied to
$\mathcal{E}$. In particular, the EM approach is readily applicable to models
for which the conditional distributions $p_{\theta}\left(x\mid y\right)$ are
available in closed form, precisely because this enables the minimization of
$q$ with $\theta$ held fixed to be carried out tractably. Recent work of Kuntz
et al.,  (2023) has sought to move beyond this class of models by modelling
the distribution $q$ nonparametrically. Drawing inspiration from the mature
literature on gradient flows (e.g., see Ambrosio et al.,  (2005)), this work
constructs several gradient flows of the functional $\mathcal{E}$ over the
joint space of extended space of parameters and probability distributions over
the latent space. The resulting flow, called particle Gradient Descent (PGD),
simultaneously evolves a parameter estimate, written $\theta_{t}$, and an
estimate of the posterior represented by a finite particle system, written
$q_{t}$. These gradient flows are then discretized in time, yielding an
algorithm qualitatively similar to gradient descent of $\mathcal{E}$.

In convex optimization, it is well understood that while the gradient descent
method has many favourable theoretical guarantees, it is also suboptimal. In
particular, under standard computational models (Nemirovskij and Yudin, ,
1983), there are practical algorithms which yield substantially better worst-
case convergence guarantees and improved practical performance (e.g., see
Nesterov,  (1983)). Among optimization algorithms which achieve these
‘accelerated’ convergence rates, a common feature is the presence of
‘momentum’ effects in the dynamics of the algorithm. Roughly speaking,
gradient descent of a function $f$ is analogous to solving the ordinary
differential equation (ODE) $\dot{\theta}_{t}=-\nabla
f\left(\theta_{t}\right),$ then a ‘momentum-type’ method might instead
resemble solving a second-order ODE like
$\ddot{\theta}_{t}+\gamma\dot{\theta}_{t}+\nabla f\left(\theta_{t}\right)=0$
for some ‘damping’ parameter $\gamma>0$ (Su et al., , 2014; Wibisono et al., ,
2016). In the context of machine learning, Sutskever et al.,  (2013)
demonstrated empirically that incorporating momentum effects into stochastic
gradient descent (SGD) could have substantial benefits, such as eliminating
the performance gap between standard SGD and competing ‘Hessian-free’
deterministic methods based on higher-order differential information (Martens,
, 2010).

Inspired by the success of momentum methods, the present work seeks to explore
whether PGD can similarly be accelerated by utilizing ‘momentum’ strategies.
Our contributions are as follows: (1) we construct a continuous-time flow that
incorporates momentum effects into PGD; (2) under suitable conditions, we
establish convergence in continuous time; (3) we derive a discretization of
the flow that can achieve better performance than PGD and compare
competitively against other methods.

The structure of this manuscript is as follows: in Section 2, we review
existing work on gradient flows on Euclidean space $\mathbb{R}^{d_{\theta}}$,
probability space $\mathcal{P}(\mathbb{R}^{d_{x}})$, and extended space
$\mathbb{R}^{d_{\theta}}\times\mathcal{P}(\mathbb{R}^{d_{x}})$. We also review
momentum methods on Euclidean space $\mathbb{R}^{d_{\theta}}$, and probability
space $\mathcal{P}(\mathbb{R}^{d_{x}})$. Then, in Section 3, we describe how
momentum can be incorporated into the PGD; the result is a momentum-enrich
dynamical system called Momentum PGD (MPGD). After, in Section 4, we establish
convergence of the continuous-time MPGD. In the Section 5, we describe how to
discretize the continuous MPGD algorithm to produce our proposed algorithm.
Finally, in Section 6, we study on the toy examples of how the choice of
momentum parameters affects the optimization process, as well as the effects
of various choices in the discretization process. As a large-scale experiment,
we compare our proposed method for training Variational Autoencoders (Kingma
and Welling, , 2014) against current methods of training latent variable
models.

##  2 Background

The main goal of the paper is to propose algorithmic improvements to the PGD
algorithm of Kuntz et al.,  (2023), which is derived as a gradient flow on an
extended space of parameters and probability distributions, by incorporating
some notion of momentum into the dynamics. Towards this goal, we first provide
an overview of standard gradient flows on these spaces while detailing
precisely how momentum ideas have been successfully incorporated in related
scenarios.

###  2.1 Gradient Flows on $\mathbb{R}^{d}$, and their Acceleration

Given a function $f:\mathbb{R}^{d}\to\mathbb{R}$ of sufficient regularity, the
(Euclidean) gradient flow of $f$ is given by the ODE:
$\dot{\theta}_{t}=-\nabla_{\theta}f\left(\theta_{t}\right)$ where
$\nabla_{\theta}f$ is the standard (Euclidean) gradient of $f$, i.e., the
vector such that $\theta\approx\tilde{\theta}\implies
f\left(\tilde{\theta}\right)=f\left(\theta\right)+\langle\nabla_{\theta}f\left(\theta\right),\tilde{\theta}-\theta\rangle+o\left(\|\tilde{\theta}-\theta\|_{2}\right).$
Discretizing the gradient flow in time with a standard forward Euler
integrator yields the celebrated gradient descent method. Thus, it is natural
to ask if “accelerated” gradient descent, such as Nesterov,  (1983)’s
Accelerated Gradient (NAG) and Polyak,  (1964)’s Heavy ball, can be viewed
from the lens of discretizations of a dynamical system. This has proved a
fruitful research question (Su et al., , 2014; Shi et al., , 2021). A
recurring pattern observed is that many successful accelerated algorithms can
be viewed as discretizations of a certain ‘momentum-enriched’ system, i.e.,
corresponding to the evolution of a kinetic particle that carries both a
position and a momentum. For example, Nesterov’s Accelerated Gradient (NAG)
method (Nesterov, , 1983) can be seen as a discretization of Wilson et al.,
(2016, Equation 7)’s system of ODEs

$\displaystyle\dot{\theta}_{t}$ $\displaystyle=\eta
m_{t},\quad\dot{m}_{t}=-\gamma\eta
m_{t}-\nabla_{\theta}f\left(\theta_{t}\right),$ (1)

where $\gamma,\eta$ are positive hyperparameters (see Section I.2). For
convex, quadratic $f$, this system reproduces the dynamics of a damped
harmonic oscillator (McCall, , 2010), and this equips the hyperparameters of
the ODE with intuitive interpretations: $\eta^{-1}$ represents the mass of the
particle, and $\gamma$ determines the strength of the damping. Additionally,
in the massless limit (i.e. sending $\eta\to\infty$ and rescaling time
appropriately), one can formally recover that the solution follows the
gradient flow $\dot{\theta}_{t}=-\nabla_{\theta}f\left(\theta_{t}\right)$.

Towards generalising this construction, it is useful to introduce the
‘Hamiltonian’ function
$F_{\eta}\left(\theta,m\right):=f\left(\theta\right)+\frac{\eta}{2}\|m\|_{2}^{2}$,
as well as the ‘damping matrix’
$\mathsf{D}_{\gamma}:=\left(\begin{array}[]{cc}0_{d}&-I_{d}\\\ I_{d}&\gamma
I_{d}\end{array}\right),$ which, upon abbreviating
$\vartheta=\left(\theta,m\right)$, allows us to express this system of ODEs as

$\displaystyle\dot{\vartheta}_{t}=-\mathsf{D}_{\gamma}\nabla_{\vartheta}F_{\eta}\left(\vartheta_{t}\right).$

This identifies the system as an instance of a ‘damped’ or ‘conformal’
Hamiltonian flow; e.g., see McLachlan and Perlmutter,  (2001); Maddison et
al.,  (2018). Stated concisely, we observe that NAG admits an interpretation
as a ‘$\gamma$-damped Hamiltonian flow of $F_{\eta}$’. Related observations
were made in Wibisono et al.,  (2016); Wilson et al.,  (2016).

###  2.2 Gradient Flows on $\mathcal{P}\left(\mathbb{R}^{d}\right)$, and their
Acceleration

At least in abstract, one can consider following a similar program for solving
optimisation tasks over the space of probability distributions on Euclidean
space, $\mathcal{P}\left(\mathbb{R}^{d}\right)$. One challenge in making this
idea practical is how to endow this space with a suitable geometric structure
compatible with relevant algorithmic primitives. That is, when considering
functionals $\mathcal{F}:\mathcal{P}\left(\mathbb{R}^{d}\right)\to\mathbb{R}$
and distributions $q,\tilde{q}$ which are ‘close’ to one another in a relevant
sense, we want to be able to make estimates of the form
$\mathcal{F}\left(\tilde{q}\right)\approx\mathcal{F}\left(q\right)+\mathrm{D}\mathcal{F}\left[q\right]\left(\tilde{q}-q\right),$
for some suitable notion of derivative $\mathrm{D}\mathcal{F}$. In our work,
we consider the geometry induced by the Wasserstein-2 metric (Jordan et al., ,
1998). Other choices for distances will generate different geometries and
algorithms (e.g., the Stein geometry (Duncan et al., , 2023; Liu, , 2017;
Sharrock and Nemeth, , 2023)).

The important objects are then as follows: given a functional
$\mathcal{F}:\mathcal{P}\left(\mathbb{R}^{d}\right)\to\mathbb{R}$ and
distributions $q,\tilde{q}\in\mathcal{P}\left(\mathbb{R}^{d}\right)$, define
the functional gradient of
$\delta_{q}\mathcal{F}:\mathcal{P}\left(\mathbb{R}^{d}\right)\to\left(\mathbb{R}^{d}\to\mathbb{R}\right)$
so that
$q\approx\tilde{q}\implies\mathcal{F}\left(\tilde{q}\right)\approx\mathcal{F}\left(q\right)+\int\delta_{q}\mathcal{F}\left[q\right]\left(x\right)\,\left(\tilde{q}\left(x\right)-q\left(x\right)\right)\,\mathrm{d}x.$
Now, let $v$ denote a vector field, let $\epsilon>0$, and write $q_{\epsilon}$
for the law of $x+\epsilon v\left(x\right)$ when $x\sim q$. Under suitable
assumptions, one can then compute that
$\mathcal{F}\left(q_{\epsilon}\right)=\mathcal{F}\left(q\right)+\epsilon\int
q\left(\mathrm{d}x\right)\left\langle\nabla_{x}\delta_{q}\mathcal{F}\left[q\right]\left(x\right),v\left(x\right)\right\rangle+o\left(\epsilon\right)$
$\text{as }\epsilon\to 0^{+}$, where
$\left\langle{\cdot},{\cdot}\right\rangle$ is the Euclidean inner product
(Ambrosio et al., , 2005, Lemma 10.4.1). Therefore, the Wasserstein gradient
of $\mathcal{F}$ is defined as
$\nabla_{W}\mathcal{F}\left(q\right)=-\nabla_{x}\cdot\left(q\nabla_{x}\delta_{q}\mathcal{F}\left[q\right]\right).$
Interpreting gradients from the perspective of steepest descent, this
expression reflects the notion that in order to reduce
$\mathcal{F}\left(q\right)$ by updating $q$ by a perturbation which is small
in the sense of the 2-Wasserstein geometry, it would be suitable to take each
$x$ in the support of $q$ and transport it towards
$x-\epsilon\nabla_{x}\delta_{q}\mathcal{F}\left[q\right]\left(x\right)$ (for
more details, see Section C.1).

Having defined these notions, the Wasserstein gradient flow (WGF) of the
functional $\mathcal{F}$ then corresponds to the partial differential equation
(PDE) $\dot{q}_{t}=-\nabla_{W}\mathcal{F}\left(q_{t}\right).$ A significant
example of such a flow is given by taking
$\mathcal{F}:q\mapsto\mathsf{KL}\left(q,p\right)$ for some fixed
$p\in\mathcal{P}\left(\mathbb{R}^{d}\right)$. In this case, one can compute
$\delta_{q}\mathcal{F}\left[q\right]=\log\frac{q}{p}$ and so, the WGF is given
by:

$\displaystyle\dot{q}_{t}=-\nabla_{W}\mathcal{F}\left(q_{t}\right)$
$\displaystyle=\nabla_{x}\cdot\left(q_{t}\nabla_{x}\log\left(\frac{q_{t}}{p}\right)\right)=-\nabla_{x}\cdot\left(q_{t}\nabla_{x}\log
p\right)+\Delta_{x}q_{t}.$

An insight of Jordan et al.,  (1998) was to observe that this PDE is precisely
the Fokker-Planck PDE corresponding to the overdamped Langevin SDE with
invariant measure $p$, i.e., $\mathrm{d}X_{t}=\nabla_{x}\log
p\left(x\right)\,\mathrm{d}t+\sqrt{2}\,\mathrm{d}W_{t}$ where $W_{t}$ denotes
the standard Wiener process. As such, we can interpret the simulation of the
overdamped Langevin SDE as gradient descent of $\mathcal{F}$. Following the
earlier discussion of momentum-enriched dynamical systems for accelerated
optimisation over $\mathbb{R}^{d}$, it is natural to ask whether, with a
suitably modified objective functional, and a suitable interpretation of
damped Hamiltonian flows, can one obtain accelerated methods for solving
optimisation problems over the space of probability distributions?

This question was answered in the affirmative by Ma et al.,  (2021), who
proceed in two steps. Firstly, instead of working with the functional
$\mathcal{F}:\mathcal{P}\left(\mathbb{R}^{d}\right)\to\mathbb{R}$ which sends
$q$ to $\mathsf{KL}\left(q,p\right)$, they fix a hyperparameter $\eta>0$ and
define a new functional on the space of momentum-enriched probability
distributions, i.e.
$\mathcal{G}_{\eta}:\mathcal{P}\left(\mathbb{R}^{d}\times\mathbb{R}^{d}\right)\to\mathbb{R}:q\mapsto\mathsf{KL}\left(q,p\otimes
r_{\eta}\right)$ where $r_{\eta}=\mathcal{N}\left(u\mid
0,\eta^{-1}I_{d}\right)$ and $p\otimes q$ is the product measure; $u$ will
denote momentum variables which are conjugate to $x$ in all that follows.
Secondly, they extend the notion of a gradient flow over
$\mathcal{P}\left(\mathbb{R}^{d}\right)$ to the notion of a $\gamma$-damped
Hamiltonian flow of $\mathcal{G}_{\eta}$ over
$\mathcal{P}\left(\mathbb{R}^{d}\times\mathbb{R}^{d}\right)$ as follows:
abbreviating $\Upsilon:=\left(x,u\right)$, define an evolution of probability
distributions
$q_{t}\in\mathcal{P}\left(\mathbb{R}^{d}\times\mathbb{R}^{d}\right)$ by
$\dot{q}_{t}=\nabla_{\Upsilon}\cdot\left(q_{t}\mathsf{D}_{\gamma}\nabla_{\Upsilon}\delta_{q}\mathcal{G}_{\eta}\left[q_{t}\right]\right).$
Thirdly, much as Jordan et al.,  (1998) observe that the WGF of $\mathcal{F}$
corresponds to the evolution of the overdamped Langevin SDE, Ma et al.,
(2021) observe that this damped Hamiltonian flow admits an analogous
correspondence with the underdamped Langevin SDE. This amounts to the claim
that the underdamped Langevin SDE admits an interpretation as a
‘$\gamma$-damped Hamiltonian flow of $\mathcal{G}_{\eta}$’ in the space of
probability distributions. Finally, Ma et al.,  (2021) shows that under
suitable conditions, this process converges to its stationary distribution
more rapidly than the corresponding overdamped Langevin SDE. As such, the
transition from gradient flow to damped Hamiltonian flow is again able to
deliver accelerated rates of convergence.

###  2.3 Gradient Flows on
$\mathbb{R}^{d_{\theta}}\times\mathcal{P}\left(\mathbb{R}^{d_{x}}\right)$

Combining the previous two settings, we now consider the optimization of
objectives defined over a product of an Euclidean space and probability
distributions over another Euclidean space (potentially of distinct
dimension). Given a functional
$\mathcal{E}:\mathbb{R}^{d_{\theta}}\times\mathcal{P}\left(\mathbb{R}^{d_{x}}\right)\to\mathbb{R}$,
the ‘extended’ gradient flow of $\mathcal{E}$ with respect to the product
Euclidean-Wasserstein geometry is given by

$\displaystyle\dot{\theta}_{t}$
$\displaystyle=-\nabla_{\theta}\mathcal{E}\left(\theta_{t},q_{t}\right),\quad\dot{q}_{t}=-\nabla_{W}\mathcal{E}\left(\theta_{t},q_{t}\right),$

where $\nabla_{W}$ denotes the Wasserstein gradient with respect to the $q$
variable.

In Kuntz et al.,  (2023), this extended gradient flow is used as a tool for
solving maximum likelihood estimation problems for latent variable models. In
particular, they seek to minimize the objective functional
$\mathcal{E}\left(\theta,q\right)$. Thus, the extended gradient flow takes the
form

$\displaystyle\dot{\theta}_{t}$ $\displaystyle=\int
q_{t}\left(\mathrm{d}x\right)\nabla_{\theta}\ell\left(\theta_{t},x\right),\quad\dot{q}_{t}=\nabla_{x}\cdot\left(q_{t}\left(x\right)\nabla_{x}\left(\log
q_{t}\left(x\right)-\ell\left(\theta_{t},x\right)\right)\right),$

where $\ell(\theta,x):=\log\rho_{\theta}(x)$ and
$\rho_{\theta}(x):=p_{\theta}(y,x)$. Following the approach of Jordan et al.,
(1998), we can interpret the gradient flow of $\mathcal{E}$ as describing the
evolution of the coupled ODE-SDE system

$\displaystyle\dot{\theta}_{t}$ $\displaystyle=\int
q_{t}\left(\mathrm{d}x\right)\nabla_{\theta}\ell\left(\theta_{t},x\right),\quad\mathrm{d}X_{t}=\nabla_{\theta}\ell\left(\theta_{t},X_{t}\right)\,\mathrm{d}t+\sqrt{2}\,\mathrm{d}W_{t},$

where $q_{t}:=\mathrm{Law}\left(X_{t}\right)$. Since the drift of $\theta_{t}$
depends on the law of $X_{t}$, this is said to be a ‘distribution-dependent’
or ‘McKean-Vlasov’ process (Kac, , 1956; McKean Jr, , 1966). PGD is obtained
via discretization in time (see Section I.1).

##  3 Momentum Particle-based Maximum Likelihood Flow

Drawing inspiration from prior works as detailed in Section 2, we endeavour to
accelerate gradient flows over the extended space of parameters and
probability distributions by i) enriching both components of the space with
momentum variables, ii) defining a suitable extension of the functional of
interest which is defined over this momentum-enriched extended space, iii)
constructing a suitable damped Hamiltonian flow on this space, and iv)
interpreting the resulting flows as a coupled system of ODEs and SDEs, which
can then be discretised. In particular, we will carry out this program to
describe a (doubly-)momentum-enriched modification of the PGD algorithm.
First, we enrich both spaces with a momentum variable, which is performed by
equipping the parameter space with momentum $m\in\mathbb{R}^{d_{\theta}}$, and
the latent space with momentum $u\in\mathbb{R}^{d_{x}}$. Accordingly, we move
from working on the space of probability distributions over the latent space
$\mathbb{R}^{d_{x}}$ to the space of probability distributions over the
momentum-enriched latent space, i.e.
$\mathcal{P}\left(\mathbb{R}^{d_{x}}\times\mathbb{R}^{d_{x}}\right)$.
Subsequently, we fix positive friction hyperparameters
$\left(\eta_{\theta},\eta_{x}\right)$ (corresponding to the parameter and
latent spaces respectively), define the (unnormalised) momentum-enriched joint
target law
$\rho_{\theta,\eta_{x}}\left(x,u\right):=p_{\theta}\left(y,x\right)r_{\eta_{x}}\left(u\right)$
(where again we denote $r_{\eta_{x}}\left(u\right)=\mathcal{N}\left(u\mid
0,\eta_{x}^{-1}I_{d_{x}}\right)$), and finally, define the doubly momentum-
enriched free energy functional
$\mathcal{F}_{\eta_{\theta},\eta_{x}}:\mathbb{R}^{d_{\theta}}\times\mathbb{R}^{d_{\theta}}\times\mathcal{P}\left(\mathbb{R}^{d_{x}}\times\mathbb{R}^{d_{x}}\right)$
as

$\displaystyle\mathcal{F}_{\eta_{\theta},\eta_{x}}\left(\theta,m,q\right):=\int
q\left(x,u\right)\log\left(\frac{q\left(x,u\right)}{\rho_{\theta,\eta_{x}}\left(x,u\right)}\right)\,\mathrm{d}x\,\mathrm{d}u+\frac{\eta_{\theta}}{2}\|m\|_{2}^{2}.$

We now define the extended $\left(\gamma_{\theta},\gamma_{x}\right)$-damped
Hamiltonian flow of $\mathcal{F}_{\eta_{\theta},\eta_{x}}$ as (abbreviating
$\vartheta=\left(\theta,m\right)$ and $\Upsilon=\left(x,u\right)$)

$\displaystyle\dot{\vartheta}_{t}$
$\displaystyle=-\mathsf{D}_{\gamma_{\theta}}\nabla_{\vartheta}\mathcal{F}_{\eta_{\theta},\eta_{x}}\left(\vartheta_{t},q_{t}\right)$
$\displaystyle\dot{q}_{t}$
$\displaystyle=\nabla_{\Upsilon}\cdot\left(q_{t}\mathsf{D}_{\gamma_{x}}\nabla_{\Upsilon}\mathcal{\delta}_{q}{F}_{\eta_{\theta},\eta_{x}}\left[\vartheta_{t},q_{t}\right]\right).$

Following a similar path to earlier derivations, this evolution corresponds to
the coupled, momentum-enriched ODE-SDE system

$\displaystyle\mathrm{d}\theta_{t}$
$\displaystyle=\eta_{\theta}m_{t}\,\mathrm{d}t$ (2a)
$\displaystyle\mathrm{d}m_{t}$
$\displaystyle=-\gamma_{\theta}\eta_{\theta}m_{t}\,\mathrm{d}t-\nabla_{\theta}\mathcal{E}(\theta_{t},q_{t,X})\mathrm{d}t$
(2b) $\displaystyle\mathrm{d}X_{t}$ $\displaystyle=\eta_{x}U_{t}\,\mathrm{d}t$
(2c) $\displaystyle\mathrm{d}U_{t}$
$\displaystyle=-\gamma_{x}\eta_{x}U_{t}\,\mathrm{d}t+\nabla_{\theta}\ell\left(\theta_{t},X_{t}\right)\,\mathrm{d}t+\sqrt{2\gamma_{x}}\,\mathrm{d}W_{t},$
(2d)

where $q_{t,X}=\mathrm{Law}\left(X_{t}\right)$. This is again a McKean-Vlasov
process. In Proposition 3.1, we establish the existence and uniqueness of
strong solutions to (2). The proof can be found in
LABEL:{proof:existence_uniqueness}.

######  Proposition 3.1 (Existence and Uniqueness of strong solutions to (2)).

Under 2, there exists a unique strong solution to (2) for any initial
condition
$(\vartheta_{0},q_{0})\in\mathbb{R}^{2d_{\theta}}\times\mathcal{P}(\mathbb{R}^{2d_{x}})$.

The choice of the hyperparameters
$\left(\eta_{\theta},\gamma_{\theta},\eta_{x},\gamma_{x}\right)$ is crucial to
the practical performance of MPGD. We observe that (as is also the case for
NAG, the underdamped Langevin SDE, and other similar systems) there are three
different qualitative regimes for the dynamics: i) the underdamped regime, in
which the parameter values oscillate, ii) the overdamped regime, in which one
recovers PGD-type behaviour, and iii) the critically-damped regime, in which
oscillations are largely suppressed, but the momentum effects are still able
to accelerate the convergence behaviour relative to PGD. However, obtaining
parameters to achieve critical damping is problem-dependent and, at the time
of writing, rigorous approaches are limited to simple targets (e.g., see
Dockhorn et al.,  (2022)).

##  4 Convergence of the Flow

In this section, we study the quantitative convergence properties of the flow
described in Section 3. Following the work of Wilson et al.,  (2016); Wibisono
et al.,  (2016); Ma et al.,  (2021), we establish convergence through a
suitable Lyapunov argument (Lyapunov, , 1992). In particular, we will exhibit
a nonnegative function $\mathcal{L}$ which majorises $\mathcal{F}$ and is
contracted at an exponential rate along the evolution of the flow. From this,
it will be straightforward to deduce estimates on the convergence of
$\mathcal{F}$ to its minimal value.

######  Assumption 1 (Extended Log Sobolev Inequality).

Suppose that $\theta^{*}$ maximizes the marginal likelihood, i.e.,
$\theta^{*}=\operatorname*{arg\,sup}_{\theta\in\mathbb{R}^{d_{\theta}}}p_{\theta}(y)$,
and that there exists a constant $C_{\mathcal{E}}>0$, such that

$\mathcal{E}(\theta,q)+\log
p_{\theta^{*}}\left(y\right)\leq\frac{1}{2C_{\mathcal{E}}}\left|\left|{\nabla\mathcal{E}(\theta,q)}\right|\right|^{2}_{q},\quad\forall\theta\in\mathbb{R}^{d_{\theta}},\enskip
q\in\mathcal{P}(\mathbb{R}^{d_{x}}),$

where

$\left|\left|{\nabla\mathcal{E}(\theta,q)}\right|\right|^{2}_{q}:=\left|\left|{\nabla_{\theta}\mathcal{E}(\theta,q)}\right|\right|^{2}+\left|\left|{\nabla_{x}\log{q}-\nabla_{x}\log{\rho_{\theta}}}\right|\right|^{2}_{q}.$
(3)

######  Assumption 2 ($K$-Lipschitz Gradient).

We assume that the potential $\ell(\theta,x)$ has a $K$-Lipschitz gradient.
More precisely, there exists some constant $K>0$ such that for all

$\left\|\nabla_{(\theta,x)}\ell(\theta,x)-\nabla_{(\theta,x)}\ell(\theta^{\prime},x^{\prime})\right\|\leq
K\left(\|x-x^{\prime}\|+\|\theta-\theta^{\prime}\|\right),\quad\forall\theta\in\mathbb{R}^{d_{\theta}},\forall
x\in\mathbb{R}^{d_{x}}.$

As a first step towards this, one can compute the time derivative of
$\mathcal{F}$ along the flow as

$\displaystyle\dot{\mathcal{F}}=-\gamma_{\theta}\left\|\nabla_{m}\mathcal{F}\right\|^{2}-\gamma_{x}\left\|\nabla_{u}\delta_{q}\mathcal{F}\right\|^{2}\leq
0;$

see Proposition E.1 for the calculations. This establishes that $\mathcal{F}$
is non-increasing along the flow, and while providing some useful a priori
assurances on the stability of the flow, is not sufficient for establishing
convergence to the minimiser of $\mathcal{F}$. This is due in part to the
degeneracy of this derivative term, noting that it contains no gradients with
respect to $\left(\theta,x\right)$; this is a common phenomenon for momentum-
enhanced dynamical systems. As such, we follow the constructions of Wilson et
al.,  (2016); Ma et al.,  (2021) and add a sort of ‘twisted’ gradient term, to
re-introduce these missing gradients. We proceed as follows: let
$\left(\tau_{\theta},\tau_{\theta
m},\tau_{m},\tau_{x},\tau_{xu},\tau_{x}\right)$ be non-negative constants such
that the matrices

$\displaystyle
T_{\left(\theta,m\right)}=K^{-1}\left(\begin{array}[]{cc}\tau_{\theta}I_{d_{\theta}}&\frac{\tau_{\theta
m}}{2}I_{d_{\theta}}\\\ \frac{\tau_{\theta
m}}{2}I_{d_{\theta}}&\tau_{m}I_{d_{\theta}}\end{array}\right),\quad
T_{\left(x,u\right)}=K^{-1}\left(\begin{array}[]{cc}\tau_{x}I_{d_{x}}&\frac{\tau_{xu}}{2}I_{d_{x}}\\\
\frac{\tau_{xu}}{2}I_{d_{x}}&\tau_{u}I_{d_{x}}\end{array}\right)$

are positive-semidefinite, with $K$ the Lipschitz gradient constant of 2.
Define now

$\displaystyle\left\|\nabla_{\left(\theta,m\right)}\mathcal{F}\right\|_{T_{\left(\theta,m\right)}}^{2}$
$\displaystyle:=\left\langle\nabla_{\left(\theta,m\right)}\mathcal{F},T_{\left(\theta,m\right)}\nabla_{\left(\theta,m\right)}\mathcal{F}\right\rangle,$
$\displaystyle\left\|\nabla_{\left(x,u\right)}\delta_{q}\mathcal{F}\right\|_{T_{\left(x,u\right)}}^{2}$
$\displaystyle:=\left\langle\nabla_{\left(x,u\right)}\delta_{q}\mathcal{F},T_{\left(x,u\right)}\nabla_{\left(x,u\right)}\delta_{q}\mathcal{F}\right\rangle,$

which are readily seen to be non-negative. Finally, define the putative
Lyapunov function

$\displaystyle\mathcal{L}\left(\theta,m,q\right):=\mathcal{F}\left(\theta,m,q\right)+\left\|\nabla_{\left(\theta,m\right)}\mathcal{F}\right\|_{T_{\left(\theta,m\right)}}^{2}+\left\|\nabla_{\left(x,u\right)}\delta_{q}\mathcal{F}\right\|_{T_{\left(x,u\right)}}^{2}.$

In LABEL:{prop:convergence_flow}, we provide sufficient conditions under which
$\mathcal{L}\left(\theta,m,q\right)$ contracts exponentially along the MPGD
flow. The conditions amount to suitable adaptations of classical conditions
like the Polyak-Lojasiewicz and Logarithmic Sobolev inequalities to the
present setting, along with a smoothness condition on the gradients of the
log-likelihood.

######  Proposition 4.1 (Exponential Convergence of the Continuous Flow).

Under the $K$-Lipschitz gradient 2 and the log-Sobolev 1 with constant
$C_{\mathcal{E}}$, given parameters
$\left(\gamma_{x},\gamma_{\theta},\eta_{x},\eta_{\theta}\right)$, if the
constants $\left(\tau_{\theta},\tau_{\theta
m},\tau_{m},\tau_{x},\tau_{xu},\tau_{u}\right)$ satisfy the inequalities in
(34) for some rate $\varphi$, then it holds that

$\displaystyle\dot{\mathcal{L}}_{t}\leq-\varphi
C_{\mathcal{E}}\mathcal{L}_{t}.$

Moreover, we have exponential convergence of $\mathcal{F}_{t}$, i.e.,
$\mathcal{F}_{t}+\log p_{\theta_{*}}\leq\mathcal{L}_{0}\exp\left(-\varphi
C_{\mathcal{E}}t\right)$.

For brevity, the assumptions, proof and inequalities are deferred to Appendix
F. In Section F.3, we demonstrate concrete choices of hyperparameters for
which the inequalities of (34) are satisfied.

##  5 Discretization

(a) Different regimes in MPGD.

(b) $\theta_{t}$ integrator comparison.

(c) Role of gradient correction.

Figure 1: Toy Hierarchical Model. (a) Different regimes that arise from
varying the momentum parameters; (b) Comparison between our Exponential (Exp)
integrator and a Nesterov-like integrator for different momentum parameters;
(c) we compare the performance of the MPGD with and without gradient
correction.

In order to realize an actionable algorithm, we need to discretize the
dynamics of Equation (2) both in space and in time. Following the approach of
Kuntz et al.,  (2023), we approximate
$q_{t,X}\approx\frac{1}{M}\sum_{i\in\left[M\right]}\delta_{X^{i,M}_{t}}=:q^{M}_{t,X}$
using a set of particles. This yields the system

$\displaystyle\mathrm{d}\theta^{M}_{t}$
$\displaystyle=\eta_{\theta}m^{M}_{t}\,\mathrm{d}t$ (4a)
$\displaystyle\mathrm{d}m^{M}_{t}$
$\displaystyle=-\gamma_{\theta}\eta_{\theta}m^{M}_{t}\,\mathrm{d}t-\nabla_{\theta}\mathcal{E}\left(\theta^{M}_{t},q^{M}_{t,X}\right)\,\mathrm{d}t$
(4b) $\displaystyle\text{for }i\in[M],\quad\mathrm{d}X_{t}^{i,M}$
$\displaystyle=\eta_{x}U_{t}^{i,M}\,\mathrm{d}t$ (4c) $\displaystyle\text{for
}i\in[M],\quad\mathrm{d}U_{t}^{i,M}$
$\displaystyle=-\gamma_{x}\eta_{x}U_{t}^{i,M}\,\mathrm{d}t+\nabla_{\theta}\ell\left(\theta_{t}^{M},X_{t}^{i,M}\right)\,\mathrm{d}t+\sqrt{2\gamma_{x}}\,\mathrm{d}W_{t}^{i}.$
(4d)

where $\\{W_{t}^{i}\\}_{i=1}^{M}$ comprises $M$ independent standard Wiener
processes. Proposition 5.1 provides asymptotic justification about using a
particle approximation for $q_{t,X}$ to approximate the flow in (2). The proof
can be found in the Appendix H.

######  Proposition 5.1.

Under 2, we have

$\lim_{M\rightarrow\infty}\mathbb{E}\sup_{t\in[0,T]}\\{\|\vartheta_{t}-\vartheta_{t}^{M}\|^{2}+W^{2}_{2}(q_{t}^{\otimes
M},q^{M}_{t})\\}=0,$

where $\vartheta_{t}:=(\theta_{t},m_{t})$, $q_{t}^{\otimes
M}:=\Pi_{i=1}^{M}q_{t}$ with $q_{t}=\mathrm{Law}(X_{t},U_{t})$ are defined by
the SDE in (2); and $\vartheta^{M}_{t}:=(\theta^{M}_{t},m_{t}^{M})$ , and
$q^{M}_{t}=\mathrm{Law}(\\{(X_{t}^{i,M},U_{t}^{i,M})\\}_{i=1}^{M})$ in (4).

As noted by Ma et al.,  (2021), naive application of an Euler-Maruyama scheme
to momentum-enriched dynamics may be insufficient for attaining an accelerated
rate of convergence. We appeal to the literature on the discretization of
underdamped Langevin dynamics to design an appropriate integrator. One
integrator that can achieve accelerated rates is (Euler) Exponential
Integrator (Cheng et al., , 2018; Sanz-Serna and Zygalakis, , 2021; Hochbruck
and Ostermann, , 2010). The main strategy is to ‘freeze’ the nonlinear
components of the dynamics in a suitable way, and then solve the resulting
linear SDE. We also incorporate a partial update inspired by Sutskever et al.,
(2013)’s interpretation of NAG. More precisely, given
$\left(\theta_{0},m_{0},\\{\left(X^{i}_{0},U_{0}^{i}\right)\\}_{i=1}^{M}\right)$,
define an approximating SDE on the time interval $t\in\left[0,h\right]$ as

$\displaystyle\mathrm{d}\tilde{\theta}_{t}$
$\displaystyle=\eta_{\theta}\tilde{m}_{t}\,\mathrm{d}t$ (5a)
$\displaystyle\mathrm{d}\tilde{m}_{t}$
$\displaystyle=-\gamma_{\theta}\eta_{\theta}\tilde{m}_{t}\,\mathrm{d}t-\nabla_{\theta}\mathcal{E}\left(\bar{\theta}_{0},\tilde{q}_{0,X}^{M}\right)\,\mathrm{d}t$
(5b) $\displaystyle\text{for }i\in[M],\quad\mathrm{d}\tilde{X}_{t}^{i}$
$\displaystyle=\eta_{x}\tilde{U}_{t}^{i}\,\mathrm{d}t$ (5c)
$\displaystyle\text{for }i\in[M],\quad\mathrm{d}\tilde{U}_{t}^{i}$
$\displaystyle=-\gamma_{x}\eta_{x}\tilde{U}_{t}^{i}\,\mathrm{d}t+\nabla_{x}\ell\left(\tilde{\theta}_{h},X_{0}^{i}\right)\,\mathrm{d}t+\sqrt{2\gamma_{x}}\,\mathrm{d}W_{t}^{i},$
(5d)

where
$\tilde{q}_{t,X}^{M}:=\frac{1}{M}\sum_{i=1}^{M}\delta_{\tilde{X}_{t}^{i}}$,
$\bar{\theta}_{0}:=\theta_{0}+\frac{\iota_{\theta}(h)}{\gamma_{\theta}}\tilde{m}_{0}$,
and
$\iota_{\theta}\left(t\right):=1-\exp\left(-\eta_{\theta}\gamma_{\theta}t\right)$
which (for reasons which will become clear when solving the SDE) can be
thought of as a partial update to $\tilde{\theta}_{h}$. This is a linear SDE,
which can then be solved analytically, yielding the following iteration (see
Section I.4 for details):

$\displaystyle\tilde{\theta}_{k}$
$\displaystyle=\tilde{\theta}_{k-1}+\frac{\iota_{\theta}\left(h\right)}{\gamma_{\theta}}\tilde{m}_{k-1}-\frac{1}{\gamma_{\theta}}\left(h-\frac{\iota_{\theta}\left(h\right)}{\gamma_{\theta}\eta_{\theta}}\right)\nabla_{\theta}\mathcal{E}\left(\bar{\theta}_{k-1},\tilde{q}_{k-1,X}^{M}\right)$
$\displaystyle\tilde{m}_{k}$
$\displaystyle=(1-\iota_{\theta}\left(h\right))\tilde{m}_{k-1}-\frac{\iota_{\theta}\left(h\right)}{\gamma_{\theta}\eta_{\theta}}\nabla_{\theta}\mathcal{E}\left(\bar{\theta}_{k-1},\tilde{q}_{k-1,X}^{M}\right)$
$\displaystyle\text{for }i\in[M],\quad\tilde{X}_{k}^{i}$
$\displaystyle=\tilde{X}_{k-1}^{i}+\frac{\iota_{x}\left(h\right)}{\gamma_{x}}\tilde{U}_{k-1}^{i}+\frac{1}{\gamma_{x}}\left(h-\frac{\iota_{x}\left(h\right)}{\gamma_{x}\eta_{x}}\right)\nabla_{x}\ell\left(\tilde{\theta}_{k},\tilde{X}_{k-1}^{i}\right)+L_{\Sigma}^{XX}\xi_{k}^{i}$
$\displaystyle\text{for }i\in[M],\quad\tilde{U}_{k}^{i}$
$\displaystyle=(1-\iota_{x}\left(h\right))\tilde{U}_{k-1}^{i}+\frac{\iota_{x}\left(h\right)}{\gamma_{x}\eta_{x}}\nabla_{x}\ell\left(\tilde{\theta}_{k},\tilde{X}_{k-1}^{i}\right)+L_{\Sigma}^{XU}\xi_{k}^{i}+L_{\Sigma}^{UU}\xi_{k}^{\prime,i},$

where $\iota_{x}\left(t\right):=1-\exp\left(-\eta_{x}\gamma_{x}t\right)$,
$\left\\{\xi_{k}^{i},\xi_{k}^{\prime,i}\right\\}_{i\in\left[M\right]}\overset{\text{i.i.d.}}{\sim}\mathcal{N}\left(0_{d_{x}},I_{d_{x}}\right)$,
$\bar{\theta}_{k-1}:=\tilde{\theta}_{k-1}+\frac{\iota_{\theta}(h)}{\gamma_{\theta}}\tilde{m}_{k-1}$,
and $\left(L_{\Sigma}^{XX},L_{\Sigma}^{XU},L_{\Sigma}^{UU}\right)$ are
suitable scalars that depend on $(\eta_{x},\gamma_{x})$, and which can be
found in the Appendix Equation 64.

The astute reader will notice that our approximating ODE-SDE (i.e., Equation
5) deviates from Cheng et al.,  (2018)’s style of discretization.
Specifically, the difference lies in where we compute the gradient in
Equations 5b and 5d which we refer to as gradient correction. This subtle
choice reflects the decision to utilize partially and fully updated
approximations to $\tilde{\theta}_{h}$, i.e., the partially updated
$\bar{\theta_{0}}$ in Equation 5b and fully updated $\tilde{\theta}_{h}$ in
Equation 5d. This is reminiscent of NAG’s use of a partial parameter update to
compute the gradient, unlike Polyak’s Heavy Ball, which does not (Sutskever et
al., , 2013, Section 2.1). In our case, we empirically found that this choice
is critical for a more stable discretization and, as a result, enables the
algorithm to take larger step sizes to allow the algorithm to travel “faster”
(see Toy HMM in Section 6).

##  6 Experiment

(a) MPGD.

(b) PGD.

(c) VI.

(d) ABP.

Figure 2: CIFAR-10. Samples generated from various algorithms.

###  6.1 Toy Hierarchical Model

As a toy example, we consider the hierarchical model proposed in Kuntz et al.,
(2023). The model is given by
$p_{\theta}(y,x):=\prod_{i=1}^{N}\mathcal{N}(y_{i}|x_{i},1)\mathcal{N}(x_{i}|\theta,1)$,
where $N=100$ is the number of data points. The dataset $y$ is sampled from a
model with $\theta=100$. In this experiment, we wish to understand the
behaviour of MPGD compared against PGD. We are particularly interested in (1)
how the momentum parameters
$\left(\eta_{\theta},\gamma_{\theta},\eta_{x},\gamma_{x}\right)$ affect the
optimization process; (2) comparing our proposed discretization to another
that follows in the style of NAG (detailed in Section I.3); and (3) the role
of the gradient correction term described in Section 5. The results are shown
in Figure 1. Further experiment details can be found in Section J.4.1.

In Figure 1(a), we show that different regimes can arise from different
choices of hyper-parameters (as discussed in Section 3). In Figure 1(b), we
compare the performance of MPGD using (our) Exponential integrator with a NAG-
like discretization for $(\theta_{t},m_{t})$-component. We vary the “momentum
coefficient”, i.e., we have $\mu\in\\{0.5,0.8,0.9\\}$ with
$\mu_{\theta}=\mu_{x}=\mu$. It can be seen that MPGD with our exponential
integrator for $\theta_{t}$ performs better than NAG-like integrator (see
Section J.3 for more discussion). In Figure 1(c), we show the effect of the
gradient correction term for three different step sizes in
$(\theta,m)$-components while keeping the step sizes in $(X,U)$ fixed. The
different lines in the figure are generated from a step size of $c\cdot
10^{-5/2}$ where $c\in\\{0.99,1,1.01\\}$. It can be seen that our proposed
method is more effective than the other discretization. In red, we show MPGD
with the absence of gradient correction in Equation 5b (i.e., when we use
$\theta_{0}$ instead of $\bar{\theta}_{0}$ in Equation 5b), and, in green, we
show the MPGD when the gradient correction is absent in both Equation 5b and
Equation 5d (i.e. when we use $\theta_{0}$ instead of $\bar{\theta}_{0}$ and
$\tilde{\theta}_{h}$ in Equation 5b and Equation 5d). It can be seen that
usage of the gradient correction term results in a more stable algorithm.

###  6.2 Image Generation

Figure 3: Posterior Cloud vs Epochs. We show the evolution of the
reconstruction of a particle for persistent methods. The particle is taken at
epoch $\\{10,20,40\\}$ on MNIST and CIFAR-10.

For this task, we consider two datasets: MNIST (LeCun et al., , 1998) and
CIFAR-10 (Krizhevsky and Hinton, , 2009). For the model, we use a variational
autoencoder (Kingma and Welling, , 2014) with a VampPrior (Tomczak and
Welling, , 2018). As baselines, we compare our proposed method MPGD against
PGD, Alternating Backpropagation (ABP) (Han et al., , 2017), Short Run (SR)
(Nijkamp et al., , 2020), and amortized variational inference (VI) (Kingma and
Welling, , 2014). ABP is the most similar to PGD. They differ in that ABP
takes multiple steps of the (unadjusted and overdamped) Langevin algorithm
(ULA) instead of PGD’s single step. They (MPGD, PGD, ABP) are also persistent,
meaning that the ULA chain starts at the previous particle location, whereas
SR restarts the chain at a random location sampled from the prior but like ABP
runs the chain for several steps. Further experiment details can be found in
Section J.4.2.

We are interested in the generative performance of the resulting models. For a
qualitative measure, we show the samples produced by the model in Figure 2 for
CIFAR (for MNIST samples, see Figure 5 in the Appendix). As a quantitative
measure, we report the Fréchet inception distance (FID) (Heusel et al., ,
2017) as shown in Table 1. It can be seen that our proposed method does well
compared against other baselines. The closest competitor is ABP. As noted by
Kuntz et al.,  (2023), ABP can reap the benefits of taking multiple ULA steps
to locate the posterior mode quickly and reduce the transient phase. This
hypothesis is further confirmed in Figure 3, where we visualize the evolution
of a single particle across various epochs. It can be seen that ABP’s and
MPGD’s distinct methods of reducing the transient phase (by taking multiple
steps or utilizing momentum) are effective. As a competitor SR does perform
badly, which we attribute to its non-persistent property: by restarting the
particles at the prior and running short chains, it is unable to overcome the
bias of short chains and locate the posterior mode. Variational inference
performs well for CIFAR but surprisingly falls slightly short in MNIST. We
hypothesise that this is due to the choice of parametrization of the
variational distribution.

Method | MPGD | PGD | ABP | SR | VI
---|---|---|---|---|---
MNIST | $\bm{48.7}\pm 1.5$ | $102.4\pm 1.1$ | $\bm{53.9}\pm 13$ | $134.34\pm 1.6$ | $72.6\pm 1.5$
CIFAR | $\bm{93.2}\pm 1.9$ | $104.2\pm 5.4$ | $96.9\pm 3.1$ | $140.6\pm 16.3$ | $\bm{93.4}\pm 0.9$
Table 1: FID scores on MNIST and CIFAR-10 after the final epoch. In bold, we
indicate the lowest two scores on average. We write $(\mu\pm\sigma)$ to
indicate the mean $\mu$ and standard deviation $\sigma$ of the FID score
calculated over three independent trials.

##  7 Conclusion and Future work

We presented momentum PGD, a method of incorporating momentum into particle
gradient descent. We establish the convergence in continuous time, and through
experiments, we showed that, for suitably chosen momentum parameters, the
resulting algorithm achieves better performance than PGD. Future work
includes: a systematic method of tuning momentum parameters
$\left(\eta_{\theta},\gamma_{\theta},\eta_{x},\gamma_{x}\right)$ (or some
justification for using the momentum coefficient heuristic), theoretical
characterization of the difference between our discretization scheme compared
with other potential schemes akin to Sanz-Serna and Zygalakis,  (2021).

## References

  * Ambrosio et al., (2005) Ambrosio, L., Gigli, N., and Savaré, G. (2005).  Gradient flows: in metric spaces and in the space of probability measures.  Springer Science & Business Media.
  * Andrieu and Thoms, (2008) Andrieu, C. and Thoms, J. (2008).  A tutorial on adaptive mcmc.  Statistics and computing, 18:343–373.
  * Bakry and Émery, (2006) Bakry, D. and Émery, M. (2006).  Diffusions hypercontractives.  In Séminaire de Probabilités XIX 1983/84: Proceedings, pages 177–206. Springer.
  * Bernton, (2018) Bernton, E. (2018).  Langevin monte carlo and jko splitting.  In Conference on learning theory, pages 1777–1798. PMLR.
  * Carmona, (2016) Carmona, R. (2016).  Lectures on BSDEs, stochastic control, and stochastic differential games with financial applications.  SIAM.
  * Chen et al., (2022) Chen, F., Ren, Z., and Wang, S. (2022).  Uniform-in-time propagation of chaos for mean field langevin dynamics.  arXiv preprint arXiv:2212.03050.
  * Cheng et al., (2018) Cheng, X., Chatterji, N. S., Bartlett, P. L., and Jordan, M. I. (2018).  Underdamped Langevin MCMC: A non-asymptotic analysis.  In Conference on Learning Theory, pages 300–323. PMLR.
  * Chizat, (2022) Chizat, L. (2022).  Mean-field langevin dynamics: Exponential convergence and annealing.  arXiv preprint arXiv:2202.01009.
  * Chizat and Bach, (2018) Chizat, L. and Bach, F. (2018).  On the global convergence of gradient descent for over-parameterized models using optimal transport.  Advances in neural information processing systems, 31.
  * De Bortoli et al., (2021) De Bortoli, V., Durmus, A., Pereyra, M., and Vidal, A. F. (2021).  Efficient stochastic optimisation by unadjusted langevin monte carlo: Application to maximum marginal likelihood and empirical bayesian estimation.  Statistics and Computing, 31:1–18.
  * Delyon et al., (1999) Delyon, B., Lavielle, M., and Moulines, E. (1999).  Convergence of a stochastic approximation version of the em algorithm.  Annals of statistics, pages 94–128.
  * Dempster et al., (1977) Dempster, A. P., Laird, N. M., and Rubin, D. B. (1977).  Maximum likelihood from incomplete data via the EM algorithm.  Journal of the Royal Statistical Society: Series B (Methodological), 39(1):1–22.
  * Diao et al., (2023) Diao, M. Z., Balasubramanian, K., Chewi, S., and Salim, A. (2023).  Forward-backward gaussian variational inference via jko in the bures-wasserstein space.  In International Conference on Machine Learning, pages 7960–7991. PMLR.
  * Dockhorn et al., (2022) Dockhorn, T., Vahdat, A., and Kreis, K. (2022).  Score-based generative modeling with critically-damped Langevin diffusion.  In International Conference on Learning Representations (ICLR).
  * Duncan et al., (2023) Duncan, A., Nüsken, N., and Szpruch, L. (2023).  On the geometry of Stein variational gradient descent.  Journal of Machine Learning Research, 24(56):1–39.
  * Garbuno-Inigo et al., (2020) Garbuno-Inigo, A., Nüsken, N., and Reich, S. (2020).  Affine invariant interacting langevin dynamics for bayesian inference.  SIAM Journal on Applied Dynamical Systems, 19(3):1633–1658.
  * Gelfand and Silverman, (2000) Gelfand, I. M. and Silverman, R. A. (2000).  Calculus of Variations.  Courier Corporation.
  * Good, (1983) Good, I. J. (1983).  Good thinking: The foundations of probability and its applications.  University of Minnesota Press.
  * Han et al., (2017) Han, T., Lu, Y., Zhu, S.-C., and Wu, Y. N. (2017).  Alternating back-propagation for generator network.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 31.
  * Heusel et al., (2017) Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. (2017).  GANs trained by a two time-scale update rule converge to a local nash equilibrium.  Advances in Neural Information Processing Systems, 30.
  * Hinton, (2002) Hinton, G. E. (2002).  Training products of experts by minimizing contrastive divergence.  Neural computation, 14(8):1771–1800.
  * Hochbruck and Ostermann, (2010) Hochbruck, M. and Ostermann, A. (2010).  Exponential integrators.  Acta Numerica, 19:209–286.
  * Hu et al., (2021) Hu, K., Ren, Z., Šiška, D., and Szpruch, Ł. (2021).  Mean-field langevin dynamics and energy landscape of neural networks.  In Annales de l’Institut Henri Poincare (B) Probabilites et statistiques, volume 57, pages 2043–2065. Institut Henri Poincaré.
  * Jordan et al., (1998) Jordan, R., Kinderlehrer, D., and Otto, F. (1998).  The variational formulation of the fokker–planck equation.  SIAM Journal on Mathematical Analysis, 29(1):1–17.
  * Kac, (1956) Kac, M. (1956).  Foundations of kinetic theory.  In Proceedings of The third Berkeley symposium on mathematical statistics and probability, volume 3, pages 171–197.
  * Kingma and Welling, (2014) Kingma, D. P. and Welling, M. (2014).  Auto-encoding variational bayes.  In Bengio, Y. and LeCun, Y., editors, 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings.
  * Krizhevsky and Hinton, (2009) Krizhevsky, A. and Hinton, G. (2009).  Learning multiple layers of features from tiny images.
  * Kruger, (2003) Kruger, A. Y. (2003).  On Fréchet subdifferentials.  Journal of Mathematical Sciences, 116(3):3325–3358.
  * Kuntz et al., (2023) Kuntz, J., Lim, J. N., and Johansen, A. M. (2023).  Particle algorithms for maximum likelihood training of latent variable models.  In Proceedings on 26th International Conference on Artificial Intelligence and Statistics (AISTATS), volume 206 of Proceedings of Machine Learning Research, pages 5134–5180.
  * Lambert et al., (2022) Lambert, M., Chewi, S., Bach, F., Bonnabel, S., and Rigollet, P. (2022).  Variational inference via wasserstein gradient flows.  Advances in Neural Information Processing Systems, 35:14434–14447.
  * LeCun et al., (1998) LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. (1998).  Gradient-based learning applied to document recognition.  Proceedings of the IEEE, 86(11):2278–2324.
  * Liu, (2017) Liu, Q. (2017).  Stein variational gradient descent as gradient flow.  Advances in Neural Information Processing Systems, 30.
  * Loève, (1977) Loève, M. (1977).  Probability Concepts, pages 151–176.  Springer New York, New York, NY.
  * Lyapunov, (1992) Lyapunov, A. M. (1992).  The general problem of the stability of motion.  International Journal of Control, 55(3):531–534.
  * Ma et al., (2021) Ma, Y.-A., Chatterji, N. S., Cheng, X., Flammarion, N., Bartlett, P. L., and Jordan, M. I. (2021).  Is there an analog of Nesterov acceleration for gradient-based MCMC?  Bernoulli, 27(3):1942 – 1992.
  * Maddison et al., (2018) Maddison, C. J., Paulin, D., Teh, Y. W., O’Donoghue, B., and Doucet, A. (2018).  Hamiltonian descent methods.  arXiv preprint arXiv:1809.05042.
  * Martens, (2010) Martens, J. (2010).  Deep learning via Hessian-free optimization.  In International Conference on Machine Learning, volume 27, pages 735–742.
  * McCall, (2010) McCall, M. W. (2010).  Classical Mechanics: From Newton to Einstein: A Modern Introduction.  John Wiley & Sons.
  * McKean Jr, (1966) McKean Jr, H. P. (1966).  A class of Markov processes associated with nonlinear parabolic equations.  Proceedings of the National Academy of Sciences, 56(6):1907–1911.
  * McLachlan and Perlmutter, (2001) McLachlan, R. and Perlmutter, M. (2001).  Conformal Hamiltonian systems.  Journal of Geometry and Physics, 39(4):276–300.
  * Mei et al., (2018) Mei, S., Montanari, A., and Nguyen, P.-M. (2018).  A mean field view of the landscape of two-layer neural networks.  Proceedings of the National Academy of Sciences, 115(33):E7665–E7671.
  * Neal and Hinton, (1998) Neal, R. M. and Hinton, G. E. (1998).  A view of the EM algorithm that justifies incremental, sparse, and other variants.  Learning in graphical models, pages 355–368.
  * Nemirovskij and Yudin, (1983) Nemirovskij, A. S. and Yudin, D. B. (1983).  Problem complexity and method efficiency in optimization.  Wiley-Interscience.
  * Nesterov, (2003) Nesterov, Y. (2003).  Introductory Lectures on Convex Optimization: A Basic Course.  Springer Science & Business Media.
  * Nesterov, (1983) Nesterov, Y. E. (1983).  A method of solving a convex programming problem with convergence rate $\mathcal{O}\bigl{(}\frac{1}{k^{2}}\bigr{)}$.  In Doklady Akademii Nauk, volume 269, pages 543–547. Russian Academy of Sciences.
  * Nijkamp et al., (2020) Nijkamp, E., Pang, B., Han, T., Zhou, L., Zhu, S.-C., and Wu, Y. N. (2020).  Learning multi-layer latent variable model via variational optimization of short run MCMC for approximate inference.  In Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part VI 16, pages 361–378. Springer.
  * Nitanda and Suzuki, (2017) Nitanda, A. and Suzuki, T. (2017).  Stochastic particle gradient descent for infinite ensembles.  arXiv preprint arXiv:1712.05438.
  * Nitanda et al., (2022) Nitanda, A., Wu, D., and Suzuki, T. (2022).  Convex analysis of the mean field langevin dynamics.  In International Conference on Artificial Intelligence and Statistics, pages 9741–9757. PMLR.
  * O’Donoghue and Candes, (2015) O’Donoghue, B. and Candes, E. (2015).  Adaptive restart for accelerated gradient schemes.  Foundations of Computational Mathematics, 15:715–732.
  * Pang et al., (2020) Pang, B., Han, T., Nijkamp, E., Zhu, S.-C., and Wu, Y. N. (2020).  Learning latent space energy-based prior model.  Advances in Neural Information Processing Systems, 33:21994–22008.
  * Platen and Bruti-Liberati, (2010) Platen, E. and Bruti-Liberati, N. (2010).  Numerical Solution of Stochastic Differential Equations with Jumps in Finance.  Springer Science & Business Media.
  * Polyak, (1964) Polyak, B. T. (1964).  Some methods of speeding up the convergence of iteration methods.  USSR Computational Mathematics and Mathematical Physics, 4(5):1–17.
  * Roberts and Rosenthal, (2009) Roberts, G. O. and Rosenthal, J. S. (2009).  Examples of adaptive mcmc.  Journal of computational and graphical statistics, 18(2):349–367.
  * Sanz-Serna and Zygalakis, (2021) Sanz-Serna, J. M. and Zygalakis, K. C. (2021).  Wasserstein distance estimates for the distributions of numerical approximations to ergodic stochastic differential equations.  The Journal of Machine Learning Research, 22(1):11006–11042.
  * Sharrock and Nemeth, (2023) Sharrock, L. and Nemeth, C. (2023).  Coin sampling: Gradient-based Bayesian inference without learning rates.  In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J., editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 30850–30882. PMLR.
  * Shi et al., (2021) Shi, B., Du, S. S., Jordan, M. I., and Su, W. J. (2021).  Understanding the acceleration phenomenon via high-resolution differential equations.  Mathematical Programming, 195:79–148.
  * Silvester, (2000) Silvester, J. R. (2000).  Determinants of block matrices.  The Mathematical Gazette, 84(501):460–467.
  * Staib et al., (2019) Staib, M., Reddi, S., Kale, S., Kumar, S., and Sra, S. (2019).  Escaping saddle points with adaptive gradient methods.  In International Conference on Machine Learning, pages 5956–5965. PMLR.
  * Su et al., (2014) Su, W., Boyd, S., and Candes, E. (2014).  A differential equation for modeling Nesterov’s accelerated gradient method: theory and insights.  Advances in Neural Information Processing Systems, 27.
  * Sutskever et al., (2013) Sutskever, I., Martens, J., Dahl, G., and Hinton, G. (2013).  On the importance of initialization and momentum in deep learning.  In International Conference on Machine Learning, pages 1139–1147. PMLR.
  * Suzuki et al., (2023) Suzuki, T., Wu, D., and Nitanda, A. (2023).  Convergence of mean-field langevin dynamics: Time and space discretization, stochastic gradient, and variance reduction.  arXiv preprint arXiv:2306.07221.
  * Tieleman and Hinton, (2012) Tieleman, T. and Hinton, G. (2012).  Lecture 6.5-rmsprop, coursera: Neural networks for machine learning.  University of Toronto, Technical Report, 6.
  * Tomczak and Welling, (2018) Tomczak, J. and Welling, M. (2018).  VAE with a VampPrior.  In International Conference on Artificial Intelligence and Statistics, pages 1214–1223. PMLR.
  * Van Handel, (2014) Van Handel, R. (2014).  Probability in high dimension.  Lecture Notes (Princeton University).
  * Villani, (2009) Villani, C. (2009).  Optimal transport: old and new, volume 338.  Springer.
  * Wibisono et al., (2016) Wibisono, A., Wilson, A. C., and Jordan, M. I. (2016).  A variational perspective on accelerated methods in optimization.  Proceedings of the National Academy of Sciences, 113(47):E7351–E7358.
  * Wilson et al., (2016) Wilson, A. C., Recht, B., and Jordan, M. I. (2016).  A Lyapunov analysis of momentum methods in optimization.  arXiv preprint arXiv:1611.02635.
  * Yao and Yang, (2022) Yao, R. and Yang, Y. (2022).  Mean field variational inference via wasserstein gradient flow.  arXiv preprint arXiv:2207.08074.

[sections] [sections]l1

##  Appendix A Notation

The following table summarizes some key notation used throughout.

$\mathcal{E}$ |  Free Energy.
---|---
$\mathcal{F}$ |  Momentum-enriched Free Energy
$z_{t}$ |  The tuple $(\theta_{t},m_{t},q_{t})\in\mathbb{R}^{d_{\theta}}\times\mathbb{R}^{d_{\theta}}\times\mathcal{P}(\mathbb{R}^{d_{x}}\times\mathbb{R}^{d_{x}})$
$\nabla_{a}f$ |  Euclidean Gradient of $f$ w.r.t. $a$
$\nabla_{(a,b)}f$ |  $\left[\nabla_{a}f,\nabla_{b}f\right]^{\top}$
$\ell$ |  $\ell(\theta,x):=\log p_{\theta}(y,x)$
$\rho_{\theta}$ |  $\rho_{\theta}(x):=p_{\theta}(y,x)$
$\nabla_{a}\cdot$ |  Divergence operator w.r.t. $a$
$\nabla^{*}_{a}$ |  Adjoint operator of $\nabla_{a}$ (see Section F.4.2).
$[n]$ |  $[n]:=\\{1,...,n\\}$
$\Delta$ |  Laplacian operator, $\Delta=\nabla\cdot\nabla$
$\mathcal{P}(\mathbb{R}^{d})$ |  The space of probability measures that are absolutely continuous w.r.t. Lebesgue measure (have densities) and possess finite second moments.
$\left\langle{\cdot},{\cdot}\right\rangle$ (and $\|\cdot\|$) |  Euclidean inner product or Frobenius inner product (and it’s inner norm)
$\left\langle{\cdot},{\cdot}\right\rangle_{\rho}$ (and $\|\cdot\|_{\rho}$) |  $L^{2}(\rho)$ inner product (and its norm)
$o(\epsilon)$ |  Bachman-Landau little-o notation
$\partial F(\mu)$ |  Fréchet subdifferential (see Definition C.1)

##  Appendix B Related Work

The present work sits at the juncture of i) deterministic gradient flows for
optimising objectives over “parameter” spaces, typically expressible through
the discretization of ODEs, and ii) stochastic gradient flows for optimising
objectives over the space of probability measures, typically expressible
through discretisation of mean-field SDEs. The former class of problems is too
vast to be properly surveyed here, effectively including a large proportion of
modern continuous optimisation problems. The latter class has seen substantial
growth over the past few years in particular, with various problems related to
sampling (Liu, , 2017; Bernton, , 2018; Garbuno-Inigo et al., , 2020; Duncan
et al., , 2023), variational inference (Yao and Yang, , 2022; Lambert et al.,
, 2022; Diao et al., , 2023), and the training of shallow neural networks (Mei
et al., , 2018; Chizat and Bach, , 2018; Nitanda and Suzuki, , 2017; Hu et
al., , 2021; Nitanda et al., , 2022; Chizat, , 2022; Chen et al., , 2022;
Suzuki et al., , 2023) being studied in this framework. While there exist
earlier works which combine optimisation with Markovian sampling (e.g.
stochastic approximation approaches to the EM algorithm (Delyon et al., ,
1999; De Bortoli et al., , 2021), training of energy-based models (Hinton, ,
2002), and hyperparameter tuning in MCMC (Andrieu and Thoms, , 2008; Roberts
and Rosenthal, , 2009)), the connection to gradient flows remains somewhat
under-developed at present. We hope that the present work can encourage
further exploration of these connections.

##  Appendix C Gradient Flow on
$\mathbb{R}^{d_{\theta}}\times\mathcal{P}(\mathcal{X})$

In this section, we will describe gradient flows on the extended space
$\mathbb{R}^{d_{\theta}}\times\mathcal{P}(\mathcal{X})$ where
$\mathcal{X}={\mathbb{R}^{d_{x}}}$. We begin with an exposition of gradient
flows in Wasserstein space, i.e. the space of distributions $\mathcal{P(X)}$
endowed with the Wasserstein-$2$ metric. We aim to give an intuitive
introduction as opposed to a rigorous one; those readers with an interest in
the latter are directed to Ambrosio et al.,  (2005). Using the ideas in
Ambrosio et al.,  (2005), we show how the notion of gradients can be
generalized to the product space
$\mathbb{R}^{d_{\theta}}\times\mathcal{P}(\mathcal{X})$.

###  C.1 Gradient Flow on $\mathcal{P(X)}$

This exposition aims to introduce the intuition and necessary objects. For a
rigorous treatment, we point the reader to (Ambrosio et al., , 2005). We begin
by describing the gradient flow on $\mathcal{P(X)}$ endowed with the
Wasserstein-$2$ metric. The Wasserstein distance is defined as

$\displaystyle
W_{2}^{2}(p,q)=\inf_{\pi\in\Pi(p,q)}\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}}\left\|x-y\right\|^{2}\pi(\mathrm{d}x\times\mathrm{d}y),$

where $\Pi(p,q)$ is the set of all couplings between $p$ and $q$.

In Ambrosio et al.,  (2005, Chapter 11), the authors discuss various
approaches for adapting gradient flows on well studied spaces (such as
Euclidean and Riemannian spaces) to the Wasserstein space. One of these
approaches proceeds by first defining suitable notions of tangent space and
subdifferential, following which the simple definition of gradient flow
modelled on Riemannian manifolds can then reproduced. In this case, the
Fréchet subdifferential (Ambrosio et al., , 2005, Definition 10.1.1) is
defined as follow:

######  Definition C.1.

[Fréchet differential on Wasserstein Space] Let
$F:\mathcal{P(X)}\rightarrow\mathbb{R}$ be a sufficiently regular function. We
say that $\xi\in L^{2}(q)$ belongs to the Fréchet subdifferential $\partial
F(q)$ if for all $q^{\prime}\in\mathcal{P(X)}$, we have

$\displaystyle
F(q^{\prime})-F(q)\geq\left\langle{\xi},{t_{q}^{q^{\prime}}-i}\right\rangle_{q}+o(W_{2}(q,q^{\prime})),$

where $t_{p}^{q}$ is the optimal map between $p$ and $q$ (Ambrosio et al., ,
2005, see (7.1.4)) and $i$ is the identity map. Furthermore, if
$\xi\in\partial F(q)$ also satisfies

$\displaystyle
F(t_{\\#}q)-F(q)\geq\left\langle{\xi},{t-i}\right\rangle_{q}+o(\|t-i\|_{q}),$

for all $t\in L^{2}(q)$, then we say that $\xi$ is a strong subdifferential.

See Ambrosio et al.,  (2005, Definition 10.1.1) for more details. The strong
subdifferential can be thought of as the (Wasserstein) “gradient” of $F$.
Equipped with this notion of gradient, we can define the gradient (descent)
flow of $F:\mathcal{P(X)}\rightarrow\mathbb{R}$ as follows:

######  Definition C.2 (Gradient Flow).

We say that a curve $p_{t}:[0,1]\rightarrow\mathcal{P(X)}$ is a gradient flow
of $F:\mathcal{P(X)}\rightarrow\mathbb{R}$ if for all $t>0$, it satisfies the
continuity equation $\partial_{t}p_{t}+\nabla_{x}\cdot(v_{t}p_{t})=0$, where
the tangent vector $v_{t}$ satisfies $-\,v_{t}\in\partial F(p_{t})$ for all
$t$.

Thus, for our application, we are interested in computing the strong
subdifferential of $F$. If $F$ is an integral of the type

$F(p)=\int f(x,p(x),\nabla_{x}p(x))\mathrm{d}x,$ (6)

where $f:\mathcal{X}\times\mathbb{R}\times\mathcal{X}\to\mathbb{R}$ is
sufficiently regular, then we will see that its strong subdifferential admits
a analytic solution.

Functionals of this form of great interest in the calculus of variations
(e.g., see Gelfand and Silverman,  (2000)). A vital quantity which is used to
study these functionals is the first variation. Writing
$\delta_{p}F[p]:\mathcal{X}\rightarrow\mathbb{R}$ for the first variation of
$F$, the unique up to constants function such that

$\left.\frac{\mathrm{d}}{\mathrm{d}\epsilon}F(q_{\epsilon})\right|_{\epsilon=0}=-\int\delta_{q}F[q](x)\,\nabla_{x}\cdot(q(x)v(x))\mathrm{d}x$

for all $v$ such that $q_{\epsilon}:=(i+\epsilon
v)_{\\#}q\in\mathcal{P}(\mathcal{X})$ for sufficiently small $\epsilon$, one
can readily establish that for (6)-typed $F$, it is given by

$\displaystyle\delta_{p}F[p](x):=\nabla_{(2)}f(x,p(x),\nabla_{x}p(x))-\nabla_{x}\cdot(\nabla_{(3)}f(x,p(x),\nabla_{x}p(x)),$

where $\nabla_{(i)}$ denotes the partial derivative w.r.t. the $i$-th argument
(Ambrosio et al., , 2005, (10.4.2)). It can be shown that for any
$\xi\in\partial F(p)$ which is a strong subdifferential, it holds that
$\xi(x)\overset{p-a.e.}{=}\nabla_{x}\delta_{p}F[p](x)$ (Ambrosio et al., ,
2005, Lemma 10.4.1).

###  C.2 Gradient flow on
$\mathbb{R}^{d_{\theta}}\times\mathcal{P}(\mathbb{R}^{d_{x}})$

We provided here a generalization of the Fréchet differential (for a broad
survey of Fréchet differentials, see Kruger,  (2003)) to the extended space
$\mathbb{R}^{d_{\theta}}\times\mathcal{P}(\mathbb{R}^{d_{x}})$:

######  Definition C.3 (Fréchet differential on
$\mathbb{R}^{d_{\theta}}\times\mathcal{P}(\mathbb{R}^{d_{x}})$).

Let
$\mathcal{F}:\mathbb{R}^{d_{\theta}}\times\mathcal{P(X)}\rightarrow\mathbb{R}$
be a sufficiently regular function. We say that
$(\xi_{\theta},\xi_{q})\in\mathbb{R}^{d_{\theta}}\times L_{2}(q)$ belongs to
the Fréchet subdifferential $\partial\mathcal{F}(\theta,q)$ if for all
$(\theta^{\prime},q^{\prime})\in\mathbb{R}^{d_{\theta}}\times\mathcal{P}(\mathbb{R}^{d_{x}})$,
it holds that

$\displaystyle\mathcal{F}(\theta^{\prime},q^{\prime})-\mathcal{F}(\theta,q)\geq\left\langle{\xi_{\theta}},{\theta^{\prime}-\theta}\right\rangle+\left\langle{\xi_{q}},{t_{q}^{q^{\prime}}-i}\right\rangle_{q}+o(W_{2}(q^{\prime},q)+\|\theta^{\prime}-\theta\|).$

Furthermore, if $(\xi_{\theta},\xi_{q})\in\partial\mathcal{F}(\theta,q)$ also
satisfies

$\displaystyle\mathcal{F}(\theta+\tau,t_{\\#}q)-\mathcal{F}(\theta,q)\geq\left\langle{\xi_{\theta}},{\tau}\right\rangle+\left\langle{\xi_{q}},{t-i}\right\rangle_{q}+o(\|t-i\|_{q}+\|\tau\|),$

for all $(\tau,t)\in\mathbb{R}^{d_{\theta}}\times L_{2}(q)$, then we say that
$(\xi_{\theta},\xi_{q})$ is a strong subdifferential.

For the remainder of the section, we shall assume that for any perturbation
$(\sigma,v)$, there exists some
$\delta_{q}F[\theta,q]:\mathcal{X}\rightarrow\mathbb{R}$ (often called the
first variation) such that the following expansion holds:

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}\epsilon}\mathcal{F}(\theta_{\epsilon},q_{\epsilon})\big{|}_{\epsilon=0}=$
$\displaystyle\left\langle{\nabla_{\theta}\mathcal{F}(\theta,q)},{\sigma}\right\rangle-\int\delta_{q}\mathcal{F}[\theta,q](x)\,\nabla_{x}\cdot(v(x)q(x))\,\mathrm{d}x,$
(7)

where $q_{\epsilon}:=(i+\epsilon v)_{\\#}q$ and
$\theta_{\epsilon}:=\theta+\epsilon\sigma$. For all $\mathcal{F}$ of interest
in this paper, this assumption holds; see for instance, see Proposition F.4.

We follow in the argument of Ambrosio et al.,  (2005, Lemma 10.4.1) to show
that if $(\xi_{\theta},\xi_{q})\in\partial\mathcal{F}(\theta,q)$ is a strong
subdifferential, then we have that
$\xi_{\theta}=\nabla_{\theta}\mathcal{F}(\theta,q)$ and
$\xi_{q}(x)\overset{q-a.e.}{=}\nabla_{x}\delta_{q}\mathcal{F}[\theta,q](x)$.
Let $(\xi_{\theta},\xi_{p})\in\partial\mathcal{F}(\theta,p)$ be a strong
subdifferential. By using the strong subdifferential property for some
perturbation $(\epsilon\sigma,\epsilon v)$ and taking left and right limits of
(C.3), we obtain that

$\displaystyle\limsup_{\epsilon\uparrow
0}\frac{\mathcal{F}(\theta_{\epsilon},q_{\epsilon})-\mathcal{F}(\theta,q)}{\epsilon}\leq\left\langle{\xi_{\theta}},{\sigma}\right\rangle+\left\langle{\xi_{q}},{v}\right\rangle_{q}\leq\liminf_{\epsilon\downarrow
0}\frac{\mathcal{F}(\theta_{\epsilon},q_{\epsilon})-\mathcal{F}(\theta,q)}{\epsilon}.$

On the other hand, from our assumption (7), we have that

$\displaystyle\lim_{\epsilon\rightarrow
0}\frac{\mathcal{F}(\theta_{\epsilon},q_{\epsilon})-\mathcal{F}(\theta,q)}{\epsilon}$
$\displaystyle=\left\langle{\nabla_{\theta}\mathcal{F}(\theta,q)},{\sigma}\right\rangle-\int\delta_{q}\mathcal{F}[\theta,q](x)\,\nabla_{x}\cdot(v(x)q(x))\,\mathrm{d}x$
$\displaystyle=\left\langle{\nabla_{\theta}\mathcal{F}(\theta,q)},{\sigma}\right\rangle+\left\langle{\nabla_{x}\delta_{q}\mathcal{F}[\theta,q]},{v}\right\rangle_{q},$

where the last inequality follows from integration by parts and the divergence
theorem. Hence, we have that
$\xi_{\theta}=\nabla_{\theta}\mathcal{F}(\theta,q)$, and
$\xi_{q}\overset{q-a.e.}{=}\nabla_{x}\delta_{q}\mathcal{F}[\theta,q]$.

###  C.3 First Variation

In this section, we derive the first variation for integral expressions taking
a certain form. In particular, we are interested in computing the strong
subdifferential of $\mathcal{F}$ of the following types:

$\displaystyle\mathcal{F}(\theta,q)$
$\displaystyle:=\int{f}(\theta,x,q(x),\nabla_{x}q(x))\mathrm{d}x,$ (VI-I)
$\displaystyle\mathcal{F}(\theta,q)$ $\displaystyle:=\int\int
f(\theta,x,y,q(x),q(y))\mathrm{d}x\mathrm{d}y.$ (VI-II)

An exampled of (VI-I)-typed $\mathcal{F}$ is when $\mathcal{F}$ is the free
energy $\mathcal{E}$. Following standard techniques from the calculus of
variations (Gelfand and Silverman, , 2000), consider a perturbation
$(\sigma,v)\in\mathbb{R}^{d_{\theta}}\times L_{2}(q)$ and define a mapping
$\Phi$ by

$\displaystyle\Phi(\epsilon):=\mathcal{F}(\theta_{\epsilon},q_{\epsilon})$

where $q_{\epsilon}:=(i+\epsilon v)_{\\#}q$ and
$\theta_{\epsilon}:=\theta+\epsilon\sigma$.

(VI-I)-typed $\mathcal{F}$. By application of the change-of-variables formula
(see for instance, see Ambrosio et al.,  (2005, Lemma 5.5.3)), we can compute
the density $q_{\epsilon}$ and its derivative
$\frac{\partial}{\partial\epsilon}q_{\epsilon}(y)$ as follows:

$\displaystyle q_{\epsilon}(y)$
$\displaystyle=\frac{q}{\mathrm{det}(I+\epsilon\nabla_{x}v)}\circ(i+\epsilon
v)^{-1}(y),\quad\frac{\partial}{\partial\epsilon}q_{\epsilon}(y)\big{|}_{\epsilon=0}=-\nabla_{x}\cdot[qv](y).$

Thus, we can compute the derivative of $\Phi$, provided that the interchange
of derivative and integral can be justified, as

$\displaystyle\frac{d}{d\epsilon}\Phi(\epsilon)\bigg{|}_{\epsilon=0}=$
$\displaystyle\int\frac{d}{d\epsilon}f(\theta_{\epsilon},x,q_{\epsilon}(x),\nabla_{x}q_{\epsilon}(x))\,\mathrm{d}x$
$\displaystyle=$
$\displaystyle\int\left\langle{\nabla_{\theta}f(\theta,x,q(x),\nabla_{x}q(x))},{\sigma}\right\rangle\,\mathrm{d}x$
$\displaystyle-\int\nabla_{(3)}f(\theta,x,q(x),\nabla_{x}q(x))\,\nabla_{x}\cdot[vq](x)\,\mathrm{d}x$
$\displaystyle-\int\left\langle{\nabla_{(4)}f(\theta,x,q(x),\nabla_{x}q(x))},{\nabla_{x}[\nabla_{x}\cdot[vq]](x)}\right\rangle\,\mathrm{d}x$

Applying integration by parts and the divergence theorem, we obtain that

$\displaystyle\frac{d}{d\epsilon}\Phi(\epsilon)\bigg{|}_{\epsilon=0}=$
$\displaystyle\left\langle{\nabla_{\theta}\mathcal{F}(\theta,x,q(x),\nabla_{x}q(x))},{\sigma}\right\rangle$
$\displaystyle-\int\nabla_{(3)}f(\theta,x,q(x),\nabla_{x}q(x))\,[\nabla_{x}\cdot(vq)](x)\,\mathrm{d}x$
$\displaystyle+\int\nabla_{x}\cdot[\nabla_{(4)}f](\theta,x,q(x),\nabla_{x}q(x))[\nabla_{x}\cdot(vq)](x)\,\mathrm{d}x$

We can hence write the first variation of (VI-I)-typed $\mathcal{F}$ as

$\delta_{q}\mathcal{F}[\theta,q](x)=\nabla_{(3)}f(\theta,x,q(x),\nabla_{x}q(x))-\nabla_{x}\cdot(\nabla_{(4)}f(\theta,x,q(x),\nabla_{x}q(x)).$
(8)

(VI-II)-typed $\mathcal{F}$. Similarly to above, we can define
$\Phi(\epsilon):=\mathcal{F}(\theta_{\epsilon},q_{\epsilon})$, whose
derivative is given by

$\displaystyle\frac{d}{d\epsilon}\Phi(\epsilon)\bigg{|}_{\epsilon=0}=$
$\displaystyle\left\langle{\nabla_{\theta}\mathcal{F}(\theta,q)},{\sigma}\right\rangle$
$\displaystyle-\int\int\nabla_{(4)}f(\theta,x,y,q(x),q(y))\,\nabla_{x}\cdot[qv](x)\,\mathrm{d}x\,\mathrm{d}y$
$\displaystyle-\int\int\nabla_{(5)}f(\theta,x,y,q(x),q(y))\,\nabla_{y}\cdot[qv](y)\,\mathrm{d}x\,\mathrm{d}y.$

Hence, the first variation is given by

$\displaystyle\delta_{q}\mathcal{F}[\theta,q](x)=$
$\displaystyle\int\left[\nabla_{(4)}f(\theta,x,y,q(x),q(y))+\nabla_{(5)}f(\theta,y,x,q(y),q(x))\right]\mathrm{d}y.$
(9)

######  Example 1 (First Variation of $\mathcal{E}$).

It can be seen that the free energy is (VI-I)-typed with
$f(\theta,x,a,g)=a\cdot\log\frac{a}{p_{\theta}(y,x)}$. Hence, by combining the
formula (8) with the fact that

$\displaystyle\nabla_{3}f(\theta,x,a,b)$
$\displaystyle=\log\frac{a}{p_{\theta}}+1$
$\displaystyle\nabla_{4}(f(\theta,x,a,g)$ $\displaystyle=0,$

we obtain the following expression for the first variation:

$\displaystyle\delta_{q}\mathcal{F}[\theta,q](x)=\log\frac{q(x)}{p_{\theta}(y,x)}+1.$

This provides an alternative derivation of this expression to that given in
Kuntz et al.,  (2023, Lemma 1).

##  Appendix D A Log Sobolev Inequality for $\mathcal{E}$ can be transferred
to $\mathcal{F}$

In this section, we show that nice properties of the functional $\mathcal{E}$
transfer to the functional $\mathcal{F}$, i.e. if $\mathcal{E}$ satisfies a
log Sobolev inequality (1), then so does $\mathcal{F}$ (with a modified
constant):

######  Proposition D.1.

If 1 holds, then upon defining
$C:=\min\\{C_{\mathcal{E}},\eta_{\theta},\eta_{x}\\}$, it holds that

$\mathcal{F}(\theta,m,q)+\log
p_{\theta^{*}}\left(y\right)\leq\frac{1}{2C}\left|\left|{\nabla\mathcal{F}(\theta,m,q)}\right|\right|_{q}^{2},\quad\forall\theta,m\in\mathbb{R}^{d_{\theta}},\enskip
q\in\mathcal{P}(\mathbb{R}^{2d_{x}}),$

where

$\left|\left|{\nabla\mathcal{F}(\theta,m,q)}\right|\right|_{q}^{2}:=\|\nabla_{(\theta,m)}\mathcal{F}(\theta,m,q)\|^{2}+\|\nabla_{(x,u)}\log
q-\nabla_{(x,u)}\log\rho_{\theta,\eta_{x}}\|^{2}_{q}.$ (10)

###### Proof.

As we show at the end of the proof, we have the fact that

$\left|\left|{\nabla_{(x,u)}\log\left(\frac{q}{\rho_{\theta}\otimes
r_{\eta_{x}}}\right)}\right|\right|_{q}^{2}\geq{\left|\left|{\nabla_{x}\log\left(\frac{q_{X}}{\rho_{\theta}}\right)}\right|\right|^{2}}_{q(\mathrm{d}x)}+{\left|\left|{\nabla_{u}\log\frac{q_{U|X}}{r_{\eta_{x}}}}\right|\right|^{2}}_{q}.$
(11)

From the definition, we have

$\left|\left|{\nabla\mathcal{F}(\theta,m,q)}\right|\right|_{q}^{2}=\left|\left|{\nabla_{\theta}\mathcal{F}(\theta,m,q)}\right|\right|^{2}+\left|\left|{\eta_{\theta}m}\right|\right|^{2}+\left|\left|{\nabla_{(x,u)}\log\frac{q}{\rho_{\theta,\eta_{x}}}}\right|\right|_{q}^{2}.$

Using (11) and the fact that

$\nabla_{\theta}\mathcal{F}(\theta,m,q)=-\int\nabla_{\theta}\log
p_{\theta}(y,x,u)q(\mathrm{d}x,\mathrm{d}u)=-\int\nabla_{\theta}\log
p_{\theta}(y,x)q(\mathrm{d}x)=\nabla_{\theta}\mathcal{E}(\theta,q),$

then it follows that

$\displaystyle\left|\left|{\nabla\mathcal{F}(\theta,m,q)}\right|\right|_{q}^{2}\geq$
$\displaystyle\left|\left|{\nabla_{\theta}\mathcal{E}(\theta,q_{X})}\right|\right|_{q}^{2}+\eta_{\theta}^{2}\left|\left|{m}\right|\right|^{2}+{\left|\left|{\nabla_{u}\log\frac{q_{U|X}}{r_{\eta}}}\right|\right|^{2}}_{q}.$

Because $C=\min\\{C_{\mathcal{E}},\eta_{\theta},\eta_{x}\\}$, the above
implies that

$\displaystyle\frac{1}{2C}\left|\left|{\nabla\mathcal{F}(\theta,m,q)}\right|\right|_{q}^{2}$
$\displaystyle\geq\frac{1}{2C_{\mathcal{E}}}\left|\left|{\nabla_{\theta}\mathcal{E}(\theta,q)}\right|\right|_{q}^{2}+\frac{\eta_{\theta}}{2}\left|\left|{m}\right|\right|^{2}+\frac{1}{2\eta_{x}}{\left|\left|{\nabla_{u}\log\frac{q_{U|X}}{r_{\eta_{x}}}}\right|\right|^{2}_{q}}.$
(12)

Since $r_{\eta_{x}}=\mathcal{N}(0,\eta_{x}^{-1}I_{d_{x}})$ is
$\eta_{x}$-strongly log-concave, by the Bakry-Émery criterion of Bakry and
Émery,  (2006), it holds that

$\frac{1}{2\eta_{x}}{\left|\left|{\nabla_{u}\log\left(\frac{q(\cdot|x)}{r_{\eta_{x}}}\right)}\right|\right|^{2}_{q(\mathrm{d}u|x)}}\geq\mathsf{KL}(q_{U|x}|r_{\eta_{x}})\quad\forall
x\in\mathcal{X},$

and so we obtain that

$\displaystyle\frac{1}{2\eta_{x}}{\left|\left|{\nabla_{u}\log\frac{q_{U|X}}{r_{\eta_{x}}}}\right|\right|^{2}}_{q}$
$\displaystyle=\frac{1}{2\eta_{x}}\mathbb{E}_{q_{X}}\left[\left|\left|{\nabla_{u}\log\frac{q(\cdot|X)}{r_{\eta_{x}}}}\right|\right|^{2}_{q(du|\cdot)}\right]$
$\displaystyle\geq\mathbb{E}_{q(\mathrm{d}x)}\left[\mathsf{KL}(q_{U|x}|r_{\eta_{x}})\right].$

Plugging the above into (12) and using the log Sobolev inequality for
$\mathcal{E}$, we obtain that for $\mathcal{F}$:

$\displaystyle\frac{1}{2C}\left|\left|{\nabla\mathcal{F}(\theta,m,q)}\right|\right|_{q}^{2}$
$\displaystyle\geq\mathcal{E}(\theta,q_{X})+\log
p_{\theta^{*}}(y)+\frac{\eta_{\theta}}{2}\|m\|_{2}^{2}+\mathbb{E}_{q(\mathrm{d}x)}\left[\mathsf{KL}(q_{U|X}|r_{\eta_{x}})\right]$
$\displaystyle\geq\mathcal{F}(\theta,m,q)+\log p_{\theta^{*}}\left(y\right).$

We have one loose end to tie up: proving (11). To do so, note that

$\displaystyle\left|\left|{\nabla_{(x,u)}\log\left(\frac{q}{\rho_{\theta}\otimes
r_{\eta_{x}}}\right)}\right|\right|_{q}^{2}=$
$\displaystyle{\left|\left|{\nabla_{(x,u)}\log\left(\frac{q_{X}}{\rho_{\theta}}\right)+\nabla_{(x,u)}\log\left(\frac{q_{U|X}}{r_{\eta_{x}}}\right)}\right|\right|^{2}_{q}}$
$\displaystyle=$
$\displaystyle{\left|\left|{\nabla_{x}\log\left(\frac{q_{X}}{\rho_{\theta}}\right)+\nabla_{x}\log
q_{U|X}}\right|\right|^{2}_{q}}$ (13)
$\displaystyle+{\left|\left|{\nabla_{u}\log\left(\frac{q_{U|X}}{r_{\eta_{x}}}\right)}\right|\right|^{2}_{q}}.$
(14)

Taking the first term and expanding the square, we have

$\displaystyle\left|\left|{\nabla_{x}\log\left(\frac{q_{X}}{\rho_{\theta}}\right)+\nabla_{x}\log
q_{U|X}}\right|\right|^{2}_{q}=$
$\displaystyle{\left|\left|{\nabla_{x}\log\left(\frac{q_{X}}{\rho_{\theta}}\right)}\right|\right|^{2}_{q}}$
$\displaystyle+2\left\langle{\nabla_{x}\log\left(\frac{q_{X}}{\rho_{\theta}}\right)},{\nabla_{x}\log
q_{U|X}}\right\rangle_{q}$ $\displaystyle+{\left|\left|{\nabla_{x}\log
q_{U|X}}\right|\right|^{2}_{q}}.$

Applying Jensen’s inequality to the final term, we obtain

$\displaystyle\left|\left|{\nabla_{x}\log\left(\frac{q_{X}}{\rho_{\theta}}\right)+\nabla_{x}\log
q_{U|X}}\right|\right|^{2}_{q}\geq$
$\displaystyle{\left|\left|{\nabla_{x}\log\left(\frac{q_{X}}{\rho_{\theta}}\right)}\right|\right|^{2}}_{q(\mathrm{d}x)}$
$\displaystyle+2{\left\langle{\nabla_{x}\log\left(\frac{q_{X}}{\rho_{\theta}}\right)},{\mathbb{E}_{q(\mathrm{d}u|\cdot)}\left[\nabla_{x}\log
q(U|\cdot)\right]}\right\rangle_{q(\mathrm{d}x)}}$
$\displaystyle+{\left|\left|{\mathbb{E}_{q(\mathrm{d}u|\cdot)}\left[\nabla_{x}\log
q(U|\cdot)\right]}\right|\right|^{2}_{q(\mathrm{d}x)}}.$

One can compute explicitly that for all $x$, it holds that
$\mathbb{E}_{q(\mathrm{d}u|x)}\left[\nabla_{x}\log q(u|X)\right]=0$; we thus
obtain that

$\displaystyle\left|\left|{\nabla_{x}\log\frac{q_{X}}{\rho_{\theta}}+\nabla_{x}\log
q_{U|X}}\right|\right|^{2}_{q}\geq\left|\left|{\nabla_{x}\log\frac{q_{X}}{\rho_{\theta}}}\right|\right|^{2}_{q(\mathrm{d}x)},$

and so (11) follows from (14). ∎

##  Appendix E $\mathcal{F}(z_{t})$ is non-increasing

In the following proposition, we show that
$\mathcal{F}_{t}:=\mathcal{F}(z_{t})$ is non-increasing in time.

######  Proposition E.1.

For any $\gamma_{x}\geq 0$ and $\gamma_{\theta}\geq 0$, it holds that

$\displaystyle\dot{\mathcal{F}}_{t}=-\gamma_{\theta}\|\nabla_{m}\mathcal{F}_{t}\|^{2}-\gamma_{x}\|\nabla_{u}\delta_{q}\mathcal{F}_{t}\|^{2}_{q_{t}}\leq
0.$

###### Proof.

We begin by computing the time derivative

$\displaystyle\dot{\mathcal{F}}_{t}=$
$\displaystyle-\left\langle\nabla_{(\theta,m)}\mathcal{F}_{t},\begin{pmatrix}0&-I_{d_{\theta}}\\\
I_{d_{\theta}}&\gamma_{\theta}I_{d_{\theta}}\end{pmatrix}\nabla_{(\theta,m)}\mathcal{F}_{t}\right\rangle$
(15)
$\displaystyle-\left\langle\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t},\begin{pmatrix}0&-I_{d_{x}}\\\
I_{d_{x}}&\gamma_{x}I_{d_{x}}\end{pmatrix}\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t}\right\rangle_{q_{t}}.$
(16)

Decomposing the matrix $\begin{pmatrix}0&-I_{d_{\theta}}\\\
I_{d_{\theta}}&\gamma_{\theta}I_{d_{\theta}}\end{pmatrix}$ into symmetric and
skew-symmetric components (write $D$, $Q$ respectively), i.e.

$\displaystyle\begin{pmatrix}0&-I_{d_{\theta}}\\\
I_{d_{\theta}}&\gamma_{\theta}I_{d_{\theta}}\end{pmatrix}=\underbrace{\begin{pmatrix}0&-I_{d_{\theta}}\\\
I_{d_{\theta}}&0\end{pmatrix}}_{=:Q}+\underbrace{\begin{pmatrix}0&0\\\
0&\gamma_{\theta}I_{d_{\theta}}\end{pmatrix}}_{=:D},$

we can simplify the RHS of (15) to

$\displaystyle\text{RHS
}\eqref{eq:dFdt_t1}=-\left\langle\nabla_{(\theta,m)}\mathcal{F}_{t},D\nabla_{(\theta,m)}\mathcal{F}_{t}\right\rangle=-\gamma_{\theta}\|\nabla_{m}\mathcal{F}_{t}\|^{2}.$

Similarly, we can show that

$\displaystyle\eqref{eq:dFdt_t2}=-\gamma_{x}\|\nabla_{u}\delta_{q}\mathcal{F}_{t}\|^{2}_{q_{t}}.$

As such, for $\gamma_{x}\geq 0$, $\gamma_{\theta}\geq 0$, the claim follows. ∎

##  Appendix F Proof of Proposition 4.1

The outline of the proof of Proposition 4.1 is:

  * •

Step 1 (Section F.1): Explicitly computing an upper bound of the time
derivative of $\mathcal{F}$ as a quadratic form.

  * •

Step 2 (Section F.2): Under the conditions specified in (34), we show that
this time derivative is bounded above by another quadratic form that allows us
to apply the log-Sobolev inequality of Proposition D.1.

  * •

Step 3: Using log-Sobolev inequality, Proposition D.1, and Grönwall’s
inequality, we obtain the desired result.

The remainder of the sections are dedicated to supporting the proof of Steps 1
and 2. This is done by developing the technical tools and carrying out
explicit computations. The particular roles of the following sections are:

  1. 1.

Section F.4.1: Computing the time derivative of $\mathcal{L}$ for Step 1.

  2. 2.

Sections F.4.3, F.4.2 and F.4.4: Introducing the adjoint, commutator and
another operator, as well as their explicit forms.

  3. 3.

Section F.4.5: Using the operators and explicit forms in Sections F.4.3, F.4.2
and F.4.4, we can upper bound terms introduced in Section F.4.1 for Step 1. We
utilize this in Step 1.

  4. 4.

Section F.4.6. Bounding the cross terms (or interaction terms) between
$\theta$ and $x$ that arises in the time derivative.

  5. 5.

Section F.4.7. Establishing sufficient conditions for a matrix with a
particular form to be positive semi-definite given in Proposition F.18. This
is used in Step 2.

Recall that the Lyapunov function is given by:

$\mathcal{L}:=\mathcal{F}+\log
p_{\theta_{*}}(y)+\left\|\nabla_{\left(\theta,m\right)}\mathcal{F}\right\|_{T_{\left(\theta,m\right)}}^{2}+\left\|\nabla_{\left(x,u\right)}\delta_{q}\mathcal{F}\right\|_{T_{\left(x,u\right)}}^{2}$

where

$\displaystyle
T_{(\theta,m)}:=\frac{1}{K}\begin{pmatrix}\tau_{\theta}I_{d_{\theta}}&\frac{\tau_{\theta
m}}{2}I_{d_{\theta}}\\\ \frac{\tau_{\theta
m}}{2}I_{d_{\theta}}&\tau_{m}I_{d_{\theta}}\end{pmatrix},\quad
T_{(x,u)}:=\frac{1}{K}\begin{pmatrix}\tau_{x}I_{d_{x}}&\frac{\tau_{xu}}{2}I_{d_{x}}\\\
\frac{\tau_{xu}}{2}I_{d_{x}}&\tau_{u}I_{d_{x}}\end{pmatrix}$ (17)

###### Proof of Proposition 4.1.

The proof is completed in the following three steps:

Step 1. In Section F.1, we show that the time derivative of
$\mathcal{L}(z_{t})$ satisfies the following upper bound:

$\displaystyle\frac{d}{dt}\mathcal{L}_{t}\leq$
$\displaystyle-\left\langle{\nabla_{(\theta,m)}\mathcal{F}_{t}},{S_{(\theta,m)}\nabla_{(\theta,m)}\mathcal{F}_{t}}\right\rangle-\left\langle{\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t}},{S_{(x,u)}\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t}}\right\rangle_{q_{t}}.$

where $S_{(x,u)}$ and $S_{(\theta,m)}$ are suitable matrices defined in (30)
and (31), respectively.

Step 2. Then, in Section F.2, when Equation 34 holds for some rate
$\varphi>0$, we have that

$\displaystyle S_{(\theta,m)}$ $\displaystyle\succeq\varphi
C\left(T_{(\theta,m)}+\frac{1}{2C}I_{d_{\theta}}\right),$ $\displaystyle
S_{(x,u)}$ $\displaystyle\succeq\varphi
C\left(T_{(x,u)}+\frac{1}{2C}I_{d_{x}}\right),$

where $C:=\min\\{C_{\mathcal{E}},\eta_{\theta},\eta_{x}\\}$.

Step 3. By Step $1$ and Step $2$, we obtain

$\displaystyle\frac{d}{dt}\mathcal{L}_{t}\leq$ $\displaystyle-\varphi
C\left(\frac{1}{2C}\left|\left|{\nabla_{(\theta,m)}\mathcal{F}_{t}}\right|\right|^{2}+\left\|\nabla_{\left(\theta,m\right)}\mathcal{F}_{t}\right\|_{T_{\left(\theta,m\right)}}^{2}\right)$
$\displaystyle-\varphi
C\left(\frac{1}{2C}{\left|\left|{\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t}}\right|\right|^{2}_{q_{t}}}+\left\|\nabla_{\left(x,u\right)}\delta_{q}\mathcal{F}_{t}\right\|_{T_{\left(x,u\right)}}^{2}\right).$

Comparing the above to (10) and applying Proposition D.1,

$\displaystyle\frac{d}{dt}\mathcal{L}_{t}\leq-\varphi C\mathcal{L}_{t},$

and application of Gronwall’s inequality yields the final result

$\displaystyle\mathcal{F}_{t}+\log
p_{\theta_{*}}\leq\mathcal{L}_{0}\exp\left(-\frac{C_{\mathcal{E}}}{10}t\right).$

∎

###  F.1 Proof of Step 1

As shown in Proposition F.2, the time derivative of the Lyapunov function
$\mathcal{L}$ is given by

$\displaystyle\frac{d}{dt}\mathcal{L}_{t}=$
$\displaystyle-\left\langle{\nabla_{(\theta,m)}\mathcal{L}_{t}},{\begin{pmatrix}0&-I_{d_{\theta}}\\\
I_{d_{\theta}}&\gamma_{\theta}I_{d_{\theta}}\end{pmatrix}\nabla_{(\theta,m)}\mathcal{F}_{t}}\right\rangle$
(18)
$\displaystyle-\left\langle{\nabla_{(x,u)}\delta_{q}\mathcal{L}_{t}},{\begin{pmatrix}0&-I_{d_{x}}\\\
I_{d_{x}}&\gamma_{\theta}I_{d_{x}}\end{pmatrix}\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t}}\right\rangle_{q_{t}},$
(19)

where we abbreviate $\mathcal{L}_{t}:=\mathcal{L}(z_{t})$ and
$\mathcal{F}_{t}:=\mathcal{F}(z_{t})$.

We deal with the inner products in (18,19) one-by-one, starting with (18): we
first compute the gradient of $\mathcal{L}$ w.r.t. to $(\theta,m)$, finding
that

$\displaystyle\nabla_{\left(\theta,m\right)}\mathcal{L}=\nabla_{(\theta,m)}\mathcal{F}+2[\nabla_{(\theta,m)}^{2}\mathcal{F}]\,T_{(\theta,m)}\nabla_{(\theta,m)}\mathcal{F}+2\mathbb{E}_{q}\left[[\nabla_{(\theta,m)}\nabla_{(x,u)}\delta_{q}\mathcal{F}]\,T_{(x,u)}\nabla_{(x,u)}\delta_{q}\mathcal{F}\right].$

It then follows that

$\displaystyle\eqref{eq:dldthetadt}$
$\displaystyle=-\left\langle{\nabla_{(\theta,m)}\mathcal{F}_{t}},{\tilde{\Gamma}_{\theta}\nabla_{(\theta,m)}\mathcal{F}_{t}}\right\rangle$
(20)
$\displaystyle-2\left\langle{\nabla_{(\theta,m)}\nabla_{(x,u)}[\delta_{q}\mathcal{F}_{t}]\,T_{(x,u)}\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t}},{\begin{pmatrix}0&-I_{d_{\theta}}\\\
I_{d_{\theta}}&\gamma_{\theta}I_{d_{\theta}}\end{pmatrix}\nabla_{(\theta,m)}\mathcal{F}_{t}}\right\rangle_{q_{t}}.$
(21)

where

$\displaystyle\tilde{\Gamma}_{\theta}:=\begin{pmatrix}0&-I_{d_{\theta}}\\\
I_{d_{\theta}}&\gamma_{\theta}I_{d_{\theta}}\end{pmatrix}+2T_{(\theta,m)}\nabla_{(\theta,m)}^{2}\mathcal{F}_{t}\begin{pmatrix}0&-I_{d_{\theta}}\\\
I_{d_{\theta}}&\gamma_{\theta}I_{d_{\theta}}\end{pmatrix}.$

By $\mathcal{F}$’s definition,

$\displaystyle\nabla^{2}_{(\theta,m)}\mathcal{F}=\begin{pmatrix}\nabla^{2}_{\theta}\mathcal{F}&0_{d_{\theta}}\\\
0_{d_{\theta}}&\eta_{\theta}I_{d_{\theta}}\end{pmatrix};$

whence we see that

$\displaystyle\nabla_{(\theta,m)}^{2}\mathcal{F}\begin{pmatrix}0_{d_{\theta}}&-I_{d_{\theta}}\\\
I_{d_{\theta}}&\gamma_{\theta}I_{d_{\theta}}\end{pmatrix}$
$\displaystyle=\begin{pmatrix}0_{d_{\theta}}&-\nabla_{\theta}^{2}\mathcal{F}\\\
\eta_{\theta}I_{d_{\theta}}&\gamma_{\theta}\eta_{\theta}I_{d_{\theta}}\end{pmatrix}$
$\displaystyle\Rightarrow
T_{(\theta,m)}\nabla_{(\theta,m)}^{2}\mathcal{F}\begin{pmatrix}0_{d_{\theta}}&-I_{d_{\theta}}\\\
I_{d_{\theta}}&\gamma_{\theta}I_{d_{\theta}}\end{pmatrix}$
$\displaystyle=\begin{pmatrix}\frac{\tau_{\theta
m}\eta_{\theta}}{2}I_{d_{\theta}}&\frac{\tau_{\theta
m}\gamma_{\theta}\eta_{\theta}}{2}I_{d_{\theta}}-\tau_{\theta}\nabla_{\theta}^{2}\mathcal{F}\\\
\tau_{m}\eta_{\theta}I_{d_{\theta}}&\tau_{m}\gamma_{\theta}\eta_{\theta}I_{d_{\theta}}-\frac{\tau_{\theta
m}}{2}\nabla_{\theta}^{2}\mathcal{F}\end{pmatrix},$
$\displaystyle\Rightarrow\tilde{\Gamma}_{\theta}$
$\displaystyle=\begin{pmatrix}\tau_{\theta
m}\eta_{\theta}I_{d_{\theta}}&[\tau_{\theta
m}\gamma_{\theta}\eta_{\theta}-1]I_{d_{\theta}}-2\tau_{\theta}\nabla_{\theta}^{2}\mathcal{F}_{t}\\\
[2\tau_{m}\eta_{\theta}+1]I_{d_{\theta}}&[2\tau_{m}\eta_{\theta}+1]\gamma_{\theta}I_{d_{\theta}}-\tau_{\theta
m}\nabla_{\theta}^{2}\mathcal{F}_{t}\end{pmatrix}.$

Given that, for any matrix $A$, we have
$\left\langle{\cdot},{A\cdot}\right\rangle=\left\langle{\cdot},{A^{sym}\cdot}\right\rangle$
where $A^{sym}=\left(A+A^{\top}\right)/2$, then with
$\Gamma_{\theta}:=\frac{\tilde{\Gamma}_{\theta}+\tilde{\Gamma}_{\theta}^{\top}}{2}$,
we obtain that

$\eqref{eq:dldthetadt_t1}=-\left\langle{\nabla_{(\theta,m)}\mathcal{F}_{t}},{\Gamma_{\theta}\nabla_{(\theta,m)}\mathcal{F}_{t}}\right\rangle,$
(22)

and hence that

$\Gamma_{\theta}=\begin{pmatrix}\tau_{\theta
m}\eta_{\theta}I_{d_{\theta}}&\left[\tau_{m}+\frac{\tau_{\theta
m}\gamma_{\theta}}{2}\right]\eta_{\theta}I_{d_{\theta}}-\tau_{\theta}\nabla_{\theta}^{2}\mathcal{F}_{t}\\\
\left[\tau_{m}+\frac{\tau_{\theta
m}\gamma_{\theta}}{2}\right]\eta_{\theta}I_{d_{\theta}}-\tau_{\theta}\nabla_{\theta}^{2}\mathcal{F}_{t}&[2\tau_{m}\eta_{\theta}+1]\gamma_{\theta}I_{d_{\theta}}-\tau_{\theta
m}\nabla_{\theta}^{2}\mathcal{F}_{t}\end{pmatrix}.$ (23)

We now turn our attention to (19). Defining
$\mathcal{H}_{t}(x,u):=\sqrt{\frac{q_{t}(x,u)}{\rho_{\theta_{t},\eta_{x}}(x,u)}}$
and Proposition F.2, we see that

$\displaystyle\text{\eqref{eq:dldqdt}}=$
$\displaystyle-4\left\langle{\nabla_{(x,u)}\log\mathcal{H}_{t}},{\begin{pmatrix}0_{d_{x}}&-I_{d_{x}}\\\
I_{d_{x}}&\gamma_{x}I_{d_{x}}\end{pmatrix}\nabla_{(x,u)}\log\mathcal{H}_{t}}\right\rangle_{q_{t}}$
(24)
$\displaystyle-8\left\langle{\nabla_{(x,u)}\frac{\nabla_{(x,u)}^{*}T_{(x,u)}\nabla_{(x,u)}\mathcal{H}_{t}}{\mathcal{H}_{t}}},{\begin{pmatrix}0_{d_{x}}&-I_{d_{x}}\\\
I_{d_{x}}&\gamma_{x}I_{d_{x}}\end{pmatrix}\nabla_{(x,u)}\log\mathcal{H}_{t}}\right\rangle_{q_{t}}$
(25)
$\displaystyle+4\left\langle{\nabla_{(x,u)}\nabla_{(\theta,m)}\left[\log{\rho_{\theta_{t}}}\right]T_{(\theta,m)}\nabla_{(\theta,m)}\mathcal{F}_{t}},{\begin{pmatrix}0_{d_{x}}&-I_{d_{x}}\\\
I_{d_{x}}&\gamma_{x}I_{d_{x}}\end{pmatrix}\nabla_{(x,u)}\log\mathcal{H}_{t}}\right\rangle_{q_{t}},$
(26)

Furthermore, we have

$\displaystyle\text{\eqref{eq:dldqdt_t1}}=-4\gamma_{x}\left|\left|{\nabla_{u}\log\mathcal{H}_{t}}\right|\right|^{2}_{q_{t}}=-\gamma_{x}\left|\left|{\nabla_{u}\log\frac{q_{t}}{\rho_{\theta_{t},\eta_{x}}}}\right|\right|^{2}_{q_{t}}.$

In Proposition F.11, we show that (25) can be further simplified as follows:

(25)
$\displaystyle=-2\gamma_{x}\left\langle\nabla_{(x,u)}\nabla_{u}\left[\log\frac{q_{t}}{\rho_{\theta_{t},\eta_{x}}}\right],T_{(x,u)}\nabla_{(x,u)}\nabla_{u}\left[\log\frac{q_{t}}{\rho_{\theta_{t},\eta_{x}}}\right]\right\rangle_{q_{t}}$
$\displaystyle-\left\langle\nabla_{(x,u)}\log\frac{q_{t}}{\rho_{\theta_{t},\eta_{x}}},\tilde{\Gamma}_{x}\nabla_{(x,u)}\log\frac{q_{t}}{\rho_{\theta_{t},\eta_{x}}}\right\rangle_{q_{t}}$

where

$\displaystyle\tilde{\Gamma}_{x}=\begin{pmatrix}\frac{\tau_{xu}\eta_{x}}{2}I_{d_{x}}&\frac{\tau_{u}+\tau_{xu}\gamma_{x}}{2}\eta_{x}I_{d_{x}}+\frac{\tau_{x}}{2}\nabla^{2}_{x}\ell\\\
\frac{\tau_{u}+\tau_{xu}\gamma_{x}}{2}\eta_{x}I+\frac{\tau_{x}}{2}\nabla^{2}_{x}\ell&2\gamma_{x}\tau_{u}\eta_{x}I_{d_{x}}+\frac{\tau_{xu}}{2}\nabla^{2}_{x}\ell\end{pmatrix}.$

We can then combine (24) and (25) to obtain

$\displaystyle\text{\eqref{eq:dldqdt_t1}}+\text{ \eqref{eq:dldqdt_t2}}=$
$\displaystyle-2\gamma_{x}\left\langle\nabla_{(x,u)}\nabla_{u}\left[\log\frac{q_{t}}{\rho_{\theta_{t},\eta_{x}}}\right],T_{(x,u)}\nabla_{(x,u)}\nabla_{u}\left[\log\frac{q_{t}}{\rho_{\theta_{t},\eta_{x}}}\right]\right\rangle_{q_{t}}$
$\displaystyle-\left\langle\nabla_{(x,u)}\log\frac{q_{t}}{\rho_{\theta_{t},\eta_{x}}},\Gamma_{x}\nabla_{(x,u)}\log\frac{q_{t}}{\rho_{\theta_{t},\eta_{x}}}\right\rangle_{q_{t}}.$

where $\Gamma_{x}=\tilde{\Gamma}_{x}+\begin{pmatrix}0_{d_{x}}&0_{d_{x}}\\\
0_{d_{x}}&\gamma_{x}I_{d_{x}}\end{pmatrix}$ is again positive semi-definite.

Given any p.s.d. matrix $M\in\mathbb{R}^{d\times d}$ and a general matrix
$A\in\mathbb{R}^{d\times d}$, it holds generically that

$\displaystyle\langle
A,MA\rangle=\left\langle{A},{LL^{\top}A}\right\rangle=\left\langle{L^{\top}A},{L^{\top}A}\right\rangle=\|L^{\top}A\|^{2}\geq
0,$

where $M=LL^{\top}$ is Cholesky decomposition of $M$. Thus, we have the
following upper bound:

$\displaystyle\text{\eqref{eq:dldqdt_t1}}+\text{ \eqref{eq:dldqdt_t2}}$
$\displaystyle\leq-\left\langle\nabla_{(x,u)}\log\frac{q_{t}}{\rho_{\theta_{t},\eta_{x}}},\Gamma_{x}\nabla_{(x,u)}\log\frac{q_{t}}{\rho_{\theta_{t},\eta_{x}}}\right\rangle_{q_{t}}$
(27)
$\displaystyle=-\left\langle\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t},\Gamma_{x}\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t}\right\rangle_{q_{t}}$
(28)

Cross Terms (21) and (26). We will now deal with the cross terms of (26) and
(21). In Proposition F.17, we show the following upper-bound: for all
$\varepsilon>0$, it holds that

$\displaystyle\text{\eqref{eq:dldqdt_t3}}+\text{\eqref{eq:dldthetadt_t2}}$
$\displaystyle\leq\varepsilon\left\langle\nabla_{(\theta,m)}\mathcal{F}_{t},\Gamma_{\times}\nabla_{(\theta,m)}\mathcal{F}_{t}\right\rangle+\frac{1}{\varepsilon}\left\langle\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t},\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t}\right\rangle_{q_{t}},$
(29)

where
$\Gamma_{\times}=K^{2}\begin{pmatrix}\tau_{\theta}^{2}I_{d_{\theta}}&\tau_{\theta}\left(\frac{\tau_{xu}+\tau_{\theta
m}}{2}\right)I_{d_{\theta}}\\\ \tau_{\theta}\left(\frac{\tau_{xu}+\tau_{\theta
m}}{2}\right)I_{d_{\theta}}&\left(\left(\frac{\tau_{xu}+\tau_{\theta
m}}{2}\right)^{2}+\tau_{x}^{2}\right)I_{d_{\theta}}\end{pmatrix}$.

Conclusion of Step 1. Combining the results in (22, 28, 29), we upper bound
the time derivative of the Lyapunov function as

$\displaystyle\frac{d}{dt}\mathcal{L}_{t}\leq$
$\displaystyle-\left\langle\nabla_{(\theta,m)}\mathcal{F}_{t},S_{(\theta,m)}\nabla_{(\theta,m)}\mathcal{F}_{t}\right\rangle-\left\langle\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t},S_{(x,u)}\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t}\right\rangle_{q_{t}},$

where

$\displaystyle S_{(x,u)}$
$\displaystyle=\Gamma_{x}-\frac{1}{\varepsilon}I_{d_{x}}=\begin{pmatrix}\left(\frac{\tau_{xu}\eta_{x}-1}{2}-\frac{1}{\varepsilon}\right)I_{d_{x}}&\frac{\tau_{u}+\tau_{xu}\gamma_{x}}{2}\eta_{x}I_{d_{x}}+\frac{\tau_{x}}{2}\nabla^{2}_{x}\ell\\\
\frac{\tau_{u}+\tau_{xu}\gamma_{x}}{2}\eta_{x}I_{d_{x}}+\frac{\tau_{x}}{2}\nabla^{2}_{x}\ell&\left(\gamma_{x}(2\tau_{u}\eta_{x}+1)-\frac{1}{\varepsilon}\right)I_{d_{x}}+\frac{\tau_{xu}}{2}\nabla^{2}_{x}\ell\end{pmatrix},$
(30) $\displaystyle S_{(\theta,m)}$
$\displaystyle=\Gamma_{\theta}-\varepsilon\Gamma_{\times}=\begin{pmatrix}s_{\theta}I_{d_{\theta}}&s_{\theta
m}I_{d_{\theta}}-\tau_{\theta}\nabla_{\theta}^{2}\mathcal{F}_{t}\\\ s_{\theta
m}I_{d_{\theta}}-\tau_{\theta}\nabla_{\theta}^{2}\mathcal{F}_{t}&s_{m}I_{d_{\theta}}-\tau_{\theta
m}\nabla_{\theta}^{2}\mathcal{F}_{t}\\\ \end{pmatrix},$ (31)

with constants defined as

$\displaystyle s_{\theta}$ $\displaystyle:=\tau_{\theta
m}\eta_{\theta}-\varepsilon K^{2}\tau_{\theta}^{2},\quad s_{\theta
m}:=\tau_{m}\eta_{\theta}+\frac{\tau_{\theta
m}\gamma_{\theta}}{2}\eta_{\theta}-\varepsilon
K^{2}\tau_{\theta}\frac{\tau_{xu}+\tau_{\theta m}}{2},$ $\displaystyle s_{m}$
$\displaystyle:=(2\tau_{m}\eta_{\theta}+1)\gamma_{\theta}-\varepsilon
K^{2}\left(\left(\frac{\tau_{xu}+\tau_{\theta
m}}{2}\right)^{2}+\tau_{x}^{2}\right).$

###  F.2 Proof of Step 2

By the definitions of $T_{(\theta,m)}$ and $T_{(x,u)}$ in (17), and those of
$S_{(\theta,m)}$ and $S_{(x,u)}$ in (30, 31),

$\displaystyle S_{(\theta,m)}-\varphi
C\left(T_{(\theta,m)}+\frac{1}{2C}I_{d_{\theta}}\right)$
$\displaystyle=\begin{pmatrix}\alpha_{\theta}I_{d_{\theta}}&\beta_{\theta}I_{d_{\theta}}-{\tau_{\theta}}\nabla_{\theta}^{2}\mathcal{F}\\\
\beta_{\theta}I_{d_{\theta}}-{\tau_{\theta}}\nabla_{\theta}^{2}\mathcal{F}&\kappa_{\theta}I_{d_{\theta}}-\tau_{\theta
m}\nabla_{\theta}^{2}\mathcal{F}\end{pmatrix},$ (32) $\displaystyle
S_{(x,u)}-\varphi C\left(T_{(x,u)}+\frac{1}{2C}I_{d_{x}}\right)$
$\displaystyle=\begin{pmatrix}\alpha_{x}I_{d_{x}}&\beta_{x}I_{d_{x}}+\frac{\tau_{x}}{2}\nabla_{x}^{2}\ell\\\
\beta_{x}I_{d_{x}}+\frac{\tau_{x}}{2}\nabla_{x}^{2}\ell&\kappa_{x}I_{d_{x}}+\frac{\tau_{xu}}{2}\nabla_{x}^{2}\ell\end{pmatrix}.$
(33)

where

$\displaystyle\alpha_{\theta}$ $\displaystyle:=\tau_{\theta
m}\eta_{\theta}-\varepsilon K^{2}\tau_{\theta}^{2}-\varphi
C\left(\tau_{\theta}+\frac{1}{2C}\right),$ $\displaystyle\beta_{\theta}$
$\displaystyle:=\left(\tau_{m}+\frac{\tau_{\theta
m}\gamma_{\theta}}{2}\right)\eta_{\theta}-\varepsilon
K^{2}\tau_{\theta}\left(\frac{\tau_{xu}+\tau_{\theta m}}{2}\right)-\varphi
C\frac{\tau_{\theta m}}{2},$ $\displaystyle\kappa_{\theta}$
$\displaystyle:=(2\tau_{m}\eta_{\theta}+1)\gamma_{\theta}-\varepsilon
K^{2}\left(\left(\frac{\tau_{xu}+\tau_{\theta
m}}{2}\right)^{2}+\tau_{x}^{2}\right)-\varphi
C\left(\tau_{m}+\frac{1}{2C}\right),$ $\displaystyle\alpha_{x}$
$\displaystyle=\frac{\tau_{xu}\eta_{x}}{2}-\frac{1}{\varepsilon}-\varphi
C\left(\tau_{x}+\frac{1}{2C}\right),$ $\displaystyle\beta_{x}$
$\displaystyle=\frac{(\tau_{u}+\tau_{xu}\gamma_{x})\eta_{x}}{2}-\varphi
C\frac{\tau_{xu}}{2},$ $\displaystyle\kappa_{x}$
$\displaystyle=\gamma_{x}(2\tau_{u}\eta_{x}+1)-\frac{1}{\varepsilon}-\varphi
C\left(\tau_{u}+\frac{1}{2C}\right).$

######  Proposition F.1.

Defining
$\left(\alpha_{\theta},\beta_{\theta},\kappa_{\theta},\alpha_{x},\beta_{x},\kappa_{x}\right)$
as above, assume that $K\geq 2C_{\mathcal{E}}$ if

$\displaystyle\tau_{\theta m}K-\alpha_{\theta}-\kappa_{\theta}$
$\displaystyle\leq 0,$ (34a)
$\displaystyle\tau_{\theta}^{2}K^{2}+(\tau_{\theta
m}\alpha_{\theta}-2\beta_{\theta}\tau_{\theta})K-\alpha_{\theta}\kappa_{\theta}+\beta^{2}_{\theta}$
$\displaystyle\leq 0,$ (34b)
$\displaystyle\tau_{\theta}^{2}K^{2}-(\tau_{\theta
m}\alpha_{\theta}-2\beta_{\theta}\tau_{\theta})K-\alpha_{\theta}\kappa_{\theta}+\beta^{2}_{\theta}$
$\displaystyle\leq 0,$ (34c)
$\displaystyle\frac{\tau_{xu}}{2}K-\alpha_{x}-\kappa_{x}$ $\displaystyle\leq
0,$ (34d)
$\displaystyle\frac{\tau_{x}^{2}}{4}K^{2}-\left(\frac{\tau_{xu}}{2}\alpha_{x}-\beta_{x}{\tau_{x}}\right)K-\alpha_{x}\kappa_{x}+\beta^{2}_{x}$
$\displaystyle\leq 0,$ (34e)
$\displaystyle\frac{\tau_{x}^{2}}{4}K^{2}+\left(\frac{\tau_{xu}}{2}\alpha_{x}-\beta_{x}{\tau_{x}}\right)K-\alpha_{x}\kappa_{x}+\beta^{2}_{x}$
$\displaystyle\leq 0,$ (34f)

hold for some rate $\varphi>0$. The following lower bounds for $S_{(x,u)}$ and
$S_{(\theta,m)}$ then hold:

$S_{(x,u)}\succeq\varphi C\left(T_{(x,u)}+\frac{1}{2C}I_{d_{x}}\right),\quad
S_{(\theta,m)}\succeq\varphi
C\left(T_{(\theta,m)}+\frac{1}{2C}I_{d_{\theta}}\right),$

where $C:=\min\\{C_{\mathcal{E}},\eta_{\theta},\eta_{x}\\}$.

###### Proof.

Note the matrices in (32, 33) have the same form as that in (47). As we show
below,

$-KI_{d_{x}}\preceq\nabla_{x}^{2}\ell\preceq KI_{d_{x}},\quad-
KI_{d_{\theta}}\preceq\nabla^{2}_{\theta}\mathcal{F}\preceq KI_{d_{\theta}}.$
(35)

Then, applying Proposition F.18 tells us that (32, 33) are positive
semidefinite if the conditions (34) are satisfied.

Now, to prove the inequalities in (35). 2 and (Nesterov, , 2003, Lemma 1.2.2)
imply the first two inequalities and

$-KI_{d_{\theta}}\preceq\nabla^{2}_{\theta}\ell\preceq KI_{d_{\theta}}.$

The other two inequalities in (35) follow from the above. To see this, for
$\nabla^{2}_{\theta}\mathcal{F}\preceq KI_{d_{\theta}}$, we have

$\left\langle{v},{(KI_{d_{\theta}}-\nabla^{2}_{\theta}\mathcal{F})v}\right\rangle=\int\left\langle{v},{(KI_{d_{\theta}}+\nabla^{2}_{\theta}\ell)v}\right\rangle
q(\mathrm{d}x\times\mathrm{d}u),$

and since
$\left\langle{v},{\nabla^{2}_{\theta}[\ell]v}\right\rangle\geq-K\|v\|^{2}$, we
have shown that
$\left\langle{v},{(KI_{d_{\theta}}-\nabla^{2}_{\theta}\mathcal{F})v}\right\rangle\geq
0$ implying the desired result of $\nabla^{2}_{\theta}\mathcal{F}\preceq
KI_{d_{\theta}}$. A similar argument can be made for
$-KI_{d_{\theta}}\preceq\nabla^{2}_{\theta}\mathcal{F}$. ∎

###  F.3 Examples for (34) holding

We verified the above with a symbolic calculator written in Mathematica for
the choices of $\eta_{x},\eta_{\theta},\gamma_{x},\gamma_{\theta}$ hold for
the following choices:

  1. 1.

Rate $\varphi=\frac{1}{10}$, momentum parameters $\gamma_{x}=\frac{3}{2}$,
$\gamma_{\theta}=3$, $\eta_{x}=2K$, $\eta_{\theta}=2K$, and elements of the
Lyapunov function $\tau_{\theta}=\frac{3}{8},\tau_{\theta
m}=\frac{5}{4},\tau_{m}=2,\tau_{x}=\frac{1}{4},\tau_{xu}=\frac{5}{4},\tau_{u}=\frac{7}{4}$,
and $\varepsilon=15$.

  2. 2.

Rate $\varphi=\frac{1}{30}$, momentum parameters $\gamma_{x}=\frac{6}{5}$,
$\gamma_{\theta}=\frac{5}{2}$, $\eta_{x}=2K$, $\eta_{\theta}=2K$, and elements
of the Lyapunov function $\tau_{\theta}=\frac{3}{8},\tau_{\theta
m}=\frac{5}{4},\tau_{m}=2,\tau_{x}=\frac{1}{4},\tau_{xu}=1,\tau_{u}=\frac{3}{2}$,
and $\varepsilon=15$.

###  F.4 Supporting Proofs and Derivations

####  F.4.1 The derivative of $\mathcal{L}_{t}$ along the flow

Here, we prove (18,19).

######  Proposition F.2.

The derivative of the Lyapunov function $\mathcal{L}$ is given by

$\displaystyle\frac{d}{dt}\mathcal{L}_{t}=$
$\displaystyle-\left\langle{\nabla_{(\theta,m)}\mathcal{L}_{t}},{\begin{pmatrix}0&-I_{d_{\theta}}\\\
I_{d_{\theta}}&\gamma_{\theta}I_{d_{\theta}}\end{pmatrix}\nabla_{(\theta,m)}\mathcal{F}_{t}}\right\rangle-\left\langle{\nabla_{(x,u)}\delta_{q}\mathcal{L}_{t}},{\begin{pmatrix}0&-I_{d_{x}}\\\
I_{d_{x}}&\gamma_{\theta}I_{d_{x}}\end{pmatrix}\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t}}\right\rangle_{q_{t}},$

where the first variation is given by

$\displaystyle\delta_{q}\mathcal{L}_{t}=$
$\displaystyle\delta_{q}\mathcal{F}_{t}+\delta_{q}\mathcal{L}^{1}_{t}+\delta_{q}\mathcal{L}^{2}_{t}$
$\displaystyle=$
$\displaystyle\log\frac{q(x,u)}{\rho_{\theta,\eta_{x}}(x,u)}+1-2\left\langle{\nabla_{(\theta,m)}\log{\rho_{\theta}(x,u)}},{T_{(\theta,m)}\nabla_{(\theta,m)}\mathcal{F}[z]}\right\rangle$
$\displaystyle+\frac{4}{\mathcal{H}}\nabla_{(x,u)}^{*}T_{(x,u)}\nabla_{(x,u)}\mathcal{H},$

with $\nabla^{*}$ and $\mathcal{H}$ defined in Proposition F.5.

###### Proof.

From the linearity of $\mathcal{L}$, we have

$\dot{\mathcal{L}}_{t}=\dot{\mathcal{F}}_{t}+\dot{\mathcal{L}^{1}_{t}}+\dot{\mathcal{L}^{2}_{t}}$

where

$\displaystyle\mathcal{L}_{t}^{1}$
$\displaystyle:=\left\langle{\nabla_{(\theta,m)}\mathcal{F}_{t}},{T_{(\theta,m)}\nabla_{(\theta,m)}\mathcal{F}_{t}}\right\rangle,$
$\displaystyle\mathcal{L}_{t}^{2}$
$\displaystyle:={\left\langle{\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t}},{T_{(x,u)}\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t}}\right\rangle}_{q_{t}}.$

From Section C.2, it can be seen that the time derivatives of
$\mathcal{F}_{t}$, $\mathcal{L}^{1}_{t}$, and $\mathcal{L}^{2}_{t}$, are given
by

$\displaystyle\dot{\mathcal{F}_{t}}$
$\displaystyle=\left\langle{\nabla_{(\theta,m)}\mathcal{F}_{t}},{{(\dot{\theta},\dot{m})}_{t}}\right\rangle+\left\langle{\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t}},{(\dot{x},\dot{u})_{t}}\right\rangle_{q_{t}}$
$\displaystyle\dot{\mathcal{L}^{1}_{t}}$
$\displaystyle=\left\langle{\nabla_{(\theta,m)}\mathcal{L}^{1}_{t}},{{(\dot{\theta},\dot{m})}_{t}}\right\rangle+\left\langle{\nabla_{(x,u)}\delta_{q}\mathcal{L}^{1}_{t}},{(\dot{x},\dot{u})_{t}}\right\rangle_{q_{t}}$
$\displaystyle\dot{\mathcal{L}^{2}_{t}}$
$\displaystyle=\left\langle{\nabla_{(\theta,m)}\mathcal{L}^{2}_{t}},{{(\dot{\theta},\dot{m})}_{t}}\right\rangle+\left\langle{\nabla_{(x,u)}\delta_{q}\mathcal{L}^{2}_{t}},{(\dot{x},\dot{u})_{t}}\right\rangle_{q_{t}}$

where ${(\dot{\theta},\dot{m})}_{t}=\begin{pmatrix}0&-I_{d_{\theta}}\\\
I_{d_{\theta}}&\gamma_{\theta}I_{d_{\theta}}\end{pmatrix}\nabla_{(\theta,m)}\mathcal{F}_{t}$
and $(\dot{x},\dot{u})_{t}=\begin{pmatrix}0&-I_{d_{x}}\\\
I_{d_{x}}&\gamma_{\theta}I_{d_{x}}\end{pmatrix}\nabla_{(x,u)}\delta_{q}\mathcal{F}_{t}$.
From Proposition F.3, Proposition F.4, and Proposition F.5, and linearity of
the inner product, we obtain as desired. ∎

######  Proposition F.3.

The first variation of $\mathcal{F}$ is given by

$\delta_{q}\mathcal{F}[z](x,u)=\log\frac{q(x,u)}{\rho_{\theta,\eta_{x}}(x,u)}+1,$

where $z=(\theta,m,q)$.

###### Proof.

It can be seen that $\mathcal{F}$ is (VI-I)-typed equation with
$f((\theta,m),(x,u),q,g)=q\log\frac{q}{\rho_{\theta,\eta_{x}}(x)}$. Thus,
using the formula (8), we obtain as its first variation

$\delta_{q}\mathcal{F}[z](x,u)=\log\frac{q(x,u)}{\rho_{\theta,\eta_{x}}(x,u)}+1.$

Hence, we have as desired. ∎

For the first variation of $\mathcal{L}_{1}$, we have the following result:

######  Proposition F.4.

The first variation of $\mathcal{L}_{1}$ is given by

$\delta_{q}\mathcal{L}_{1}[z](x,u)=-2\left\langle{\nabla_{(\theta,m)}\log{\rho_{\theta}(x,u)}},{T_{(\theta,m)}\nabla_{(\theta,m)}\mathcal{F}[z]}\right\rangle,$

where $z=(\theta,m,q)$.

###### Proof.

It can be seen that $\mathcal{L}_{1}$ is a variational integral of type (VI-
II) with

$\displaystyle f((\theta,m),(x,u),(x^{\prime},u^{\prime}),q,q^{\prime})=$
$\displaystyle
qq^{\prime}\left\langle{\nabla_{\theta}\log{\rho(x,u)}},{T_{\theta\theta}\nabla_{\theta}\log{\rho(x^{\prime},u^{\prime})}}\right\rangle$
$\displaystyle-q\left\langle{\nabla_{\theta}\log{\rho(x,u)}},{T_{\theta
m}m}\right\rangle$
$\displaystyle-q^{\prime}\left\langle{m},{T_{m\theta}\nabla_{\theta}\log{\rho(x^{\prime},u^{\prime})}}\right\rangle$
$\displaystyle+\left\langle{m},{T_{mm}m}\right\rangle,$

where we write $T_{\theta\theta}$ for the upper left block of
$T_{(\theta,m)}$, and similarly for $T_{\theta m}$, and $T_{mm}$.

Using the formula (9) and the fact that

$\displaystyle\nabla_{(4)}f((\theta,m),(x,u),(x^{\prime},u^{\prime}),q,q^{\prime})=$
$\displaystyle
q^{\prime}\left\langle{\nabla_{\theta}\log{\rho(x,u)}},{T_{\theta\theta}\nabla_{\theta}\log{\rho(x^{\prime},u^{\prime})}}\right\rangle$
$\displaystyle-\left\langle{\nabla_{\theta}\log{\rho(x,u)}},{T_{\theta
m}m}\right\rangle$

For the first term of (9), we have

$\displaystyle\int\nabla_{(4)}f((\theta,m),(x,u),(x^{\prime},u^{\prime}),q(x,u),q(x^{\prime},u^{\prime}))\mathrm{d}x^{\prime}\mathrm{d}u^{\prime}$
$\displaystyle=-\left\langle{\nabla_{(\theta,m)}\log{\rho(x,u)}},{T_{(\theta,m)}\nabla_{(\theta,m)}\mathcal{F}[z]}\right\rangle,$

and by symmetry, it follows similarly for the last term of (9). We have
obtained as desired. ∎

######  Proposition F.5.

The first variation of $\mathcal{L}_{2}$ is given by

$\delta_{q}\mathcal{L}_{2}[z]=\frac{4}{\mathcal{H}}\nabla_{(x,u)}^{*}T_{(x,u)}\nabla_{(x,u)}\mathcal{H},$

where $\mathcal{H}(x,u):=\sqrt{\frac{q(x,u)}{\rho_{\theta,\eta_{x}}(x,u)}}$
and the operator $\nabla_{(x,u)}^{*}$ is given as

$\nabla_{(x,u)}^{*}(f)(x,u):=-\left\langle{\nabla_{(x,u)}\log\rho_{\theta,\eta_{x}}(x,u)},{f(x,u)}\right\rangle-\nabla_{(x,u)}\cdot[f](x,u),$

where $\nabla_{x}\cdot$ is the divergence operator w.r.t. $x$.

###### Proof.

First note that $\mathcal{L}_{2}$ can be written equivalently as

$\left\langle{\frac{\nabla_{(x,u)}q}{q}-\nabla_{(x,u)}\log\rho_{\theta,\eta_{x}}},{T_{(x,u)}\left[\frac{\nabla_{(x,u)}q}{q}-\nabla_{(x,u)}\log\rho_{\theta}\right]}\right\rangle_{q}$

It can be seen that $\mathcal{L}_{2}$ is a variational integral of type (VI-I)
with

$f((\theta,m),(x,u),q,g)=q\left\langle{\frac{g}{q}-\nabla_{(x,u)}\log\rho_{\theta,\eta_{x}}(x,u)},{T_{(x,u)}\left[\frac{g}{q}-\nabla_{(x,u)}\log\rho_{\theta,\eta_{x}}(x,u)\right]}\right\rangle.$

This follows from since

$\displaystyle\nabla_{(3)}f((\theta,m),(x,u),q,g)=$
$\displaystyle\left\langle{\frac{g}{q}-\nabla_{(x,u)}\log\rho_{\theta,\eta_{x}}(x,u)},{T_{(x,u)}\left[\frac{g}{q}-\nabla_{(x,u)}\log\rho_{\theta,\eta_{x}}(x,u)\right]}\right\rangle$
$\displaystyle-2q\left\langle{\frac{g}{q^{2}}},{T_{(x,u)}\left[\frac{g}{q}-\nabla_{(x,u)}\log\rho_{\theta,\eta_{x}}(x,u)\right]}\right\rangle$
$\displaystyle=$
$\displaystyle-\left\langle{\frac{g}{q}+\nabla_{(x,u)}\log\rho_{\theta,\eta_{x}}(x,u)},{T_{(x,u)}\left[\frac{g}{q}-\nabla_{(x,u)}\log\rho_{\theta,\eta_{x}}(x,u)\right]}\right\rangle$

and

$\displaystyle\nabla_{(4)}f((\nabla,m),(x,u),q,g)=$ $\displaystyle
2T_{(x,u)}\left[\frac{g}{q}-\nabla_{(x,u)}\log\rho_{\theta,\eta_{x}}(x,u)\right]$

Using formula (8), we obtain

$\displaystyle\delta_{q}\mathcal{L}_{2}=$
$\displaystyle-\left\langle{\nabla_{(x,u)}\log{q(x,u)}+\nabla_{(x,u)}\log{\rho_{\theta,\eta_{x}}(x,u)}},{T_{(x,u)}\nabla_{(x,u)}\log\frac{q(x,u)}{\rho_{\theta,\eta_{x}}(x,u)}}\right\rangle$
(36a)
$\displaystyle-\nabla_{(x,u)}\cdot\left(2T_{(x,u)}\nabla_{(x,u)}\log\frac{q(x,u)}{\rho_{\theta,\eta_{x}}(x,u)}\right)$
(36b)

To obtain the desired result, note that we can rewrite

$\displaystyle\eqref{eq:l2_fv_t1}$
$\displaystyle=-\left\langle{\nabla_{(x,u)}\log\sqrt{q(x,u)}+\nabla_{(x,u)}\log\sqrt{\rho_{\theta,\eta_{x}}(x,u)}},{\frac{4}{\mathcal{H}}T_{(x,u)}\nabla_{(x,u)}\log\mathcal{H}}\right\rangle$
$\displaystyle\eqref{eq:l2_fv_t2}$
$\displaystyle=-\nabla_{(x,u)}\cdot\left(\frac{4}{\mathcal{H}}T_{(x,u)}\nabla_{(x,u)}\mathcal{H}\right),$
$\displaystyle=-\left\langle{\nabla_{(x,u)}\frac{4}{\mathcal{H}}},{T_{(x,u)}\nabla_{(x,u)}\mathcal{H}}\right\rangle-\frac{4}{\mathcal{H}}\nabla_{(x,u)}\cdot(T_{(x,u)}\nabla_{(x,u)}\mathcal{H})$
$\displaystyle=\left\langle{\nabla_{(x,u)}\log\mathcal{H}},{\frac{4}{\mathcal{H}}T_{(x,u)}\nabla_{(x,u)}\mathcal{H}}\right\rangle-\frac{4}{\mathcal{H}}\nabla_{(x,u)}\cdot(T_{(x,u)}\nabla_{(x,u)}\mathcal{H}).$

Hence, we have obtained the desired result. ∎

####  F.4.2 The Adjoint of $\nabla$

Throughout the proofs, we use $\nabla^{*}$ to denote the linear map that is an
adjoint of $\nabla$ (the Euclidean gradient operator). For this section, we
write $\rho:=\rho_{\theta,\eta_{x}}$.

In the view of $\nabla$ as a linear map from
$(L^{2}_{\rho}(\mathcal{X};\mathbb{R}),\left\langle{\cdot},{\cdot}\right\rangle_{\rho})$
to
$(L^{2}_{\rho}(\mathcal{X};\mathcal{X}),\left\langle{\cdot},{\cdot}\right\rangle_{\rho})$
where
$\mathcal{X}\in\\{\mathbb{R}^{d_{x}},\mathbb{R}^{d_{x}}\times\mathbb{R}^{d_{x}}\\}$
and

$\displaystyle L^{2}_{\rho}(\mathcal{X};\mathcal{X}):=\\{\text{sufficiently
regular }v:\mathcal{X}\to\mathcal{X}\text{ such that
}\left|\left|{v}\right|\right|_{\rho}<\infty\\}.$

Then its adjoint $\nabla^{*}$ is given by

$\nabla^{*}v:=-\frac{1}{\rho}\nabla\cdot[\rho
v]=-\left\langle{\nabla\log\rho},{v}\right\rangle-\nabla\cdot v.$

where $\nabla\cdot$ denotes the divergence operator. Whatever the gradient
$\nabla$ is taken with respect to, the adjoint $\nabla^{*}$ is taken with
respect to the corresponding quantity. In order to see the adjoint property,
we have

$\displaystyle\left\langle{\nabla f},{v}\right\rangle_{\rho}$
$\displaystyle=\int\left\langle{\nabla
f(x,u)},{v(x,u)}\right\rangle\rho(\mathrm{d}x\times du)$
$\displaystyle=\int\left\langle{\nabla
f(x,u)},{\rho(x,u)v(x,u)}\right\rangle\,\mathrm{d}x\times du$
$\displaystyle=-\int f(x,u)\nabla\cdot[\rho(x,u)v(x,u)]\,\mathrm{d}x\times
du=\left\langle{f},{\nabla^{*}v}\right\rangle_{\rho},$

where we used integration by parts combined with the divergence theorem and
the fact that $\rho$ must vanish at the boundary to obtain the second
equation.

Another case of interest is in the view of $\nabla$ as a linear map from
$(L^{2}_{\rho}(\mathcal{X};\mathcal{X}^{\prime}),\left\langle{\cdot},{\cdot}\right\rangle_{\rho})$
to
$(L^{2}_{\rho}(\mathcal{X};\mathcal{X}\times\mathcal{X}^{\prime}),\left\langle{\cdot},{\cdot}\right\rangle_{\rho})$,
where the underlying inner product on the latter space is induced by the
Frobenius inner product. In this case, the adjoint $\nabla^{*}$ is defined a
similar way: Let $M\in
L^{2}(\mathcal{X},\mathcal{X}\times\mathcal{X}^{\prime})$ and $v\in
L^{2}(\mathcal{X},\mathcal{X}^{\prime})$

$\displaystyle\left\langle{\nabla v},{M}\right\rangle_{\rho}$
$\displaystyle=\int\sum_{i=1}^{d_{\mathcal{X}}}\sum_{j=1}^{d_{\mathcal{X^{\prime}}}}(\nabla
v)_{ij}M_{ij}\,\rho(\mathrm{d}x\times\mathrm{d}u)$
$\displaystyle=\int\sum_{i=1}^{d_{\mathcal{X}}}\sum_{j=1}^{d_{\mathcal{X}^{\prime}}}\partial_{i}[v_{j}]\,M_{ij}\,\rho(\mathrm{d}x\times\mathrm{d}u)=\int\sum_{j=1}^{d_{\mathcal{X}^{\prime}}}\left\langle{\nabla
v_{j}},{Me_{j}}\right\rangle\,\rho(\mathrm{d}x\times\mathrm{d}u)$
$\displaystyle=-\int\sum_{j=1}^{d_{\mathcal{X}^{\prime}}}(v_{j})\nabla\cdot(\rho
Me_{j})\,\mathrm{d}x\times\mathrm{d}u$
$\displaystyle=-\int\sum_{j=1}^{d_{\mathcal{X}^{\prime}}}(v_{j})[\left\langle{\nabla\log\rho},{Me_{j}}\right\rangle+\nabla\cdot(Me_{j})]\,\rho(\mathrm{d}x\times\mathrm{d}u)$
$\displaystyle=\left\langle{\nabla^{*}M},{v}\right\rangle_{\rho}.$

where, for the second line, we apply integration by parts combined with the
divergence theorem and the fact that $\rho$ vanish at the boundary. Each
element of $(\nabla^{*}M)_{i}\in\mathcal{X^{\prime}}$ for
$i\in[{d_{\mathcal{X}^{\prime}}}]$, is defined to be

$\displaystyle(\nabla^{*}M)_{i}$
$\displaystyle:=-\left\langle{\nabla\log\rho},{M^{\top}e_{i}}\right\rangle-\nabla\cdot(M^{\top}e_{i}).$

Or, more succinctly, we have

$\displaystyle(\nabla^{*}M)$
$\displaystyle:=-M^{\top}\nabla\log\rho-\nabla\cdot M,$

where $\nabla\cdot M$ denotes the divergence operator for a matrix field
defined as $(\nabla\cdot M)_{i}=\partial_{j}M_{ji}$ in Einstein’s notation.

####  F.4.3 Commutator $[\cdot,\cdot]$

Another quantity widely used in the proof is the commutator denoted by
$[\cdot,\cdot]$. It can be thought of as an indicator for if two operators
commute. It is defined as

$[A,B]f=(AB)f-(BA)f.$

If $[A,B]f=0$, then $A$ and $B$ is said to commute. We list the following
propositions and their proofs that will be useful for the proof of 4.1. Again
for this section, we write $\rho:=\rho_{\theta,\eta_{x}}$.

######  Proposition F.6.

Let $f:\mathbb{R}^{d_{x}}\rightarrow\mathbb{R}$, and $f\in C^{3}$, then we
have

$[\nabla_{u},\nabla^{*}_{u}]\nabla_{u}f=\eta_{x}\nabla_{u}f.$

###### Proof.

Expanding out the commutator, we obtain

$[\nabla_{u},\nabla^{*}_{u}]\nabla_{u}f=-\nabla_{u}\left\langle{\nabla_{u}\log\rho},{\nabla_{u}f}\right\rangle-\nabla_{u}[\nabla_{u}\cdot(\nabla_{u}f)]+\nabla^{2}_{u}f\,\nabla_{u}\log\rho+\nabla_{u}\cdot(\nabla_{u}^{2}f),$

and since
$\nabla_{u}[\nabla_{u}\cdot(\nabla_{u}f)]=\partial_{u_{i}}\partial_{u_{j}}\partial_{u_{j}}f=\partial_{u_{j}}\partial_{u_{j}}\partial_{u_{i}}f=\nabla_{u}\cdot(\nabla_{u}^{2}f)$.
We have that

$[\nabla_{u},\nabla^{*}_{u}]\nabla_{u}f=-\nabla_{u}^{2}\log\rho\,\nabla_{u}f=\eta_{x}\nabla_{u}f.$

∎

######  Proposition F.7.

Let $f:\mathbb{R}^{d_{x}}\rightarrow\mathbb{R}$ and $f\in C^{3}$, we have that
$\nabla^{*}_{x}$ and $\nabla_{u}$ commutes for $\nabla_{u}f$ . In other words,
we have

$[\nabla^{*}_{x},\nabla_{u}]\nabla_{u}f=0.$

###### Proof.

We begin from the definition

$\displaystyle[\nabla^{*}_{x},\nabla_{u}]\nabla_{u}f=$
$\displaystyle\nabla^{*}_{x}\nabla_{u}^{2}f-\nabla_{u}\nabla^{*}_{x}\nabla_{u}f$
$\displaystyle=$
$\displaystyle-\nabla_{u}^{2}f\,{\nabla_{x}\log\rho}-\nabla_{x}\cdot(\nabla_{u}^{2}f)$
$\displaystyle+\nabla_{u}\left\langle{\nabla_{x}\log\rho},{\nabla_{u}f}\right\rangle+\nabla_{u}\nabla_{x}\cdot(\nabla_{u}f)$
$\displaystyle=$ $\displaystyle 0.$

This follows as

$\displaystyle\nabla_{x}\cdot(\nabla_{u}^{2}f)$
$\displaystyle=\partial_{x_{j}}\partial_{u_{j}}\partial_{u_{i}}f=\partial_{u_{i}}\partial_{x_{j}}\partial_{u_{j}}f=\nabla_{u}\nabla_{x}\cdot(\nabla_{u}f),$
$\displaystyle\nabla_{u}\left\langle{\nabla_{x}\log\rho},{\nabla_{u}f}\right\rangle$
$\displaystyle=\nabla_{u}^{2}f\,\nabla_{x}\log\rho+\nabla_{u}\nabla_{x}[\log\rho]\,\nabla_{u}f=\nabla_{u}^{2}f\,\nabla_{x}\log\rho.$

∎

######  Proposition F.8.

For $f:\mathbb{R}^{d_{x}}\rightarrow\mathbb{R}$ and $f\in C^{3}$, we have that
$\nabla_{x}$ and $\nabla_{u}^{*}\nabla_{u}$ commutes, i.e., we have

$\displaystyle[\nabla_{x},\nabla_{u}^{*}\nabla_{u}]f=0.$

###### Proof.

We can see that

$[\nabla_{x},\nabla_{u}^{*}\nabla_{u}]f=\nabla_{x}\nabla_{u}^{*}\nabla_{u}f-\nabla_{u}^{*}\nabla_{u}\nabla_{x}f=0.$

This follows since the first term can be written as

$\displaystyle\nabla_{x}\nabla_{u}^{*}\nabla_{u}f=\nabla_{x}\left\langle{\nabla_{u}\log\rho},{\nabla_{u}f}\right\rangle+\nabla_{x}\nabla_{u}\cdot(\nabla_{u}f)$
$\displaystyle=\nabla_{x}{\nabla_{u}[f]}\,\nabla_{u}\log\rho+\nabla_{x}\nabla_{u}\cdot(\nabla_{u}f),$

and, for the second term, we have

$\displaystyle\nabla_{u}^{*}\nabla_{u}\nabla_{x}f=\nabla_{u}\nabla_{x}[f]\,\nabla_{u}\log\rho+\nabla_{u}\cdot(\nabla_{u}\nabla_{x}f).$

Since we have

$\displaystyle\nabla_{x}\nabla_{u}\cdot(\nabla_{u}f)=\partial_{x_{i}}\partial_{u_{j}}\partial_{u_{j}}f=\partial_{u_{j}}\partial_{u_{j}}\partial_{x_{i}}f=\nabla_{u}\cdot(\nabla_{u}\nabla_{x}f),$

we have as desired. ∎

####  F.4.4 Another operator $B$

For this section, we write $\rho:=\rho_{\theta,\eta_{x}}$. Let
$f:\mathbb{R}^{2d_{x}}\rightarrow\mathbb{R}$, another operator denoted by $B$
that we are interested in is defined as follow:

$\displaystyle B[f]:=$
$\displaystyle\nabla_{(x,u)}^{*}\left[fQ\begin{pmatrix}-\nabla_{x}\log\rho\\\
-\nabla_{u}\log\rho\end{pmatrix}\right]=\nabla_{(x,u)}^{*}\left[f\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}\right]$

This can be written equivalently as

$\displaystyle B[f]=$
$\displaystyle-\left\langle{\nabla_{(x,u)}\log\rho},{f\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}}\right\rangle-\nabla_{(x,u)}\cdot\left(f\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}\right)$ $\displaystyle=$
$\displaystyle-\left\langle{\nabla_{(x,u)}f},{\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}}\right\rangle=-\left\langle{\nabla_{x}f},{\nabla_{u}\log\rho}\right\rangle+\left\langle{\nabla_{x}\log\rho},{\nabla_{u}f}\right\rangle.$

One can generalize the operator $B$ to vector-valued functions
$f:\mathbb{R}^{2d_{x}}\rightarrow\mathbb{R}^{d_{x^{\prime}}}$ as follows (its
equivalence when $d_{x^{\prime}}=1$ is easy to see):

$\displaystyle B[f]:=$
$\displaystyle\nabla_{(x,u)}^{*}\left[Q\begin{pmatrix}-\nabla_{x}\log\rho\\\
-\nabla_{u}\log\rho\end{pmatrix}\otimes
f\right]=\nabla_{(x,u)}^{*}\left[\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}\otimes f\right].$

This can be equivalently written as

$\displaystyle B[f]=$
$\displaystyle-\left[f\otimes\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho_{\theta}\end{pmatrix}\right]\,\nabla_{(x,u)}\log\rho-\nabla_{(x,u)}\cdot\left(\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}\otimes f\right)$ $\displaystyle=$
$\displaystyle\left\langle{\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}},{\nabla_{(x,u)}\log\rho}\right\rangle
f-\nabla_{(x,u)}f^{\top}\,\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}$ $\displaystyle=$
$\displaystyle\nabla_{(x,u)}f^{\top}\,\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}$

where we use the fact that

$\displaystyle\nabla_{(x,u)}\cdot\left(\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}\otimes f\right)$
$\displaystyle=\partial_{(x,u)_{j}}\left[f_{i}\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}_{j}\right]$
$\displaystyle=\partial_{(x,u)_{j}}[f_{i}]\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}_{j}+f_{i}\partial_{(x,u)_{j}}\left[\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}_{j}\right]$
$\displaystyle=\partial_{(x,u)_{j}}[f_{i}]\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}_{j}=\nabla_{(x,u)}f^{\top}\,\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}.$

In the following Proposition, we show that the operator is anti-symmetric.

######  Proposition F.9.

For all $f,g\in\mathbb{R}^{2d_{x}}\rightarrow\mathbb{R}^{d_{x^{\prime}}}$, we
have
$\left\langle{B[f]},{g}\right\rangle_{\rho}=-\left\langle{f},{B[g]}\right\rangle_{\rho}$.

###### Proof.

From the definition, we have

$\displaystyle\left\langle{B[f]},{g}\right\rangle_{\rho}$
$\displaystyle=\left\langle{\nabla_{(x,u)}^{*}\left[\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}\otimes
f\right]},{g}\right\rangle_{\rho}=\left\langle{\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}\otimes
f},{\nabla_{(x,u)}g}\right\rangle_{\rho}$
$\displaystyle=\left\langle{f},{\nabla_{(x,u)}g^{\top}\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}}\right\rangle_{\rho}$

where we use the adjoint operator of $\nabla^{*}$. Hence, we have shown that

$\left\langle{B[f]},{g}\right\rangle_{\rho}=-\left\langle{f},{B[g]}\right\rangle_{\rho}.$

∎

######  Proposition F.10.

For $f:\mathbb{R}^{2d_{x}}\rightarrow\mathbb{R}$, we have

$[\nabla_{u},B]f=\eta_{x}\nabla_{x}f$

and

$[\nabla_{x},B]f=\nabla_{x}^{2}\log\rho\,\nabla_{u}f$

###### Proof.

For the first equality, we begin by expanding out the commutator

$[\nabla_{u},B]f=\nabla_{u}B[f]-B\nabla_{u}[f].$

Since we have

$\nabla_{u}B[f]=-\nabla_{u}\nabla_{x}f\,\nabla_{u}\log\rho-\nabla_{u}^{2}\log\rho\,\nabla_{x}f+\nabla_{u}^{2}f\,\nabla_{x}\log\rho,$

and,

$B\nabla_{u}[f]=\nabla_{u}\nabla_{(x,u)}f\,\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix},$

then, we must have

$[\nabla_{u},B]f=-\nabla_{u}^{2}\log\rho\,\nabla_{x}f=\eta_{x}\nabla_{x}f.$

For the second inequality, we begin similarly

$[\nabla_{x},B]f=\nabla_{x}B[f]-B\nabla_{x}[f].$

For the first term, we have

$\nabla_{x}B[f]=-\nabla_{x}^{2}f\,\nabla_{u}\log\rho+\nabla^{2}_{x}\log\rho\,\nabla_{u}f+\nabla_{u}\nabla_{x}f\,\nabla_{x}\log\rho$

and, for the second term,

$B\nabla_{x}[f]=\nabla_{x}\nabla_{(x,u)}f\,\begin{pmatrix}\nabla_{u}\log\rho\\\
-\nabla_{x}\log\rho\end{pmatrix}.$

Hence, we have

$[\nabla_{x},B]f=[\nabla_{x},B]f=\nabla_{x}^{2}\log\rho\,\nabla_{u}f.$

∎

####  F.4.5 Simplifying (25)

For this section, we drop the subscript $\rho_{\theta,\eta_{x}}$ and write
$\rho$.

######  Proposition F.11 (Simplifying (25)).

We have that

$\displaystyle\eqref{eq:dldqdt_t2}=$
$\displaystyle-2\gamma_{x}\left\langle\nabla_{(x,u)}\nabla_{u}\left[\log\frac{q_{t}}{\rho_{t}}\right],T_{(x,u)}\nabla_{(x,u)}\nabla_{u}\left[\log\frac{q_{t}}{\rho_{t}}\right]\right\rangle_{q_{t}}$
$\displaystyle-\left\langle\nabla_{(x,u)}\log\frac{q_{t}}{\rho_{t}},\tilde{\Gamma}_{x}\nabla_{(x,u)}\log\frac{q_{t}}{\rho_{t}}\right\rangle_{q_{t}},$

where $\rho:=\rho_{\theta,\eta_{x}}$ (also similarly for $\rho_{t}$) and

$\tilde{\Gamma}_{x}=\frac{1}{2}\begin{pmatrix}\tau_{xu}\eta_{x}I_{d_{x}}&\left(\tau_{u}\eta_{x}+\tau_{xu}\gamma_{x}\eta_{x}\right)I_{d_{x}}+\tau_{x}\nabla^{2}_{x}[\ell]\\\
\left(\tau_{u}\eta_{x}+\tau_{xu}\gamma_{x}\eta_{x}\right)I_{d_{x}}+\tau_{x}\nabla^{2}_{x}[\ell]&4\gamma_{x}\tau_{u}\eta_{x}I_{d_{x}}+\tau_{xu}\nabla^{2}_{x}[\ell]\end{pmatrix}.$

###### Proof.

Recall that $\mathcal{H}_{t}:=\sqrt{\frac{q_{t}}{p_{\theta_{t}}}}$. We begin
by applying the quotient rule to obtain the fact that

$\displaystyle\nabla_{(x,u)}\left(\frac{\nabla_{(x,u)}^{*}T_{(x,u)}\nabla_{(x,u)}\mathcal{H}_{t}}{\mathcal{H}_{t}}\right)$
$\displaystyle=$
$\displaystyle\frac{\nabla_{(x,u)}\nabla_{(x,u)}^{*}\left[T_{(x,u)}\nabla_{(x,u)}\mathcal{H}_{t}\right]}{\mathcal{H}_{t}}-\frac{\nabla_{(x,u)}\mathcal{H}_{t}\nabla_{(x,u)}^{*}\left[T_{(x,u)}\nabla_{(x,u)}\mathcal{H}_{t}\right]}{\mathcal{H}_{t}^{2}}$

Then, we can write (25) as

$\displaystyle\eqref{eq:dldqdt_t2}=$
$\displaystyle-8\left\langle\frac{\nabla_{(x,u)}\nabla_{(x,u)}^{*}\left[T_{(x,u)}\nabla_{(x,u)}\mathcal{H}_{t}\right]}{\mathcal{H}_{t}},\begin{pmatrix}0_{d_{x}}&-I_{d_{x}}\\\
I_{d_{x}}&\gamma_{x}I_{d_{x}}\end{pmatrix}\nabla_{(x,u)}\log\mathcal{H}_{t}\right\rangle_{q_{t}}$
(37)
$\displaystyle+8\left\langle\frac{\nabla_{(x,u)}\mathcal{H}_{t}\nabla_{(x,u)}^{*}\left[T_{(x,u)}\nabla_{(x,u)}\mathcal{H}_{t}\right]}{\mathcal{H}^{2}_{t}},\begin{pmatrix}0_{d_{x}}&-I_{d_{x}}\\\
I_{d_{x}}&\gamma_{x}I_{d_{x}}\end{pmatrix}\nabla_{(x,u)}\log\mathcal{H}_{t}\right\rangle_{q_{t}}$
(38)

We can simplify the terms individually given in Proposition Proposition F.13
and Proposition F.12 for (37) and (38) Hence, we can write (25) equivalently
the following:

$\displaystyle\eqref{eq:dldqdt_t2}$
$\displaystyle=16\gamma_{x}\left\langle\nabla_{(x,u)}\nabla_{u}\left[\mathcal{H}_{t}\right],T_{(x,u)}\frac{\nabla_{(x,u)}\mathcal{H}_{t}}{\mathcal{H}_{t}}\otimes\nabla_{u}\mathcal{H}_{t}\right\rangle_{\rho_{t}}$
(39a)
$\displaystyle-8\gamma_{x}\left\langle\frac{\nabla_{(x,u)}\mathcal{H}_{t}}{\mathcal{H}_{t}}\otimes\nabla_{u}\mathcal{H}_{t},T_{(x,u)}\frac{\nabla_{(x,u)}\mathcal{H}_{t}}{\mathcal{H}_{t}}\otimes\nabla_{u}\mathcal{H}_{t}\right\rangle_{\rho_{t}}$
(39b)
$\displaystyle-8\gamma_{x}\left\langle\nabla_{(x,u)}\nabla_{u}\left[\mathcal{H}_{t}\right],T_{(x,u)}\nabla_{(x,u)}\nabla_{u}\left[\mathcal{H}_{t}\right]\right\rangle_{\rho_{t}}$
(39c)
$\displaystyle-2\tau_{xu}\eta_{x}\|\nabla_{x}\mathcal{H}_{t}\|^{2}_{\rho_{t}}$
(39d)
$\displaystyle-4\left\langle\nabla_{x}\mathcal{H}_{t},\left[\left(\tau_{u}\eta_{x}+\tau_{xu}\gamma_{x}\eta_{x}\right)I_{d_{x}}+\tau_{x}\nabla^{2}_{x}\left[\ell\right]\right]\nabla_{u}\mathcal{H}_{t}\right\rangle_{\rho_{t}}$
(39e)
$\displaystyle-8\left\langle\nabla_{u}\mathcal{H}_{t},\left[\tau_{u}\eta_{x}\gamma_{x}I_{d_{x}}+\frac{\tau_{xu}}{4}\nabla^{2}_{x}\left[\ell\right]\right]\nabla_{u}\mathcal{H}_{t}\right\rangle_{\rho_{t}}.$
(39f)

We can rewrite the sum of (39a), (39b), (39c) as

$\displaystyle\eqref{eq:g_t1}+\eqref{eq:g_t2}+\eqref{eq:g_t3}$
$\displaystyle=-8\gamma_{x}\left\langle\nabla_{(x,u)}\nabla_{u}\left[\mathcal{H}_{t}\right]-\frac{\nabla_{(x,u)}\mathcal{H}_{t}}{\mathcal{H}_{t}}\otimes\nabla_{u}\mathcal{H}_{t},T_{(x,u)}\left(\nabla_{(x,u)}\nabla_{u}\left[\mathcal{H}_{t}\right]-\frac{\nabla_{(x,u)}\mathcal{H}_{t}}{\mathcal{H}_{t}}\otimes\nabla_{u}\mathcal{H}_{t}\right)\right\rangle_{\rho_{t}}.$

Since we have that
$\partial_{x}\partial_{u}\mathcal{H}_{t}=\partial_{x}(\mathcal{H}_{t}\partial_{u}\log\mathcal{H}_{t})=(\partial_{x}\mathcal{H}_{t})(\partial_{u}\log\mathcal{H}_{t})+\mathcal{H}_{t}(\partial_{x}\partial_{u}\log\mathcal{H}_{t})$,
we can write the above as

$\displaystyle\eqref{eq:g_t1}+\eqref{eq:g_t2}+\eqref{eq:g_t3}$
$\displaystyle=-8\gamma_{x}\left\langle\nabla_{(x,u)}\nabla_{u}\left[\log\mathcal{H}_{t}\right],T_{(x,u)}\nabla_{(x,u)}\nabla_{u}\left[\log\mathcal{H}_{t}\right]\right\rangle_{q_{t}}$
$\displaystyle=-2\gamma_{x}\left\langle\nabla_{(x,u)}\nabla_{u}\left[\log\frac{q_{t}}{\rho_{t}}\right],T_{(x,u)}\nabla_{(x,u)}\nabla_{u}\left[\log\frac{q_{t}}{\rho_{t}}\right]\right\rangle_{q_{t}}$

As for the other terms, we have that

$\displaystyle\eqref{eq:g_t4}+\eqref{eq:g_t5}+\eqref{eq:g_t6}=$
$\displaystyle-4\left\langle\nabla_{(x,u)}\log\mathcal{H}_{t},\tilde{\Gamma}_{x}\nabla_{(x,u)}\log\mathcal{H}_{t}\right\rangle_{q_{t}}$
$\displaystyle=$
$\displaystyle-\left\langle\nabla_{(x,u)}\log\frac{q_{t}}{\rho_{t}},\tilde{\Gamma}_{x}\nabla_{(x,u)}\log\frac{q_{t}}{\rho_{t}}\right\rangle_{q_{t}},$

where

$\tilde{\Gamma}_{x}:=\frac{1}{2}\begin{pmatrix}\tau_{xu}\eta_{x}I_{d_{x}}&\left(\tau_{u}\eta_{x}+\tau_{xu}\gamma_{x}\eta_{x}\right)I_{d_{x}}+\tau_{x}\nabla^{2}_{x}\ell\\\
\left(\tau_{u}\eta_{x}+\tau_{xu}\gamma_{x}\eta_{x}\right)I_{d_{x}}+\tau_{x}\nabla^{2}_{x}\ell&4\gamma_{x}\tau_{u}\eta_{x}I_{d_{x}}+\tau_{xu}\nabla^{2}_{x}\ell\end{pmatrix}.$

We have as desired. ∎

######  Proposition F.12 (Simplifying (38)).

We have

$\displaystyle\eqref{eq:quotient_t2}=$ $\displaystyle
16\gamma_{x}\left\langle\nabla_{(x,u)}\nabla_{u}\left[\mathcal{H}_{t}\right],T_{(x,u)}\frac{\nabla_{(x,u)}\mathcal{H}}{\mathcal{H}_{t}}\otimes\nabla_{u}\mathcal{H}_{t}\right\rangle_{\rho_{t}}$
$\displaystyle-8\gamma_{x}\left\langle\frac{\nabla_{(x,u)}\mathcal{H}_{t}}{\mathcal{H}_{t}}\otimes\nabla_{u}\mathcal{H}_{t},T_{(x,u)}\frac{\nabla_{(x,u)}\mathcal{H}_{t}}{\mathcal{H}_{t}}\otimes\nabla_{u}\mathcal{H}_{t}\right\rangle_{\rho_{t}}$

###### Proof.

We have

$\displaystyle\eqref{eq:quotient_t2}=$ $\displaystyle
8\left\langle\frac{\nabla_{(x,u)}^{*}\left[T_{(x,u)}\nabla_{(x,u)}\mathcal{H}_{t}\right]}{\mathcal{H}_{t}},\left\langle\nabla_{(x,u)}\mathcal{H}_{t},\begin{pmatrix}0&-I\\\
I&\gamma_{x}I\end{pmatrix}\nabla_{(x,u)}\mathcal{H}_{t}\right\rangle\right\rangle_{\rho_{t}}$
$\displaystyle=$ $\displaystyle
8\gamma_{x}\left\langle\frac{\nabla_{(x,u)}^{*}\left[T_{(x,u)}\nabla_{(x,u)}\mathcal{H}_{t}\right]}{\mathcal{H}_{t}},\left\|\nabla_{u}\mathcal{H}_{t}\right\|_{2}^{2}\right\rangle_{\rho_{t}}.$

Then, using the adjoint of operator $\nabla_{(x,u)}^{*}$ as was described in
Section F.4.2, we obtain

$\displaystyle\eqref{eq:quotient_t2}=$ $\displaystyle
8\gamma_{x}\left\langle\nabla_{(x,u)}\frac{\left\|\nabla_{u}\mathcal{H}_{t}\right\|^{2}_{2}}{\mathcal{H}_{t}},T_{(x,u)}\nabla_{(x,u)}\mathcal{H}_{t}\right\rangle_{\rho_{t}}.$

Then, using the quotient rule, we have that