# Dynamic Flows on Curved Space Generated by Labeled Data

Xinru Hua1 Truyen Nguyen2 Tam Le3&Jose Blanchet1 Viet Anh Nguyen4 1 Stanford
University
2The University of Akron
3The Institute of Statistical Mathematics / RIKEN AIP
4Chinese University of Hong Kong

###### Abstract

The scarcity of labeled data is a long-standing challenge for many machine
learning tasks. We propose our gradient flow method to leverage the existing
dataset (i.e., source) to generate new samples that are close to the dataset
of interest (i.e., target). We lift both datasets to the space of probability
distributions on the feature-Gaussian manifold, and then develop a gradient
flow method that minimizes the maximum mean discrepancy loss. To perform the
gradient flow of distributions on the curved feature-Gaussian space, we
unravel the Riemannian structure of the space and compute explicitly the
Riemannian gradient of the loss function induced by the optimal transport
metric. For practical applications, we also propose a discretized flow, and
provide conditional results guaranteeing the global convergence of the flow to
the optimum. We illustrate the results of our proposed gradient flow method on
several real-world datasets and show our method can improve the accuracy of
classification models in transfer learning settings.

##  1 Introduction

A major challenge in many data science applications is the scarcity of labeled
data. Data augmentation methods have been studied in the literature; see for
example, the noise injection methods Moreno-Barea et al. (2018), generative
models Yi et al. (2019), and Shorten and Khoshgoftaar (2019) for a survey. We
consider a setting where one domain has only a few labeled samples for each
class, so we cannot train a well-performing classifier with the available
data. To alleviate the data scarcity problem in this setting, we propose to
enrich the target dataset by generating additional labeled samples. Using
generative models is not possible in our setting because they usually require
more than a few samples for each class to learn and generate high-quality new
samples Gao et al. (2018). In our work, we choose a source dataset with
extensive labeled data and then flow the labeled data to the target dataset.
Precisely, we introduce a novel data augmentation methodology based on a
gradient flow approach that minimizes the maximum mean discrepancy
($\mathrm{MMD}$) distance between the target and the augmented data.
Therefore, when minimizing the $\mathrm{MMD}$ distance, we are able to obtain
an efficient scheme which generates additional labeled data from the target
distribution. Our scheme is model-independent and can be applied to any
datasets regardless of the number of classes or dimensionality.

Mathematically, we consider a feature space $\mathcal{X}={\mathbb{R}}^{m}$ and
a categorical label space $\mathcal{Y}$. We have a source domain dataset
consisting of $N$ samples $(x_{i},y_{i})\in\mathcal{X}\times\mathcal{Y}$ for
$i=1,\ldots,N$, and a target domain dataset of $M$ samples
$(\bar{x}_{j},\bar{y}_{j})\in\mathcal{X}\times\mathcal{Y}$ for
$j=1,\ldots,M(M\ll N)$. The ultimate goal of this paper is to generate new
samples in the target domain, and we aim to generate new samples whose
distribution is as close as possible to the distribution that governs the
target domain.

We here introduce a gradient flow method Arbel et al. (2019); Mroueh et al.
(2019) to synthesize new, unseen data samples. Gradient flow is a continuous
flow along the path where a considered loss function decreases its value.
Because we have extensive source domain samples, it is possible to flow each
source sample towards the target data while minimizing the loss function. The
terminal product of the flow will be new samples that can sufficiently
approximate the distribution of the target domain. Thus, gradient flow is an
approach to synthesize new target domain samples, and is a complement to data
augmentation methods, like adding random noise.

Unfortunately, formulating a gradient flow algorithm for labeled data with
categorical set $\mathcal{Y}$ is problematic. Indeed, there is no clear metric
structure on $\mathcal{Y}$ in order to define the topological neighborhood,
this in turn leads to the difficulty of forming the gradients with respect to
the categorical component. To overcome this difficulty, we lift each
individual label to a richer structure. For example, a label such as “0” is
replaced by a mean vector and a covariance matrix based on the whole
distribution of the information associated to this particular label. Then it
will be much more natural to apply gradient flow algorithms in the space of
the lifted representation. A gradient flow on the dataset space with this idea
was recently proposed in Alvarez-Melis and Fusi (2021) by leveraging a new
notion of distance between datasets in Alvarez-Melis and Fusi (2020); Courty
et al. (2017); Damodaran et al. (2018). The main idea behind this approach is
to reparametrize the categorical space $\mathcal{Y}$ using the conditional
distribution of the features, which is assumed to be Gaussian, and then
construct a gradient flow on the feature-Gaussian space. Nevertheless, the
theoretical analysis in Alvarez-Melis and Fusi (2021) focuses solely on the
gradients with respect to the feature with no treatment of the flow with
respect to the Gaussian component. In fact, the space of Gaussian
distributions is not a (flat) vector space, and extracting gradient
information depends on the choice of the metric imposed on this Gaussian
space. On the other hand, our method computes the full gradient with respect
to the Gaussian component (the mean and covariance matrix component that
correspond to the label component).

Paper | Dataset |  | On curved Riemannian space
---
| Gradient has mean
---
and covariance component
Alvarez-Melis and Fusi (2021) | synthetic, *NIST, and CIFAR10 |  |
Arbel et al. (2019) | synthetic |  |
Ours | synthetic, *NIST, and TinyImageNet | ✓ | ✓

Table 1:  To the best of our knowledge, we provide the _first_ results on the
full gradient of the features and lifted labels on a curved Riemannian space.
We also conduct numerical experiments on the highest-dimension real-world
datasets.

Our gradient flows minimize the $\mathrm{MMD}$ loss function, thus it belongs
to the family of $\mathrm{MMD}$ gradient flows that was pioneered in Mroueh et
al. (2019); Arbel et al. (2019), and further extended in Mroueh and Nguyen
(2021). The $\mathrm{MMD}$ function compares two distributions via their
kernel mean embeddings on a flat reproducing kernel Hilbert space (RKHS). In
contrast to the Kullback-Leibler divergence flow, the $\mathrm{MMD}$ flow can
employ a sample approximation for the target distribution Liu (2017). Further,
the squared $\mathrm{MMD}$ possesses unbiased sample gradients Bińkowski et
al. (2018); Bellemare et al. (2017). However, existing literature on
$\mathrm{MMD}$ flows focus on distributions on flat Euclidean spaces. The flow
developed in our paper here is for distributions on a curved Riemannian
feature-Gaussian space. Moreover, our approach is distinctive from the flow in
Alvarez-Melis and Fusi (2021) because we impose a specific metric on the
Gaussian component, and we compute explicitly the Riemannian gradient of the
$\mathrm{MMD}$ loss function with respect to this metric to formulate our
flow. Table 1 compares our work with two recent papers on gradient flow in
theory and numerical experiments.

Recently, generative models Rezende et al. (2016); Wang et al. (2021) are
successful in generating image samples from given distributions. The most
important difference with our method is that generative models learn a prior
distribution from massive data that are similar to the target data and
generate new target samples conditioning on the prior distribution Wang et al.
(2020); Gao et al. (2018). Comparatively, our algorithm can transfer between
two non-similar and non-related distributions, for example, from random
Gaussian noise to MNIST in Supplementary B.11. Another benefit of our method
is that we provide conditions for global convergence of our algorithms in
Section 4, whereas generative models or more specific, generative adversarial
networks (GANs), currently do not guarantee global convergence Wiatrak et al.
(2019).

The application of our gradient flow is few-shot transfer learning, where we
want to train classifiers with limited labeled data in the target domain. The
numerical experiments in Section 5 demonstrate that our gradient flows can
effectively augment the target data, and thus can significantly boost the
accuracy in the classification task in the few-shot learning setting.
Moreover, we run experiments on Tiny ImageNet datasets to highlight that our
algorithm is scalable to higher-dimensional image data, that is higher than
recent gradient flow works Alvarez-Melis and Fusi (2021); Fan and Alvarez-
Melis (2022). We also compare our method with Alvarez-Melis and Fusi (2021),
mixup method Zhang et al. (2017), and traditional data augmentation methods.
Results in Supplementary B.8–B.10 show that our method improves the accuracy
in transfer learning more than these methods.

Some related works study nonparametric gradient flows using the
$2$-Wasserstein distance between distributions Ambrosio et al. (2008); Jordan
et al. (1998); Otto (2001); Villani (2008); Santambrogio (2015, 2017); Frogner
and Poggio (2020), but only for distributions on Euclidean spaces and a
different distance. Nonparametric gradient flows with other metrics include
Sliced-Wasserstein Descent Liutkus et al. (2019), Stein Descent Liu (2017);
Liu and Wang (2016), and Sobolev Descent Mroueh et al. (2019). However, they
also only consider distributions on Euclidean spaces. In particular, Liu
(2017) introduce Riemannian structures for the Stein geometry on flat spaces,
while ours is for an optimal transport metric on a curved space. Parametric
flows for training GANs are studied in Chizat and Bach (2018); Chen and Li
(2018); Arbel et al. (2020); Mroueh and Nguyen (2021).

#### Contributions.

We study a gradient flow approach to synthesize new labeled samples related to
the target domain. To construct this flow, we consider the space of
probability distributions on the feature-Gaussian manifold, and we are
metrizing this space with an optimal transport distance. We summarize the
contributions of this paper as follows.

  * •

We study in details the Riemannian structure of the feature-Gaussian manifold
in Section 3, as well as the Riemannian structure of the space of probability
measures supported on this manifold in Supplementary A.1.

  * •

We consider a gradient flow that minimizes the squared $\mathrm{MMD}$ loss
function to the target distribution. We describe explicitly the (Riemannian)
gradient of the squared $\mathrm{MMD}$ in Lemma 5, and we provide a partial
differential equation describing the evolution of the gradient flow that
follows the (Riemannian) steepest descent direction.

  * •

We propose two discretized schemes to approximate the continuous gradient flow
equation in Section 4.1 and 4.2. We provide conditions guaranteeing the global
convergence of our gradient flows to the optimum in both schemes.

  * •

In Section 5, we demonstrate numerical results with our method on real-world
image datasets. We show that our method can generate high-fidelity images and
improve the classification accuracy in transfer learning settings.

#### Notations.

We use ${\mathbb{S}}^{n}$ to denote the set of $n\times n$ real and symmetric
matrices, and ${\mathbb{S}}^{n}_{++}\subset{\mathbb{S}}^{n}$ consists of all
positive definite matrices. For $A\in{\mathbb{S}}^{n}$,
$\text{tr}(A)\coloneqq\sum_{i}A_{ii}$. We use $\langle\cdot,\cdot\rangle$ and
$\|\cdot\|_{2}$ to denote the standard inner product and norm on Euclidean
spaces. Let ${\mathcal{P}}(X)$ be the collection of all probability
distributions with finite second moment on metric space $X$. If $\varphi:X\to
Y$ is a Borel map and $\nu\in{\mathcal{P}}(X)$, then the push-forward
$\varphi_{\\#}\nu$ is the distribution on $Y$ given by
$\varphi_{\\#}\nu(E)=\nu(\varphi^{-1}(E))$ for all Borel sets $E\subset Y$.
For a function $f$ of the continuous time variable $t$, $f_{t}$ denotes the
value of $f$ at $t$ while $\partial_{t}f$ denotes the standard derivative of
$f$ w.r.t. $t$. Also, $\delta_{z}$ denotes the Dirac delta measure at $z$.

All proofs are provided in the Supplementary material.

##  2 Labeled Data Synthesis via Gradient Flows of Lifted Distributions

In this section, we describe our approach to synthesize target domain samples
using gradient flows. A holistic view of our method is presented in Fig. 1.

Figure 1: Schematic view of our approach: The source and target datasets are
first lifted to distributions $\rho^{0}$ and $\varrho$ on the feature-Gaussian
space (left box). We then run a gradient flow for $T$ iterations to get a
terminal distribution $\rho^{T}$ (middle). Atoms of $\rho^{T}$ are projected
to get labeled target samples (right).

In the first step, we would need to lift the feature-label space
$\mathcal{X}\times\mathcal{Y}$ to a higher dimensional space where a metric
can be defined. Consider momentarily the source data samples
$(x_{i},y_{i})_{i=1}^{N}$. Notice that this data can be represented as an
empirical distribution $\nu$ on $\mathcal{X}\times\mathcal{Y}$. More
precisely, we have $\nu=N^{-1}\sum_{i=1}^{N}\delta_{(x_{i},y_{i})}$. As
$\mathcal{Y}$ is discrete, the law of conditional probabilities allows us to
dis-integrate $\nu$ into the conditional distributions $\nu_{y}$ of $X|Y=y$
satisfying $\nu(E\times F)=\int_{F}\nu_{y}(E)\nu^{2}(\mathrm{d}y)$ for every
$E\subset\mathcal{X}$ and $F\subset\mathcal{Y}$, where $\nu^{2}\coloneqq
N^{-1}\sum_{i=1}^{N}\delta_{y_{i}}$ is the second marginal of $\nu$ (Ambrosio
et al., 2008, Theorem 5.3.1). The lifting procedure is obtained by employing a
pre-determined mapping $\phi:\mathcal{X}\to{\mathbb{R}}^{n}$, and any
categorical value $y\in\mathcal{Y}$ can now be represented as an
$n$-dimensional distribution $\phi_{\\#}\nu_{y}$. Using this lifting, any
source sample $(x_{i},y_{i})\in\mathcal{X}\times\mathcal{Y}$ is lifted to a
point
$(x_{i},\phi_{\\#}\nu_{y_{i}})\in\mathcal{X}\times\mathcal{P}({\mathbb{R}}^{n})$
and the source dataset is representable as an empirical distribution of the
form $N^{-1}\sum_{i=1}^{N}\delta_{(x_{i},\phi_{\\#}\nu_{y_{i}})}$.

The lifted representation of a categorical value $y\in\mathcal{Y}$ as an
$n$-dimensional distribution
$\phi_{\\#}\nu_{y}\in\mathcal{P}({\mathbb{R}}^{n})$ is advantageous because
$\mathcal{P}({\mathbb{R}}^{n})$ is metrizable, for example, using the
$2$-Wasserstein distance. The downside is that $\mathcal{P}({\mathbb{R}}^{n})$
is infinite dimensional, and encoding the datasets in this lifted
representation is not efficient. To resolve this issue, we assume that
$\phi_{\\#}\nu_{y}$ is Gaussian for all $y\in\mathcal{Y}$, and thus any
distribution $\phi_{\\#}\nu_{y}$ can be characterized by the mean vector
$\mu_{y}\in{\mathbb{R}}^{n}$ and covariance matrix
$\Sigma_{y}\in{\mathbb{S}}_{++}^{n}$ defined as
$\mu_{y}=\int_{\mathcal{X}}\phi(x)\nu_{y}(\mathrm{d}x)\mbox{ and
}\Sigma_{y}=\int_{\mathcal{X}}\big{[}\phi(x)-\mu_{y}\big{]}\big{[}\phi(x)-\mu_{y}\big{]}^{\top}\nu_{y}(\mathrm{d}x)$
for all $y\in\mathcal{Y}$, where ⊤ denotes the transposition of a vector. In
real-world settings, the conditional moments of $\phi(X)|Y$ are sufficiently
different for $y\neq y^{\prime}$, and thus the representations using
$(\mu_{y},\Sigma_{y})$ will unlikely lead to any loss of label information.
With this lifting, the source data thus can be represented as an empirical
distribution $\rho^{0}$ on
${\mathbb{R}}^{m}\times{\mathbb{R}}^{n}\times{\mathbb{S}}_{++}^{n}$ via
$\rho^{0}=N^{-1}\sum\nolimits_{i=1}^{N}\delta_{(x_{i},\mu_{y_{i}},\Sigma_{y_{i}})}$.
By an analogous construction to compute $\bar{\mu}_{y}$ and $\bar{\Sigma}_{y}$
using the target data, the target domain data
$(\bar{x}_{j},\bar{y}_{j})_{j=1}^{M}$ can be represented as another empirical
distribution
$\varrho=M^{-1}\sum\nolimits_{j=1}^{M}\delta_{(\bar{x}_{j},\bar{\mu}_{\bar{y}_{j}},\bar{\Sigma}_{\bar{y}_{j}})}$.
Let us denote the shorthand
$\mathcal{Z}={\mathbb{R}}^{m}\times{\mathbb{R}}^{n}\times{\mathbb{S}}_{++}^{n}$,
then $\rho^{0}$ and $\varrho$ are both probability measures on $\mathcal{Z}$.
We refer to $\rho^{0}$ and $\varrho$ as the feature-Gaussian representations
of the source and target datasets.

We now consider the gradient flow associated with the optimization problem

$\min\limits_{\rho\in\mathcal{P}(\mathcal{Z})}\leavevmode\nobreak\
\Big{\\{}\mathcal{F}(\rho)\coloneqq\frac{1}{2}\mathrm{MMD}(\rho,\varrho)^{2}\Big{\\}}$

under the initialization $\rho=\rho^{0}$. The objective function
$\mathcal{F}(\rho)$ quantifies how far an incumbent solution $\rho$ is from
the target distribution $\varrho$, measured using the $\mathrm{MMD}$ distance.
In Sections 3 and 4, we will provide the necessary ingredients to construct
this flow.

Suppose that after $T$ iterations of the discretized gradient flow algorithm,
we obtain a distribution $\rho^{T}\in\mathcal{P}(\mathcal{Z})$ that is
sufficiently close to $\varrho$, i.e., $\mathcal{F}(\rho^{T})$ is close to
zero. Then we can recover new target labels by projecting the samples of the
distribution $\rho^{T}$ to the locations on $\mathcal{X}\times\mathcal{Y}$.
This projection can be computed efficiently by solving a linear optimization
problem, as discussed in Supplementary B.3.

######  Remark 1 (Reduction of dimensions).

If $m=n$ and $\phi$ is the identity map, then our lifting procedure coincides
with that proposed in Alvarez-Melis and Fusi (2020). However, a large $n$ is
redundant, especially when the cardinality of $\mathcal{Y}$ is low. If $n\ll
m$, then $\phi$ offers significant reduction in the number of dimensions, and
will speed up the gradient flow algorithms.

######  Remark 2 (Generalization to elliptical distributions).

Our framework can be extended to the symmetric elliptical distributions
because the Bures distance for elliptical distributions admits the same
closed-form as for the Gaussian distributions Gelbrich (1990). In this paper,
we use $\phi$ as the t-SNE embedding. According to van der Maaten and Hinton
(2008), t-SNE’s low-dimensional embedded space forms a Student-t distribution,
which is an elliptical distribution.

##  3 Riemannian Geometry of $\mathcal{Z}$ and $\mathcal{P}(\mathcal{Z})$

If we opt to measure the distance between two Gaussian distributions using the
2-Wasserstein metric, then this choice would induce a natural distance $d$ on
the space
$\mathcal{Z}={\mathbb{R}}^{m}\times{\mathbb{R}}^{n}\times{\mathbb{S}}_{++}^{n}$
prescribed as

$\displaystyle
d\big{(}(x_{1},\mu_{1},\Sigma_{1}),(x_{2},\mu_{2},\Sigma_{2})\big{)}$
$\displaystyle\hskip
10.00002pt\coloneqq\big{[}\|x_{1}-x_{2}\|_{2}^{2}+\|\mu_{1}-\mu_{2}\|_{2}^{2}+\mathbb{B}(\Sigma_{1},\Sigma_{2})^{2}\big{]}^{\frac{1}{2}},$
(3.1)

where $\mathbb{B}$ is the Bures metric on ${\mathbb{S}}^{n}_{++}$ given by
$\mathbb{B}(\Sigma_{1},\Sigma_{2})\coloneqq\big{[}\text{tr}(\Sigma_{1}+\Sigma_{2}-2[\Sigma_{1}^{\frac{1}{2}}\Sigma_{2}\Sigma_{1}^{\frac{1}{2}}]^{\frac{1}{2}})\big{]}^{\frac{1}{2}}$.

As $\mathbb{B}$ is a metric on ${\mathbb{S}}^{n}_{+}$ (Bhatia et al., 2019,
p.167), $d$ is hence a product metric on $\mathcal{Z}$. In this section,
first, we study the non-Euclidean geometry of $\mathcal{Z}$ under the ground
metric $d$. Second, we investigate the Riemannian structure on
$\mathcal{P}(\mathcal{Z})$, the space of all distributions supported on
$\mathcal{Z}$ and with finite second moment, that is induced by the optimal
transport distance. These Riemannian structures are required to define the
Riemannian gradients of any loss functionals on $\mathcal{P}(\mathcal{Z})$,
and will remain important in our development of the gradient flow for the
squared $\mathrm{MMD}$.

The space $\mathcal{Z}$ is not a linear vector space. In this section, we
reveal the Riemannian structure on $\mathcal{Z}$ associated to the ground
metric $d$. As we shall see, $\mathcal{Z}$ is a curved space as its geodesics
are not straight lines and involve solutions to the Lyapunov equation. For any
positive definite matrix $\Sigma\in{\mathbb{S}}^{n}_{++}$ and any symmetric
matrix $V\in{\mathbb{S}}^{n}$, the Lyapunov equation

$H\Sigma+\Sigma H=V$ (3.2)

has a unique solution $H\in{\mathbb{S}}^{n}$ (Bhatia, 1997, Theorem VII.2.1).
Let ${\mathrm{L}}_{\Sigma}[V]$ denote this unique solution $H$.

The space ${\mathbb{S}}^{n}_{++}$ is a Riemannian manifold with the Bures
metric $\mathbb{B}$ as the associated distance function, see (Takatsu, 2011,
Proposition A). Since $\mathcal{Z}$ is the product of two Euclidean spaces and
${\mathbb{S}}^{n}_{++}$, this gives rise to the following geometric structure
for $\mathcal{Z}$.

######  Proposition 3 (Geometry of $\mathcal{Z}$).

The space $\mathcal{Z}$ is a Riemannian manifold: at each point
$z=(x,\mu,\Sigma)\in\mathcal{Z}$, the tangent space is
$\mathrm{T}_{z}\mathcal{Z}={\mathbb{R}}^{m}\times{\mathbb{R}}^{n}\times{\mathbb{S}}^{n}$
and the Riemannian metric is

$\displaystyle\big{\langle}(w_{1},v_{1},V_{1}),(w_{2},v_{2},V_{2})\big{\rangle}_{z}$
$\displaystyle\hskip 50.00008pt\coloneqq\langle w_{1},w_{2}\rangle+\langle
v_{1},v_{2}\rangle+\langle V_{1},V_{2}\rangle_{\Sigma}$ (3.3)

for two tangent vectors $(w_{1},v_{1},V_{1})$ and $(w_{2},v_{2},V_{2})$ in
${\mathbb{R}}^{m}\times{\mathbb{R}}^{n}\times{\mathbb{S}}^{n}$, where $\langle
V_{1},V_{2}\rangle_{\Sigma}\coloneqq\text{tr}\Big{(}{\mathrm{L}}_{\Sigma}[V_{1}]\,\Sigma\,{\mathrm{L}}_{\Sigma}[V_{2}]\Big{)}$.
Moreover, the distance function corresponding to this Riemannian metric
coincides with the distance $d$ given by (3).

As $\mathcal{Z}$ is a product Riemannian manifold, any geodesic in
$\mathcal{Z}$ is of the form $(\theta,\gamma,\Gamma)$ with $\theta$, $\gamma$
being the Euclidean geodesics (straight lines) and $\Gamma$ being a geodesic
in the Riemannian manifold ${\mathbb{S}}^{n}_{++}$. More precisely, for each
$\Sigma\in{\mathbb{S}}^{n}_{++}$ and each tangent vector
$V\in{\mathbb{S}}^{n}$, the geodesic in the manifold ${\mathbb{S}}^{n}_{++}$
emanating from $\Sigma$ with direction $V$ is given by

$\Gamma(t)=(I+t{\mathrm{L}}_{\Sigma}[V])\Sigma(I+t{\mathrm{L}}_{\Sigma}[V])\quad\mbox{for
}t\in J^{*},$ (3.4)

where $J^{*}$ is the open interval about the origin given by
$J^{*}=\\{t\in{\mathbb{R}}:\,I+t{\mathrm{L}}_{\Sigma}[V]\in{\mathbb{S}}^{n}_{++}\\}$
Malagò et al. (2018). As a consequence, for each point
$(x,\mu,\Sigma)\in\mathcal{Z}$ and each tangent vector
$(w,v,V)\in{\mathbb{R}}^{m}\times{\mathbb{R}}^{n}\times{\mathbb{S}}^{n}$, the
Riemannian exponential map in $\mathcal{Z}$ for $t\in J^{*}$ is given by

$\exp_{(x,\mu,\Sigma)}(t(w,v,V))\coloneqq(\theta(t),\gamma(t),\Gamma(t)).$
(3.5)

where $\theta(t)\coloneqq x+tw$, $\gamma(t)\coloneqq\mu+tv$, and $\Gamma(t)$
is defined by (3.4). By definition, $t\mapsto\exp_{(x,\mu,\Sigma)}(t(w,v,V))$
is the geodesic emanating from $(x,\mu,\Sigma)$ with direction $(w,v,V)$.

Given the Riemannian metric (3), one can define the corresponding notion of
gradient and divergence Lee (2003). For a differentiable function
$\varphi:\mathcal{Z}\to{\mathbb{R}}$, its gradient
$\widetilde{\nabla}_{d}\varphi(z)$ w.r.t. the metric $d$ defined by (3) is the
unique element in the tangent space
${\mathbb{R}}^{m}\times{\mathbb{R}}^{n}\times{\mathbb{S}}^{n}$ satisfying

$\big{\langle}\widetilde{\nabla}_{d}\varphi(z),(w,v,V)\big{\rangle}_{z}=D\varphi_{z}(w,v,V)$

for all
$(w,v,V)\in{\mathbb{R}}^{m}\times{\mathbb{R}}^{n}\times{\mathbb{S}}^{n}$ with
$D\varphi_{z}(w,v,V)$ denoting the standard directional derivative of
$\varphi$ at $z$ in the direction $(w,v,V)$. By exploiting the special form of
$\langle\cdot,\cdot\rangle_{z}$ in (3), we can compute
$\widetilde{\nabla}_{d}\varphi(z)$ explicitly:

######  Lemma 4 (Gradients).

For a differentiable function $\varphi:\mathcal{Z}\to{\mathbb{R}}$, we have
for $z=(x,\mu,\Sigma)$ that

$\widetilde{\nabla}_{d}\varphi(z)=\big{(}\nabla_{x}\varphi(z),\leavevmode\nobreak\
\nabla_{\mu}\varphi(z),\leavevmode\nobreak\
2[\nabla_{\Sigma}\varphi(z)]\Sigma+2\Sigma[\nabla_{\Sigma}\varphi(z)]\big{)},$

(3.6)

where $(\nabla_{x},\nabla_{\mu},\nabla_{\Sigma})$ are the standard (Euclidean)
gradients of the respective components.

The last component in formula (3.6) for $\widetilde{\nabla}_{d}\varphi$
reflects the curved geometry of $\mathcal{Z}$, and can be interpreted as the
Riemannian gradient of the function $\Sigma\mapsto\varphi(x,\mu,\Sigma)$
w.r.t. the Bures distance $\mathbb{B}$.

For a continuous vector field
$\Phi:\mathcal{Z}\to{\mathbb{R}}^{m}\times{\mathbb{R}}^{n}\times{\mathbb{S}}^{n}$
and a distribution $\rho\in{\mathcal{P}}(\mathcal{Z})$, the divergence
$\mathrm{div}_{d}(\rho\Phi)$ is the signed measure on $\mathcal{Z}$ satisfying
the integration by parts formula

$\int_{\mathcal{Z}}\varphi(z)\,\mathrm{div}_{d}(\rho\Phi)(\mathrm{d}z)=-\int_{\mathcal{Z}}\langle\Phi(z),\widetilde{\nabla}_{d}\varphi(z)\rangle_{z}\,\,\rho(\mathrm{d}z)$

for every differentiable function $\varphi:\mathcal{Z}\to{\mathbb{R}}$ with
compact support. In case $\rho$ has a density w.r.t. the Riemannian volume
form on $\mathcal{Z}$, then this definition coincides with the standard
divergence operator induced by Riemannian metric (3). The optimal transport
distance and its induced Riemannian metric on the space
${\mathcal{P}}(\mathcal{Z})$ are relegated to Supplementary A.1

##  4 Gradient Flow for Maximum Mean Discrepancy

As ${\mathcal{P}}(\mathcal{Z})$ is an infinite dimensional curved space, many
machine learning methods based on finite dimensional or linear structure
cannot be directly applied to this manifold. To circumvent this problem, we
use a positive definite kernel to map ${\mathcal{P}}(\mathcal{Z})$ to a RKHS
and then perform our analysis on it. Let $k$ be a positive definite kernel on
$\mathcal{Z}$, and let $\mathcal{H}$ be the RKHS generated by $k$. The inner
product on $\mathcal{H}$ is denoted by
$\langle\cdot,\cdot\rangle_{\mathcal{H}}$, and the kernel mean embedding
$\rho\in{\mathcal{P}}(\mathcal{Z})\longmapsto\mathbf{m}_{\rho}(\cdot)\in\mathcal{H}$
is given by
$\mathbf{m}_{\rho}(z)\coloneqq\int_{\mathcal{Z}}k(z,w)\,\rho(\mathrm{d}w)$ for
$z$ in $\mathcal{Z}$. The $\mathrm{MMD}$ Gretton et al. (2012) between
$\rho\in\mathcal{P}(\mathcal{Z})$ and the target $\varrho$ is defined as the
maximum of the mean difference between the two distributions over all test
functions in the unit ball of $\mathcal{H}$ (see Supplementary A.3). Moreover,
it can be expressed by
$\mathrm{MMD}(\rho,\varrho)=\|\mathbf{m}_{\rho}-\mathbf{m}_{\varrho}\|_{\mathcal{H}}$.
When $k$ is characteristic, the kernel mean embedding
$\rho\mapsto\mathbf{m}_{\rho}$ is injective and therefore,
$\mathrm{MMD}(\rho,\varrho)=0$ if and only if $\rho=\varrho$.

Consider the loss function
${\mathcal{F}}[\rho]\coloneqq\frac{1}{2}\mathrm{MMD}(\rho,\varrho)^{2}=\frac{1}{2}\|\mathbf{m}_{\rho}-\mathbf{m}_{\varrho}\|_{\mathcal{H}}^{2}$.
As explained in the introduction, there are three advantages of $\mathrm{MMD}$
over Kullback-Leibler divergence: its associated gradient flow can employ a
sample approximation for the target distribution, the input distribution
$\rho$ does not have to be absolutely continuous w.r.t. the target
distribution $\varrho$, and the squared $\mathrm{MMD}$ possesses unbiased
sample gradients. For each $\rho$, the Riemannian gradient
$\mathrm{grad}\,{\mathcal{F}}[\rho]$ is defined as the unique element in
$\mathrm{T}_{\rho}{\mathcal{P}}(\mathcal{Z})$ satisfying
$g_{\rho}(\mathrm{grad}\,{\mathcal{F}}[\rho],\zeta)=\frac{\mathrm{d}}{\mathrm{d}t}\Big{|}_{t=0}{\mathcal{F}}[\rho_{t}]$
for every differentiable curve $t\mapsto\rho_{t}\in{\mathcal{P}}(\mathcal{Z})$
passing through $\rho$ at $t=0$ with tangent vector
$\partial_{t}\rho_{t}|_{t=0}=\zeta$. By using the Riemannian metric tensor
(A.3), we can compute explicitly this gradient.

######  Lemma 5 (Gradient formula).

The Riemannian gradient of $\mathcal{F}$ satisfies
$\mathrm{grad}\,{\mathcal{F}}[\rho]=-\mathrm{div}_{d}\left(\rho\widetilde{\nabla}_{d}[\mathbf{m}_{\rho}-\mathbf{m}_{\varrho}]\right).$

The Riemannian gradient $\mathrm{grad}\,{\mathcal{F}}$ on
${\mathcal{P}}(\mathcal{Z})$ depends not only on the gradient operator
$\widetilde{\nabla}_{d}$ but also on the divergence operator. Using Lemma 5,
we can rewrite the gradient flow equation
$\partial_{t}\rho_{t}=-\mathrm{grad}\,{\mathcal{F}}[\rho_{t}]$ explicitly as

$\partial_{t}\rho_{t}=\mathrm{div}_{d}\big{(}\rho_{t}\widetilde{\nabla}_{d}[\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}]\big{)}\quad\mbox{for}\quad
t\geq 0.$ (4.1)

The next result exhibits the rate at which ${\mathcal{F}}$ decreases its value
along the flow.

######  Proposition 6 (Rate of decrease).

Along the gradient flow $t\mapsto\rho_{t}\in{\mathcal{P}}(\mathcal{Z})$ given
by (4.1), we have

$\frac{\mathrm{d}}{\mathrm{d}t}{\mathcal{F}}[\rho_{t}]=-\int_{\mathcal{Z}}\big{\|}\widetilde{\nabla}_{d}[\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}]\big{\|}_{z}^{2}\,\rho_{t}(\mathrm{d}z)\quad\mbox{for}\quad
t\geq 0.$

Proposition 6 implies that
$\frac{\mathrm{d}}{\mathrm{d}t}{\mathcal{F}}[\rho_{t}]=0$ if and only if
$\widetilde{\nabla}_{d}[\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}](z)=0$ for
every $z$ in the support of the distribution $\rho_{t}$. Thus, the objective
function will decrease whenever the gradient
$\widetilde{\nabla}_{d}[\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}]$ is not
identically zero.

###  4.1 Riemannian Forward Euler Scheme

We propose the Riemannian version of the forward Euler scheme to discretize
continuous flow (4.1):

$\displaystyle\rho^{\tau+1}=\exp(s_{\tau}\Phi^{\tau})_{\\#}\rho^{\tau}\qquad$
(4.2) $\displaystyle\mbox{with}\leavevmode\nobreak\ \leavevmode\nobreak\
\Phi^{\tau}\coloneqq-\widetilde{\nabla}_{d}[\mathbf{m}_{\rho^{\tau}}-\mathbf{m}_{\varrho}],$

where $s_{\tau}>0$ is the step size. Here, for a vector field
$\Phi=(\Phi_{1},\Phi_{2},\Phi_{3}):\mathcal{Z}\to{\mathbb{R}}^{m}\times{\mathbb{R}}^{n}\times{\mathbb{S}}^{n}$
and for $\varepsilon\geq 0$, $\exp(\varepsilon\Phi):\mathcal{Z}\to\mathcal{Z}$
is the Riemannian exponential map induced by (3.5), i.e., for
$z=(x,\mu,\Sigma)\in\mathcal{Z}$:

$\displaystyle\exp_{z}(\varepsilon\Phi(z))=\begin{pmatrix}x+\varepsilon\Phi_{1}(z)\\\
\mu+\varepsilon\Phi_{2}(z)\\\
(I+\varepsilon{\mathrm{L}}_{\Sigma}[\Phi_{3}(z)])\Sigma(I+\varepsilon{\mathrm{L}}_{\Sigma}[\Phi_{3}(z)])\end{pmatrix}.$

Notice in the above equation that the input $z$ affects simultaneously the
bases of the exponential map $\exp_{z}$ as well as the direction $\Phi(z)$.
This map is the $\varepsilon$-perturbation of the identity map along geodesics
with directions $\Phi$. When
$\rho^{\tau}=N^{-1}\sum_{i=1}^{N}\delta_{z^{\tau}_{i}}$ is an empirical
distribution, scheme (4.2) flows each particle $z^{\tau}_{i}$ to the new
position $z^{\tau+1}_{i}=\exp_{z^{\tau}_{i}}(s_{\tau}\Phi(z^{\tau}_{i}))$. The
next lemma shows that $\Phi^{\tau}$ is the steepest descent direction for
${\mathcal{F}}$ w.r.t. the exponential map among all directions in the space
$\mathbb{L}^{2}(\rho^{\tau})$, which is the collection of all vector fields
$\Phi$ on $\mathcal{Z}$ satisfying
$\|\Phi\|_{\mathbb{L}^{2}(\rho^{\tau})}^{2}\coloneqq\int_{\mathcal{Z}}\|\Phi(z)\|_{z}^{2}\rho^{\tau}(\mathrm{d}z)<\infty$.

######  Lemma 7 (Steepest descent direction).

Fix a distribution $\rho^{\tau}\in\mathcal{P}(\mathcal{Z})$. For any vector
field
$\Phi:\mathcal{Z}\to{\mathbb{R}}^{m}\times{\mathbb{R}}^{n}\times{\mathbb{S}}^{n}$,
we have

$\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}{\mathcal{F}}[\exp(\varepsilon\Phi)_{\\#}\rho^{\tau}]=\int_{\mathcal{Z}}\langle\widetilde{\nabla}_{d}[\mathbf{m}_{\rho^{\tau}}-\mathbf{m}_{\varrho}](z),\Phi(z)\rangle_{z}\leavevmode\nobreak\
\rho^{\tau}(\mathrm{d}z).$

If $\hat{\Phi}^{\tau}$ is the unit vector field (w.r.t. the
$\|\cdot\|_{\mathbb{L}^{2}(\rho^{\tau})}$ norm) in the direction of
$\Phi^{\tau}$ given in (4.2), then

$\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\big{|}_{\varepsilon=0}{\mathcal{F}}[\exp(\varepsilon\hat{\Phi}^{\tau})_{\\#}\rho^{\tau}]=-\|\widetilde{\nabla}_{d}[\mathbf{m}_{\rho^{\tau}}-\mathbf{m}_{\varrho}]\|_{\mathbb{L}^{2}(\rho^{\tau})}$

and this is the fastest decay rate among all unit directions $\Phi$ in
$\mathbb{L}^{2}(\rho^{\tau})$.

It follows from Lemma 7 that the discrete scheme (4.2) satisfies the
Riemannian gradient descent property: if
$\widetilde{\nabla}_{d}[\mathbf{m}_{\rho^{\tau}}-\mathbf{m}_{\varrho}]$ is
nonzero and if $s_{\tau}>0$ is chosen sufficiently small, then
${\mathcal{F}}[\rho^{\tau+1}]<{\mathcal{F}}[\rho^{\tau}]$. In Proposition 14
in the Supplementary, we quantify the amount of decrease of $\mathcal{F}$ at
each iteration. Algorithm 1 implements the flow (4.2) iteratively. Each
iteration in Algorithm 1 has complexity $O(N(Nm+n^{3}))$, where $m$ is the
feature’s dimension, $n$ is the reduced dimension ($n\ll m$), $N$ is the
number of particles.

Algorithm 1 Discretized Gradient Flow Algorithm for Scheme (4.2)

1: Input: a source distribution
$\rho^{0}=N^{-1}\sum_{i=1}^{N}\delta_{z^{0}_{i}}$, a target distribution
$\varrho=M^{-1}\sum_{j=1}^{M}\delta_{\bar{z}_{j}}$, a number of iterations
$T$, a sequence of step sizes $s_{\tau}>0$ with $\tau=0,1,...,T$ and a kernel
$k$

2: Initialization: Compute
$\bar{\Psi}(z)=M^{-1}\sum_{j=1}^{M}\widetilde{\nabla}_{d}^{1}k(z,\bar{z}_{j})$
with $\widetilde{\nabla}_{d}^{1}k(z,\bar{z}_{j})$ is $\widetilde{\nabla}_{d}$
of $z\mapsto k(z,\bar{z}_{j})$

3: repeat for each $\tau=0,\ldots,T-1$:

4: Compute
$\Psi^{\tau}(z)=N^{-1}\sum_{i=1}^{N}\widetilde{\nabla}_{d}^{1}k(z,z^{\tau}_{i})$

5: for $i=1,\ldots,N$

6: do
$z^{\tau+1}_{i}\leftarrow\exp_{z^{\tau}_{i}}\big{(}s_{\tau}(\bar{\Psi}-\Psi^{\tau})(z^{\tau}_{i})\big{)}$

7: end for

8: Output: $\rho^{T}=N^{-1}\sum_{i=1}^{N}\delta_{z^{T}_{i}}$

#### Convergence.

We now study the (weak) convergence of the solution $\rho_{t}$ of the
continuous gradient flow (4.1), as well as the discretized counterpart
$\rho^{\tau}$ of flow (4.2), to the target distribution $\varrho$. When the
kernel $k$ is characteristic, this convergence is equivalent to
$\lim_{t\to\infty}\mathrm{MMD}(\rho_{t},\varrho)=0$. Because the objective
function ${\mathcal{F}}$ is not displacement convex (Arbel et al., 2019,
Section 3.1), the convergent theory for gradient flows in Ambrosio et al.
(2008) does not apply even in the case of Euclidean spaces. In general, there
is a possibility that $\mathrm{MMD}(\rho_{t},\varrho)$ does not decrease to
zero as $t\rightarrow\infty$. In view of Proposition 6, this happens if the
solutions $\rho_{t}$ are trapped inside the set
$\big{\\{}\rho:\,\int_{\mathcal{Z}}\big{\|}\widetilde{\nabla}_{d}[\mathbf{m}_{\rho}-\mathbf{m}_{\varrho}]\big{\|}_{z}^{2}\,\rho(\mathrm{d}z)=0\big{\\}}$.
For each distribution $\rho$ on $\mathcal{Z}$, we define in Supplementary A.3
a symmetric linear and positive operator
$\mathbb{K}_{\rho}:\mathcal{H}\to\mathcal{H}$ with the property that
$\langle\mathbb{K}_{\rho}[\mathbf{m}_{\rho}-\mathbf{m}_{\varrho}],\mathbf{m}_{\rho}-\mathbf{m}_{\varrho}\rangle_{\mathcal{H}}=\int_{\mathcal{Z}}\big{\|}\widetilde{\nabla}_{d}[\mathbf{m}_{\rho}-\mathbf{m}_{\varrho}]\big{\|}_{z}^{2}\,\rho(\mathrm{d}z)$
We further show in Proposition 16 that $\rho_{t}$ globally converges in
$\mathrm{MMD}$ if the minimum eigenvalue $\lambda_{t}$ of the operator
$\mathbb{K}_{\rho_{t}}$ satisfies an integrability condition.

###  4.2 Noisy Riemannian Forward Euler Scheme

The analysis in Section 4.1 reveals that the gradient flows suffer from
convergence issues if the residual
$\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}$ belongs to the null space of the
operator $\mathbb{K}_{\rho_{t}}$. To resolve this, we employ graduated
optimization Arbel et al. (2019); Gulcehre et al. (2016, 2017); Hazan et al.
(2016) used for non-convex optimization in Euclidean spaces. Specifically, we
modify algorithm (4.2) by injecting Gaussian noise into the exponential map at
each iteration $\tau$ to obtain

$\displaystyle\rho^{\tau+1}=\exp(s_{\tau}\Phi^{\tau})_{\\#}\rho^{\tau,\beta_{\tau}}\quad$
(4.3) $\displaystyle\text{with}\hskip
1.99997ptf^{\beta_{\tau}}:(z,u)\mapsto\exp_{z}(\beta_{\tau}u),\,\,\,\rho^{\tau,\beta_{\tau}}\coloneqq{f^{\beta_{\tau}}}_{\\#}(\rho^{\tau}\otimes
g).$

Here $g$ is a Gaussian measure with distribution
$\mathcal{N}_{{\mathbb{R}}^{m}}(0,1)\otimes\mathcal{N}_{{\mathbb{R}}^{n}}(0,1)\otimes\mathcal{N}_{{\mathbb{S}}^{n}}(0,1)$
on the tangent space and $\mathcal{N}_{{\mathbb{S}}^{n}}(0,1)$ denotes an
$n$-by-$n$ symmetric matrix whose upper triangular elements are i.i.d.
standard Gaussian random variables. When
$\rho^{\tau}=N^{-1}\sum_{i=1}^{N}\delta_{z^{\tau}_{i}}$, scheme (4.3) flows
each particle $z^{\tau}_{i}$ first to
$z^{\tau,\beta_{\tau}}_{i}\coloneqq\exp_{z^{\tau}_{i}}(\beta_{\tau}U)$ with
noise $U\sim g$ and then to
$z^{\tau+1}_{i}=\exp_{z^{\tau,\beta_{\tau}}_{i}}(s_{\tau}\Phi(z^{\tau,\beta_{\tau}}_{i}))$.
Our next result extends Proposition 8 in Arbel et al. (2019) for the standard
quadratic cost on the Euclidean space to the nonstandard cost function $d^{2}$
on the curved Riemannian manifold $\mathcal{Z}_{++}$. It demonstrates that
scheme (4.3) achieves the global minimum of ${\mathcal{F}}$ provided that $k$
is a Lipschitz-gradient kernel and both the noise level $\beta_{\tau}$ and the
step size $s_{\tau}$ are well controlled. The proof of Proposition 8 is given
in Supplementary A.3 and relies on arguments that are different from that of
Arbel et al. (2019).

######  Proposition 8 (Objective value decay for noisy scheme).

Suppose that $k$ is a Lipschitz-gradient kernel111See Definition A.3 for the
technical definition with constant $L$, and the noise level $\beta_{\tau}$
satisfies

$\lambda\beta_{\tau}^{2}{\mathcal{F}}[\rho^{\tau}]\leq\int_{\mathcal{Z}}\|\Phi^{\tau}(z)\|_{z}^{2}\,\rho^{\tau,\beta_{\tau}}(\mathrm{d}z)$
(4.4)

for some constant $\lambda>0$. Then for $\rho^{\tau+1}$ obtained from scheme
(4.3), we have

${\mathcal{F}}[\rho^{\tau+1}]\leq{\mathcal{F}}[\rho^{0}]\,\exp\Big{(}-\lambda\sum\nolimits_{i=0}^{\tau}[s_{i}\big{(}1-2Ls_{i}\big{)}\beta_{i}^{2}]\Big{)}.$

In particular, ${\mathcal{F}}[\rho^{\tau}]$ tends to zero if the sequence
$\sum_{i=0}^{\tau}s_{i}\big{(}1-2Ls_{i}\big{)}\beta_{i}^{2}$ goes to positive
infinity. For an adaptive step size $s_{\tau}\leq 1/4L$, this condition is met
if, for example, $\beta_{\tau}$ is chosen of the form $(\tau
s_{\tau})^{-\frac{1}{2}}$ while still satisfying (4.4). The noise perturbs the
direction of descent, whereas the step size determines how far to move along
this perturbed direction. The noise level needs to be adjusted so that the
gradient is not too blurred, but it does not necessarily decrease at each
iteration. When the incumbent distribution $\rho^{\tau}$ is close to a local
optimum, it is helpful to increase the noise level to escape the local
optimum. We demonstrate in Lemma 13 in the Supplementary that any positive
definite kernel $k$ with bounded Hessian w.r.t. distance $d$ is a Lipschitz-
gradient kernel. Algorithm 3 in the Supplementary describes (4.3) in details.

##  5 Numerical Experiments

We evaluate the proposed gradient flow on real-world datasets and then
illustrate its applications in transfer learning. We augment samples for the
target dataset, where only a few samples in the dataset are available. We
consider three datasets: the MNIST (M) LeCun and Cortes (2010), Fashion-MNIST
(F) Xiao et al. (2017), Kuzushiji-MNIST (K) Clanuwat et al. (2018). To satisfy
the Gaussianity assumption of the conditional distributions, we cluster all
the images from each class of the datasets and keep the largest cluster for
each class. To demonstrate the scalability of our algorithm to higher-
dimensional images, we run experiments on Tiny ImageNet (TIN) Russakovsky et
al. (2015) and upscaled SVHN Netzer et al. (2011) datasets, where images are
of $3\times 64\times 64$ size.

Our mapping $\phi$ is from $\mathbb{R}^{m}$ to $\mathbb{R}^{2}$ in the lifting
procedure. To compute the $\mathrm{MMD}$ distance using kernel embeddings, we
use a tensor kernel $k$ on $\mathcal{Z}$ composed from three standard Gaussian
kernels corresponding for each component of the feature space
${\mathbb{R}}^{m}$, the mean space ${\mathbb{R}}^{2}$ and the covariance
matrix space ${\mathbb{S}}^{2}_{++}$. As a consequence, $k$ is a
characteristic kernel by (Szabó and Sriperumbudur, 2018, Theorem 4).

Figure 2: Sample path visualizations for five pairs of source-target domain.
The original image and additional results are in the supplementary.

#### Experiment: Gradient Flow between Datasets.

We visualize the path travelled by each sample from the source domain to the
target domain, as depicted in Fig. 2. We draw randomly $N=200$ images equally
for 10 classes of the source domain, and $M=50$ images equally for 10 classes
of the target domain ($M=10$ for the TIN and SVHN datasets). In each
subfigure, each column represents a snapshot of a certain time-step and the
samples flow from the source (left) to the target (right) as the number of
steps increases. The first column in Fig. 2 are the images from the source
domain, where the gradient flows start. Empirically, the algorithm converges
after step 140 for *NIST datasets and step 6000 for TIN and SVHN. The
experiments are run on a C5.4xlarge AWS instance (a CPU instance) and all
finish within one hour.

Figure 3: Average target domain accuracy on the test split for transfer
learning with one-shot (left) and five-shot (right). Results are taken over 10
independent replications, and the range of accuracy is displayed by the error
bars.

###  5.1 Application in Transfer Learning

Our gradient flow can alleviate the problem of insufficient labeled data by
synthesizing new samples to augment the target dataset. In this section, we
demonstrate that the generated target domain samples can improve the accuracy
in one-shot and five-shot transfer learning tasks.

First, we fix a source domain and pretrain a classifier $P$ on this domain. We
draw randomly $N$ samples from the source domain to form the source dataset
$(x_{i},y_{i})_{i=1}^{N}$. Next, we pick a target domain and draw randomly a
few samples from this target domain: for example, in 1-shot learning, only 1
image per class from the target domain is selected to form the target dataset
$D=(\bar{x}_{j},\bar{y}_{j})_{j=1}^{M}$. We then perform a noisy gradient flow
scheme (4.3) from the source dataset to the target dataset to get N new
samples $S_{T}=(x_{i}^{T},y_{i}^{T})_{i=1}^{N}$. With the target dataset $D$
and new samples $S_{T}$, we can retrain the classifier $P$. Similarly, we can
also train new classifiers from scratch using datasets $D$ and $D\cup S_{T}$.
Finally, we test the classifiers on the test set of the target domain.

Fig. 3 presents the accuracy of five transfer learning strategies on four
pairs of source and target domain. For the labels above the plot, labels
without $P$ mean training a new classifier from scratch, whereas labels with
$P$ mean transferring the pretrained classifier. $D$ and $S_{T}$ represent the
samples in the target domain and our flowed samples. We observe a common trend
that the addition of the flowed samples $S_{T}$ always improves the accuracy
of the classifiers, as we compare $D\cup S_{T}$ with $D$ and compare $P\cup
D\cup S_{T}$ with $P\cup D$. Moreover, the data augmentation with $S_{T}$
leads to a higher increase of accuracy for the 1-shot learning, where the data
scarcity problem is more severe. The transfer learning results for SVHN and
TIN datasets are provided in the Supplementary B.7. Although Few-shot learning
is more challenging due to the high complexity of the datasets, the addition
of $S_{T}$ still improves the accuracy. We also compare with Alvarez-Melis and
Fusi (2021)222The only gradient flow work that has experiments on *NIST
datasets, but it does not run experiments on TIN and SVHN., mixup method Zhang
et al. (2017), and image augmentation methods, see results in Supplementary
B.8–B.10.

#### Conclusions.

This paper focuses on a gradient flow approach to generate new labeled data
samples in the target domain. To overcome the discrete nature of the labels,
we represent datasets as distributions on the feature-Gaussian space, and the
flow is formulated to minimize the $\mathrm{MMD}$ loss function under an
optimal transport metric. Contrary to existing gradient flows on linear
structure, our flows are developed on the curved Riemannian manifold of
Gaussian distributions. We provide explicit formula for the Riemannian
gradient of the $\mathrm{MMD}$ loss, and analyze in details the flow equations
and the convergence properties of both continuous and discretized forms. The
numerical experiments demonstrate that our method can efficiently generate
high-fidelity labeled training data for real-world datasets, and improve the
classification accuracy in few-shot learning. The main limitation exists in
the assumption that the data of one label forms an elliptical distribution.

## Ethical Statement

Our work has positive societal impacts, because it can help reduce repetitive
data collection and labeling work. It does not have possible negative societal
impacts in the current stage.

## References

  * Alvarez-Melis and Fusi [2020] David Alvarez-Melis and Nicolo Fusi.  Geometric dataset distances via optimal transport.  In Advances in Neural Information Processing Systems, volume 33, pages 21428–21439, 2020.
  * Alvarez-Melis and Fusi [2021] David Alvarez-Melis and Nicolò Fusi.  Dataset dynamics via gradient flows in probability space.  In Proceedings of the 38th International Conference on Machine Learning, volume 139, pages 219–230, 2021.
  * Ambrosio et al. [2008] L. Ambrosio, N. Gigli, and G. Savaré.  Gradient Flows in Metric Spaces and in the Space of Probability Measures.  Birkhäuser Verlag, 2008.
  * Arbel et al. [2019] Michael Arbel, Anna Korba, Adil Salim, and Arthur Gretton.  Maximum mean discrepancy gradient flow.  In Advances in Neural Information Processing Systems, volume 32, pages 6481–6491, 2019.
  * Arbel et al. [2020] M Arbel, A Gretton, W Li, and G Montufar.  Kernelized Wasserstein natural gradient.  In International Conference on Learning Representations, 2020.
  * Bellemare et al. [2017] Marc G. Bellemare, Ivo Danihelka, Will Dabney, Shakir Mohamed, Balaji Lakshminarayanan, Stephan Hoyer, and Remi Munos.  The Cramer distance as a solution to biased Wasserstein gradients.  arXiv preprint arXiv:1705.10743, 2017.
  * Benamou and Brenier [2000] Jean-David Benamou and Yann Brenier.  A computational fluid mechanics solution to the Monge-Kantorovich mass transfer problem.  Numerische Mathematik, 2000.
  * Bhatia et al. [2019] R. Bhatia, T. Jain, and Y. Lim.  On the Bures-Wasserstein distance between positive definite matrices.  Expositiones Mathematicae, 37(2):165–191, 2019.
  * Bhatia [1997] Rajendra Bhatia.  Matrix Analysis.  Springer, 1997.
  * Bińkowski et al. [2018] Mikołaj Bińkowski, Danica J. Sutherland, Michael Arbel, and Arthur Gretton.  Demystifying MMD GANs.  In International Conference on Learning Representations, 2018.
  * Chen and Li [2018] Yifan Chen and Wuchen Li.  Optimal transport natural gradient for statistical manifolds with continuous sample space.  In Information Geometry, 2018.
  * Chizat and Bach [2018] Lénaïc Chizat and Francis Bach.  On the global convergence of gradient descent for over-parameterized models using optimal transport.  In Advances in Neural Information Processing Systems, volume 31, 2018.
  * Clanuwat et al. [2018] Tarin Clanuwat, Mikel Bober-Irizar, Asanobu Kitamoto, Alex Lamb, Kazuaki Yamamoto, and David Ha.  Deep learning for classical Japanese literature.  arXiv preprint arXiv:1812.01718, 2018.
  * Courty et al. [2017] Nicolas Courty, Rémi Flamary, Amaury Habrard, and Alain Rakotomamonjy.  Joint distribution optimal transportation for domain adaptation.  In Advances in Neural Information Processing Systems, volume 30, 2017.
  * Damodaran et al. [2018] Bharath Bhushan Damodaran, Benjamin Kellenberger, Rémi Flamary, Devis Tuia, and Nicolas Courty.  DeepJDOT: Deep joint distribution optimal transport for unsupervised domain adaptation.  In Proceedings of the European Conference on Computer Vision (ECCV), pages 447–463, 2018.
  * Fan and Alvarez-Melis [2022] Jiaojiao Fan and David Alvarez-Melis.  Generating synthetic datasets by interpolating along generalized geodesics.  In NeurIPS 2022 Workshop on Synthetic Data for Empowering ML Research, 2022.
  * Frogner and Poggio [2020] Charlie Frogner and Tomaso Poggio.  Approximate inference with Wasserstein gradient flows.  In International Conference on Artificial Intelligence and Statistics, pages 2581–2590. PMLR, 2020.
  * Gao et al. [2018] Hang Gao, Zheng Shou, Alireza Zareian, Hanwang Zhang, and Shih-Fu Chang.  Low-shot learning via covariance-preserving adversarial augmentation networks.  Advances in Neural Information Processing Systems, 31, 2018.
  * Gelbrich [1990] M. Gelbrich.  On a formula for the ${L}^{2}$ Wasserstein metric between measures on Euclidean and Hilbert spaces.  Mathematische Nachrichten, 147(1):185–203, 1990.
  * Gretton et al. [2012] Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola.  A kernel two-sample test.  Journal of Machine Learning Research, 13(25):723–773, 2012.
  * Gulcehre et al. [2016] Caglar Gulcehre, Marcin Moczulski, Misha Denil, and Yoshua Bengio.  Noisy activation functions.  In Proceedings of The 33rd International Conference on Machine Learning, pages 3059–3068, 2016.
  * Gulcehre et al. [2017] Caglar Gulcehre, Marcin Moczulski, Francesco Visin, and Yoshua Bengio.  Mollifying networks.  In 5th International Conference on Learning Representations, 2017\.
  * Hazan et al. [2016] Elad Hazan, Kfir Yehuda Levy, and Shai Shalev-Shwartz.  On graduated optimization for stochastic non-convex problems.  In Proceedings of The 33rd International Conference on Machine Learning, pages 1833–1841, 2016.
  * Hinton et al. [2012] Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky.  Neural networks for machine learning, Lecture 6a: Overview of mini-batch gradient descent.  2012\.
  * Jordan et al. [1998] R. Jordan, D. Kinderlehrer, and F. Otto.  The variational formulation of the Fokker–Planck equation.  SIAM Journal on Mathematical Analysis, 29:1–17, 1998.
  * LeCun and Cortes [2010] Yann LeCun and Corinna Cortes.  MNIST handwritten digit database.  2010\.
  * Lee [2003] John Lee.  Introduction to Smooth Manifolds.  Springer-Verlag, 2003.
  * Liu and Wang [2016] Qiang Liu and Dilin Wang.  Stein variational gradient descent: A general purpose Bayesian inference algorithm.  In Advances in Neural Information Processing Systems, volume 29, 2016.
  * Liu [2017] Qiang Liu.  Stein variational gradient descent as gradient flow.  In Advances in Neural Information Processing Systems, volume 30, 2017.
  * Liutkus et al. [2019] Antoine Liutkus, Umut Simsekli, Szymon Majewski, Alain Durmus, and Fabian-Robert Stöter.  Sliced-Wasserstein flows: Nonparametric generative modeling via optimal transport and diffusions.  In Proceedings of the 36th International Conference on Machine Learning, 2019.
  * Malagò et al. [2018] L. Malagò, L. Montrucchio, and G. Pistone.  Wasserstein Riemannian geometry of Gaussian densities.  Information Geometry, 1(2):137–179, 2018.
  * Moreno-Barea et al. [2018] Francisco J. Moreno-Barea, Fiammetta Strazzera, José M. Jerez, Daniel Urda, and Leonardo Franco.  Forward noise adjustment scheme for data augmentation.  In 2018 IEEE Symposium Series on Computational Intelligence (SSCI), pages 728–734, 2018.
  * Mroueh and Nguyen [2021] Youssef Mroueh and Truyen Nguyen.  On the convergence of gradient descent in GANs: MMD GAN as a gradient flow.  In International Conference on Artificial Intelligence and Statistics, 2021.
  * Mroueh et al. [2019] Youssef Mroueh, Tom Sercu, and Anant Raj.  Sobolev descent.  In Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics, pages 2976–2985, 2019.
  * Netzer et al. [2011] Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng.  Reading digits in natural images with unsupervised feature learning.  2011\.
  * Otto [2001] F. Otto.  The geometry of dissipative evolution equations: the porous medium equation.  Comm. Partial Differential Equations, 26:101–174, 2001.
  * Rezende et al. [2016] Danilo Rezende, Ivo Danihelka, Karol Gregor, Daan Wierstra, et al.  One-shot generalization in deep generative models.  In International conference on machine learning, pages 1521–1529. PMLR, 2016.
  * Russakovsky et al. [2015] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al.  Imagenet large scale visual recognition challenge.  International journal of computer vision, 115(3):211–252, 2015\.
  * Santambrogio [2015] Filippo Santambrogio.  Optimal Transport for Applied Mathematicians: Calculus of Variations, PDEs and Modeling.  Birkhäuser, 2015.
  * Santambrogio [2017] Filippo Santambrogio.  Euclidean, metric, and Wasserstein gradient flows: An overview.  Bullentin of Mathematical Sciences, 7:87–154, 2017.
  * Shorten and Khoshgoftaar [2019] Connor Shorten and Taghi M Khoshgoftaar.  A survey on image data augmentation for deep learning.  Journal of big data, 6(1):1–48, 2019.
  * Szabó and Sriperumbudur [2018] Zoltán Szabó and Bharath K. Sriperumbudur.  Characteristic and universal tensor product kernels.  Journal of Machine Learning Research, 18(233):1–29, 2018.
  * Takatsu [2011] Asuka Takatsu.  Wasserstein geometry of Gaussian measures.  Osaka Journal of Mathematics, 48(4):1005–1026, 2011.
  * van der Maaten and Hinton [2008] Laurens van der Maaten and Geoffrey Hinton.  Visualizing data using t-SNE.  Journal of Machine Learning Research, 9:2579–2605, 2008.
  * Villani [2008] C. Villani.  Optimal Transport: Old and New.  Springer Science & Business Media, 2008.
  * Wang et al. [2020] Yaqing Wang, Quanming Yao, James T Kwok, and Lionel M Ni.  Generalizing from a few examples: A survey on few-shot learning.  ACM computing surveys (csur), 53(3):1–34, 2020.
  * Wang et al. [2021] Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, and Can Yang.  Deep generative learning via Schrödinger bridge.  In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 10794–10804. PMLR, 18–24 Jul 2021.
  * Wiatrak et al. [2019] Maciej Wiatrak, Stefano V Albrecht, and Andrew Nystrom.  Stabilizing generative adversarial networks: A survey.  arXiv preprint arXiv:1910.00927, 2019.
  * Xiao et al. [2017] Han Xiao, Kashif Rasul, and Roland Vollgraf.  Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms.  arXiv preprint arXiv:1708.07747, 2017.
  * Yi et al. [2019] Xin Yi, Ekta Walia, and Paul Babyn.  Generative adversarial network in medical imaging: A review.  Medical Image Analysis, 58:101552, 2019.
  * Zhang et al. [2017] Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz.  mixup: Beyond empirical risk minimization.  arXiv preprint arXiv:1710.09412, 2017.

The Supplementary Material is organized into two parts. In Section A, we
provide the proofs and further discussions of the results in the main paper.
Section B includes implementation details as well as additional numerical
results. All the models and data used to create the results in the paper are
in the supplementary file.

##  Appendix A Proofs

###  A.1 Optimal Transport and Riemannian Structure on
${\mathcal{P}}(\mathcal{Z})$

To define a gradient flow for probability distributions on $\mathcal{Z}$, it
is essential to have a concept of gradients for functionals defined on
$\mathcal{P}(\mathcal{Z})$. This requires a meaningful Riemannian structure on
${\mathcal{P}}(\mathcal{Z})$, and here, we adopt a Riemannian structure
generated by the optimal transport on ${\mathcal{P}}(\mathcal{Z})$ with ground
cost $d^{2}$. The optimal transport metric $\mathbb{W}(\rho_{0},\rho_{1})$
between any two distributions $\rho_{0}$,
$\rho_{1}\in{\mathcal{P}}(\mathcal{Z})$ is defined by formula (A.7) of
Appendix A.2. As $(\mathcal{Z},d)$ is a Riemannian manifold by Proposition 3,
it follows from the celebrated Benamou-Brenier formula Benamou and Brenier
[2000] that $\mathbb{W}$ can be expressed in terms of a dynamic formulation.
Precisely,

$\displaystyle\mathbb{W}(\rho_{0},\rho_{1})^{2}=\inf_{(\rho,\phi)\in\mathcal{A}(\rho_{0},\,\rho_{1})}\int_{0}^{1}\int_{\mathcal{Z}}\|\widetilde{\nabla}_{d}\phi_{t}(z)\|_{z}^{2}\rho_{t}(\mathrm{d}z)\mathrm{d}t,$
(A.1)

where $\mathcal{A}(\rho_{0},\,\rho_{1})$ is the collection of all pairs
$(\rho,\phi)$ of curve $\rho:[0,1]\to{\mathcal{P}}(\mathcal{Z})$ with
endpoints $\rho_{0}$ and $\rho_{1}$, and function
$\phi:[0,1]\times\mathcal{Z}\to{\mathbb{R}}$ that satisfies the continuity
equation

$\partial_{t}\rho+\mathrm{div}_{d}(\rho_{t}\widetilde{\nabla}_{d}\phi_{t})=0$
(A.2)

in the sense of distributions in $(0,1)\times\mathcal{Z}$. For brevity,
hereafter $\rho_{t}$ and $\phi_{t}$ denote functions in the $z$ variable
defined by $\rho_{t}(z)=\rho(t,z)$ and $\phi_{t}(z)=\phi(t,z)$.

Riemannian metric on ${\mathcal{P}}(\mathcal{Z})$. The formulation (A.1) gives
rise to the following Riemannian structure on ${\mathcal{P}}(\mathcal{Z})$
induced by distance $\mathbb{W}$. First, the continuity equation enables us to
identify a tangent vector $\partial_{t}\rho$ with the divergence
$-\mathrm{div}_{d}(\rho_{t}\widetilde{\nabla}_{d}\phi_{t})$. Thus the tangent
space of ${\mathcal{P}}(\mathcal{Z})$ at a distribution $\rho$ can be defined
as
$\mathrm{T}_{\rho}{\mathcal{P}}(\mathcal{Z})\coloneqq\big{\\{}-\mathrm{div}_{d}(\rho\widetilde{\nabla}_{d}\varphi):\,\varphi\,\mbox{is
a differentiable function with compact support on}\,\mathcal{Z}\big{\\}}$.
Second, we let
$g_{\rho}:\mathrm{T}_{\rho}{\mathcal{P}}(\mathcal{Z})\times\mathrm{T}_{\rho}{\mathcal{P}}(\mathcal{Z})\longrightarrow{\mathbb{R}}$
be the Riemannian metric tensor given by

$g_{\rho}(\zeta_{1},\zeta_{2})\coloneqq\int_{\mathcal{Z}}\langle\widetilde{\nabla}_{d}\varphi_{1}(z),\widetilde{\nabla}_{d}\varphi_{2}(z)\rangle_{z}\,\,\rho(\mathrm{d}z)$
(A.3)

for $\zeta_{1}=-\mathrm{div}_{d}(\rho\widetilde{\nabla}_{d}\varphi_{1})$ and
$\zeta_{2}=-\mathrm{div}_{d}(\rho\widetilde{\nabla}_{d}\varphi_{2})$. With
this definition and due to (A.2), formula (A.1) can be rewritten using the
metric tensor as

$\mathbb{W}(\rho_{0},\rho_{1})^{2}=\inf_{(\rho,\phi)\in\mathcal{A}(\rho_{0},\,\rho_{1})}\int_{0}^{1}g_{\rho_{t}}(\partial_{t}\rho,\partial_{t}\rho)\,\mathrm{d}t.$

The metric tensor (A.3) allows us to define a notion of Riemannian gradients
for functionals on ${\mathcal{P}}(\mathcal{Z})$. In the next section we shall
compute this gradient explicitly for the squared $\mathrm{MMD}$ gradient flow.

###  A.2 Proofs and Results related to Section 3

#### For Proposition 3.

Recall that the Bures distance defined on ${\mathbb{S}}_{++}^{n}$ is

$\mathbb{B}(\Sigma_{1},\Sigma_{2})\coloneqq\big{[}\text{tr}(\Sigma_{1}+\Sigma_{2}-2[\Sigma_{1}^{\frac{1}{2}}\Sigma_{2}\Sigma_{1}^{\frac{1}{2}}]^{\frac{1}{2}})\big{]}^{\frac{1}{2}},$
(A.4)

and $\widetilde{\nabla}_{d}\varphi(z)$ is the unique element in the tangent
space ${\mathbb{R}}^{m}\times{\mathbb{R}}^{n}\times{\mathbb{S}}^{n}$
satisfying

$\big{\langle}\widetilde{\nabla}_{d}\varphi(z),(w,v,V)\big{\rangle}_{z}=D\varphi_{z}(w,v,V)\quad\mbox{for
all }(w,v,V)\in{\mathbb{R}}^{m}\times{\mathbb{R}}^{n}\times{\mathbb{S}}^{n}.$
(A.5)

The proof of Proposition 3 relies on the following result from [Takatsu, 2011,
Proposition A] (see also [Bhatia et al., 2019, Theorem 5] and [Malagò et al.,
2018, Proposition 6]).

######  Proposition 9.

The space ${\mathbb{S}}^{n}_{++}$ is a Riemannian manifold with the following
structures: at each point $\Sigma\in{\mathbb{S}}^{n}_{++}$, the tangent space
is $\mathrm{T}_{\Sigma}{\mathbb{S}}^{n}_{++}={\mathbb{S}}^{n}$ and the
Riemannian metric is given by

$\langle
X,Y\rangle_{\Sigma}\coloneqq\text{tr}\Big{(}{\mathrm{L}}_{\Sigma}[X]\,\Sigma\,{\mathrm{L}}_{\Sigma}[Y]\Big{)}=\frac{1}{2}\langle{\mathrm{L}}_{\Sigma}[X],Y\rangle\quad\mbox{for}\quad
X,Y\in{\mathbb{S}}^{n}.$

Moreover, the distance function corresponding to this Riemannian metric
coincides with the Bures distance $\mathbb{B}$ given by (A.4).

We are now ready to prove Proposition 3.

###### Proof of Proposition 3.

As a consequence of Proposition 9, $\mathcal{Z}$ is the product Riemannian
manifold with tangent space
$\mathrm{T}_{(x,\mu,\Sigma)}\mathcal{Z}=\mathrm{T}_{x}{\mathbb{R}}^{m}\times\mathrm{T}_{\mu}{\mathbb{R}}^{n}\times\mathrm{T}_{\Sigma}{\mathbb{S}}^{n}_{++}$
and with the standard product Riemmanian metric (3). The result then follows.
We note that if we let
${\mathbb{D}}((x_{1},\mu_{1},\Sigma_{1}),(x_{2},\mu_{2},\Sigma_{2}))$ denote
the distance function corresponding to the Riemannian metric
$\langle\cdot,\cdot\rangle_{z}$ on $\mathcal{Z}$, then its square
${\mathbb{D}}((x_{1},\mu_{1},\Sigma_{1}),(x_{2},\mu_{2},\Sigma_{2}))^{2}$ is
the sum of the square of the distance function w.r.t. standard metric
$\langle\cdot,\cdot\rangle$ on ${\mathbb{R}}^{m}$, the square of the distance
function w.r.t. standard metric $\langle\cdot,\cdot\rangle$ on
${\mathbb{R}}^{n}$, and the square of the distance function w.r.t. metric
$\langle\cdot,\cdot\rangle_{\Sigma}$ on ${\mathbb{S}}^{n}_{++}$. As a result
and by Proposition 9, we have
${\mathbb{D}}((x_{1},\mu_{1},\Sigma_{1}),(x_{2},\mu_{2},\Sigma_{2}))^{2}=\|x_{1}-x_{2}\|_{2}^{2}+\|\mu_{1}-\mu_{2}\|_{2}^{2}+\mathbb{B}(\Sigma_{1},\Sigma_{2})^{2}$.
So, ${\mathbb{D}}$ is the same as $d$. ∎

#### For Lemma 4

###### Proof of Lemma 4.

Let us express
$\widetilde{\nabla}_{d}\varphi(z)=(\Phi_{1}(z),\Phi_{2}(z),\Phi_{3}(z))$ with
$\Phi_{1}(z)\in{\mathbb{R}}^{m},\,\Phi_{2}(z)\in{\mathbb{R}}^{n}$ and
$\Phi_{3}(z)\in{\mathbb{S}}^{n}$. Then by using the definition of Riemannian
metric $\langle\cdot,\cdot\rangle_{z}$ in (3), we can rewrite equation (A.5)
as

$\langle\Phi_{1}(z),v\rangle+\langle\Phi_{2}(z),w\rangle+\langle\frac{1}{2}{\mathrm{L}}_{\Sigma}[\Phi_{3}(z)],V\rangle=\langle\nabla\varphi(z),(v,w,V)\rangle.$

This is equivalent to

$\langle\Phi_{1}(z),v\rangle+\langle\Phi_{2}(z),w\rangle+\langle\frac{1}{2}{\mathrm{L}}_{\Sigma}[\Phi_{3}(z)],V\rangle=\langle\nabla_{x}\varphi(z),v\rangle+\langle\nabla_{\mu}\varphi(z),w\rangle+\langle\nabla_{\Sigma}\varphi(z),V\rangle,$
(A.6)

where
$\nabla\varphi(z)=\Big{(}\nabla_{x}\varphi(z),\nabla_{\mu}\varphi(z),\nabla_{\Sigma}\varphi(z)\Big{)}$
denotes the standard Euclidean gradient. Equation (A.6) is obviously satisfied
if $\Phi_{1}(z)=\nabla_{x}\varphi(z)$, $\Phi_{2}(z)=\nabla_{\mu}\varphi(z)$,
and ${\mathrm{L}}_{\Sigma}[\Phi_{3}(z)]=2\nabla_{\Sigma}\varphi(z)$. By the
definition of operator ${\mathrm{L}}_{\Sigma}$ right after (3.2), the third
identity is the same as
$\Phi_{3}(z)=2[\nabla_{\Sigma}\varphi(z)]\Sigma+2\Sigma[\nabla_{\Sigma}\varphi(z)]$.
Due to uniqueness of the gradient, we therefore infer that
$\widetilde{\nabla}_{d}\varphi(z)$ is given by the formula:

$\widetilde{\nabla}_{d}\varphi(z)=\Big{(}\nabla_{x}\varphi(z),\nabla_{\mu}\varphi(z),2[\nabla_{\Sigma}\varphi(z)]\Sigma+2\Sigma[\nabla_{\Sigma}\varphi(z)]\Big{)}.$

This completes the proof. ∎

In this paper, the optimal transport metric between any two distributions
$\rho_{0}$, $\rho_{1}\in{\mathcal{P}}(\mathcal{Z})$ is defined by

$\mathbb{W}(\rho_{0},\rho_{1})^{2}\coloneqq\inf_{\pi\in\Pi(\rho_{0},\rho_{1})}\iint_{\mathcal{Z}\times\mathcal{Z}}d(z_{0},z_{1})^{2}\,\,\pi(\mathrm{d}z_{0},\mathrm{d}z_{1}),$
(A.7)

where $\Pi(\rho_{0},\rho_{1})$ is the set of all probability distributions on
$\mathcal{Z}\times\mathcal{Z}$ whose marginals are $\rho_{0}$ and $\rho_{1}$,
respectively.

###  A.3 Proofs and Results related to Section 4

The maximum mean discrepancy (MMD) between a distribution
$\rho\in\mathcal{P}(\mathcal{Z})$ and the target distribution $\varrho$ is
defined as

$\mathrm{MMD}(\rho,\varrho)\coloneqq\sup_{f\in\mathcal{H}:\|f\|_{\mathcal{H}}\leq
1}{\Big{\\{}\int_{\mathcal{Z}}f(z)\,\rho(\mathrm{d}z)-\int_{\mathcal{Z}}f(z)\,\varrho(\mathrm{d}z)\Big{\\}}}.$

It is well-known that the $\mathrm{MMD}$ admits the following closed-form
formula [Gretton et al., 2012, Lemmas 4 and 6].

######  Lemma 10.

We have
$\mathrm{MMD}(\rho,\varrho)=\|\mathbf{m}_{\rho}-\mathbf{m}_{\varrho}\|_{\mathcal{H}}$.
As a consequence,

$\displaystyle\mathrm{MMD}(\rho,\varrho)^{2}=\int_{\mathcal{Z}}\int_{\mathcal{Z}}k(z,w)\rho(\mathrm{d}z)\rho(\mathrm{d}w)-2\int_{\mathcal{Z}}\mathbf{m}_{\varrho}(z)\rho(\mathrm{d}z)+\|\mathbf{m}_{\varrho}\|_{\mathcal{H}}^{2}.$

###### Proof of Lemma 10.

For any $f\in\mathcal{H}$, we have $f(z)=\langle
f,k(\cdot,z)\rangle_{\mathcal{H}}$. Therefore,

$\int_{\mathcal{Z}}f(z)\,\rho(\mathrm{d}z)=\Big{\langle}f,\int_{\mathcal{Z}}k(\cdot,z)\,\rho(\mathrm{d}z)\Big{\rangle}_{\mathcal{H}}=\langle
f,\mathbf{m}_{\rho}\rangle_{\mathcal{H}}\quad\mbox{for all}\quad
f\in\mathcal{H}.$ (A.8)

It follows that
$\mathrm{MMD}(\rho,\varrho)=\sup_{f\in\mathcal{H}:\|f\|_{\mathcal{H}}\leq
1}\langle
f,\mathbf{m}_{\rho}-\mathbf{m}_{\varrho}\rangle_{\mathcal{H}}=\|\mathbf{m}_{\rho}-\mathbf{m}_{\varrho}\|_{\mathcal{H}}$.
Using this closed-form formula and identity (A.8), we also obtain

$\displaystyle\mathrm{MMD}(\rho,\varrho)^{2}$
$\displaystyle=\|\mathbf{m}_{\rho}-\mathbf{m}_{\varrho}\|_{\mathcal{H}}^{2}=\langle\mathbf{m}_{\rho},\mathbf{m}_{\rho}\rangle_{\mathcal{H}}-2\langle\mathbf{m}_{\rho},\mathbf{m}_{\varrho}\rangle_{\mathcal{H}}+\langle\mathbf{m}_{\varrho},\mathbf{m}_{\varrho}\rangle_{\mathcal{H}}$
$\displaystyle=\int\mathbf{m}_{\rho}(z)\rho(\mathrm{d}z)-2\int\mathbf{m}_{\varrho}(z)\rho(\mathrm{d}z)+\|\mathbf{m}_{\varrho}\|_{\mathcal{H}}^{2}$
$\displaystyle=\iint
k(z,w)\rho(\mathrm{d}z)\rho(\mathrm{d}w)-2\int\mathbf{m}_{\varrho}(z)\rho(\mathrm{d}z)+\|\mathbf{m}_{\varrho}\|_{\mathcal{H}}^{2}.$

This completes the proof. ∎

#### For Lemma 5

###### Proof of Lemma 5.

We recall that $\mathrm{grad}\,{\mathcal{F}}[\rho]$ is defined as the unique
element in $\mathrm{T}_{\rho}{\mathcal{P}}(\mathcal{Z})$ satisfying

$g_{\rho}\Big{(}\mathrm{grad}\,{\mathcal{F}}[\rho],\partial_{t}\rho_{t}|_{t=0}\Big{)}=\frac{\mathrm{d}}{\mathrm{d}t}\Big{|}_{t=0}{\mathcal{F}}[\rho_{t}]$

for every differentiable curve $t\mapsto\rho_{t}\in{\mathcal{P}}(\mathcal{Z})$
passing through $\rho$ at $t=0$. Let
$t\mapsto\rho_{t}\in{\mathcal{P}}(\mathcal{Z})$ be such a curve. Then since
$\partial_{t}\rho_{t}|_{t=0}\in\mathrm{T}_{\rho}{\mathcal{P}}(\mathcal{Z})$,
we can write
$\partial_{t}\rho_{t}|_{t=0}=-\mathrm{div}_{d}(\rho\widetilde{\nabla}_{d}\varphi)$
for some differentiable function $\varphi$ on $\mathcal{Z}$. Then by using
Lemma 10 and $k(z,w)=k(w,z)$ we have

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}\Big{|}_{t=0}{\mathcal{F}}[\rho_{t}]$
$\displaystyle=\frac{1}{2}\frac{\mathrm{d}}{\mathrm{d}t}\Big{|}_{t=0}\left[\iint
k(z,w)\rho_{t}(\mathrm{d}z)\rho_{t}(\mathrm{d}w)-2\int\mathbf{m}_{\varrho}(z)\rho_{t}(\mathrm{d}z)\right]$
$\displaystyle=\frac{1}{2}\iint
k(z,w)\partial_{t}\rho_{t}|_{t=0}(\mathrm{d}z)\rho(\mathrm{d}w)+\frac{1}{2}\iint
k(z,w)\partial_{t}\rho_{t}|_{t=0}(\mathrm{d}w)\rho(\mathrm{d}z)$
$\displaystyle\quad-\int\mathbf{m}_{\varrho}(z)\,\partial_{t}\rho_{t}|_{t=0}(\mathrm{d}z)$
$\displaystyle=-\int_{\mathcal{Z}}\int_{\mathcal{Z}}k(z,w)\mathrm{div}_{d}(\rho\widetilde{\nabla}_{d}\varphi)(\mathrm{d}z)\,\rho(\mathrm{d}w)+\int_{\mathcal{Z}}\mathbf{m}_{\varrho}(z)\,\mathrm{div}_{d}(\rho\widetilde{\nabla}_{d}\varphi)(\mathrm{d}z).$

Let $\widetilde{\nabla}_{d}^{1}k(z,w)$ denote the gradient
$\widetilde{\nabla}_{d}$ of the function $z\mapsto k(z,w)$. It then follows
from the definition of the divergence operator
$\mathrm{div}_{d}(\rho\widetilde{\nabla}_{d}\varphi)$ at the end of Section
3.1 that

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}\Big{|}_{t=0}{\mathcal{F}}[\rho_{t}]$
$\displaystyle=\int_{\mathcal{Z}}\int_{\mathcal{Z}}\langle\widetilde{\nabla}_{d}^{1}k(z,w),\widetilde{\nabla}_{d}\varphi(z)\rangle_{z}\,\rho(\mathrm{d}z)\rho(\mathrm{d}w)-\int_{\mathcal{Z}}\langle\widetilde{\nabla}_{d}\mathbf{m}_{\varrho}(z),\widetilde{\nabla}_{d}\varphi(z)\rangle_{z}\rho(\mathrm{d}z)$
$\displaystyle=\int\left[\Big{\langle}\int\widetilde{\nabla}_{d}^{1}k(z,w)\,\rho(\mathrm{d}w),\widetilde{\nabla}_{d}\varphi(z)\Big{\rangle}_{z}\right]\rho(\mathrm{d}z)-\int_{\mathcal{Z}}\langle\widetilde{\nabla}_{d}\mathbf{m}_{\varrho}(z),\widetilde{\nabla}_{d}\varphi(z)\rangle_{z}\rho(\mathrm{d}z)$
$\displaystyle=\int\left[\Big{\langle}\widetilde{\nabla}_{d}\int
k(z,w)\,\rho(\mathrm{d}w),\widetilde{\nabla}_{d}\varphi(z)\Big{\rangle}_{z}\right]\rho(\mathrm{d}z)-\int_{\mathcal{Z}}\langle\widetilde{\nabla}_{d}\mathbf{m}_{\varrho}(z),\widetilde{\nabla}_{d}\varphi(z)\rangle_{z}\rho(\mathrm{d}z)$
$\displaystyle=\int_{\mathcal{Z}}\langle\widetilde{\nabla}_{d}[\mathbf{m}_{\rho}-\mathbf{m}_{\varrho}](z),\widetilde{\nabla}_{d}\varphi(z)\rangle_{z}\rho(\mathrm{d}z).$

By the definition of the Riemannian metric tensor $g_{\rho}$ given in (A.3)
and due to
$\partial_{t}\rho_{t}|_{t=0}=-\mathrm{div}_{d}(\rho\widetilde{\nabla}_{d}\varphi)$,
we thus obtain

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}\Big{|}_{t=0}{\mathcal{F}}[\rho_{t}]$
$\displaystyle=g_{\rho}\Big{(}-\mathrm{div}_{d}\big{(}\rho\widetilde{\nabla}_{d}[\mathbf{m}_{\rho}-\mathbf{m}_{\varrho}]\big{)},-\mathrm{div}_{d}(\rho\widetilde{\nabla}_{d}\varphi)\Big{)}$
$\displaystyle=g_{\rho}\Big{(}-\mathrm{div}_{d}\big{(}\rho\widetilde{\nabla}_{d}[\mathbf{m}_{\rho}-\mathbf{m}_{\varrho}]\big{)},\partial_{t}\rho_{t}|_{t=0}\Big{)}.$

Therefore, we infer that
$\mathrm{grad}\,{\mathcal{F}}[\rho]=-\mathrm{div}_{d}\big{(}\rho\widetilde{\nabla}_{d}[\mathbf{m}_{\rho}-\mathbf{m}_{\varrho}]\big{)}$
as desired. ∎

#### For Proposition 6

###### Proof of Proposition 6.

The proof is similar to that of Lemma 5 and with the same notation for
$\widetilde{\nabla}_{d}^{1}k(z,w)$. Indeed, by the same computation at the
beginning of the proof of Lemma 5 we have

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}{\mathcal{F}}[\rho_{t}]$
$\displaystyle=\frac{1}{2}\iint
k(z,w)\partial_{t}\rho_{t}(\mathrm{d}z)\rho_{t}(\mathrm{d}w)+\frac{1}{2}\iint
k(z,w)\partial_{t}\rho_{t}(\mathrm{d}w)\rho_{t}(\mathrm{d}z)-\int\mathbf{m}_{\varrho}(z)\,\partial_{t}\rho_{t}(\mathrm{d}z)$
$\displaystyle=\iint
k(z,w)\partial_{t}\rho_{t}(\mathrm{d}z)\rho_{t}(\mathrm{d}w)-\int\mathbf{m}_{\varrho}(z)\,\partial_{t}\rho_{t}(\mathrm{d}z).$

This together with the gradient flow equation (4.1) gives

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}{\mathcal{F}}[\rho_{t}]$
$\displaystyle=\int_{\mathcal{Z}}\int_{\mathcal{Z}}k(z,w)\mathrm{div}_{d}\big{(}\rho_{t}\widetilde{\nabla}_{d}[\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}]\big{)}(\mathrm{d}z)\rho_{t}(\mathrm{d}w)$
$\displaystyle\quad-\int_{\mathcal{Z}}\mathbf{m}_{\varrho}(z)\,\mathrm{div}_{d}\big{(}\rho_{t}\widetilde{\nabla}_{d}[\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}]\big{)}(\mathrm{d}z).$

Using the definition of the divergence operator $\mathrm{div}_{d}$ at the end
of Section 3.1, we further obtain

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}{\mathcal{F}}[\rho_{t}]$
$\displaystyle=-\int_{\mathcal{Z}}\int_{\mathcal{Z}}\langle\widetilde{\nabla}_{d}^{1}k(z,w),\widetilde{\nabla}_{d}[\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}](z)\rangle_{z}\,\rho_{t}(\mathrm{d}z)\rho_{t}(\mathrm{d}w)$
$\displaystyle\quad+\int_{\mathcal{Z}}\langle\widetilde{\nabla}_{d}\mathbf{m}_{\varrho}(z),\widetilde{\nabla}_{d}[\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}](z)\rangle_{z}\rho_{t}(\mathrm{d}z)$
$\displaystyle=-\int\left[\Big{\langle}\int\widetilde{\nabla}_{d}^{1}k(z,w)\,\rho_{t}(\mathrm{d}w),\widetilde{\nabla}_{d}[\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}](z)\Big{\rangle}_{z}\right]\rho_{t}(\mathrm{d}z)$
$\displaystyle\quad+\int_{\mathcal{Z}}\langle\widetilde{\nabla}_{d}\mathbf{m}_{\varrho}(z),\widetilde{\nabla}_{d}[\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}](z)\rangle_{z}\rho_{t}(\mathrm{d}z)$
$\displaystyle=-\int_{\mathcal{Z}}\langle\widetilde{\nabla}_{d}\mathbf{m}_{\rho_{t}}(z),\widetilde{\nabla}_{d}[\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}](z)\rangle_{z}\rho_{t}(\mathrm{d}z)$
$\displaystyle\quad+\int_{\mathcal{Z}}\langle\widetilde{\nabla}_{d}\mathbf{m}_{\varrho}(z),\widetilde{\nabla}_{d}[\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}](z)\rangle_{z}\rho_{t}(\mathrm{d}z)$
$\displaystyle=-\int_{\mathcal{Z}}\|\widetilde{\nabla}_{d}[\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}](z)\|_{z}^{2}\rho_{t}(\mathrm{d}z).$

This yields the desired result. ∎

#### For Lemma 7

###### Proof of Lemma 7.

From the formula for $\exp_{z}(\varepsilon\Phi(z))$ given at the beginning of
Section 4.1, we observe that

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}\exp_{z}(\varepsilon\Phi(z))$
$\displaystyle=\Big{(}\Phi_{1}(z),\Phi_{2}(z),{\mathrm{L}}_{\Sigma}[\Phi_{3}(z)]\,\Sigma+\Sigma\,{\mathrm{L}}_{\Sigma}[\Phi_{3}(z)]\Big{)}$
$\displaystyle=\Big{(}\Phi_{1}(z),\Phi_{2}(z),\Phi_{3}(z)\Big{)}=\Phi(z),$
(A.9)

where the second equality is due to the definition of
${\mathrm{L}}_{\Sigma}[V]$ given at the beginning of Section 3.1.

We obtain from Lemma 10 that

$\displaystyle\mathrm{MMD}(\exp(\varepsilon\Phi)_{\\#}\rho^{\tau},\varrho)^{2}$
$\displaystyle=\iint
k(z,w)\exp(\varepsilon\Phi)_{\\#}\rho^{\tau}(\mathrm{d}z)\exp(\varepsilon\Phi)_{\\#}\rho^{\tau}(\mathrm{d}w)-2\int\mathbf{m}_{\varrho}(z)\exp(\varepsilon\Phi)_{\\#}\rho^{\tau}(\mathrm{d}z)+\|\mathbf{m}_{\varrho}\|_{\mathcal{H}}^{2}$
$\displaystyle=\iint
k\Big{(}\exp_{z}(\varepsilon\Phi(z)),\exp_{w}(\varepsilon\Phi(w))\Big{)}\rho^{\tau}(\mathrm{d}z)\rho^{\tau}(\mathrm{d}w)-2\int\mathbf{m}_{\varrho}\Big{(}\exp_{z}(\varepsilon\Phi(z))\Big{)}\rho^{\tau}(\mathrm{d}z)+\|\mathbf{m}_{\varrho}\|_{\mathcal{H}}^{2}.$

Moreover, we have

$\displaystyle\iint\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}\Big{[}k(\exp_{z}(\varepsilon\Phi(z)),\exp_{w}(\varepsilon\Phi(w)))\Big{]}\rho^{\tau}(\mathrm{d}z)\rho^{\tau}(\mathrm{d}w)$
$\displaystyle=\iint\left\\{\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}\Big{[}k(\exp_{z}(\varepsilon\Phi(z)),w)\Big{]}+\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}\Big{[}k(z,\exp_{w}(\varepsilon\Phi(w)))\Big{]}\right\\}\rho^{\tau}(\mathrm{d}z)\rho^{\tau}(\mathrm{d}w)$
$\displaystyle=\int\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}\Big{[}\int
k(\exp_{z}(\varepsilon\Phi(z)),w)\rho^{\tau}(\mathrm{d}w)\Big{]}\rho^{\tau}(\mathrm{d}z)+\int\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}\Big{[}\int
k(z,\exp_{w}(\varepsilon\Phi(w)))\rho^{\tau}(\mathrm{d}z)\Big{]}\rho^{\tau}(\mathrm{d}w)$
$\displaystyle=\int\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}\Big{[}\mathbf{m}_{\rho^{\tau}}(\exp_{z}(\varepsilon\Phi(z)))\Big{]}\rho^{\tau}(\mathrm{d}z)+\int\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}\Big{[}\mathbf{m}_{\rho^{\tau}}(\exp_{w}(\varepsilon\Phi(w)))\Big{]}\rho^{\tau}(\mathrm{d}w)$
$\displaystyle=2\int\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}\Big{[}\mathbf{m}_{\rho^{\tau}}(\exp_{z}(\varepsilon\Phi(z)))\Big{]}\rho^{\tau}(\mathrm{d}z).$

Thus, it follows that

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}\mathrm{MMD}(\exp(\varepsilon\Phi)_{\\#}\rho^{\tau},\varrho)^{2}$
$\displaystyle=2\int\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}\Big{[}\mathbf{m}_{\rho^{\tau}}(\exp_{z}(\varepsilon\Phi(z)))\Big{]}\rho^{\tau}(\mathrm{d}z)-2\int\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}\Big{[}\mathbf{m}_{\varrho}(\exp_{z}(\varepsilon\Phi(z)))\Big{]}\rho^{\tau}(\mathrm{d}z)$
$\displaystyle=2\int
D[\mathbf{m}_{\rho^{\tau}}-\mathbf{m}_{\varrho}]_{z}\Big{(}\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}\exp_{z}(\varepsilon\Phi(z))\Big{)}\rho^{\tau}(\mathrm{d}z)$

with $D\varphi_{z}(w,,v,V)$ denoting the standard directional derivative of
$\varphi$ at $z$ in the direction $(w,v,V)$. Using the definition of
${\mathcal{F}}$ together with (A.3) and the definition of gradient
$\widetilde{\nabla}_{d}$ in (A.5), we obtain

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}{\mathcal{F}}[\exp(\varepsilon\Phi)_{\\#}\rho^{\tau}]$
$\displaystyle=\int
D[\mathbf{m}_{\rho^{\tau}}-\mathbf{m}_{\varrho}]_{z}(\Phi(z))\,\rho^{\tau}(\mathrm{d}z)$
$\displaystyle=\int\langle\widetilde{\nabla}_{d}[\mathbf{m}_{\rho^{\tau}}-\mathbf{m}_{\varrho}](z),\Phi(z)\rangle_{z}\rho^{\tau}(\mathrm{d}z).$
(A.10)

This yields the first conclusion of the lemma.

Now let
$\hat{\Phi}^{\tau}\coloneqq\frac{\Phi^{\tau}}{\|\Phi^{\tau}\|_{\mathbb{L}^{2}(\rho^{\tau})}}$
be the unit vector field in the direction of
$\Phi^{\tau}\coloneqq-\widetilde{\nabla}_{d}[\mathbf{m}_{\rho^{\tau}}-\mathbf{m}_{\varrho}]$.
Then by (A.3), we have

$\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}{\mathcal{F}}[\exp(\varepsilon\hat{\Phi}^{\tau})_{\\#}\rho^{\tau}]=-\|\Phi^{\tau}\|_{\mathbb{L}^{2}(\rho^{\tau})}^{-1}\int\|\Phi^{\tau}(z)\|_{z}^{2}\rho^{\tau}(\mathrm{d}z)=-\|\Phi^{\tau}\|_{\mathbb{L}^{2}(\rho^{\tau})}\leq
0.$

On the other hand, for any unit direction $\Phi$ in
$\mathbb{L}^{2}(\rho^{\tau})$ we obtain from (A.3) and Hölder inequality that

$\displaystyle\Big{|}\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}$
$\displaystyle{\mathcal{F}}[\exp(\varepsilon\Phi)_{\\#}\rho^{\tau}]\Big{|}\leq\int\|\Phi^{\tau}(z)\|_{z}\|\Phi(z)\|_{z}\rho^{\tau}(\mathrm{d}z)$
$\displaystyle\leq\Big{(}\int\|\Phi^{\tau}(z)\|_{z}^{2}\rho^{\tau}(\mathrm{d}z)\Big{)}^{\frac{1}{2}}\Big{(}\int\|\Phi(z)\|_{z}^{2}\rho^{\tau}(\mathrm{d}z)\Big{)}^{\frac{1}{2}}=\|\Phi^{\tau}\|_{\mathbb{L}^{2}(\rho^{\tau})}.$

Therefore, we conclude further that

$\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}{\mathcal{F}}[\exp(\varepsilon\hat{\Phi}^{\tau})_{\\#}\rho^{\tau}]\leq\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}{\mathcal{F}}[\exp(\varepsilon\Phi)_{\\#}\rho^{\tau}]$

for any unit direction $\Phi$ in $\mathbb{L}^{2}(\rho^{\tau})$. These give the
last conclusion of the lemma. ∎

######  Definition 11 (Lipschitz-gradient kernel).

Let $L>0$. A differentiable kernel $k$ on $\mathcal{Z}$ is called a Lipschitz-
gradient kernel with constant $L$ if there exists a number
$\varepsilon_{0}\in(0,1)$ such that

$\displaystyle\Big{|}k(\exp_{z}(\varepsilon\Phi(z)),\exp_{w}(\delta\Phi(w)))$
$\displaystyle-k(z,w)-\big{[}\langle\widetilde{\nabla}_{d}^{1}k(z,w),\varepsilon\Phi(z)\rangle_{z}+\langle\widetilde{\nabla}_{d}^{2}k(z,w),\delta\Phi(w)\rangle_{w}\big{]}\Big{|}$
$\displaystyle\leq
L\Big{[}\|\varepsilon\Phi(z)\|_{z}^{2}+\|\delta\Phi(w)\|_{w}^{2}\Big{]}$
(A.11)

for every $\varepsilon,\,\delta\in[0,\varepsilon_{0}]$ and every bounded
vector field
$\Phi:\mathcal{Z}\to{\mathbb{R}}^{m}\times{\mathbb{R}}^{n}\times{\mathbb{S}}^{n}$.
Hereafter, $\widetilde{\nabla}_{d}^{1}k(z,w)$ and
$\widetilde{\nabla}_{d}^{2}k(z,w)$ denote respectively the gradient
$\widetilde{\nabla}_{d}$ of the function $z\mapsto k(z,w)$ and the function
$w\mapsto k(z,w)$.

######  Remark 12.

The right hand side of condition (11) can be expressed in terms of the $d$
distance as

$d\big{(}\exp_{z}(\varepsilon\Phi(z)),z\big{)}^{2}+d\big{(}\exp_{w}(\delta\Phi(w)),w\big{)}^{2}.$

Thus condition (11) can be interpreted as the gradient
$\widetilde{\nabla}_{d}k$ is Lipschitz w.r.t. the distance $d$.

Condition (11) is motivated by the following observation in the Euclidean
space. Assume that $G:{\mathbb{R}}^{d}\times{\mathbb{R}}^{d}\to{\mathbb{R}}$
is a differentiable function such that its Euclidean gradient $\nabla
G(z,w)\coloneqq(\nabla^{1}G(z,w),\nabla^{2}G(z,w))$ satisfies the standard
Lipschitz condition

$\|\nabla G(z_{1},w_{1})-\nabla G(z_{2},w_{2})\|_{2}\leq
L\|(z_{1},w_{1})-(z_{2},w_{2})\|_{2}\quad\forall(z_{1},w_{1}),\,(z_{2},w_{2})\in{\mathbb{R}}^{d}\times{\mathbb{R}}^{d}.$

Then for any point $(z,w)\in{\mathbb{R}}^{d}\times{\mathbb{R}}^{d}$ and any
tangent vector $(u,v)\in{\mathbb{R}}^{d}\times{\mathbb{R}}^{d}$, we have by
using the mean value theorem that $G(z+u,w+v)-G(z,w)=\langle\nabla
G(z_{0},w_{0}),(u,v)\rangle$ for some point $(z_{0},w_{0})$ in the line
segment in ${\mathbb{R}}^{d}\times{\mathbb{R}}^{d}$ connecting the points
$(z,w)$ and $(z+u,w+v)$. As a consequence, we obtain

$\displaystyle\Big{|}G(z+u,w+v)-G(z,w)-\big{[}\langle\nabla^{1}G(z,w),u\rangle+\langle\nabla^{2}G(z,w),v\rangle\big{]}\Big{|}$
$\displaystyle=\Big{|}\langle\nabla G(z_{0},w_{0}),(u,v)\rangle-\langle\nabla
G(z,w),(u,v)\rangle\Big{|}$ $\displaystyle=\Big{|}\big{\langle}\nabla
G(z_{0},w_{0})-\nabla G(z,w),(u,v)\big{\rangle}\Big{|}\leq\|\nabla
G(z_{0},w_{0})-\nabla G(z,w)\|_{2}\|(u,v)\|_{2}.$

Then we can use the Lipschitz condition for $\nabla G$ to imply further that

$\displaystyle\Big{|}G(z+u,w+v)$
$\displaystyle-G(z,w)-\big{[}\langle\nabla^{1}G(z,w),u\rangle+\langle\nabla^{2}G(z,w),v\rangle\big{]}\Big{|}$
$\displaystyle\leq L\|(z_{0},w_{0})-(z,w)\|_{2}\|(u,v)\|_{2}\leq
L\|(u,v)\|_{2}^{2},$

which is the same as

$\displaystyle\Big{|}G(z+u,w+v)-G(z,w)-\big{[}\langle\nabla^{1}G(z,w),u\rangle+\langle\nabla^{2}G(z,w),v\rangle\big{]}\Big{|}\leq
L\big{[}\|u\|_{2}^{2}+\|v\|_{2}^{2}\big{]}.$

Condition (11) is the Riemannian version of this last inequality for the
Euclidean space, which is a consequence of the standard Lipschitz condition
for the gradient.

#### Bounded Hessian kernels are Lipschitz-gradient.

The following lemma gives a sufficient condition for a kernel to be Lipschitz-
gradient.

######  Lemma 13.

Let $k$ be a positive definite kernel such that its Hessian w.r.t. distance
$d$ is bounded. Then $k$ is a Lipschitz-gradient kernel.

###### Proof of Lemma 13.

Let $H^{1}_{d}k(z,w)$ and $H^{2}_{d}k(z,w)$ denote respectively the Hessian
w.r.t. distance $d$ of the function $z\mapsto k(z,w)$ and the function
$w\mapsto k(z,w)$. Let $\varepsilon,\,\delta>0$, and
$\Phi:\mathcal{Z}\to{\mathbb{R}}^{m}\times{\mathbb{R}}^{n}\times{\mathbb{S}}^{n}$
be a bounded vector field. Define $\gamma_{z}(t)\coloneqq\exp_{z}(t\Phi(z))$
and $\theta_{w}(t)\coloneqq\exp_{w}(\frac{\delta}{\varepsilon}t\Phi(w))$ for
$t\in[0,\varepsilon]$. Then we have

$\displaystyle
k(\exp_{z}(\varepsilon\Phi(z)),\exp_{w}(\delta\Phi(w)))-k(z,w)=\int_{0}^{\varepsilon}\frac{\mathrm{d}}{\mathrm{d}t}[k(\gamma_{z}(t),\theta_{w}(t))]\,\mathrm{d}t$
$\displaystyle=\int_{0}^{\varepsilon}\Big{[}\langle\widetilde{\nabla}^{1}_{d}k(\gamma_{z}(t),\theta_{w}(t)),\dot{\gamma}_{z}(t)\rangle_{\gamma_{z}(t)}+\langle\widetilde{\nabla}^{2}_{d}k(\gamma_{z}(t),\theta_{w}(t)),\dot{\theta}_{w}(t)\rangle_{\theta_{w}(t)}]\Big{]}\mathrm{d}t.$

This together with the facts that $\dot{\gamma}_{z}(0)=\Phi(z)$ and
$\dot{\theta}_{w}(0)=\frac{\delta}{\varepsilon}\Phi(z)$ yields

$\displaystyle A$ $\displaystyle\coloneqq
k(\exp_{z}(\varepsilon\Phi(z)),\exp_{w}(\delta\Phi(w)))-k(z,w)-\big{[}\langle\widetilde{\nabla}_{d}^{1}k(z,w),\varepsilon\Phi(z)\rangle_{z}+\langle\widetilde{\nabla}_{d}^{2}k(z,w),\delta\Phi(w)\rangle_{w}\big{]}$
$\displaystyle=\int_{0}^{\varepsilon}\Big{[}\langle\widetilde{\nabla}^{1}_{d}k(\gamma_{z}(t),\theta_{w}(t)),\dot{\gamma}_{z}(t)\rangle_{\gamma_{z}(t)}-\langle\widetilde{\nabla}^{1}_{d}k(\gamma_{z}(0),\theta_{w}(0)),\dot{\gamma}_{z}(0)\rangle_{\gamma_{z}(0)}\Big{]}dt$
$\displaystyle\quad+\int_{0}^{\varepsilon}\Big{[}\langle\widetilde{\nabla}^{2}_{d}k(\gamma_{z}(t),\theta_{w}(t)),\dot{\theta}_{w}(t)\rangle_{\theta_{w}(t)}-\langle\widetilde{\nabla}^{2}_{d}k(\gamma_{z}(0),\theta_{w}(0)),\dot{\theta}_{w}(0)\rangle_{\theta_{w}(0)}]\Big{]}\mathrm{d}t$
$\displaystyle=\int_{0}^{\varepsilon}\int_{0}^{t}\frac{\mathrm{d}}{\mathrm{d}s}\Big{[}\langle\widetilde{\nabla}^{1}_{d}k(\gamma_{z}(s),\theta_{w}(s)),\dot{\gamma}_{z}(s)\rangle_{\gamma_{z}(s)}\Big{]}\mathrm{d}s\mathrm{d}t$
$\displaystyle\quad+\int_{0}^{\varepsilon}\int_{0}^{t}\frac{\mathrm{d}}{\mathrm{d}s}\Big{[}\langle\widetilde{\nabla}^{2}_{d}k(\gamma_{z}(s),\theta_{w}(s)),\dot{\theta}_{w}(s)\rangle_{\theta_{w}(s)}\Big{]}\mathrm{d}s\mathrm{d}t$
$\displaystyle=\int_{0}^{\varepsilon}\int_{0}^{t}\Big{[}\langle
H^{1}_{d}k(\gamma_{z}(s),\theta_{w}(s))\dot{\gamma}_{z}(s),\dot{\gamma}_{z}(s)\rangle_{\gamma_{z}(s)}+\langle\widetilde{\nabla}^{1}_{d}k(\gamma_{z}(s),\theta_{w}(s)),\ddot{\gamma}_{z}(s)\rangle_{\gamma_{z}(s)}\Big{]}\mathrm{d}s\mathrm{d}t$
$\displaystyle\quad+\int_{0}^{\varepsilon}\int_{0}^{t}\Big{[}\langle
H^{2}_{d}k(\gamma_{z}(s),\theta_{w}(s))\dot{\theta}_{w}(s),\dot{\theta}_{w}(s)\rangle_{\theta_{w}(s)}+\langle\widetilde{\nabla}^{2}_{d}k(\gamma_{z}(s),\theta_{w}(s)),\ddot{\theta}_{w}(s)\rangle_{\theta_{w}(s)}\Big{]}\mathrm{d}s\mathrm{d}t.$

Since the curve $s\mapsto\gamma_{z}(s)$ is a geodesic, its acceleration
$\ddot{\gamma}_{z}(s)$ is orthogonal to $\mathcal{Z}$ (that is,
$\ddot{\gamma}_{z}(s)$ is orthogonal to every tangent vector in
$\mathrm{T}_{z}\mathcal{Z}$). This implies that
$\langle\widetilde{\nabla}^{1}_{d}k(\gamma_{z}(s),\theta_{w}(s)),\ddot{\gamma}_{z}(s)\rangle_{\gamma_{z}(s)}=0$.
Likewise, we also have
$\langle\widetilde{\nabla}^{2}_{d}k(\gamma_{z}(s),\theta_{w}(s)),\ddot{\theta}_{w}(s)\rangle_{\theta_{w}(s)}=0$.
Thanks to these, we deduce from the above identity that

$\displaystyle A$ $\displaystyle=\int_{0}^{\varepsilon}\int_{0}^{t}\langle
H^{1}_{d}k(\gamma_{z}(s),\theta_{w}(s))\dot{\gamma}_{z}(s),\dot{\gamma}_{z}(s)\rangle_{\gamma_{z}(s)}\mathrm{d}s\mathrm{d}t$
$\displaystyle\quad+\int_{0}^{\varepsilon}\int_{0}^{t}\langle
H^{2}_{d}k(\gamma_{z}(s),\theta_{w}(s))\dot{\theta}_{w}(s),\dot{\theta}_{w}(s)\rangle_{\theta_{w}(s)}\mathrm{d}s\mathrm{d}t.$

By using the assumption that the Hessians $H^{1}_{d}$ and $H^{2}_{d}$ are
bounded, we then obtain

$\displaystyle|A|$ $\displaystyle\leq
M\int_{0}^{\varepsilon}\int_{0}^{t}\Big{[}\|\dot{\gamma}_{z}(s)\|_{\gamma_{z}(s)}^{2}+\|\dot{\theta}_{w}(s)\|_{\theta_{w}(s)}^{2}\Big{]}\mathrm{d}s\mathrm{d}t,$

where $M$ is the sup norm of the Hessian of $k$. But as $\gamma_{z}(s)$ and
$\theta_{w}(s)$ are geodesic, they have constant speeds. Therefore,
$\|\dot{\gamma}_{z}(s)\|_{\gamma_{z}(s)}=\|\dot{\gamma}_{z}(0)\|_{\gamma_{z}(0)}=\|\Phi(z)\|_{z}$
and
$\|\dot{\theta}_{w}(s)\|_{\theta_{w}(s)}=\|\dot{\theta}_{w}(0)\|_{\theta_{w}(0)}=\|\frac{\delta}{\varepsilon}\Phi(z)\|_{z}$.
Using these, we infer further that

$\displaystyle|A|$ $\displaystyle\leq
M\int_{0}^{\varepsilon}\int_{0}^{t}\Big{[}\|\Phi(z)\|_{z}^{2}+(\frac{\delta}{\varepsilon})^{2}\|\Phi(z)\|_{z}^{2}\Big{]}\mathrm{d}s\mathrm{d}t=\frac{M}{2}\Big{[}\|\varepsilon\Phi(z)\|_{z}^{2}+\|\delta\Phi(w)\|_{w}^{2}\Big{]}.$

According to Definition 11, we thus conclude that $k$ is a Lipschitz-gradient
kernel with constant $M/2$. ∎

#### Quantified estimate of decrease for the Riemannian forward Euler scheme
(4.2).

The next result quantifies the amount that the value of ${\mathcal{F}}$
decreases after each iteration.

######  Proposition 14 (Quantified estimate of decrease).

Suppose that $k$ is a Lipschitz-gradient kernel with constant $L$. Then for
$\rho^{\tau+1}$ given by (4.2) with $s_{\tau}\in(0,\varepsilon_{0}]$, we have

${\mathcal{F}}[\rho^{\tau+1}]-{\mathcal{F}}[\rho^{\tau}]\leq-
s_{\tau}\big{(}1-2Ls_{\tau}\big{)}\int_{\mathcal{Z}}\|\widetilde{\nabla}_{d}[\mathbf{m}_{\rho^{\tau}}-\mathbf{m}_{\varrho}](z)\|_{z}^{2}\,\rho^{\tau}(\mathrm{d}z).$

###### Proof of Proposition 14.

Let
$\Phi^{\tau}\coloneqq-\widetilde{\nabla}_{d}[\mathbf{m}_{\rho^{\tau}}-\mathbf{m}_{\varrho}]$.
Then from the computation at the beginning of the proof of Lemma 7 and by
using Lemma 10, we obtain

$\displaystyle{\mathcal{F}}[\rho^{\tau+1}]-{\mathcal{F}}[\rho^{\tau}]$
$\displaystyle=\frac{1}{2}\left[\mathrm{MMD}(\exp(s_{\tau}\Phi^{\tau})_{\\#}\rho^{\tau},\varrho)^{2}-\mathrm{MMD}(\rho^{\tau},\varrho)^{2}\right]$
$\displaystyle=\frac{1}{2}\iint\Big{\\{}k\Big{(}\exp_{z}(s_{\tau}\Phi^{\tau}(z)),\exp_{w}(s_{\tau}\Phi^{\tau}(w))\Big{)}-k(z,w)\Big{\\}}\rho^{\tau}(\mathrm{d}z)\rho^{\tau}(\mathrm{d}w)$
$\displaystyle\quad-\iint\Big{\\{}k\Big{(}\exp_{z}(s_{\tau}\Phi^{\tau}(z)),w\Big{)}-k(z,w)\Big{\\}}\rho^{\tau}(\mathrm{d}z)\varrho(\mathrm{d}w).$

Moreover, we have

$\displaystyle\int\langle\widetilde{\nabla}_{d}[\mathbf{m}_{\rho^{\tau}}-\mathbf{m}_{\varrho}](z),\Phi^{\tau}(z)\rangle_{z}\rho^{\tau}(\mathrm{d}z)$
$\displaystyle=\iint\langle\widetilde{\nabla}_{d}^{1}k(z,w),\Phi^{\tau}(z)\rangle_{z}\rho^{\tau}(\mathrm{d}z)\rho^{\tau}(\mathrm{d}w)-\iint\langle\widetilde{\nabla}_{d}^{1}k(z,w),\Phi^{\tau}(z)\rangle_{z}\rho^{\tau}(\mathrm{d}z)\varrho(\mathrm{d}w)$
$\displaystyle=\frac{1}{2}\left[\iint\langle\widetilde{\nabla}_{d}^{1}k(z,w),\Phi^{\tau}(z)\rangle_{z}\rho^{\tau}(\mathrm{d}z)\rho^{\tau}(\mathrm{d}w)+\iint\langle\widetilde{\nabla}_{d}^{2}k(z,w),\Phi^{\tau}(w)\rangle_{w}\rho^{\tau}(\mathrm{d}z)\rho^{\tau}(\mathrm{d}w)\right]$
$\displaystyle\quad-\iint\langle\widetilde{\nabla}_{d}^{1}k(z,w),\Phi^{\tau}(z)\rangle_{z}\rho^{\tau}(\mathrm{d}z)\varrho(\mathrm{d}w),$

where the last equality is due to the symmetry of $k$ and relation (3.6). Here
$\widetilde{\nabla}_{d}^{1}k(z,w)$ and $\widetilde{\nabla}_{d}^{2}k(z,w)$
respectively denote the gradient $\widetilde{\nabla}_{d}$ of the function
$z\mapsto k(z,w)$ and $w\mapsto k(z,w)$. Therefore, it follows that

$\displaystyle{\mathcal{F}}[\rho^{\tau+1}]$
$\displaystyle-{\mathcal{F}}[\rho^{\tau}]-s_{\tau}\int\langle\widetilde{\nabla}_{d}[\mathbf{m}_{\rho^{\tau}}-\mathbf{m}_{\varrho}](z),\Phi^{\tau}(z)\rangle_{z}\rho^{\tau}(\mathrm{d}z)$
$\displaystyle=\frac{1}{2}\iint\left\\{k\Big{(}\exp_{z}(s_{\tau}\Phi^{\tau}(z)),\exp_{w}(s_{\tau}\Phi^{\tau}(w))\Big{)}-k(z,w)\right.$
$\displaystyle\qquad\qquad\left.-\Big{[}\langle\widetilde{\nabla}_{d}^{1}k(z,w),s_{\tau}\Phi^{\tau}(z)\rangle_{z}+\langle\widetilde{\nabla}_{d}^{2}k(z,w),s_{\tau}\Phi^{\tau}(w)\rangle_{w}\Big{]}\right\\}\rho^{\tau}(\mathrm{d}z)\rho^{\tau}(\mathrm{d}w)$
$\displaystyle\quad-\iint\Big{\\{}k\Big{(}\exp_{z}(s_{\tau}\Phi^{\tau}(z)),w\Big{)}-k(z,w)-\langle\widetilde{\nabla}_{d}^{1}k(z,w),s_{\tau}\Phi^{\tau}(z)\rangle_{z}\Big{\\}}\rho^{\tau}(\mathrm{d}z)\varrho(\mathrm{d}w).$

As $s_{\tau}\in(0,\varepsilon_{0}]$, we can now use the assumption that $k$ is
a Lipschitz-gradient kernel with constant $L$ to obtain

$\displaystyle{\mathcal{F}}[\rho^{\tau+1}]-{\mathcal{F}}[\rho^{\tau}]+s_{\tau}\int_{\mathcal{Z}}\|\Phi^{\tau}(z)\|_{z}^{2}\rho^{\tau}(\mathrm{d}z)$
$\displaystyle\leq\frac{L}{2}\iint\Big{[}\|s_{\tau}\Phi^{\tau}(z)\|_{z}^{2}+\|s_{\tau}\Phi^{\tau}(w)\|_{w}^{2}\Big{]}\rho^{\tau}(\mathrm{d}z)\rho^{\tau}(\mathrm{d}w)+L\iint\|s_{\tau}\Phi^{\tau}(z)\|_{z}^{2}\rho^{\tau}(\mathrm{d}z)\varrho(\mathrm{d}w)$
$\displaystyle=2Ls_{\tau}^{2}\int\|\Phi^{\tau}(z)\|_{z}^{2}\rho^{\tau}(\mathrm{d}z).$

This gives

${\mathcal{F}}[\rho^{\tau+1}]-{\mathcal{F}}[\rho^{\tau}]\leq\big{(}-s_{\tau}+2Ls_{\tau}^{2}\big{)}\int_{\mathcal{Z}}\|\Phi^{\tau}(z)\|_{z}^{2}\rho^{\tau}(\mathrm{d}z),$

and the conclusion of the proposition follows. ∎

#### Convergence guarantees.

For each distribution $\rho$ on $\mathcal{Z}$, let
$\mathbb{K}_{\rho}:\mathcal{H}\to\mathcal{H}$ be the linear operator defined
by
$\mathbb{K}_{\rho}f(w_{1})\coloneqq\langle\tilde{\mathbb{K}}_{\rho}(w_{1},\cdot),f(\cdot)\rangle_{\mathcal{H}}$
with $\tilde{\mathbb{K}}_{\rho}:\mathcal{Z}\times\mathcal{Z}\to{\mathbb{R}}$
being given by

$\tilde{\mathbb{K}}_{\rho}(w_{1},w_{2})=\int_{\mathcal{Z}}\langle\widetilde{\nabla}_{d}^{1}k(z,w_{1}),\widetilde{\nabla}_{d}^{1}k(z,w_{2})\rangle_{z}\,\rho(\mathrm{d}z)\quad\mbox{for}\quad
w_{1},\,w_{2}\in\mathcal{Z}.$

The next result gives some basic properties of the operator
$\mathbb{K}_{\rho}$.

######  Lemma 15.

For a differentiable kernel $k$ and for $\rho\in{\mathcal{P}}(\mathcal{Z})$,
we have

  1. i)

$\mathbb{K}_{\rho}f(w)=\int_{\mathcal{Z}}\langle\widetilde{\nabla}_{d}^{1}k(z,w),\widetilde{\nabla}_{d}f(z)\rangle_{z}\,\rho(\mathrm{d}z)$
for $f\in\mathcal{H}$.

  2. ii)

$\langle\mathbb{K}_{\rho}f,g\rangle_{\mathcal{H}}=\int_{\mathcal{Z}}\langle\widetilde{\nabla}_{d}f,\widetilde{\nabla}_{d}g\rangle_{z}\,\rho(\mathrm{d}z)$
for every $f,\,g\in\mathcal{H}$. Consequently, the operator
$\mathbb{K}_{\rho}$ is symmetric and positive, and hence its spectrum is
contained in $[0,+\infty)$.

###### Proof of Lemma 15.

By using the definition of the Riemannian metric
$\langle\cdot,\cdot\rangle_{z}$ given in (3), it can be verified for
$f\in\mathcal{H}$ that

$\Big{\langle}\langle\widetilde{\nabla}_{d}^{1}k(z,w),\widetilde{\nabla}_{d}^{1}k(z,\cdot)\rangle_{z},f(\cdot)\Big{\rangle}_{\mathcal{H}}=\Big{\langle}\widetilde{\nabla}_{d}^{1}k(z,w),\langle\widetilde{\nabla}_{d}^{1}k(z,\cdot),f(\cdot)\rangle_{\mathcal{H}}\Big{\rangle}_{z}.$

As $f(z)=\langle k(z,\cdot),f(\cdot)\rangle_{\mathcal{H}}$, we moreover have
$\widetilde{\nabla}_{d}f(z)=\langle\widetilde{\nabla}_{d}^{1}k(z,\cdot),f(\cdot)\rangle_{\mathcal{H}}$.
Therefore,

$\Big{\langle}\langle\widetilde{\nabla}_{d}^{1}k(z,w),\widetilde{\nabla}_{d}^{1}k(z,\cdot)\rangle_{z},f(\cdot)\Big{\rangle}_{\mathcal{H}}=\big{\langle}\widetilde{\nabla}_{d}^{1}k(z,w),\widetilde{\nabla}_{d}f(z)\big{\rangle}_{z}.$
(A.12)

Using the definition of $\mathbb{K}_{\rho}$ and (A.12), we obtain

$\displaystyle\mathbb{K}_{\rho}f(w)$
$\displaystyle=\Big{\langle}\int\langle\widetilde{\nabla}_{d}^{1}k(z,w),\widetilde{\nabla}_{d}^{1}k(z,\cdot)\rangle_{z}\,\rho(\mathrm{d}z),f(\cdot)\Big{\rangle}_{\mathcal{H}}$
$\displaystyle=\int\Big{\langle}\langle\widetilde{\nabla}_{d}^{1}k(z,w),\widetilde{\nabla}_{d}^{1}k(z,\cdot)\rangle_{z},f(\cdot)\Big{\rangle}_{\mathcal{H}}\,\rho(\mathrm{d}z)$
$\displaystyle=\int\big{\langle}\widetilde{\nabla}_{d}^{1}k(z,w),\widetilde{\nabla}_{d}f(z)\big{\rangle}_{z}\,\rho(\mathrm{d}z),$

which gives i). Now for $f,\,g\in\mathcal{H}$, we can use part i) and similar
arguments leading to (A.12) to obtain

$\displaystyle\langle\mathbb{K}_{\rho}f,g\rangle_{\mathcal{H}}$
$\displaystyle=\Big{\langle}\int\big{\langle}\widetilde{\nabla}_{d}^{1}k(z,\cdot),\widetilde{\nabla}_{d}f(z)\big{\rangle}_{z}\,\rho(\mathrm{d}z),g(\cdot)\Big{\rangle}_{\mathcal{H}}$
$\displaystyle=\int\Big{\langle}\langle\widetilde{\nabla}_{d}^{1}k(z,\cdot),\widetilde{\nabla}_{d}f(z)\rangle_{z},g(\cdot)\Big{\rangle}_{\mathcal{H}}\,\rho(\mathrm{d}z)$
$\displaystyle=\int\Big{\langle}\langle\widetilde{\nabla}_{d}^{1}k(z,\cdot),g(\cdot)\rangle_{\mathcal{H}},\widetilde{\nabla}_{d}f(z)\Big{\rangle}_{z}\,\rho(\mathrm{d}z)=\int\big{\langle}\widetilde{\nabla}_{d}g(z),\widetilde{\nabla}_{d}f(z)\big{\rangle}_{z}\,\rho(\mathrm{d}z).$

This implies in particular that the operator $\mathbb{K}_{\rho}$ is symmetric
(i.e.
$\langle\mathbb{K}_{\rho}f,g\rangle_{\mathcal{H}}=\langle\mathbb{K}_{\rho}g,f\rangle_{\mathcal{H}}$
for $f,\,g\in\mathcal{H}$) and positive (i.e.
$\langle\mathbb{K}_{\rho}f,f\rangle_{\mathcal{H}}\geq 0$ for
$f\in\mathcal{H}$). Since any symmetric, positive, and linear operator must
have nonnegative eigenvalues, we have completed the proof. ∎

Our next result gives a quantified decay rate for the objective function.

######  Proposition 16 (Objective value decay).

There hold:

  1. i)

Let $\rho_{t}$ be given by (4.1), and let $\lambda_{t}\geq 0$ be any constant
satisfying

$\langle\mathbb{K}_{\rho_{t}}f_{t},f_{t}\rangle_{\mathcal{H}}\geq\lambda_{t}\|f_{t}\|_{\mathcal{H}}^{2}\quad\mbox{with}\quad
f_{t}\coloneqq\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}.$ (A.13)

Then
${\mathcal{F}}[\rho_{t}]\leq{\mathcal{F}}[\rho_{0}]\exp{\big{(}-2\int_{0}^{t}\lambda_{s}\mathrm{d}s\big{)}}$
for any $t\geq 0$. In particular,
$\lim_{t\to\infty}\mathrm{MMD}(\rho_{t},\varrho)=0$ if
$\int_{0}^{\infty}\lambda_{t}\,\mathrm{d}t=+\infty$.

  2. ii)

Let $\rho^{\tau}$ be given by scheme (4.2), and $\lambda_{\tau}\geq 0$ be any
constant satisfying

$\langle\mathbb{K}_{\rho_{t}}f^{\tau},f^{\tau}\rangle_{\mathcal{H}}\geq\lambda_{\tau}\|f^{\tau}|_{\mathcal{H}}^{2}\quad\mbox{with}\quad
f^{\tau}\coloneqq\mathbf{m}_{\rho^{\tau}}-\mathbf{m}_{\varrho}.$

Assume that $k$ is a Lipschitz-gradient kernel and step size $s_{\tau}$
satisfies $s_{\tau}\lambda_{\tau}<1$, then we have
${\mathcal{F}}[\rho^{\tau+1}]\leq{\mathcal{F}}[\rho^{0}]\exp\big{(}-{\sum_{i=0}^{\tau}}s_{i}\lambda_{i}\big{)}$
for any $\tau\geq 0$. In particular,
$\lim_{\tau\to\infty}\mathrm{MMD}(\rho^{\tau},\varrho)=0$ if
$\sum_{\tau=0}^{\infty}s_{\tau}\lambda_{\tau}=+\infty$.

Condition $\int_{0}^{\infty}\lambda_{t}\,\mathrm{d}t=+\infty$ guaranteeing the
convergence in $\mathrm{MMD}$ holds true for example if $\lambda_{t}\geq
c\,t^{-1}$ for some constant $c>0$ and for large $t$. We note also that
Condition (A.13) is satisfied if $\lambda_{t}$ is chosen to be the minimum
eigenvalue of operator $\mathbb{K}_{\rho_{t}}$. Thus Proposition 16 implies in
particular that $\rho_{t}$ globally converges in $\mathrm{MMD}$ if the minimum
eigenvalue $\lambda_{t}$ of operator $\mathbb{K}_{\rho_{t}}$ satisfies the
integrability condition $\int_{0}^{\infty}\lambda_{t}\,\mathrm{d}t=+\infty$.
The proof of Proposition 16 relies on the following proposition, which shows
that the dynamic of the mean embedding is governed by the equation
$\partial_{t}(\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho})=-\mathbb{K}_{\rho_{t}}(\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho})$.

######  Proposition 17 (Dynamic of the mean embedding).

Let $t\in[0,\infty)\longmapsto\rho_{t}$ be the gradient flow given by equation
(4.1). For each $t\geq 0$, take
$f_{t}\coloneqq\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}$. Then $f_{t}$ is a
solution of the linear partial differential equation

$\partial_{t}f_{t}=-\mathbb{K}_{\rho_{t}}f_{t}\quad\mbox{in}\quad[0,\infty)\times\mathcal{Z}.$
(A.14)

###### Proof of Proposition 17.

From the definition of the mean embedding and by using equation (4.1), we have

$\displaystyle\partial_{t}f_{t}(w)=\partial_{t}\mathbf{m}_{\rho_{t}}(w)$
$\displaystyle=\partial_{t}\int_{\mathcal{Z}}k(z,w)\,\rho_{t}(\mathrm{d}z)=\int_{\mathcal{Z}}k(z,w)\,\partial_{t}\rho_{t}(\mathrm{d}z)$
$\displaystyle=\int_{\mathcal{Z}}k(z,w)\,\mathrm{div}_{d}(\rho_{t}\widetilde{\nabla}_{d}f_{t})(\mathrm{d}z).$

Using the definition of the divergence operator $\mathrm{div}_{d}$ at the end
of Section 3.1, we further obtain

$\displaystyle\partial_{t}f_{t}(w)$
$\displaystyle=-\int_{\mathcal{Z}}\langle\widetilde{\nabla}_{d}^{1}k(z,w),\widetilde{\nabla}_{d}f_{t}(z)\rangle_{z}\,\rho_{t}(\mathrm{d}z).$

It then follows from part i) of Lemma 15 that
$\partial_{t}f_{t}(w)=-\mathbb{K}_{\rho_{t}}f_{t}(w)$. This completes the
proof. ∎

We are now ready to present the proof of Proposition 16.

###### Proof of Proposition 16.

Let $f_{t}\coloneqq\mathbf{m}_{\rho_{t}}-\mathbf{m}_{\varrho}$. Then we have
from Proposition 6 and part ii) of Lemma 15 that
$\partial_{t}\|f_{t}\|_{\mathcal{H}}^{2}=-2\langle\mathbb{K}_{\rho_{t}}f_{t},f_{t}\rangle_{\mathcal{H}}$.
But as

$\langle\mathbb{K}_{\rho_{t}}f_{t},f_{t}\rangle_{\mathcal{H}}\geq\lambda_{t}\|f_{t}\|_{\mathcal{H}}^{2}$

by Condition A.13, we infer that
$\partial_{t}\|f_{t}\|_{\mathcal{H}}^{2}\leq-2\lambda_{t}\|f_{t}\|_{\mathcal{H}}^{2}$,
and hence
$\partial_{t}\Big{(}\log{\|f_{t}\|_{\mathcal{H}}^{2}}\Big{)}\leq-2\lambda_{t}$.
By integrating from $0$ to $t$, one gets
$\log{\|f_{t}\|_{\mathcal{H}}^{2}}-\log{\|f_{0}\|_{\mathcal{H}}^{2}}\leq-2\int_{0}^{t}\lambda_{s}\,\mathrm{d}s$.
We next take exponential to obtain

$\|f_{t}\|_{\mathcal{H}}^{2}\leq\|f_{0}\|_{\mathcal{H}}^{2}\,\exp{\Big{(}-2\int_{0}^{t}\lambda_{s}\,\mathrm{d}s\Big{)}}.$

This can be rewritten as
${\mathcal{F}}[\rho_{t}]\leq{\mathcal{F}}[\rho_{0}]\exp{\big{(}-2\int_{0}^{t}\lambda_{s}\mathrm{d}s\big{)}}$
for $t\geq 0$. In particular, ${\mathcal{F}}[\rho_{t}]$ (and hence
$\mathrm{MMD}(\rho_{t},\varrho)$) tends to zero if
$\int_{0}^{\infty}\lambda_{t}\,\mathrm{d}t=+\infty$. This completes the proof
for part i).

To prove ii), let
$f^{\tau}\coloneqq\mathbf{m}_{\rho^{\tau}}-\mathbf{m}_{\varrho}$. Notice that
in contrast to the continuous case, upper indices are used for $f^{\tau}$ and
$\rho^{\tau}$ in the discrete case. Then by using Proposition 14 together with
part ii) of Lemma 15 and the assumption $s_{\tau}\in(0,\frac{1}{4L}]$ we have

${\mathcal{F}}[\rho^{\tau+1}]-{\mathcal{F}}[\rho^{\tau}]\leq-\frac{1}{2}s_{\tau}\langle\mathbb{K}_{\rho^{\tau}}f^{\tau},f^{\tau}\rangle_{\mathcal{H}}.$

But as
$\langle\mathbb{K}_{\rho^{\tau}}f^{\tau},f^{\tau}\rangle_{\mathcal{H}}\geq\lambda_{\tau}\|f^{\tau}\|_{\mathcal{H}}^{2}$
due to our assumption, we obtain
${\mathcal{F}}[\rho^{\tau+1}]-{\mathcal{F}}[\rho^{\tau}]\leq-
s_{\tau}\lambda_{\tau}{\mathcal{F}}[\rho^{\tau}]$, or

${\mathcal{F}}[\rho^{\tau+1}]\leq(1-s_{\tau}\lambda_{\tau}){\mathcal{F}}[\rho^{\tau}]$

for every $\tau\geq 0$. As $1-s_{\tau}\lambda_{\tau}>0$, it follows by
iteration that
${\mathcal{F}}[\rho^{\tau+1}]\leq{\mathcal{F}}[\rho^{0}]{\displaystyle\prod_{i=0}^{\tau}(1-s_{i}\lambda_{i})}$.
Due to $1-x\leq\exp(-x)$ for every $x\geq 0$, we infer that
${\mathcal{F}}[\rho^{\tau+1}]\leq{\mathcal{F}}[\rho^{0}]\exp\big{(}-{\displaystyle\sum_{i=0}^{\tau}}s_{i}\lambda_{i}\big{)}$
for $\tau\geq 0$. In particular, ${\mathcal{F}}[\rho^{\tau}]$ (and hence
$\mathrm{MMD}(\rho^{\tau},\varrho)$) tends to zero if
${\displaystyle\sum_{\tau=0}^{\infty}}s_{\tau}\lambda_{\tau}=+\infty$. ∎

#### For Proposition 8

###### Proof of Proposition 8.

Let $h(z)\coloneqq\exp_{z}(s_{\tau}\Phi^{\tau}(z))$ for
$z\in\mathcal{\mathcal{missing}}Z$. Then $\rho^{\tau+1}$ can be expressed as

$\rho^{\tau+1}=h_{\\#}\rho^{\tau,\beta_{\tau}}=(h\circ
f^{\beta_{\tau}})_{\\#}(\rho^{\tau}\otimes g).$

By the computation at the beginning of the proof of Lemma 7 using Lemma 10, we
obtain

$\displaystyle{\mathcal{F}}[\rho^{\tau+1}]$
$\displaystyle-{\mathcal{F}}[\rho^{\tau}]=\frac{1}{2}\left[\mathrm{MMD}\Big{(}(h\circ
f^{\beta_{\tau}})_{\\#}(\rho^{\tau}\otimes
g),\varrho\Big{)}^{2}-\mathrm{MMD}(\rho^{\tau},\varrho)^{2}\right]$
$\displaystyle=\frac{1}{2}\iiiint\Big{\\{}k\Big{(}h(f^{\beta_{\tau}}(z,u)),h(f^{\beta_{\tau}}(w,v))\Big{)}-k(z,w)\Big{\\}}\rho^{\tau}(\mathrm{d}z)g(\mathrm{d}u)\rho^{\tau}(\mathrm{d}w)g(\mathrm{d}v)$
$\displaystyle\quad-\iiint\Big{\\{}k\Big{(}h(f^{\beta_{\tau}}(z,u)),w\Big{)}-k(z,w)\Big{\\}}\rho^{\tau}(\mathrm{d}z)g(\mathrm{d}u)\varrho(\mathrm{d}w).$

Moreover, we have

$\displaystyle I$
$\displaystyle\coloneqq\int\langle\widetilde{\nabla}_{d}[\mathbf{m}_{\rho^{\tau}}-\mathbf{m}_{\varrho}](z),\Phi^{\tau}(z)\rangle_{z}\rho^{\tau,\beta_{\tau}}(\mathrm{d}z)$
$\displaystyle=\iint\langle\widetilde{\nabla}_{d}^{1}k(z,w),\Phi^{\tau}(z)\rangle_{z}\rho^{\tau,\beta_{\tau}}(\mathrm{d}z)\rho^{\tau}(\mathrm{d}w)-\iint\langle\widetilde{\nabla}_{d}^{1}k(z,w),\Phi^{\tau}(z)\rangle_{z}\rho^{\tau,\beta_{\tau}}(\mathrm{d}z)\varrho(\mathrm{d}w)$
$\displaystyle=\frac{1}{2}\iint\langle\widetilde{\nabla}_{d}^{1}k(z,w),\Phi^{\tau}(z)\rangle_{z}\rho^{\tau,\beta_{\tau}}(\mathrm{d}z)\rho^{\tau}(\mathrm{d}w)+\frac{1}{2}\iint\langle\widetilde{\nabla}_{d}^{2}k(z,w),\Phi^{\tau}(w)\rangle_{w}\rho^{\tau,\beta_{\tau}}(\mathrm{d}w)\rho^{\tau}(\mathrm{d}z)$
$\displaystyle\quad-\iint\langle\widetilde{\nabla}_{d}^{1}k(z,w),\Phi^{\tau}(z)\rangle_{z}\rho^{\tau,\beta_{\tau}}(\mathrm{d}z)\varrho(\mathrm{d}w)$
$\displaystyle=\frac{1}{2}\iiint\Big{\langle}\widetilde{\nabla}_{d}^{1}k(f^{\beta_{\tau}}(z,u),w),\Phi^{\tau}(f^{\beta_{\tau}}(z,u))\Big{\rangle}_{f^{\beta_{\tau}}(z,u)}\rho^{\tau}(\mathrm{d}z)g(\mathrm{d}u)\rho^{\tau}(\mathrm{d}w)$
$\displaystyle\quad+\frac{1}{2}\iiint\Big{\langle}\widetilde{\nabla}_{d}^{2}k(z,f^{\beta_{\tau}}(w,v)),\Phi^{\tau}(f^{\beta_{\tau}}(w,v))\Big{\rangle}_{f^{\beta_{\tau}}(w,v)}\rho^{\tau}(\mathrm{d}w)g(\mathrm{d}v)\rho^{\tau}(\mathrm{d}z)$
$\displaystyle\quad-\iiint\Big{\langle}\widetilde{\nabla}_{d}^{1}k(f^{\beta_{\tau}}(z,u),w),\Phi^{\tau}(f^{\beta_{\tau}}(z,u))\Big{\rangle}_{f^{\beta_{\tau}}(z,u)}\rho^{\tau}(\mathrm{d}z)g(\mathrm{d}u)\varrho(\mathrm{d}w),$

where the third equality is due to the symmetry of $k$ and relation (3.6).
Therefore, it follows that

$\displaystyle{\mathcal{F}}[\rho^{\tau+1}]-{\mathcal{F}}[\rho^{\tau}]-s_{\tau}I$
$\displaystyle=\frac{1}{2}\iiiint\left\\{k\Big{(}h(f^{\beta_{\tau}}(z,u)),h(f^{\beta_{\tau}}(w,v))\Big{)}-k(z,w)\right.$
$\displaystyle\qquad\qquad\qquad\left.-\Big{[}\langle\widetilde{\nabla}_{d}^{1}k(f^{\beta_{\tau}}(z,u),w),s_{\tau}\Phi^{\tau}(f^{\beta_{\tau}}(z,u))\rangle_{f^{\beta_{\tau}}(z,u)}\right.$
$\displaystyle\left.\qquad\qquad\qquad\quad+\langle\widetilde{\nabla}_{d}^{2}k(z,f^{\beta_{\tau}}(w,v)),s_{\tau}\Phi^{\tau}(f^{\beta_{\tau}}(w,v))\rangle_{f^{\beta_{\tau}}(w,v)}\Big{]}\right\\}\rho^{\tau}(\mathrm{d}z)g(\mathrm{d}u)\rho^{\tau}(\mathrm{d}w)g(\mathrm{d}v)$
$\displaystyle\quad-\iiint\Big{\\{}k\Big{(}h(f^{\beta_{\tau}}(z,u)),w\Big{)}-k(z,w)-\langle\widetilde{\nabla}_{d}^{1}k(f^{\beta_{\tau}}(z,u),w),s_{\tau}\Phi^{\tau}(f^{\beta_{\tau}}(z,u))\rangle_{f^{\beta_{\tau}}(z,u)}\Big{\\}}$
$\displaystyle\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\quad\rho^{\tau}(\mathrm{d}z)g(\mathrm{d}u)\varrho(\mathrm{d}w).$

As $h(z)=\exp_{z}(s_{\tau}\Phi^{\tau}(z))$ and
$s_{\tau}\in(0,\varepsilon_{0}]$, we can now use the Lipschitz-gradient
condition (11) for $k$ to obtain

$\displaystyle{\mathcal{F}}[\rho^{\tau+1}]-{\mathcal{F}}[\rho^{\tau}]-s_{\tau}I$
$\displaystyle\leq\frac{L}{2}\iiiint\Big{[}\|s_{\tau}\Phi^{\tau}(f^{\beta_{\tau}}(z,u))\|_{f^{\beta_{\tau}}(z,u)}^{2}+\|s_{\tau}\Phi^{\tau}(f^{\beta_{\tau}}(w,v))\|_{f^{\beta_{\tau}}(w,v)}^{2}\Big{]}\rho^{\tau}(\mathrm{d}z)g(\mathrm{d}u)\rho^{\tau}(\mathrm{d}w)g(\mathrm{d}v)$
$\displaystyle\quad+L\iiint\|s_{\tau}\Phi^{\tau}(f^{\beta_{\tau}}(z,u))\|_{f^{\beta_{\tau}}(z,u)}^{2}\rho^{\tau}(\mathrm{d}z)\varrho(\mathrm{d}w)g(\mathrm{d}u)$
$\displaystyle=2Ls_{\tau}^{2}\iint\|\Phi^{\tau}(f^{\beta_{\tau}}(z,u))\|_{f^{\beta_{\tau}}(z,u)}^{2}\rho^{\tau}(\mathrm{d}z)g(\mathrm{d}u).$

Using the definition
$\rho^{\tau,\beta_{\tau}}=f^{\beta_{\tau}}_{\\#}(\rho^{\tau}\otimes g)$ and
the fact
$I=-\int_{\mathcal{Z}}\|\Phi^{\tau}(z)\|_{z}^{2}\,\rho^{\tau,\beta_{\tau}}(\mathrm{d}z)$,
we can rewrite this more compactly as

$\displaystyle{\mathcal{F}}[\rho^{\tau+1}]-{\mathcal{F}}[\rho^{\tau}]$
$\displaystyle\leq-
s_{\tau}\big{(}1-2Ls_{\tau}\big{)}\int_{\mathcal{Z}}\|\Phi^{\tau}(z)\|_{z}^{2}\,\rho^{\tau,\beta_{\tau}}(\mathrm{d}z)$
$\displaystyle=-s_{\tau}\big{(}1-2Ls_{\tau}\big{)}\int_{\mathcal{Z}}\|\widetilde{\nabla}_{d}[\mathbf{m}_{\rho^{\tau}}-\mathbf{m}_{\varrho}](z)\|_{z}^{2}\,\rho^{\tau,\beta_{\tau}}(\mathrm{d}z).$

This together with condition (4.4) gives

${\mathcal{F}}[\rho^{\tau+1}]\leq(1-a_{\tau}){\mathcal{F}}[\rho^{\tau}]\quad\mbox{with}\quad
a_{\tau}\coloneqq\lambda s_{\tau}\big{(}1-2Ls_{\tau}\big{)}\beta_{\tau}^{2}.$

In particular, we must have $a_{i}\leq 1$. By iterating this estimate, we
obtain

${\mathcal{F}}[\rho^{\tau+1}]\leq{\mathcal{F}}[\rho^{0}]\,\prod_{i=0}^{\tau}(1-a_{i}).$
(A.15)

Due to $1-x\leq\exp(-x)$ for every number $x\geq 0$, we get
$\prod_{i=0}^{\tau}(1-a_{i})\leq\exp(-\sum_{i=0}^{\tau}a_{i})$. This together
with (A.15) yields the conclusion of the proposition. ∎

##  Appendix B Implementation Details And Additional Results

We use $\nabla_{\mathbb{B}}^{1}k(x,\mu,\Sigma,w)$ to denote the last component
in (3.6) for the gradient $\widetilde{\nabla}_{d}$ of the function
$(x,\mu,\Sigma)\mapsto k(x,\mu,\Sigma,w)$. Precisely,

$\nabla_{\mathbb{B}}^{1}k(x,\mu,\Sigma,w)\coloneqq
2[\nabla_{\Sigma}k(x,\mu,\Sigma,w)]\Sigma+2\Sigma[\nabla_{\Sigma}k(x,\mu,\Sigma,w)].$

###  B.1 Algorithms

Algorithm 2 Discretized Gradient Flow Algorithm for Scheme (4.2) – Detailed
Version of Algorithm 1

1: Input: a source distribution
$\rho^{0}=\frac{1}{N}\sum_{i=1}^{N}\delta_{(x^{0}_{i},\mu^{0}_{i},\Sigma^{0}_{i})}$,
a sample
$\frac{1}{M}\sum_{j=1}^{M}\delta_{(\bar{x}_{j},\bar{\mu}_{j},\bar{\Sigma}_{j})}$
for the target distribution $\varrho$, a number $T$ of iterations for
training, a sequence of step sizes $s_{\tau}>0$ with $\tau=0,1,...,T$, and a
kernel $k$.

2: Initialization:

3: Compute
$\displaystyle(\bar{\Psi}_{1},\bar{\Psi}_{2},\bar{\Psi}_{3})(x,\mu,\Sigma)=\frac{1}{M}\sum_{j=1}^{M}(\nabla_{x},\nabla_{\mu},\nabla_{\mathbb{B}}^{1})k(x,\mu,\Sigma,\bar{x}_{j},\bar{\mu}_{j},\bar{\Sigma}_{j})$

4: $\tau\leftarrow 0$

5: while $\tau<T$ do

6: Compute
$\displaystyle(\Psi_{1}^{\tau},\Psi_{2}^{\tau},\Psi_{3}^{\tau})(x,\mu,\Sigma)=\frac{1}{N}\sum_{i=1}^{N}(\nabla_{x},\nabla_{\mu},\nabla_{\mathbb{B}}^{1})k(x,\mu,\Sigma,x^{\tau}_{i},\mu^{\tau}_{i},\Sigma^{\tau}_{i})$

7: for $i=1,\ldots,N$ do

8: $x^{\tau+1}_{i}\leftarrow
x^{\tau}_{i}+s_{\tau}(\bar{\Psi}_{1}-\Psi_{1}^{\tau})(x^{\tau}_{i},\mu^{\tau}_{i},\Sigma^{\tau}_{i})$

9:
$\mu^{\tau+1}_{i}\leftarrow\mu^{\tau}_{i}+s_{\tau}(\bar{\Psi}_{2}-\Psi_{2}^{\tau})(x^{\tau}_{i},\mu^{\tau}_{i},\Sigma^{\tau}_{i})$

10:
$\Sigma^{\tau+1}_{i}\leftarrow\Big{(}I+s_{\tau}{\mathrm{L}}_{\Sigma^{\tau}_{i}}\big{[}(\bar{\Psi}_{3}-\Psi_{3}^{\tau})(x^{\tau}_{i},\mu^{\tau}_{i},\Sigma^{\tau}_{i})\big{]}\Big{)}\Sigma^{\tau}_{i}\Big{(}I+s_{\tau}{\mathrm{L}}_{\Sigma^{\tau}_{i}}[(\bar{\Psi}_{3}-\Psi_{3}^{\tau})(x^{\tau}_{i},\mu^{\tau}_{i},\Sigma^{\tau}_{i})]\Big{)}$

11: end for

12: Set $\tau\leftarrow\tau+1$

13: end while

14: Output:
$\displaystyle\rho^{T}=\frac{1}{N}\sum_{i=1}^{N}\delta_{(x^{T}_{i},\mu^{T}_{i},\Sigma^{T}_{i})}$

Algorithm 3 Discretized Gradient Flow Algorithm for Scheme (4.3)

1: Input: a source distribution
$\rho^{0}=\frac{1}{N}\sum_{i=1}^{N}\delta_{(x_{i},\mu_{i},\Sigma_{i})}$, a
target distribution
$\varrho=\frac{1}{M}\sum_{j=1}^{M}\delta_{(\bar{x}_{j},\bar{\mu}_{j},\bar{\Sigma}_{j})}$,
number of iterations $T$, step sizes $s_{\tau}>0$, noise levels
$\beta_{\tau}$, and a kernel $k$.

2: Initialization:

3: Compute
$\displaystyle(\bar{\Psi}_{1},\bar{\Psi}_{2},\bar{\Psi}_{3})(x,\mu,\Sigma)=\frac{1}{M}\sum_{j=1}^{M}(\nabla_{x},\nabla_{\mu},\nabla_{\mathbb{B}}^{1})k(x,\mu,\Sigma,\bar{x}_{j},\bar{\mu}_{j},\bar{\Sigma}_{j})$

4: $\tau\leftarrow 0$

5: while $\tau<T$ do

6: Compute
$\displaystyle(\Psi_{1}^{\tau},\Psi_{2}^{\tau},\Psi_{3}^{\tau})(x,\mu,\Sigma)=\frac{1}{N}\sum_{j=1}^{N}(\nabla_{x},\nabla_{\mu},\nabla_{\mathbb{B}}^{1})k(x,\mu,\Sigma,x^{\tau}_{j},\mu^{\tau}_{j},\Sigma^{\tau}_{j})$

7: for $i=1,\ldots,N$ do

8: Perturb $x_{i}^{\tau,p}\leftarrow
x_{i}^{\tau}+\beta_{\tau}\mathcal{N}_{{\mathbb{R}}^{m}}(0,1)$ and
$\mu_{i}^{\tau,p}\leftarrow\mu_{i}^{\tau}+\beta_{\tau}\mathcal{N}_{{\mathbb{R}}^{n}}(0,1)$

9: Set $S\leftarrow\beta_{\tau}\mathcal{N}_{{\mathbb{S}}^{n}}(0,1)$ and
perturb
$\Sigma_{i}^{\tau,p}\leftarrow(I+{\mathrm{L}}_{\Sigma_{i}^{\tau}}[S])\Sigma_{i}^{\tau}(I+{\mathrm{L}}_{\Sigma_{i}^{\tau}}[S])$

10: $x^{\tau+1}_{i}\leftarrow
x_{i}^{\tau,p}+s_{\tau}(\bar{\Psi}_{1}-\Psi_{1}^{\tau})(x^{\tau,p}_{i},\mu^{\tau,p}_{i},\Sigma^{\tau,p}_{i})$

11:
$\mu^{\tau+1}_{i}\leftarrow\mu_{i}^{\tau,p}+s_{\tau}(\bar{\Psi}_{2}-\Psi_{2}^{\tau})(x^{\tau,p}_{i},\mu^{\tau,p}_{i},\Sigma^{\tau,p}_{i})$

12:
$\Sigma^{\tau+1}_{i}\leftarrow\Big{(}I+s_{\tau}{\mathrm{L}}_{\Sigma_{i}^{\tau,p}}\big{[}(\bar{\Psi}_{3}-\Psi_{3}^{\tau})(x^{\tau,p}_{i},\mu^{\tau,p}_{i},\Sigma^{\tau,p}_{i})\big{]}\Big{)}\Sigma_{i}^{\tau,p}\Big{(}I+s_{\tau}{\mathrm{L}}_{\Sigma_{i}^{\tau,p}}[(\bar{\Psi}_{3}-\Psi_{3}^{\tau})(x^{\tau,p}_{i},\mu^{\tau,p}_{i},\Sigma^{\tau,p}_{i})]\Big{)}$

13: end for

14: Set $\tau\leftarrow\tau+1$

15: end while

16: Output:
$\displaystyle\rho^{T}=\frac{1}{N}\sum_{i=1}^{N}\delta_{(x^{T}_{i},\mu^{T}_{i},\Sigma^{T}_{i})}$

###  B.2 Kernel and Its Gradient for Implementation

We use the kernel $k$ given by:

$k\left((x,\mu,\Sigma),(\bar{x},\bar{\mu},\bar{\Sigma})\right)\coloneqq\exp\left(-\alpha\|x-\bar{x}\|_{2}^{2}\right)\exp\left(-\beta\|\mu-\bar{\mu}\|_{2}^{2}\right)\exp\left(-\gamma\|\Sigma-\bar{\Sigma}\|_{2}^{2}\right),$

where $\alpha$, $\beta$ and $\gamma$ are parameters (bandwidth) of the kernel.
We note that this kernel is characteristic by [Szabó and Sriperumbudur, 2018,
Theorem 4]. Then its standard Euclidean gradient is given by

$\displaystyle\nabla_{(x,\mu,\Sigma)}k\left((x,\mu,\Sigma),(\bar{x},\bar{\mu},\bar{\Sigma})\right)=-2\exp\left(-\alpha\|x-\bar{x}\|_{2}^{2}-\beta\|\mu-\bar{\mu}\|_{2}^{2}-\gamma\|\Sigma-\bar{\Sigma}\|_{2}^{2}\right)\begin{bmatrix}\alpha(x-\bar{x})\\\
\beta(\mu-\bar{\mu})\\\ \gamma(\Sigma-\bar{\Sigma})\end{bmatrix}.$

Thus by plugging into formula (3.6), we obtain

$\displaystyle\widetilde{\nabla}_{d}^{1}k\left((x,\mu,\Sigma),(\bar{x},\bar{\mu},\bar{\Sigma})\right)$
$\displaystyle\quad=-2\exp\left(-\alpha\|x-\bar{x}\|_{2}^{2}-\beta\|\mu-\bar{\mu}\|_{2}^{2}-\gamma\|\Sigma-\bar{\Sigma}\|_{2}^{2}\right)\begin{bmatrix}\alpha(x-\bar{x})\\\
\beta(\mu-\bar{\mu})\\\
2\gamma(2\Sigma^{2}-\Sigma\bar{\Sigma}-\bar{\Sigma}\Sigma)\end{bmatrix}.$

That is,

$\displaystyle\nabla_{x}k\left((x,\mu,\Sigma),(\bar{x},\bar{\mu},\bar{\Sigma})\right)=-2\exp\left(-\alpha\|x-\bar{x}\|_{2}^{2}-\beta\|\mu-\bar{\mu}\|_{2}^{2}-\gamma\|\Sigma-\bar{\Sigma}\|_{2}^{2}\right)\alpha(x-\bar{x}),$
$\displaystyle\nabla_{\mu}k\left((x,\mu,\Sigma),(\bar{x},\bar{\mu},\bar{\Sigma})\right)=-2\exp\left(-\alpha\|x-\bar{x}\|_{2}^{2}-\beta\|\mu-\bar{\mu}\|_{2}^{2}-\gamma\|\Sigma-\bar{\Sigma}\|_{2}^{2}\right)\beta(\mu-\bar{\mu}),$
$\displaystyle\nabla_{\mathbb{B}}^{1}k\left((x,\mu,\Sigma),(\bar{x},\bar{\mu},\bar{\Sigma})\right)$
$\displaystyle\quad=-2\exp\left(-\alpha\|x-\bar{x}\|_{2}^{2}-\beta\|\mu-\bar{\mu}\|_{2}^{2}-\gamma\|\Sigma-\bar{\Sigma}\|_{2}^{2}\right)2\gamma(2\Sigma^{2}-\Sigma\bar{\Sigma}-\bar{\Sigma}\Sigma).$

###  B.3 Label Projection

We here propose an approach to recover new samples in the feature-label space
from an empirical distribution in the feature-Gaussian space. Consider that
after $T$ iterations of the gradient algorithms, we arrive at a distribution
$\rho^{T}=\frac{1}{N}\sum_{i=1}^{N}\delta_{(x^{T}_{i},\mu^{T}_{i},\Sigma^{T}_{i})}$.
We would like to recover a distribution
$\nu^{T}\in\mathcal{P}(\mathcal{X}\times\mathcal{Y})$ which is induced by
$\rho^{\tau}$. As such, we would like to find a distribution $\nu^{T}$ of the
form

$\nu^{T}=\frac{1}{N}\sum_{i=1}^{N}\delta_{(x_{i}^{T},y_{i}^{T})}$

which corresponds to new target samples $(x_{i}^{T},y_{i}^{T})_{i=1}^{N}$.
Moreover, we are interested in recovering labels within the target domain. To
this end, let $\mathcal{Y}_{\rm target}=\\{y\in\mathcal{Y}:\exists
j\in[M]\text{ such that }\bar{y}_{j}=y\\}$ be the set of labels in the target
dataset, and remind that for any $y\in\mathcal{Y}_{\rm target}$,
$(\bar{\mu}_{y},\bar{\Sigma}_{y})\in{\mathbb{R}}^{n}\times{\mathbb{S}}_{+}^{n}$
is the mean vector and the covariance matrix of the distribution of $\phi(X)$
given $Y=y$. Notice that the mean-covariance embeddings
$(\bar{\mu}_{y},\bar{\Sigma}_{y})$ for $y\in\mathcal{Y}_{\rm target}$ depend
only on the target domain data, and it does not depend on the incumbent
distribution $\rho^{T}$, nor does it depend on the source dataset. Moreover,
we can also compute $\bar{N}_{y}$ as the number of samples from the target
dataset with label $y$.

Because $(\bar{\mu}_{y},\bar{\Sigma}_{y})$ is readily computed, we can
consider $(\bar{\mu}_{y},\bar{\Sigma}_{y})$ as the centroids and simply find
an assignment that minimizes the sum of distances from
$(\mu^{T}_{i},\Sigma^{T}_{i})$ to these centroids. We thus can assign each
sample from $\rho^{T}$ to the the target labels by solving the linear program

$\begin{array}[]{cl}\min&\displaystyle\sum_{i=1}^{N}\sum_{y\in\mathcal{Y}_{\rm
target}}\theta_{iy}\sqrt{\|\mu^{T}_{i}-\bar{\mu}_{y}\|_{2}^{2}+\mathbb{B}(\Sigma^{T}_{i},\bar{\Sigma}_{y})^{2}}\\\
\text{s.t.}&\displaystyle\sum_{y\in\mathcal{Y}_{\rm
target}}\theta_{iy}=\frac{1}{N}\quad\forall
i=1,\ldots,N,\qquad\sum_{i=1}^{N}\theta_{iy}=\frac{\bar{N}_{y}}{N}\quad\forall
y\in\mathcal{Y}_{\rm target},\quad\theta\in[0,1]^{N\times|\mathcal{Y}_{\rm
target}|},\end{array}$ (B.1)

Notice that the assignment problem above does not utilize the information from
the covariate $x_{i}^{T}$. Let $\theta^{\star}$ be the optimal solution of the
above optimization problem. Then the dataset $(x_{i}^{T},z_{i}^{T})_{i=1}^{N}$
recovered from $\rho^{\tau}$ is

$\nu^{T}=\frac{1}{N}\sum_{i=1}^{N}\delta_{(x_{i}^{T},y_{i}^{T})},\qquad
y_{i}^{T}=\sum_{y\in\mathcal{Y}_{\rm
target}}y\mathbbm{1}(\theta_{iy}^{\star}=\max\\{\theta_{i}^{\star}\\})\quad\forall
i=1,\ldots,N.$

We used the POT library to solve the label recovery problem (B.1).

###  B.4 Results on Mixture of Gaussians

We test our algorithm on a toy example: a mixture of Gaussian distributions to
another mixture of Gaussian distributions.

The source distribution and target distribution are:

$\displaystyle p_{s}(x)$
$\displaystyle=\frac{1}{4}\mathcal{N}\left(\begin{pmatrix}2.0\\\
-0.3\end{pmatrix},\begin{pmatrix}0.14&-0.00\\\
-0.00&0.22\end{pmatrix}\right)+\frac{1}{4}\mathcal{N}\left(\begin{pmatrix}2.0\\\
0.3\end{pmatrix},\begin{pmatrix}0.43&0.18\\\ 0.18&0.26\end{pmatrix}\right)$
$\displaystyle+\frac{1}{4}\mathcal{N}\left(\begin{pmatrix}-0.3\\\
2.0\end{pmatrix},\begin{pmatrix}0.66&0.02\\\
0.02&0.63\end{pmatrix}\right)+\frac{1}{4}\mathcal{N}\left(\begin{pmatrix}0.3\\\
-2.0\end{pmatrix},\begin{pmatrix}0.39&-0.02\\\ -0.02&0.13\end{pmatrix}\right)$
$\displaystyle p_{t}(x)$
$\displaystyle=\frac{1}{4}\mathcal{N}\left(\begin{pmatrix}2.9\\\
0.1\end{pmatrix},\begin{pmatrix}0.16&0.03\\\
0.03&0.20\end{pmatrix}\right)+\frac{1}{4}\mathcal{N}\left(\begin{pmatrix}0.9\\\
0.5\end{pmatrix},\begin{pmatrix}0.22&0.16\\\ 0.16&0.46\end{pmatrix}\right)$
$\displaystyle+\frac{1}{4}\mathcal{N}\left(\begin{pmatrix}0.8\\\
2.2\end{pmatrix},\begin{pmatrix}0.63&0.02\\\
0.02&0.66\end{pmatrix}\right)+\frac{1}{4}\mathcal{N}\left(\begin{pmatrix}1.4\\\
-1.8\end{pmatrix},\begin{pmatrix}0.18&0.10\\\ 0.10&0.36\end{pmatrix}\right)$

From each distribution, we sample 25 particles and flow the particles’
positions, means, and covariance simultaneously using Alg. 1. After the
algorithm converges, we recover the particles’ label in the feature-label
space by solving problem (B.1).

Figure 4: The results of flowing a mixture of 4 Gaussian distributions to a
mixture of 4 Gaussian distributions. We demonstrate the initialization (left),
the trace of particles in first 200 steps (middle), and the results at step
1000 (right).

We test how our algorithm deals with flowing a mixture of 2 Gaussian
distributions to a mixture of 4 Gaussian distributions. From the trace of
first 200 steps, we demonstrate that each source Gaussian distribution splits
into 2 Gaussian distributions. The source distribution and target distribution
are:

$\displaystyle p_{s}(x)$
$\displaystyle=\frac{1}{2}\mathcal{N}\left(\begin{pmatrix}0.0\\\
0.0\end{pmatrix},\begin{pmatrix}0.18&-0.24\\\
-0.24&0.70\end{pmatrix}\right)+\frac{1}{2}\mathcal{N}\left(\begin{pmatrix}5.8\\\
0.0\end{pmatrix},\begin{pmatrix}0.44&0.00\\\ 0.00&0.87\end{pmatrix}\right)$
$\displaystyle p_{t}(x)$
$\displaystyle=\frac{1}{4}\mathcal{N}\left(\begin{pmatrix}2.0\\\
0.7\end{pmatrix},\begin{pmatrix}0.63&-0.30\\\
-0.30&0.26\end{pmatrix}\right)+\frac{1}{4}\mathcal{N}\left(\begin{pmatrix}2.2\\\
-0.8\end{pmatrix},\begin{pmatrix}0.77&-0.18\\\ -0.18&0.55\end{pmatrix}\right)$
$\displaystyle+\frac{1}{4}\mathcal{N}\left(\begin{pmatrix}7.0\\\
0.8\end{pmatrix},\begin{pmatrix}0.63&-0.30\\\
-0.30&0.26\end{pmatrix}\right)+\frac{1}{4}\mathcal{N}\left(\begin{pmatrix}7.7\\\
-0.8\end{pmatrix},\begin{pmatrix}0.77&-0.18\\\ -0.18&0.55\end{pmatrix}\right)$

Figure 5: The results of flowing a mixture of 2 Gaussian distributions to a
mixture of 4 Gaussian distributions. We demonstrate the initialization (left),
the trace of particles in first 200 steps (middle), and the results at step
2400 (right). We use method in Section B.3 to relabel the flowed samples and
the labels correspond to their positional belongings correctly.

###  B.5 Implementation Details

Our algorithms and experiments are implemented in PyTorch. The *NIST images
are resized to $20\times 20$, thus the feature space is of dimension $m=400$.
The SVHN and Tiny ImageNet images are of size $3\times 64\times 64$. The
classifiers for *NIST datasets are of LeNet-5 architecture. The classifiers
for SVHN and Tiny ImageNet are of a ResNet-18 network. For the experiments in
Figure 3, we train the classifiers with the set of data 10 epochs with Adam
optimizer and learning rate $2\times 10^{-3}$.

When flowing images in *NIST datasets and flowing a mixture of Gaussians, we
use the parameters and methods described in Table 2. As we tune the
parameters, we notice the algorithm converges with a range of parameters and
we report one setting in Table 2.

|  SVHN&TIN |  *NIST |  Gaussian (Figure 4) |  Gaussian (Figure 5)
---|---|---|---|---
$\alpha$ | 0.002 | 0.001 | 0.3 | 0.3
$\beta$ | 0.01 | 0.002 | 0.15 | 0.1
$\gamma$ | 1.0 | 100 | 1.0 | 0.5
initial $s_{\tau}$ | 1.0 | 0.3 | 0.05 | 0.03
noise level | 0.1 | 0.01 | 0 | 0.1
$T$ | 6000 | 150 | 2000 | 2500
Optimizer |  RMSprop Hinton et al. [2012] | RMSprop | RMSprop | RMSprop

Table 2: Parameters and Optimizer

Our method assumes images of each class form one Gaussian distribution. In
reality, the data can be a mixture of Gaussian. To satisfy the Gaussianity
assumption, in the preprocessing step, we use a clustering method ($k$-nearest
neighbors) and pick only data from one mode for each class. As a consequence,
the data used in the experiment satisfies the conditional Gaussian assumption.
For example, the images of the digit 1 can have two modes: slanted left or
slanted right. In this case, we can generate two labels (1L, 1R), and the
methodology developed in this paper can be applied in a straightforward
manner. When testing our transfer learning scheme, we apply the same
clustering method on the test dataset, so our test set is within the same mode
as our training set.

We store the preprocessed data and apply dimension reduction method on the
data’s means and covariance matrices, so the Lyapunov equation is much faster
to solve. We use the cluster’s mean and covariance matrix to approximate the
1-shot and 5-shor data’s mean and covariance matrix. In 1-shot learning, the
covariance matrix is an identity matrix. All the code and data are available
in the supplementary file. We use $k$-nearest neighbors algorithm to solve the
labels of the flowed data, as it performs better with the noisy scheme.

###  B.6 Additional Results on *NIST Datasets

We conduct additional experiments of flowing between KMNIST and FashionMNIST
datasets. The results of our flows are illustrated in the same fashion as
Figure 2 and are in the supplementary folder. In each subfigure, each column
represents a snapshot of a certain time-step and the samples flow from the
source (left) to the target (right). To check if the algorithm is converging,
we compute the MMD between the source dataset and the target dataset. Figure 6
is an example of MMD decreasing in transferring the *NIST datasets:

Figure 6: MMD keeps decreasing as we flow the samples from the source domain
to the target domain

From FashionMNIST dataset to KMNIST dataset and from KMNIST dataset to
FashionMNIST dataset, we also conduct transfer learning experiments. We use
the same model architecture and training settings as in Fig. 3. We illustrate
the accuracy and error bars of the 1-shot learning and 5-shot learning in Fig.
8. Our flowed samples ($S_{T}$) increase the accuracy of the transferred
classifiers in both 1-shot and 5-shot learning.

Figure 7: Average target domain accuracy on the test split for transfer
learning with one-shot (left) and ﬁve-shot (right). Results are taken over 10
independent replications, and the range of accuracy is displayed by the error
bars.

###  B.7 Transfer Learning Results on SVHN and Tiny ImageNet

We run the same transfer learning experiments as section 5.1 on SVHN and Tiny
ImageNet datasets and we report the accuracy on the test split333The Tiny
ImageNet dataset does not have labels for the test split, so we use the
validation split.. From the results, we can see $S_{T}$ always improve the
accuracy.

On the higher-dimensional data, it is a lot more difficult to train a strong
classifier from only 1 or 5 samples each class, so the accuracy is not so high
compared to the *NIST datasets. Another reason is there are only 50 samples in
each class of the validation dataset of Tiny ImageNet, so the number of test
samples might be too small.

Figure 8: Average target domain accuracy on the test split for transfer
learning with one-shot (left) and ﬁve-shot (right). Results are taken over 10
independent replications, and the range of accuracy is displayed by the error
bars. The $\bm{1}$ and $\bm{5}$ in the parenthesis mean 1-shot and 5-shot
transfer learning respectively.

###  B.8 Comparison with baseline

We compare with Alvarez-Melis and Fusi [2021]’s approach as the baseline. We
adopted the same values of parameters (number of steps, step size) from the
paper of Alvarez-Melis and Fusi [2021] and experimented with different values
of entropy regularization $\lambda$. Entropy regularization $\lambda$ is a
hidden parameter in their code and they reported using the value of
$\lambda=100$ in transfer learning experiments. Also, we experimented with
different methods in their code and found that “xyaugm” gives the best
qualitative gradient flow results.

We test their algorithm using the same transfer learning setting as ours
(using the same clustered data as our experiments), see Table 3 and Table 5
for results. The results show that their method doesn’t improve the
classification accuracy after adding the transferred data $S_{T}$, if we
compare the accuracy $D\cup S_{T}$ and $P\cup D\cup S_{T}$ with $D$ and $P\cup
D$ correspondingly.

Accuracy | $D\cup S_{T}$ | $P\cup D\cup S_{T}$
---|---|---
$\lambda=0.001$ | 0.1941 | 0.2478
$\lambda=0.01$ | 0.1175 | 0.2998
$\lambda=1.0$ | 0.1739 | 0.2951
$\lambda=100$ | 0.2093 | 0.3025
Table 3: KMNIST $\shortrightarrow$ MNIST Accuracy | $D\cup S_{T}$ | $P\cup D\cup S_{T}$
---|---|---
$\lambda=0.001$ | 0.4488 | 0.3083
$\lambda=0.01$ | 0.3915 | 0.2897
$\lambda=1.0$ | 0.5268 | 0.4627
$\lambda=100$ | 0.2917 | 0.3307
Table 4: FMNIST $\shortrightarrow$ MNIST

For runtime comparison, the default device for Alvarez-Melis and Fusi [2021]’s
code is on the CPU. As we run their code on the GPU, the kernel crashed
without giving any informative errors. Thus, we will compare our codes’
runtime per step on the CPU. While our approach takes about $0.512$ second,
the approach in Alvarez-Melis and Fusi [2021] requires $74.78$ seconds.

###  B.9 Comparison with Mixup Method

We run mixup augmentation on F$\shortrightarrow$M, K$\shortrightarrow$M,
M$\shortrightarrow$F, and M$\shortrightarrow$K on the same one-shot learning
setting and same data as ours using the method in Zhang et al. [2017] with
$\alpha=0.2$ (suggested by Zhang et al. [2017]). The results are in Table 5.
When comparing $D$ with $D\cup S_{T}$ and $P\cup D$ with $P\cup D\cup S_{T}$,
the added mixup samples ($S_{T}$) decrease the accuracy. A reason could be
that there are too few examples in the target domain, thus mixup does not add
much complexity to the target domain. The mixup experiments in [A] show
improvements $<1\%$, which is much smaller than ours in Fig.3.

Last, we want to emphasize that our gradient flow method aims to increase the
number of samples of one distribution, so our method is a complement to
augmentation methods.Both gradient flow and augmentation methods are applied
in practice at the same time.

Training Data | F$\shortrightarrow$M | K$\shortrightarrow$M | M$\shortrightarrow$F | M$\shortrightarrow$K
---|---|---|---|---