# Global Convergence of SGD For Logistic Loss on Two Layer Neural Nets

Pulkit Gopalani<EMAIL_ADDRESS>
Computer Science and Engineering
University of Michigan, Ann Arbor Samyak Jha<EMAIL_ADDRESS>
Department of Mathematics
Indian Institute of Technology, Bombay Anirbit Mukherjee
<EMAIL_ADDRESS>
Department of Computer Science
The University of Manchester

###### Abstract

In this note, we demonstrate a first-of-its-kind provable convergence of SGD
to the global minima of appropriately regularized logistic empirical risk of
depth $2$ nets – for arbitrary data and with any number of gates with
adequately smooth and bounded activations like sigmoid and tanh. We also prove
an exponentially fast convergence rate for continuous time SGD that also
applies to smooth unbounded activations like SoftPlus. Our key idea is to show
the existence of Frobenius norm regularized logistic loss functions on
constant-sized neural nets which are “Villani functions” and thus be able to
build on recent progress with analyzing SGD on such objectives.

##  1 Introduction

Modern developments in artificial intelligence have been significantly been
driven by the rise of deep-learning. The highly innovative engineers who have
ushered in this A.I. revolution have developed a vast array of heuristics that
work to get the neural net to perform “human like” tasks. Most such successes,
can mathematically be seen to be solving the function optimization/“risk
minimization” question,
$\min_{n\in{\mathcal{N}}}\mathbb{E}_{{\bm{z}}\in{\mathcal{D}}}[{\rm\ell}(n,{\bm{z}})]$
where members of ${\mathcal{N}}$ are continuous functions representable by
neural nets and $\ell:{\mathcal{N}}\times{\rm
Support}({\mathcal{D}})\rightarrow[0,\infty)$ is called a “loss function” and
the algorithm only has sample access to the distribution ${\mathcal{D}}$. The
successful neural experiments can be seen as suggesting that there are many
available choices of $\ell,~{}{\mathcal{N}}~{}\&~{}{\mathcal{D}}$ for which
highly accurate solutions to this seemingly extremely difficult question can
be easily found. This is a profound mathematical mystery of our times

This work is about developing our understanding of some of the most ubiquitous
methods of training nets. In particular, we shed light on how regularization
can aid the analysis and help prove convergence to global minima for
stochastic gradient methods for neural nets in hitherto unexplored and
realistic parameter regimes.

In the last few years, there has been a surge in the literature on provable
training of various kinds of neural nets in certain regimes of their widths or
depths, or for very specifically structured data, like noisily realizable
labels. Motivated by the abundance of experimental studies it has often been
surmised that Stochastic Gradient Descent (SGD) on neural net losses – with
proper initialization and learning rate – converges to a low–complexity
solution, one that generalizes – when it exists Zhang et al. (2018).

But, to the best of our knowledge a convergence result for any stochastic
training algorithm applied to the logistic loss for even depth $2$ nets (one
layer of activations with any kind of non–linearity), without either an
assumption on the width or the data, has remained elusive so far. We recall
that this is the most common way to train classifiers facing binary class
labelled data.

In this work, we not only take a step towards addressing the above question in
the theory of neural networks but we also do so while keeping to a standard
algorithm, the Stochastic Gradient Descent (SGD). In light of the above, our
key message can be summarily stated as follows,

######  Theorem 1.1 (Informal Statement of the Main Results, Theorem 3.2).

If the initial weights are sampled from an appropriate class of distributions,
then for nets with a single layer of sigmoid or tanh gates – for arbitrary
data and size of the net – SGD on appropriately regularized logistic loss,
while using constant steps of size ${\mathcal{O}}(\epsilon)$, will converge in
${\mathcal{O}}(\frac{1}{\epsilon})$ steps to weights at which the expected
regularized loss would be $\epsilon$–close to its global minimum.

We note that the threshold amount of regularization needed in the above would
be independent of the width of the nets. Further, this threshold would be
shown to scale s.t it can either naturally turn out to be proportionately
small if the norms of the training data are small or can be made arbitrarily
small by choosing outer layer weights to be small. Our above result is made
possible by the crucial observation informally stated in the following lemma -
which infact holds for more general nets than what is encompassed by the above
theorem,

######  Lemma 1.2.

It is possible to add a constant amount of Frobenius norm regularization on
the weights, to the logistic loss on depth-$2$ nets with activations like
SoftPlus, sigmoid and tanh gates s.t with no assumptions on the data or the
size of the net, the regularized loss would be a Villani function.

Since our result stated above does not require any assumptions on the data, or
the neural net width, we posit that this significantly improves on previous
work in this direction. To the best of our knowledge, similar convergence
guarantees in the existing literature either require some minimum neural net
width – growing quickly w.r.t. inverse accuracy and the training set size (NTK
regime Chizat et al. (2018); Du et al. (2018b)), infinite width (Mean Field
regime Chizat & Bach (2018); Chizat (2022); Mei et al. (2018)) or other
assumptions on the data when the width is parametric (e.g. realizable data, Ge
et al. (2019); Zhou et al. (2021)).

In contrast to all these, we show that with appropriate $\ell_{2}$
regularization, SGD on logistic loss on 2–layer sigmoid / tanh nets converges
to the global infimum of the loss. Our critical observation towards this proof
is that the above standard losses on 2–layer nets – for a broad class of
activation functions — are a “Villani function”. Our proof get completed by
leveraging the relevant results in Shi et al. (2020).

#### Organization

In Section 2 we shall give a literature review of existing proofs about
guaranteed training of neural nets. In Section 3 we present our primary
results – Theorem 3.2 which show the global convergence of SGD on regularized
logistic loss with $\pm 1$ labelled data and and for gates like sigmoid and
tanh. Additionally, in Theorem 3.5 we also point out that for our
architecture, if using the SoftPlus activation, we can show that the
underlying SDE converges in expectation to the global minimizer in linear
time. In Section 4, we give a brief overview of the methods in Shi et al.
(2020), leading up to the proof of Theorem 3.2 in Section 5. In Section 3.2 we
discuss some experimental demonstrations, which show that there exist nets
trained on the loss function considered such that they have high binary
classification accuracy near the critical value of the regularizer considered
for the proof. We end in Section 6 with a discussion of various open questions
that our work motivates. Appendix A onwards one can find the calculations
needed in the main theorems’ proofs.

##  2 Related Work

Firstly, we note that in recent times major advances have been made about
understanding the statistical properties of doing binary classification by
neural nets. In Zhou & Huo (2023), the authors consider $\\{\pm 1\\}$ labelled
data distributed as a Gaussian Mixture Model and the labels satisfying a
Tsybakov-type noise conditions with the noise exponent being $q$. The authors
obtain that with probability $1-\delta$ of sampling $n$ data, the difference
between the population risk of the empirical risk minimizer and the Bayes’
optimal risk is upperbounded by $C_{q,d}\log(\frac{2}{\delta})(\log
n)^{4}{(\frac{1}{n})^{\frac{q+1}{q+2}}}$, where $C_{q,d}$ is some constant
depending on $q$ and data dimension $d$. To appreciate this, we note that
earlier in Shen et al. (2022), similar bounds for CNNs with logistic loss were
obtained. However, unlike the result in Shen et al. (2022), the excess risk
bound obtained in Zhou & Huo (2023) for the hinge loss doesn’t blow up with
respect to the increasing smoothness of the minimizer of the risk over all
measurable functions.

In the setting of finite–width neural nets trained on logistic loss for binary
classification, Chatterji et al. (2021), had shown that if one has (a) small
initial loss (${\rm poly}\left(\frac{1}{n}\right)$, where $n$ is the number of
training data) and (b) ‘smoothly approximately ReLU activation function’, then
the loss converges at a rate of $O\left(\frac{1}{t}\right)$ over $t$ steps of
gradient descent. But to ensure the smallness of the initial loss, this result
needs to assume a large width which scales polylogarithmically with inverse of
the confidence parameter. In that limited sense this can be seen to be
belonging to the larger framework of proofs at asymptotically wide nets which
we review as follows.

#### Review of the NTK Approach To Provable Neural Training :

One of the most popular parameter zones for theory of provable training of
nets has been the so–called “NTK” (Neural Tangent Kernel) regime – where the
width is a high degree polynomial in the training set size and inverse
accuracy (a somewhat unrealistic regime) and the net’s last layer weights are
scaled inversely with width as the width goes to infinity, Du et al. (2018a);
Su & Yang (2019); Allen-Zhu et al. (2019b; a); Du & Lee (2018); Arora et al.
(2019b); Li et al. (2019); Arora et al. (2019a); Chizat et al. (2018); Du et
al. (2018b). The core insight in this line of work can be summarized as
follows: for large enough width, SGD with certain initializations converges to
a function that fits the data perfectly, with minimum norm in the RKHS defined
by the neural tangent kernel – which gets specified entirely by the
initialization (which is such that the initial output is of order one). A key
feature of this regime is that the net’s matrices do not travel outside a
constant radius ball around the starting point – a property that is often not
true for realistic neural net training scenarios.

In particular, for the case of depth $2$ nets – with similarly smooth gates as
we focus on – in Song et al. (2021) global convergence of gradient descent was
shown using number of gates scaling sub-quadratically in the number of data -
which, to the best of our knowledge, is the smallest known width requirement
for such a convergence in a classification setup. On the other hand, for the
special case of training depth $2$ nets with $\mathop{\mathrm{ReLU}}$ gates on
cross-entropy loss for doing binary classification, in Ji & Telgarsky (2020)
it was shown that one needs to blow up the width only poly-logarithmically
with target accuracy to get global convergence for SGD.

#### Review of the Mean-Field Approach To Provable Neural Net Training :

In a separate direction of attempts towards provable training of neural nets,
works like Chizat & Bach (2018) showed that a Wasserstein gradient flow limit
of the dynamics of discrete time algorithms on shallow nets, converges to a
global optimizer – if the convergence of the flow is assumed. We note that
such an assumption is very non-trivial because the dynamics being analyzed in
this setup is in infinite dimensions – a space of probability measures on the
parameters of the net. Similar kind of non–asymptotic convergence results in
this so–called ‘mean–field regime’ were also obtained in Mei et al. (2018);
Fang et al. (2021); Chizat & Bach (2018); Nguyen & Pham (2020); Sirignano &
Spiliopoulos (2022); Ren & Wang (2022). The key idea in the mean–field regime
is to replace the original problem of neural training which is a non-convex
optimization problem in finite dimensions by a convex optimization problem in
infinite dimensions – that of probability measures over the space of weights.
The mean–field analysis necessarily require the probability measures (whose
dynamics is being studied) to be absolutely–continuous and thus de facto it
only applies to nets in the limit of them being infinitely wide.

We note that the results in the NTK regime hold without regularization while
in many cases the mean–field results need it, Mei et al. (2018); Chizat
(2022); Tzen & Raginsky (2020).

In the next subsection we shall give a brief overview of some of the attempts
that have been made to get provable deep-learning at parametric width.

#### Need And Attempts To Go Beyond Large Width Limits of Nets

The essential proximity of the NTK regime to kernel methods and it being less
powerful than finite nets has been established from multiple points of view,
Allen-Zhu & Li (2019); Wei et al. (2019).

In He & Su (2020), the authors had given a a very visibly poignant way to see
that the NTK limit is not an accurate representation of a lot of the usual
deep-learning scenarios. Their idea was to define a notion of “local
elasticity” – when doing a SGD update on the weights using a data say
${\bm{x}}$, it measures the fractional change in the value of the net at a
point ${\bm{x}}^{\prime}$ as compared to ${\bm{x}}$. It’s easy to see that
this is a constant function for linear regression - as is what happens at the
NTK limit (Theorem 2.1 Lee et al. (2019)). But it has been shown in Dan et al.
(2021) that this local-elasticity function indeed has non-trivial time-
dynamics (particularly during the early stages of training) when a moderately
large neural net is trained on logistic loss.

Specific to depth-2 nets – as we consider here – there is a stream of
literature where analytical methods have been honed to this setup to get good
convergence results without width restrictions - while making other structural
assumptions about the data or the net. Janzamin et al. (2015) was one of the
earliest breakthroughs in this direction and for the restricted setting of
realizable labels they could provably get arbitrarily close to the global
minima. For non-realizable labels they could achieve the same while assuming a
large width but in all cases they needed access to the score function of the
data distribution which is a computationally hard quantity to know. In a more
recent development, Awasthi et al. (2021) have improved over the above to
include $\mathop{\mathrm{ReLU}}$ gates while being restricted to the setup of
realizable data and its marginal distribution being Gaussian.

One of the first proofs of gradient based algorithms doing neural training for
depth$-2$ nets appeared in Zhong et al. (2017). In Ge et al. (2019)
convergence was proven for training depth-$2$ $\mathop{\mathrm{ReLU}}$ nets
for data being sampled from a symmetric distribution and the training labels
being generated using a ‘ground truth’ neural net of the same architecture as
being trained – the so-called “Teacher–Student” setup. For similar
distributional setups, some of the current authors had in Karmakar et al.
(2020) identified classes of depth–$2$ $\mathop{\mathrm{ReLU}}$ nets where
they could prove linear-time convergence of training – and they also gave
guarantees in the presence of a label poisoning attack. The authors in Zhou et
al. (2021) consider another Teacher–Student setup of training depth $2$ nets
with absolute value activations. In this work, authors can get convergence in
poly$(d,\frac{1}{\epsilon})$ time, in a very restricted setup of assuming
Gaussian data, initial loss being small enough, and the teacher neurons being
norm bounded and ‘well–separated’ (in angle magnitude). Cheridito et al.
(2022) get width independent convergence bounds for Gradient Descent (GD) with
ReLU nets, however at the significant cost of having the restrictions of being
only an asymptotic guarantee and assuming an affine target function and
one–dimensional input data. While being restricted to the Gaussian data and
the realizable setting for the labels, an intriguing result in Chen et al.
(2021) showed that fully poly-time learning of arbitrary depth 2 ReLU nets is
possible if one is in the “black-box query model”.

#### Related Work on Provable Training of Neural Networks Using Regularization

Using a regularizer is quite common in deep-learning practice and in recent
times a number of works have appeared which have established some of these
benefits rigorously. In particular, Wei et al. (2019) show that for a specific
classification task (noisy–XOR) definable in any dimension $d$, no NTK based 2
layer neural net can succeed in learning the distribution with low
generalization error in $o(d^{2})$ samples, while in $O(d)$ samples one can
train the neural net using Frobenius/$\ell_{2}-$norm regularization. Nakkiran
et al. (2021) show that for a specific optimal value of the $\ell_{2}$\-
regularizer the double descent phenomenon can be avoided for linear nets - and
that similar tuning is possible even for real world nets.

In the seminal work Raginsky et al. (2017), it was pointed out that one can
add a regularization to a gradient Lipschitz loss and make it satisfy the
dissipativity condition so that Stochastic Gradient Langevin Dynamics (SGLD)
provably converges to its global minima. But SGLD is seldom used in practice,
and to the best of our knowledge it remains unclear if the observation in
Raginsky et al. (2017) can be used to infer the same about SGD. Also it
remains open if there exists neural net losses which satisfy all the
assumptions needed in the above result. We note that the convergence time in
Raginsky et al. (2017) for SGLD is
$\mathcal{O}\left(\frac{1}{\epsilon^{5}}\right)$ using an
$\mathcal{O}\left(\epsilon^{4}\right)$ learning rate, while in our Theorem 3.2
SGD converges in expectation to the global infimum of the regularized neural
loss in time, $\mathcal{O}\left(\frac{1}{\epsilon}\right)$ using a
$\mathcal{O}\left(\epsilon\right)$ step-length.

In summary, to the best of our knowledge, it has remained an unresolved
challenge to show convergence of SGD for logistic loss on any neural
architecture with a constant number of gates while not constraining the
distribution of the data to a specific functional form. In this work, we
exploit the use of some regularization to be able to resolve this optimization
puzzle in our key result, Theorem 3.2 – and in our experiments we show that
the regularization needed may not harm downstream classification performance.
Thus we take a step towards bridging this important lacuna in the existing
theory of stochastic optimization for neural nets in general.

##  3 Setup and Main Results

We start with defining the neural net architecture, the loss function and the
algorithm for which we will prove our convergence results.

######  Definition 1 (Constant Step-Size SGD On Depth-2 Nets).

Let, $\sigma:\mathbb{R}\rightarrow\mathbb{R}$ (applied elementwise for vector
valued inputs) be atleast once differentiable activation function.
Corresponding to it, consider the width $p$, depth $2$ neural nets with fixed
outer layer weights ${\bm{a}}\in\mathbb{R}^{p}$ and trainable weights
${\bm{W}}\in\mathbb{R}^{p\times d}$ as,

$\mathbb{R}^{d}\ni{\bm{x}}\mapsto
f({\bm{x}};\,{\bm{a}},{\bm{W}})={\bm{a}}^{\top}\sigma({\bm{W}}{\bm{x}})\in\mathbb{R}$

Then, corresponding to a given set of $n$ binary training data
$({\bm{x}}_{i},y_{i})\in\mathbb{R}^{d}\times\\{+1,-1\\}$, with
$\norm{{\bm{x}}_{i}}_{2}\leq B_{x},~{}i=1,\ldots,n$ define the individual data
logistic losses
$\tilde{L}_{i}({\bm{W}})\coloneqq\log(1+e^{-y_{i}f_{i}\left(W\right)})$. Then
for any $\lambda>0$ let the regularized logistic empirical risk be,

$\tilde{L}({\bm{W}})\coloneqq\frac{1}{n}\sum_{i=1}^{n}\tilde{L}_{i}({\bm{W}})+\frac{\lambda}{2}\norm{{\bm{W}}}_{F}^{2}$
(1)

Correspondingly, we consider SGD with step-size $s>0$ as,

${\bm{W}}^{k+1}={\bm{W}}^{k}-\frac{s}{b}\sum_{i\in{\mathcal{B}}_{k}}\nabla\tilde{L}_{i}({\bm{W}}^{k})-s\lambda{\bm{W}}^{k}$

where ${\mathcal{B}}_{k}$ is a randomly sampled mini-batch of size $b$.

######  Definition 2 (Properties of the Activation $\sigma$).

Let the $\sigma$ used in Definition 1 be bounded s.t.
$\absolutevalue{\sigma(x)}\leq B_{\sigma}$, $C^{\infty}$, $L-$Lipschitz and
$L_{\sigma}^{\prime}-$smooth. Further assume that $\exists$ a constant vector
${\bm{c}}$ and positive constants $B_{\sigma},M_{D}$ and $M_{D}^{\prime}$ s.t
$\sigma(\mathbf{0})={\bm{c}}$ and $\forall
x\in\mathbb{R},\absolutevalue{\sigma^{\prime}(x)}\leq
M_{D},\absolutevalue{\sigma^{\prime\prime}(x)}\leq M_{D}^{\prime}$ .

In terms of the above constants we can now quantify the smoothness of the
empirical loss as follows,

######  Lemma 3.1 (for Classification with Logistic Loss).

In the setup of binary classification as contained in Definition 1, and the
definition $M_{D}$ and $L$ as given in Definition 2 above, there exists a
constant $\lambda_{c}=\frac{M_{D}LB_{x}^{2}\norm{a}^{2}_{2}}{2}$ s.t
$\forall\lambda>\lambda_{c}$ and $s>0$ the Gibbs measure
$\sim\textup{exp}\left(-\frac{2\tilde{L}}{s}\right)$ satisfies a Poincaré-type
inequality with the corresponding constant $\lambda_{s}$. Moreover, if the
activation satisfies the conditions of Definition 2 then
$\exists\textup{gLip}\left(\tilde{L}\right)>0$ such that the empirical loss is
$\textup{gLip}\left(\tilde{L}\right)$-smooth and we can bound the smoothness
coefficient of the empirical loss as,

$\textup{gLip}(\tilde{L})\leq\sqrt{p}\left(\frac{\sqrt{p}\norm{{\bm{a}}}_{2}M_{D}^{2}B_{x}}{4}+\left(\frac{2+\norm{c}_{2}+\norm{{\bm{a}}}_{2}B_{\sigma}}{4}\right)M^{\prime}_{D}B_{x}p+\lambda\right)$
(2)

The precise form of the Poincaré-type inequalities used above is detailed in
Theorem 4.1.

######  Theorem 3.2 (Global Convergence of SGD on Sigmoid and Tanh Neural Nets
of 2 Layers for Any Width and Any Data - for Binary Classification With
Logistic Loss).

We continue in the setup of logistic loss from Definitions 1 and 2 and Lemma
3.1. For any $s>0$, define the probability measure
$\mu_{s}\coloneqq\frac{1}{Z_{s}}\exp\left({-\frac{2\tilde{L}({\bm{W}})}{s}}\right)$,
$Z_{s}$ being the normalization factor. Then, $\forall~{}T>0,$ and desired
accuracy, $\epsilon>0$, $\exists$ constants $A(\tilde{L})$, $B(T,\tilde{L})$
and $C(s,\tilde{L})$ s.t if the above SGD is executed at a constant step-size

$s=s^{*}\coloneqq\min\left(\frac{1}{{\rm
gLip}(\tilde{L})},\frac{\epsilon}{2\cdot(A(\tilde{L})+B(T,\tilde{L}))}\right)$

with the weights ${\bm{W}}^{0}$ initialized from a distribution with p.d.f
$\rho_{\rm initial}\in L^{2}(\frac{1}{\mu_{s^{*}}})$ and $\norm{\rho_{\rm
initial}-\mu_{s^{*}}}_{\mu_{s^{*}}^{-1}}\leq\frac{\epsilon}{2\cdot
C(s^{*},\tilde{L})}\cdot e^{\lambda_{s^{*}}\cdot T}$ – then, in expectation,
the regularized empirical risk of the net, $\tilde{L}$ would converge to its
global infimum, with the rate of convergence given as,

${\mathbb{E}}\tilde{L}({\bm{W}}^{\frac{T}{s^{*}}})-\inf_{{\bm{W}}}\tilde{L}({\bm{W}})\leq\epsilon.$

For clarity, in the following lemma, we also point out, that a guarantee on
the quality of the weights obtained by SGD on this regularized loss, can also
be given such that it is explicitly parametric in the number of steps taken as
well as the initialization.

######  Lemma 3.3 (Bounds on Error for Arbitrary Initialization.).

We continue in the setup of logistic loss from Definitions 1 and 2 and Lemma
3.1. For any $s>0$, define the probability measure
$\mu_{s}\coloneqq\frac{1}{Z_{s}}\exp\left({-\frac{2\tilde{L}({\bm{W}})}{s}}\right)$,
$Z_{s}$ being the normalization factor. Then, $\forall~{}T>0$ and desired
accuracy $\epsilon>0$, $\exists$ constants $A(\tilde{L})$, $B(T,\tilde{L})$
and $C(s,\tilde{L})$ s.t if the above SGD is executed at a constant step-size

$s=s^{*}(\epsilon,T)\coloneqq\min\left(\frac{1}{{\rm
gLip}(\tilde{L})},\epsilon\cdot\frac{1}{(A(\tilde{L})+B(T,\tilde{L}))}\right)$

with the weights ${\bm{W}}^{0}$ initialized from any distribution with p.d.f
$\rho_{\rm initial}\in L^{2}(\frac{1}{\mu_{s^{*}}})$ and then, the error at
the end of having taken $k=\frac{T}{s^{*}}$ SGD steps can bounded as,

${\mathbb{E}}\tilde{L}({\bm{W}}^{k})-\min_{{\bm{W}}}\tilde{L}({\bm{W}})\leq\epsilon+C(s^{*},\tilde{L})\norm{\rho_{\rm
initial}-\mu_{s^{*}}}_{\mu_{s^{*}}^{-1}}e^{-s^{*}\lambda_{s^{*}}\cdot k}$

The proof of Theorem 3.2 is given in Section 5 and the proof of Lemma 3.1 can
be read off from the calculations done as a part of the proof of Theorem 3.2.
The proof for Lemma 3.3 follows from the calculations done for proof of
Theorem 3.2 and hence it’s omitted.

We make a few quick remarks about the nature of the above guarantees, Firstly,
we note that the “time horizon” $T$ above is a free parameter - which in turn
parameterizes the choice of the step-size and the initial weight distribution.
Choosing a larger $T$ makes the constraints on the initial weight distribution
weaker at the cost of making the step-size smaller and the required number of
SGD steps larger. But for any value of $T$, the above theorem guarantees that
SGD, initialized from weights sampled from a certain class of distributions,
converges in expectation to the global minima of the regularized empirical
loss for our nets for any data and width, in time
${\mathcal{O}}(\frac{1}{\epsilon})$ using a learning rate of
${\mathcal{O}}(\epsilon)$.

Secondly, we note that the phenomenon of a lower bound on the regularization
parameter being needed for certain nice learning theoretic property to emerge
has been seen in kernel settings too, Yang et al. (2017).

Also, to put into context the emergence of a critical value of the regularizer
for nets as in the above theorem, we recall the standard result, that there
exists an optimal value of the $\ell_{2}-$regularizer at which the excess risk
of the similarly penalized linear regression becomes dimension free
(Proposition 3.8, Bach (2022)). However, we recall that the quantities
required for computing this “optimal” regularizer are not knowable while
training and hence it is not practically implementable. Thus, we see that for
binary classification, one can define a notion of an “optimal” regularizer and
it remains open to investigate if such a similar threshold of regularization
also exists for nets. Our above theorem can be seen as a step in that
direction.

Thirdly, we note that the lowerbounds on training time of neural nets proven
in works like Goel et al. (2020) do not apply here since these are proven for
SQ algorithms and SGD is not of this type.

Finally, note that the threshold values of regularization computed above,
$\lambda_{c}$, do not explicitly depend on the training data or the neural
architecture, consistent with observations in Anthony & Bartlett (2009); Zhang
et al. (2021). It depends on the activation and scales with the norm of the
input data and the outer layer of weights.

For intuition, suppose in the binary classification setting we set the outer
layer weights s.t we have $\norm{{\bm{a}}}_{2}\cdot B_{x}=1$. This leads to
$\lambda_{c}=\frac{M_{D}L}{2}.$ For the sigmoid activation,
$\sigma_{\beta}(x)$, by calculations as above we would get, $\lambda_{c}$ in
this case (say $\lambda_{c}^{si,\beta}$) to be $=\frac{\beta^{2}}{32}.$ Since
$\beta=1$ is the most widely used setting for the above sigmoid activation,
this results in,

$\lambda_{c}^{si,1}\approx 0.03125$ (3)

###  3.1 Global Convergence of Continuous Time SGD on Nets with SoftPlus Gates

In, Shi et al. (2020) the authors had established that, if the loss
$\tilde{L}$ is gradient Lipschitz, then over any fixed time horizon $T>0$, as
$s\rightarrow 0$, the dynamics of the SGD in Definition 1 is arbitrarily well
approximated (in expectation) by the unique global solution that exists for
the Stochastic Differential Equation (SDE),

$\displaystyle\differential{{\mathbf{W}}_{s}(t)}=-\nabla\tilde{L}(\textrm{W}_{s}(t))\differential{t}+\sqrt{s}\differential{{\mathbf{B}}(t)}\quad\quad\text{(SGD
--SDE)}$ (4)

where ${\mathbf{B}}(t)$ is the standard Brownian motion. The SGD convergence
proven in the last section critically uses this mapping to a SDE. In Shi et
al. (2020) it was further pointed out that if we only want to get a non-
asymptotic convergence rate for the continuous time dynamics, the smoothness
of the loss function is not needed and only the Villani condition suffices. In
this short section we shall exploit this to show convergence of continuous
time SGD on $\tilde{L}$ with the activation function being the unbounded
‘SoftPlus’. Also, in contrast to the guarantee about SGD in the previous
subsection here we shall see that the SDE converges exponentially faster i.e
at a linear rate.

######  Definition 3 (SoftPlus activation).

For $\beta>0,$ $x\in\mathbb{R},$ define the SoftPlus activation function as

${\rm SoftPlus}_{\beta}(x)=\frac{1}{\beta}\log_{e}{\left(1+\exp(\beta
x)\right)}$

######  Remark.

Note that $\lim_{\beta\rightarrow\infty}{\rm SoftPlus}_{\beta}(x)={\rm
ReLU}(x).$ Also note that for $f(x)=$ SoftPlus${}_{\beta}(x)$,
$f^{\prime}(x)=\sigma_{\beta}(x)$ (sigmoid function as defined above) and
hence $\absolutevalue{f^{\prime}(x)}\leq M_{D}$ for $M_{D}=1$ and $f(x)$ is
$L-$Lipschitz for $L=1$.

Recall the following fact that was proven as a part of Lemma 3.1,

######  Lemma 3.4.

There exists a constant
$\lambda_{c}:=\frac{M_{D}LB_{x}^{2}\norm{{\bm{a}}}_{2}^{2}}{2}$ s.t
$\forall~{}\lambda>\lambda_{c}~{}\&~{}s>0$, the Gibbs’ measure
$\mu_{s}\coloneqq\frac{1}{Z_{s}}\exp\left({-\frac{2\tilde{L}({\bm{W}})}{s}}\right)$,
$Z_{s}$ being the normalization factor, satisfies a Poincaré-type inequality
with the corresponding constant $\lambda_{s}$.

######  Theorem 3.5 (Continuous Time SGD Converges to Global Minima of
SoftPlus Nets in Linear Time).

We consider the SDE as given in Equation 4 on a Frobenius norm regularized
logistic empirical loss on depth$-2$ neural nets as specified in Equation 1,
while using $\sigma(x)={\rm SoftPlus}_{\beta}(x)$ for $\beta>0$, the
regularization threshold being s.t
$\lambda>\lambda_{c}=\frac{M_{D}LB_{x}^{2}\norm{{\bm{a}}}_{2}^{2}}{2}$ and
with the weights ${\bm{W}}_{0}$ being initialized from a distribution with
p.d.f $\rho_{\rm initial}\in L^{2}(\frac{1}{\mu_{s}})$.

Then, for any $S>0$, $\exists~{}G(S,\tilde{L})$ and $C(s,\tilde{L})$, an
increasing function of $s$, s.t for any step size
$0<s\leq\min\left\\{\frac{\epsilon}{2G(S,\tilde{L})},S\right\\}$ and for
$t\geq\frac{1}{\lambda_{s}}\log{\left(\frac{2\,C(s,\tilde{L})\norm{\rho_{\rm
initial}-\mu_{s}}_{\mu_{s}^{-1}}}{\epsilon}\right)}$ we have that,

$\mathbb{E}\,\tilde{L}({\bm{W}}(t))-\min_{{\bm{W}}}\tilde{L}({\bm{W}})\leq\epsilon.$

###### Proof.

The SoftPlus function is Lipschitz, hence using the same analysis as in
(Appendix 5), we can claim that for $\lambda>\lambda_{c}$ the loss function in
Definition 1 with SoftPlus activations is a Villani function (and hence
confining, by definition).

Then, from Proposition $3.1$ of Shi et al. (2020) it follows that,
$\exists~{}C(s,\tilde{L})$, an increasing function of $s$, that satisfies,

$\absolutevalue{\mathbb{E}\tilde{L}({\bm{W}}_{s}(t))-\mathbb{E}\tilde{L}({\bm{W}}_{s}(\infty))}\leq
C(s,\tilde{L})\norm{\rho_{\rm
initial}-\mu_{s}}_{\mu_{s}^{-1}}\,e^{-\lambda_{s}t}.$

From Proposition $3.2$ of Shi et al. (2020) it follows that, for any $S>0$,
for $s\in(0,S)$, $\exists~{}G(S,\tilde{L})$ that quantifies the excess risk at
the stationary point of the SDE as,

$\tilde{L}({\bm{W}}(\infty))-\min_{{\bm{W}}}\tilde{L}\leq G(S,\tilde{L})\,s$

Combining the above, the final result claimed follows as in Corollary $3.3$ in
Shi et al. (2020). ∎

###  3.2 An Experimental Demonstration of the Maintenance of Classification
Accuracy At Various Regularizations at Different Widths

For further illustration of the ramifications of the novel convergence
theorems shown above, in here we present some experimental studies of doing
binary classification by training depth $2$ sigmoid activated nets with the
regularized loss considered in the above convergence proofs. And we will be
using the normalizations that correspond to the theoretically needed threshold
value of the regularizer being $\lambda_{c}=0.03125$ (Equation 3).

We sampled the data from a clearly separable dataset with a margin – $n$ data
vectors in $d-$dimensions were sampled as a $n\times d$ normally distributed
matrix whereby after sampling each row vector was normalized to have unit
norm. The data whose last coordinate were $>0.2$ were assigned label $+1$ and
where the last coordinate was $<-0.2$ was assigned label $-1$ and the rest of
the data were discarded. In our experimental setting, we fixed to $d=10$. The
data were split into being $20\%$ test data and the rest used for training.

Then we simulate SGD based training on the above data for multiple neural nets
at various values of $\lambda$ in $[0,\lambda_{C}]$ and for neural net widths
$p$. The step-length in the experiments is constant across all widths and
lambda, across all settings.

The elements of the (trainable) weight Matrix ${\bm{W}}_{0}$ of dimension
$p\times d$ were initialized from a standard normal distribution, i.e
${\bm{W}}_{0}\sim{\bf\mathcal{N}}(0,{\bm{I}}_{p\times d})$. Likewise, the
elements of the (fixed) outer layer ${\bm{a}}$, of dimension $1\times p$ were
sampled as ${\bm{a}}\sim{\bf\mathcal{N}}(0,{\bm{I}}_{1\times p})$ and then
normalized.

For all experimental settings the neural networks were trained for $500$
epochs, and with a mini-batch size of $256$. At the end of training we
measured the test accuracy of classification - as the downstream metric of
measuring the goodness of training.

As shown in the experimental graph (Figure 1), we demonstrate examples of nets
trained on regularized logistic loss showing remarkable accuracy for
regularization $\lambda\leq\lambda_{c}$ – even though the model was not been
exposed to the $0-1$ criteria during training.

The code for the experiments can be found at this Colaboratory File.

Figure 1: Test Accuracy across various widths $p$ and regularizer $\lambda$

Thus we have demonstrated that the threshold amount of regularization that was
needed for the proof of convergence may not at all harm the downstream
performance metric of classification.

##  4 Overview of Shi et al. (2020)

In Section 5, we will give the proof for our main result (Theorem 3.2). As
relevant background for the proof, in here we shall give a brief overview of
the framework in Shi et al. (2020), which can be summarized as follows :
suppose one wants to minimize the function
$\tilde{L}({\bm{W}})\coloneqq\frac{1}{n}\sum_{i=1}^{n}\tilde{L}_{i}({\bm{W}})$,
where $i$ indexes the training data, ${\bm{W}}$ is in the parameter space (the
optimization space) of the loss function and $\tilde{L}_{i}$ is the loss
evaluated on the $i^{th}-$datapoint. On this objective, a constant step-size
mini-batch implementation of the Stochastic Gradient Descent (SGD) consists of
doing the following iterates,
${\bm{W}}_{k+1}={\bm{W}}_{k}-\frac{s}{b}\sum_{i}\nabla\tilde{L}_{i}({\bm{W}}_{k})$,
where the sum is over a mini-batch (a randomly sampled subset of the training
data) of size $b$ and $s$ is the fixed step-length. In, Shi et al. (2020) the
authors established that over any fixed time horizon $T>0$, as $s\rightarrow
0$, the dynamics of this SGD is arbitrarily well approximated (in expectation)
by the Stochastic Differential Equation (SDE),

$\displaystyle\differential{{\bm{W}}_{s}(t)}=-\nabla\tilde{L}({\bm{W}}_{s}(t))\differential{t}+\sqrt{s}\differential{{\bm{B}}(t)}\quad\quad\text{(SGD
--SDE)}$ (5)

where ${\bm{B}}(t)$ is the standard Brownian motion. We recall that the Markov
semigroup operator $P_{t}$ for a stochastic process $X_{t}$ and its
infinitesimal generator ${\mathcal{L}}$ are given as,
$P_{t}f(x)\coloneqq\mathbb{E}[f(X_{t})\mid X_{0}=x]$ and
${\mathcal{L}}f\coloneqq\lim_{t\,\downarrow\,0}\,\frac{P_{t}f-f}{t}$.

Invoking the Forward Kolmogorov equation ${\partial_{t}f}={\mathcal{L}}^{*}f$,
one obtains the following Fokker–Planck–Smoluchowski PDE governing the
evolution of the density of the SDE,

$\displaystyle\partialderivative{\rho_{s}}{t}$
$\displaystyle=\langle\nabla\rho_{s},\nabla\tilde{L}\rangle+\rho_{s}\Delta\tilde{L}+\frac{s}{2}\Delta\rho_{s}\quad\quad\text{(FPS)}$
(6)

Further, under appropriate conditions on $\tilde{L}$ the above implies that
the density $\rho_{s}(t)$ converges exponentially fast to the Gibbs’ measure
corresponding to the objective function i.e the distribution with p.d.f

$\mu_{s}=\frac{1}{Z_{s}}{\exp\left(-\frac{2\tilde{L}({\bm{W}})}{s}\right)}$

where $Z_{s}$ is the normalization factor. The sufficient conditions on
$\tilde{L}$ that were shown to be needed to achieve this “mixing” and to know
a rate for it, are that of $\tilde{L}$ being a “Villani Function” as defined
below,

######  Definition 4 (Villani Function (Villani (2009); Shi et al. (2020))).

A map $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ is called a Villani function if
it satisfies the following conditions,

  1. 1.

$f\in C^{\infty}$

  2. 2.

$\lim_{\norm{{\bm{x}}}\rightarrow\infty}f({\bm{x}})=+\infty$

  3. 3.

${\displaystyle\int_{\mathbb{R}^{d}}}\exp\left({-\frac{2f({\bm{x}})}{s}}\right)\differential{{\bm{x}}}<\infty~{}\forall
s>0$

  4. 4.

$\lim_{\norm{{\bm{x}}}\rightarrow\infty}\left(-\Delta
f({\bm{x}})+\frac{1}{s}\cdot\norm{\nabla
f({\bm{x}})}^{2}\right)=+\infty~{}\forall s>0$

Further, any $f$ that satisfies conditions 1 – 3 is said to be “confining”.

From Lemma $5.2$ Shi et al. (2020), the empirical or the population risk,
$\tilde{L}$, being confining is sufficient for the FPS PDE (equation 6) to
evolve the density of SGD–SDE (equation 4) to the said Gibbs’ measure.

But, to get non-asymptotic guarantees of convergence (Corollary $3.3$, Shi et
al. (2020)) – even for the SDE, we need a Poincaré–type inequality to be
satisfied (as defined below) by the aforementioned Gibbs’ measure $\mu_{s}$. A
sufficient condition for this Poincaré–type inequality to be satisfied is if a
confining loss function $\tilde{L}$ also satisfied the last condition in
definition 4 (and is consequently a Villani function).

######  Theorem 4.1 (Poincaré–type Inequality (Shi et al. (2020))).

Given a $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ which is a Villani Function
(Definition 4), for any given $s>0$, define a measure with the density,
$\mu_{s}({\bm{x}})=\frac{1}{Z_{s}}{\exp\left(-\frac{2f({\bm{x}})}{s}\right)}$,
where $Z_{s}$ is a normalization factor. Then this (normalized) Gibbs’ measure
$\mu_{s}$ satisfies a Poincare-type inequality i.e $\exists~{}\lambda_{s}>0$
(determined by $f$) s.t $\forall h\in C_{c}^{\infty}({\mathbb{R}}^{d})$ we
have,

${\rm
Var}_{\mu_{s}}[h]\leq\frac{s}{2\lambda_{s}}\cdot\mathbb{E}_{\mu_{s}}[\norm{\nabla
h}^{2}]$

The approach of Shi et al. (2020) has certain key interesting differences from
many other contemporary uses of SDEs to prove the convergence of discrete time
stochastic algorithms. Instead of focusing on the convergence of parameter
iterates ${\bm{W}}^{k}$, they instead look at the dynamics of the expected
error i.e $\mathbb{E}[\tilde{L}({\bm{W}}^{k})]$, for $\tilde{L}$ the empirical
or the population risk. This leads to a transparent argument for the
convergence of $\mathbb{E}[\tilde{L}({\bm{W}}^{k})]$ to
$\inf_{{\bm{W}}}\tilde{L}({\bm{W}})$, by leveraging standard results which
help one pass from convergence guarantees on the SDE to a convergence of the
SGD.

We note that Shi et al. (2020) achieve this conversion of guarantees from SDE
to SGD by additionally assuming gradient smoothness of $\tilde{L}$ – and we
would show that this assumption holds for the natural neural net loss
functions that we consider.

##  5 Proof of Theorem 3.2

###### Proof.

Note that $\tilde{L}$ being a confining function can be easily read off from
Definition 4. Further, as shown in Appendix A, the following inequalities
hold,

$\displaystyle\norm{\nabla_{{\bm{W}}}L\left({\bm{W}}\right)}^{2}$
$\displaystyle\geq\left(\lambda^{2}-\frac{\lambda\norm{{\bm{a}}}^{2}_{2}M_{D}B_{x}^{2}L}{2}\right)\norm{{\bm{W}}}^{2}_{F}-\lambda\norm{{\bm{W}}}_{F}\norm{{\bm{a}}}_{2}M_{D}B_{x}\left(1+\frac{\norm{{\bm{a}}}_{2}\norm{{\bm{c}}}_{2}}{2}\right)$
$\displaystyle\Delta_{{\bm{W}}{\bm{W}}}\tilde{L}$ $\displaystyle\leq
p\left[M_{d}^{2}B_{x}^{2}\norm{{\bm{a}}}_{2}^{2}+{\norm{{\bm{a}}}_{2}}\left[\left(B_{y}+\norm{{\bm{a}}}_{2}\left(\norm{{{\bm{c}}}}_{2}+LB_{x}\norm{{\bm{W}}}_{F}\right)\right)\left(M_{D}^{\prime}B_{x}^{2}\right)\right]+\lambda
d\right]$ (7)

Combining the above two inequalities we can conclude that, $\exists$ functions
$g_{1},g_{2},g_{3}$ such that,

$\displaystyle\frac{1}{s}\norm{\nabla_{{\bm{W}}}\tilde{L}}^{2}-\Delta_{{\bm{W}}{\bm{W}}}\tilde{L}$
$\displaystyle\geq
g_{1}(\lambda,s)\norm{{\bm{W}}}_{F}^{2}-g_{2}(\lambda,s)\norm{{\bm{W}}}_{F}+g_{3}(\lambda,s)$

where in particular,

$g_{1}(\lambda,s)=\lambda^{2}-2\lambda\cdot
M_{D}LB_{x}^{2}\norm{{\bm{a}}}_{2}^{2}.$

Hence we can conclude that for $\lambda>\lambda_{c}\coloneqq
2M_{D}LB_{x}^{2}\norm{{\bm{a}}}_{2}^{2},\forall s>0,$
$\frac{1}{s}\norm{\nabla_{{\bm{W}}}\tilde{L}}^{2}-\Delta_{{\bm{W}}{\bm{W}}}\tilde{L}$
diverges as $\norm{{\bm{W}}}\rightarrow+\infty$, since $g_{1}(\lambda,s)>0.$
The key aspect of the above analysis being that the bound on
$\Delta_{{\bm{W}}{\bm{W}}}$ does not depend on $\norm{{\bm{W}}}_{F}^{2}.$ Thus
we have, that the following limit holds,

$\displaystyle\lim_{\norm{{\bm{W}}}_{F}\rightarrow+\infty}\left(\frac{1}{s}\norm{\nabla_{{\bm{W}}}\tilde{L}}^{2}-\Delta_{{\bm{W}}{\bm{W}}}\tilde{L}\right)=+\infty$

for the range of $\lambda$ as given in the theorem, hence proving that
$\tilde{L}$ is a Villani function.

Towards getting an estimate of the step-length as given in the theorem
statement, we also show in Appendix B that the loss function $\tilde{L}$ is
gradient–Lipschitz with the smoothness coefficient being upperbounded as,

$\textup{gLip}(\tilde{L})\leq\sqrt{p}\left(\frac{\sqrt{p}\norm{{\bm{a}}}_{2}M_{D}^{2}B_{x}}{4}+\left(\frac{2+\norm{{\bm{c}}}_{2}+\norm{{\bm{a}}}_{2}B_{\sigma}}{4}\right)M^{\prime}_{D}B_{x}p+\lambda\right)$

Now we can invoke Theorem $3$ (Part 1), Shi et al. (2020) with appropriate
choices of $s$ and initialization to get the main result as given in Theorem
3.2. In Appendix C one can find a discussion of the computation of the
specific constants involved in the expression for the suggested step-length
$s^{*}$ and the class of initial weight distribution p.d.fs $\rho_{\rm
initial}$. ∎

##  6 Conclusion

To the best of our knowledge, in this work, we have shown the first proof of
convergence of SGD to the global minima of logistic loss on a neural net
whilst not making any assumptions about the data or the width of the net. Our
result relied on the convergence of discrete time algorithms like SGD to their
continuous time counterpart (SDEs) - a theme which has lately been an active
field of research.

We posit that trying to reproduce our Theorem 3.2 using a direct analysis of
the dynamics of SGD could be a fruitful venture leading to interesting
insights. Additionally, our results motivate a new direction of pursuit in
deep-learning theory, centered around understanding the nature of the Poincaré
constant of Gibbs’ measures induced by neural nets.

Our experiments further demonstrated that there exists neural nets and binary
class labelled data such that optimizing on our provably good smooth loss
functions also does highly accurate classification. This motivates a
fascinating direction of future pursuit about the phenomenon of classification
calibration or Bayes consistency i.e to be able to analytically identify cases
where convergence guarantees as given here could be extended to guarantees on
the classification error.

Lastly, given the new method of proving neural training by SGD that has been
opened up in this work, it naturally begets the question if these insights
could also resolve similar mysteries for more exotic loss functions and neural
architectures that are in use.

## References

  * Allen-Zhu & Li (2019) Zeyuan Allen-Zhu and Yuanzhi Li.  What can resnet learn efficiently, going beyond kernels?  In _Advances in Neural Information Processing Systems_ , pp. 9015–9025, 2019.
  * Allen-Zhu et al. (2019a) Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang.  Learning and generalization in overparameterized neural networks, going beyond two layers.  In _Advances in neural information processing systems_ , pp. 6155–6166, 2019a.
  * Allen-Zhu et al. (2019b) Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song.  A convergence theory for deep learning via over-parameterization.  In _International Conference on Machine Learning_ , pp. 242–252, 2019b.
  * Anthony & Bartlett (2009) Martin Anthony and Peter L Bartlett.  Neural network learning: Theoretical foundations, 2009.
  * Arora et al. (2019a) Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang.  Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks.  In _International Conference on Machine Learning_ , pp. 322–332, 2019a.
  * Arora et al. (2019b) Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang.  On exact computation with an infinitely wide neural net.  In _Advances in Neural Information Processing Systems_ , pp. 8139–8148, 2019b.
  * Awasthi et al. (2021) Pranjal Awasthi, Alex Tang, and Aravindan Vijayaraghavan.  Efficient algorithms for learning depth-2 neural networks with general relu activations.  In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), _Advances in Neural Information Processing Systems_ , volume 34, pp. 13485–13496. Curran Associates, Inc., 2021.  URL https://proceedings.neurips.cc/paper/2021/file/700fdb2ba62d4554dc268c65add4b16e-Paper.pdf.
  * Bach (2022) Francis Bach.  _Learning Theory from First Principles_.  2022\.
  * Chatterji et al. (2021) Niladri S Chatterji, Philip M Long, and Peter Bartlett.  When does gradient descent with logistic loss interpolate using deep networks with smoothed relu activations?  In _Conference on Learning Theory_ , pp. 927–1027. PMLR, 2021.
  * Chen et al. (2021) Sitan Chen, Adam R Klivans, and Raghu Meka.  Efficiently learning any one hidden layer relu network from queries, 2021\.  URL https://arxiv.org/abs/2111.04727.
  * Cheridito et al. (2022) Patrick Cheridito, Arnulf Jentzen, and Florian Rossmannek.  Gradient descent provably escapes saddle points in the training of shallow relu networks, 2022.  URL https://arxiv.org/abs/2208.02083.
  * Chizat (2022) Lénaïc Chizat.  Mean-field langevin dynamics : Exponential convergence and annealing.  _Transactions on Machine Learning Research_ , 2022.  URL https://openreview.net/forum?id=BDqzLH1gEm.
  * Chizat & Bach (2018) Lenaic Chizat and Francis Bach.  On the global convergence of gradient descent for over-parameterized models using optimal transport.  In _Advances in neural information processing systems_ , pp. 3036–3046, 2018.
  * Chizat et al. (2018) Lenaic Chizat, Edouard Oyallon, and Francis Bach.  On lazy training in differentiable programming.  2018\.  doi: 10.48550/ARXIV.1812.07956.  URL https://arxiv.org/abs/1812.07956.
  * Dan et al. (2021) Soham Dan, Phanideep Gampa, and Anirbit Mukherjee.  Investigating the locality of neural network training dynamics, 2021.  URL https://arxiv.org/abs/2111.01166.
  * Du & Lee (2018) Simon Du and Jason Lee.  On the power of over-parametrization in neural networks with quadratic activation.  In _International Conference on Machine Learning_ , pp. 1329–1338, 2018.
  * Du et al. (2018a) Simon S. Du, Jason D. Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai.  Gradient descent finds global minima of deep neural networks, 2018a.
  * Du et al. (2018b) Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh.  Gradient descent provably optimizes over-parameterized neural networks, 2018b.  URL https://arxiv.org/abs/1810.02054.
  * Evans (2010) L.C. Evans.  _Partial Differential Equations_.  Graduate studies in mathematics. American Mathematical Society, 2010.  ISBN 9780821849743.  URL https://books.google.co.in/books?id=Xnu0o_EJrCQC.
  * Fang et al. (2021) Cong Fang, Jason Lee, Pengkun Yang, and Tong Zhang.  Modeling from features: a mean-field framework for over-parameterized deep neural networks.  In Mikhail Belkin and Samory Kpotufe (eds.), _Proceedings of Thirty Fourth Conference on Learning Theory_ , volume 134 of _Proceedings of Machine Learning Research_ , pp. 1887–1936. PMLR, 15–19 Aug 2021.  URL https://proceedings.mlr.press/v134/fang21a.html.
  * Ge et al. (2019) Rong Ge, Rohith Kuditipudi, Zhize Li, and Xiang Wang.  Learning two-layer neural networks with symmetric inputs.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=H1xipsA5K7.
  * Goel et al. (2020) Surbhi Goel, Aravind Gollakota, Zhihan Jin, Sushrut Karmalkar, and Adam Klivans.  Superpolynomial lower bounds for learning one-layer neural networks using gradient descent.  In Hal Daumé III and Aarti Singh (eds.), _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 3587–3596. PMLR, 13–18 Jul 2020.  URL https://proceedings.mlr.press/v119/goel20a.html.
  * He & Su (2020) Hangfeng He and Weijie Su.  The local elasticity of neural networks.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=HJxMYANtPH.
  * Janzamin et al. (2015) Majid Janzamin, Hanie Sedghi, and Anima Anandkumar.  Beating the perils of non-convexity: Guaranteed training of neural networks using tensor methods, 2015.  URL https://arxiv.org/abs/1506.08473.
  * Ji & Telgarsky (2020) Ziwei Ji and Matus Telgarsky.  Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=HygegyrYwH.
  * Karmakar et al. (2020) Sayar Karmakar, Anirbit Mukherjee, and Theodore Papamarkou.  Depth-2 neural networks under a data-poisoning attack, 2020.  URL https://arxiv.org/abs/2005.01699.
  * Lee et al. (2019) Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington.  Wide neural networks of any depth evolve as linear models under gradient descent.  _Advances in neural information processing systems_ , 32, 2019.
  * Li et al. (2019) Zhiyuan Li, Ruosong Wang, Dingli Yu, Simon S Du, Wei Hu, Ruslan Salakhutdinov, and Sanjeev Arora.  Enhanced convolutional neural tangent kernels.  _arXiv preprint arXiv:1911.00809_ , 2019.
  * Mei et al. (2018) Song Mei, Andrea Montanari, and Phan-Minh Nguyen.  A mean field view of the landscape of two-layer neural networks.  _Proceedings of the National Academy of Sciences_ , 115(33):E7665–E7671, 2018.  doi: 10.1073/pnas.1806579115.  URL https://www.pnas.org/doi/abs/10.1073/pnas.1806579115.
  * Nakkiran et al. (2021) Preetum Nakkiran, Prayaag Venkat, Sham M. Kakade, and Tengyu Ma.  Optimal regularization can mitigate double descent.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=7R7fAoUygoa.
  * Nguyen & Pham (2020) Phan-Minh Nguyen and Huy Tuan Pham.  A rigorous framework for the mean field limit of multilayer neural networks.  _CoRR_ , abs/2001.11443, 2020.  URL https://arxiv.org/abs/2001.11443.
  * Raginsky et al. (2017) Maxim Raginsky, Alexander Rakhlin, and Matus Telgarsky.  Non-convex learning via stochastic gradient langevin dynamics: a nonasymptotic analysis.  In Satyen Kale and Ohad Shamir (eds.), _Proceedings of the 2017 Conference on Learning Theory_ , volume 65 of _Proceedings of Machine Learning Research_ , pp. 1674–1703. PMLR, 07–10 Jul 2017.  URL https://proceedings.mlr.press/v65/raginsky17a.html.
  * Ren & Wang (2022) Zhenjie Ren and Songbo Wang.  Entropic fictitious play for mean field optimization problem, 2022.  URL https://arxiv.org/abs/2202.05841.
  * Shen et al. (2022) Guohao Shen, Yuling Jiao, Yuanyuan Lin, and Jian Huang.  Approximation with cnns in sobolev space: with applications to classification.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), _Advances in Neural Information Processing Systems_ , volume 35, pp. 2876–2888. Curran Associates, Inc., 2022.  URL https://proceedings.neurips.cc/paper_files/paper/2022/file/136302ea7874e2ff96d517f9a8eb0a35-Paper-Conference.pdf.
  * Shi et al. (2020) Bin Shi, Weijie J. Su, and Michael I. Jordan.  On learning rates and schrödinger operators, 2020.  URL https://arxiv.org/abs/2004.06977.
  * Sirignano & Spiliopoulos (2022) Justin A. Sirignano and Konstantinos Spiliopoulos.  Mean field analysis of deep neural networks.  _Math. Oper. Res._ , 47(1):120–152, 2022.  doi: 10.1287/moor.2020.1118.  URL https://doi.org/10.1287/moor.2020.1118.
  * Song et al. (2021) Chaehwan Song, Ali Ramezani-Kebrya, Thomas Pethick, Armin Eftekhari, and Volkan Cevher.  Subquadratic overparameterization for shallow neural networks.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), _Advances in Neural Information Processing Systems_ , 2021.  URL https://openreview.net/forum?id=NhbFhfM960.
  * Su & Yang (2019) Lili Su and Pengkun Yang.  On learning over-parameterized neural networks: A functional approximation perspective.  In _Advances in Neural Information Processing Systems_ , pp. 2637–2646, 2019.
  * Tzen & Raginsky (2020) Belinda Tzen and Maxim Raginsky.  A mean-field theory of lazy training in two-layer neural nets: entropic regularization and controlled mckean-vlasov dynamics.  _CoRR_ , abs/2002.01987, 2020.  URL https://arxiv.org/abs/2002.01987.
  * Villani (2009) Cédric Villani.  Hypocoercivity.  2009\.
  * Wei et al. (2019) Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma.  Regularization matters: Generalization and optimization of neural nets vs their induced kernel.  In _Advances in Neural Information Processing Systems_ , pp. 9709–9721, 2019.
  * Yang et al. (2017) Yun Yang, Mert Pilanci, and Martin J. Wainwright.  Randomized sketches for kernels: Fast and optimal nonparametric regression.  _The Annals of Statistics_ , 45(3):991 – 1023, 2017.  doi: 10.1214/16-AOS1472.  URL https://doi.org/10.1214/16-AOS1472.
  * Zhang et al. (2018) Chiyuan Zhang, Qianli Liao, Alexander Rakhlin, Brando Miranda, Noah Golowich, and Tomaso Poggio.  Theory of deep learning iib: Optimization properties of sgd, 2018.  URL https://arxiv.org/abs/1801.02254.
  * Zhang et al. (2021) Lin Zhang, Ling Feng, Kan Chen, and Choy Heng Lai.  Edge of chaos as a guiding principle for modern neural network training, 2021.  URL https://arxiv.org/abs/2107.09437.
  * Zhong et al. (2017) Kai Zhong, Zhao Song, Prateek Jain, Peter L. Bartlett, and Inderjit S. Dhillon.  Recovery guarantees for one-hidden-layer neural networks, 2017.  URL https://arxiv.org/abs/1706.03175.
  * Zhou et al. (2021) Mo Zhou, Rong Ge, and Chi Jin.  A local convergence theory for mildly over-parameterized two-layer neural network.  In Mikhail Belkin and Samory Kpotufe (eds.), _Proceedings of Thirty Fourth Conference on Learning Theory_ , volume 134 of _Proceedings of Machine Learning Research_ , pp. 4577–4632. PMLR, 15–19 Aug 2021.  URL https://proceedings.mlr.press/v134/zhou21b.html.
  * Zhou & Huo (2023) Tian-Yi Zhou and Xiaoming Huo.  Classification of data generated by gaussian mixture models using deep relu networks.  _arXiv preprint arXiv:2308.08030v1_ , 2023.

##  Appendix A Towards Establishing the Villani condition for the Empirical
Logistic Loss

We start with the following observation,

######  Lemma A.1.

Letting ${\bm{w}}_{j}$ denote the $j^{\rm th}$ row of ${\bm{W}}$ we have,

$\nabla_{{\bm{w}}_{j}}f_{i}\left({\bm{W}}\right)=a_{j}\sigma^{\prime}\left({\bm{w}}_{j}^{\top}\bm{x}_{i}\right){\bm{x}}_{i}$

In the following, $\norm{{\bm{W}}}$ for a matrix ${\bm{W}}$ denotes its
spectral or operator norm. We recall that the regularized logistic loss for
training the given neural net on data
${\mathbb{D}}=\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{n}$ is,

$\tilde{L}({\bm{W}})\coloneqq\frac{1}{n}\sum_{i=1}^{n}\tilde{L}_{i}({\bm{W}})+R_{\lambda}({{\bm{W}}})\coloneqq\frac{1}{n}\sum_{i=1}^{n}\left\\{\frac{1}{2}\left(y_{i}-f({\bm{x}}_{i};{\bm{a}},{\bm{W}})\right)^{2}\right\\}+\frac{\lambda}{2}\norm{{\bm{W}}}_{F}^{2}$

Explicitly stated, the SGD iterates with step-length $s>0$ that we analyze on
the above loss are,

$\displaystyle{\bm{W}}^{k+1}$
$\displaystyle={\bm{W}}^{k}-s\left[\frac{1}{b}\sum_{i\in{\mathcal{B}}_{k}}\nabla\tilde{L}_{i}({\bm{W}}^{k})+\nabla
R_{\lambda}({{\bm{W}}^{k}})\right]$
$\displaystyle={\bm{W}}^{k}-\frac{s}{b}\sum_{i\in{\mathcal{B}}_{k}}\Big{(}-\left(y_{i}-f({\bm{x}}_{i};{\bm{a}},{\bm{W}}^{k})\right)\cdot\nabla_{{\bm{W}}^{k}}f({\bm{x}}_{i};{\bm{a}},{\bm{W}}^{k})\Big{)}-s\lambda{\bm{W}}^{k}$
$\displaystyle=(1-s\lambda){\bm{W}}^{k}+\frac{s}{b}\sum_{i\in{\mathcal{B}}_{k}}\left[(y_{i}-f({\bm{x}}_{i};{\bm{a}},{\bm{W}}^{k}))\cdot\nabla_{{\bm{W}}^{k}}f({\bm{x}}_{i};{\bm{a}},{\bm{W}}^{k})\right]$

We also note the following observation,

######  Lemma A.2.

$\frac{1}{1+e^{y_{i}f_{i}\left({\bm{W}}\right)}}\leq\frac{1}{2}+\frac{\absolutevalue{f_{i}({\bm{W}})}}{4}$
(8)

###### Proof.

$\frac{1}{1+e^{y_{i}f_{i}\left({\bm{W}}\right)}}\leq\frac{1}{1+e^{-\absolutevalue{f_{i}\left({\bm{W}}\right)}}}=\frac{e^{\absolutevalue{f_{i}\left({\bm{W}}\right)}}}{1+e^{\absolutevalue{f_{i}\left({\bm{W}}\right)}}}$

Consider the function $g(z)=\frac{e^{z}}{1+e^{z}},\ z\in[0,\infty)$

$g(0)=\frac{1}{2},\ \
g^{\prime}(z)=\frac{e^{z}}{(1+e^{z})^{2}}\leq\frac{1}{4}$

Hence, we have

$g(z)-g(0)=\int_{0}^{z}g^{\prime}(z)\leq\int_{0}^{z}\frac{1}{4}=\frac{z}{4}$

Therefore,

$g(z)\leq\frac{1}{2}+\frac{z}{4}$ (9)

∎

### Lower-bounding the Norm of the Gradient of the Empirical Logistical Loss

Recall that we are using the following empirical loss function,

$L\left({\bm{W}}\right)=\frac{1}{n}\sum_{i=1}^{n}\ell\left(y_{i}f_{i}\left({\bm{W}}\right)\right)+\frac{\lambda}{2}\norm{{\bm{W}}}^{2}_{2}$
(10)

where $\ell\left(z\right)=\log\left(1+e^{-z}\right)$ is the logistic loss
function. It follows that,

$\nabla_{{\bm{w}}_{j}}L\left({\bm{W}}\right)=\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}\right)\right)y_{i}\nabla_{{\bm{w}}_{j}}f\left({\bm{W}}\right)+\lambda{\bm{w}}_{j}$
(11)

Therefore,

$\nabla_{{\bm{w}}_{j}}L\left({\bm{W}}\right)=\frac{1}{n}\sum_{i=1}^{n}\frac{-y_{i}}{1+e^{y_{i}f_{i}({\bm{W}})}}\nabla_{{\bm{w}}_{j}}f\left({\bm{W}}\right)+\lambda{\bm{w}}_{j}$

And that implies,

$\begin{split}\norm{\nabla_{{\bm{w}}_{j}}L\left({\bm{W}}\right)}_{2}^{2}&=\lambda^{2}\norm{{\bm{w}}_{j}}_{2}^{2}-2\left<\lambda{\bm{w}}_{j},\
\frac{1}{n}\sum_{i=1}^{n}\frac{y_{i}}{1+e^{y_{i}f_{i}({\bm{W}})}}\nabla_{{\bm{w}}_{j}}f\left({\bm{W}}\right)\right>+\norm{\frac{1}{n}\sum_{i=1}^{n}\frac{-y_{i}}{1+e^{y_{i}f_{i}({\bm{W}})}}\nabla_{{\bm{w}}_{j}}f\left({\bm{W}}\right)}^{2}\\\
&\geq\lambda^{2}\norm{{\bm{w}}_{j}}_{2}^{2}-2\left<\lambda{\bm{w}}_{j},\frac{1}{n}\sum_{i=1}^{n}\frac{y_{i}}{1+e^{y_{i}f_{i}({\bm{W}})}}\nabla_{{\bm{w}}_{j}}f\left({\bm{W}}\right)\right>\\\
&\geq\lambda^{2}\norm{{\bm{w}}_{j}}^{2}_{2}-2\lambda\norm{{\bm{w}}_{j}}_{2}\left(\frac{1}{n}\sum_{i=1}^{n}\frac{\norm{\nabla_{{\bm{w}}_{j}}f_{i}\left({\bm{W}}\right)}}{1+e^{y_{i}f_{i}({\bm{W}})}}\right)\\\
&\geq\lambda^{2}\norm{{\bm{w}}_{j}}^{2}_{2}-2\lambda\norm{{\bm{w}}_{j}}_{2}\left(\frac{1}{n}\sum_{i=1}^{n}\frac{\absolutevalue{a_{j}}\absolutevalue{\sigma^{\prime}\left({\bm{w}}_{j}^{\top}\bm{x}_{i}\right)}\norm{\bm{x_{i}}}_{2}}{1+e^{y_{i}f_{i}({\bm{W}})}}\right)\end{split}$
(12)

In the last line above we have invoked Lemma A.1. Now invoking Lemma A.2 and
recalling the definition of $M_{D}$ and $B_{x}$ in the above, we can claim
that,

$\begin{split}\norm{\nabla_{{\bm{w}}_{j}}L\left({\bm{W}}\right)}_{2}^{2}&\geq\lambda^{2}\norm{{\bm{w}}_{j}}^{2}_{2}-2\lambda\norm{{\bm{w}}_{j}}_{2}\absolutevalue{a_{j}}M_{D}\norm{{\bm{x}}_{i}}\left(\frac{1}{n}\sum_{i=1}^{n}\frac{1}{2}+\frac{\absolutevalue{f_{i}\left({\bm{W}}\right)}}{4}\right)\\\
&\geq\lambda^{2}\norm{{\bm{w}}_{j}}^{2}_{2}-2\lambda\norm{{\bm{w}}_{j}}_{2}\absolutevalue{a_{j}}M_{D}B_{x}\left(\frac{1}{2}+\frac{\norm{{\bm{a}}}_{2}(LB_{x}\norm{{\bm{W}}}_{2}+\norm{{\bm{c}}}_{2})}{4}\right)\\\
\end{split}$

Summing the above over all $j$ and using Cauchy-Schwartz inequality,
$\sum_{j=1}^{p}\norm{{\bm{w}}_{j}}_{2}\cdot\absolutevalue{a_{j}}\leq\norm{{\bm{W}}}_{F}\cdot\norm{{\bm{a}}}$
and $\norm{{\bm{W}}}_{2}\leq\norm{{\bm{W}}}_{F}$, we get,

$\begin{split}\norm{\nabla_{{\bm{W}}}L\left({\bm{W}}\right)}^{2}\geq\left(\lambda^{2}-\frac{\lambda\norm{{\bm{a}}}^{2}_{2}M_{D}B_{x}^{2}L}{2}\right)\norm{{\bm{W}}}^{2}_{F}-\lambda\norm{{\bm{W}}}_{F}\norm{{\bm{a}}}_{2}M_{D}B_{x}\left(1+\frac{\norm{{\bm{a}}}_{2}\norm{{\bm{c}}}_{2}}{2}\right)\end{split}$
(13)

### Analyzing the Laplacian of the Empirical Logistic Loss

We begin with observing that,

$\begin{split}\absolutevalue{\nabla_{{\bm{w}}_{j}}\cdot\left(\nabla_{{\bm{w}}_{j}}L\left({\bm{W}}\right)\right)}&=\absolutevalue{\nabla_{{\bm{w}}_{j}}\cdot\left(\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}\right)\right)y_{i}\nabla_{{\bm{w}}_{j}}f\left({\bm{W}}\right)\right)+\lambda
d}\\\
&\leq\absolutevalue{\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime\prime}\left(f_{i}\left({\bm{W}}\right)\right)\norm{\nabla_{{\bm{w}}_{j}}f\left({\bm{W}}\right)}_{2}^{2}}+\absolutevalue{\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}\right)\right)y_{i}\Delta_{{\bm{w}}_{j}}f\left({\bm{W}}\right)}+\lambda
d\end{split}$

In above we have defined,
$\Delta_{{\bm{w}}_{j}}f=\nabla_{{\bm{w}}_{j}}\cdot(\nabla_{{\bm{w}}_{j}}f)$ We
recall that,
$\ell^{\prime\prime}(z)=\frac{e^{z}}{(1+e^{z})^{2}}\leq\frac{1}{4}$, as show
in Lemma A.2. Further by invoking Lemma A.1, the definition of $M_{D}$ and
$B_{x}$ we have,

$\begin{split}\absolutevalue{\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime\prime}\left(y_{i}f_{i}\left({\bm{W}}\right)\right)\norm{\nabla_{{\bm{w}}_{j}}f\left({\bm{W}}\right)}_{2}^{2}}&\leq\absolutevalue{\frac{1}{n}\sum_{i=1}^{n}\frac{1}{4}\norm{\nabla_{{\bm{w}}_{j}}f\left({\bm{W}}\right)}_{2}^{2}}\\\
&\leq\absolutevalue{\frac{1}{n}\sum_{i=1}^{n}\frac{1}{4}\absolutevalue{a_{j}}^{2}(\sigma^{\prime}({\bm{w}}_{j}^{\top}x_{i}))^{2}\norm{{\bm{x}}_{i}}^{2}}\\\
&\leq\frac{M_{D}^{2}B_{x}^{2}\norm{{\bm{a}}}_{2}^{2}}{4}\end{split}$ (14)

Recalling the definition of $M_{D}^{\prime}$ we have,

$\begin{split}\absolutevalue{\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}\right)\right)y_{i}\Delta_{{\bm{w}}_{j}}f\left({\bm{W}}\right)}&\leq\absolutevalue{\frac{1}{n}\sum_{i=1}^{n}\left(\frac{1}{2}+\frac{\absolutevalue{f_{i}\left({\bm{W}}\right)}}{4}\right)y_{i}\Delta_{{\bm{w}}_{j}}f\left({\bm{W}}\right)}\\\
&\leq\left(\frac{2+\norm{c}_{2}}{4}+\frac{\norm{{\bm{a}}}_{2}LB_{x}\norm{{\bm{W}}}_{F}}{4}\right)B_{x}^{2}M_{D}^{\prime}\norm{{\bm{a}}}_{2}\end{split}$
(15)

And hence,

$\absolutevalue{\nabla_{{\bm{w}}_{j}}\cdot\left(\nabla_{{\bm{w}}_{j}}L\left({\bm{W}}\right)\right)}\leq\left(\frac{2+\norm{c}_{2}}{4}+\frac{\norm{{\bm{a}}}_{2}LB_{x}\norm{{\bm{W}}}_{F}}{4}\right)B_{x}^{2}M_{D}^{\prime}\norm{{\bm{a}}}_{2}+\frac{M_{D}^{2}B_{x}^{2}\norm{{\bm{a}}}_{2}^{2}}{4}+\lambda
d$ (16)

Summing over all j, we get,

$\Delta_{{\bm{W}}}L\left({\bm{W}}\right)\leq\sum_{j=1}^{p}\absolutevalue{\nabla_{{\bm{w}}_{j}}\cdot\left(\nabla_{{\bm{w}}_{j}}L\left({\bm{W}}\right)\right)}\leq
p\left[\left(\frac{2+\norm{c}_{2}}{4}+\frac{\norm{{\bm{a}}}_{2}LB_{x}\norm{{\bm{W}}}_{F}}{4}\right)B_{x}^{2}M_{D}^{\prime}\norm{{\bm{a}}}_{2}+\frac{M_{D}^{2}B_{x}^{2}\norm{{\bm{a}}}_{2}^{2}}{4}+\lambda
d\right]$ (17)

##  Appendix B Bounding the Gradient Lipschitzness Coefficient of the
Empirical Logistic Loss

Towards the upcoming computation, we define the following function,

$g_{j}\left({\bm{W}}\right)\coloneqq\nabla_{{\bm{w}}_{j}}L=\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}\right)\right)y_{i}\nabla_{{\bm{w}}_{j}}f_{i}\left({\bm{W}}\right)+\lambda{\bm{w}}_{j}$
(18)

Consequently, corresponding to two values of the weight matrix ${\bm{W}}_{1}$
and ${\bm{W}}_{2}$, we obtain that,

$\begin{split}&\norm{g_{j}\left({\bm{W}}_{2}\right)-g_{j}\left({\bm{W}}_{1}\right)}_{2}=\\\
&\norm{\frac{1}{n}\sum_{i=1}^{n}y_{i}\left(\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}_{2}\right)\right)\nabla_{{\bm{w}}_{2,j}}f_{i}\left({\bm{W}}_{2}\right)-\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}_{1}\right)\right)\nabla_{{\bm{w}}_{1,j}}f_{i}\left({\bm{W}}_{1}\right)\right)+\lambda{\bm{w}}_{2,j}-\lambda{\bm{w}}_{1,j}}_{2}\\\
&\leq\norm{\frac{1}{n}\sum_{i=1}^{n}y_{i}\left(\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}_{2}\right)\right)\nabla_{{\bm{w}}_{2,j}}f_{i}\left({\bm{W}}_{2}\right)-\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}_{1}\right)\right)\nabla_{{\bm{w}}_{1,j}}f_{i}\left({\bm{W}}_{1}\right)\right)}_{2}+\norm{\lambda{\bm{w}}_{2,j}-\lambda{\bm{w}}_{1,j}}_{2}\\\
&\leq\frac{B_{x}\norm{{\bm{a}}}_{2}}{n}\sum_{i=1}^{n}\absolutevalue{\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}_{2}\right)\right)\sigma^{\prime}({\bm{w}}^{\top}_{2,j}{\bm{x}}_{i})-\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}_{1}\right)\right)\sigma^{\prime}({\bm{w}}^{\top}_{1,j}{\bm{x}}_{i})}+\norm{\lambda{\bm{w}}_{2,j}-\lambda{\bm{w}}_{1,j}}_{2}\\\
\end{split}$ (19)

In the last line above we have invoked Lemma A.1 and the fact that
$y_{i}\in\\{1,-1\\}$ for $i=1,\ldots,n$. Hence, the problem simplifies to
finding the Lipschitz constant of
$\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}\right)\right)\sigma^{\prime}({\bm{w}}^{\top}_{j}{\bm{x}}_{i})$
Define $h_{k}$ as follows:

${\bm{h}}_{k}({\bm{W}}):=\nabla_{{\bm{w}}_{k}}\left(\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}\right)\right)\sigma^{\prime}({\bm{w}}^{\top}_{j}{\bm{x}}_{i})\right)$
(20)
$\begin{split}\norm{{\bm{h}}_{k}({\bm{W}})}_{2}&=\norm{\nabla_{{\bm{w}}_{k}}\left(\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}\right)\right)\sigma^{\prime}({\bm{w}}^{\top}_{j}{\bm{x}}_{i})\right)}_{2}\\\
&=\norm{a_{k}\ell^{\prime\prime}\left(y_{i}f_{i}\left({\bm{W}}\right)\right)\sigma^{\prime}({\bm{w}}^{\top}_{j}{\bm{x}}_{i})\sigma^{\prime}({\bm{w}}^{\top}_{k}{\bm{x}}_{i}){\bm{x}}_{i}+\bm{1}_{k=j}\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}\right)\right)\sigma^{\prime\prime}({\bm{w}}^{\top}_{j}{\bm{x}}_{i}){\bm{x}}_{i}}_{2}\\\
&\leq\norm{a_{k}\ell^{\prime\prime}\left(y_{i}f_{i}\left({\bm{W}}\right)\right)\sigma^{\prime}({\bm{w}}^{\top}_{j}{\bm{x}}_{i})\sigma^{\prime}({\bm{w}}^{\top}_{k}{\bm{x}}_{i}){\bm{x}}_{i}}+\norm{\bm{1}_{k=j}\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}\right)\right)\sigma^{\prime\prime}({\bm{w}}^{\top}_{j}{\bm{x}}_{i}){\bm{x}}_{i}}_{2}\\\
&\leq\frac{\norm{{\bm{a}}}_{2}M_{D}^{2}B_{x}}{4}+\norm{\bm{1}_{k=j}\left(\frac{1}{2}+\frac{\absolutevalue{f_{i}({\bm{W}})}}{4}\right)\sigma^{\prime\prime}({\bm{w}}^{\top}_{j}{\bm{x}}_{i}){\bm{x}}_{i}}_{2}\\\
&\leq\frac{\norm{{\bm{a}}}_{2}M_{D}^{2}B_{x}}{4}+\left(\frac{2+\norm{c}_{2}}{4}+\frac{\norm{{\bm{a}}}_{2}B_{\sigma}}{4}\right)M^{\prime}_{D}B_{x}\sqrt{p}=L_{\rm
prod}\end{split}$

where in the second term the $\sqrt{p}$ factor comes in from using Cauchy-
Schwarz inequality. We concatenate these functions along the indices
$k=1,2,\ldots,p$, to get

${\bm{h}}({\bm{W}})\coloneqq\nabla_{{\bm{W}}}\left[f({\bm{x}}_{i};\,{\bm{a}},{\bm{W}})\sigma^{\prime}({\bm{w}}_{j}^{\top}{\bm{x}}_{i})\right]=\left[{{\bm{h}}}_{1}({\bm{W}}),{{\bm{h}}}_{2}({\bm{W}}),\ldots,{{\bm{h}}}_{p}({\bm{W}})\right]$

And hence,

$\norm{{\bm{h}}({\bm{W}})}_{2}\leq\sqrt{p}L_{\rm prod}$

Therefore, the Lipschitz constant of
$\ell^{\prime}\left(y_{i}f_{i}\left({\bm{W}}\right)\right)\sigma^{\prime}({\bm{w}}^{\top}_{j}{\bm{x}}_{i})$
is,

$\frac{\norm{{\bm{a}}}_{2}M_{D}^{2}B_{x}\sqrt{p}}{4}+\left(\frac{2+\norm{c}_{2}}{4}+\frac{\norm{{\bm{a}}}_{2}B_{\sigma}}{4}\right)M^{\prime}_{D}B_{x}p$

Hence the Lipschitz constant of $g_{j}\left({\bm{W}}\right)$ is

$\frac{\norm{{\bm{a}}}^{2}_{2}M_{D}^{2}B_{x}^{2}\sqrt{p}}{4}+\left(\frac{2+\norm{c}_{2}}{4}+\frac{\norm{{\bm{a}}}_{2}B_{\sigma}}{4}\right)M^{\prime}_{D}B_{x}^{2}\norm{{\bm{a}}}_{2}p+\lambda$

Proceeding as in the case of ${\bm{h}}({\bm{W}})$ above, we now concatenate
the above gradients ${\bm{g}}_{j}$ w.r.t the index $j$ in a vector form (of
dimension $pd$) to get the following $pd-$dimensional gradient vector of the
empirical loss,

$\nabla_{{\bm{W}}}\tilde{L}={\bm{g}}({\bm{W}})\coloneqq\left[{\bm{g}}_{1}({\bm{W}}),{\bm{g}}_{2}({\bm{W}}),\ldots,{\bm{g}}_{p}({\bm{W}})\right]$

and the Lipschitz constant of ${\bm{g}}$ \- and hence the gradient Lipschitz
constant for $\tilde{L}$ to be bounded as,

$\textup{gLip}(\tilde{L})\leq\sqrt{p}\left(\frac{\sqrt{p}\norm{{\bm{a}}}_{2}M_{D}^{2}B_{x}}{4}+\left(\frac{2+\norm{c}_{2}}{4}+\frac{\norm{{\bm{a}}}_{2}B_{\sigma}}{4}\right)M^{\prime}_{D}B_{x}p+\lambda\right)$

Thus we get the expression as required in the proof in Section 5.

##  Appendix C Estimating the Time Needed to Reach $\epsilon$ Error

###  C.1 Recalling the Constants in the Bound

Invoking Theorem 3 from Shi et al. (2020) with a “time horizon” parameter
$T>0$, $f=\tilde{L}$, the convergence guarantee for running $k$ steps of SGD
at a step-size $s$ as given in Theorem 3.2 for $k\cdot s\in(0,T)$ and
$s\in(0,1/{\rm gLip}(\tilde{L}))$ can be given as,

$\mathbb{E}\tilde{L}({\bm{W}}^{k})-\min_{{\bm{W}}}\tilde{L}({\bm{W}})\leq(A(\tilde{L})+B(T,\tilde{L}))s+C(s,\tilde{L})\norm{\rho-\mu_{s}}_{\mu_{s}^{-1}}e^{-s\lambda_{s}k}$

where $A,B,C$ are constants as mentioned in Shi et al. (2020). In particular,
the $C$ therein was defined as follows

$C(s,\tilde{L})=\left(\int_{{\mathbb{R}}^{p\times
d}}(\tilde{L}({\bm{W}})-\min_{{\bm{W}}}\tilde{L}({\bm{W}}))^{2}\mu_{s}({\bm{W}})\,d{\bm{W}}\right)^{1/2}$

for $\mu_{s}({\bm{W}})$ being the Gibbs’ measure,
$\mu_{s}({\bm{W}})=\frac{1}{Z_{s}}\exp\left({-\frac{2\tilde{L}({\bm{W}})}{s}}\right)$
with $Z_{s}$ being the normalization factor.

For determining $\lambda_{s}$, Shi et al. (2020) consider the function
$V_{s}({\bm{W}})=\norm{\nabla\tilde{L}}^{2}/s-\Delta\tilde{L}.$ Let
$R_{0,s}>0$ be large enough such that $V_{s}({\bm{W}})>0$ for
$\norm{{\bm{W}}}_{F}\geq R_{0,s}.$ For $R_{s}>R_{0,s},$ Shi et al. (2020)
define $\epsilon(R_{s})$ as

$\epsilon(R_{s})=\frac{1}{\inf\\{V_{s}({\bm{W}})\,:\,\norm{{\bm{W}}}_{F}\geq
R_{s}\\}}$

where $R_{s}$ is assumed large enough such that $\int_{\norm{{\bm{W}}}_{F}\leq
R_{s}}d\mu_{s}\geq 1/2.$ For $B_{R_{s}}$ as the ball of radius ${R_{s}}$
centered at origin in ${\mathbb{R}}^{p\times d}$, Shi et al. (2020) define

$\mu_{s,R_{s}}=\left[\int_{\norm{{\bm{W}}}_{F}\leq
R_{s}}d\mu_{s}({\bm{W}})\right]^{-1}\mu_{s}({\bm{W}})\mathbf{1}_{\norm{{\bm{W}}}_{F}\leq
R_{s}}.$

Using the Poincaré inequality in a bounded domain[Evans (2010), Theorem 1,
Chapter 5.8], Shi et al. (2020) define the constant $C(R_{s})$ to be s.t the
the following holds $\forall h\in C_{c}^{\infty}({\mathbb{R}}^{d})$,

$\int_{{\bm{W}}\in{\mathbb{R}}^{p\times d}}h^{2}d\mu_{s,R_{s}}\leq s\cdot
C(R_{s})\int_{{\bm{W}}\in{\mathbb{R}}^{p\times d}}\norm{\nabla
h}^{2}\mu_{s,R_{s}}d\mu_{s,R_{s}}+\left(\int_{{\bm{W}}\in{\mathbb{R}}^{p\times
d}}h\,d\mu_{s,R_{s}}\right)^{2}$

Then the key quantity $\lambda_{s}$ occurring in the aforementioned
convergence guarantee for SGD was shown to be,

$\lambda_{s}=\frac{1+3s\left(\inf_{{\bm{W}}\in{\mathbb{R}}^{p\times
d}}V_{s}({\bm{W}})\right)\epsilon(R_{s})}{2(C(R_{s})+3\epsilon(R_{s}))}$

###  C.2 Time Estimation

From above, it follows that we can write the error of running SGD at a
constant step-size $s$ for $k=\frac{T}{s}$ iterations as,

${\mathbb{E}}\tilde{L}({\bm{W}}^{\frac{T}{s}})-\min_{{\bm{W}}}\tilde{L}({\bm{W}})\leq(A(\tilde{L})+B(T,\tilde{L}))\cdot
s+C(s,\tilde{L})\norm{\rho-\mu_{s}}_{\mu_{s}^{-1}}e^{-\lambda_{s}\cdot T}$

Hence one way to get the above error below any given $\epsilon>0$ is to choose
the step-size as,

$s=s^{*}(\epsilon,T)\coloneqq\min\left(\frac{1}{{\rm
gLip}(\tilde{L})},\frac{\epsilon}{2}\cdot\frac{1}{(A(\tilde{L})+B(T,\tilde{L}))}\right)$

and initialize the weights by sampling from any distribution with p.d.f
$\rho=\rho_{\rm initial}~{}s.t$

$\norm{\rho_{\rm
initial}-\mu_{s^{*}}}_{\mu_{s^{*}}^{-1}}\leq\frac{\epsilon}{2}\cdot\frac{e^{\lambda_{s^{*}}\cdot
T}}{C(s^{*},\tilde{L})}$

We recall that in above the time horizon $T$ is a free parameter - which in
the above two equations we see to be determining both the step-size $s^{*}$
and the $\rho_{\rm initial}$.

From above the accuracy guarantee as given in the statement of Theorem 3.2
follows for the constant step-size $s^{*}(\epsilon,T)$ SGD for
$k=\frac{T}{s^{*}(\epsilon,T)}$ steps starting from weights sampled from the
distribution $\rho_{\rm initial}$.