*arg min Danil Akhtiamov<EMAIL_ADDRESS>
Department of Computing and Mathematical Sciences
California Institute of Technology
Pasadena, CA 91125 and Babak Hassibi<EMAIL_ADDRESS>
Department of Electrical Engineering
California Institute of Technology
Pasadena, CA 91125

# The Generalization Error of Stochastic Mirror Descent on Over-Parametrized
Linear Models

###### Abstract

Despite being highly over-parametrized, and having the ability to fully
interpolate the training data, deep networks are known to generalize well to
unseen data. It is now understood that part of the reason for this is that the
training algorithms used have certain implicit regularization properties that
ensure interpolating solutions with “good” properties are found. This is best
understood in linear over-parametrized models where it has been shown that the
celebrated stochastic gradient descent (SGD) algorithm finds an interpolating
solution that is closest in Euclidean distance to the initial weight vector.
Different regularizers, replacing Euclidean distance with Bregman divergence,
can be obtained if we replace SGD with stochastic mirror descent (SMD).
Empirical observations have shown that in the deep network setting, SMD
achieves a generalization performance that is different from that of SGD (and
which depends on the choice of SMD’s potential function. In an attempt to
begin to understand this behavior, we obtain the generalization error of SMD
for over-parametrized linear models for a binary classification problem where
the two classes are drawn from a Gaussian mixture model. We present simulation
results that validate the theory and, in particular, introduce two data
models, one for which SMD with an $\ell_{2}$ regularizer (i.e., SGD)
outperforms SMD with an $\ell_{1}$ regularizer, and one for which the reverse
happens.

###### keywords:

Stochastic gradient descent, stochastic mirror descent, generalization error,
binary classification, convex Gaussian min-max theorem

##  1 Introduction

Stochastic gradient descent (SGD), along with its variants, is the workhorse
of modern machine learning. Among these variants is stochastic mirror descent
(SMD) [Nemirovski and Yudin(1983)], which differs from SGD in that, instead of
updating the weight vector, one updates the gradient of a so-called
“potential” function of the weight vector, along the negative direction of the
instantaneous gradient of the loss function. The potential function is what
defines a particular instantiation SMD. It is required to be differentiable
and strictly convex. When the potential is the squared Euclidean norm we get
SGD.

In deep learning the models are over-parameterized, typically with a number of
parameters that is orders of magnitude larger than the size of the training
set. In such a setting, there are uncountably many weight vectors that
perfectly interpolate the data. And so, it is not clear which will generalize
well on unseen data—some may, some may not [Zhang et al.(2016)Zhang, Bengio,
Hardt, Recht, and Vinyals]. One of the open questions in deep learning is why
SGD almost invariably finds solutions that generalize well? In an attempt to
understand this question [Gunasekar et al.(2017)Gunasekar, Woodworth,
Bhojanapalli, Neyshabur, and Srebro] showed that, in over-parameterized linear
models, GD finds an interpolating weight vector that minimizes its Euclidean
distance from the initial weight vector. In particular, if we initialize with
zero (or rather very close to zero in practice) , it finds an interpolating
weight vector with minimum 2-norm. This is what is called ”implicit
regularization” and is what makes the solution obtained by GD different from
other interpolating solutions. [Gunasekar et al.(2017)Gunasekar, Woodworth,
Bhojanapalli, Neyshabur, and Srebro] further showed that mirror descent does
the same, except that it finds an interpolating solution that minimizes its
Bregman divergence from the initial weight vector. This observation allows one
to impose different regularizations on the interpolating weight vector. In
[Azizan and Hassibi(2019b)] these results were extended to SGD and SMD, and
then informally extended to nonlinear models, such as deep networks, in
[Azizan et al.(2021)Azizan, Lale, and Hassibi].

In [Azizan et al.(2021)Azizan, Lale, and Hassibi] it was empirically observed
that the generalization error of SMD, for deep networks initialized with the
same weight vector and trained on the same training set, varied with the
choice of potential function. In particular, for a ResNet-18 network (with 11
million weights) trained on the CIFAR-10 dataset for the different potentials
$\ell_{1}$, $\ell_{2}$, $\ell_{3}$, and $\ell_{10}$-norms, the generalization
error varied quite noticeably. Surprisingly, the $\ell_{1}$ regularizer
yielded the worst generalization performance and the $\ell_{10}$ regularizer
the best. The $\ell_{2}$ (corresponding to SGD) and $\ell_{3}$ regularizers
straddled a midway generalization performance. This paper is concerned with
studying the generalization error of SMD for different potential functions. As
a first step in this direction, we will look at the problem of binary
classification in over-parameterized linear models where the two classes are
drawn from a Gaussian mixture model. In this setting, we obtain the
generalization error of SMD for general potentials and study them in more
detail for the $\ell_{2}$ (i.e., SGD) and $\ell_{1}$ cases. We introduce two
data models, one for which SGD outperforms SMD with an $\ell_{1}$ regularizer,
and one for which the reverse happens. In both cases, the empirical results
well match the theory.

The hope is that the results obtained here will guide us to the analysis of
nonlinear over-parameterized models and, ultimately, to understanding the
generalization behavior of deep networks under different training algorithms.

The remainder of the paper is organized as follows. Section 2 gives some
preliminary descriptions of the SMD algorithm and its implicit regularization
property, introduces the binary classification problem for Gaussian mixtures,
reviews the CGMT framework (the main tool used for our analysis), and
introduces the two explicit data models that will be studied and analyzed.
Section 3 gives general expressions for the generalization error for linear
classification of binary Gaussian mixture models which are the main results of
the paper. Section 4 gives explicit expressions for the specific models
considered and Section 5 gives numerical results collaborating the theory and
showcasing the relative merits of SGD and $\ell_{1}$-SMD. The paper concludes
with Section 6.

##  2 Preliminaries

In this section, we provide a brief overview of SMD, of the binary
classification model we will be studying, and the of Convex Gaussian Min-max
Theorem (CGMT) which is fundamental to our analysis.

###  2.1 Stochastic Mirror Descent

Let $L(w)$ be a separable loss function of some unknown weight (parameter)
vector $w\in\mathbb{R}^{d}$, $L(w)=\sum_{i=1}^{n}L_{i}(w)$, where $L_{i}(w)$
is usually denoted as local loss function. To minimize $L(\cdot)$ over $w$,
one can use classical methods such as SGD, which iteratively updates the
weight vector estimate along the negative direction of the instantaneous
gradient [Robbins and Monro(1951)]. SMD is a family of optimization algorithms
which includes SGD as a special case [Nemirovski and Yudin(1983)]. SMD uses a
strictly convex differentiable potential function $\psi(\cdot)$ such that the
weight vector updates are done in the “mirrored” domain determined by
$\nabla\psi(\cdot)$

$\nabla\psi(w_{i})=\nabla\psi(w_{i-1})-\eta\nabla L_{i}(w_{i-1}),\quad i\geq
1,$ (1)

where $\eta>0$ is the learning rate. Due to strict convexity,
$\nabla\psi(\cdot)$ defines an invertible transformation. It is designed to
exploit the geometrical structure of the optimization problem with the
appropriate choice of potential function. In particular, the update rule in
eq:MD can be equivalently written as

$w_{i}=\mbox{arg}\min_{w}D_{\psi}(w,w_{i-1})+\eta w_{t}^{\top}\nabla
L(w_{i-1}),$ (2)

where $D_{\psi}(\cdot,\cdot)$ is the Bregman divergence with respect to
$\psi(\cdot)$:

$D_{\psi}(w,w_{i-1})=\psi(w)-\psi(w_{i-1})-\nabla\psi(w_{i-1})^{\top}(w-w_{i-1}).$

Note that $D_{\psi}(\cdot,\cdot)$ is non-negative, convex in its first
argument and $D_{\psi}(w,w^{\prime})=0$ iff $w=w^{\prime}$, due to strict
convexity. Due to this construction, different choices of the potential
function $\psi(\cdot)$ yield different optimization algorithms, e.g.
$\psi(w)=\frac{1}{2}\|w\|_{2}^{2}$ gives SGD.

Recently, an array of works has documented and studied the implicit
regularization induced by the MD and SMD algorithms used for optimization
[Gunasekar et al.(2018b)Gunasekar, Lee, Soudry, and Srebro, Azizan et
al.(2020)Azizan, Lale, and Hassibi, Gunasekar et al.(2018a)Gunasekar, Lee,
Soudry, and Srebro, Azizan and Hassibi(2019a), Azizan et al.(2021)Azizan,
Lale, and Hassibi, Azizan et al.(2022)Azizan, Lale, and Hassibi]. These works
considered the setting of modern learning problems which are highly
overparameterized, i.e., the number of parameters are significantly larger
than the number of training data points. In particular, they consider a
training set $\mathcal{D}=\\{(x_{i},y_{i}):i=1,\dots,n\\}$ where
$x_{i}\in\mathbb{R}^{d}$ are the inputs, and $y_{i}\in\mathbb{R}$ are the
outputs obtained from an underlying distribution. The learning problem is to
fit a model $f(x_{i},w)$ (linear or nonlinear) that explains the data in
$\mathcal{D}$ with some unknown weight vector $w\in\mathbb{R}^{d}$. In the
overparameterized (interpolating) regime, the problem setting often has $d\gg
n$, which results in a manifold of (uncountably infinitely many) solutions,
$\mathcal{W}$ that interpolate the training data, i.e.,
$\mathcal{W}=\\{w^{\prime}\in\mathbb{R}^{d}\ |\
f(x_{i},w^{\prime})=y_{i},(x_{i},y_{i})\in\mathcal{D}\\}$.

Defining a loss function on individual data points
$L_{i}(w)=\ell(y_{i}-f(x_{i},w))$, for some differentiable non-negative
function $\ell(\cdot)$ with $\ell(0)=0$, the aforementioned works showed that
SMD converges to the solution of

$\aligned&\underset{w}{\text{min}}D_{\psi}(w,w_{0})\\\
\text{s.t.}y_{i}=f(x_{i},w),\quad i=1,\dots,n,$ (3)

for any initialization $w_{0}$ if $f$ is linear, i.e., $f(x,w)=x^{\top}w$,
while for a nonlinear $f$, SMD converges to a point on $\mathcal{W}$ which is
very close to the solution of opt_implicit_bregman. Further, if
$w_{0}=\argmin_{w}\psi(w)$, SMD solves the problem in opt_implicit_bregman for
$\psi(w)$ instead of $D_{\psi}(w,w_{0})$. This implicit regularization is
clearly observed in practice (see [Azizan et al.(2021)Azizan, Lale, and
Hassibi]) where the solutions of SMD display significantly different
generalization performance on the unseen data. The relationship between the
potential chosen for SMD and the generalization error is unclear. In this
work, we take a step towards understanding this relationship by focusing on
the simpler case of linear over-parametrized models and considering a
particular binary classification problem that we describe below.

###  2.2 Binary Classification for a Gaussian Mixture Model

We consider a binary classification problem with two classes, where for class
1 the feature vector $x\in\mathbb{R}^{d}$ is drawn at random from
$\mathcal{N}(\mu_{1},\Sigma_{1})$, with $\mu_{1}\in\mathbb{R}^{d}$ the mean
and $\Sigma_{1}\in\mathbb{R}^{d\times d}$ the covariance matrix, and where the
label is chosen as $y=1$. Similarly, for class 2 the regressor is drawn from
$\mathcal{N}(\mu_{2},\Sigma_{2})$ and has label $y=-1$.

We will consider a linear classifier given by a weight vector
$w\in\mathbb{R}^{d}$. In other words for a given feature vector $x$, we will
declare that $x$ belongs to class 1 if $x^{T}w>0$ and to class 2 if
$x^{T}w<0$. It is then straightforward to show the following result.

######  Lemma 2.1.

Given a weight vector $w$, and assuming the feature vectors are equally likely
to be drawn from class 1 or class 2, the corresponding generalization error
for the Gaussian mixture model with means $\mu_{1}$ and $\mu_{2}$ and
covariance matrices $\Sigma_{1},\Sigma_{2}$ is given by

$E(w)=\frac{1}{2}Q(\frac{\mu_{1}^{T}w}{\sqrt{w^{T}\Sigma_{1}w}})+\frac{1}{2}Q(-\frac{\mu_{2}^{T}w}{\sqrt{w^{T}\Sigma_{2}w}})$

where $Q(\cdot)$ is the integral of the tail of the standard normal
distribution.

######  Proof 2.2.

See the Appendix.

Now assume half the training data is drawn from class 1, i.e.,

$x_{i}\sim\mathcal{N}(\mu_{1},\Sigma_{1}),~{}~{}~{}y_{i}=1,~{}~{}~{}~{}~{}i=1,\ldots\frac{n}{2}$

and the other half from class 2:

$x_{i}\sim\mathcal{N}(\mu_{2},\Sigma_{2}),~{}~{}~{}y_{i}=-1,~{}~{}~{}~{}~{}i=\frac{n}{2}+1,\ldots
n$

Since we are in the over-parametrized regime, we will assume that $n<d$. In
addition, we will assume that both $n,d\rightarrow\infty$.

Defining the matrix of features and the vector of labels

$X=\left[\begin{array}[]{cccc}x_{1}&x_{2}&\ldots&x_{n}\end{array}\right]~{}~{}~{},~{}~{}~{}y=\left[\begin{array}[]{c}1_{\frac{n}{2}}\\\
-1_{\frac{n}{2}}\end{array}\right]$

where $1_{\frac{n}{2}}$ is the all-one vector of length $\frac{n}{2}$, it is
easy to see that, if initialized with $w_{0}=\argmin_{w}\psi(w)$, SMD returns
the weight vector that solves

$\aligned&\underset{w}{\text{min}}\psi(w)\\\ \text{s.t.}X^{T}w=y$ (4)

In other words, SMD returns a weight vector $w$ that minimizes the potential
$\psi(\cdot)$ among all weight vectors that interpolate the training data.

The goal of this paper is to compute and characterize the generalization error
of SMD using different potentials for the linear binary classifier with
Gaussian mixture model. As can be seen from Lemma 2.1, this requires us to
characterize the four quantities

$\mu_{1}^{T}w~{}~{},~{}~{}\mu_{2}^{T}w~{}~{},~{}~{}w^{T}\Sigma_{1}w~{}~{},~{}~{}w^{T}\Sigma_{2}w$

In fact, in much of the subsequent analysis, we shall assume
$\Sigma_{1}=\sigma_{1}^{2}I$ and $\Sigma_{1}=\sigma_{1}^{2}I$, which implies
we need only characterize the following three quantities

$\mu_{1}^{T}w~{}~{},~{}~{}\mu_{2}^{T}w~{}~{},~{}~{}\|w\|^{2}.$

Since the data model that we are considering is a Gaussian mixture, we shall
make use of the Convex Gaussian Min-Max Theorem (CGMT) [Thrampoulidis et
al.(2015b)Thrampoulidis, Oymak, and Hassibi], which is a tight and extended
version of a classical Gaussian comparison inequality [Gordon(1985)].

###  2.3 Convex Gaussian Min-max Theorem (CGMT)

The CGMT framework has been developed to analyze the properties of the
solutions to non-smooth regularized convex optimization problems and has been
successfully applied to characterize the precise performance in numerous
applications such as $M$-estimators, generalized lasso, massive MIMO, phase
retrieval, regularized logistic regression, adversarial training, and max-
margin classifiers [Stojnic(2013), Thrampoulidis et al.(2018)Thrampoulidis,
Abbasi, and Hassibi, Salehi et al.(2019)Salehi, Abbasi, and Hassibi,
Thrampoulidis et al.(2015a)Thrampoulidis, Abbasi, and Hassibi, Abbasi et
al.(2019)Abbasi, Salehi, and Hassibi, Salehi et al.(2018)Salehi, Abbasi, and
Hassibi, Miolane and Montanari(2021), Taheri et al.(2021)Taheri, Pedarsani,
and Thrampoulidis, Aubin et al.(2020)Aubin, Krzakala, Lu, and Zdeborová,
Javanmard and Soltanolkotabi(2022), Montanari et al.(2019)Montanari, Ruan,
Sohn, and Yan, Salehi et al.(2020)Salehi, Abbasi, and Hassibi]. In this
framework, a given challenging optimization problem denoted as the primary
optimization (PO) problem, is associated with a simplified auxiliary
optimization (AO) problem from which the optimal solution can be tightly
inferred. Specifically, the (PO) and (AO) problems are defined as follows:

$\displaystyle\Phi(\mathbf{G})$
$\displaystyle:=\min_{\mathbf{w}\in\mathcal{S}_{\mathbf{w}}}\max_{\mathbf{u}\in\mathcal{S}_{\mathbf{u}}}\mathbf{u}^{\top}\mathbf{Gw}+\psi(\mathbf{w},\mathbf{u})\\!\\!$
(PO) $\displaystyle\phi(\mathbf{g},\mathbf{h})$
$\displaystyle\\!:=\\!\\!\min_{\mathbf{w}\in\mathcal{S}_{\mathbf{w}}}\max_{\mathbf{u}\in\mathcal{S}_{\mathbf{u}}}\|\mathbf{w}\|_{2}\mathbf{g}^{\top}\mathbf{u}\\!+\\!\|\mathbf{u}\|_{2}\mathbf{h}^{\top}\mathbf{w}\\!+\\!\psi(\mathbf{w},\mathbf{u})\\!\\!$
(AO)

where $\mathbf{G}\in\mathbb{R}^{m\times
n},\mathbf{g}\in\mathbb{R}^{m},\mathbf{h}\in\mathbb{R}^{n},\mathcal{S}_{\mathbf{w}}\subset\mathbb{R}^{n},\mathcal{S}_{\mathbf{u}}\subset\mathbb{R}^{m}$
and $\psi:\mathbb{R}^{n}\times\mathbb{R}^{m}\rightarrow\mathbb{R}$. Denoting
any optimal minimizers of (PO) and (AO) as
$\mathbf{w}_{\Phi}:=\mathbf{w}_{\Phi}(\mathbf{G})$ and
$\mathbf{w}_{\phi}:=\mathbf{w}_{\phi}(\mathbf{g},\mathbf{h})$, respectively,
CGMT result states the following.

######  Theorem 2.3 (CGMT [Thrampoulidis et al.(2018)Thrampoulidis, Abbasi,
and Hassibi]).

In (PO) and (AO), let $\mathcal{S}_{\mathbf{w}},\mathcal{S}_{\mathbf{u}}$ be
convex compact sets, $\psi$ be continuous and convex-concave on
$\mathcal{S}_{\mathbf{w}}\times\mathcal{S}_{\mathbf{u}}$, and,
$\mathbf{G},\mathbf{g}$ and $\mathbf{h}$ all have entries iid standard normal.
Let $\mathcal{S}$ be an arbitrary open subset of $\mathcal{S}_{\mathbf{w}}$
and $\mathcal{S}^{c}:=\mathcal{S}_{\mathbf{w}}\setminus\mathcal{S}$. Denote by
$\Phi_{\mathcal{S}^{c}}(\mathbf{G})$ and
$\phi_{\mathcal{S}^{c}}(\mathbf{g},\mathbf{h})$ the optimal costs of (PO) and
(AO) respectively when $\mathbf{w}$ is minimized over $\mathcal{S}^{c}$. If
there exist constants $\bar{\phi}<\bar{\phi}_{\mathcal{S}^{c}}$ such that
$\phi(\mathbf{g},\mathbf{h})\stackrel{{\scriptstyle
p}}{{\longrightarrow}}\bar{\phi}$, and
$\phi_{\mathcal{S}^{c}}(\mathbf{g},\mathbf{h})\stackrel{{\scriptstyle
p}}{{\longrightarrow}}\bar{\phi}_{\mathcal{S}^{c}}$, (converge in
probability), then
$\lim_{n\rightarrow\infty}\mathbb{P}\left(\mathbf{w}_{\Phi}(\mathbf{G})\in\mathcal{S}\right)=1$.

The probabilities in Theorem 2.3 are with respect to the randomness of
$\mathbf{G},\mathbf{g}$, and $\mathbf{h}$. Notice that from the assumptions in
the theorem statement, we know that
$\mathbf{w}_{\phi}(\mathbf{g},\mathbf{h})\in\mathcal{S}$ with probability
approaching $1$ for the (AO) problem. However, Theorem 2.3 gives a stronger
result and concludes the same characterization for the solution of the
seemingly different optimization problem (PO). Appropriate choices of $S$ then
allow us to conclude that the desired values of $\mu_{1}^{T}w$,
$\mu_{2}^{T}w$, $\|w\|^{2}$ for (PO) concentrate in the same domain as the
same values for (AO) as well as that empirical distributions of $w$ coincide
for (PO) and (AO) provided that $n$ is big enough ($n\geq 100$ suffices in
practice). In our analysis, we use it to characterize the empirical
distributions of weights identified by SMD algorithms and to determine the
desired generalization errors.

###  2.4 Two Explicit Models

As we shall subsequently see, the performance of SMD for various potentials
will highly depend on the parameters $\mu_{1}$, $\mu_{2}$, $\Sigma_{1}$,
$\Sigma_{2}$ of the Gaussian mixture model. In what follows we shall consider
the following two explicit models.

  * •

Model 1: In this model, we will assume that $\mu_{1}$ is an iid standard
normal vector, $\mu_{2}=\sqrt{1-\epsilon^{2}}\mu_{1}+\epsilon v$, where $v$ is
another independent iid standard normal vector. This implies that both mean
vectors have a length (roughly) equal to $\sqrt{d}$ and a relative angle

$\theta=\cos^{-1}(\sqrt{1-\epsilon^{2}})=\sin^{-1}(\epsilon).$

For simplicity, we will further take $\sigma_{1}=\sigma_{2}=1$ and fix
$\epsilon=0.1$.

The parameter $\epsilon$ will allow us to control the angle between the two
mean vectors and thereby the difficulty in separating the two classes. The
difference between the two mean vectors is spread homogeneously across the
entries of the vectors. As we shall subsequently see, through the theory and
empirical results, this model gives better generalization results on SGD than
$\ell_{1}$-SMD.

  * •

Model 2: In this model, we take $\mu_{1i}=\mu_{2i}$ and iid standard normal
for $i>1$, and $\mu_{11}=-\mu_{2i}=t=2$, $\sigma_{1}=\sigma_{2}=1$. In other
words, the mean vectors of the two classes differ in only a single component.
As expected, linear classification for this model is much more conducive to a
sparsifying regularizer and both the theory and empirical results will show
that, in terms of generalization performance, $\ell_{1}$-SMD significantly
outperforms SGD.

###  2.5 Some Useful Lemmas

The following lemmas will be of use for making calculations specific to models
$1$ and $2$.

######  Lemma 2.4.

Let $X\sim\mathcal{N}(0,\sigma^{2})$. Then
$\mathbb{E}[(|X|-1)^{2}\mathbbm{1}_{|X|>1}]=2(\sigma^{2}+1)Q(\frac{1}{\sigma})-\frac{2\sigma}{\sqrt{2\pi}}e^{-\frac{1}{2\sigma^{2}}}$.

######  Proof 2.5.

See Appendix.

######  Lemma 2.6.

The following equality holds for any $x>0$:

$\sqrt{x}=\min_{\beta>0}\frac{1}{2\beta}+\frac{\beta x}{2}$

######  Proof 2.7.

See Appendix.

##  3 Main Results

We begin with a theorem that holds for arbitrary mirror $\psi$. It will later
be used and specialized to study $\ell_{1}$-SMD and SGD.

######  Theorem 3.1.

The empirical distribution of the parameters identified by SMD with a mirror
$\psi$ applied to the Gaussian mixture model with means $\mu_{1},\mu_{2}$ and
covariance matrices $\sigma_{1}^{2}I,\sigma_{2}^{2}I$ matches the empirical
distribution of $\hat{w}$ obtained by solving the following optimization
problem for $w$:

$\max_{\alpha\geq 0}\min_{w,\beta\geq
0}\max_{\gamma_{1},\gamma_{2}}\psi(w)+\alpha
g^{T}w+\frac{\alpha}{2\beta}+\frac{\alpha\beta
n}{2}\|w\|_{2}^{2}+\frac{\gamma_{1}(\mu_{1}^{T}w-1)}{\sigma_{1}}+\frac{\gamma_{2}(\mu_{2}^{T}w+1)}{\sigma_{2}}-\frac{\gamma_{1}^{2}+\gamma_{2}^{2}}{\alpha\beta
n}$

where $g\in\mathbb{R}^{d}$ is a vector of iid standard normal entries. The
values of $\|w\|^{2}$, $\mu_{1}^{T}w$ and $\mu_{2}^{T}w$ inferred from this
optimization problem coincide with the same values for parameters found by
SMD.

######  Proof 3.2.

See Appendix.

The point of the above theorem is that once the distribution of $w$ is
identified from the optimization, the quantities $\|w\|^{2}$, $\mu_{1}^{T}w$,
and $\mu_{2}^{T}w$, necessary to obtain the generalization error, can be
computed and the histogram for $w$ will recover the weight histogram for SMD.

The next two theorems arise as applications of Theorem 3.1 when we specialize
to $\psi(w)=\|w\|_{2}^{2}$ and $\psi(w)=\|w\|_{1}$ and study the optimization
problem in more detail. These more detailed analyses will allow us to
determine the generalization errors of SGD and $\ell_{1}$-SMD, respectively.

######  Theorem 3.3.

The empirical distribution of the parameters identified by SGD applied to the
Gaussian mixture model with means $\mu_{1},\mu_{2}$ and covariance matrices
$\sigma_{1}^{2}I,\sigma_{2}^{2}I$ matches the empirical distribution of
$\hat{w}$ given by

$\displaystyle\hat{w}$ $\displaystyle=$ $\displaystyle-\frac{\alpha
g}{2+\alpha\beta n}+\frac{\alpha\beta n}{4\Delta}(\frac{\alpha\beta
n}{4}(\|\mu_{2}\|_{2}^{2}+\mu_{1}^{T}\mu_{2})+\sigma_{1}^{2}(\frac{\alpha\beta
n}{2}+1))\mu_{1}$ $\displaystyle-\frac{\alpha\beta
n}{4\Delta}(\frac{\alpha\beta
n}{4}(\|\mu_{1}\|_{2}^{2}+\mu_{1}^{T}\mu_{2})+\sigma_{2}^{2}(\frac{\alpha\beta
n}{2}+1))\mu_{2}$

where $g\in\mathbb{R}^{d}$ is a vector of iid standard normal entries and
$\alpha$ and $\beta$ are defined as solutions to the following two-dimensional
scalar optimization problem:

$\max_{\alpha\geq 0}\min_{\beta\geq
0}-\frac{\alpha^{2}d}{4(1+\frac{\alpha\beta n}{2})}-\frac{(\alpha\beta
n)^{2}}{16\Delta}(\|\mu_{1}\|^{2}+\|\mu_{2}\|^{2}-(\frac{\sigma_{1}^{2}}{\sigma_{2}^{2}}+\frac{\sigma_{2}^{2}}{\sigma_{1}^{2}})\mu_{2}^{T}\mu_{1})-$
$-\frac{(\alpha\beta
n)^{3}}{32\Delta}(\|\mu_{1}\|^{2}+\|\mu_{2}\|^{2}-(\frac{\sigma_{1}^{2}}{\sigma_{2}^{2}}+\frac{\sigma_{2}^{2}}{\sigma_{1}^{2}})\mu_{2}^{T}\mu_{1}+(\frac{1}{2\sigma_{1}^{2}}+\frac{1}{2\sigma_{2}^{2}})(\|\mu_{1}\|^{2}\|\mu_{2}\|^{2}-(\mu_{2}^{T}\mu_{1})^{2}))$

Here $\Delta$ is also a function of $\alpha$ and $\beta$ and is defined as

$\displaystyle\Delta$ $\displaystyle=$ $\displaystyle(\frac{\alpha\beta
n}{4})^{2}(4\sigma_{1}^{2}\sigma_{2}^{2}+\|\mu_{1}\|^{2}\|\mu_{2}\|^{2}-(\mu_{2}^{T}\mu_{1})^{2}+2(\sigma_{1}^{2}\|\mu_{1}\|^{2}+\sigma_{2}^{2}\|\mu_{2}\|^{2}))$
$\displaystyle+\frac{\alpha\beta
n}{4}(2\sigma_{1}^{2}+2\sigma_{2}^{2}+\sigma_{1}^{2}\|\mu_{1}\|^{2}+\sigma_{2}^{2}\|\mu_{2}\|^{2})+\sigma_{1}^{2}\sigma_{2}^{2}$

The values of $\|w\|^{2}$, $\mu_{1}^{T}w$ and $\mu_{2}^{T}w$ inferred from
this optimization problem coincide with the same values for parameters found
by SGD.

######  Proof 3.4.

See Appendix.

Note that ${\hat{w}}$ is simply a non-zero mean Gaussian vector and so
$\|w\|^{2}$, $\mu_{1}^{T}w$, and $\mu_{2}^{T}w$ can be readily computed,
thereby allowing the evaluation of the generalization error via Lemma 2.1.

######  Theorem 3.5.

The empirical distribution of the parameters identified by $\ell_{1}$-SMD
applied to the Gaussian mixture model with means $\mu_{1},\mu_{2}$ and
covariance matrices $\sigma_{1}^{2}I,\sigma_{2}^{2}I$ matches the empirical
distribution of $\hat{w}$ given by

$\hat{w}_{i}=-(\alpha\beta
n)^{-1}sign(\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i}+\alpha
g_{i})\max(0,|\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}+\frac{\gamma_{2}}{\sigma_{2}}\gamma_{2}\mu_{2i}+\alpha
g_{i}|-1),$

where the $g_{i}$ are iid standard normal and $\gamma_{1},\gamma_{2},\alpha$
and $\beta$ are defined as solutions of the following four-dimensional
optimization problem:

$\max_{\alpha\geq 0}\min_{\beta\geq
0}\max_{\gamma_{1},\gamma_{2}}\frac{\gamma_{2}}{\sigma_{2}}-\frac{\gamma_{1}}{\sigma_{1}}+\frac{\alpha}{2\beta}-\frac{\gamma_{1}^{2}+\gamma_{2}^{2}}{\alpha\beta
n}-\sum_{i}\frac{|\alpha
g_{i}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i}|-1}{2\alpha\beta
n}\max(0,|\alpha
g_{i}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i}|-1)$

The values of $\|w\|^{2}$, $\mu_{1}^{T}w$ and $\mu_{2}^{T}w$ inferred from
this optimization problem coincide with the same values for parameters found
by $\ell_{1}$\- SMD.

######  Proof 3.6.

See Appendix.

We should comment that in Theorem 3.3 the optimization for the parameters
$\alpha$ and $\beta$ is deterministic. However, in Theorem 3.5, the
optimization for $\alpha$, $\beta$, $\gamma_{1}$, and $\gamma_{2}$ is
stochastic. However, if we make some statistical assumptions on the means
$\mu_{1}$ and $\mu_{2}$, then by the law of large numbers the term

$\sum_{i}\frac{|\alpha
g_{i}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i}|-1}{2\alpha\beta
n}\max(0,|\alpha
g_{i}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i}|-1)$

will concentrate. In fact, this is why we consider the explicit Models 1 and 1
described earlier.

##  4 Specific Results

We now specialize the previous theorems to the Models 1 and 2 described
earlier. This will allow us to get explicit expressions for the generalization
error and to compare the performances of SGD and $\ell_{1}$-SMD.

###  4.1 Model 1

Recall here that $\mu_{1}$ is standard normal,
$\mu_{2}=\sqrt{1-\epsilon^{2}}\mu_{1}+\epsilon v$, where $v$ is an independent
standard normal vector. We further assume $\sigma_{1}=\sigma_{2}=1$.

######  Lemma 4.1.

Denote
$\sigma^{2}=\gamma_{1}^{2}+\gamma_{2}^{2}+\alpha^{2}+2\gamma_{1}\gamma_{2}\sqrt{1-\epsilon^{2}}$.
The following equality holds:

$\mathbb{E}_{\mu_{1i},v_{i},g_{i}}[(|\gamma_{1}\mu_{1i}+\gamma_{2}\mu_{2i}+\alpha
g_{i}|-1)\max(0,|\gamma_{1}\mu_{1i}+\gamma_{2}\mu_{2i}+\alpha
g_{i}|-1)]=2(\sigma^{2}+1)Q(\frac{1}{\sigma})-\frac{2\sigma}{\sqrt{2\pi}}e^{-\frac{1}{2\sigma^{2}}}$

######  Proof 4.2.

Note that

$\gamma_{1}\mu_{1i}+\gamma_{2}\mu_{2i}+\alpha
g_{i}=(\gamma_{1}+\sqrt{1-\epsilon^{2}}\gamma_{2})\mu_{1i}+\epsilon\gamma_{2}v+\alpha
g_{i}\sim\mathcal{N}(0,\gamma_{1}^{2}+\gamma_{2}^{2}+\alpha^{2}+2\gamma_{1}\gamma_{2}\sqrt{1-\epsilon^{2}})$

Denote $X=\gamma_{1}\mu_{1i}+\gamma_{2}\mu_{2i}+\alpha
g_{i}\sim\mathcal{N}(0,\sigma^{2})$.

The initial expectation can then be rewritten in the following way and found
using Lemma 2.4:

$\mathbb{E}[(|X|-1)^{2}\mathbbm{1}_{|X|>1}]=2(\sigma^{2}+1)Q(\frac{1}{\sigma})-\frac{2\sigma}{\sqrt{2\pi}}e^{-\frac{1}{2\sigma^{2}}}$

######  Lemma 4.3.

Under the terminology from Theorem 3.1, the following equalities hold:

$\gamma_{1}=\frac{\alpha\beta
n}{2\sigma_{1}}(\mu_{1}^{T}w-1),\gamma_{2}=\frac{\alpha\beta
n}{2\sigma_{2}}(\mu_{2}^{T}w+1)$

######  Proof 4.4.

Follows immediately from taking derivatives by $\gamma_{1}$ and $\gamma_{2}$
and equating them to $0$.

######  Remark 4.5.

Since we work in the asymptotic regime $d\to\infty$, we can replace

$\sum_{i}\frac{(|\gamma_{1}\mu_{1i}+\gamma_{2}\mu_{2i}+\alpha
g_{i}|-1)}{2\alpha\beta n}\max(0,|\gamma_{1}\mu_{1i}+\gamma_{2}\mu_{2i}+\alpha
g_{i}|-1)$

from the objective of Theorem 3.5 by $\frac{d(\sigma^{2}+1)}{\alpha\beta
n}Q(\frac{1}{\sigma})-\frac{d\sigma}{\sqrt{2\pi}\alpha\beta
n}e^{-\frac{1}{2\sigma^{2}}}$. Note that this expression is invariant to the
transformation $(\gamma_{1},\gamma_{2})\to(-\gamma_{2},-\gamma_{1})$ and so is
the rest of the objective from Theorem 3.5. Since this objective is strictly
concave in $\gamma_{1}$ and $\gamma_{2}$, we conclude that the optimal
parameters must satisfy $\gamma_{2}=-\gamma_{1}$. All this being said, we
experiment with the following three-dimensional optimization problem in the
numerical part of the work related to model 1:

$\max_{\alpha\geq 0}\min_{\beta\geq
0}\max_{\gamma_{1}}\frac{\alpha}{2\beta}-\frac{2\gamma_{1}^{2}}{\alpha\beta
n}-2\gamma_{1}-\frac{d(\sigma^{2}+1)}{\alpha\beta
n}Q(\frac{1}{\sigma})+\frac{d\sigma}{\sqrt{2\pi}\alpha\beta
n}e^{-\frac{1}{2\sigma^{2}}}$
$\text{where}\sigma^{2}=2\gamma_{1}^{2}(1-\sqrt{1-\epsilon^{2}})+\alpha^{2}$

The same lemma suggests us an approximation for $\|w\|_{2}^{2}$, since

$w_{i}^{2}=\frac{(|\gamma_{1}\mu_{1i}+\gamma_{2}\mu_{2i}+\alpha
g_{i}|-1)^{2}}{(\alpha\beta
n)^{2}}\mathbbm{1_{|\gamma_{1}\mu_{1i}+\gamma_{2}\mu_{2i}+\alpha g_{i}|>1}}$

we will approximate $\|w\|_{2}^{2}=\sum_{i}w_{i}^{2}$ as

$\|w\|_{2}^{2}\approx\frac{2d(\sigma^{2}+1)}{(\alpha\beta
n)^{2}}Q(\frac{1}{\sigma})-\frac{2d\sigma}{\sqrt{2\pi}(\alpha\beta
n)^{2}}e^{-\frac{1}{2\sigma^{2}}}.$

Finally, we find dot products $\mu_{1}^{T}w$ and $\mu_{2}^{T}w$ using Lemma
4.3:

$\mu_{1}^{T}w=\frac{2\gamma_{1}}{\alpha\beta
n}+1,\mu_{2}^{T}w=-\frac{2\gamma_{1}}{\alpha\beta n}-1$

###  4.2 Model 2

Recall, as before that $\mu_{1i}=\mu_{2i}$ is iid standard normal for $i>1$
and that $\mu_{11}=-\mu_{2i}=t=2$. For simplicity, we take
$\sigma_{1}=\sigma_{2}=1$.

######  Remark 4.6.

Analogously to Remark 4.5, we will use Lemma 4.1 to approximate the objective
by a simpler expression. In this case, the first term of the sum

$\sum_{i}\frac{(|\gamma_{1}\mu_{1i}+\gamma_{2}\mu_{2i}+\alpha
g_{i}|-1)}{2\alpha\beta n}\max(0,|\gamma_{1}\mu_{1i}+\gamma_{2}\mu_{2i}+\alpha
g_{i}|-1)$

cannot be replaced by anything rather than
$\frac{(|\gamma_{1}t-\gamma_{2}t+\alpha g_{1}|-1)}{2\alpha\beta
n}\max(0,|\gamma_{1}t-\gamma_{2}t+\alpha g_{1}|-1)$ itself, so we just leave
it this way. Note that for $i>1$ we have
$\gamma_{1}\mu_{1i}+\gamma_{2}\mu_{2i}+\alpha
g_{i}=(\gamma_{1}+\gamma_{2})\mu_{2i}+\alpha
g_{i}\sim\mathcal{N}(0,(\gamma_{1}+\gamma_{2})^{2}+\alpha^{2})$. Thus, we
replace the sum of the other terms by $\frac{(d-1)(\sigma^{2}+1)}{\alpha\beta
n}Q(\frac{1}{\sigma})-\frac{(d-1)\sigma}{\sqrt{2\pi}\alpha\beta
n}e^{-\frac{1}{2\sigma^{2}}}$ following the same reasoning as in Remark 4.5,
where $\sigma^{2}=(\gamma_{1}+\gamma_{2})^{2}+\alpha^{2}$ this time. Note that
this again makes the entire objective invariant to the same transformation
$(\gamma_{1},\gamma_{2})\to(-\gamma_{2},-\gamma_{1})$. We conclude that
$\gamma_{2}=-\gamma_{1}$ and $\sigma=\alpha$, which leads us to:

$\max_{\alpha\geq 0}\min_{\beta\geq
0}\max_{\gamma_{1}}\frac{\alpha}{2\beta}-\frac{2\gamma_{1}^{2}}{\alpha\beta
n}-2\gamma_{1}-\frac{(|2\gamma_{1}t+\alpha g_{1}|-1)}{2\alpha\beta
n}\max(0,|2\gamma_{1}t+\alpha g_{1}|-1)-$
$-\frac{(d-1)(\alpha^{2}+1)}{\alpha\beta
n}Q(\frac{1}{\alpha})+\frac{(d-1)}{\sqrt{2\pi}\beta
n}e^{-\frac{1}{2\alpha^{2}}}$

Again analogously to Remark 4.5, we obtain:

$\|w\|_{2}^{2}\approx\frac{(|2\gamma_{1}t+\alpha g_{1}|-1)^{2}}{(\alpha\beta
n)^{2}}\mathbbm{1}_{|2\gamma_{1}t+\alpha
g_{1}|>1}+\frac{2(d-1)(\alpha^{2}+1)}{(\alpha\beta
n)^{2}}Q(\frac{1}{\alpha})-\frac{2(d-1)}{\sqrt{2\pi}\alpha(\beta
n)^{2}}e^{-\frac{1}{2\alpha^{2}}}$
$\mu_{1}^{T}w=\frac{2\gamma_{1}}{\alpha\beta
n}+1,\mu_{2}^{T}w=-\frac{2\gamma_{1}}{\alpha\beta n}-1$

##  5 Numerical simulations

This section provides a comparison between classification errors obtained by
training linear models using SGD and $\ell_{1}$-SMD and evaluating the
corresponding performances empirically to the classification errors predicted
by CGMT. We used code from a publically available repository
https://github.com/SahinLale/StochasticMirrorDescent provided by authors of
[Azizan et al.(2021)Azizan, Lale, and Hassibi] with minor changes for
training. CGMT predictions were calculated numerically by solving the
corresponding optimization problems via a grid search and then using Remarks
4.5 and 4.6 along with Lemma 2.1 to evaluate the error. In the tables
presented below, CGMT $\ell_{1}$-SMD stands for the classification error
predicted by CGMT for $\ell_{1}$ stochastic mirror descent, empirical
$\ell_{1}$ stands for the test error evaluated for a trained $\ell_{1}$-SMD
initialized near $0$, CGMT SGD and empirical SGD signify the same values, but
for SGD. The prediction of $\ell_{1}$-SMD and the empirical results for
$\ell_{1}$-SMD appear to not depend too dramatically on the realizations of
$\mu_{1}$ and $\mu_{2}$. That is, they seem to be well-concentrated for model
$1$. However, they were less so for model $2$, so we averaged both over $5$
evaluations each. The prediction of CGMT SGD and the empirical performance of
SGD were observed to be well-concentrated for both models. As the reader can
see, the match between the empirical and CGMT-predicted SGD performances is
better than between the same quantities for $\ell_{1}$-SMD. We believe that
this arises because the latter is more challenging numerically, as the
corresponding expression for CGMT involves a $3$ \- dimensional optimization
instead of $2$ -dimensional and is more sensitive to parameter changes. Apart
from that, evaluating the performance of $\ell_{1}$-SMD empirically is also
more challenging because it requires more iterations to converge. That is,
there always is a chance that the generalization errors could match more
closely if the algorithm was run for more iterations. In either event, the
match between the theoretical and empirical generalization errors is quite
good in all cases.

As can be seen from the Tables, for Model 1, SGD has slightly superior
performance compared to $\ell_{1}$-SMD. This is reasonable, since the
difference between the two mean vectors is spread homogeneously across the
entries of the vectors.

However, for Model 2, $\ell_{1}$-SMD has significantly better performance.
Again, this is expected because the mean vectors of the two classes differ in
only a single component. Therefore linear classification for this model is
much more conducive to a sparsifying regularizer.

These results clearly demonstrate that the generalization performance of
linear classifiers on binary Gaussian mixture models tangibly depends on the
mirror used by the training algorithm and on the model the data obeys. We
believe this general principle to hold for deep networks as well, although it
will merit a much more difficult and detailed analysis.

###  5.1 Model 1

n | d | CGMT $\ell_{1}$ | Empirical $\ell_{1}$ | CGMT SGD | Empirical SGD
---|---|---|---|---|---
500 | 1000 | 0.242 | 0.275 | 0.202 | 0.191
200 | 1000 | 0.315 | 0.306 | 0.199 | 0.194
100 | 1000 | 0.370 | 0.346 | 0.253 | 0.249
1000 | 10000 | 0.023 | 0.012 | 0 | 0

###  5.2 Model 2

n | d | CGMT $\ell_{1}$ | Empirical $\ell_{1}$ | CGMT SGD | Empirical SGD
---|---|---|---|---|---
100 | 1000 | 0.056 | 0.059 | 0.155 | 0.152
1000 | 10000 | 0.045 | 0.051 | 0.152 | 0.150
500 | 10000 | 0.048 | 0.034 | 0.211 | 0.218

##  6 Conclusion

In this paper we studied the problem of linear classification of binary
Gaussian mixture models using SMD training with general potentials. Using a
CGMT analysis we are able to find expressions for the generalization error.
Numerical simulations show a good agreement between the theory and empirical
results. In particular, we observe that the generalization performance depends
heavily on the mirror used in SMD, as well as on the data model. We exhibited
two models, one for which SGD was superior and one for which $\ell_{1}$-SMD is
so. There are several directions in which these results can be extended. One
is to find explicit expressions for the generalization performance of other
mirrors, most notably $\ell_{\infty}$. Another is to extend the classification
problem beyond the binary case to more complicated Gaussian mixtures. Finally,
we consider the work performed here to be a small step in the direction of
understanding the generalization performance of deep networks.

We are grateful to Sahin Lale for helping us with numerical experiments.

## References

  * [Abbasi et al.(2019)Abbasi, Salehi, and Hassibi] Ehsan Abbasi, Fariborz Salehi, and Babak Hassibi.  Performance analysis of convex data detection in mimo.  In _ICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pages 4554–4558. IEEE, 2019.
  * [Aubin et al.(2020)Aubin, Krzakala, Lu, and Zdeborová] Benjamin Aubin, Florent Krzakala, Yue Lu, and Lenka Zdeborová.  Generalization error in high-dimensional perceptrons: Approaching bayes error with convex optimization.  _Advances in Neural Information Processing Systems_ , 33:12199–12210, 2020.
  * [Azizan and Hassibi(2019a)] Navid Azizan and Babak Hassibi.  A characterization of stochastic mirror descent algorithms and their convergence properties.  In _ICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pages 5167–5171. IEEE, 2019a.
  * [Azizan and Hassibi(2019b)] Navid Azizan and Babak Hassibi.  Stochastic gradient/mirror descent: Minimax optimality and implicit regularization.  In _International Conference on Learning Representations (ICLR)_ , 2019b.
  * [Azizan et al.(2020)Azizan, Lale, and Hassibi] Navid Azizan, Sahin Lale, and Babak Hassibi.  A study of generalization of stochastic mirror descent algorithms on overparameterized nonlinear models.  In _ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pages 3132–3136. IEEE, 2020.
  * [Azizan et al.(2021)Azizan, Lale, and Hassibi] Navid Azizan, Sahin Lale, and Babak Hassibi.  Stochastic mirror descent on overparameterized nonlinear models.  _IEEE Transactions on Neural Networks and Learning Systems_ , 33(12):7717–7727, 2021.
  * [Azizan et al.(2022)Azizan, Lale, and Hassibi] Navid Azizan, Sahin Lale, and Babak Hassibi.  Explicit regularization via regularizer mirror descent.  _arXiv preprint arXiv:2202.10788_ , 2022.
  * [Gordon(1985)] Yehoram Gordon.  Some inequalities for gaussian processes and applications.  _Israel Journal of Mathematics_ , 50:265–289, 1985.
  * [Gunasekar et al.(2017)Gunasekar, Woodworth, Bhojanapalli, Neyshabur, and Srebro] Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro.  Implicit regularization in matrix factorization.  In _Advances in Neural Information Processing Systems_ , pages 6152–6160, 2017.
  * [Gunasekar et al.(2018a)Gunasekar, Lee, Soudry, and Srebro] Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro.  Characterizing implicit bias in terms of optimization geometry.  In _International Conference on Machine Learning_ , pages 1827–1836, 2018a.
  * [Gunasekar et al.(2018b)Gunasekar, Lee, Soudry, and Srebro] Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro.  Implicit bias of gradient descent on linear convolutional networks.  _arXiv preprint arXiv:1806.00468_ , 2018b.
  * [Javanmard and Soltanolkotabi(2022)] Adel Javanmard and Mahdi Soltanolkotabi.  Precise statistical analysis of classification accuracies for adversarial training.  _The Annals of Statistics_ , 50(4):2127–2156, 2022.
  * [Miolane and Montanari(2021)] Léo Miolane and Andrea Montanari.  The distribution of the lasso: Uniform control over sparse balls and adaptive parameter tuning.  _The Annals of Statistics_ , 49(4):2313–2335, 2021.
  * [Montanari et al.(2019)Montanari, Ruan, Sohn, and Yan] Andrea Montanari, Feng Ruan, Youngtak Sohn, and Jun Yan.  The generalization error of max-margin linear classifiers: High-dimensional asymptotics in the overparametrized regime.  _arXiv preprint arXiv:1911.01544_ , 2019.
  * [Nemirovski and Yudin(1983)] Arkadii Nemirovski and David Borisovich Yudin.  Problem complexity and method efficiency in optimization.  1983\.
  * [Robbins and Monro(1951)] Herbert Robbins and Sutton Monro.  A stochastic approximation method.  _The annals of mathematical statistics_ , pages 400–407, 1951.
  * [Salehi et al.(2018)Salehi, Abbasi, and Hassibi] Fariborz Salehi, Ehsan Abbasi, and Babak Hassibi.  A precise analysis of phasemax in phase retrieval.  In _2018 IEEE International Symposium on Information Theory (ISIT)_ , pages 976–980. IEEE, 2018.
  * [Salehi et al.(2019)Salehi, Abbasi, and Hassibi] Fariborz Salehi, Ehsan Abbasi, and Babak Hassibi.  The impact of regularization on high-dimensional logistic regression.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * [Salehi et al.(2020)Salehi, Abbasi, and Hassibi] Fariborz Salehi, Ehsan Abbasi, and Babak Hassibi.  The performance analysis of generalized margin maximizers on separable data.  In _International conference on machine learning_ , pages 8417–8426. PMLR, 2020.
  * [Stojnic(2013)] Mihailo Stojnic.  A framework to characterize performance of lasso algorithms.  _arXiv preprint arXiv:1303.7291_ , 2013.
  * [Taheri et al.(2021)Taheri, Pedarsani, and Thrampoulidis] Hossein Taheri, Ramtin Pedarsani, and Christos Thrampoulidis.  Fundamental limits of ridge-regularized empirical risk minimization in high dimensions.  In _International Conference on Artificial Intelligence and Statistics_ , pages 2773–2781. PMLR, 2021.
  * [Thrampoulidis et al.(2015a)Thrampoulidis, Abbasi, and Hassibi] Christos Thrampoulidis, Ehsan Abbasi, and Babak Hassibi.  Lasso with non-linear measurements is equivalent to one with linear measurements.  _Advances in Neural Information Processing Systems_ , 28, 2015a.
  * [Thrampoulidis et al.(2015b)Thrampoulidis, Oymak, and Hassibi] Christos Thrampoulidis, Samet Oymak, and Babak Hassibi.  Regularized linear regression: A precise analysis of the estimation error.  In _Conference on Learning Theory_ , pages 1683–1709. PMLR, 2015b.
  * [Thrampoulidis et al.(2018)Thrampoulidis, Abbasi, and Hassibi] Christos Thrampoulidis, Ehsan Abbasi, and Babak Hassibi.  Precise error analysis of regularized $m$-estimators in high dimensions.  _IEEE Transactions on Information Theory_ , 64(8):5592–5628, 2018.
  * [Zhang et al.(2016)Zhang, Bengio, Hardt, Recht, and Vinyals] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals.  Understanding deep learning requires rethinking generalization.  _arXiv preprint arXiv:1611.03530_ , 2016.

##  Appendix A Technical proofs

######  Proof A.1.

of Lemma 2.1 By definition,

$E(w)=\frac{1}{2}\mathbb{P}_{x\sim\mathcal{N}(\mu_{1},\Sigma_{1}})(w^{T}x<0)+\frac{1}{2}\mathbb{P}_{x\sim\mathcal{N}(\mu_{2},\Sigma_{2})}(w^{T}x\geq
0)$

Rewrite $x=\mu_{1}+y_{1}$ for $x\sim\mathcal{N}(\mu_{1},\Sigma_{1})$ and
$x=\mu_{2}+y_{2}$ for $x\sim\mathcal{N}(\mu_{2},\Sigma_{2})$. Note that
$y_{1}\sim\mathcal{N}(0,\Sigma_{1})$ and $y_{2}\sim\mathcal{N}(0,\Sigma_{2})$.
We obtain:

$E(w)=\frac{1}{2}\mathbb{P}_{y_{1}\sim\mathcal{N}(0,\Sigma_{1})}(w^{T}y_{1}<-\mu_{1}^{T}w)+\frac{1}{2}\mathbb{P}_{y_{2}\sim\mathcal{N}(0,\Sigma_{2})}(w^{T}y_{2}\geq-\mu_{2}^{T}w)$

Since $z_{1}=w^{T}y_{1}\sim\mathcal{N}(0,w^{T}\Sigma_{1}w)$ and
$z_{2}=w^{T}y_{2}\sim\mathcal{N}(0,w^{T}\Sigma_{2}w)$ we have:

$E(w)=\frac{1}{2}\mathbb{P}_{z_{1}\sim\mathcal{N}(0,w^{T}\Sigma_{1}w)}(z_{1}<-\mu_{1}^{T}w)+\frac{1}{2}\mathbb{P}_{z_{2}\sim\mathcal{N}(0,w^{T}\Sigma_{2}w)}(z_{2}\geq-\mu_{2}^{T}w)=$
$=\frac{1}{2}\mathbb{P}_{z^{\prime}_{1}\sim\mathcal{N}(0,1)}(z^{\prime}_{1}<-\frac{\mu_{1}^{T}w}{\sqrt{w^{T}\Sigma_{1}w}})+\frac{1}{2}\mathbb{P}_{z^{\prime}_{2}\sim\mathcal{N}(0,1)}(z^{\prime}_{2}\geq-\frac{\mu_{2}^{T}w}{\sqrt{w^{T}\Sigma_{2}w}})=$
$=\frac{1}{2}\mathbb{P}_{z^{\prime}_{1}\sim\mathcal{N}(0,1)}(z^{\prime}_{1}>\frac{\mu_{1}^{T}w}{\sqrt{w^{T}\Sigma_{1}w}})+\frac{1}{2}\mathbb{P}_{z^{\prime}_{2}\sim\mathcal{N}(0,1)}(z^{\prime}_{2}\geq-\frac{\mu_{2}^{T}w}{\sqrt{w^{T}\Sigma_{2}w}})=$
$=\frac{1}{2}Q(\frac{\mu_{1}^{T}w}{\sqrt{w^{T}\Sigma_{1}w}})+\frac{1}{2}Q(-\frac{\mu_{2}^{T}w}{\sqrt{w^{T}\Sigma_{2}w}})$

######  Proof A.2.

of Lemma 2.4 Denote $Y=\frac{X}{\sigma}\sim\mathcal{N}(0,1)$.

$\mathbb{E}[(|X|-1)^{2}\mathbbm{1}_{|X|>1}]=\mathbb{E}[X^{2}\mathbbm{1}_{|X|>1}]-2\mathbb{E}[|X|\mathbbm{1}_{|X|>1}]+\mathbb{E}[\mathbbm{1}_{|X|>1}]=2\mathbb{E}[X^{2}\mathbbm{1}_{X>1}]-4\mathbb{E}[X\mathbbm{1}_{X>1}]+2\mathbb{E}[\mathbbm{1}_{X>1}]=$
$=2\sigma^{2}\mathbb{E}[Y^{2}\mathbbm{1}_{Y>\frac{1}{\sigma}}]-4\sigma\mathbb{E}[Y\mathbbm{1}_{Y>\frac{1}{\sigma}}]+2\mathbb{E}[\mathbbm{1}_{Y>\frac{1}{\sigma}}]=\frac{2\sigma^{2}}{\sqrt{2\pi}}\int^{+\infty}_{\frac{1}{\sigma}}t^{2}e^{-\frac{t^{2}}{2}}dt-\frac{4\sigma}{\sqrt{2\pi}}\int^{+\infty}_{\frac{1}{\sigma}}te^{-\frac{t^{2}}{2}}dt+2Q(\frac{1}{\sigma})=$
$=-\frac{2\sigma^{2}}{\sqrt{2\pi}}\int^{+\infty}_{\frac{1}{\sigma}}tde^{-\frac{t^{2}}{2}}-\frac{4\sigma}{\sqrt{2\pi}}\int^{+\infty}_{\frac{1}{\sigma}}e^{-\frac{t^{2}}{2}}d\frac{t^{2}}{2}+2Q(\frac{1}{\sigma})=-\frac{2\sigma^{2}}{\sqrt{2\pi}}e^{-\frac{t^{2}}{2}}t_{|_{\frac{1}{\sigma}}^{+\infty}}+\frac{2\sigma^{2}}{\sqrt{2\pi}}\int^{+\infty}_{\frac{1}{\sigma}}e^{-\frac{t^{2}}{2}}dt-$
$-\frac{4\sigma}{\sqrt{2\pi}}\int^{+\infty}_{\frac{1}{2\sigma^{2}}}e^{-z}dz+2Q(\frac{1}{\sigma})=\frac{2\sigma}{\sqrt{2\pi}}e^{-\frac{1}{2\sigma^{2}}}+2\sigma^{2}Q(\frac{1}{\sigma})-\frac{4\sigma}{\sqrt{2\pi}}e^{-\frac{1}{2\sigma^{2}}}+2Q(\frac{1}{\sigma})=$
$=2(\sigma^{2}+1)Q(\frac{1}{\sigma})-\frac{2\sigma}{\sqrt{2\pi}}e^{-\frac{1}{2\sigma^{2}}}$

######  Proof A.3.

of Lemma 2.6 Differentiate the objective from the right hand side by $\beta$:

$\frac{d}{d\beta}(\frac{1}{2\beta}+\frac{\beta
x}{2})=\frac{1}{-2\beta^{2}}+\frac{x}{2}$

We conclude that $\frac{1}{2\beta}+\frac{\beta x}{2}$ is minimized at
$\beta=\frac{1}{\sqrt{x}}$. The value the objective takes at this point is

$\frac{1}{2\beta}+\frac{\beta
x}{2}=\frac{1}{\frac{2}{\sqrt{x}}}+\frac{\frac{1}{\sqrt{x}}x}{2}=\sqrt{x}$

######  Proof A.4.

of Theorem 3.1

Using equation opt_implicit it is straightforward to see that SMD with mirror
$\psi$ converges to $\hat{w}$ solving the following optimization problem for
$w$:

$\min_{w}\max_{\lambda}\psi(w)+\lambda^{T}(X^{T}w-y)$

Denote by $M$ the $d\times n$ matrix satisfying $M_{ij}=\mu_{1i}$ if $i\leq
n/2$ and $M_{ij}=\mu_{2i}$ otherwise. In words, $M$ is the matrix whose first
$n/2$ columns are $\mu_{1}$ and whose last $n/2$ columns are $\mu_{2}$. Denote
the random matrix with independent standard Gaussian entries by $\tilde{X}$.
Also denote the vectors consisting of the first $n/2$ and last $n/2$
coordinates of $\lambda$ by $\lambda_{1}$ and $\lambda_{2}$ respectively so
that $\lambda=\pmatrix{\lambda}_{1}\\\ \lambda_{2}$. Finally, define
$\tilde{\lambda_{i}}=\sigma_{i}\lambda_{i},i=0,1$ and
$\tilde{\lambda}=\pmatrix{\tilde{\hfil}}{\lambda_{1}}\\\ \tilde{\lambda_{2}}$.
We then have

$\lambda^{T}X^{T}w-\lambda^{T}y=\tilde{\lambda}^{T}\tilde{X}^{T}w+\lambda^{T}M^{T}w-\lambda^{T}y=\tilde{\lambda}^{T}\tilde{X}^{T}w+\pmatrix{\lambda}_{1}\\\
\lambda_{2}^{T}\pmatrix{\mu}_{1}^{T}w\mathbbm{1}_{\frac{n}{2}}^{T}&\mu_{2}^{T}w\mathbbm{1}_{\frac{n}{2}}^{T}{}^{T}-\lambda^{T}y=$
$=\tilde{\lambda}^{T}\tilde{X}^{T}w+\tilde{\lambda}^{T}\pmatrix{}{\mu_{1}^{T}w}{\sigma_{1}}\mathbbm{1}_{\frac{n}{2}}^{T}&\frac{\mu_{2}^{T}w}{\sigma_{2}}\mathbbm{1}_{\frac{n}{2}}^{T}{}^{T}-\tilde{\lambda}^{T}(\frac{\mathbbm{1}_{\frac{n}{2}}^{T}}{\sigma_{1}},\frac{\mathbbm{1}_{\frac{n}{2}}^{T}}{\sigma_{2}})$

Plugging it in in the optimization problem above and denoting
$\tilde{\mu_{i}}=\frac{\mu_{i}}{\sigma_{i}},i=1,2$,
$m=(\tilde{\mu_{1}}^{T}w-\sigma_{1}^{-1},\dots,\tilde{\mu_{1}}^{T}w-\sigma_{1}^{-1},\tilde{\mu_{2}}^{T}w-\sigma_{2}^{-1},\dots,\tilde{\mu_{2}}^{T}w-\sigma_{2}^{-1})$
, we obtain:

$\min_{w}\max_{\lambda}\psi(w)+\lambda^{T}(X^{T}w-y)=\min_{w}\max_{\tilde{\lambda}}\tilde{\lambda}^{T}\tilde{X}w+\psi(w)+\tilde{\lambda}^{T}m$

Since $\phi(w,\tilde{\lambda})=\psi(w)+\tilde{\lambda}^{T}m$ is convex in $w$
and is concave (linear) in $\tilde{\lambda}$ and $\tilde{X}$ is standard
normal, we can replace this PO problem by the corresponding AO, which is known
to yield solutions with the same empirical distribution and the same
distribution of $\mu_{1}^{T}w,\mu_{2}^{T}w,\|w\|_{2}^{2}$ according to Theorem
2.3:

$\min_{w}\max_{\tilde{\lambda}}\|\tilde{\lambda}\|g^{T}w+\tilde{\lambda}^{T}h\|w\|_{2}+\psi(w)+\tilde{\lambda}^{T}m$

Write $\tilde{\lambda}=\alpha u$, where $\alpha=|\tilde{\lambda}|\geq 0$ and
$u$ is unit. We then have:

$\min_{w}\max_{\alpha\geq 0,u}\alpha
g^{T}w+\tilde{\lambda}^{T}(h\|w\|_{2}+m)+\psi(w)$

Since this is clearly maximized when $u$ is aligned with $h\|w\|_{2}+m$, we
simplify the expression:

$\min_{w}\max_{\alpha\geq 0}\alpha g^{T}w+\alpha\|h\|w\|_{2}+m\|_{2}+\psi(w)$

We will simplify $\|h\|w\|_{2}+m\|_{2}$ before proceeding further. Note that

$\|h\|w\|_{2}+m\|_{2}=\sqrt{h^{T}h\|w\|_{2}^{2}+2h^{T}m\|w\|_{2}+m^{T}m}$

Recall that $h$ is standard normal. Thus, $h^{T}h$ is almost equal to $n$,
because we work in the asymptotic regime. The second term $2h^{T}m\|w\|_{2}$
is negligible compared to the first for almost any $h$ because $m$ is always
in the span of two vectors $(1,\dots,1,0,\dots,0)^{T}$ and
$(0,\dots,0,1,\dots,1)^{T}$ and the projection of $w$ onto this $2$\-
dimensional span is negligible for almost any $h$. Finally,
$m^{T}m=\frac{n}{2\sigma_{1}^{2}}(\mu_{1}^{T}w-1)^{2}+\frac{n}{2\sigma_{2}^{2}}(\mu_{2}^{T}w+1)^{2}$.
Incorporating all these observations into the objective and switching the
order of optimization using the convex-concativity of the terms we get:

$\max_{\alpha\geq 0}\min_{w}\alpha
g^{T}w+\alpha\sqrt{n\|w\|_{2}^{2}+\frac{n}{2\sigma_{1}^{2}}(\mu_{1}^{T}w-1)^{2}+\frac{n}{2\sigma_{2}^{2}}(\mu_{2}^{T}w+1)^{2}}+\psi(w)$

To get rid of the square root, we use Lemma 2.6 and arrive to:

$\max_{\alpha\geq 0}\min_{w,\beta\geq 0}\alpha
g^{T}w+\frac{\alpha}{2\beta}+\frac{\alpha\beta
n}{2}(\|w\|_{2}^{2}+\frac{1}{2\sigma_{1}^{2}}(\mu_{1}^{T}w-1)^{2}+\frac{1}{2\sigma_{2}^{2}}(\mu_{2}^{T}w+1)^{2})+\psi(w)$

Substituting $\frac{\mu_{1}^{T}w-1}{\sigma_{1}}$ and
$\frac{\mu_{2}w+1}{\sigma_{2}}$ by $a_{1}$ and $a_{2}$ respectively and adding
two more scalar variables $\gamma_{1},\gamma_{2}$ we deduce:

$\max_{\alpha\geq 0}\min_{w,\beta\geq
0}\max_{\gamma_{1},\gamma_{2},a_{1},a_{2}}\alpha
g^{T}w+\frac{\alpha}{2\beta}+\frac{\alpha\beta
n}{2}(\|w\|_{2}^{2}+\frac{1}{2}a_{1}^{2}+\frac{1}{2}a_{2}^{2})+\gamma_{1}(\frac{\mu_{1}^{T}w-1}{\sigma_{1}}-a_{1})+\gamma_{2}(\frac{\mu_{2}w+1}{\sigma_{2}}-a_{2})+\psi(w)$

Taking the derivatives by $a_{i},i=1,2$ and equating them to $0$ leads to
$a_{i}=\frac{2\gamma_{i}}{\alpha\beta n},i=1,2$. Plugging these in and
simplifying and regrouping the terms we obtain the desired optimization
problem:

$\max_{\alpha\geq 0}\min_{w,\beta\geq
0}\max_{\gamma_{1},\gamma_{2}}\psi(w)+\alpha
g^{T}w+\frac{\alpha}{2\beta}+\frac{\alpha\beta
n}{2}\|w\|_{2}^{2}+\frac{\gamma_{1}(\mu_{1}^{T}w-1)}{\sigma_{1}}+\frac{\gamma_{2}(\mu_{2}^{T}w+1)}{\sigma_{2}}-\frac{\gamma_{1}^{2}+\gamma_{2}^{2}}{\alpha\beta
n}$

######  Proof A.5.

of Theorem 3.3

Put $\psi(w)=\|w\|^{2}_{2}$ in the objective of Theorem 3.1:

$\max_{\alpha\geq 0}\min_{\beta\geq
0}\max_{\gamma_{1},\gamma_{2}}\min_{w}\|w\|^{2}_{2}+\alpha
g^{T}w+\frac{\alpha}{2\beta}+\frac{\alpha\beta
n}{2}\|w\|_{2}^{2}+\frac{\gamma_{1}(\mu_{1}^{T}w-1)}{\sigma_{1}}+\frac{\gamma_{2}(\mu_{2}^{T}w+1)}{\sigma_{2}}-\frac{\gamma_{1}^{2}+\gamma_{2}^{2}}{\alpha\beta
n}$

Denote $\tilde{\gamma_{i}}=\frac{\gamma_{i}}{\sigma_{i}},i=1,2$ and rewrite it
in the following way:

$\max_{\alpha\geq 0}\min_{\beta\geq
0}\max_{\tilde{\gamma_{1}},\tilde{\gamma_{2}}}\min_{w}\tilde{\gamma_{2}}-\tilde{\gamma_{1}}+\frac{\alpha}{2\beta}-\frac{\sigma_{1}^{2}\tilde{\gamma_{1}}^{2}+\sigma_{2}^{2}\tilde{\gamma_{2}}^{2}}{\alpha\beta
n}+\sum_{i}(1+\frac{\alpha\beta n}{2})w_{i}^{2}+w_{i}(\alpha
g_{i}+\tilde{\gamma_{1}}\mu_{1i}+\tilde{\gamma_{2}}\mu_{2i})$

Thus, the minimization over $w$ reduces to minimization over $w_{i}$ for each
$w_{i}$ separately. The latter is straightforward because the objective of the
minimization is just a quadratic polynomial. Therefore, the optimal
$w_{i}=-\frac{\alpha
g_{i}+\tilde{\gamma_{1}}\mu_{1i}+\tilde{\gamma_{2}}\mu_{2i}}{2+\alpha\beta n}$
for each $i$ and thus $w=-\frac{\alpha
g+\tilde{\gamma_{1}}\mu_{1}+\tilde{\gamma_{2}}\mu_{2}}{2+\alpha\beta n}$.
Hence, we obtain the following optimization problem:

$\max_{\alpha\geq 0}\min_{\beta\geq
0}\max_{\tilde{\gamma_{1}},\tilde{\gamma_{2}}}\tilde{\gamma_{2}}-\tilde{\gamma_{1}}+\frac{\alpha}{2\beta}-\frac{\sigma_{1}^{2}\tilde{\gamma_{1}}^{2}+\sigma_{2}^{2}\tilde{\gamma_{2}}^{2}}{\alpha\beta
n}-\frac{1}{4+2\alpha\beta n}\sum_{i}(\alpha
g_{i}+\tilde{\gamma_{1}}\mu_{1i}+\tilde{\gamma_{2}}\mu_{2i})^{2}$

We will simplify the sum before proceeding further with the expression. First,
note that

$\sum_{i}(\alpha
g_{i}+\tilde{\gamma_{1}}\mu_{1}+\tilde{\gamma_{2}}\mu_{2})^{2}=\sum_{i}\alpha^{2}g_{i}^{2}+\tilde{\gamma_{1}}^{2}\mu_{1i}^{2}+\tilde{\gamma_{2}}^{2}\mu_{2i}^{2}+2\alpha(\tilde{\gamma_{1}}\mu_{1i}g_{i}+\tilde{\gamma_{2}}\mu_{2i}g_{i})+2\tilde{\gamma_{1}}\tilde{\gamma_{2}}\mu_{1i}\mu_{2i}=$
$=\alpha^{2}\|g\|_{2}^{2}+\tilde{\gamma_{1}}^{2}\|\mu_{1}\|_{2}^{2}+\tilde{\gamma_{2}}^{2}\|\mu_{2}\|_{2}^{2}+2\alpha\tilde{\gamma_{1}}\mu_{1}^{T}g+2\alpha\tilde{\gamma_{2}}\mu_{2}^{T}g+2\tilde{\gamma_{1}}\tilde{\gamma_{2}}\mu_{1}^{T}\mu_{2}$

Since $g$ is standard normal and $\mu_{1},\mu_{2}$ are two fixed vectors we
can ignore the $2\alpha\tilde{\gamma_{1}}\mu_{1}^{T}g$ and
$2\alpha\tilde{\gamma_{2}}\mu_{2}^{T}g$ terms and replace $\|g\|_{2}^{2}$ by
$d$ asymptotically. Hence, we can replace the sum with:

$\alpha^{2}d+\tilde{\gamma_{1}}^{2}\|\mu_{1}\|_{2}^{2}+\tilde{\gamma_{2}}^{2}\|\mu_{2}\|_{2}^{2}+2\tilde{\gamma_{1}}\tilde{\gamma_{2}}\mu_{1}^{T}\mu_{2}$

Plugging it back into the main objective we have:

$\max_{\alpha\geq 0}\min_{\beta\geq
0}\max_{\tilde{\gamma_{1}},\tilde{\gamma_{2}}}\tilde{\gamma_{2}}-\tilde{\gamma_{1}}+\frac{\alpha}{2\beta}-\frac{\sigma_{1}^{2}\tilde{\gamma_{1}}^{2}+\sigma_{2}^{2}\tilde{\gamma_{2}}^{2}}{\alpha\beta
n}-\frac{1}{4+2\alpha\beta
n}(\alpha^{2}d+\tilde{\gamma_{1}}^{2}\|\mu_{1}\|_{2}^{2}+\tilde{\gamma_{2}}^{2}\|\mu_{2}\|_{2}^{2}+2\tilde{\gamma_{1}}\tilde{\gamma_{2}}\mu_{1}^{T}\mu_{2})$

Take the derivatives by $\tilde{\gamma_{1}},\tilde{\gamma_{2}}$ and equate
them to zero:

$-1-\frac{2\sigma_{1}^{2}\tilde{\gamma_{1}}}{\alpha\beta
n}-\frac{\|\mu_{1}\|_{2}^{2}\tilde{\gamma_{1}}+\mu_{1}^{T}\mu_{2}\tilde{\gamma_{2}}}{2+\alpha\beta
n}=0$ $1-\frac{2\sigma_{2}^{2}\tilde{\gamma_{2}}}{\alpha\beta
n}-\frac{\|\mu_{2}\|_{2}^{2}\tilde{\gamma_{2}}+\mu_{1}^{T}\mu_{2}\tilde{\gamma_{1}}}{2+\alpha\beta
n}=0$

Denote

$\tilde{\tilde{\gamma_{i}}}=\frac{4\tilde{\gamma}_{i}}{\alpha\beta
n(\alpha\beta
n+2)},\tilde{\tilde{\gamma}}=(\tilde{\tilde{\gamma_{1}}},\tilde{\tilde{\gamma_{2}}})^{T}\text{and}M=\pmatrix{}{\alpha\beta
n}{4}\|\mu_{1}\|_{2}^{2}+\sigma_{1}^{2}(\frac{\alpha\beta
n}{2}+1)&\frac{\alpha\beta n}{4}\mu_{1}^{T}\mu_{2}\\\ \frac{\alpha\beta
n}{4}\mu_{1}^{T}\mu_{2}\frac{\alpha\beta
n}{4}\|\mu_{2}\|_{2}^{2}+\sigma_{2}^{2}(\frac{\alpha\beta n}{2}+1)$

The linear system of equations in $\tilde{\gamma}_{1}$ and
$\tilde{\gamma}_{2}$ then translates as

$M\tilde{\tilde{\gamma}}=\pmatrix{-}1\\\ 1$

Note that $det(M)=\Delta$, where $\Delta$ is defined in the statement of the
theorem. Hence, we deduce:

$\tilde{\tilde{\gamma}}=\frac{1}{\Delta}\pmatrix{}{\alpha\beta
n}{4}\|\mu_{2}\|_{2}^{2}+\sigma_{1}^{2}(\frac{\alpha\beta
n}{2}+1)&-\frac{\alpha\beta n}{4}\mu_{1}^{T}\mu_{2}\\\ -\frac{\alpha\beta
n}{4}\mu_{1}^{T}\mu_{2}\frac{\alpha\beta
n}{4}\|\mu_{1}\|_{2}^{2}+\sigma_{2}^{2}(\frac{\alpha\beta
n}{2}+1)\pmatrix{-}1\\\ 1$

Which gives us

$\tilde{\tilde{\gamma_{1}}}=-\frac{1}{\Delta}(\frac{\alpha\beta
n}{4}(\|\mu_{2}\|_{2}^{2}+\mu_{1}^{T}\mu_{2})+\sigma_{1}^{2}(\frac{\alpha\beta
n}{2}+1))$ $\tilde{\tilde{\gamma_{2}}}=\frac{1}{\Delta}(\frac{\alpha\beta
n}{4}(\|\mu_{1}\|_{2}^{2}+\mu_{1}^{T}\mu_{2})+\sigma_{2}^{2}(\frac{\alpha\beta
n}{2}+1))$

Recover $\tilde{\gamma_{1}},\tilde{\gamma_{2}}$:

$\tilde{\gamma_{1}}=-\frac{\alpha\beta n(\alpha\beta
n+2)}{4\Delta}(\frac{\alpha\beta
n}{4}(\|\mu_{2}\|_{2}^{2}+\mu_{1}^{T}\mu_{2})+\sigma_{1}^{2}(\frac{\alpha\beta
n}{2}+1))$ $\tilde{\gamma_{2}}=\frac{\alpha\beta n(\alpha\beta
n+2)}{4\Delta}(\frac{\alpha\beta
n}{4}(\|\mu_{1}\|_{2}^{2}+\mu_{1}^{T}\mu_{2})+\sigma_{2}^{2}(\frac{\alpha\beta
n}{2}+1))$

We can find the optimal $w$ using $\tilde{\gamma_{1}}$ and
$\tilde{\gamma_{2}}$:

$w=-\frac{\alpha
g+\tilde{\gamma_{1}}\mu_{1}+\tilde{\gamma_{2}}\mu_{2}}{2+\alpha\beta n}=$
$=-\frac{\alpha g}{2+\alpha\beta n}+\frac{\alpha\beta
n}{4\Delta}(\frac{\alpha\beta
n}{4}(\|\mu_{2}\|_{2}^{2}+\mu_{1}^{T}\mu_{2})+\sigma_{1}^{2}(\frac{\alpha\beta
n}{2}+1))\mu_{1}-\frac{\alpha\beta n}{4\Delta}(\frac{\alpha\beta
n}{4}(\|\mu_{1}\|_{2}^{2}+\mu_{1}^{T}\mu_{2})+\sigma_{2}^{2}(\frac{\alpha\beta
n}{2}+1))\mu_{2}$

Instead of directly inserting $\tilde{\gamma_{1}}$ and $\tilde{\gamma_{2}}$
into the objective now, which appears to be a horrendous task, we will
remember what the optimal $w$ is but will return to the initial objective and
change the order of optimization first using that the objective is convex in
$w$ and concave in $\tilde{\gamma_{1}},\tilde{\gamma_{2}}$:

$\max_{\alpha\geq 0}\min_{\beta\geq
0}\min_{w}\max_{\tilde{\gamma_{1}},\tilde{\gamma_{2}}}\|w\|^{2}_{2}+\alpha
g^{T}w+\frac{\alpha}{2\beta}+\frac{\alpha\beta
n}{2}\|w\|_{2}^{2}+\tilde{\gamma_{1}}(\mu_{1}^{T}w-1)+\tilde{\gamma_{2}}(\mu_{2}^{T}w+1)-\frac{\sigma_{1}^{2}\tilde{\gamma_{1}}^{2}+\sigma_{2}^{2}{\gamma_{2}}^{2}}{\alpha\beta
n}$

Differentiating by $\tilde{\gamma_{1}}$ and $\tilde{\gamma_{2}}$ and equating
to $0$ again we immediately see that $\tilde{\gamma_{1}}=\frac{\alpha\beta
n}{2\sigma_{1}^{2}}(\mu_{1}^{T}w-1)$ and $\tilde{\gamma_{2}}=\frac{\alpha\beta
n}{2\sigma_{2}^{2}}(\mu_{2}^{T}w+1)$. Incorporating this remark into the
objective we get:

$\max_{\alpha\geq 0}\min_{\beta\geq 0}\min_{w}\|w\|^{2}_{2}+\alpha
g^{T}w+\frac{\alpha}{2\beta}+\frac{\alpha\beta
n}{2}\|w\|_{2}^{2}+\frac{\alpha\beta
n}{4\sigma_{1}^{2}}(\mu_{1}^{T}w-1)^{2}+\frac{\alpha\beta
n}{4\sigma_{2}^{2}}(\mu_{2}^{T}w+1)^{2}$

Note that this is a quadratic function in $w$ whose linear term is equal to

$(\alpha g+\frac{\alpha\beta
n}{2}(\frac{\mu_{2}}{\sigma^{2}_{2}}-\frac{\mu_{1}}{\sigma^{2}_{1}}))^{T}w$

.

Hence, the value of the objective at the optimal parameter $w$ equals
$(\frac{\alpha}{2}g+\frac{\alpha\beta
n}{4}(\frac{\mu_{2}}{\sigma^{2}_{2}}-\frac{\mu_{1}}{\sigma^{2}_{1}}))^{T}w$.
Use the expression for the optimal $w$ we derived earlier to evaluate it at
the optimal $w$. We will deem the cross-terms negligible because for a random
standard Gaussian $g$ its dot products with $\mu_{1}$ and $\mu_{2}$ are
negligible and will also replace $g^{T}g$ by $d$:

$(\frac{\alpha}{2}g+\frac{\alpha\beta
n}{4}(\frac{\mu_{2}}{\sigma^{2}_{2}}-\frac{\mu_{1}}{\sigma^{2}_{1}}))^{T}w=-\frac{\alpha^{2}d}{4+2\alpha\beta
n}+\frac{(\alpha\beta n)^{2}}{16\Delta}(\frac{\alpha\beta
n}{4}(\|\mu_{2}\|_{2}^{2}+\mu_{1}^{T}\mu_{2})+\sigma_{1}^{2}(\frac{\alpha\beta
n}{2}+1))(\frac{\mu_{1}^{T}\mu_{2}}{\sigma_{2}^{2}}-\frac{\|\mu_{1}\|_{2}^{2}}{\sigma_{1}^{2}})-$
$-\frac{(\alpha\beta n)^{2}}{16\Delta}(\frac{\alpha\beta
n}{4}(\|\mu_{1}\|_{2}^{2}+\mu_{1}^{T}\mu_{2})+\sigma_{2}^{2}(\frac{\alpha\beta
n}{2}+1))(\frac{\|\mu_{2}\|_{2}^{2}}{\sigma_{2}^{2}}-\frac{\mu_{1}^{T}\mu_{2}}{\sigma_{1}^{2}})=$
$=-\frac{\alpha^{2}d}{4(1+\frac{\alpha\beta n}{2})}-\frac{(\alpha\beta
n)^{2}}{16\Delta}(\|\mu_{1}\|^{2}+\|\mu_{2}\|^{2}-(\frac{\sigma_{1}^{2}}{\sigma_{2}^{2}}+\frac{\sigma_{2}^{2}}{\sigma_{1}^{2}})\mu_{2}^{T}\mu_{1})-$
$-\frac{(\alpha\beta
n)^{3}}{32\Delta}(\|\mu_{1}\|^{2}+\|\mu_{2}\|^{2}-(\frac{\sigma_{1}^{2}}{\sigma_{2}^{2}}+\frac{\sigma_{2}^{2}}{\sigma_{1}^{2}})\mu_{2}^{T}\mu_{1}+(\frac{1}{2\sigma_{1}^{2}}+\frac{1}{2\sigma_{2}^{2}})(\|\mu_{1}\|^{2}\|\mu_{2}\|^{2}-(\mu_{2}^{T}\mu_{1})^{2}))$

######  Proof A.6.

of Theorem 3.5

Plug in $\psi(w)=\|w\|_{1}$ in the objective of Theorem 3.1:

$\max_{\alpha\geq 0}\min_{\beta\geq
0}\max_{\gamma_{1},\gamma_{2}}\min_{w}\|w\|_{1}+\alpha
g^{T}w+\frac{\alpha}{2\beta}+\frac{\alpha\beta
n}{2}\|w\|_{2}^{2}+\frac{\gamma_{1}(\mu_{1}^{T}w-1)}{\sigma_{1}}+\frac{\gamma_{2}(\mu_{2}^{T}w+1)}{\sigma_{2}}-\frac{\gamma_{1}^{2}+\gamma_{2}^{2}}{\alpha\beta
n}$

Note that this expression can be split in $i$:

$\max_{\alpha\geq 0}\min_{\beta\geq
0}\max_{\gamma_{1},\gamma_{2}}\min_{w}\frac{\gamma_{2}}{\sigma_{2}}-\frac{\gamma_{1}}{\sigma_{1}}+\frac{\alpha}{2\beta}-\frac{\gamma_{1}^{2}+\gamma_{2}^{2}}{\alpha\beta
n}+\sum_{i}|w_{i}|+\alpha g_{i}w_{i}+\frac{\alpha\beta
n}{2}w_{i}^{2}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}w_{i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i}w_{i}$

Therefore, minimizing the entire expression in $w$ is equivalent to minimizing
the corresponding summand for each $i$:

$\min_{w_{i}}|w_{i}|+\alpha g_{i}w_{i}+\frac{\alpha\beta
n}{2}w_{i}^{2}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}w_{i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i}w_{i}$

Denoting $u_{i}=|w_{i}|,\epsilon_{i}=sign(w_{i})$ we rewrite it as:

$\min_{u_{i}\geq 0,\epsilon_{i}=\pm 1}u_{i}+u_{i}\epsilon_{i}(\alpha
g_{i}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i})+\frac{\alpha\beta
n}{2}u_{i}^{2}$

It is clear now that this is minimized when $\epsilon_{i}=-sign(\alpha
g_{i}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i})$
and the problem reduces to:

$\min_{u_{i}\geq 0}u_{i}(1-|\alpha
g_{i}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i}|)+\frac{\alpha\beta
n}{2}u_{i}^{2}$

The latter is just a quadratic problem with a constraint $u_{i}\geq 0$ and
therefore the solution is $u_{i}=\max(0,\frac{|\alpha
g_{i}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i}|-1}{\alpha\beta
n})$. The corresponding value of the objective is:

$-\frac{|\alpha
g_{i}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i}|-1}{2\alpha\beta
n}\max(0,|\alpha
g_{i}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i}|-1)$

Recover the corresponding $w_{i}$:

$w_{i}=\epsilon_{i}u_{i}=-sign(\alpha
g_{i}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i})\max(0,\frac{|\alpha
g_{i}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i}|-1}{\alpha\beta
n})$

Replacing each $|w_{i}|+\alpha g_{i}w_{i}+\frac{\alpha\beta
n}{2}w_{i}^{2}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}w_{i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i}w_{i}$
in the main objective from the beginning of the proof with the value of the
objective at the optimal $w_{i}$ we just found, we derive the desired four -
dimensional optimization problem:

$\max_{\alpha\geq 0}\min_{\beta\geq
0}\max_{\gamma_{1},\gamma_{2}}\frac{\gamma_{2}}{\sigma_{2}}-\frac{\gamma_{1}}{\sigma_{1}}+\frac{\alpha}{2\beta}-\frac{\gamma_{1}^{2}+\gamma_{2}^{2}}{\alpha\beta
n}-\sum_{i}\frac{|\alpha
g_{i}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i}|-1}{2\alpha\beta
n}\max(0,|\alpha
g_{i}+\frac{\gamma_{1}}{\sigma_{1}}\mu_{1i}+\frac{\gamma_{2}}{\sigma_{2}}\mu_{2i}|-1)$


# Riemannian Diffusion Schrödinger Bridge

James Thornton  Michael Hutchinson  Emile Mathieu  Valentin De Bortoli  Yee
Whye Teh  Arnaud Doucet

###### Abstract

Score-based generative models exhibit state of the art performance on density
estimation and generative modeling tasks. These models typically assume that
the data geometry is flat, yet recent extensions have been developed to
synthesize data living on Riemannian manifolds. Existing methods to accelerate
sampling of diffusion models are typically not applicable in the Riemannian
setting and Riemannian score-based methods have not yet been adapted to the
important task of interpolation of datasets. To overcome these issues, we
introduce _Riemannian Diffusion Schrödinger Bridge_. Our proposed method
generalizes Diffusion Schrödinger Bridge introduced in (De Bortoli et al.,
2021) to the non-Euclidean setting and extends Riemannian score-based models
beyond the first time reversal. We validate our proposed method on synthetic
data and real Earth and climate data.

Diffusion processes, Generative modeling, Riemannian manifold, Score-based
generative models, Schrödinger bridge

##  1 Background

###  1.1 Score Based Generative Modeling

In Euclidean spaces, Score-based Generative Modeling (SGM) (Song & Ermon,
2019; Song et al., 2021b) consists of two main components. The first is a
forward _noising_ process $(\mathbf{X}_{t})_{t\geq 0}$ defined via the
stochastic differential equation (SDE) (1) and the initial distribution
$\mathbf{X}_{0}\sim p_{\textup{data}}$, targeting an easy-to-sample prior
$p_{\textup{prior}}$. The second component is a backward _denoising_ process
$(\mathbf{Y}_{t})_{t\geq 0}=(\mathbf{X}_{T-t})_{t\in\left[0,T\right]}$ defined
by the time-reversal of the noising SDE (1) (Cattiaux et al., 2021; Haussmann
& Pardoux, 1986) from $p_{\textup{prior}}=p_{T}$ to $p_{\textup{data}}=p_{0}$
. Here $f$ is the drift, $g$ is the time-dependent volatility, while
$(\mathbf{B}_{t})_{t\geq 0}$ is a d-dimensional Brownian motion and
$\mathbb{P}_{t}$ is the distribution of $\mathbf{X}_{t}$ with corresponding
density $p_{t}$. The denoising process defines a generative model by sampling
$\mathbf{Y}_{0}\sim p_{\textup{prior}}$

$\mathrm{d}\mathbf{X}_{t}=f(t,\mathbf{X}_{t})\mathrm{d}t+g(t)\mathrm{d}\mathbf{B}_{t},\qquad\mathrm{d}\mathbf{Y}_{t}=\\{-f(t,\mathbf{Y}_{t})+g^{2}(t)\nabla\log
p_{T-t}(\mathbf{Y}_{t})\\}\mathrm{d}t+g(t)\mathrm{d}\mathbf{B}_{t}.$ (1)

The intractable score term, $\nabla\log p_{T-t}(\mathbf{Y}_{t})$, may be
approximated by using the following score matching identity
$\textstyle{\nabla_{x_{t}}\log
p_{t}(x_{t})=\int_{\mathbb{R}^{d}}\nabla_{x_{t}}\log
p_{t|0}(x_{t}|x_{0})~{}p_{0|t}(x_{0}|x_{t})\mathrm{d}x_{0}}$, with tractable
transition density $p_{t|0}$. This is then used to train a neural network
$\bm{s}_{\theta}:\mathbb{R}\times\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ with
regression objective
$\bm{s}_{\theta^{*}}=\arg\min_{\theta}\mathbb{E}_{\mathbf{X}_{t},\mathbf{X}_{0}}\|\nabla_{x}\log
p_{t|0}(\mathbf{X}_{t}|\mathbf{X}_{0})-\bm{s}_{\theta}(t,\mathbf{X}_{t})\|^{2}$.
Once trained, one can generate samples which are approximately distributed
according to $p_{\textup{data}}$ via (1) (e.g. considering the Euler–Maruyama
discretization of this SDE) by substituting
$\bm{s}_{\theta^{*}}(t,x_{t})\approx\nabla_{x_{t}}\log p_{t}(x_{t})$.

###  1.2 Riemannian Score Based Generative Modeling

De Bortoli et al. (2022) extended SGM to compact Riemannian manifolds, denoted
$\mathcal{M}$, henceforth abbreviated RSGM for Riemannian Score-based
Generative Modeling. Given a $\mathcal{M}$-valued diffusion process (2, left),
where $\mathbf{B}_{t}^{\mathcal{M}}$ denotes Brownian motion on $\mathcal{M}$;
the time reversal process may be written as (2, right) (De Bortoli et al.,
2022, Theorem 1)

$\displaystyle\mathrm{d}\mathbf{X}_{t}$
$\displaystyle=f(t,\mathbf{X}_{t})\mathrm{d}t+g(t)\mathrm{d}\mathbf{B}_{t}^{\mathcal{M}},$
$\displaystyle\mathrm{d}\mathbf{Y}_{t}$
$\displaystyle=\\{-f(t,\mathbf{Y}_{t})+g^{2}(t)\nabla\log
p_{T-t}(\mathbf{Y}_{t})\\}\mathrm{d}t+g(t)\mathrm{d}\tilde{\mathbf{B}}_{t}^{\mathcal{M}}.\vspace{-0.1cm}$
(2)

For a more thorough background on Riemannian geometry and time-reversal on
manifolds see De Bortoli et al. (2022, App. B and G). Simulating a diffusion
on a manifold is crucial to this approach. In contrast to Euclidean SGMs,
closed-form sampling schemes are generally not available for diffusions on
manifolds, hence approximation schemes such as the Geodesic Random Walk from
Algo. 1 are required. Note that $\mathbf{B}_{t}^{\mathcal{M}}$ may be
simulated using Algo 1 for $g=1,f=0$. This consists of applying an
Euler–Maruyama step in the tangent space, whereby Gaussian noise has also been
projected to the tangent space using the projection operator $\mathrm{P}$,
then projecting back to the manifold using the exponential mapping.

ForwardForwardBackwardBackward$t$$t=0$$t=T$RSGMRDSB Figure 1: Earthquake data:
empirical data in red, generated samples in green using $N=10$ diffusion
steps. Top forward/backward pair: RSGM. Bottom pair: RDSB, with $5$ IPF steps.

1: Input: step size: $\gamma$ , initial state $X_{0}$

2: for $k\in\\{0,\dots,N-1\\}$ do

3:
$\bar{\mathbf{Z}}_{k+1}\sim\mathcal{N}(0,I_{p}),\quad\mathbf{Z}_{k+1}=\mathrm{P}(\mathbf{X}_{k})\bar{\mathbf{Z}}_{k+1}$

4: $\mathbf{W}_{k+1}=\gamma
f(k\gamma,\mathbf{X}_{k})+\sqrt{\gamma}g(k\gamma)\mathbf{Z}_{k+1}$

5: $\mathbf{X}_{k+1}=\exp_{\mathbf{X}_{k}}\left(\mathbf{W}_{k+1}\right)$

6: end for

7: return $\\{\mathbf{X}_{k}\\}_{k=0}^{N-1}$

Algorithm 1 Simulating diffusions on a manifold

Score matching may be used to approximate the time-reversal (2) in the
Riemannian setting (De Bortoli et al., 2022). Unlike for Euclidean SGM,
$\nabla\log p_{t|0}(x_{t}|x_{0})$ is not available in closed-form. Instead,
one may use the score identity 111Here all gradients are considered w.r.t. the
Riemannian metric of $\mathcal{M}$. $\textstyle{\nabla_{x_{t}}\log
p_{t}(x_{t})=\int_{\mathcal{M}}\nabla_{x_{t}}\log
p_{t|s}(x_{t}|x_{s})\mathbb{P}_{s|t}(x_{t},\mathrm{d}x_{s})}$ for $s\approx
t$, where again $\mathbb{P}$ is the distribution of $\mathbf{X}$ and $p_{t}$
corresponds to the density of $\mathbb{P}_{t}$ with respect to the uniform
distribution on $\mathcal{M}$, then $\nabla\log
p_{t}(x_{t})=\operatorname*{arg\,min}_{\mathbf{s}_{t}}\int_{\mathcal{M}^{2}}\|\nabla_{x}\log
p_{t|s}(x_{t}|x_{s})-\mathbf{s}_{t}(x_{t})\|^{2}\mathrm{d}\mathbb{P}_{s,t}(x_{s},x_{t})$.

###  1.3 The Schrödinger Bridge Problem

A Schrödinger bridge extension of SGM has been introduced to reduce the number
of diffusion steps for SGM by learning _both_ forward and backward diffusions
(De Bortoli et al., 2021). We briefly recall the notion of dynamical
Schrödinger bridge (Léonard, 2012; Chen et al., 2016; Vargas et al., 2021; De
Bortoli et al., 2021; Chen et al., 2022). We consider a reference path
probability measure
$\mathbb{P}\in\mathcal{P}(\mathrm{C}(\left[0,T\right]),\mathcal{M})$ where
$\mathcal{P}(\mathrm{C}(\left[0,T\right]),\mathcal{M})$ is the space measures
on continuous paths $(\mathbf{X}_{t})_{t\in[0,1]}$ in $\mathcal{M}$ . In
practice, we set $\mathbb{P}$ to be the distribution of the Brownian motion
$(\mathbf{B}_{t}^{\mathcal{M}})_{t\in\left[0,T\right]}$ initialized from
$p_{\textup{data}}$, i.e. $\mathbf{B}_{0}^{\mathcal{M}}$ has distribution
$p_{\textup{data}}$. We consider the _dynamical Schrödinger bridge problem_

$\mathbb{Q}^{\star}=\operatorname*{arg\,min}\\{\operatorname{KL}\left(\mathbb{Q}|\mathbb{P}\right)\,:\;\mathbb{Q}\in\mathcal{P}(\mathrm{C}(\left[0,T\right]),\mathcal{M}),\
\mathbb{Q}_{0}=p_{\textup{data}},\ \mathbb{Q}_{T}=p_{\textup{prior}}\\}.$ (3)

The idealised solution $\mathbb{Q}^{\star}$ is called the Schrödinger Bridge
(SB). Given a backward process
$(\mathbf{Y}_{t}^{\star})_{t\in\left[0,T\right]}$ associated to
$\mathbb{Q}^{\star}$, one can obtain a generative model as follows. First
sample from $\mathbf{Y}^{\star}_{T}\sim p_{\textup{prior}}$ and then follow
the (backward) dynamics of $(\mathbf{Y}^{\star}_{t})_{t\in\left[0,T\right]}$.
By definition, we obtain that $\mathbf{Y}^{\star}_{0}\sim p_{\textup{data}}$,
the data distribution.

In practice, the solution of the SB problem is approximated using the
Iterative Proportional Fitting (IPF) algorithm, which coincides with the
Sinhkorn algorithm in discrete space (Sinkhorn, 1967; Peyré & Cuturi, 2019).
IPF defines a sequence of path probability measures
$(\mathbb{Q}^{n})_{n\in\mathbb{N}}\in(\mathcal{P}(\mathrm{C}(\left[0,T\right],\mathcal{M})))^{\mathbb{N}}$,
such that $\mathbb{Q}^{0}=\mathbb{P}$ and for any $n\in\mathbb{N}$

$\displaystyle\mathbb{Q}^{2n+1}=\operatorname*{arg\,min}\\{\operatorname{KL}\left(\mathbb{Q}|\mathbb{Q}^{2n}\right)\,:\;\mathbb{Q}\in\mathcal{P}(\mathrm{C}(\left[0,T\right]),\mathcal{M}),\mathbb{Q}_{T}=p_{\textup{prior}}\\},$
(4)
$\displaystyle\mathbb{Q}^{2n+2}=\operatorname*{arg\,min}\\{\operatorname{KL}\left(\mathbb{Q}|\mathbb{Q}^{2n+1}\right)\,:\;\mathbb{Q}\in\mathcal{P}(\mathrm{C}(\left[0,T\right]),\mathcal{M}),\mathbb{Q}_{0}=p_{\textup{data}}\\}.$
(5)

Under mild assumptions on $\mathbb{P}$, $p_{\textup{data}}$ and
$p_{\textup{prior}}$, we have that $(\mathbb{Q}^{n})_{n\in\mathbb{N}}$
converges towards $\mathbb{Q}^{\star}$ (Nutz & Wiesel, 2022).

##  2 Riemannian Diffusion Schrödinger Bridge

In Euclidean state spaces, De Bortoli et al. (2021) proposed Diffusion
Schrödinger Bridge (DSB), an algorithm to approximate the solution to the SB
problem based on time-reversal, using score matching to approximate the IPF
iterates. We propose Riemannian Diffusion Schrödinger Bridge (RDSB), an
extension of DSB to approximate solutions of the SB problem for compact
Riemannian manifolds.

We first connect the IPF iterates $(\mathbb{Q}^{n})_{n\in\mathbb{N}}$ with
time reversal of diffusion processes on $\mathcal{M}$. To simplify notation,
we rewrite Equation 2 in terms of drift functions $f^{n}$ and $b^{n}$ in
Equations 6 and 7. Here $\mathbf{Y}^{n}$ corresponds to the time-reversal of
$\mathbf{X}^{n}$, initialized at $\mathbf{Y}^{n}_{T}\sim p_{\textup{prior}}$.
$\mathbf{X}^{0}$ is the Brownian motion on $\mathcal{M}$, $f^{n}=0$, and
$\mathbf{X}^{n+1}$ denotes the time-reversal of the diffusion
$\mathbf{Y}^{n}$, initialized at $\mathbf{X}^{n}_{0}\sim p_{\textup{data}}$.

$\mathrm{d}\mathbf{X}^{n}_{t}=f^{n}(t,\mathbf{X}_{t}^{n})\mathrm{d}t+g(t)\mathrm{d}\mathbf{B}_{t}^{\mathcal{M}},$
(6)
$\mathrm{d}\mathbf{Y}^{n}_{t}=b^{n}(T-t,\mathbf{Y}_{t}^{n})\mathrm{d}t+g(t)\mathrm{d}\tilde{\mathbf{B}}_{t}^{\mathcal{M}}.$
(7)

1: for $n\in\\{0,\dots,L\\}$ do

2: while not converged do

3: Sample $t_{i}\sim\textrm{Uniform}([0,T])$

4: Simulate $\\{\mathbf{X}^{i}_{t_{i}}\\}_{i=0}^{B}$, where
$\mathbf{X}^{i}_{0}\sim p_{\textup{data}}$

5: Compute $\hat{\ell}^{b}_{n}(\phi^{n})$ using (8)

6: $\phi^{n}\leftarrow\textrm{Gradient Step}(\hat{\ell}^{b}_{n}(\phi^{n}))$

7: end while

8: while not converged do

9: Sample $t_{i}\sim\textrm{Uniform}([0,T])$

10: Simulate $\\{\mathbf{Y}^{i}_{t_{i}}\\}_{i=0}^{B}$, where
$\mathbf{Y}^{i}_{T}\sim p_{\textup{prior}}$

11: Compute $\hat{\ell}^{f}_{n+1}(\theta^{n+1})$ using (9)

12: $\theta^{n+1}\leftarrow\textrm{Gradient
Step}(\hat{\ell}^{f}_{n+1}(\theta^{n+1}))$

13: end while

14: end for

15: Output:  $(\theta^{L+1},\phi^{L})$

Algorithm 2 RDSB

######  Proposition 2.1.

Let $\mathbb{P}$ be the path measure of the Brownian motion initialized at
$p_{\textup{prior}}$ and IPF iterates $(\mathbb{Q}^{n})_{n}$ be as defined in
Section 1.3. Assume that for any $n\in\mathbb{N}$,
$\operatorname{KL}\left(\mathbb{Q}^{n}|\mathbb{P}\right)<+\infty$ and that for
any $t\in\left[0,T\right]$ and $n\in\mathbb{N}$, $\mathbb{Q}^{n}_{t}$ admits a
smooth positive density w.r.t. $p_{\textup{prior}}$. Then, for any
$n\in\mathbb{N}$: $\mathbb{Q}^{2n}$ and $R(\mathbb{Q}^{2n+1})$ solve the time-
reversal for (6) and (7) respectively. $R(\mathbb{P})_{t}=\mathbb{P}_{T-t}$
denotes the reverse time and for any $n\in\mathbb{N}$, $t\in\left[0,T\right]$
and $x\in\mathcal{M}$, $b^{n}(t,x)=-f^{n}(t,x)+g(t)^{2}\nabla\log
p^{n}_{t}(x)$, $f^{n+1}(t,x)=-b^{n}(t,x)+g(t)^{2}\nabla\log q^{n}_{t}(x)$,
with $f^{0}(t,x)=0$, and $p^{n}_{t}$, $q_{t}^{n}$ the densities of
$\mathbb{Q}^{2n}_{t}$ and $\mathbb{Q}_{t}^{2n+1}$.

Proof The proof follows De Bortoli et al. (2021, Proposition 6) using De
Bortoli et al. (2022, Theorem 1) instead of Cattiaux et al. (2021, Theorem
4.19)

In particular, we have that $\mathbb{Q}^{1}$ is the diffusion process
associated with RSGM, i.e. the time-reversal of the Brownian motion
initialized at $p_{\textup{prior}}$. Hence, $\mathbb{Q}^{2n+1}$ for
$n\in\mathbb{N}$ with $n\geq 1$ can be seen as a refinement of
$\mathbb{Q}^{1}$. In the next proposition, we show that the drift term of the
diffusion processes associated with $(\mathbb{Q}^{n})_{n\in\mathbb{N}}$ can be
approximated leveraging score-based techniques.

######  Proposition 2.2.

Let $(\mathbf{X}_{t})_{t\in\left[0,T\right]}$ be a $\mathcal{M}$-valued
process with distribution
$\mathbb{P}\in\mathcal{P}(\mathrm{C}(\left[0,T\right]),\mathcal{M})$ such that
for any $t\in\left[0,T\right]$, $\mathbf{X}_{t}$ admits a positive density
$p_{t}\in\mathrm{C}^{\infty}(\mathcal{M})$ w.r.t. $p_{\textup{prior}}$. For
any $t\in\left[0,T\right]$ and $x\in\mathcal{M}$, let
$b(t,x)=-f(t,x)+g(t)^{2}\nabla\log p_{t}(x)$. Then, for any
$t\in\left[0,T\right]$, we have that

$b(t,\cdot)=\operatorname*{arg\,min}_{r\in\mathrm{L}^{2}(\mathbb{P}_{t})}\mathbb{E}[\tfrac{1}{2}\|f(t,\mathbf{X}_{t})+r(\mathbf{X}_{t})\|_{2}^{2}+g(t)^{2}\mathrm{div}(r)(\mathbf{X}_{t})].$

Proof. For any $x\in\mathcal{M}$, $t\in\left[0,T\right]$,
$b(t,\cdot)=\operatorname*{arg\,min}_{r\in\mathcal{X}(\mathcal{M})}\mathbb{E}[\|r(t,\mathbf{X}_{t})-\\{-f(t,\mathbf{X}_{t})+g(t)^{2}\nabla\log
p_{t}(\mathbf{X}_{t})\\}\|]_{2}^{2}$. Expanding the quadratic, dropping terms
not dependent on $r$ and using the divergence theorem, (see Lee, 2018, p.51),
$\mathbb{E}[\langle r(\mathbf{X}_{t}),\nabla\log
p_{t}(\mathbf{X}_{t})\rangle_{\mathcal{M}}]=-\mathbb{E}[\mathrm{div}(r)(\mathbf{X}_{t})]$
concludes the proof.

$\displaystyle\hat{\ell}^{b}_{n}(\phi)$
$\displaystyle\textstyle{=\sum_{i=1}^{B}\tfrac{1}{2}\|f^{n}_{\theta}(t_{i},x^{i}_{t_{i}})+b^{n}_{\phi}(t_{i},x^{i}_{t_{i}}))\|_{2}^{2}+g^{2}(t_{i})\mathrm{div}(b^{n}_{\phi})(t_{i},x^{i}_{t_{i}})}$
(8) $\displaystyle\hat{\ell}^{f}_{n}(\theta)$
$\displaystyle\textstyle{=\sum_{i=1}^{B}\tfrac{1}{2}\|b^{n-1}_{\phi}(t_{i},x^{i}_{t_{i}})+f^{n}_{\theta}(t_{i},x^{i}_{t_{i}})\|_{2}^{2}+g^{2}(t_{i})\mathrm{div}(f^{n}_{\theta})(t_{i},x^{i}_{t_{i}})}$
(9)

In practice, for $L$ IPF steps, $(f^{n})_{n=1}^{L}$ and $(b^{n})_{n=0}^{L}$
are approximated using neural networks, $(f^{n}_{\theta^{n}})_{n=1}^{L}$,
$(b^{n}_{\phi^{n}})_{n=1}^{L}$ with parameters $(\theta^{n},\phi^{n})_{n}$.
Approximating drifts or means of the SDE is computationally cheaper than
storing an evaluating $2L$ separate score networks. Proposition 2.2 provides
loss functions. The divergence terms may be estimated using automatic
differentiation or Hutchinson’s trace estimator (Hutchinson, 1989).

(a) Earthquake

(b) Fire

(c) Volcano

(d) Flood

Figure 2: Climate data: $N=10$ steps, $5$ IPF iterations. Red points are
training data and green are generated points.

Likelihood Computation. The Ordinary Differential Equation (ODE),
$\mathrm{d}\hat{\mathbf{X}}_{t}=\\{f(t,\hat{\mathbf{X}}_{t})-\tfrac{1}{2}g(t)^{2}\nabla\log
p_{T-t}(\hat{\mathbf{X}}_{t})\\}\mathrm{d}t$, has the same marginal
probabilities as SDE (1), in particular $p_{\textup{data}}$ and hence may be
used for likelihood computation using an adaptive ODE solver and trained score
(Song et al., 2021b). Similar results hold for the Schrödinger Bridge (De
Bortoli et al., 2021; Chen et al., 2022) and in the Riemannian setting (De
Bortoli et al., 2022). We construct the appropriate probability flow ODE using
drift approximations as
$\mathrm{d}\hat{\mathbf{X}}_{t}=\frac{1}{2}\\{f_{\theta}(t,\hat{\mathbf{X}}_{t})-b_{\phi}(t,\hat{\mathbf{X}}_{t})\\}\mathrm{d}t$.
The computed likelihood assumes convergence of the forward noising process to
be valid, in particular that $p_{T}=p_{\textup{prior}}$. RDSB enforces this
convergence.

Other acceleration methods. Leading SGM acceleration techniques (Xiao et al.,
2022; Song et al., 2021a) use an implicit approach to denoising by estimating
$x_{0}$ from $x_{t}$, then sampling $\mathbf{X}_{s}|x_{t},x_{0}$ for $s<t$.
These techniques are not applicable in the Riemannian setting as it is not
typically possible to sample $\mathbf{X}_{s}|x_{t},x_{0}$ for $t>s$ or
$\mathbf{X}_{0}|x_{t}$ for large jumps $t\not\approx s$, $t\not\approx 0$.

##  3 Experiments

In each of the experiments, $4$-layer, fully connected networks with hidden
width $512$ are used for both $f_{\theta}$, $b_{\phi}$. $N=10$ diffusion steps
of size $1/N$ was used. Triangular schedule $g^{2}(t)$ was selected, linearly
interpolated from a peak of $0.05$ at $t=T/2$ and low of $0.001$ at $t=0,T$.
Other experimental details follow De Bortoli et al. (2022).

###  3.1 Earth and Climate Data

We validate RDSB on empirical distributions of occurrences of real Earth and
climate science events including earthquakes (NGDC/WDS, 2022a); wild fires
(EOSDIS, 2020); volcanic eruptions (NGDC/WDS, 2022b), and floods (Brakenridge,
2017). Figure 2 illustrates that generated samples are visually similar to the
empirical datasets.

Close inspection of Figure 1 shows RDSB with $5$ IPF iterations exhibits
better convergence than RSGM (equivalent to RDSB with $1$ IPF iteration) for
$N=10$ diffusion steps on the earthquake data. In particular, RDSB generates
fewer outlying samples to true data samples.

Figure 3: Interpolation between spherical harmonic datasets $(l=2,m=4)$ and
$(l=2,m=4)$ after $2$ IPF iterations

###  3.2 Dataset Interpolation on a Manifold

RDSB enables interpolation between datasets by choosing $p_{\textup{prior}}$
to be another dataset, and not necessary $p_{T}$, the terminal distribution of
the inital noising process. Figure 3 illustrates RDSB applied to samples on
the sphere taken with probability proportional to the real component of
spherical harmonic function with parameters $(l=2,m=4)$ and $(l=2,m=6)$.

##  4 Future Directions

We introduce novel methodology for generative modeling and interpolation on
compact Riemannian manifolds, which accelerates and generalizes RSGM. In
future work, we will apply RDSB to more challenging settings such as
interpolation in robotics and for protein modeling, see (Jing et al., 2022).
From a theoretical point of view, the restriction to compact spaces allows us
to leverage tools from Optimal Transport (Peyré & Cuturi, 2019) to prove the
geometric convergence of RDSB to an approximation of the Schrödinger Bridge
and quantify this bias.

## References

  * Brakenridge (2017) Brakenridge, G.  Global active archive of large flood events.  http://floodobservatory.colorado.edu/Archives/index.html, 2017.
  * Cattiaux et al. (2021) Cattiaux, P., Conforti, G., Gentil, I., and Léonard, C.  Time reversal of diffusion processes under a finite entropy condition.  _arXiv preprint arXiv:2104.07708_ , 2021.
  * Chen et al. (2022) Chen, T., Liu, G.-H., and Theodorou, E. A.  Likelihood training of Schrödinger bridge using forward-backward SDEs theory.  In _International Conference on Learning Representations_ , 2022.
  * Chen et al. (2016) Chen, Y., Georgiou, T., and Pavon, M.  Entropic and displacement interpolation: a computational approach using the Hilbert metric.  _SIAM Journal on Applied Mathematics_ , 76(6):2375–2396, 2016.
  * De Bortoli et al. (2021) De Bortoli, V., Thornton, J., Heng, J., and Doucet, A.  Diffusion Schrödinger bridge with applications to score-based generative modeling.  In _Advances in Neural Information Processing Systems_ , 2021.
  * De Bortoli et al. (2022) De Bortoli, V., Mathieu, E., Hutchinson, M., Thornton, J., Teh, Y. W., and Doucet, A.  Riemannian score-based generative modeling.  _arXiv preprint arXiv:2202.02763_ , 2022.
  * EOSDIS (2020) EOSDIS.  Land, atmosphere near real-time capability for eos (lance) system operated by NASA’s earth science data and information system (esdis).  https://earthdata.nasa.gov/earth-observation-data/near-real-time/firms/active-fire-data, 2020\.
  * Haussmann & Pardoux (1986) Haussmann, U. G. and Pardoux, E.  Time reversal of diffusions.  _The Annals of Probability_ , 14(4):1188–1205, 1986.
  * Hutchinson (1989) Hutchinson, M. F.  A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines.  _Communications in Statistics-Simulation and Computation_ , 18(3):1059–1076, 1989.
  * Jing et al. (2022) Jing, B., Corso, G., Barzilay, R., and Jaakkola, T. S.  Torsional diffusion for molecular conformer generation.  In _ICLR Workshop on Machine Learning for Drug Discovery_ , 2022.
  * Lee (2018) Lee, J. M.  _Introduction to Riemannian manifolds_.  Springer, 2018.
  * Léonard (2012) Léonard, C.  From the Schrödinger problem to the Monge–Kantorovich problem.  _Journal of Functional Analysis_ , 262(4):1879–1920, 2012.
  * NGDC/WDS (2022a) NGDC/WDS.  Ncei/wds global significant earthquake database.  https://www.ncei.noaa.gov/access/metadata/landing-page/bin/iso?id=gov.noaa.ngdc.mgg.hazards:G012153, 2022a.
  * NGDC/WDS (2022b) NGDC/WDS.  Ncei/wds global significant volcanic eruptions database.  https://www.ncei.noaa.gov/access/metadata/landing-page/bin/iso?id=gov.noaa.ngdc.mgg.hazards:G10147, 2022b.
  * Nutz & Wiesel (2022) Nutz, M. and Wiesel, J.  Stability of Schrödinger potentials and convergence of Sinkhorn’s algorithm.  _arXiv preprint arXiv:2201.10059_ , 2022.
  * Peyré & Cuturi (2019) Peyré, G. and Cuturi, M.  Computational optimal transport.  _Foundations and Trends® in Machine Learning_ , 11(5-6):355–607, 2019.
  * Sinkhorn (1967) Sinkhorn, R.  Diagonal equivalence to matrices with prescribed row and column sums.  _The American Mathematical Monthly_ , 74(4):402–405, 1967.
  * Song et al. (2021a) Song, J., Meng, C., and Ermon, S.  Denoising diffusion implicit models.  In _International Conference on Learning Representations_ , 2021a.
  * Song & Ermon (2019) Song, Y. and Ermon, S.  Generative modeling by estimating gradients of the data distribution.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Song et al. (2021b) Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B.  Score-based generative modeling through stochastic differential equations.  In _International Conference on Learning Representations_ , 2021b.
  * Vargas et al. (2021) Vargas, F., Thodoroff, P., Lamacraft, A., and Lawrence, N.  Solving Schrödinger bridges via maximum likelihood.  _Entropy_ , 23(9):1134, 2021.
  * Xiao et al. (2022) Xiao, Z., Kreis, K., and Vahdat, A.  Tackling the generative learning trilemma with denoising diffusion GANs.  In _International Conference on Learning Representations_ , 2022.