# Stochastic Mirror Descent in Average Ensemble Models

Taylan Kargin
Department of Electrical Engineering
California Institute of Technology
Pasadena, CA 91125
&Fariborz Salehi
Google
Seattle, WA 98103
&Babak Hassibi
Department of Electrical Engineering
California Institute of Technology
Pasadena, CA 91125

###### Abstract

The stochastic mirror descent (SMD) algorithm is a general class of training
algorithms, which includes the celebrated stochastic gradient descent (SGD),
as a special case. It utilizes a mirror potential to influence the implicit
bias of the training algorithm. In this paper we explore the performance of
the SMD iterates on mean-field ensemble models. Our results generalize earlier
ones obtained for SGD on such models. The evolution of the distribution of
parameters is mapped to a continuous time process in the space of probability
distributions. Our main result gives a nonlinear partial differential equation
to which the continuous time process converges in the asymptotic regime of
large networks. The impact of the mirror potential appears through a
multiplicative term that is equal to the inverse of its Hessian and which can
be interpreted as defining a gradient flow over an appropriately defined
Riemannian manifold. We provide numerical simulations which allow us to study
and characterize the effect of the mirror potential on the performance of
networks trained with SMD for some binary classification problems.

_K_ eywords Stochastic Mirror Descent $\cdot$ Ensemble Models $\cdot$ Mean-
field limit

##  1 Introduction

Machine learning models have been successfully used in a wide variety of
applications ranging from spam detection, to image classification, to natural
language processing. Despite this great success, our theoretical understanding
on why various machine learning methods demonstrate the performances they do
is still at an early stage.

To obtain some understanding of the behavior of machine learning algorithms we
shall focus on the structure/architecture of the model, as well as the
training mechanism that generates a suitable model with respect to the
available data. With regards to the architecture, a common theme among most
modern machine learning models (e.g. artificial/convolutional/recurrent neural
networks) is that they repetitively use simple blocks such as perceptrons in
ANNs, and convolution filters in CNNs. The idea of combining multiple simple
models has its roots in classical statistics and is known as boosting. The
justification is that a simple model often possesses a low variance while its
bias can be potentially high; therefore, combining multiple such models can
reduce the bias while keeping the variance low.

The optimization algorithm maps the training data to model parameters. Its
role is especially critical in modern setups where the number of parameters
can be overwhelmingly larger than the number of inputs. In such
overparameterized settings, there are often many models that perfectly fit the
training data. Thus, the specific choice of the optimization algorithm
determines the resulting model in the interpolating regime. In fact,
understanding the connection between the optimization algorithm and the
generalization performance of the resulting model is an important open
problem.

In this paper we attempt to study these aspects by analyzing the convergence
behavior of the stochastic mirror descent (SMD) iterates on a general class of
models, known as _average ensemble models_ where the output is generated by
taking an average of similar models that only differ in their choice of
parameters. Two-layer neural networks, e.g., can be viewed as a special case
of the average ensemble.

In particular, we address the impact of the potential function used in SMD on
the parameters of the average ensemble model. Introduced by Nemirovskii and
Yudin [1983] for convex optimization problems with underlying geometric
structure, SMD is as a generalization of stochastic gradient descent (SGD)
where the updates are _mirrored_ to a dual domain using a _potential
function_. It has been shown that choice of the potential function plays a
pivotal role in determining the implicit bias of the SMD in the interpolating
regime Azizan et al. [2019], Gunasekar et al. [2020a].

###  1.1 Contributions and Prior Work

In this paper, we develop a mean-field type result for average ensemble models
trained with SMD. We consider a general class of convex loss functions that
includes the squared loss as a special case. In particular, we derive a
partial differential equation that describes the "time evolution" of the
distribution of the parameters in the network as training progresses. This
distribution can be used to compute the risk, or generalization error, of the
trained model. Our result generalizes similar mean-field results for SGD, with
the difference being that the Hessian of the potential function appears. We
further give a geometric interpretation at the "distributional" level by
showing an equivalent formulation of our result as a geometric flow on an
appropriate Riemannian manifold. We finally provide numerical results to
illustrate the applicability of the method and to showcase the effect of the
potential function on the performance and implicit bias of the resulting
trained network.

A mean-field characterization of two-layer neural networks trained with SGD
under quadratic loss was studied by Mei et al. [2018], Sirignano and
Spiliopoulos [2020], and Rotskoff and Vanden-Eijnden [2019]. A similar line of
work is given by Chizat and Bach [2018] where more general loss functions are
considered for gradient descent algorithm. Mean field results with similar
flavor to ours have been given in Raginsky and Bouvrie [2012] and Borovykh et
al. [2020] in the context of distributed convex optimization. These papers
consider the problem of minimizing a fixed objective function with noisy
information by many "interacting optimizers" coupled through an interaction
matrix. Despite being insightful, the setting presented in these papers does
not fit into the case of learning with average ensemble models. An important
connection between mirror descent and natural gradient descent on dual
(Riemannian) space was established by Raskutti and Mukherjee [2014]. Works by
Gunasekar et al. [2020a] and Gunasekar et al. [2020b] advanced this geometric
viewpoint by showing a correspondence between continuous-time mirror descent
and Riemannian gradient flow on the primal domain governed by the Hessian
matrix of the potential function.

##  2 Preliminaries

###  2.1 Notations

We gather here the basic notations used. Vectors are presented with bold
lower-case letters, and bold upper letters are reserved for matrices. For a
vector $\mathbf{v}$, $v_{i}$ denotes its $i^{\text{th}}$ entry, and
$\|\mathbf{v}\|_{p}$ is its $\ell_{p}$ norm. $\mathbb{S}^{d}_{++}$ is the set
of symmetric and positive-definite matrices with dimension $d$.

For a set $\mathcal{S}$ in $\mathbb{R}^{d}$ and $k\in\mathbb{N}$,
$\mathcal{C}_{b}^{k}(\mathcal{S})$ denotes the space of bounded and $k$-times
differentiable functions with continuous $k^{th}$ derivative,
$\mathcal{P}(\mathcal{S})$ denotes the space of Borel probability measures on
$\mathcal{S}$ and $\mathcal{P}_{m}(\mathcal{S})$ denotes probability measures
with bounded $m^{\text{th}}$ moment.

For a random vector $\mathbf{Z}$ in $\mathbb{R}^{p}$, $\mu(d\mathbf{z})$
indicates its probability measure. $\delta_{\mathbf{x}}$ is a Dirac mass at
point $\mathbf{x}$. Integral of a function
$f:\mathcal{S}\rightarrow\mathbb{R}$ with respect to a measure $\rho$ on
$\mathcal{S}$ is denoted as,

$\langle\rho,f\rangle:=\int f(\mathbf{x})\rho(d\mathbf{x}).$ (1)

Let $\\{\rho_{n}\\}_{n\in\mathbb{N}}$ be a sequence of probability measures
defined on $\mathcal{S}\subseteq\mathbb{R}^{d}$. The sequence is said to be
weakly converging to the measure $\rho^{\star}$ if and only if,
$\underset{n\rightarrow\infty}{\lim}|\langle\rho_{n},f\rangle-\langle\rho^{\star},f\rangle|=0$,
for all $f\in\mathcal{C}_{b}(\mathcal{S})$.

###  2.2 Problem Setup

Consider an online learning setting where we sequentially observe a fresh
sample of data, i.e., $(\mathbf{x}^{k},y^{k})$ for $k=1,2,\ldots$. Here
$\mathbf{x}^{k}\in\mathcal{X}\subseteq\mathbb{R}^{p}$ is the feature vector,
and $y^{k}\in\mathcal{Y}\subseteq\mathbb{R}$ is the label. As a shorthand, we
will interchangeably use $\mathbf{z}:=(\mathbf{x},y)$. Our goal is to learn a
model that generates an estimate of the label given the feature vector as its
input. A common approach for modeling is known as the average ensemble, where
the output is generated by taking the average of simpler models that only
differ in their parameters. This approach is advantageous from the bias-
variance perspective. Since simpler models have a low variance (yet
potentially a high bias), the average ensemble generates a low-variance
estimate of the labels with a reasonably good bias.
Given a parameter space $\mathcal{S}\subset\mathbb{R}^{d}$ and
$\sigma:\mathcal{X}\times\mathcal{S}\rightarrow\mathcal{Y}$, consider the
following parametric family of model functions,

$\mathcal{H}_{\sigma}=\\{\sigma(\cdot,\bm{\theta}):\mathcal{X}\rightarrow\mathcal{Y};\;\bm{\theta}\in\mathcal{S}\\}.$
(2)

This parametric representation of model functions has applications in various
settings. For instance, by setting $d=p$ and defining
$\sigma(\mathbf{v}_{1},\mathbf{v}_{2})=\sigma^{\star}(\mathbf{v}_{1}^{T}\mathbf{v}_{2})$,
the class $\mathcal{H}_{\sigma}$ represents the generalized linear models with
nonlinearity $\sigma^{\star}:\mathbb{R}\rightarrow\mathbb{R}$.

The average ensemble model generates an estimate by taking the average of $n$
models in $\mathcal{H}_{\sigma}$. In other words, for
$\mathbf{x}\in\mathcal{X}$ the estimate is computed as follows,

$\hat{y}=h_{n}(\mathbf{x};\,\bm{\Theta}):=\frac{1}{n}\sum_{i=1}^{n}\sigma(\mathbf{x},{\bm{\theta}}_{i}),$
(3)

where $\bm{\Theta}=\begin{pmatrix}\bm{\theta}_{1},\leavevmode\nobreak\
\bm{\theta}_{2},\ldots,\leavevmode\nobreak\
\bm{\theta}_{n}\end{pmatrix}\in\mathcal{S}^{n}$.

As explained above, average ensemble models are used in many practical
settings. A popular instance of such models is the two-layer neural network,
where
${\bm{\theta}}_{i}=\begin{bmatrix}\bar{\bm{\theta}}_{i},\leavevmode\nobreak\
w_{i}\end{bmatrix}^{T}$, and the output is computed as,

$\hat{y}_{NN}=\sum_{i=1}^{n}\frac{w_{i}}{n}\bar{\sigma}(\mathbf{x},\bar{\bm{\theta}}_{i}),$
(4)

where $\mathbf{w}\in\mathbb{R}^{n}$ represents the weights of the last layer
of the network and $\bar{\sigma}(\cdot)$ is the activation function.

Given a loss function $\ell:\mathbb{R}^{2}\rightarrow\mathbb{R}$ and a
training dataset $\\{(\mathbf{x}^{k},y^{k})\\}_{k=1}^{K}$, our goal is to
minimize the empirical risk

$\inf_{\Theta\in\mathcal{S}^{n}}\left\\{\hat{R}_{n}(\Theta):=\frac{1}{K}\sum_{k=1}^{K}\ell(y^{k},\,h_{n}(\mathbf{x}^{k};\,\Theta))\right\\}\leavevmode\nobreak\
,$ (5)

over parameters $\bm{\Theta}$. We assume that the training data points are
generated independently from a distribution, i.e.,
$(\mathbf{x}^{k},y^{k})\sim\mu(d\mathbf{z})$. The loss function
$\ell:\mathbb{R}^{2}\rightarrow\mathbb{R}$ is assumed to be convex w.r.t. its
second argument. Important examples are the quadratic loss,
$\ell(u,v)=\frac{1}{2}(u-v)^{2}$, and the logistic loss,
$\ell(u,v)=\log\big{(}1+\exp(-u\times v)\big{)}$. However, the problem is non-
convex whenever $\sigma$ is. One can rewrite the average ensemble model as an
integral over the empirical distribution of the parameters,
$\hat{\rho}_{n}(d\bm{\theta}):=\sum_{i=1}^{n}\frac{1}{n}\delta_{\bm{\theta}_{i}}(d\bm{\theta})$,
i.e.,

$h_{n}(\mathbf{x};\,\bm{\Theta})=\langle\hat{\rho}_{n},\sigma(\mathbf{x},.)\rangle.$
(6)

This new form leads us to define ensemble averages by any probability
distribution on $\mathcal{S}$. Namely, we define

$h(\mathbf{x};\,\rho)=\langle\rho,\sigma(\mathbf{x},.)\rangle,$ (7)

for any $\rho\in\mathcal{P}(\mathcal{S})$. We call this the mean-field
ensemble model and the former the finite ensemble model. It is easy to check
that $h_{n}(\mathbf{x};\,\bm{\Theta})=h(\mathbf{x};\,\hat{\rho}_{n})$. Note
that whenever the loss function $\ell$ is convex, the mean-field ensemble
model function
$h(\mathbf{x};\,\cdot):\mathcal{P}(\mathcal{S})\rightarrow\mathcal{Y}$ is also
convex in the space of probability measures $\mathcal{P}(\mathcal{S})$
regardless of $\sigma(\cdot)$.

####  2.2.1 Stochastic Mirror Descent Updates

Mirror descent algorithms introduced by Nemirovskii and Yudin [1983], and
their stochastic variants, are commonly used iterative methods that exploit a
potential function to impose certain attributes in the resulting model. Assume
$\mathcal{X}=\mathbb{R}^{p}$, $\mathcal{Y}=\mathbb{R}$, and
$\mathcal{S}=\mathbb{R}^{d}$. Here, we consider the mirror descent updates
initialized at $\bm{\Theta}^{0}\in\mathcal{S}^{n}$ where, for
$i=1,2,\ldots,n$, the updates are defined as,

$\nabla\psi(\bm{\theta}_{i}^{k+1})=\nabla\psi(\bm{\theta}_{i}^{k})-\frac{\tau}{n}\nabla_{\bm{\theta}_{i}}\ell(y^{k+1},\,h_{n}(\mathbf{x}^{k+1};\,\Theta^{k})),$
(8)

and $k$ runs from $0$ to $K-1$. Here $\psi(\cdot)$ is a strongly convex and
differentiable potential function, $\tau$ a fixed scaling of the step-size,
and $\bm{\Theta}^{k}$ the parameters after the $k^{\text{th}}$ update. Strong
convexity of $\psi(\cdot)$ ensures the invertibility of its gradient which
makes the updates well-defined (see Bubeck [2015]).

After evaluating the gradient of the loss function (w.r.t. $\bm{\theta}_{i}$)
and using the identity (6), we can rewrite the SMD updates for
$i=1,2,\ldots,n$ and $k=0,\dots,K-1$ as

$\begin{cases}\hat{\rho}^{k}_{n}=\frac{1}{n}\sum_{i=1}^{n}\delta_{\bm{\theta}_{i}^{k}}\leavevmode\nobreak\
,\\\
\nabla\psi(\bm{\theta}_{i}^{k+1})=\nabla\psi(\bm{\theta}_{i}^{k})+\frac{\tau}{n}\mathbf{F}(\bm{\theta}_{i}^{k},\,\hat{\rho}^{k},\,\mathbf{z}^{k+1})\leavevmode\nobreak\
,\end{cases}$ (9)

where ${\hat{\rho}}^{k}$ denotes the empirical distribution of the parameters
after the $k^{\text{th}}$ update. The function $\mathbf{F}$ is the gradient of
the loss w.r.t. $\bm{\theta}$, which is defined as,

$\mathbf{F}(\bm{\theta},{\rho},\mathbf{z}):=-\partial_{2}\ell\big{(}y;\langle\rho,\leavevmode\nobreak\
\sigma(\mathbf{x},\cdot)\rangle\big{)}\nabla_{\bm{\theta}}\sigma(\mathbf{x},\bm{\theta}),$
(10)

where $\partial_{2}$ indicates the derivative w.r.t. the second argument. It
is worth noting that the potential function, $\psi(\cdot)$, is often chosen to
enforce some structure on the resulting parameter. Setting
$\psi(\cdot)=\frac{1}{2}\|\cdot\|_{2}^{2}$ gives the SGD iterates.

_Our goal is to compute the (converging) distribution of the parameters in the
asymptotic regime where $n,K$ go to infinity at a fixed ratio._ This
distribution plays a pivotal role in understanding the generalization behavior
of the resulting model. In particular, the expected risk can be computed as,

$R(\rho):=\mathbb{E}_{(\mathbf{x},y)\sim\mu(d\mathbf{z})}\big{[}\ell\big{(}y;\langle\rho,\leavevmode\nobreak\
\sigma(\mathbf{x},\cdot)\rangle\big{)}\big{]}\leavevmode\nobreak\ ,$ (11)

where $\rho:=\rho(d\bm{\theta})$ denotes the (empirical) distribution.

##  3 Main Result

In this section, we present the main result of the paper, i.e., the
characterization of the converging distribution of the SMD iterates. The
resulting distribution is given as the solution to a continuous-time nonlinear
PDE. We consider the converging distribution in the asymptotic regime where
$K,n\rightarrow\infty$ at a fixed ratio,
$\delta:=\frac{K}{n}\in\mathbb{R}_{++}$.

In Section 3.1, we explain how to map SMD updates to a continuous-time process
on probability distributions in $\mathbb{R}^{d}$. Consequently, a PDE will be
derived in Section 3.2 whose evolution is captured by the Hessian of the
potential function, and the function $\mathbf{F}$ defined in (10). Finally,
Section 3.3 incorporates the main result of the paper which indicates that the
distribution of the mirror descent updates weakly converges to the solution of
the PDE.

We require the following set of assumptions.

######  Assumption 1.

  * (i)

$\sigma:\mathbb{R}^{p}\times\mathbb{R}^{d}\rightarrow\mathbb{R}$ is bounded
and $\sigma(\mathbf{x},\,\cdot)\in\mathcal{C}_{b}^{2}(\mathbb{R}^{d})$ for all
$\mathbf{x}\in\mathbb{R}^{p}$.

  * (ii)

The loss function $\ell:\mathbb{R}\times\mathbb{R}\rightarrow\mathbb{R}$ is
convex and $\mathcal{C}^{2}(\mathbb{R})$ w.r.t. its second argument.

  * (iii)

The mirror function $\psi:\mathbb{R}^{d}\rightarrow\mathbb{R}$ is
$\lambda-$strongly convex and $\psi\in\mathcal{C}^{2}(\mathbb{R}^{d})$.

  * (iv)

The initial parameters
$\bm{\Theta}^{0}=\begin{pmatrix}\bm{\theta}_{1}^{0},\leavevmode\nobreak\
\ldots,\leavevmode\nobreak\ \bm{\theta}_{n}^{o}\end{pmatrix}$ are sampled
i.i.d. from a distribution $\rho^{0}\in\mathcal{P}_{4}(\mathbb{R}^{d})$.

  * (v)

The data distribution has finite $4^{\text{th}}$-order moments, i.e.,
$\mu\in\mathcal{P}_{4}(\mathbb{R}^{p+1})$

Assumptions $(i)$ and $(v)$ are required for well-definiteness of expectations
and together with $(iv)$, they will guarantee our convergence result.
Convexity of loss function $\ell$ ensures convexity of the expected risk of
the mean-field ensemble which model (11). Assumptions $(ii)$ and $(iv)$
guarantee uniqueness of the resulting PDE. Finally, assumption $(iii)$ is
needed for non-singularity of the Hessian of the mirror function.

###  3.1 Continuous Time Viewpoint

Recall from the previous section ${\hat{\rho}}^{k}_{n}$, for
$k=0,1,\ldots,K-1$, is defined as the empirical measure of the parameters
after $k^{\text{th}}$ mirror update, i.e,
${\hat{\rho}}^{k}_{n}(d\bm{\theta}):=\frac{1}{n}\sum_{i=1}^{n}\delta_{\bm{\theta}_{i}^{k}}(d\bm{\theta})$.

To characterize the converging distribution of SMD, we exploit a similar
approach as Wang et al. [2017] and Sirignano and Spiliopoulos [2020] by
viewing $\\{\hat{\rho}^{k}_{n}\\}_{k=0}^{K-1}$ as a _Markov process in the
state space_ $\mathcal{P}(\mathbb{R}^{d})$. The Markov property follows from
(9), where $\\{\bm{\Theta}^{k}\\}_{k=0}^{K}$ forms an exchangeable Markov
process, i.e., any permutation of the indices
$\begin{pmatrix}\bm{\theta}_{1}^{k},\leavevmode\nobreak\
\ldots,\leavevmode\nobreak\ \bm{\theta}_{n}^{k}\end{pmatrix}$ will leave their
joint distribution intact.

We introduce $\bar{\rho}_{n}(t)$ which is a continuous-time embedding of the
discrete-time process $\\{\hat{\rho}^{k}_{n}\\}_{k=0}^{K-1}$, defined as,

$\bar{\rho}_{n}(t):=\hat{\rho}_{n}^{\lfloor{nt/\tau}\rfloor}\leavevmode\nobreak\
,\qquad t\in[0,T],$ (12)

where $T:=K\times\frac{\tau}{n}=\delta\tau$ is a positive constant.
$\bar{\rho}_{n}(t)$ is a piecewise constant process that is right-continuous
with left limits (RCLL). This process can be viewed as an element in the
Skorokhod space $\mathcal{D}([0,T];\mathcal{P}(\mathcal{S}))$111For a set
$\mathcal{X}$ and $T>0$,the Skorokhod space $\mathcal{D}([0,T],\mathcal{X})$
is the space of càdlàg functions from the interval $[0,T]$ to $\mathcal{X}$
which are right-continuous with left-limits.. We will show in Section 3.3 that
$\\{\bar{\rho}_{n}(t)\\}_{n\in\mathbb{N}}$ has a limit as
$n\rightarrow\infty$.

###  3.2 A Partial Differential Equation

Here we introduce a non-linear partial differential equation (PDE) whose
solution will determine the converging distribution of SMD updates. Let
$\mathbf{H}:\mathbb{R}^{d}\rightarrow\mathbb{S}^{d}_{++}$ be the Hessian
operator of the mirror potential function $\psi(\cdot)$, i.e.,

$\mathbf{H}(\bm{\theta}):=\frac{\partial^{2}}{\partial{\bm{\theta}}^{2}}\psi(\bm{\theta})\leavevmode\nobreak\
\leavevmode\nobreak\ ,\leavevmode\nobreak\ \leavevmode\nobreak\
\bm{\theta}\in\mathbb{R}^{d}.$ (13)

Note that $\mathbf{H}(\bm{\theta})\succeq\lambda\mathbf{I}_{d}$ due to the
strong convexity of the mirror potential. Consider the following PDE on
probability measures with the initial condition
$\rho_{0}\in\mathcal{P}(\mathbb{R}^{d})$,

$\partial_{t}\rho_{t}=-\nabla_{\bm{\theta}}\cdot\left(\rho_{t}\mathbf{H}(\bm{\theta})^{-1}\mathbf{v}(\bm{\theta},\rho_{t})\right),$
(14)

where $\nabla_{\bm{\theta}}\cdot\mathbf{J}$ denotes the divergence of the
vector field $\mathbf{J}$222For a vector field
$\mathbf{F}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$, defined as
$\mathbf{F}(\mathbf{x})=\begin{bmatrix}F_{1}(\mathbf{x}),F_{2}(\mathbf{x}),\ldots,F_{d}(\mathbf{x})\end{bmatrix}^{T}$,
the divergence is defined as,
$\nabla_{\mathbf{x}}\cdot\mathbf{F}=\frac{\partial
F_{1}}{\partial{x_{1}}}+\frac{\partial
F_{2}}{\partial{x_{2}}}+\ldots+\frac{\partial F_{d}}{\partial{x_{d}}}.$ , and
$\mathbf{v}:\mathbb{R}^{d}\times\mathcal{P}(\mathbb{R}^{d})\rightarrow\mathbb{R}^{d}$
is defined as,

$\displaystyle\mathbf{v}(\bm{\theta},\rho)$
$\displaystyle:=\mathbb{E}_{(\mathbf{x},y)\sim\mu(d\mathbf{z})}\left[\mathbf{F}(\bm{\theta},{\rho},\mathbf{x},y)\right],$
(15) $\displaystyle\leavevmode\nobreak\
=\int\mathbf{F}(\bm{\theta},{\rho},\mathbf{x},y)\mu(d\mathbf{z})\leavevmode\nobreak\
,$

where $\mathbf{F}$ is the negative gradient of the loss function defined in
(10). An alternative but equivalent definition for the velocity field
$\mathbf{v}$ is given as

$\mathbf{v}(\bm{\theta},\rho)=-\nabla_{\bm{\theta}}\frac{\delta
R}{\delta\rho}(\bm{\theta},\rho),$ (16)

where $\frac{\delta R}{\delta\rho}$ denotes the functional (Fréchet)
derivative of the expected risk functional $R(\rho)$ (11).

We note that the PDE (14) should be understood in the weak sense, i.e.,
$\rho_{t}$ is a solution of (14) if it solves the following equation for any
test function $f\in\mathcal{C}_{b}^{2}(\mathcal{S})$

$\displaystyle\frac{d}{d\,t}\langle\rho_{t},f\rangle$
$\displaystyle=\langle\rho_{t},\nabla
f(\cdot)^{T}\mathbf{H}(\cdot)^{-1}\mathbf{v}(\cdot,\rho_{t})\rangle,$ (17)
$\displaystyle=\int\nabla
f(\bm{\theta})^{T}\mathbf{H}(\bm{\theta})^{-1}\mathbf{v}(\bm{\theta},\rho_{t})\rho_{t}(d\bm{\theta}).$

The PDE (14) can be interpreted as a continuity equation describing the flow
of probability distribution driven by the "distorted" drift/velocity field
$\mathbf{H}^{-1}(\bm{\theta})\mathbf{v}(\bm{\theta},\rho)$. The distortion of
the drift term is due to inverse Hessian of the mirror potential while the
impact of the data distribution is captured by $\mathbf{v}$. In Section 4, we
provide an equivalent formulation for (14) in which the impact of mirror is
totally captured by a newly defined Riemannian metric. This leads to a
coordinate-free and abstract reformulation of (14).

######  Remark 1.

The PDE in (14) is presented in a general form. For specific choices of the
loss function, the drift term $\mathbf{v}$ can be simplified. There are
variety of approaches to solve this nonlinear PDE; a common one is fixed-point
iteration starting from a (reasonable) guess for the solution in the path
space and making iterative updates according to (14). Discussion of methods
for solving PDEs is beyond the scope of this paper, and we refer the reader to
Sznitman [1991] for more details.

####  3.2.1 Example: Squared loss

The PDE introduced in (14) applies to any convex loss function. The most
popular choice in regression problems is the squared loss. Here, we present a
simplified version of the velocity field $\mathbf{v}$ for the squared loss. To
this end, define the following two auxiliary functions.

$\begin{cases}V(\bm{\theta})=-\mathbb{E}_{(\mathbf{x},y)\sim\mu(d\mathbf{z})}\left[y\cdot\sigma(\mathbf{x},\bm{\theta})\right]\leavevmode\nobreak\
,\\\
U(\bm{\theta},\bm{\theta}^{\prime})=\mathbb{E}_{(\mathbf{x},y)\sim\mu(d\mathbf{z})}\left[\sigma(\mathbf{x},\bm{\theta})\cdot\sigma(\mathbf{x},\bm{\theta}^{\prime})\right]\leavevmode\nobreak\
.\end{cases}$ (18)

Then velocity field $\mathbf{v}(\bm{\theta},\rho)$ can be written as,

$\mathbf{v}(\bm{\theta},\rho)=-\nabla_{\bm{\theta}}V(\bm{\theta})-\langle\rho,\nabla_{\bm{\theta}}U(\bm{\theta},\cdot)\rangle.$
(19)

One main advantage of having this equivalent formulation is that the auxiliary
functions can be computed for a broad class of distributions $\nu$, and
nonlinear functions $\sigma$. For our numerical simulations in Section 5 we
compute $U$ and $V$ for
$\sigma(\mathbf{x},\bm{\theta})=\text{erf}\left(\mathbf{x}^{T}{\bm{\theta}}\right).$333For
$z\in\mathbb{R}$, $\text{erf}(z)$ is defined as,
$\text{erf}(z):=\frac{2}{\sqrt{\pi}}\int_{0}^{z}e^{-x^{2}}dx.$ Analytical
expressions for $V$ and $U$ for single-hidden-layer networks, Gaussian mixture
data models, and general $\sigma$ can be derived, though we shall not include
them here for lack of space.

###  3.3 Converging distribution

Our main result is Theorem 1, which states that the time evolution of the
empirical distribution of the parameters converges to the solution of the PDE
(14) (initialized with ${\rho}^{0}$). Equivalently, the empirical distribution
after $k=r\cdot n$ updates can be approximated the solution of the PDE at time
$t=r\cdot\tau$.

######  Theorem 1 (Converging Distribution of SMD).

Let $\\{(\mathbf{x}^{k},y^{k})\\}_{k=0}^{K-1}$ denote a dataset where
$(\mathbf{x}^{k},y^{k})\sim\mu(d\mathbf{z})$. Consider the stochastic mirror
descent updates in (9) and assume that Assumption 1 holds. As
$n,K\rightarrow\infty$ at a fixed ratio $\delta=\frac{K}{n}$, the continuous-
time process $\\{\bar{\rho}_{n}(t):t\in[0,T]\\}$ defined in (12) converges
weakly to $\\{\rho_{t}^{\star}:t\in[0,T]\\}$ in the Skorokhod path space
$\mathcal{D}([0,T];\mathcal{P}(\mathcal{S}))$ where $\rho_{t}^{\star}$ is the
solution of the following PDE for $t\in[0,T]$,

$\partial_{t}{\rho_{t}^{\star}}=-\nabla_{\bm{\theta}}\cdot\left({\rho_{t}^{\star}}\mathbf{H}(\bm{\theta})^{-1}\mathbf{v}(\bm{\theta},{\rho_{t}^{\star}})\right),$

with $\rho_{0}^{\star}={\hat{\rho}}^{0}$. In other words, the empirical
distribution of parameters resulting from SMD iterates converges to the
solution of the partial differential equation (14).

######  Remark 2.

As a side technical note, the convergence in Theorem 1 is in Skorokhod
distance, which is a distance defined on càdlàg processes. For a given
$f\in\mathcal{C}_{b}^{2}(\mathbb{R}^{d})$, denote by
$\bar{\rho}_{n}[f]:=\\{\langle\bar{\rho}_{n}(t),f\rangle\,:\,t\in[0,T]\\}$ and
$\rho^{\star}[f]:=\\{\langle{\rho_{t}^{\star}},f\rangle\,:\,t\in[0,T]\\}$ the
processes generated by projecting function $f$ onto
$\mathcal{D}([0,T];\mathcal{\mathbb{R}})$. Theorem 1 states that the
(Skorokhod) distance between these two càdlàg processes converges to zero in
probability as $n\rightarrow\infty$, i.e.,

$d_{\mathcal{D}[0,T]}\left(\bar{\rho}_{n}[f],\,\rho^{\star}[f]\right)\overset{\mathbb{P}}{\longrightarrow}\leavevmode\nobreak\
0\leavevmode\nobreak\ ,$ (20)

where $d_{\mathcal{D}[0,T]}$ is the Skorokhod distance on
$\mathcal{D}([0,T];\mathcal{\mathbb{R}})$. We refer the interested reader to
Billingsley [1999] and Kallenberg [2010] for a thorough discussion on
convergence of probability measures and random processes.

An outline of the proof is provided in Section 6.
Notice that when the mirror potential is chosen as
$\psi(\cdot)=\frac{1}{2}\|\cdot\|_{2}^{2}$, its Hessian becomes
$\mathbf{H}(\bm{\theta})=\mathbf{I}_{d}$, the identity matrix and the PDE (14)
takes the following simpler form.

$\partial_{t}\rho_{t}=-\nabla_{\bm{\theta}}\cdot\left(\rho_{t}\mathbf{v}(\bm{\theta},\rho_{t})\right).$
(21)

This corresponds to the specific case of stochastic gradient descent (SGD). A
neat interpretation of (21) as a Wasserstein gradient flow444For a given a
metric space $(\mathcal{X},d)$, the Wasserstein-2 distance,
$W_{2}:\mathcal{P}(\mathcal{X})\times\mathcal{P}(\mathcal{X})\rightarrow\mathbb{R}$
on $\mathcal{P}(\mathcal{X})$ is defined as
$W_{2}(\rho_{1},\rho_{2})=\left(\inf_{\gamma\in\Gamma(\rho_{1},\rho_{2})}\int
d(x_{1},x_{2})^{2}\gamma(dx_{1},dx_{2})\right)^{1/2}$ where
$\Gamma(\rho_{1},\rho_{2})$ is the set of probability distributions on
$\mathcal{P}(\mathcal{X}\times\mathcal{X})$ with marginals $\rho_{1}$ and
$\rho_{2}$. on the space of probability distributions
$\mathcal{P}(\mathbb{R}^{d})$ is given by Mei et al. [2018], Sirignano and
Spiliopoulos [2020], Chizat and Bach [2018], and Rotskoff and Vanden-Eijnden
[2019].

At a high level picture, this means that the solution of (21),
$t\mapsto\rho_{t}$ follows a trajectory on $\mathcal{P}(\mathbb{R}^{d})$ that
minimizes the expected risk functional $R(\rho)$ as time passes. This view
also gives a clear meaning to velocity field $\mathbf{v}(\bm{\theta},\rho)$.
Recalling the alternative definition given in (16), we can view
$\mathbf{v}(\bm{\theta},\rho)=-\nabla_{\bm{\theta}}\frac{\delta
R}{\delta\rho}(\bm{\theta},\rho)$ as the direction of "steepest descent" for
the expected risk functional $R(\rho)$ in the Euclidean sense. We refer the
reader to Ambrosio et al. [2008] for a comprehensive study on gradient flows.

When a general mirror potential is considered, the inverse Hessian
$\mathbf{H}^{-1}$ in the main equation $\eqref{eq:PDE}$ is posing a challenge
to this view since $\mathbf{H}^{-1}(\bm{\theta})\mathbf{v}(\bm{\theta},\rho)$
is not the steepest descent direction in the Euclidean way. In the next
section, we introduce an appropriately defined Riemannian manifold in order to
obtain a similar steepest descent interpretation.

##  4 Riemannian Formulation

In this section, we give an informal interpretation of equation (14) in a
geometric sense. The main idea is to reformulate PDE (14) on a Riemannian
manifold on which the metric tensor is defined by $\mathbf{H}(\bm{\theta})$.

######  Theorem 2.

Assume the entries of
$[\mathbf{H}(\bm{\theta})]_{ij}:\mathbb{R}^{d}\rightarrow\mathbb{R}$ be smooth
functions of $\bm{\theta}$, and equip the parameter space
$\mathcal{S}=\mathbb{R}^{d}$ with a Riemannian metric
$g_{ij}(\bm{\theta}):=[\mathbf{H}(\bm{\theta})]_{ij}$. Assume, in addition,
that the initial distribution admits a density with respect to volume measure
on $(\mathcal{S},g)$ as $\rho_{0}(\bm{\theta})$ =
$p(\bm{\theta},0)dV_{g}(\bm{\theta})$. Then, PDE (14) takes the following form
on the Riemannian manifold $(\mathcal{S},g)$

$\frac{\partial p(\bm{\theta},t)}{\partial
t}=\operatorname{div}_{g}\left(p(\bm{\theta},t)\operatorname{grad}_{g}\frac{\delta
R}{\delta\rho}(\bm{\theta},p_{t})\right).$ (22)

where $\operatorname{div}_{g}$ and $\operatorname{grad}_{g}$ denote the
divergence and gradient operators on $(\mathcal{S},g)$, respectively.

The insight behind this theorem is to define a new geometry on
$\mathcal{S}=\mathbb{R}^{d}$ from the Hessian $\mathbf{H}(\bm{\theta})$ of the
potential. Defining Riemannian metrics from Hessians of smooth functions is a
widely-used technique (Shima and Yagi [1997]).

For a comparison, we rewrite the PDE (21) by inserting the definition of
$\mathbf{v}$ as

$\partial_{t}\rho_{t}=\nabla_{\bm{\theta}}\cdot\left(\rho_{t}\nabla_{\bm{\theta}}\frac{\delta
R}{\delta\rho}(\bm{\theta},\rho_{t})\right).$ (23)

When contrasted with the above equation derived for Euclidean mirror
potential, the PDE (22) has the same form where divergence and gradient
operators on Euclidean space are replaced by their Riemannian counterparts.

In lieu of this observation, we give an interpretation of (22) as a geometric
flow describing the time evolution of parameter distribution on the Riemannian
manifold defined by the Hessian of the potential. In this new geometry,
steepest descent directions no longer follow straight lines, but rather follow
geodesics555Geodesics are the shortest distance paths connecting points on a
Riemannian manifold. imposed by the mirror potential.

###### Proof.

We start by defining the gradient operator on $(\mathcal{S},g)$. For a smooth
function $f\in\mathcal{C}^{\infty}(\mathcal{S})$, its gradient at a point
$\bm{\theta}\in\mathcal{S}$ with respect to metric $g$ is defined as

$\operatorname{grad}_{g}f:=\mathbf{H}(\bm{\theta})^{-1}\nabla_{\bm{\theta}}f(\bm{\theta})$
(24)

We thus obtain the following identity for $\mathbf{H}^{-1}\mathbf{v}$

$\mathbf{H}(\bm{\theta})^{-1}\mathbf{v}(\bm{\theta},\rho)=-\operatorname{grad}_{g}\frac{\delta
R}{\delta\rho}(\bm{\theta},\rho).$ (25)

Inserting this back to equation (14) gives

$\partial_{t}\rho_{t}=\nabla_{\bm{\theta}}\cdot\left(\rho_{t}\operatorname{grad}_{g}\frac{\delta
R}{\delta\rho}(\bm{\theta},\rho_{t})\right).$ (26)

For a smooth vector field $\mathbf{J}(\bm{\theta})$, the Riemannian divergence
operator on $(\mathcal{S},g)$ is defined as

$\operatorname{div}_{g}\mathbf{J}:=\frac{1}{\sqrt{\operatorname{det}{\mathbf{H}(\bm{\theta})}}}\nabla_{\bm{\theta}}\cdot\left(\sqrt{\operatorname{det}{\mathbf{H}(\bm{\theta})}}\mathbf{J}(\bm{\theta})\right)$
(27)

Note that $\sqrt{\operatorname{det}{\mathbf{H}(\bm{\theta})}}$ also appears in
the definition of the volume measure $dV_{g}(\bm{\theta})$ on
$(\mathcal{S},g)$ as,

$dV_{g}(\theta)=\sqrt{\operatorname{det}{\mathbf{H}(\bm{\theta})}}d\bm{\theta}.$
(28)

In other words, infinitesimal volumes on $(\mathcal{S},g)$ are distorted by a
factor of $\sqrt{\operatorname{det}{\mathbf{H}(\bm{\theta})}}$. Inserting the
new divergence operator on $(\mathcal{S},g)$ to (26), we rewrite it as

$\partial_{t}\frac{\rho_{t}}{\sqrt{\operatorname{det}{\mathbf{H}(\bm{\theta})}}}=\operatorname{div}_{g}\left(\frac{\rho_{t}}{\sqrt{\operatorname{det}{\mathbf{H}(\bm{\theta})}}}\operatorname{grad}_{g}\frac{\delta
R}{\delta\rho}(\bm{\theta},\rho_{t})\right).$ (29)

Using the assumption on the initial distribution, the solution of at any time
step can be expressed similarly

$\displaystyle\rho_{t}(d\bm{\theta})$
$\displaystyle=p(\bm{\theta},t)dV_{g}(\bm{\theta}),$ (30)
$\displaystyle=p(\bm{\theta},t)\sqrt{\operatorname{det}{\mathbf{H}(\bm{\theta})}}d\bm{\theta}.$

Then, the resulting (22) follows immediately. ∎

##  5 Numerical Simulations

In order to illustrate the applicability of our theoretical results, we have
performed several simulations by numerically solving the PDE (14). Due to lack
of space, we only include a single illustrative example here.

Consider a binary classification problem with two-dimensional inputs and a
data distribution given by

$\begin{cases}y=+1,\;\mathbf{x}\sim\text{GM}_{1},\quad\text{with probability
$\frac{1}{2}$},\\\ y=-1,\;\mathbf{x}\sim\text{GM}_{2},\quad\text{with
probability $\frac{1}{2}$}\end{cases}$ (31)

where $\text{GM}_{1}$ and $\text{GM}_{2}$ are two different Gaussian mixtures,
with three centers each, and isotropic covariance matrices

$\text{GM}_{i}(\mathbf{x})=\frac{1}{3}\sum_{j=1}^{3}\operatorname{\mathcal{N}}(\mathbf{x}\,|\,\mathbf{m}_{i}^{j},\,s^{2}\mathbf{I}_{2}),\;\;i=1,2$
(32)

Here $s=0.1$. The exact position of the centers
$\\{\mathbf{m}_{i}^{1},\mathbf{m}_{i}^{2},\mathbf{m}_{i}^{3}\\}$ are shown in
Figure 1 and are chosen so that they form two equilateral triangles (one for
each class) with geometric centers at the origin. Most importantly, centers of
the second class (blue) are exactly the midpoints of the edges of the triangle
formed by the centers of the first class (red). Due to the symmetry in the
data distribution, it is possible to separate the two pairs of centers for
each class using three lines through the origin.

Figure 1: The centers of the Gaussian mixture model for class 1 (red) and
class 2 (blue). Note that the three centers for each class form an equilateral
triangle with geometrical center at the origin. The centers for class 2 are
the midpoint of the edges connecting the centers of class 1. Finally, the
centers of the two classes can be separated by three lines with relative
angles $120^{o}$ passing through the origin.

We set the function
$\sigma(\mathbf{x},\bm{\theta})=\text{erf}(\theta^{T}\mathbf{x})$ and choose
the quadratic loss function $\ell(y,\,\hat{y})=\frac{1}{2}(y-\hat{y})^{2}$.
This corresponds to a single hidden-layered neural network with non-linearity
given by the error function erf. The choice of the error function allows a
simple closed form expression for the functions $V(\bm{\theta})$ and
$U(\bm{\theta},\,\bm{\theta}^{\prime})$ introduced in section 3.2.1, viz.,

$\displaystyle\nabla
V(\bm{\theta})=-\frac{b(\bm{\theta})^{-1/2}}{3\sqrt{\pi}}\sum_{i=1}^{2}y_{i}\sum_{j=1}^{3}a_{ij}(\bm{\theta})\hat{\mathbf{m}}_{i}^{j}$
(33)

where $b(\bm{\theta}):=1+2s^{2}\|\bm{\theta}\|_{2}^{2}$, and $y_{1}=+1$,
$y_{2}=-1$ as class labels and

$\displaystyle
a_{ij}(\bm{\theta}):=e^{\frac{(\bm{\theta}^{T}\mathbf{m}_{i}^{j})^{2}}{b(\bm{\theta})}}$
(34)
$\displaystyle\hat{\mathbf{m}}_{i}^{j}:=\mathbf{m}_{i}^{j}-\frac{2s^{2}\bm{\theta}^{T}\mathbf{m}_{i}^{j}}{b(\bm{\theta})}\bm{\theta}$

We also obtained a somewhat similar expression for
$\nabla_{1}U(\bm{\theta},\bm{\theta}^{\prime})$.

Finally, we considered the SMD algorithm for this data distribution for two
different mirror functions, $\|\cdot\|_{p}^{2}$ for $p=2$ and $p=1.5$. Note
that $p=2$ gives us the standard SGD. We solved the PDE (21) for these two
mirrors and show the results in Figure 2.

The top row corresponds to mirror potential $\|\cdot\|_{2}^{2}$, and the
bottom row corresponds to mirror potential $\|\cdot\|_{1.5}^{2}$. The first
three columns correspond to the distribution of the two-dimensional parameter
vector at times $t=0,2,4$, where $t=0$ represents the beginning of training
and $t=4$ the end of training. For both SMDs the two-dimensional weight vector
was initialized with a zero-mean Gaussian distribution. The evolution of the
distribution of the weights can be seen at times $t=2$ and $t=4$. In the case
of SGD, the weights appear to converge to three point masses that represent
the three lines that separate the centers of the two data classes. In the case
of $p=1.5$, the results are more interesting. There appears to be more than
three types of nonlinear units; in fact, there is mass along the $\theta_{1}$
and $\theta_{2}$ axes, signifying that some of fraction of the weights are
zero. This seems reasonable behavior since compared to $p=2$, the mirror
potential with $p=1.5$ will encourage sparsity—since it has implicit bias
towards minimizing the $\ell_{1.5}$ norm Azizan and Hassibi [2019], Gunasekar
et al. [2020a].

The fourth column represents the output of the network for various inputs
after training has ended ($t=4$). For $p=2$, we see that the two classes have
been separated by three lines passing through the origin. For $p=1.5$, the
structure of the output is more interesting. The classes are separated by more
than three line segments, and not all go through the origin. Clearly, the two
mirrors have solved the classification problem in two different ways.

Despite this, the end performance of both trained models is nearly identical:
the expected risk at the completion of training is 0.4136 for the $\ell_{2}$
norm and 0.4118 for the $\ell_{1.5}$ norm. Interestingly, $p=1.5$ performs
slightly better.

Figure 2: The results of solving the PDE (21). The top row corresponds to
mirror potential $\|\cdot\|_{2}^{2}$, yielding SGD, and the bottom row
corresponds to mirror potential $\|\cdot\|_{1.5}^{2}$. The first three columns
correspond to the distribution of the two-dimensional parameter vector at
times $t=0,2,4$, where $t=0$ represents the beginning of training and $t=4$
the end of training. The fourth column represents the output of the network
for various inputs after training has ended ($t=4$). The expected risk at the
completion of training is 0.4136 for the $\ell_{2}$ norm and 0.4118 for the
$\ell_{1.5}$ norm.

##  6 Outline of the Proof

Here we give a brief outline of the proof of Theorem 1. The techniques we use
are similar to those of Wang et al. [2017] and Sirignano and Spiliopoulos
[2020].
We start by introducing equivalent mirror domain counterparts of SMD iterates
at (8) and proceed with analyzing the exchangeable Markov process in the
mirror domain. For $\bm{\theta}\in\mathbb{R}^{d}$,
$\bm{\omega}:=\nabla\psi(\bm{\theta})$ denotes its mirror domain counterpart,
and $\pi:=\rho\circ\nabla\psi^{-1}$ denotes probability measure in the mirror
domain666For a probability measure $\rho$ on a set $\mathcal{S}$ and a
function $f:\mathcal{S}\rightarrow\mathcal{S}$, the push-forward of $\rho$
under $f$ is defined as $(\rho\circ f^{-1})(B):=\rho(f^{-1}(B))$ for all Borel
$B\subset\mathcal{S}$.. We rewrite SMD iterations with this new set of
notations as follows

$\begin{cases}\hat{\pi}_{n}^{k}=\frac{1}{n}\sum_{i=1}^{n}\delta_{\bm{\omega}_{i}^{k}}\leavevmode\nobreak\
,\\\
\bm{\omega}_{i}^{k+1}=\bm{\omega}_{i}^{k}+\frac{\tau}{n}\mathbf{G}(\bm{\omega}_{i}^{k},\hat{\pi}_{n}^{k},\mathbf{z}^{k+1})\leavevmode\nobreak\
,\end{cases}$ (35)

where
$\mathbf{G}(\bm{\omega},\pi,\mathbf{z}):=\mathbf{F}(\nabla\psi^{-1}(\bm{\omega}),\pi\circ\nabla\psi,\mathbf{z})$.
We continue by characterizing the asymptotic limit of mirror domain empirical
distributions $\\{\hat{\pi}_{n}^{k}\\}_{k=0}^{K}$, and consequently convert
the results to the original parameter domain by a change of variables. To
characterize the evolution of the empirical measure of parameters, we analyze
the change in integral of a test function. Let
$f\in\mathcal{C}_{b}^{2}(\mathbb{R}^{d})$, and consider the the following
consecutive difference,

$\displaystyle\langle$
$\displaystyle\hat{\pi}_{n}^{k+1},f\rangle-\langle\hat{\pi}_{n}^{k},f\rangle=\frac{1}{n}\sum_{i=1}^{n}\left[f(\bm{\omega}_{i}^{k+1})-f(\bm{\omega}_{i}^{k})\right],$
(36) $\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\left[\frac{\tau}{n}\nabla
f(\bm{\omega}_{i}^{k})^{T}\mathbf{G}(\bm{\omega}_{i}^{k},\hat{\bm{\pi}}_{n}^{k},\mathbf{z}^{k+1})+r_{i}^{k}\right],$

where the last equality is simply the first-order Taylor expansion and
$r_{i}^{k}$ is the remainder term which is given as,

$r_{i}^{k}=\frac{1}{2}(\bm{\omega}_{i}^{k+1}-\bm{\omega}_{i}^{k})^{T}\nabla^{2}f(\bar{\bm{\omega}}_{i}^{k})(\bm{\omega}_{i}^{k+1}-\bm{\omega}_{i}^{k})$
(37)

for a point $\bar{\bm{\omega}}_{i}^{k}$ that lies on the line segment
$[\bm{\omega}_{i}^{k},\bm{\omega}_{i}^{k+1}]$. Proceeding on-wards, we
introduce the following three terms,

$\displaystyle d^{k}$ $\displaystyle:=\frac{\tau}{n^{2}}\sum_{i=1}^{n}\nabla
f(\bm{\omega}_{i}^{k})^{T}\mathbb{E}_{\mathbf{z}\sim\mu(d\mathbf{z})}\left[\mathbf{G}(\bm{\omega}_{i}^{k},\hat{\pi}_{n}^{k},\mathbf{z})\right]$
(38) $\displaystyle m^{k}$
$\displaystyle:=\frac{\tau}{n^{2}}\sum_{i=1}^{n}\nabla
f(\bm{\omega}_{i}^{k})^{T}\mathbf{G}(\bm{\omega}_{i}^{k},\hat{\pi}_{n}^{k},\mathbf{z}^{k+1})-d^{k}$
(39) $\displaystyle r^{k}$ $\displaystyle:=\frac{1}{n}\sum_{i=1}^{n}r_{i}^{k}$
(40)

Rewriting (36) in terms of these variables gives,

$\langle\hat{\pi}_{n}^{k+1},f\rangle-\langle\hat{\pi}_{n}^{k},f\rangle=d^{k}+m^{k}+r^{k}.$
(41)

Notice that using the definition of the drift function (15),
$\mathbb{E}_{\mathbf{z}\sim\mu(d\mathbf{z})}\left[\mathbf{G}(\bm{\omega},\hat{\pi},\mathbf{z})\right]$
can be written as,

$\mathbb{E}_{\mathbf{z}\sim\mu(d\mathbf{z})}\left[\mathbf{G}(\bm{\omega},\hat{\pi},\mathbf{z})\right]=\mathbf{v}(\nabla\psi^{-1}(\bm{\omega}),\,\hat{\pi}\circ\nabla\psi).$
(42)

We use the shorthand notation
$\mathbf{u}(\bm{\omega},\,\pi):=\mathbf{v}(\nabla\psi^{-1}(\bm{\omega}),\,\pi\circ\nabla\psi)$
to denote the mirror domain counterpart of $\mathbf{v}$. Consider a
continuous-time process $\\{\bar{\pi}_{n}(t)\,:\,t\in[0,T]\\}$ similar to our
earlier definition in (12). Taking the sum of (41) for $0\leq
k<\lfloor{nt/\tau}\rfloor$,

$\displaystyle\langle\bar{\pi}_{n}(t),f\rangle$
$\displaystyle-\langle\hat{\pi}^{0},f\rangle=\sum_{k=0}^{\lfloor{nt/\tau}\rfloor-1}\left(\langle\hat{\pi}_{n}^{k+1},f\rangle-\langle\hat{\pi}_{n}^{k},f\rangle\right)$
(43)
$\displaystyle=\sum_{k=0}^{\lfloor{nt/\tau}\rfloor-1}d^{k}+\sum_{k=0}^{\lfloor{nt/\tau}\rfloor-1}m^{k}+\sum_{k=0}^{\lfloor{nt/\tau}\rfloor-1}r^{k}.$

The first sum on the RHS can be written as,

$\displaystyle\sum_{k=0}^{\lfloor{nt/\tau}\rfloor-1}d^{k}$
$\displaystyle=\frac{\tau}{n}\sum_{l=0}^{\lfloor{nt/\tau}\rfloor-1}\langle\hat{\pi}_{n}^{k},\nabla
f(\cdot)^{T}\mathbf{u}(\cdot,\,\hat{\pi}_{n}^{k})\rangle,$ (44)
$\displaystyle=\int_{0}^{[\,t\,]_{n}}\langle\bar{\pi}_{n}(s),\,\nabla
f(\cdot)^{T}\mathbf{u}(\cdot,\bar{\pi}_{n}(s))\rangle ds.$

where $[\,t\,]_{n}:=\lfloor{\frac{nt}{\tau}}\rfloor\frac{\tau}{n}$, and the
last equality derived from the definition of the process $\bar{\pi}_{n}(t)$.
By replacing (44) in (43) we have,

$\displaystyle\langle\bar{\pi}_{n}(t),\,f\rangle=\langle\hat{\pi}^{0},\,f\rangle$
$\displaystyle+\int_{0}^{[\,t\,]_{n}}\langle\bar{\pi}_{n}(s),\,\nabla
f(\cdot)^{T}\mathbf{u}(\cdot,\bar{\pi}_{n}(s))\rangle ds$ (45)
$\displaystyle+\sum_{l=0}^{\lfloor{nt/\tau}\rfloor-1}m^{l}+\sum_{l=0}^{\lfloor{nt/\tau}\rfloor-1}r^{l}.$

The proof concludes by showing the two sums on the RHS of (45) converge to
zero as $n\rightarrow\infty$ and $\frac{\tau}{n}=O(1/n)$. The convergence of
the first sum comes from the fact that $\\{m^{l}\\}_{\ell>=0}$ can be viewed
as deviations of a martingale form its mean, and therefore can be bounded by
martingale concentration inequalities (e.g. see Lemma A.1 in Mei et al.
[2018]). The last sum also approaches to zero due to our assumptions on the
boundedness of the second derivatives of $\sigma$ and $\ell$. Thus, we have
the following:

$\langle\pi(t),\,f\rangle=\langle\pi^{0},\,f\rangle+\int_{0}^{t}\langle\pi(s),\,\nabla
f(\cdot)^{T}\mathbf{u}(\cdot\pi(s))\rangle ds.$ (46)

The result obtained above can be rewritten in terms of primal domain
variables, namely $\mathbf{v}$, and $\rho(t)$ by performing a change of
variables as defined earlier in this section. Then, Hessian of the mirror
function comes as the Jacobian of this change of variables process.

##  7 Conclusion

In this paper we explored the performance of mirror descent iterates through a
mean-field lens for ensemble average models. We characterized the time
evolution of the distribution of the parameters through a nonlinear PDE and
gave a Riemannian interpretation. It would be interesting to further study the
effect of the mirror potential on the performance of SMD-trained models for
different learning problems via analytic and numerical study of this PDE. It
would also be interesting to generalize these results to other network models
and learning algorithms.

## References

  * Ambrosio et al. [2008] Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré.  _Gradient flows in metric spaces and in the space of probability measures_.  Lectures in mathematics ETH Zürich. Birkhäuser, Basel, 2. ed edition, 2008.  ISBN 978-3-7643-8721-1 978-3-7643-8722-8.  OCLC: 254181287.
  * Azizan and Hassibi [2019] Navid Azizan and Babak Hassibi.  A Stochastic Interpretation of Stochastic Mirror Descent: Risk-Sensitive Optimality.  In _2019 IEEE 58th Conference on Decision and Control (CDC)_ , pages 3960–3965, Nice, France, December 2019. IEEE.  ISBN 978-1-72811-398-2.  doi: 10.1109/CDC40024.2019.9030229.  URL https://ieeexplore.ieee.org/document/9030229/.
  * Azizan et al. [2019] Navid Azizan, Sahin Lale, and Babak Hassibi.  Stochastic Mirror Descent on Overparameterized Nonlinear Models: Convergence, Implicit Regularization, and Generalization.  _arXiv:1906.03830 [cs, math, stat]_ , June 2019.  URL http://arxiv.org/abs/1906.03830.  arXiv: 1906.03830.
  * Billingsley [1999] Patrick Billingsley.  _Convergence of probability measures_.  Wiley series in probability and statistics. Probability and statistics section. Wiley, New York, 2nd ed edition, 1999.  ISBN 978-0-471-19745-4.
  * Borovykh et al. [2020] Anastasia Borovykh, Nikolas Kantas, Panos Parpas, and Grigorios A. Pavliotis.  On stochastic mirror descent with interacting particles: convergence properties and variance reduction.  _arXiv:2007.07704 [math, stat]_ , November 2020.  arXiv: 2007.07704.
  * Bubeck [2015] Sébastien Bubeck.  Convex Optimization: Algorithms and Complexity.  _arXiv:1405.4980 [cs, math, stat]_ , November 2015.  arXiv: 1405.4980.
  * Chizat and Bach [2018] Lenaic Chizat and Francis Bach.  On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport.  _arXiv:1805.09545 [cs, math, stat]_ , October 2018.  URL http://arxiv.org/abs/1805.09545.  arXiv: 1805.09545.
  * Gunasekar et al. [2020a] Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro.  Characterizing Implicit Bias in Terms of Optimization Geometry.  _arXiv:1802.08246 [cs, stat]_ , June 2020a.  URL http://arxiv.org/abs/1802.08246.  arXiv: 1802.08246.
  * Gunasekar et al. [2020b] Suriya Gunasekar, Blake Woodworth, and Nathan Srebro.  Mirrorless Mirror Descent: A More Natural Discretization of Riemannian Gradient Flow.  _arXiv:2004.01025 [cs, math, stat]_ , April 2020b.  URL http://arxiv.org/abs/2004.01025.  arXiv: 2004.01025.
  * Kallenberg [2010] Olav Kallenberg.  _Foundations of Modern Probability_.  Probability and its Applications. Springer New York, New York, NY, 2\. ed edition, 2010.  ISBN 978-1-4419-2949-5.  OCLC: 865092618.
  * Mei et al. [2018] Song Mei, Andrea Montanari, and Phan-Minh Nguyen.  A Mean Field View of the Landscape of Two-Layers Neural Networks.  _arXiv:1804.06561 [cond-mat, stat]_ , August 2018.  URL http://arxiv.org/abs/1804.06561.  arXiv: 1804.06561.
  * Nemirovskii and Yudin [1983] A. S. Nemirovskii and D. B. Yudin.  _Problem complexity and method efficiency in optimization_.  Wiley-Interscience series in discrete mathematics. Wiley, Chichester ; New York, 1983.  ISBN 978-0-471-10345-5.
  * Raginsky and Bouvrie [2012] Maxim Raginsky and Jake Bouvrie.  Continuous-time stochastic Mirror Descent on a network: Variance reduction, consensus, convergence.  In _2012 IEEE 51st IEEE Conference on Decision and Control (CDC)_ , pages 6793–6800, Maui, HI, USA, December 2012. IEEE.  ISBN 978-1-4673-2066-5 978-1-4673-2065-8 978-1-4673-2063-4 978-1-4673-2064-1.  doi: 10.1109/CDC.2012.6426639.  URL http://ieeexplore.ieee.org/document/6426639/.
  * Raskutti and Mukherjee [2014] Garvesh Raskutti and Sayan Mukherjee.  The Information Geometry of Mirror Descent.  _arXiv:1310.7780 [cs, stat]_ , April 2014.  URL http://arxiv.org/abs/1310.7780.  arXiv: 1310.7780.
  * Rotskoff and Vanden-Eijnden [2019] Grant M. Rotskoff and Eric Vanden-Eijnden.  Trainability and accuracy of neural networks: An interacting particle system approach, 2019.
  * Shima and Yagi [1997] Hirohiko Shima and Katsumi Yagi.  Geometry of hessian manifolds.  _Differential Geometry and its Applications_ , 7(3):277–290, 1997.  ISSN 0926-2245.  doi: https://doi.org/10.1016/S0926-2245(96)00057-5.  URL https://www.sciencedirect.com/science/article/pii/S0926224596000575.
  * Sirignano and Spiliopoulos [2020] Justin Sirignano and Konstantinos Spiliopoulos.  Mean Field Analysis of Neural Networks: A Law of Large Numbers.  _SIAM Journal on Applied Mathematics_ , 80(2):725–752, January 2020.  ISSN 0036-1399, 1095-712X.  doi: 10.1137/18M1192184.  URL https://epubs.siam.org/doi/10.1137/18M1192184.
  * Sznitman [1991] Alain-Sol Sznitman.  Topics in propagation of chaos.  In Paul-Louis Hennequin, editor, _Ecole d’Eté de Probabilités de Saint-Flour XIX — 1989_ , volume 1464, pages 165–251. Springer Berlin Heidelberg, Berlin, Heidelberg, 1991.  ISBN 978-3-540-53841-7 978-3-540-46319-1.  doi: 10.1007/BFb0085169.  URL http://link.springer.com/10.1007/BFb0085169.  Series Title: Lecture Notes in Mathematics.
  * Wang et al. [2017] Chuang Wang, Jonathan Mattingly, and Yue M Lu.  Scaling limit: Exact and tractable analysis of online learning algorithms with applications to regularized regression and pca.  _arXiv preprint arXiv:1712.04332_ , 2017.