Sharper Guarantees for Learning Neural Network Classifiers with Gradient Methods

## Abstract

Abstract In this paper, we study the data-dependent convergence and generalization behavior of gradient methods for neural networks with smooth activation. Our first result is a novel bound on the excess risk of deep networks trained by the logistic loss, via an alogirthmic stability analysis. Compared to previous works, our results improve upon the shortcomings of the well-established Rademacher complexity-based bounds. Importantly, the bounds we derive in this paper are tighter, hold even for neural networks of small width, do not scale unfavorably with width, are algorithm-dependent, and consequently capture the role of initialization on the sample complexity of gradient descent for deep nets. Specialized to noiseless data separable with margin γ 𝛾 \gamma italic_γ by neural tangent kernel (NTK) features of a network of width Ω ⁢ ( poly ⁢ ( log ⁡ ( n ) ) ) Ω poly 𝑛 \Omega(\mathrm{poly}(\log(n))) roman_Ω ( roman_poly ( roman_log ( italic_n ) ) ) , we show the test-error rate to be e O ⁢ ( L ) / γ 2 ⁢ n superscript 𝑒 𝑂 𝐿 superscript 𝛾 2 𝑛 e^{O(L)}/{\gamma^{2}n} italic_e start_POSTSUPERSCRIPT italic_O ( italic_L ) end_POSTSUPERSCRIPT / italic_γ start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT italic_n , where n 𝑛 n italic_n is the training set size and L 𝐿 L italic_L denotes the number of hidden layers. This is an improvement in the test loss bound compared to previous works while maintaining the poly-logarithmic width conditions. We further investigate excess risk bounds for deep nets trained with noisy data, establishing that under a polynomial condition on the network width, gradient descent can achieve the optimal excess risk. Finally, we show that a large step-size significantly improves upon the NTK regime’s results in classifying the XOR distribution. In particular, we show for a one-hidden-layer neural network of constant width m 𝑚 m italic_m with quadratic activation and standard Gaussian initialization that mini-batch SGD with linear sample complexity and with a large step-size η = m 𝜂 𝑚 \eta=m italic_η = italic_m reaches the perfect test accuracy after only ⌈ log ⁡ ( d ) ⌉ 𝑑 \lceil\log(d)\rceil ⌈ roman_log ( italic_d ) ⌉ iterations, where d 𝑑 d italic_d is the data dimension.

## 1 Introduction

### 1.1 Overview

Neural networks, with their vast capacity for capturing intricate patterns in data, have triggered a paradigm shift in machine learning. Despite the power of these networks in modeling complex relationships, the interplay between their optimization and generalization behaviors (that is the gap between training and test errors) continues to be a compelling area of research. In practice, training neural networks using gradient-based optimization methods often leads to interpolation. That is, deep networks can meticulously fit the training data, driving empirical loss to near-zero and training error to perfect classification. However, these networks also demonstrate the capability to generalize well to unseen data.
Various recent research endeavors have explored the training and test error guarantees of deep networks, with a focus on the Neural Tangent Kernel (NTK) regime [, ]. One prominent feature of such works is that during gradient descent iterates the network’s weights are constrained to move at most a constant distance with respect to overparameterization i.e., $\|w^{\star}-w_{0}\|=O_{m}(1)$, where $w^{\star}\in\mathbb{R}^{p}$ denotes the vector of target weights, $w_{0}$ is the initial weight vector, and $m$ is the network width [, , ].

Yet, even for the relatively simple setting of learning deep nets in the kernel regime, the existing generalization bounds are still sub-optimal. Moreover the boundaries of the kernel regime are still largely unknown and an active area of research [, , ]. While the kernel regime can partially demonstrate the behavior of neural networks, the resulting guarantees often require large width, small step-size or large iteration and sample complexities. There is increasing evidence in recent years that for certain class of data distributions neural networks can overcome these limitations by using large step-sizes which allows the network’s parameters to move a long distance from initialization often leading to better sample and computational complexities [, , ].

In this work, we study the generalization and convergence behavior of gradient-based algorithms in neural nets with smooth activation functions for a wide class of data distributions. Our first result characterizes the test and train loss rates for classification problems under the condition that deviation from initialization is bounded depending on the network’s width. In particular, for $L$-hidden layer networks our results hold under $\|w^{\star}-w_{0}\|\lesssim m^{O(1/L)}$, allowing the network’s weights to move from initialization a distance increasing with $m$. This shows that kernel regime continues to hold for a wider range of setups than previous results for which the deviations are constrained to be constant in $m$. The key reason for this improvement is exploiting the objective’s Hessian structure in the gradient descent (GD) path. More importantly, using the Hessian information enables us to develop, for the first time, *algorithm-dependent* generalization bounds of deep neural networks. As it will be discussed throughout the paper, the bounds we derive via algorithmic-stability are tighter than previous relevant bounds in literature (Table 1).

**Table 1: Comparison of our results on learning deep nets with GD under NTK separability condition to related prior results. Here $m:$ network width, $L:$ network depth, $\gamma:$ NTK-margin, $n:$ number of samples and $T:$ number of iterations.**
|  | Activation | Width | Train Loss | Test Loss |
| --- | --- | --- | --- | --- |
| [, ] | ReLU | $\Omega(\mathrm{poly}(\frac{\log(n)}{\gamma}))$ | $\;\widetilde{O}(\frac{1}{\gamma^{2}T})$ | $\widetilde{O}(\frac{e^{O(L)}}{\gamma^{2}}\sqrt{\frac{m}{n}}+\frac{1}{\gamma^{2<br>}T})$ |
| Thm. 2.1, Cor. 1 | Smooth | $\Omega(\mathrm{poly}(\frac{\log(n)}{\gamma}))$ | $\widetilde{O}(\frac{1}{\gamma^{2}T})$ | $\widetilde{O}(\frac{e^{O(L)}}{\gamma^{2}n}+\frac{1}{\gamma^{2}T})$ |

**Table 2: Comparison of our findings on learning the $d$-dimensional XOR distribution with SGD to relevant prior results.**
|  | Activation | Width | Iteration | Sample |
| --- | --- | --- | --- | --- |
| [, , ],Cor.1 | all | $\Omega(\mathrm{poly}(d))$ | $d^{2}$ | $\widetilde{O}(d^{2})$ |
| [] | ReLU | $\Omega(\mathrm{poly}(\log(d)))$ | $\mathrm{poly}(\log(d))$ | $\widetilde{O}(d)$ |
| Thm. 2.4 | Quadratic | $\Omega(1)$ | $\log(d)$ | $\widetilde{O}(d)$ |

Our analysis is inspired by the recent progress in bounding the multi-layer networks’ Hessian spectral norm [, ] as well as the progress in understanding the loss landscape of neural net classifiers during gradient descent iterates [, ]. In contrast with previous works, we then specialize these results to a well-known NTK separability condition tailored for noiseless data and show that our results substantially improve the upper-bounds on the test error performance while still allowing the width to be small, specifically poly-logarithmic on sample size (Table 1).
A more detailed comparison is deferred to Section 1.2. We also consider the case of noisy data distributions and show that deep models are consistent, i.e., they can achieve the optimal test loss in the presence of noise as sample size grows.

While these results improve upon the existing bounds within the NTK regime, we show in Section 2.2 that using large step-size can further improve both the computational and sample complexities. In particular, we show for the stylized setup of data following the XOR distribution, a two-layer neural network with quadratic activation reaches zero test error after only $\log(d)$ steps of mini-batch Stochastic Gradient Descent (SGD) with an aggressive step-size $\eta=m$. A comparison of our findings with the guarantees of the kernel regime with both GD (as per Corollary 1; see also [, ]) and one-pass SGD [, ] and the most relevant work in the feature learning regime [] is summarized in Table 2. Below is a summary of our contributions.

- $\bullet$
In Theorem 2.1, we develop sufficient conditions for the global convergence of gradient descent in deep and smooth networks and show that if $m=\Omega(\|w^{\star}-w_{0}\|^{6L+4})$, the training loss is bounded by $O(\frac{\|w^{\star}-w_{0}\|^{2}}{\eta T})$, where $\eta$ is the step-size and $w^{\star}$ can be any choice of network weights that achieves small training loss. Under similar conditions on $m$, we show the generalization error is bounded by $O(\frac{\|w^{\star}-w_{0}\|^{2}G_{0}^{2}}{n})$ where $G_{0}$ is the Lipschitz parameter of the network at initialization.
- $\bullet$
In Corollary 1, we interpret these results by specializing them to a commonly-used margin-based NTK separability condition. The results of the corollary and comparison to previous works in literature are summarized in Table 1. A promising feature of our approach is that the test loss bound does not have an unfavorable dependence on the width while still maintaining minimum poly-logarithmic width conditions, which is new in the context of deep learning. To the best of our knowledge this is the tightest test error bound for deep nets trained by GD in the NTK regime.
- $\bullet$
We consider the more general case of noisy data with non-vanishing optimal test loss in Theorem 2.3 and show that under a polynomial growth condition on network width, GD achieves a convergence rate of $1/\sqrt{n}$ to the optimal loss after $T=\sqrt{n}$ iterations.
- $\bullet$
In Section 2.2, we consider the $d$-dimensional XOR distribution and show that a one-hidden layer network of constant width after exactly $\log(d)$ iterations of SGD with step-size $\eta=m$ achieves perfect test accuracy with $n=\widetilde{O}(d)$ samples, considerably improving kernel regime’s limitations.

### 1.2 Prior works

##### Generalization of deep nets.

Among prior works on the generalization capabilities of deep networks, the only initialization dependent bounds were provided in
[] obtaining bounds of order $O(\frac{\mathcal{R}}{n})$ where the Rademacher complexity is derived as $\mathcal{R}:=(\prod_{i=1}^{L}\|W_{i}\|_{2})(\sum_{i=1}^{L}\frac{\|W_{i}^{\top}
-M_{i}^{\top}\|_{2,1}^{2/3}}{\|W_{i}\|_{2}^{2/3}})^{3/2}.$ Here $\|W_{i}\|_{2}$ is the spectral norm of the weight matrix of layer $i$ (typically a constant) and $M_{i}$s are any data-independent matrices. Thus one can choose $M_{i}=W_{i,0}$, i.e., the initialization weight matrix. In fact, the above bound resembles the bound that we obtained via an optimization-dependent stability analysis. However, note that $\mathcal{R}$ depends on the distance traversed by weights through the $\ell_{2,1}$ norm which is always larger than the Frobenius norm, and in the worst case, the gap can be significantly large depending on the width. To see this, note that for a matrix $V\in\mathbb{R}^{m\times m}$ it holds that $\|V\|_{2,1}\leq\sqrt{m}\|V\|_{F}$. We note that “initialization-independent” bounds (e.g., [, ]) that are usually proportional to $\|w_{t}\|$ (rather than $\|w_{t}-w_{0}\|$) are strictly looser than the bound we obtain. This is primarily due to the fact that $\|w_{t}-w_{0}\|$ can be much smaller than $\|w_{t}\|$ and in fact as our experiments show $\|w_{t}-w_{0}\|$ is of constant order and can even decrease with width. Whereas, $\|w_{t}\|$ (or $\|w_{t}\|/\sqrt{m}$ due to the normalization in our setup) grows by increasing $m$, making the initialization-independent bounds potentially grow with width at the rate $O(\sqrt{m})$, despite lacking an explicit dependence on $m$. Hence, for wide networks, prior generalization bounds of deep neural nets based on Rademacher complexity are larger than the bound we derive in Theorem 2.1.

##### Test rates under the NTK separability condition.

Other works that provide generalization bounds and optimization guarantees for neural nets include [, , , , , , ]. In particular, [] derived the width condition $m=\Omega(\frac{\mathrm{poly}(\log(n))}{\gamma^{8}})$ for achieving the $\frac{1}{\gamma^{2}\sqrt{n}}$-test error rate in two-layer nets via a uniform-convergence argument [].
This bound was extended to deep networks in [] with a generalization gap of order

$$ $\widetilde{O}\left(\frac{4^{L}}{\gamma^{2}}\sqrt{\frac{m}{n}}\wedge\left(\frac {L^{3/2}}{\gamma^{2}\sqrt{n}}+\frac{L^{11/3}}{\gamma^{2}m^{1/6}}\right)\right),$ $$

where $\wedge$ takes the minimum of two quantities. As discussed earlier, this bound is dependent on width since it is derived essentially by taking the minimum of two generalization bounds based on Rademacher complexity derived in [] and []. Importantly, in the small width regime the bound simplifies into $\tilde{O}(\frac{4^{L}}{\gamma^{2}}\sqrt{\frac{m}{n}})$, which has an undesirable dependence on the width. In this paper, we improve the generalization gap to $\tilde{O}(\frac{e^{O(L)}}{\gamma^{2}n})$ under the width condition $m=\Omega(\mathrm{poly}(\log(n)/\gamma))$. To the best of our knowledge, these are the smallest generalization bound and width condition in literature up to now for learning deep neural nets.
The key reason behind the improvement is leveraging the Hessian structure of the objective throughout the gradient descent iterates. The improved generalization guarantees result from the algorithmic dependency of our bounds and in fact the bounds can even be expressed such that they solely depend on the cumulative training loss (c.f. Eq. 5). As the training loss captures the role of initialization and is independent of width, the resulting generalization bounds share the same favorable properties.

##### Feature learning and the XOR distribution.

Some recent works have pointed out the limitations of the kernel regime in understanding the full power of neural networks []. In particular, as in the kernel regime, the networks weights are bounded not to move significantly from initialization, the learned features are not considerably different from those learned at initialization. On the other hand, it is hypothesized that neural networks can learn the true underlying features of the data distribution if the network weights are allowed (by large step-sizes or avoiding early-stopping) to move a large distance from initialization. This phenomenon was first proved for specific regression tasks where the labels essentially only depend on a small number of features, such as when $y=g(Ux)$ for $U\in\mathbb{R}^{k\times d}$ where $k\ll d$ [, , , ], using one large gradient-descent step leading to superior sample complexities compared to the kernel regime. For classification tasks, some focus has been on the XOR distribution(a.k.a. parities) []. Recent works have studied the problem of learning the $d$-dimensional XOR distribution using neural networks in both NTK and feature learning settings [, , , , ]. Specifically, it has been shown that under NTK with a sufficiently small step size, a polynomially wide network requires $d^{2}$ GD steps and $d^{2}$ sample size. Some studies have achieved linear sample complexity for learning XOR [, , ]; but these methods involve more computational effort compared to our results. The work most related to ours is [], which demonstrated that with a particular Gaussian initialization, a ReLU network requires $\mathrm{poly}(\log(d))$ large SGD steps and $\mathrm{poly}(\log(d))$ neurons to learn XOR with linear sample complexity. In contrast, we show that by using a quadratic activation, learning this distribution requires only $\log(d)$ steps with a *constant*-width network, while maintaining the same linear sample complexity.

### Notation

Probability and expectation with respect to the randomness in random variable $x$ are denoted by $\Pr_{x}(\cdot)$ and $\mathbb{E}_{x}[\cdot]$. We use the standard complexity notation $\lesssim,o(\cdot),O(\cdot),\Theta(\cdot),\Omega(\cdot)$ and use $\tilde{o}(\cdot),\tilde{O}(\cdot),\tilde{\Theta}(\cdot),\tilde{\Omega}(\cdot)$ to hide polylogarithmic factors. We denote $a\wedge b:=\min\{a,b\}$. The Gradient and Hessian of the model $\Phi:\mathbb{R}^{p\times d}\rightarrow\mathbb{R}$ with respect to the first input (i.e., weights) are denoted by $\nabla\Phi$ and $\nabla^{2}\Phi$, respectively. The minimum eigenvalue of a symmetric matrix is denoted by $\lambda_{\min}(\cdot)$. We use $\|\cdot\|$ for the Euclidean norm of vectors and $\|\cdot\|_{2}$ for the spectral norm of matrices. We denote $[w_{1},w_{2}]:=\{w\,:\,w=\alpha w_{1}+(1-\alpha)w_{2},\alpha\in[0,1]\}$ the line between $w_{1},w_{2}\in\mathbb{R}^{p}$.

## 2 Main results

Throughout the paper, we consider the following unregularized objective for a neural network classifier parameterized with $w\in\mathbb{R}^{p},$

$$ $\displaystyle\min_{w\in\mathbb{R}^{p}}\widehat{F}(w):=\frac{1}{n}\sum_{i=1}^{n }f\left(y_{i}\Phi(w,x_{i})\right),$ (1) $$

with data points satisfying $\|x\|\leq 1$, the binary labels $y_{i}\in\{\pm 1\}$ and $f(\cdot)$ is a loss function for classification tasks such as the logistic loss, $f(t):=\log(1+e^{-t})$ and $\Phi(\cdot,x)$ is the network’s output. We also define the test loss as $F(w):=\mathbb{E}_{x,y}[f(y\Phi(w,x))].$

### 2.1 Train and test loss bounds in deep nets

In our first theorem, we establish conditions for the width and target weights of the network that guarantees the training loss decays to zero if the network can interpolate the training set. We consider gradient descent update rule where at any iteration $t\leq T$: $w_{t+1}=w_{t}-\eta\nabla\widehat{F}(w_{t}).$
Before stating the theorem, we note that this result is valid under the standard descent-lemma condition for the step-size, as stated in Lemma 3 in the appendix. In particular, the descent lemma holds if the step-size satisfies the standard Eq. 34 in the appendix.

###### Theorem 2.1 (Train & Test loss of deep nets) .

Consider $L$-layer neural network model with the model defined as $\Phi(w,x):=\frac{1}{\sqrt{m}}W_{L+1}^{\top}(\frac{1}{\sqrt{m}}\sigma(W_{L}^{
\top}\cdots\frac{1}{\sqrt{m}}\sigma(W_{1}^{\top}x)\cdots)$ where $\sigma$ is a $1$-smooth and $1$-Lipschitz activation function such that $\sigma(0)=0$. Moreover, let $\beta_{L}$ be a constant that only depends on $L$, let all parameters of the network be initialized as i.i.d. standard Gaussian and assume the step-size satisfies the condition of the descent lemma. Fix $T$ and assume the target weights vector $w^{\star}\in\mathbb{R}^{p}$ that obtain small training loss such that

$$ $\displaystyle\rho^{\star}\geq\max\left\{\sqrt{\eta T\widehat{F}(w^{\star})}, \sqrt{\eta\widehat{F}(w_{0})}\right\}.$ (2) $$

where $\rho^{\star}:=\|w^{\star}-w_{0}\|$. Moreover, assume the width $m$ is large enough such that it satisfies,

$$ $\displaystyle m\geq 4\beta_{L}^{2}(6{\rho^{\star}})^{6L+4}$ (3) $$

Then, $\|w_{t}-w_{0}\|=O(\rho^{\star})$ and the training loss satisfies with high probability over initialization,

$$ $\displaystyle\widehat{F}(w_{T})\leq\frac{4{\rho^{\star}}^{2}}{\eta T}.$ (4) $$

Moreover, the expected generalization gap satisfies with high probability over initialization,

$$ $\displaystyle\mathbb{E}_{\mathcal{S}}\left[{F}(w_{T})-\widehat{F}(w_{T})\right ]\leq 2.2\,\frac{\eta(G_{0}+1/4)^{2}}{n}\,\mathbb{E}_{\mathcal{S}}\left[\sum_{ t=0}^{T-1}\widehat{F}(w_{t})\right],$ (5) $$

where $G_{0}$ is the Lipschitz parameter of network at initialization i.e., $\|\nabla\Phi(w_{0},\cdot)\|\leq G_{0}$ and the expectation is over the data sampling.

In words, the main condition of the theorem is the existence of network weights denoted by $w^{\star}$ that achieves small training error (Eq. 2) and its distance from initialization is at most $O(m^{1/(6L+4))})$ as implied by Eq. 3. Under these conditions, the training loss is controlled solely by $\|w^{\star}-w_{0}\|$ and has no explicit dependence on the width or depth of the network.
As it will be stated in Corollary 1, in the NTK regime with margin $\gamma$ it holds $\|w^{\star}-w_{0}\|=O(\log(n)/\gamma)$, leading to the width condition $m=\Omega(\mathrm{poly}(\log(n)/\gamma))$. This is unlike previous results which either required polynomial width (such as [, ]) or led to sub-optimal bounds (e.g., [, ]).

In general, the theorem is valid for any feasible minimizer $w^{\star}\in\mathbb{R}^{p}$. Thus, we can choose $w^{\star}$ with smallest value for $\|w^{\star}-w_{0}\|$ to optimize the bounds. With such choice of $w^{\star}$, the distance the weights obtained by GD travel is also minimized as $\forall t\in[T]:\|w_{t}-w_{0}\|=O(\|w^{\star}-w_{0}\|)$. Thus, gradient descent tends toward solutions which attain small loss and lie at minimum possible distance from initialization. This is in line with related prior observations in several other works such as [, ].

The theorem also shows the sample complexity and iteration complexity of learning deep networks with gradient descent and demonstrates the role of initialization and weight’s norms on the test error. Note that by replacing our training loss guarantees, the generalization gap simplifies into:

$$ $\displaystyle\mathbb{E}_{\mathcal{S}}\left[{F}(w_{T})-\widehat{F}(w_{T})\right ]\leq 9\frac{{\rho^{\star}}^{2}(G_{0}+1/4)^{2}}{n}\,.$ (6) $$

Hence, the test loss after $T=\Theta(n)$ iterations takes the form of

$$ $\displaystyle\mathbb{E}_{\mathcal{S}}\Big{[}{F}(w_{T})\Big{]}=O\left(\frac{ \Big{\|}w^{\star}-w_{0}\Big{\|}^{2}G_{0}^{2}}{n}+\frac{\Big{\|}w^{\star}-w_{0} \Big{\|}^{2}}{\eta n}\right),$ (7) $$

where the first term is the generalization error and the second term is the training loss. Remarkably, Eq. 7 shows the tight correlation between the two terms, as the generalization gap is virtually the optimization error scaled by the squared Lipschitz constant $G_{0}^{2}$. This is indeed the consequence of Eq. 5 which bounds the generalization gap based on the cumulative optimization error.

At a high-level, the test loss essentially depends on two quantities: (i) the Euclidean distance between the target weights and the initialization and (ii) the Lipschitz parameter of network at initialization. Due to the algorithmic-dependent nature of our generalization bounds and unlike the prior bounds in literature, the bound in Eq. 7 captures the role of initialization on the test error: smaller deviations from initialization lead to smaller test error bounds and in particular the bound approaches zero as $\rho^{\star}$ goes to zero. As discussed earlier, gradient descent favors such solutions with small deviations from initialization.
In addition to “distance from initialization”, the test error bounds also depend on the squared Lipschitz parameter of the network. For standard Gaussian initialization, it can be shown (cf. Section A.2) that $G_{0}\lesssim e^{O(L)}$ which introduces an exponential dependence on depth to the generalization bound. We remark that this dependence also appears in the corresponding bounds derived via uniform convergence and Rademacher complexity (e.g., [, , ]) through the term $\prod_{i=1}^{L}\|W_{i}\|_{2}$.

An interesting feature of our approach is the algorithmic dependent bound in Eq. 5. This bound is generally tighter than the bound in Eq. 6. With the descent lemma condition on the step-size (c.f. Lemma 3) it holds that $\eta<1/(G_{0}^{2}+1/4)$ which simplifies the bound into:

$$ $\displaystyle\mathbb{E}_{\mathcal{S}}\left[{F}(w_{T})-\widehat{F}(w_{T})\right ]\leq\frac{2.2}{n}\,\mathbb{E}_{\mathcal{S}}\left[\sum_{t=0}^{T-1}\widehat{F}( w_{t})\right].$ (8) $$

Hence, we have a bound which only depends on the training performance and the number of training samples. In our experiments in Section 3, we compute this bound for real-world data and compare it with the empirical results for generalization and test loss.

###### Remark 2.2 .

Our analysis relies on the recent progress in characterising the spectral norm of the deep net’s Hessian during GD updates [, ]. In particular, [] proved that with standard Gaussian initialization, the model’s Hessian is bounded with high probability by $\|\nabla^{2}\Phi(w,x)\|=O(\frac{R^{3L}}{\sqrt{m}})$ if $\|w-w_{0}\|\leq R$. These guarantees can also be used to study the convergence rate of deep networks trained by quadratic loss as done by [, ] but their approach leads to excessively large width conditions. In contrast, here we consider classification tasks with an improved poly-logarithmic width requirement and also study the generalization performance.

Recently, the algorithmic stability has been employed for two-layer neural nets in [, ]which is an improved analysis of the stability-based approach typically used for convex objectives in [, , ]. Here, we essentially extend the stability analysis to deep networks. Compared to the two-layer nets, here in every iterate of gradient descent, the Hessian’s norm guarantees depend on the network’s weights. In particular, the analysis has to take into account that both $\|w_{t}-w_{0}\|$ and $\|w_{t}-w^{\star}\|$ remain small during GD updates. This is necessary in order to ensure the Hessian’s norm guarantees and the approximate quasi convexity property hold during all GD iterates. We do this by an induction based argument which bounds these terms based on the fixed quantity $\|w^{\star}-w_{0}\|$, which is guaranteed to be bounded based on width by assumption.

###### Theorem 2.1 (Train & Test loss of deep nets) .

###### Remark 2.2 .

#### 2.1.1 Specializing to the NTK-separablity condition

The results in the last section can be specifically applied to a class of data distributions that includes the XOR distribution which will be discussed more through the rest of the paper. Before stating our result in the corollary, we state the neural tangent kernel (NTK) separability assumption.

###### Assumption 1 (NTK-separability [ Nitanda et al., 2019 ] ) .

Assume the tangent kernel of the model at initialization separates the data with margin $\gamma$ i.e., for a unit-norm vector $w\in\mathbb{R}^{p}$ it holds for all $i\in[n]:$
$y_{i}\left\langle\nabla\Phi\left(w_{0},x_{i}\right),w\right\rangle\geq\gamma$.

In words, the above assumption implies that the features learned by the gradient at initialization can be linearly separated by some weights $w$. This assumption is commonly used in deep learning literature for studying classification tasks [, ].

###### Corollary 1 (NTK results) .

Consider the same setup as Theorem 2.1 and let Assumption 1 hold. Define constant $B>0$ that bounds the model’s output at initialization i.e., $\forall i\in[n]:|\Phi(w_{0},x_{i})|<B$. Assume the width is large enough such that
$m\geq\beta_{L}^{2}\left(\frac{2B+\log(1/{\epsilon})}{\gamma}\right)^{6L+4}.$
Then there exists $w^{\star}\in\mathbb{R}^{p}$ such that $F(w^{\star})\leq{\epsilon}$ and $w^{\star}$ lies at a bounded Euclidean distance from initialization such that $\|w^{\star}-w_{0}\|\leq\frac{1}{\gamma}(2B+\log(1/{\epsilon}))$.

To interpret this result, we apply this result to Theorem 2.1 and fix ${\epsilon}=1/T$. First note that the output of network at initialization is constant with high-probability over initialization (e.g., see [, Lemma F.4], [, Lemma 4.4]) which implies $B=\tilde{O}(1)$. Overall, with the stopping condition $T=\Theta(n)$ and recalling that $G_{0}\leq e^{O(L)}$, Corollary 1 yields the expected test error rate of order

$$ $\widetilde{O}\left(\frac{e^{O(L)}}{\gamma^{2}n}+\frac{1}{\gamma^{2}T}\right),$ $$

under the condition that $m=\Omega(\mathrm{poly}(\log(n))\cdot\beta_{L}/\gamma^{6L+4})$. We remark the test error bound does not have any explicit dependence on $m$. This in itself is not surprising as one intuitively expects the bound not to scale unfavorably with width. Although recent works have derived width independent generalization bounds for two-layer networks [, , ], we are not aware of any prior work proving width-independent bounds for learning multi-layer networks with GD. As discussed earlier in the introduction, the bound derived in the closely related work [] scales unfavorably with $m$ as it grows at the rate $\sqrt{m/n}$. In fact, the authors refer to deriving width-independent bounds as an open problem in [, Sec 3.1].

###### Assumption 1 (NTK-separability [ Nitanda et al., 2019 ] ) .

###### Corollary 1 (NTK results) .

#### 2.1.2 Consistency of GD with noisy data

The results of the last section mainly apply to data settings when the network can find the optimal solution within a bounded distance from initialization that depends on the network’s width. This setting was specially tailored to the noiseless case where achieving vanishing test loss was possible. We discuss next the case of learning deep nets by noisy data and show that achieving optimal test loss might be feasible in this setting.

###### Theorem 2.3 (Test error for noisy data) .

Consider the same setup and notation as Theorem 2.1 and assume the width of the network satisfies $m\geq\beta_{L}^{2}n^{3L+3}$ and the step-size satisfies the conditions of the descent lemma. Then, with high probability over initialization the expected test loss at iteration $T=\sqrt{n}$ is bounded as:

$$ $\displaystyle\mathbb{E}_{\mathcal{S}}\Big{[}F(w_{T})\Big{]}\leq\left(1+\frac{4 }{\sqrt{n}}\right)\left(F(w^{\star})+\frac{{\rho^{\star}}^{2}}{\eta\sqrt{n}}+ \frac{{\rho^{\star}}^{2}}{n\sqrt{n}}+\frac{1}{\sqrt{n}}\right),$ (9) $$

where $\rho^{\star}:=\|w^{\star}-w_{0}\|$ and $w^{\star}$ is the minimizer of the population loss i.e., $w^{\star}=\arg\min_{w\in\mathbb{R}^{p}}F(w)$.

The result above shows that GD reaches optimal level given that $\sqrt{n}\gg\|w^{\star}-w_{0}\|^{2}$; as for large $n$, it leads to the simplified expression

$$ $\mathbb{E}_{\mathcal{S}}\Big{[}F(w_{T})\Big{]}-F(w^{\star})=O\left(\frac{\Big{ \|}w^{\star}-w_{0}\Big{\|}^{2}+F(w^{\star})}{\sqrt{n}}\right).$ $$

Thus, even in non-interpolating regime, GD can still achieve the optimal solution of over-parameterized deep networks.

The bound in Eq. 9 is derived via a stability argument which bounds the generalization gap based on the training performance. However, contrary to the conditions of Theorem 2.1 here we do not have the interpolation condition as $F(w^{\star})$ is not vanishing. This comes at the expense of a larger width condition where the width is polynomial in $n$ whereas poly-logarithmic width was sufficient in Theorem 2.1. The condition on early stopping further guarantees that the test loss reach the Bayes error given $n$ is sufficiently large. It is worth noting that the setting above still is operating almost in the NTK regime. This can be verified by the observation that for the bound to be meaningful it should hold that $\|w^{\star}-w_{0}\|\ll\sqrt{n}\lesssim m^{O(1/L)}$, as per the theorem’s width condition. Finally, we remark that the result of Thm 2.3 nicely connects to recent empirical and theoretical results [] which show that with early-stopping, GD can find a good solution for clustered data with label noise.

###### Theorem 2.3 (Test error for noisy data) .

### 2.2 Overcoming NTK limitations for XOR dataset

Next, we show the previous NTK guarantees can be improved by using large step-sizes. We consider the stylized set-up of learning the $d-$dimensional XOR data distribution with two layer neural networks with quadratic activation which takes the following form:

$$ $\displaystyle\Phi(w,x)=\frac{1}{2m}\sum_{i=1}^{m}a_{i}(x^{\top}w_{i})^{2}.$ $$

In the above, $a_{i}\in\{\pm 1\}$ are fixed during training and satisfy $\sum_{i}a_{i}=0$ and only the first layer weights $w_{i}\in\mathbb{R}^{d}$ are trained. For initialization of first layer’s weights we have $\forall i\leq m,j\leq d:w_{0,ij}\mathrel{\overset{\text{\small{iid}}}{
\scalebox{1.5}[1.0]{$\sim$}}}N(0,\frac{1}{d})$. The data points $(x,y)\in\{\pm 1\}^{d}\times\{\pm 1\}$ are uniformly drawn from the resulting $d$-dimensional distribution of $2^{d}$ points where the labels are determined as $y=x(1)\cdot x(2).$ For this dataset it is known that the NTK-margin satisfies $\gamma\gtrsim 1/d$ [, , ]. Therefore, by kernel regime’s guarantees (e.g., see Corollary 1) GD can learn this distribution with $d^{2}$ samples in polynomial time with networks of polynomial width. Our next result shows improved guarantees for learning this distribution.
We assume the loss function is the linear loss $f(t)=-t$ and consider mini-batch SGD with the update rule $w^{t+1}=w^{t}-\eta\widehat{F}(w^{t})$, where at each step $n$ data points are drawn i.i.d. from the distribution to form $\widehat{F}(\cdot)$.
The next theorem shows the computational and sample complexities of learning this dataset for the aforementioned setup.

###### Theorem 2.4 (Improved guarantees for learning XOR) .

Assume mini-batch SGD with batch-size $n\geq d\cdot\log^{14}(d)$ and step-size $\eta=m$. Then, after $T=\log(d)$ iterations the test accuracy satisfies with probability at least $1-e^{\log(m)-\log^{2}(d)}-2e^{-\frac{m}{16}}-o_{d}(1)$ over the randomness of initialization and data sampling,

$$ $\displaystyle\mathbb{E}_{x,y}\Big{[}\mathbf{1}_{\{y\Phi(w^{T},x)>0\}}\Big{]} \geq 1-o_{d}(1).$ (10) $$

Hence, with logarithmic number of SGD steps we use $\widetilde{O}(d)$ samples in total to reach almost perfect test accuracy. We remark based on the guarantees of Theorem 2.4, the network’s width can be constant and at most must be polynomial in $d$ in order for Eq. 10 to hold with high probability. This aligns with our experiments in Fig. 4, demonstrating that the network’s width can be independent of data dimension as a network of constant width (where $m=20$) suffices for learning arbitrary high-dimensional XOR data. We also note the considerable gains in iteration, sample and computational complexities by Theorem 2.4 resulting from escaping the NTK regime with our step-size selection. Contrary to Theorem 2.1 which required small step-size in accordance with the descent lemma, here the step-size is proportional to width. For a comparison with recent results for this dataset we refer to Table 2. To the best of our knowledge these are the best complexities on iteration and network width for this setup.

###### Remark 2.5 .

The proof of Theorem 2.4 (provided in Appendix E) is nuanced and involves computing expected weights and their corresponding error terms due to SGD sampling noise for each parameter of the network. It is then showed that signal strength (i.e., the strength of important features) grows as $\frac{2^{t}}{\sqrt{d}}$ whereas the error terms due to sampling noise and initialization grow at most at the rate $(1+\frac{1}{\mathrm{poly}(\log(d))})^{t}\frac{\sqrt{d}}{\mathrm{poly}(\log(d))
}+\frac{2^{t}t}{\sqrt{d}\cdot\mathrm{poly}(\log(d))}$. Therefore, after $\log(d)$ steps the noise strength reaches $\frac{\sqrt{d}}{\mathrm{poly}(\log(d))}$ whereas the signal’s magnitude is at least $\sqrt{d}$, letting the signal outgrow the noise and leading to the network classifying every point correctly.

###### Theorem 2.4 (Improved guarantees for learning XOR) .

###### Remark 2.5 .

## 3 Numerical Results

##### Experiments on learning under NTK with small step-size.

In this section, we present numerical results on the behavior of the generalization bound derived in Theorem 2.1 for real-world data (FashionMNIST and MNIST datasets) and compare it with the empirical generalization gap.

For demonstrating our theoretical results, we are interested in the algorithmic-based generalization bound (Eq. 8) derived as,

$$ $\displaystyle\mathbb{E}_{\mathcal{S}}\Big{[}{F}(w_{T})$ $\displaystyle-\widehat{F}(w_{T})\Big{]}\leq\frac{2.2}{n}\,\mathbb{E}_{\mathcal {S}}\left[\sum_{t=0}^{T-1}\widehat{F}(w_{t})\right].$ (11) $$

Figure: Figure 1: Iteration-based distance from initialization ($\|w_{t}-w_{0}\|$), training loss, test loss and generalization gap (i.e., test loss – train loss) for training a two hidden-layer neural network with *FashionMNIST* dataset and two choices of step-size. Here $n=12\times 10^{3},m=500,$ and total number of parameters $p\approx 6\times 10^{5}$.
Refer to caption: x1.png

Figure: Figure 4: Left: Misclassification error based on iteration in learning the $d-$dimensional XOR distribution with SGD. Right: Total number of SGD steps based on data dimension to reach approximately zero test error.
Refer to caption: extracted/5920941/xor.png

We find it helpful to note that the bound above requires the width condition in the theorem (i.e., $m=\Omega(\|w^{\star}-w_{0}\|^{6L+4})$) to be valid. However, verifying this condition is not feasible in general. Moreover, the bound is valid for *expected* generalization gap where the expectation is taken over data sampling. Therefore computing the exact values of both sides of the above inequality is computationally exhaustive. For our experiments we consider one realization to estimate these values.
Due to both of these reasons, the theoretical test loss and generalization loss that we present in this section should only be taken as approximations of the general behavior of the bound and not as an actual verified bound on the generalization. However, in order to reduce these impacts we conduct several experiments with varying network’s width.

We consider binary classification with a 2-hidden layer network with softplus activation ($\sigma(t)=\log(1+e^{t})$) trained by the logistic loss function. Figure 3 presents train, test and generalization behavior of GD for learning a such a model with FashionMNIST dataset. The two lines in each figure correspond to $\eta=0.01,0.1$. In the two rightmost plots, the resulting theoretical generalization and test loss curves derived from Eq. 11 are compared with the empirical values. Note the good alignment between theoretical and empirical behavior for the generalization and test loss.

A similar behavior is observed in Figures 3-3. In Figure 3, we consider a similar setup but we reduce the sample size to $4000$ training data in order to allow larger test-loss behavior. The resulting plots show the training and test loss performance for two choices of $m=250$ and $m=500.$ In Figure 3, we consider the MNIST dataset for $m=300$ and $m=600$ and let the sample size be $n=2000.$ It is noteworthy that the findings in both figures bear similarities, with the theoretical bounds providing non-vacuous and accurate approximations for the test loss and generalization gap.

##### Experiments on learning the XOR distribution.

Figure 4 demonstrates the test error curves associated with learning the XOR distribution according to the setting of Theorem 2.4. In particular, we fix $n=6d$, $\eta=m=20$ and set the total number of SGD steps as $T=\lceil\log(d)\rceil$. Note that the number of iterations required to reach perfect accuracy grows with $d$. The right side of Figure 4 provides further insight into the relationship between dimensionality and convergence rate. It displays the total number of SGD steps required to reach a test error below 0.01 for different values of $d$ using $n=3d$, $m=\eta=20$. The results are averages over five independent experiments and highlight the logarithmic dependence of the total SGD iterations on data dimension for achieving near-zero test error.

## 4 Conclusions and Future Directions

We explored the convergence and generalization of smooth neural networks trained with gradient methods. Our first goal in this paper was to derive generalization bounds through a new stability-based approach which had not been discussed in the vast literature of deep learning. These findings represent an improvement over previous results that either required substantial over-parameterization or provided suboptimal generalization rates that depended on the network width. For general noisy data distributions, we also derived generalization guarantees showing that GD can reach the optimal test loss. We also showed that orders of magnitude improvements in sample and computational complexity are possible by surpassing the NTK limitations and using SGD with large step-size. Several directions remain open to future research:

- $\bullet$
It remains open to explore whether the minimum width conditions of Theorems 2.1-2.3 can be improved in terms of $\gamma$ or $n$.
- $\bullet$
It is also interesting to extend the XOR analysis to the noisy setting where a fraction of the data points have corrupted labels.
- $\bullet$
The feature learning phenomenon in multi-index classification tasks remains largely unexplored. While we believe the XOR setup is a must-take first step, extensions to other multi-index models can shed further light on the strengths and limitations of neural networks.
- $\bullet$
We also aim to understand the potential benefits of network depth in either the NTK regime or feature learning i.e., whether adding a single layer can improve sample complexity or reduce the total number of SGD iterations.

## Appendix

## Appendix A Preliminaries : Hessian and Gradient norm

### A.1 Hessian norm

It is essential to our analysis to obtain Hessian’s spectral norm throughout the iterates of GD. With our weights’ initialization in place, one can bound the Hessian’s spectral norm based on the euclidean distance of the weights from initialization and network’s size as in the following assumption.

###### Assumption 2 (Hessian’s Spectral Norm) .

The spectral norm of the model’s Hessian for any $w\in\mathbb{R}^{p},\rho>0$ such that $\|w-w_{0}\|\leq\rho$ is bounded as $\|\nabla^{2}\Phi(w,\cdot)\|_{2}\leq\beta_{L}\frac{\rho^{3L}}{\sqrt{m}}$ for some constant $\beta_{L}$ that depends only on $L$.

It can be shown [, Lemma 1] that for standard Gaussian initialization and bounded data, Assumption 2 holds with high-probability over initialization for $\beta_{L}=e^{O(L)}$.

For the spectral norm of the model’s Hessian which is needed by our analysis, we use the following bound which bounds the spectral norm based on the distance from initialization.

###### Lemma 1 (Hessian’s spectral norm [ Liu et al., 2020 ] ) .

Assume the defined $L$ layer neural network architecture with standard Gaussian initialization i.e., $w_{{}_{0,{ij}}}^{(\ell)}\mathrel{\overset{\text{\small{iid}}}{\scalebox{1.5}[1
.0]{$\sim$}}}N(0,1)$. Let $\|w-w_{0}\|\leq\rho$. Then with high probability

$$ $\displaystyle\|\nabla^{2}\Phi(w,x)\|_{2}\leq\beta_{L}\frac{\rho^{3L}}{\sqrt{m}}$ $$

for some constant $\beta_{L}$ that only depends on $L.$

As mentioned in the main body, the result above guarantees that Assumption 2 holds with high probability.

###### Assumption 2 (Hessian’s Spectral Norm) .

###### Lemma 1 (Hessian’s spectral norm [ Liu et al., 2020 ] ) .

### A.2 Gradient norm at initialization

Recall the model for Theorems 2.1-2.3,

$$ $\displaystyle\Phi(w,x)$ $\displaystyle:=\frac{1}{\sqrt{m}}\langle W^{(L+1)},x^{(L)}\rangle$ $\displaystyle x_{i}^{(\ell+1)}$ $\displaystyle:=\frac{1}{\sqrt{m}}\sigma(\langle W_{i}^{(\ell+1)},x^{(\ell)} \rangle)\;\;\;\forall\ell\in\{1,\cdots,L-1\},\forall i\in\{1,\cdots,m\}$ $\displaystyle x_{i}^{(1)}$ $\displaystyle:=\sigma(\langle x,W^{(1)}_{i}\rangle)\;\;\;\;\;\forall i\in\{1, \cdots,m\}$ $$

Let us consider the standard Gaussian initialization. We have

$$ $\displaystyle\frac{\partial\Phi(w,x)}{\partial W^{(\ell+1)}_{ij}}=\frac{1}{ \sqrt{m}}x^{(\ell)}_{j}\sigma^{\prime}(\langle W_{i}^{(\ell+1)},x^{(\ell)} \rangle)\frac{\partial\Phi(w,x)}{\partial x^{(\ell+1)}_{i}}$ (13) $$

Therefore,

$$ $\displaystyle\|\frac{\partial\Phi(w,x)}{\partial W^{(\ell+1)}}\|^{2}=\sum_{ij} \|\frac{\partial\Phi(w,x)}{\partial W^{(\ell+1)}_{ij}}\|^{2}\leq\frac{1}{m}\|x ^{(\ell)}\|^{2}\|\frac{\partial\Phi(w,x)}{\partial x^{(\ell+1)}}\|^{2}$ (14) $$

by [, Lemmas F.3 and F.5] :

$$ $\displaystyle\|x^{(\ell)}\|\leq c_{0}^{\ell}\sqrt{m},~{}~{}~{}~{}\|\frac{ \partial\Phi}{\partial x^{(\ell)}}\|\leq c_{0}^{L-\ell+1}$ (15) $$

Therefore $\|\frac{\partial\Phi(w_{0},x)}{\partial W^{(\ell+1)}}\|^{2}\leq c_{0}^{2L}$ and

$$ $\|\nabla\Phi(w_{0},x)\|=(\sum_{\ell=0}^{L}\|\frac{\partial\Phi(w_{0},x)}{ \partial W^{(\ell+1)}}\|^{2})^{1/2}\leq\sqrt{L+1}c_{0}^{L}.$ $$

where $c_{0}$ is any constant that bounds the spectral norm of initialization weights with high probability, i.e., $\frac{\|W_{0}\|_{2}}{\sqrt{m}}\leq c_{0}.$ Specially, well-known bounds on the spectral norm of Gaussian matrices guarantee that $c_{0}$ is constant with high proabability [].

Moreover, based on [], if $\|w-w_{0}\|=\rho$ the spectral norm is bounded by,

$$ $\displaystyle\|\nabla^{2}\Phi(w,x)\|_{2}\leq\beta_{L}\frac{\rho^{3L}}{\sqrt{m}}$ (16) $$

where $\beta_{L}$ is a depth-only dependent variable derived to be $\beta_{L}=e^{O(L)}$.

Now, we use the above relation on Hessian to obtain a bound on the gradient:

$$ $\displaystyle\frac{\|\nabla\Phi(w^{*},x)-\nabla\Phi(w_{0},x)\|}{\|w^{*}-w_{0} \|}\leq\max_{w\in[w_{0},w^{*}]}\|\nabla^{2}\Phi(w,x)\|_{2}$ (17) $$

therefore if $w^{\star}\in\mathcal{B}_{\rho}(w_{0})$,

$$ $\displaystyle\|\nabla\Phi(w^{*},x)-\nabla\Phi(w_{0},x)\|\leq\frac{\beta_{L} \rho^{3L+1}}{\sqrt{m}}.$ (18) $$

Thus,

$$ $\displaystyle\|\nabla\Phi(w^{*},x)\|-\|\nabla\Phi(w_{0},x)\|\leq\|\nabla\Phi(w ^{*},x)-\nabla\Phi(w_{0},x)\|\leq\frac{\beta_{L}\rho^{3L+1}}{\sqrt{m}}.$ (19) $$

Thus w.h.p.

$$ $\displaystyle\max_{i\in[n]}\|\nabla\Phi(w^{*},x_{i})\|\leq\max_{i\in[n]}\| \nabla\Phi(w_{0},x_{i})\|+\frac{\beta_{L}\rho^{3L+1}}{\sqrt{m}}$ (20) $$

Let us denote the gradient and hessian bounds on $\Phi$ with $C_{1}$ and $C_{2}$, respectively. Note that based on our derivations above it holds that $C_{1}\leq\frac{\beta_{L}\rho^{3L+1}}{\sqrt{m}}+G_{0}$
and $C_{2}\leq\frac{\beta_{L}\rho^{3L}}{\sqrt{m}},$ where $G_{0}$ is the Lipschtiz parameter at initialization. With these parameters we can now obtain the Lipschitz and Smoothness parameter as well as the smallest eigenvalue of the objective along the gradient descent path. This is derived in the next lemma.

### A.3 Objective’s Gradient and Hessian along the GD path

###### Lemma 2 .

The following are true for the loss gradient and Hessian:

- 1.
$\|\nabla\widehat{F}(w)\|\leq C_{1}\,\widehat{F}(w).$
- 2.
$\|\nabla^{2}\widehat{F}(w)\|_{2}\leq(C_{2}+(C_{1})^{2})\widehat{F}(w).$
- 3.
$\lambda_{\min}\left(\nabla^{2}\widehat{F}(w)\right)\geq-C_{2}\widehat{F}(w)$.

where we define $C_{1}:=G_{0}+\beta_{L}\|w-w_{0}\|^{3L+1}/\sqrt{m},$ $C_{2}:=\beta_{L}\|w-w_{0}\|^{3L}/\sqrt{m}$ and $G_{0}$ is the Lipschitz parameter of the network at initialization i.e., $\|\nabla\Phi(w_{0},\cdot)\|\leq G_{0}.$

###### Proof.

The loss gradient is derived as follows,

$$ $\displaystyle\nabla\widehat{F}(w)$ $\displaystyle=\frac{1}{n}\sum_{i=1}^{n}f^{\prime}(y_{i}\Phi(w,x_{i}))y_{i} \nabla\Phi(w,x_{i})$ $$

Recalling that $y_{i}\in\{\pm 1\}$, we can write

$$ $\displaystyle\Big{\|}\nabla\widehat{F}(w)\Big{\|}$ $\displaystyle=\frac{1}{n}\Big{\|}\sum_{i=1}^{n}f^{\prime}(y_{i}\Phi(w,x_{i}))y _{i}\nabla\Phi(w,x_{i})\Big{\|}$ $\displaystyle\leq\frac{1}{n}\sum_{i=1}^{n}|f^{\prime}(y_{i}\Phi(w,x_{i}))|\Big {\|}\nabla\Phi(w,x_{i})\Big{\|}.$ $\displaystyle\leq C_{1}\,\widehat{F}(w).$ (21) $$

For the Hessian of loss, note that

$$ $\displaystyle\nabla^{2}\widehat{F}(w)$ $\displaystyle=\frac{1}{n}\sum_{i=1}^{n}f^{\prime\prime}(y_{i}\Phi(w,x_{i})) \nabla\Phi(w,x_{i})\nabla\Phi(w,x_{i})^{\top}+f^{\prime}(y_{i}\Phi(w,x_{i}))y_ {i}\nabla^{2}\Phi(w,x_{i}).$ (22) $$

It follows that

$$ $\displaystyle\Big{\|}\nabla^{2}\widehat{F}(w)\Big{\|}_{2}$ $\displaystyle=\left\|\frac{1}{n}\sum_{i=1}^{n}f^{\prime}(y_{i}\Phi(w,x_{i}))y_ {i}\nabla^{2}\Phi(w,x_{i})+f^{\prime\prime}(y_{i}\Phi(w,x_{i}))\nabla\Phi(w,x_ {i})\nabla_{1}\Phi(w,x_{i})^{\top}\right\|$ $\displaystyle\leq\frac{1}{n}\sum_{i=1}^{n}|f^{\prime}(y_{i}\Phi(w,x_{i}))|\Big {\|}\nabla^{2}\Phi(w,x_{i})\Big{\|}+|f^{\prime\prime}(y_{i}\Phi(w,x_{i}))|\Big {\|}\nabla\Phi(w,x_{i})\nabla\Phi(w,x_{i})^{\top}\Big{\|}$ $\displaystyle\leq\frac{1}{n}\sum_{i=1}^{n}|f^{\prime}(y_{i}\Phi(w,x_{i}))|\Big {\|}\nabla^{2}\Phi(w,x_{i})\Big{\|}+|f^{\prime\prime}(y_{i}\Phi(w,x_{i}))|\Big {\|}\nabla\Phi(w,x_{i})\Big{\|}^{2}$ $\displaystyle\leq C_{2}\widehat{F}(w)+(C_{1})^{2}\widehat{F}(w).$ (23) $$

To lower-bound the minimum eigenvalue of Hessian, note that the loss function $f$ is convex and thus $f^{\prime\prime}(\cdot)\geq 0$. Therefore the first term in Eq. 22 is positive semi-definite and the second term can be lower-bounded as follows,

$$ $\displaystyle\lambda_{\min}(\nabla^{2}\widehat{F}(w))$ $\displaystyle\geq-\left\|\frac{1}{n}\sum_{i=1}^{n}y_{i}f^{\prime}(y_{i}\Phi(w, x_{i}))\nabla^{2}\Phi(w,x_{i})\right\|$ $\displaystyle\geq-\frac{1}{n}\sum_{i=1}^{n}|y_{i}f^{\prime}(y_{i}\Phi(w,x_{i}) )|\Big{\|}\nabla^{2}\Phi(w,x_{i})\Big{\|}$ $\displaystyle\geq-C_{2}\widehat{F}(w).$ $$

∎

###### Lemma 2 .

###### Proof.

## Appendix B Proof of Theorem 2.1

With the bounds on Hessian and Gradient from the last section, we are ready to prove Theorem 2.1 for obtaining the training and generalization rates of deep nets. We start with proving the training guarantees in the following theorem.

### B.1 Proof for training loss

###### Theorem B.1 (Convergence of GD in deep nets) .

Let the descent lemma hold. Fix a $T$ and assume the target weights vector $w^{\star}\in\mathbb{R}^{p}$ that obtain small training loss such that

$$ $\displaystyle\rho^{\star}\geq\max\left\{\sqrt{\eta T\widehat{F}(w^{\star})}, \sqrt{\eta\widehat{F}(w_{0})}\right\}.$ $$

where $\rho^{\star}:=\|w^{\star}-w_{0}\|$. Moreover, assume the width $m$ is large enough such that it satisfies,

$$ $\displaystyle m\geq 4\beta_{L}^{2}(3{\rho^{\star}})^{6L+4}.$ $$

The the training loss satisfies with high probability over initialization,

$$ $\displaystyle\frac{1}{T}\sum_{t=1}^{T}\widehat{F}\left(w_{t}\right)\leq 2 \widehat{F}(w^{\star})+\frac{2{\rho^{\star}}^{2}}{\eta T}.$ $$

To begin the proof, we first state our descent lemma.

###### Lemma 3 (Descent lemma) .

Consider the notation as in Lemma 2. Assume for all $t\in[T]$ we have $\eta\leq 1/((C_{1})^{2}+C_{2})$, then

$$ $\widehat{F}\left(w_{t+1}\right)\leq\widehat{F}\left(w_{t}\right)-\frac{\eta}{2 }\left\|\nabla\widehat{F}\left(w_{t}\right)\right\|^{2}.$ $$

###### Proof.

By Taylor expansion and the bound on smoothness parameter of the objective, we immediately deduce:

$$ $\displaystyle\widehat{F}\left(w_{t+1}\right)$ $\displaystyle\leq\widehat{F}\left(w_{t}\right)-\eta\left\|\nabla\widehat{F} \left(w_{t}\right)\right\|^{2}+\eta^{2}\frac{(C_{1})^{2}+C_{2}}{2}\left\| \nabla\widehat{F}\left(w_{t}\right)\right\|^{2}$ $\displaystyle\leq\widehat{F}\left(w_{t}\right)-\frac{\eta}{2}\left\|\nabla \widehat{F}\left(w_{t}\right)\right\|^{2}$ $$

This completes the proof of the lemma.
∎

Define $\rho(t):=\|w_{t}-w_{0}\|$ and $\rho^{*}=\|w^{\star}-w_{0}\|$. Also, define $\widetilde{\rho}$ as any constant such that

$$ $\tilde{\rho}>\max\{\rho(T-1),\rho(T-2),\cdots,\rho(1),\rho^{\star}\}.$ $$

Then for any $w\in[w_{t},w^{\star}]$ such that $w=\alpha w_{t}+(1-\alpha)w^{\star}$ it holds that

$$ $\|w-w_{0}\|\leq\|\alpha w_{t}+(1-\alpha)w^{\star}-w_{0}\|\leq\alpha\|w_{t}-w_{ 0}\|+(1-\alpha)\|w^{\star}-w_{0}\|\leq\tilde{\rho},$ $$

and thus for any $w\in[w_{t},w^{\star}]$

$$ $\displaystyle\lambda_{\min}\left(\nabla^{2}\widehat{F}(w)\right)\geq- \widetilde{C}\widehat{F}(w)$ (24) $$

where $\widetilde{C}:=\frac{\beta_{L}\tilde{\rho}^{3L}}{\sqrt{m}}$, i.e., the smoothness parameter with respect to the distance $\tilde{\rho}$ and as computed by Lemma 2. Therefore, by Taylor remainder theorem, for any fixed $w^{\star}$ there exists $w^{\prime}_{t}\in[w_{t},w^{\star}]$ such that .

$$ $\displaystyle\widehat{F}(w^{\star})$ $\displaystyle\geq\widehat{F}(w_{t})+\left\langle\nabla\widehat{F}(w_{t}),w^{ \star}-w_{t}\right\rangle+\frac{1}{2}\lambda_{\min}\left(\nabla^{2}\widehat{F} (w_{t}^{\prime})\right)\|w^{\star}-w_{t}\|^{2}$ $\displaystyle\geq\widehat{F}(w_{t})+\left\langle\nabla\widehat{F}(w_{t}),w^{ \star}-w_{t}\right\rangle-\frac{1}{2}\frac{\beta_{L}\tilde{\rho}^{3L}}{{\sqrt{ m}}}\widehat{F}(w_{t}^{\prime})\,\|w^{\star}-w_{t}\|^{2}.$ $$

As a direct result of the above inequality, we have derived that for any $w^{\star}$,

$$ $\displaystyle\widehat{F}(w^{\star})\geq\widehat{F}(w_{t})+\left\langle\nabla \widehat{F}(w_{t}),w^{\star}-w_{t}\right\rangle-\frac{1}{2}\frac{\beta_{L} \tilde{\rho}^{3L}}{{\sqrt{m}}}\max_{w^{\prime}_{t}\in[w_{t},w^{\star}]} \widehat{F}(w_{t}^{\prime})\,\|w-w_{t}\|^{2}.$ $$

Putting this in the relation from descent lemma (Lemma 3) followed by completion of squares using $w_{t+1}-w_{t}=-\eta\nabla\widehat{F}(w_{t})$ gives

$$ $\displaystyle\widehat{F}(w_{t+1})$ $\displaystyle\leq\widehat{F}(w^{\star})-\left\langle\nabla\widehat{F}(w_{t}),w ^{\star}-w_{t}\right\rangle-\frac{\eta}{2}\left\|\nabla\widehat{F}(w_{t}) \right\|^{2}+\frac{1}{2}\frac{\beta_{L}\tilde{\rho}^{3L}}{{\sqrt{m}}}\max_{w^{ \prime}_{t}\in[w_{t},w^{\star}]}\widehat{F}(w_{t}^{\prime})\,\|w^{\star}-w_{t} \|^{2}$ $\displaystyle=\widehat{F}(w^{\star})+\frac{1}{\eta}\left(\|w^{\star}-w_{t}\|^{ 2}-\|w^{\star}-w_{t+1}\|^{2}\right)+\frac{1}{2}\frac{\beta_{L}\tilde{\rho}^{3L }}{{\sqrt{m}}}\max_{w^{\prime}_{t}\in[w_{t},w^{\star}]}\widehat{F}(w_{t}^{ \prime})\,\|w^{\star}-w_{t}\|^{2}.$ (25) $$

Telescoping the above display for $t=0,\ldots,{T-1}$, we deduce that,

$$ $\displaystyle\frac{1}{T}\sum_{t=1}^{T}\widehat{F}\left(w_{t}\right)\leq \widehat{F}(w^{\star})+\frac{\left\|w^{\star}-w_{0}\right\|^{2}}{\eta T}+\frac {\beta_{L}\tilde{\rho}^{3L}}{2T}\sum_{t=0}^{T-1}\max_{w^{\prime}_{t}\in[w_{t}, w^{\star}]}\widehat{F}\left(w_{t}^{\prime}\right)\left\|w^{\star}-w_{t}\right \|^{2}.$ (26) $$

In order to proceed with the above relation, it is necessary to further simplify the terms $\max_{w^{\prime}_{t}\in[w_{t},w^{\star}]}\widehat{F}\left(w_{t}^{\prime}\right)$. In doing so, we utilize the following Generalized Local Quasi-convexity (GLQC) property from [].

###### Proposition 1 (Prop. 8 in [ Taheri and Thrampoulidis, 2024 ] ) .

Let $w_{1},w_{2}\in\mathbb{R}^{p}$ be two arbitrary points. Suppose $\widehat{F}:\mathbb{R}^{p}\rightarrow\mathbb{R}$ satisfies the self-bounded weak convexity property with parameter $\kappa$ along the line $[w_{1},w_{2}]$. Assume $\left\|w_{1}-w_{2}\right\|\leq D$ where $D<\sqrt{2/\kappa}$. Then,

$$ $\displaystyle\max_{v\in[w_{1},w_{2}]}\widehat{F}(v)\leq\tau\cdot\max\{\widehat {F}(w_{1}),\widehat{F}(w_{2})\}.$ (27) $$

where $\tau:=\left(1-\kappa D^{2}/2\right)^{-1}$.

For our objective we have the minimum eigenvalue of hessian being lower bounded by $\kappa=\widetilde{C}:=\beta_{L}\tilde{\rho}^{3L}$. Therefore for $w_{t},w^{\star}$, if $\|w^{\star}-w_{t}\|^{2}<D^{2}<1/2\widetilde{C}$, then the GLQC property holds with $\tau=4/3$. Recall that by the defnition of $\tilde{\rho}$ it holds that $\|w^{\star}-w_{t}\|\leq\tilde{\rho}$. Thus, to guarantee the conditions of the proposition above are satisfied it suffices to hold that

$$ $\displaystyle\sqrt{m}\geq 2\beta_{L}\tilde{\rho}^{3L+2}.$ (28) $$

which is exactly the width condition of our theorem. With this condition on width and the resulting simplifications from the proposition, we can write:

$$ $\displaystyle\frac{1}{T}\sum_{t=1}^{T}\widehat{F}\left(w_{t}\right)$ $\displaystyle\leq\widehat{F}(w^{\star})+\frac{\left\|w^{\star}-w_{0}\right\|^{ 2}}{\eta T}+\frac{\beta_{L}\tilde{\rho}^{3L}}{\sqrt{m}}\frac{2}{3T}\sum_{t=0}^ {T-1}\max\{\widehat{F}\left(w^{\star}\right),\widehat{F}(w_{t})\}\left\|w^{ \star}-w_{t}\right\|^{2},$ $\displaystyle\leq\widehat{F}(w^{\star})+\frac{\left\|w^{\star}-w_{0}\right\|^{ 2}}{\eta T}+\frac{1}{3T}\sum_{t=0}^{T-1}\max\{\widehat{F}\left(w^{\star}\right ),\widehat{F}(w_{t})\}$ $\displaystyle\leq\frac{4}{3}\widehat{F}(w^{\star})+\frac{\left\|w^{\star}-w_{0 }\right\|^{2}}{\eta T}+\frac{1}{3T}\sum_{t=0}^{T-1}\widehat{F}(w_{t})$ $$

Thus,

$$ $\displaystyle\frac{1}{T}\sum_{t=1}^{T}\widehat{F}\left(w_{t}\right)\leq 2 \widehat{F}(w^{\star})+\frac{3\left\|w^{\star}-w_{0}\right\|^{2}}{2\eta T}+ \frac{1}{2T}\widehat{F}(w_{0})$ (29) $$

Thus far, we have characterized the convergence under the condition that $\sqrt{m}\geq 2\beta_{L}\tilde{\rho}^{3L+2}=2\beta_{L}(\max\{\|w_{T-1}-w_{0}\|,
\cdots,\|w_{1}-w_{0}\|,\|w^{\star}-w_{0}\|\})^{3L+2}.$ It remains to characterize the above condition on width independent of the iteration number. In fact, in what follows we are able to show it is sufficient to reduce this condition such that it depends on only $\|w^{\star}-w_{0}\|$. We do it by an iterate-wise norm bound on $\|w^{\star}-w_{t}\|$ based on $\|w^{\star}-w_{0}\|$.

Recall that we had,

$$ $\displaystyle\widehat{F}(w_{t+1})\leq\widehat{F}(w^{\star})+\frac{1}{\eta} \left(\|w^{\star}-w_{t}\|^{2}-\|w^{\star}-w_{t+1}\|^{2}\right)+\frac{1}{2} \widetilde{C}\max_{w^{\prime}_{t}\in[w_{t},w^{\star}]}\widehat{F}(w_{t}^{ \prime})\,\|w^{\star}-w_{t}\|^{2}.$ (30) $$

where recall that $\widetilde{C}=\frac{\beta_{L}\tilde{\rho}^{3L}}{\sqrt{m}},$ and $\widetilde{\rho}:=\max\{\rho(1),\rho(2),...\rho(T-1),\rho^{\star}\}$.

Hence,

$$ $\displaystyle\|w_{t+1}-w^{\star}\|^{2}\leq\|w_{t}-w^{\star}\|^{2}+\eta\widehat {F}(w^{\star})+\frac{\eta}{2}\frac{\beta_{L}\tilde{\rho}^{3L}}{\sqrt{m}}\max_{ w^{\prime}_{t}\in[w_{t},w^{\star}]}\widehat{F}(w_{t}^{\prime})\,\|w^{\star}-w_ {t}\|^{2}.$ (31) $$

The following lemma proves through an induction-based argument that choosing $\tilde{\rho}=3\rho^{\star}$ and the corresponding width condition based on this $\tilde{\rho}$ is sufficient to ensure that $\rho(t)\leq\tilde{\rho}$ for all iterates $t\in[T]$; which in turn guarantees the above argument for convergence to be valid.

###### Lemma 4 .

Fix any $T\geq 0$
and assume any $w^{\star}$ such that

$$ $\displaystyle\|w^{\star}-w_{0}\|^{2}\geq\max\{\eta T\widehat{F}(w^{\star}), \eta\widehat{F}(w_{0})\}.$ (32) $$

Pick $\widetilde{\rho}=3\rho^{\star}=3\|w^{\star}-w_{0}\|$ and assume $m$ is large enough such that $\widetilde{\rho}^{2}<1/2\widetilde{C}$ where $\widetilde{C}=\frac{\beta_{L}\tilde{\rho}^{3L}}{\sqrt{m}}$. Then $\rho(t):=\|w_{t}-w_{0}\|\leq\widetilde{\rho}$ and $\|w_{t}-w^{\star}\|\leq\widetilde{\rho}$ for all $t\in[T]$.

###### Proof.

The statements of the lemma are correct for $t=0$. Now assume for $t\leq T-1$ we have $\rho(t)\leq\widetilde{\rho}$ and $\|w_{t}-w^{\star}\|\leq\widetilde{\rho}.$ Recall, by 31

$$ $\displaystyle\|w_{T}-w^{\star}\|^{2}$ $\displaystyle\leq\|w_{T-1}-w^{\star}\|^{2}+\eta\widehat{F}(w^{\star})+\frac{ \eta}{2}\frac{\beta_{L}\tilde{\rho}^{3L}}{\sqrt{m}}\|w_{T-1}-w^{\star}\|^{2} \max_{w^{\prime}_{T-1}\in[w_{T-1},w^{\star}]}\widehat{F}(w^{\prime}_{T-1})$ (33) $$

Note that by the assumption of the lemma and the induction assumption it holds that

$$ $\max\{\|w_{T-1}-w_{0}\|,\|w^{\star}-w_{0}\|\}\leq\widetilde{\rho}.$ $$

As a result, for any $w\in[w_{T-1},w^{\star}]$, it holds that $\|w-w^{\star}\|\leq\tilde{\rho}$ and thus $\|\nabla\widehat{F}(w)\|\geq\frac{\beta_{L}\tilde{\rho}^{3L}}{\sqrt{m}}
\widehat{F}(w)$. Moreover, note that as per the lemma’s assumption $\|w_{T-1}-w^{\star}\|^{2}\leq\widetilde{\rho}^{2}\leq 1/2\widetilde{C}$. Thus, the conditions of Proposition 1 hold with $\tau=4/3$ and we have

$$ $\displaystyle\|w_{T}-w^{\star}\|^{2}$ $\displaystyle\leq\|w_{T-1}-w^{\star}\|^{2}+\eta\widehat{F}(w^{\star})+\frac{2 \eta}{3}\frac{\beta_{L}\tilde{\rho}^{3L}}{\sqrt{m}}\|w_{T-1}-w^{\star}\|^{2} \cdot\max\{\widehat{F}(w_{T-1}),\widehat{F}(w^{\star})\}$ $$

Noting again that $\|w_{T-1}-w^{\star}\|^{2}\leq\widetilde{\rho}^{2}\leq 1/2\widetilde{C}$ yields,

$$ $\displaystyle\|w_{T}-w^{\star}\|^{2}$ $\displaystyle\leq\|w_{T-1}-w^{\star}\|^{2}+\eta\widehat{F}(w^{\star})+\frac{ \eta}{3}\max\{\widehat{F}(w_{T-1}),\widehat{F}(w^{\star})\}$ $\displaystyle\leq\|w_{T-1}-w^{\star}\|^{2}+\frac{4\eta}{3}\widehat{F}(w^{\star })+\frac{\eta}{3}\widehat{F}(w_{T-1})$ $$

Since the conditions of the lemma hold for all $t\in[T-1]$, we can telescope the summation to obtain,

$$ $\displaystyle\|w_{T}-w^{\star}\|^{2}$ $\displaystyle\leq\|w_{0}-w^{\star}\|^{2}+\frac{4\eta}{3}T\widehat{F}(w^{\star} )+\frac{\eta}{3}\sum_{t=0}^{T-1}\widehat{F}(w_{t})$ $\displaystyle\leq\|w_{0}-w^{\star}\|^{2}+\frac{4\eta}{3}T\widehat{F}(w^{\star} )+\frac{\eta}{3}\widehat{F}(w_{0})+\frac{\eta}{3}\sum_{t=1}^{T-1}\widehat{F}(w _{t})$ $\displaystyle\leq\|w_{0}-w^{\star}\|^{2}+\frac{4\eta}{3}T\widehat{F}(w^{\star} )+\frac{\eta}{3}\widehat{F}(w_{0})+\frac{\eta}{3}\Big{(}2T\widehat{F}(w^{\star })+\frac{3\left\|w^{\star}-w_{0}\right\|^{2}}{2\eta}+\frac{1}{2}\widehat{F}(w_ {0})\Big{)}$ $\displaystyle=\|w_{0}-w^{\star}\|^{2}+2\eta T\widehat{F}(w^{\star})+\frac{\eta }{2}\widehat{F}(w_{0})+\frac{\|w^{\star}-w_{0}\|^{2}}{2}$ $$

Therefore by the conditions of the lemma and noting that $A_{0}=\|w_{0}-w^{\star}\|^{2}={(\rho^{\star})}^{2}$:

$$ $\displaystyle\|w_{T}-w^{\star}\|^{2}\leq\|w_{0}-w^{\star}\|^{2}+2\|w_{0}-w^{ \star}\|^{2}+\frac{1}{2}\|w^{\star}-w_{0}\|^{2}+\frac{1}{2}\|w^{\star}-w_{0}\| ^{2}=4(\rho^{\star})^{2}$ $$

Thus, $\|w_{T}-w^{\star}\|^{2}\leq 9(\rho^{\star})^{2}=\widetilde{\rho}^{2}.$ It remains to prove that $\rho(T):=\|w_{T}-w_{0}\|\leq\widetilde{\rho}$. To obtain this, note that by triangle inequality,

$$ $\displaystyle\|w_{T}-w_{0}\|$ $\displaystyle\leq\|w_{T}-w^{\star}\|+\|w^{\star}-w_{0}\|$ $\displaystyle\leq 2\rho^{\star}+\rho^{\star}=3\rho^{\star}=\widetilde{\rho}.$ $$

The proof is complete.
∎

The lemma above shows that controlling $\rho^{\star}=\|w^{\star}-w_{0}\|$ is enough to guarantee that GLQC property (Proposition 1) holds throughout the gradient path. In particular, if the width is large enough such that for $\widetilde{\rho}=3\rho^{\star}$: $\widetilde{\rho}^{2}\leq\frac{1}{2\widetilde{C}}=\frac{\sqrt{m}}{2\beta_{L}
\tilde{\rho}^{3L}}$
or equivalently

$$ $\sqrt{m}\geq 2\beta_{L}\tilde{\rho}^{3L+2}.$ $$

then the conditions of the lemma hold and the bounds on the training loss are valid. Note that we need to ensure the conditions of the descent lemma (Lem 3) hold throughout the GD iterates. This is guaranteed by taking

$$ $\displaystyle\eta\leq\frac{1}{(\widetilde{C}_{1})^{2}+\tilde{C}}$ (34) $$

where note that $\tilde{C}_{1}=G_{0}+\beta_{L}\frac{(3\|w^{\star}-w_{0}\|)^{3L+1}}{\sqrt{m}},$ and $\tilde{C}=\beta_{L}\frac{(3\|w^{\star}-w_{0}\|)^{3L}}{\sqrt{m}}$. As per Lemma 2, this condition ensures that step-size is smaller than the inverse of objective’s smoothness parameter throughout all iterates since $\|w_{t}-w_{0}\|\leq 3\|w^{\star}-w_{0}\|$ for all $t\in[T]$.

This completes the proof of Theorem B.1.

###### Theorem B.1 (Convergence of GD in deep nets) .

###### Lemma 3 (Descent lemma) .

###### Proof.

###### Proposition 1 (Prop. 8 in [ Taheri and Thrampoulidis, 2024 ] ) .

###### Lemma 4 .

###### Proof.

### B.2 Proof for Generalization loss

###### Theorem B.2 (Generalization loss of GD in deep nets) .

Let the descent lemma hold. Fix a $T$ and assume the target weights vector $w^{\star}\in\mathbb{R}^{p}$ that separates the data distribution such that

$$ $\displaystyle\rho^{\star}\geq\max\left\{\sqrt{\eta T\widehat{F}(w^{\star})}, \sqrt{\eta\widehat{F}(w_{0})}\right\}.$ (35) $$

where $\rho^{\star}:=\|w^{\star}-w_{0}\|$. Moreover, assume the width $m$ is large enough such that it satisfies,

$$ $\displaystyle\sqrt{m}\geq 4\beta_{L}\bar{\rho}^{3L+2}$ (36) $$

where $\bar{\rho}=6\rho^{\star}$.
Then the expected generalization error satisfies with high probability,

$$ $\displaystyle\mathbb{E}_{\mathcal{S}}\Big{[}{F}(w_{T})-\widehat{F}(w_{T})\Big{ ]}\leq 9\frac{(G_{0}+1/4)^{2}{\rho^{\star}}^{2}}{n},$ (37) $$

where $G_{0}$ is the Lipschitz parameter of network at initialization i.e., $\|\nabla\Phi(w_{0},\cdot)\|\leq G_{0}.$

As per the last theorem, we know if $m$ is large enough based on $\rho^{\star}$, then the hessian during the iterates of GD satisfies $\|\nabla^{2}\Phi(w,x)\|\leq C_{2}$ and thus as per Lemma 2 we have,

- 1.
$\|\nabla\widehat{F}(w)\|_{F}\leq C_{1}\,\widehat{F}(w).$
- 2.
$\|\nabla^{2}\widehat{F}(w)\|_{2}\leq(C_{2}+(C_{1})^{2})\widehat{F}(w).$
- 3.
$\lambda_{\min}\left(\nabla^{2}\widehat{F}(w)\right)\geq-C_{2}\widehat{F}(w)$.

throughout this section and with a slight abuse of notation, we define $C_{1}=G_{0}+\beta_{L}(3\|w-w_{0}\|)^{3L+1}/\sqrt{m},$ $C_{2}=\beta_{L}(3\|w-w_{0}\|)^{3L}/\sqrt{m}$ and $G_{0}$ is the Lipschitz parameter of the network at initialization i.e., $\|\nabla\Phi(w_{0},\cdot)\|\leq G_{0}.$ Recall that in obtaining these quantities we used the guarantees of Thm 2.1 that along the gradient descent path $\max\{\|w_{t}-w_{0}\|,\|w_{t}-w^{\star}\|\}\leq 3\|w^{\star}-w_{0}\|$

We derive the generalization bound through leave-one-out error estimates. In particular, let $w_{t}^{\neg i}$ denote the output of $t$-th iteration of GD on the leave-one-out loss $F^{\neg i}(\cdot):\mathbb{R}^{p}\rightarrow\mathbb{R}$ defined as $F^{\neg i}(w):=$ $\frac{1}{n}\sum_{j\neq i}\widehat{F}_{j}(w)$. Then, a simple calculation based on the Lipschitz property of the objective relates the generalization loss to the leave-one-out error as stated in the following lemma.

###### Lemma 5 (Lem. B.3 in [ Lei and Ying, 2020 ] ) .

Suppose each sample loss $f(\cdot,z)$ is $G$-Lipschitz for almost surely all data points $z\sim\mathcal{D}$. Then, the following relation holds between expected generalization loss and model stability at any iterate $T$,

$$ $\displaystyle\mathbb{E}_{\mathcal{S}}\Big{[}{F}(w_{T})-\widehat{F}(w_{T})\Big{ ]}\leq 2G\cdot\;\mathbb{E}_{\mathcal{S}}\Big{[}\frac{1}{n}\sum_{i=1}^{n}\|w_{T }-w_{T}^{\neg i}\|\Big{]}.$ (38) $$

We will use the fact that under the descent lemma $\eta<1/(C_{2}+C_{1}^{2})$, based on [, Lem. 9]:

$$ $\displaystyle\left\|\left(w-\eta\nabla\widehat{F}(w)\right)-\left(w^{\prime}- \eta\nabla\widehat{F}\left(w^{\prime}\right)\right)\right\|\leq\max_{\tilde{w} \in[w,w^{\prime}]}(1+\eta C_{2}\widehat{F}(\tilde{w}))\left\|w-w^{\prime} \right\|.$ (39) $$

Then, we can bound the leave-one-out error as follows:

$$ $\displaystyle\left\|w_{t+1}-w_{t+1}^{\neg i}\right\|$ $\displaystyle\leq\left\|\left(w_{t}-\eta\nabla\widehat{F}^{\neg i}\left(w_{t} \right)\right)-\left(w_{t}^{\neg i}-\eta\nabla\widehat{F}^{\neg i}\left(w_{t}^ {\neg i}\right)\right)\right\|+\frac{\eta}{n}\left\|\nabla\widehat{F}_{i}\left (w_{t}\right)\right\|$ $\displaystyle\leq\left\|\left(w_{t}-\eta\nabla\widehat{F}^{\neg i}\left(w_{t} \right)\right)-\left(w_{t}^{\neg i}-\eta\nabla\widehat{F}^{\neg i}\left(w_{t}^ {\neg i}\right)\right)\right\|+\frac{\eta C_{1}}{n}\widehat{F}_{i}\left(w_{t}\right)$ $\displaystyle\leq\left(1+\eta C_{2}\max_{w_{t}^{\prime}\in[w_{t},w_{t}^{\neg i }]}\widehat{F}^{\neg i}\left(w^{\prime}_{t}\right)\right)\left\|w_{t}-w_{t}^{ \neg i}\right\|+\frac{\eta C_{1}}{n}\widehat{F}_{i}\left(w_{t}\right),$ $$

By the aforementioned property from Proposition 1, we have:

$$ $\displaystyle\max_{w_{t}^{\prime}\in[w_{t},w_{t}^{\neg i}]}\widehat{F}^{\neg i }\left(w^{\prime}_{t}\right)$ $\displaystyle\leq\frac{4}{3}\max\{\widehat{F}^{{\neg i}}(w_{t}),\widehat{F}^{{ \neg i}}(w_{t}^{\neg i})\}$ $\displaystyle\leq\frac{4}{3}\max\{\widehat{F}(w_{t}),\widehat{F}^{{\neg i}}(w_ {t}^{\neg i})\}$ $\displaystyle\leq\frac{4}{3}(\widehat{F}(w_{t})+\widehat{F}^{{\neg i}}(w_{t}^{ \neg i}))$ $$

Note that for $w_{t}^{^{\prime}}$ we have

$$ $\displaystyle\|w_{t}^{\prime}-w_{0}\|\leq\alpha\|w_{t}-w_{0}\|+(1-\alpha)\|w_{ t}^{\neg i}-w_{0}\|\leq\widetilde{\rho}$ (40) $$

therefore the self-bounded weak convexity holds with $\kappa=C_{2}$.
Moreover,

$$ $\|w_{t}-w_{t}^{\neg i}\|\leq\|w_{t}-w_{0}\|+\|w_{t}^{\neg i}-w_{0}\|\leq 3\rho ^{\star}+3\rho^{\star}=6\rho^{\star}=2\widetilde{\rho}$ $$

Denote $\bar{\rho}:=6\rho^{\star}(=2\widetilde{\rho})$.
If $m$ is large enough such that $\bar{\rho}^{2}<1/2\bar{C}_{2}$ where recall $\bar{C}_{2}=\frac{\beta_{L}\bar{\rho}^{3L}}{\sqrt{m}},$
then the GLQC property holds with $\tau=4/3$ and we have,

$$ $\displaystyle\left\|w_{t+1}-w_{t+1}^{\neg i}\right\|$ $\displaystyle\leq\left(1+\frac{4}{3}\eta\bar{C}_{2}(\widehat{F}(w_{t})+ \widehat{F}^{{\neg i}}(w_{t}^{\neg i}))\right)+\frac{\eta C_{1}}{n}\widehat{F} _{i}\left(w_{t}\right)$ $\displaystyle\leq\exp(\frac{4}{3}\eta\bar{C}_{2}(\widehat{F}(w_{t})+\widehat{F }^{{\neg i}}(w_{t}^{\neg i})))+\frac{\eta C_{1}}{n}\widehat{F}_{i}\left(w_{t}\right)$ $\displaystyle\leq\frac{\eta C_{1}}{n}\sum_{k=0}^{t}\exp\left(\sum_{\tau=1}^{t} \frac{4}{3}\eta\bar{C}_{2}(\widehat{F}(w_{\tau})+\widehat{F}^{{\neg i}}(w_{ \tau}^{\neg i}))\right)\widehat{F}_{i}(w_{k})$ $$

Fix a $T$ and assume $\bar{C}_{2}$ is small enough such that

$$ $\displaystyle\frac{4}{3}\eta\bar{C}_{2}\sum_{\tau=1}^{T}(\widehat{F}(w_{\tau}) +\widehat{F}^{{\neg i}}(w_{\tau}^{\neg i}))\leq\frac{16}{3}\bar{C}_{2}(\eta T \widehat{F}(w^{\star})+{\rho^{\star}}^{2})\leq 0.05$ (41) $$

then,

$$ $\displaystyle\left\|w_{T}-w_{T}^{\neg i}\right\|$ $\displaystyle\leq\frac{\eta C_{1}\exp(0.05)}{n}\sum_{t=0}^{T}\widehat{F}_{i}(w _{t})$ $\displaystyle\leq 1.1\frac{\eta C_{1}}{n}\sum_{t=0}^{T}\widehat{F}_{i}(w_{t}).$ $$

The choice of 0.05 in Eq. 41 is to ensure that the resulting constant in the generalization bound is close to 1. By averaging over $i\in[n]$ and by our training loss guarantees from Theorem 2.1:

$$ $\displaystyle\frac{1}{n}\sum_{i=1}^{n}\left\|w_{T}-w_{T}^{\neg i}\right\|$ $\displaystyle\leq 1.1\frac{\eta C_{1}}{n}\sum_{t=0}^{T}\widehat{F}(w_{t})$ $\displaystyle\leq 2.2\frac{C_{1}}{n}(\eta T\widehat{F}(w^{\star})+{\rho^{\star }}^{2})$ $$

based on Lemma 5 and noting that $G=C_{1}$, we find that

$$ $\displaystyle\mathbb{E}_{\mathcal{S}}\Big{[}{F}(w_{T})-\widehat{F}(w_{T})\Big{ ]}\leq\frac{4.4}{n}\,\mathbb{E}_{\mathcal{S}}\left[C_{1}^{2}\eta T\widehat{F}( w^{\star})+C_{1}^{2}{\rho^{\star}}^{2}\right]$ (42) $$

Recall from Lemma 2 that $C_{1}\leq G_{0}+\beta_{L}\|w_{t}-w_{0}\|^{3L+1}/{\sqrt{m}}$ for all $t\in[T]$. As we require $\sqrt{m}\geq 4\beta_{L}(6\|w^{\star}-w_{0}\|^{3L+2})$ and from convergence guarantees it holds that $\|w_{t}-w_{0}\|\leq 3\|w^{\star}-w_{0}\|$ we conclude that

$$ $C_{1}\leq G_{0}+\frac{1}{2^{3L+2}\|w^{\star}-w_{0}\|}\leq G_{0}+1/4,$ $$

where in the last step we assumed without loss of generality that $\|w^{\star}-w_{0}\|\geq 1.$ This completes the proof for the generalization error in Theorem B.2. The proof of Theorem 2.1 is complete.

###### Theorem B.2 (Generalization loss of GD in deep nets) .

###### Lemma 5 (Lem. B.3 in [ Lei and Ying, 2020 ] ) .

## Appendix C Proof of Corollary 1

###### Corollary 2 (Restatement of Corollary 1 ) .

Assume the tangent kernel of the model at initialization separates the data with margin $\gamma$ i.e., for a unit-norm $w\in\mathbb{R}^{p}$ it holds for all $i\in[n]:$
$y_{i}\left\langle\nabla\Phi\left(w_{0},x_{i}\right),w\right\rangle\geq\gamma$. Define constant $B>0$ that bounds the model’s output at initialization i.e., $\forall i\in[n]:|\Phi(w_{0},x_{i})|<B$. Assume the width is large enough such that $m\geq\beta_{L}^{2}(\frac{2B+\log(1/{\epsilon})}{\gamma})^{6L+4}$. Then there exists $w^{\star}\in\mathbb{R}^{p}$ such that $F(w^{\star})\leq{\epsilon}$ and $w^{\star}$ lies at a Euclidean distance of order $O(1/\gamma)$ from initialization where $\|w^{\star}-w_{0}\|\leq\frac{1}{\gamma}(2B+\log(1/{\epsilon}))$.

###### Proof of Corollary 1 .

By Taylor there exists $w^{\prime}\in\left[w^{\star},w_{0}\right]$ such that,

$$ $y_{i}\Phi\left(w^{\star},x_{i}\right)=y_{i}\Phi\left(w_{0},x_{i}\right)+y_{i} \left\langle\nabla\Phi\left(w_{0},x_{i}\right),w^{\star}-w_{0}\right\rangle+ \frac{1}{2}y_{i}\left\langle w^{\star}-w_{0},\nabla^{2}\Phi\left(w^{\prime},x_ {i}\right)\left(w^{\star}-w_{0}\right)\right\rangle$ $$

.

Pick $w^{\star}=w_{0}+\frac{w}{\gamma}(2B+\log(1/\varepsilon))$. Since $\left\|w\right\|=1$, we derive the desired for $\left\|w^{\star}-w_{0}\right\|$. we have $\left\|\nabla^{2}\Phi\left(w^{\prime},x_{i}\right)\right\|\leq O(\frac{\rho^{3
L}}{\sqrt{m}})$. Choosing $\rho=\|w^{\star}-w_{0}\|=\frac{1}{\gamma}(2B+\log(1/{\epsilon}))$ and noting the given condition on $m$ implies $m\geq\beta_{L}^{2}\rho^{6L+4}$ we obtain $\frac{\beta_{L}\rho^{3L+2}}{\sqrt{m}}\leq 1$.
Therefore with the given assumption on $m\geq\frac{\beta_{L}^{2}}{\gamma^{6L+4}}(2B+\log(1/{\epsilon}))^{6L+4}$ leads to $\frac{1}{2}\left\|\nabla^{2}\Phi\left(w^{\prime},x_{i}\right)\right\|\left\|w^
{\star}-w_{0}\right\|^{2}\leq 1.$ Then,

$$ $\displaystyle y_{i}\Phi\left(w^{\star},x_{i}\right)$ $\displaystyle\geq-\left|y_{i}\Phi\left(w_{0},x_{i}\right)\right|+y_{i}\left \langle\nabla_{1}\Phi\left(w_{0},x_{i}\right),w^{\star}-w_{0}\right\rangle- \frac{1}{2}\left\|\nabla_{1}^{2}\Phi\left(w^{\prime},x_{i}\right)\right\|\left \|w^{\star}-w_{0}\right\|^{2}$ $\displaystyle\geq-B+2B+\log(1/\varepsilon)-B$ $\displaystyle=\log(1/\varepsilon)$ $$

where we assumed without loss of generality that $B\geq 1.$
After applying the logistic loss function on $y_{i}\Phi(w^{\star},x_{i})$, the inequality above gives $\widehat{F}(w)\leq\varepsilon$. This completes the proof.
∎

###### Corollary 2 (Restatement of Corollary 1 ) .

###### Proof of Corollary 1 .

## Appendix D Consistency of GD in deep nets: Proof of Theorem 2.3

###### Theorem D.1 (Consistency) .

Assume the width of the network satisfies $m\geq\beta_{L}^{2}n^{3L+3}$ and the step-size satisfies $\eta<1/L_{F}$. Then with high probability the expected test loss at iteration $T\geq\sqrt{n}$ is bounded as:

$$ $\displaystyle\mathbb{E}_{\mathcal{S}}\left[F(w_{T})\right]\leq\left(1+\frac{4} {\sqrt{n}}\right)\left(F(w^{\star})+\frac{{\rho^{\star}}^{2}}{\eta\sqrt{n}}+ \frac{{\rho^{\star}}^{2}}{n\sqrt{n}}+\frac{1}{\sqrt{n}}\right)$ (43) $$

where $w^{\star}$ is the minimizer of the population loss i.e., $w^{\star}=\arg\min_{w\in\mathbb{R}^{p}}F(w)$.

Recall that based on 39

$$ $\left\|(w-\eta\nabla\widehat{F}(w))-\left(w^{\prime}-\eta\nabla\widehat{F} \left(w^{\prime}\right)\right)\right\|\leq 1+\eta\max_{\alpha}C_{2}\left\|w-w^ {\prime}\right\|,$ (44) $$

where $C_{2}:=O(\frac{\beta_{L}\rho^{3L}}{\sqrt{m}})$ and $\rho:=\|w-w_{0}\|$

In particular, note that for GD, we can obtain for $\|w-w_{0}\|$:

$$ $\displaystyle\|w_{T}-w_{0}\|$ $\displaystyle\leq\eta\sum_{t=0}^{T-1}\|\nabla\widehat{F}(w_{t})\|$ $\displaystyle\leq\eta\sum_{t=0}^{T-1}C_{1}$ $\displaystyle=\eta TC_{1},$ (45) $$

where the first step follows by update rule of GD and in the second step we used the Lipschitz property of objective during GD iterates that ensures $\|\nabla\widehat{F}(w_{t})\|$ remains bounded by $C_{1}$ for all $t\in[T]$.
Thus, we can thus show by induction that $\|w_{T}-w_{0}\|\leq\eta TC_{1}$. Then for $m=\Omega(T^{2})$ and $w_{T}^{\neg i},w_{T}:$

$$ $\displaystyle\left\|(w_{T}-\eta\nabla\widehat{F}(w_{T}))-\left(w^{\neg i}_{T}- \eta\nabla\widehat{F}\left(w_{T}^{\neg i}\right)\right)\right\|\leq(1+\eta C_{ 2})\left\|w_{T}-w_{T}^{\neg i}\right\|,$ (46) $$

where $C_{2}=\beta_{L}\tilde{\rho}^{3L}/\sqrt{m}$ is the smoothness parameter of the objective throughout the GD updates $w_{t}$ for all $t$. Then,

$$ $\displaystyle\left\|w_{t+1}-w_{t+1}^{\neg i}\right\|$ $\displaystyle\leq\left\|\left(w_{t}-\eta\nabla\widehat{F}^{\neg i}\left(w_{t} \right)\right)-\left(w_{t}^{\neg i}-\eta\nabla\widehat{F}^{\neg i}\left(w_{t}^ {\neg i}\right)\right)\right\|+\frac{\eta}{n}\left\|\nabla\widehat{F}_{i}\left (w_{t}\right)\right\|$ $\displaystyle\leq\left\|\left(w_{t}-\eta\nabla\widehat{F}^{\neg i}\left(w_{t} \right)\right)-\left(w_{t}^{\neg i}-\eta\nabla\widehat{F}^{\neg i}\left(w_{t}^ {\neg i}\right)\right)\right\|+\frac{\eta C_{1}}{n}\widehat{F}_{i}\left(w_{t}\right)$ $\displaystyle\leq\left(1+\eta C_{2}\right)\left\|w_{t}-w_{t}^{\neg i}\right\|+ \frac{\eta C_{1}}{n}\widehat{F}_{i}\left(w_{t}\right),$ $$

Recall that in the above $\widehat{F}_{i}(\cdot)$ refers to the objective at the $i$th data sample and thus has the same Lipschits constant as $\widehat{F}(w)=\frac{1}{n}\sum_{i=1}^{n}\widehat{F}_{i}(w)$.

Noting that $1+x\leq e^{x}$ for $x\geq 0$, we can proceed with the above inequality as follows by unrolling the iterates of $\|w_{t}-w_{t}^{\neg i}\|$ and the fact that $w_{0}=w_{0}^{\neg i}$:

$$ $\displaystyle\left\|w_{t+1}-w_{t+1}^{\neg i}\right\|$ $\displaystyle\leq\exp(\eta C_{2})\|w_{t}-w_{t}^{\neg i}\|+\frac{\eta C_{1}}{n} \widehat{F}_{i}\left(w_{t}\right)$ $\displaystyle\leq\frac{\eta C_{1}}{n}\sum_{k=0}^{t}\exp(\eta C_{2}t)\widehat{F }_{i}(w_{k}).$ $$

Recall that, $C_{2}=\beta_{L}\tilde{\rho}^{3L}/\sqrt{m}$, and $\tilde{\rho}=\max\{\|w_{t}-w_{0}\|,\|w_{t-1}-w_{0}\|,\cdots,\|w_{1}-w_{0}\|\}$.
Based on Eq. 45, $\|w_{t}-w_{0}\|\leq\eta tC_{1}$ which implies $\tilde{\rho}\leq\eta tC_{1}.$ As a result, for some fixed $T$:

$$ $\displaystyle\|w_{T}-w_{T}^{\neg i}\|\leq\frac{\eta C_{1}}{n}\sum_{t=0}^{T-1} \exp\left(\frac{\eta\beta_{L}T}{\sqrt{m}}(\eta TC_{1})^{3L}\right)\widehat{F}_ {i}(w_{t}).$ $$

Thus for $m$ larger enough such that $\sqrt{m}\geq 2\beta_{L}(\eta C_{1}T)^{3L+1}$, the exponential term is smaller than $2$ and we conclude that,

$$ $\displaystyle\left\|w_{T}-w_{T}^{\neg i}\right\|$ $\displaystyle\leq\frac{2\eta C_{1}}{n}\sum_{t=0}^{T-1}\widehat{F}_{i}(w_{t}).$ (47) $$

Therefore, by averaging over $i\in[n]$ we derive the on-average leave-one-out:

$$ $\displaystyle\frac{1}{n}\sum_{i}\left\|w_{T}-w_{T}^{\neg i}\right\|\leq\frac{2 \eta C_{1}}{n}\sum_{t=0}^{T-1}\widehat{F}(w_{t})$ (48) $$

The only remaining part is characterising the right hand side of the above inequality. This is done as follows.

###### Theorem D.1 (Consistency) .

### Proof of optimization error

Define $\rho(t):=\|w_{t}-w_{0}\|$ and $\rho^{*}=\|w^{\star}-w_{0}\|$. if we define $\widetilde{\rho}>\max\{\rho(t),\rho^{\star}\}$ for all $t\in[T]$ then for any $w\in[w_{t},w^{\star}]$ it holds that $\|w-w_{0}\|\leq\widetilde{\rho}$ and thus based on Lemma 2

$$ $\displaystyle\lambda_{\min}(\nabla^{2}\widehat{F}(w))\geq-\widetilde{C} \widehat{F}(w),$ (49) $$

where $\widetilde{C}:=\frac{\beta_{L}\tilde{\rho}^{3L}}{\sqrt{m}}$.
Therefore by Eq. 26,

$$ $\displaystyle\widehat{F}(w_{T})\leq\frac{1}{T}\sum_{t=1}^{T}\widehat{F}\left(w _{t}\right)$ $\displaystyle\leq\widehat{F}(w^{\star})+\frac{\left\|w^{\star}-w_{0}\right\|^{ 2}}{\eta T}+\widetilde{C}\frac{1}{2T}\sum_{t=0}^{T-1}\max_{w^{\prime}_{t}\in[w _{t},w^{\star}]}\widehat{F}\left(w_{t}^{\prime}\right)\left\|w^{\star}-w_{t} \right\|^{2},$ $\displaystyle\leq\widehat{F}(w^{\star})+\frac{\left\|w^{\star}-w_{0}\right\|^{ 2}}{\eta T}+\widetilde{C}\frac{1}{2T}\sum_{t=0}^{T-1}\left\|w^{\star}-w_{t} \right\|^{2}$ $\displaystyle\leq\widehat{F}(w^{\star})+\frac{\left\|w^{\star}-w_{0}\right\|^{ 2}}{\eta T}+\widetilde{C}\left\|w^{\star}-w_{0}\right\|^{2}+\widetilde{C}\frac {1}{T}\sum_{t=0}^{T-1}\|w_{t}-w_{0}\|^{2}.$ $$

For the last term, we have a simple bound:

$$ $\displaystyle\|w_{t}-w_{0}\|=\eta\|\sum_{k=0}^{t-1}\nabla\widehat{F}(w_{k})\| \leq\eta\sum_{k=0}^{t-1}\|\nabla\widehat{F}(w_{k})\|\leq\eta tC_{1}$ $$

where $C_{1}$ denote the lipschitz parameter bound along the GD iterates and the last inequality follows by descent lemma condition.

Then,

$$ $\displaystyle\widehat{F}(w_{T})\leq\widehat{F}(w^{\star})+\frac{\left\|w^{ \star}-w_{0}\right\|^{2}}{\eta T}+\widetilde{C}\left\|w^{\star}-w_{0}\right\|^ {2}+\eta^{2}\widetilde{C}C_{1}^{2}T^{2}.$ $$

Pick large enough width such that $m\geq\beta_{L}^{2}n^{3L+3}(\eta C_{1})^{6L}$ and choose $T=\sqrt{n}$. Then $\tilde{\rho}<(\eta C_{1}n)^{3L/2}$ and

$$ $\displaystyle\widehat{F}(w_{T})$ $\displaystyle\leq\widehat{F}(w^{\star})+\frac{{\rho^{\star}}^{2}}{\eta\sqrt{n} }+\frac{\beta_{L}\tilde{\rho}^{3L}}{\sqrt{m}}({\rho^{\star}}^{2}+\eta^{2}C_{1} ^{2}n)$ $\displaystyle\leq\widehat{F}(w^{\star})+\frac{{\rho^{\star}}^{2}}{\eta\sqrt{n} }+\frac{{\rho^{\star}}^{2}+\eta^{2}C_{1}^{2}n}{n\sqrt{n}}$ $$

In particular, in the above one can choose $w^{\star}:=\arg\min_{w\in\mathbb{R}^{p}}F(w)$. Moreover, from the generalization guarantees,

$$ $\displaystyle\mathbb{E}_{\mathcal{S}}\left[F(w_{T})\right]$ $\displaystyle\leq\mathbb{E}_{\mathcal{S}}[\widehat{F}(w_{T})]+\frac{4\eta C_{1 }^{2}}{n}\sum_{t=1}^{T-1}\mathbb{E}_{\mathcal{S}}[\widehat{F}(w_{t})]$ $\displaystyle\leq\left(1+\frac{4\eta C_{1}^{2}}{\sqrt{n}}\right)\left(\mathbb{ E}_{\mathcal{S}}[\widehat{F}(w^{\star})]+\frac{{\rho^{\star}}^{2}}{\eta\sqrt{n }}+\frac{{\rho^{\star}}^{2}+\eta^{2}C_{1}^{2}n}{n\sqrt{n}}\right)$ $\displaystyle=\left(1+\frac{4\eta C_{1}^{2}}{\sqrt{n}}\right)\left(F(w^{\star} )+\frac{{\rho^{\star}}^{2}}{\eta\sqrt{n}}+\frac{{\rho^{\star}}^{2}+\eta^{2}C_{ 1}^{2}n}{n\sqrt{n}}\right).$ $$

Note that $\eta\leq\frac{1}{C_{1}^{2}+C_{2}}$ by the descent lemma condition (Lemma 3). Therefore
$\eta\leq 1/C_{1}^{2}.$ Without loss of generality we can assume $C_{1}\geq 1$. This simplifies the above inequality as follows:

$$ $\displaystyle\mathbb{E}_{\mathcal{S}}\left[F(w_{T})\right]$ $\displaystyle\leq\left(1+\frac{4}{\sqrt{n}}\right)\left(F(w^{\star})+\frac{{ \rho^{\star}}^{2}}{\eta\sqrt{n}}+\frac{{\rho^{\star}}^{2}}{n\sqrt{n}}+\frac{1} {\sqrt{n}}\right)$ $\displaystyle=F(w^{\star})+O\left(\frac{F(w^{\star})+{\rho^{\star}}^{2}}{\sqrt {n}}\right).$ $$

This completes the proof of Theorem 2.3.

## Appendix E Proof of Theorem 2.4

We recall our setup and notation for Theorem 2.4. We consider in this section the one-hidden layer network with quadratic activation where $x\in\mathbb{R}^{d},w_{i}\in\mathbb{R}^{d}$:

$$ $\displaystyle\Phi(w,x)=\frac{1}{2m}\sum_{i=1}^{m}a_{i}(x^{\top}w_{i})^{2}$ $$

in the above, $a_{i}=\pm 1$ are fixed and half of the neurons have $a_{i}=1$. For initialization of first layer weights $w_{i}\mathrel{\overset{\text{\small{iid}}}{\scalebox{1.5}[1.0]{$\sim$}}}N(0,
\frac{I_{d}}{d})$ that is each $w_{ij}\mathrel{\overset{\text{\small{iid}}}{\scalebox{1.5}[1.0]{$\sim$}}}N(0,
\frac{1}{d})$. The data points $x$ are uniformly drawn from the Rademacher distribution (of size $2^{d}$) that is $x(k)=\pm 1$ i.i.d and $w.p.\;1/2$ and the labels $y=x(1)\cdot x(2).$ Consider the linear loss where $f(u)=-u$ and $\widehat{F}(w):=\frac{1}{n}\sum_{j=1}^{n}f(y_{j}\Phi(x_{j},w))$. Consider stochastic gradient descent with the update rule as follows for each $w_{i}$:

$$ $\displaystyle w_{i}^{t+1}$ $\displaystyle=w_{i}^{t}-\eta\nabla\widehat{F}_{i}(w^{t})$ $\displaystyle=w_{i}^{t}-\frac{\eta a_{i}}{nm}\sum_{j=1}^{n}(x_{j}^{\top}w_{i}^ {t})x_{j}y_{j}f^{\prime}(y_{j}\Phi(x_{j},w))$ $\displaystyle=w_{i}^{t}+\frac{\eta a_{i}}{nm}\sum_{j=1}^{n}(x_{j}^{\top}w_{i}^ {t})x_{j}y_{j}$ $$

Note that in the above the superscripts denote the iteration number and the subscript $i$ indicates the vector $w_{i}\in\mathbb{R}^{d}$ is associated with neuron $i.$
Fixing the initialization, define $\mathbb{E}_{x}[w_{i}^{t}]$ as the outcome of GD after $t$ iterations on *population gradient*. Formally, $\mathbb{E}_{x}[w_{i}^{t}]$ is defined recursively as follows:

$$ $\displaystyle\mathbb{E}_{x}[w_{i}^{t+1}]:=\mathbb{E}_{x}[w_{i}^{t}]-\eta\nabla F _{i}(\mathbb{E}_{x}[w^{t}]),$ $$

where $F(w):=\mathbb{E}_{x}[\widehat{F}(w)]$ is the test loss and define $\mathbb{E}_{x}[w^{0}]:=w^{0}$. Based on previous derivations the update rule for the expected weights $\mathbb{E}_{x}[W_{i}^{t}]$ obeys the following:

$$ $\displaystyle\mathbb{E}_{x}[w_{i}^{t+1}]=\mathbb{E}_{x}[w_{i}^{t}]+\frac{\eta} {m}a_{i}\mathbb{E}_{x}\left[\langle x,\mathbb{E}_{x}[w_{i}^{t}]\rangle xy \right].$ $$

for $t=0:$

$$ $\displaystyle\mathbb{E}_{x}[w_{i}^{1}]$ $\displaystyle=w_{i}^{0}+\frac{\eta}{m}a_{i}\mathbb{E}_{x}\left[\langle x,w_{i} ^{0}\rangle xy\right].$ $$

For $t=1:$

$$ $\displaystyle\mathbb{E}_{x}[w_{i}^{2}]=\mathbb{E}_{x}[w_{i}^{1}]+\frac{\eta}{m }a_{i}\mathbb{E}_{x}\left[\langle x,\mathbb{E}_{x}[w_{i}^{1}]\rangle xy\right].$ $$

and so on.

### E.1 Calculating 𝔼 x ⁢ [ w i 1 ] subscript 𝔼 𝑥 delimited-[] superscript subscript 𝑤 𝑖 1 \mathbb{E}_{x}[w_{i}^{1}] blackboard_E start_POSTSUBSCRIPT italic_x end_POSTSUBSCRIPT [ italic_w start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT start_POSTSUPERSCRIPT 1 end_POSTSUPERSCRIPT ]

For the expected gradient $g_{i}\in\mathbb{R}^{d}$ we have by noting $y=x(1)x(2)$:

$$ $\displaystyle g_{i}=-\mathbb{E}_{x}[\langle x,w_{i}^{0}\rangle\cdot x\cdot x(1 )x(2)]$ $$

for $g_{i}^{3}$ (we drop $i$ and 0 for simplicity and denote the $k^{\prime}$th element of $W_{i}^{0}$ with $W^{k}$ ):

$$ $\displaystyle-g_{i}^{3}$ $\displaystyle=\mathbb{E}_{x}[(x(3)w^{3}+\sum_{k\neq 3}x(k)w^{k})\cdot x(3) \cdot x(1)x(2)]$ $\displaystyle=\mathbb{E}_{x}[w^{3}\cdot x(1)x(2)+x(3)x(1)x(2)\sum_{k\neq 3}x(k )w^{k}]$ $\displaystyle=\mathbb{E}_{x}[w^{3}\cdot x(1)x(2)]+\mathbb{E}_{x}[x(3)x(1)x(2) \sum_{k\neq 3}x(k)w^{k}]$ $\displaystyle=0+\mathbb{E}_{x}[x(3)\sum_{k\neq 3}x(k)w^{k}]\,\mathbb{E}_{x}[x( 1)x(2)]$ $\displaystyle=0.$ $$

Therefore for any $k\neq 1,2$ we have $g_{i}^{k}=0.$

for $k=1$ :

$$ $\displaystyle-g_{i}^{1}$ $\displaystyle=\mathbb{E}_{x}[\langle x,w\rangle\cdot x(1)\cdot x(1)x(2)]$ $\displaystyle=\mathbb{E}_{x}[w^{1}x(1)x(2)+x(1)x(2)x(1)\sum_{k\neq 1}w^{k}x(k)]$ $\displaystyle=\mathbb{E}_{x}[{w^{1}}x(1)x(2)]+\mathbb{E}_{x}[x(1)x(2)x(1)w^{2} x(2)]+\mathbb{E}_{x}[x(1)x(2)x(1)\sum_{k\neq 1,2}w^{k}x(k)]$ $\displaystyle={w^{1}}\mathbb{E}_{x}[x(1)x(2)]+w^{2}\mathbb{E}_{x}[x(1)x(2)x(1) x(2)]+\mathbb{E}_{x}[x(1)x(2)x(1)]\cdot\mathbb{E}_{x}[\sum_{k\neq 1,2}w^{k}x(k)]$ $\displaystyle=w^{2}\mathbb{E}_{x}[(x(1)x(2))^{2}]$ $\displaystyle=w^{2}.$ $$

Similarly, we derive that

$$ $\displaystyle-g_{i}^{2}=w_{i}^{0}(1)$ $$

Thus, the expected weight after one step is

$$ $\displaystyle\mathbb{E}_{x}[w_{i}^{1}]$ $\displaystyle=w_{i}^{0}+\frac{\eta a_{i}}{m}[w_{i}^{0}(2),w_{i}^{0}(1),0, \cdots,0]$ $\displaystyle=\Big{[}w_{i}^{0}(1)+\frac{\eta a_{i}}{m}w_{i}^{0}(2),w_{i}^{0}(2 )+\frac{\eta a_{i}}{m}w_{i}^{0}(1),w_{i}^{0}(3),\cdots,w_{i}^{0}(d)\Big{]}$ $$

### E.2 Calculating 𝔼 x ⁢ [ w i 2 ] subscript 𝔼 𝑥 delimited-[] superscript subscript 𝑤 𝑖 2 \mathbb{E}_{x}[w_{i}^{2}] blackboard_E start_POSTSUBSCRIPT italic_x end_POSTSUBSCRIPT [ italic_w start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT ]

For $k=1:$

$$ $\displaystyle\mathbb{E}_{x}[w_{i}^{2}(1)]$ $\displaystyle=\mathbb{E}_{x}[w_{i}^{1}(1)]+\frac{\eta a_{i}}{m}\mathbb{E}_{x}[ \langle x,\mathbb{E}_{x}[w_{i}^{1}]\rangle x(2)]$ $\displaystyle=\mathbb{E}_{x}[w_{i}^{1}(1)]+\frac{\eta a_{i}}{m}\mathbb{E}_{x}[ w_{i}^{1}(2)]$ $\displaystyle=w_{i}^{0}(1)+\frac{\eta a_{i}}{m}w_{i}^{0}(2)+\frac{\eta a_{i}}{ m}(w_{i}^{0}(2)+\frac{\eta a_{i}}{m}w_{i}^{0}(1))$ $$

Similarly we obtain for $k=2:$

$$ $\displaystyle\mathbb{E}_{x}[w_{i}^{2}(2)]$ $\displaystyle=\mathbb{E}_{x}[w_{i}^{1}(2)]+\frac{\eta a_{i}}{m}\mathbb{E}_{x}[ \langle x,\mathbb{E}_{x}[w_{i}^{1}]\rangle x(1)]$ $\displaystyle=\mathbb{E}_{x}[w_{i}^{1}(2)]+\frac{\eta a_{i}}{m}\mathbb{E}_{x}[ w_{i}^{1}(1)]$ $\displaystyle=w_{i}^{0}(2)+\frac{\eta a_{i}}{m}w_{i}^{0}(1)+\frac{\eta a_{i}}{ m}(w_{i}^{0}(1)+\frac{\eta a_{i}}{m}w_{i}^{0}(2))$ $$

for $k\neq 1,2$:

$$ $\displaystyle\mathbb{E}_{x}[w_{i}^{2}(k)]$ $\displaystyle=\mathbb{E}_{x}[w_{i}^{1}(k)]+\frac{\eta a_{i}}{m}\mathbb{E}_{x}[ \langle x,\mathbb{E}_{x}[w_{i}^{1}]\rangle x(k)x(1)x(2)]$ $\displaystyle=\mathbb{E}_{x}[w_{i}^{1}(k)]$ $$

### E.3 Calculating 𝔼 x ⁢ [ w i t ] subscript 𝔼 𝑥 delimited-[] superscript subscript 𝑤 𝑖 𝑡 \mathbb{E}_{x}[w_{i}^{t}] blackboard_E start_POSTSUBSCRIPT italic_x end_POSTSUBSCRIPT [ italic_w start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_t end_POSTSUPERSCRIPT ]

In general, we observe the following pattern for arbitrary $t$:
for $k\neq 1,2:$

$$ $\displaystyle\mathbb{E}_{x}[w_{i}^{t+1}(k)]=\mathbb{E}_{x}[w_{i}^{t}(k)]$ $$

and for $k=1,2$:

$$ $\displaystyle\mathbb{E}_{x}[w_{i}^{t+1}(1)]$ $\displaystyle=\mathbb{E}_{x}[w_{i}^{t}(1)]+\frac{\eta a_{i}}{m}\mathbb{E}_{x}[ w_{i}^{t}(2)]$ $\displaystyle\mathbb{E}_{x}[w_{i}^{t+1}(2)]$ $\displaystyle=\mathbb{E}_{x}[w_{i}^{t}(2)]+\frac{\eta a_{i}}{m}\mathbb{E}_{x}[ w_{i}^{t}(1)]$ $$

Therefore for general $t$, defining $\gamma=\eta a_{i}/m$, $\alpha(t):=\mathbb{E}_{x}[w_{i}^{t}(1)]$ and $\beta(t):=\mathbb{E}_{x}[w_{i}^{t}(2)]$ we have the following equations:

$$ $\displaystyle\alpha(t+1)$ $\displaystyle=\alpha(t)+\gamma\beta(t)$ $\displaystyle\beta(t+1)$ $\displaystyle=\beta(t)+\gamma\alpha(t).$ $$

It can be verified that the solution to the above equations take the following form:

$$ $\displaystyle\alpha(t)$ $\displaystyle=\left(\sum_{r\in[t],r:even}\binom{t}{r}\gamma^{r}\right)\alpha(0 )+\left(\sum_{r\in[t],r:odd}\binom{t}{r}\gamma^{r}\right)\beta(0)$ $\displaystyle\beta(t)$ $\displaystyle=\left(\sum_{r\in[t],r:even}\binom{t}{r}\gamma^{r}\right)\beta(0) +\left(\sum_{r\in[t],r:odd}\binom{t}{r}\gamma^{r}\right)\alpha(0)$ $$

Therefore with replacement:

$$ $\displaystyle\mathbb{E}_{x}[w_{i}^{t}(1)]$ $\displaystyle=\left(\sum_{r\in[t],r:even}\binom{t}{r}(\frac{\eta a_{i}}{m})^{r }\right)w_{i}^{0}(1)+\left(\sum_{r\in[t],r:odd}\binom{t}{r}(\frac{\eta a_{i}}{ m})^{r}\right)w_{i}^{0}(2)$ $\displaystyle\mathbb{E}_{x}[w_{i}^{t}(2)]$ $\displaystyle=\left(\sum_{r\in[t],r:even}\binom{t}{r}(\frac{\eta a_{i}}{m})^{r }\right)w_{i}^{0}(2)+\left(\sum_{r\in[t],r:odd}\binom{t}{r}(\frac{\eta a_{i}}{ m})^{r}\right)w_{i}^{0}(1)$ $$

### E.4 Empirical Gradient for multiple GD steps

The results above hold when GD is applied on the entire data distribution. Let us consider the case of SGD on $n$ data points (on the loss $F(\cdot)$) for each iteration and bound the resulting noise.

Recall that

$$ $\displaystyle w^{1}$ $\displaystyle=w^{0}-\eta\nabla\hat{F}(w^{0}),$ $\displaystyle\mathbb{E}_{x}[w^{1}]$ $\displaystyle=w^{0}-\eta\mathbb{E}_{x}[\nabla\hat{F}(w^{0})],$ $\displaystyle w^{1}$ $\displaystyle=\mathbb{E}_{x}[w^{1}]-\eta(\nabla\hat{F}(w^{0})-\mathbb{E}_{x}[ \nabla\hat{F}(w^{0})]).$ $$

Moreover, by how we defined $\mathbb{E}_{x}[w^{t}]$ for $t\geq 2$, we have

$$ $\displaystyle w^{t+1}$ $\displaystyle=w^{t}-\eta\nabla\hat{F}(w^{t}),$ $\displaystyle\mathbb{E}_{x}[w^{t+1}]$ $\displaystyle=\mathbb{E}_{x}[w^{t}]-\eta\mathbb{E}_{x}[\nabla\hat{F}(\mathbb{E }_{x}[w^{t}])],$ $\displaystyle w^{t+1}$ $\displaystyle=\mathbb{E}_{x}[w^{t}]+(w^{t}-\mathbb{E}_{x}[w^{t}])-\eta(\nabla \hat{F}(w^{t})-\mathbb{E}_{x}[\nabla\hat{F}(\mathbb{E}_{x}[w^{t}])])-\eta \mathbb{E}_{x}[\nabla\hat{F}(\mathbb{E}_{x}[w^{t}])].$ $$

Overall, from the last two equations we have:

$$ $\displaystyle w^{t+1}-\mathbb{E}_{x}[w^{t+1}]$ $\displaystyle=-\eta\sum_{\tau=0}^{t}\nabla\hat{F}(w^{\tau})-\mathbb{E}_{x} \left[\nabla\hat{F}(\mathbb{E}_{x}[w^{\tau}])\right]$ $\displaystyle=-\eta\sum_{\tau=0}^{t}\underbrace{\nabla\hat{F}(w^{\tau})-\nabla \hat{F}(\mathbb{E}_{x}[w^{\tau}])}_{\text{Term I}}+\underbrace{\nabla\hat{F}( \mathbb{E}_{x}[w^{\tau}])-\mathbb{E}_{x}[\nabla\hat{F}(\mathbb{E}_{x}[w^{\tau} ])]}_{\text{Term II}}.$ $$

We bound each term separately.

### E.5 Term II:

Starting with the second term, recall that:

$$ $\displaystyle-\nabla\hat{F}_{i}(\mathbb{E}_{x}[w^{t}])=\frac{a_{i}}{nm}\sum_{j =1}^{n}\langle x_{j},\mathbb{E}_{x}[w^{t}_{i}]\rangle x_{j}y_{j}$ $$

##### Case 1:

For $k\neq 1,2$ :

$$ $\displaystyle-\nabla\hat{F}_{i}^{k}(\mathbb{E}_{x}[w^{t}])$ $\displaystyle=\frac{a_{i}}{nm}\sum_{j=1}^{n}\langle x_{j},\mathbb{E}_{x}[w_{i} ^{t}]\rangle x_{j}(k)x_{j}(1)x_{j}(2)$ $\displaystyle=\frac{a_{i}}{nm}\left\langle\sum_{j=1}^{n}x_{j}\cdot x_{j}(k) \cdot x_{j}(1)\cdot x_{j}(2),\mathbb{E}_{x}[w_{i}^{t}]\right\rangle$ $$

Denoting $v_{j}:=x_{j}\cdot x_{j}(k)\cdot x_{j}(1)\cdot x_{j}(2)$ note that $v_{j}(r)$ are mutually independent for all $r\in d$ and $j\in[n]$. Moreover each $v_{j}(r)$ is distributed according to a Rademacher distribution and $Var(v_{j}(r))=1$.

Therefore, recalling the expression for $\mathbb{E}_{x}[w^{t}_{i}]=[\alpha_{2}^{t}w_{i}^{0}(1)+\alpha_{1}^{t}w_{i}^{0}(
2),\alpha_{1}^{t}w_{i}^{0}(1)+\alpha_{2}^{t}w_{i}^{0}(2),w_{i}^{0}(3),\cdots,w
_{i}^{0}(d)]$
where $\alpha_{2}^{t}:=\sum_{r\in[t],r:even}\binom{t}{r}(\frac{\eta a_{i}}{m})^{r}$ and $\alpha_{1}^{t}:=\sum_{r\in[t],r:odd}\binom{t}{r}(\frac{\eta a_{i}}{m})^{r}$
by expanding the summation above:

$$ $\displaystyle-\nabla F_{i}^{k}(\mathbb{E}_{x}[w^{t}])$ $\displaystyle=\frac{a_{i}}{nm}\Big{(}(\alpha_{2}^{t}\sum_{j}v_{j}(1)+\alpha_{1 }^{t}\sum_{j}v_{j}(2))w_{i}^{0}(1)$ $\displaystyle\hskip 36.135pt+(\alpha_{1}^{t}\sum_{j}v_{j}(1)+\alpha_{2}^{t} \sum_{j}v_{j}(2))w_{i}^{0}(2)+\sum_{k\neq 1,2}\sum_{j}v_{j}(k)w_{i}^{0}(k)\Big {)}.$ $$

We use Hoeffding concentration inequality for bounding $\sum_{j}v_{j}(k)=O(\sqrt{n}\log(d))$ uniformly for every $k\neq 1,2$. To be more concrete, by a union bound and Hoeffding it follows that uniformly over all $k\leq d$ and $t\leq T$:

$$ $\displaystyle\Pr\left(\sum_{j}v_{j}(k)\leq\sqrt{n}\log(dT)\right)\geq 1-\exp \left(\log(dT)-2\log^{2}(dT)\right)=:1-p_{x}.$ $$

Moreover, recall that $w_{i}^{0}(k)\sim N(0,1/d)$ and :

$$ $\displaystyle\Pr\left(w_{i}^{0}(k)\leq\frac{\log(dm)}{\sqrt{d}}\right)\geq 1- \frac{\exp(-\frac{1}{2}\log^{2}(dm))}{\sqrt{2\pi}\log(dm)}$ $$

By union bound we have for all $i\leq m,k\leq d$:

$$ $\displaystyle\Pr\left(w_{i}^{0}(k)\leq\frac{\log(dm)}{\sqrt{d}}\right)\geq 1- \frac{\exp(-\frac{1}{2}\log^{2}(dm)+\log(dm))}{\sqrt{2\pi}\log(dm)}=:1-p_{w}.$ $$

Therefore w.p. $1-p_{x}-p_{w}$ we have for all $t\leq T,i\leq m,k\leq d$: $\sum_{j}v_{j}(k)w_{i}^{0}(k)\leq\frac{\log(dm)\log(dT)\sqrt{n}}{\sqrt{d}}$
By Hoeffding w.p. $1-p_{x}-p_{w}$ over data sampling and initialization:

$$ $\displaystyle\Pr\left(\sum_{k}\sum_{j}v_{j}(k)w_{i}^{0}(k)\leq\log(dm)\log(dT) \log(d)\sqrt{n}\right)\geq 1-\exp(-\log^{2}(d))=:1-p_{1}.$ $$

Similarly, we obtain $\Pr\left(\sum_{j}v_{j}(1)w_{i}^{0}(1)\leq\frac{\sqrt{n}}{\sqrt{d}}\log(dT)\log
(dm)\right)\geq 1-p_{x}-p_{w}$
Overall, the above calculations show that for some constant $C$ with probability $1-C(p_{1}-p_{x}-p_{w})$ over initialization and data sampling it holds uniformly over all $k,i,t$:

$$ $\displaystyle|\nabla\hat{F}_{i}^{k}(\mathbb{E}_{x}[w^{t}])|$ $\displaystyle\lesssim\frac{\log(dT)\log(dm)}{nm}\left(\frac{2\sqrt{n}(|\alpha_ {1}^{t}|+\alpha_{2}^{t})}{\sqrt{d}}+\sqrt{n}\log(d)\right)$ $\displaystyle=\log(dT)\log(dm)\frac{|\alpha_{1}^{t}|+\alpha_{2}^{t}+\sqrt{d} \log(d)}{m\sqrt{nd}}.$ $$

##### Case 2:

Now assume $k=1$ then in a similar way as above

$$ $\displaystyle-\nabla\hat{F}_{i}^{1}(\mathbb{E}_{x}[w^{t}])$ $\displaystyle=\frac{a_{i}}{nm}\sum_{j=1}^{n}\langle x_{j},\mathbb{E}_{x}[w_{i} ^{t}]\rangle x_{j}(2)$ $\displaystyle=\frac{a_{i}}{m}\mathbb{E}_{x}[w_{i}^{t}(2)]+\frac{a_{i}}{nm}\sum _{j=1}^{n}\sum_{j^{\prime}\neq 2}^{d}\mathbb{E}_{x}[w_{i}^{t}(j^{\prime})]x_{j }(j^{\prime})x_{j}(2)$ $\displaystyle=\frac{a_{i}}{m}\mathbb{E}_{x}[w_{i}^{t}(2)]+O(\log(dm)\log(dT) \frac{|\alpha_{1}^{t}|+\alpha_{2}^{t}+\sqrt{d}\log(d)}{m\sqrt{nd}}).$ $$

for $k=2:$

$$ $\displaystyle-\nabla\hat{F}_{i}^{2}(\mathbb{E}_{x}[w^{t}])=\frac{a_{i}}{m} \mathbb{E}_{x}[w_{i}^{t}(1)]+O(\log(dm)\log(dT)\frac{|\alpha_{1}^{t}|+\alpha_{ 2}^{t}+\sqrt{d}\log(d)}{m\sqrt{nd}}).$ (50) $$

Therefore for the weights entering the $i$’th hidden neuron and for all $i\leq m,t\leq T$, the following concentration bound holds w.p. $1-C(p_{1}-p_{x}-p_{w})$:

$$ $\displaystyle\left\|\nabla\hat{F}_{i}(\mathbb{E}_{x}[w^{t}])-\mathbb{E}_{x}[ \nabla F_{i}(\mathbb{E}_{x}[w^{t}])]\right\|\lesssim\log(dm)\log(dT)\frac{| \alpha_{1}^{t}|+\alpha_{2}^{t}+\sqrt{d}\log(d)}{m\sqrt{n}}.$ $$

Therefore, as a result of Term II calculations, we deduce that with probability $1-C(p_{1}-p_{x}-p_{w})$,

$$ $\displaystyle\forall i,t:\left\|\nabla\hat{F}_{i}(\mathbb{E}_{x}[w^{t}])- \mathbb{E}_{x}[\nabla\hat{F}_{i}(\mathbb{E}_{x}[w^{t}])]\right\|\lesssim\log( dm)\log(dT)\frac{|\alpha_{1}^{t}|+\alpha_{2}^{t}+\sqrt{d}\log(d)}{m\sqrt{n}},$ $$

where $x^{t}$ denotes the data chosen at iteration $t$ of SGD and recall $w^{0}$ denotes the initialization weights.

### E.6 Term I:

##### Case 1:

We first consider $k\neq 1,2$. By applying Hoeffding’s concentration inequality we obtain $\sum_{j=1}^{n}x_{j}(\ell)x_{j}(k)x_{j}(1)x_{j}(2)<\sqrt{n}\log(dT)$ uniformly for every $\ell\in[d],t\leq T$ w.p. $1-\exp(\log(dT)-2\log^{2}(dT))=:1-p_{x}.$ Applying Hoeffding again yields,

$$ $\displaystyle\nabla\hat{F}_{i}^{k}(w^{t})-\nabla\hat{F}_{i}^{k}(\mathbb{E}_{x} [w^{t}])$ $\displaystyle=\frac{a_{i}}{nm}\sum_{j=1}^{n}\langle x_{j},\mathbb{E}_{x}[w_{i} ^{t}]-w_{i}^{t}\rangle x_{j}(k)x_{j}(1)x_{j}(2)$ $\displaystyle=\frac{a_{i}}{nm}\left\langle\sum_{j=1}^{n}x_{j}x_{j}(k)x_{j}(1)x _{j}(2),\mathbb{E}_{x}[w_{i}^{t}]-w_{i}^{t}\right\rangle$ $\displaystyle\leq\frac{\|\mathbb{E}_{x}[w_{i}^{t}]-w_{i}^{t}\|}{\sqrt{n}m}\log (dT)\log(dmT),$ $$

w.p. $1-\exp(\log(dmT)-2\log^{2}(dmT))-\exp(\log(dT)-2\log^{2}(dT))=:1-p_{2}-p_{x}$ uniformly over all $k\leq d,i\leq m,t\leq T$ where in the above we assumed $\mathbb{E}_{x}[w_{i}^{t}]-w_{i}^{t}$ as a fixed vector.

Denote $v^{t}:=\mathbb{E}_{x}[w_{i}^{t}]-w_{i}^{t}$. Concretely, we have the following bound for $k\neq 1,2$ on Term I:

$$ $\displaystyle\Pr_{x^{t}}\left(\forall i,k,t:\nabla\hat{F}_{i}^{k}(w^{t})- \nabla\hat{F}_{i}^{k}(\mathbb{E}_{x}[w^{t}])\leq\frac{\|v^{t}\|}{\sqrt{n}m} \log(dT)\log(dmT)\Big{|}v^{t}\right)\geq 1-p_{2}-p_{x}.$ $$

Therefore,

$$ $\displaystyle\Pr_{x^{t},v^{t}}\left(\forall i,k,t:\nabla\hat{F}_{i}^{k}(w^{t}) -\nabla\hat{F}_{i}^{k}(\mathbb{E}_{x}[w^{t}])\leq\frac{\|v^{t}\|}{\sqrt{n}m} \log(dT)\log(dmT)\right)\geq 1-p_{2}-p_{x}.$ $$

##### Case 2:

For $k=1:$

$$ $\displaystyle\nabla\hat{F}_{i}^{k}(w^{t})-\nabla\hat{F}_{i}^{k}(\mathbb{E}_{x} [w^{t}])$ $\displaystyle=\frac{a_{i}}{nm}\sum_{j=1}^{n}\langle x_{j},\mathbb{E}_{x}[w_{i} ^{t}]-w_{i}^{t}\rangle x_{j}(2)$ $\displaystyle=\frac{a_{i}}{nm}\langle\sum_{j=1}^{n}x_{j}x_{j}(2),\mathbb{E}_{x }[w_{i}^{t}]-w_{i}^{t}\rangle$ $\displaystyle=\frac{a_{i}}{m}v^{t}(2)+\frac{a_{i}}{m\sqrt{n}}\|v^{t}(\ell\neq 2 )\|\cdot\log(dT)\log(mT),$ $$

w.p. $1-p_{2}-p_{x}$.

Similarly for $k=2$ with the same probability it holds,

$$ $\displaystyle\nabla\hat{F}_{i}^{k}(w^{t})-\nabla\hat{F}_{i}^{k}(\mathbb{E}_{x} [w^{t}])$ $\displaystyle=\frac{a_{i}}{nm}\sum_{j=1}^{n}\langle x_{j},\mathbb{E}_{x}[w_{i} ^{t}]-w_{i}^{t}\rangle x_{j}(1)$ $\displaystyle=\frac{a_{i}}{nm}\langle\sum_{j=1}^{n}x_{j}x_{j}(1),\mathbb{E}_{x }[w_{i}^{t}]-w_{i}^{t}\rangle$ $\displaystyle=\frac{a_{i}}{m}v^{t}(1)+\frac{a_{i}}{m\sqrt{n}}\|v^{t}(\ell\neq 1 )\|\cdot\log(dT)\log(mT)$ $$

### E.7 Combining two terms

Recall we defined $v^{t}=\mathbb{E}_{x}[w_{i}^{t}]-w_{i}^{t}$ where $v^{t}\in\mathbb{R}^{d}$. Denote its entries by $v^{t}(k)$ for $k\leq d$.
Define $v^{0}=0$. Then, from the last two sections we obtain the following system of equations which hold w.p. $1-C(p_{1}-p_{x}-p_{2}-p_{w})$ for all $t\geq 1$:

$$ $\displaystyle|v^{t}(1)|$ $\displaystyle\leq\eta\sum_{\tau=0}^{t-1}\frac{|v^{\tau}(2)|}{m}+\frac{\|v^{ \tau}(\ell\neq 2)\|\log(dT)\log(mT)}{m\sqrt{n}}+\frac{|\alpha_{1}^{\tau}|+ \alpha_{2}^{\tau}+\sqrt{d}\log(d)}{m\sqrt{nd}}\log(dm)\log(dT)$ $\displaystyle|v^{t}(2)|$ $\displaystyle\leq\eta\sum_{\tau=0}^{t-1}\frac{|v^{\tau}(1)|}{m}+\frac{\|v^{ \tau}(\ell\neq 1)\|\log(dT)\log(mT)}{m\sqrt{n}}+\frac{|\alpha_{1}^{\tau}|+ \alpha_{2}^{\tau}+\sqrt{d}\log(d)}{m\sqrt{nd}}\log(dm)\log(dT)$ $\displaystyle|v^{t}(k)|$ $\displaystyle\leq\eta\sum_{\tau=0}^{t-1}\frac{\|v^{\tau}\|\log(dT)\log(dmT)}{m \sqrt{n}}+\frac{|\alpha_{1}^{\tau}|+\alpha_{2}^{\tau}+\sqrt{d}\log(d)}{m\sqrt{ nd}}\log(dm)\log(dT),\;\;\;\forall k\neq 1,2$ $$

Note that $|\alpha_{1}^{\tau}|+\alpha_{2}^{\tau}=2^{\tau}$. Simplify the above by assuming $\eta=m$ and denoting $\gamma_{\tau}:=\frac{2^{\tau}+\sqrt{d}\log(d)}{\sqrt{nd}}\log(dm)\log(dT)$. Then the equations above simplify to:

$$ $\displaystyle|v^{t}(1)|$ $\displaystyle\leq\sum_{\tau=0}^{t-1}|v^{\tau}(2)|+\frac{\|v^{\tau}\|\log^{2}( dmT)}{\sqrt{n}}+\sum_{\tau=0}^{t-1}\gamma_{\tau}$ $\displaystyle|v^{t}(2)|$ $\displaystyle\leq\sum_{\tau=0}^{t-1}|v^{\tau}(1)|+\frac{\|v^{\tau}\|\log^{2}( dmT)}{\sqrt{n}}+\sum_{\tau=0}^{t-1}\gamma_{\tau}$ $\displaystyle|v^{t}(k)|$ $\displaystyle\leq\sum_{\tau=0}^{t-1}\frac{\|v^{\tau}\|\log^{2}(dmT)}{\sqrt{n}} +\sum_{\tau=0}^{t-1}\gamma_{\tau},\;\;\;\forall k\neq 1,2$ $$

Define $z^{t}:=\max_{k}|v^{t}(k)|$. Then since the RHS in the third equations is smaller than the two first equations we deduce that:

$$ $\displaystyle z^{t}$ $\displaystyle\leq\sum_{\tau=0}^{t-1}{z^{\tau}}+\frac{z^{\tau}\sqrt{d}\log^{2}( dmT)}{\sqrt{n}}+\frac{2^{\tau}+\sqrt{d}\log(d)}{\sqrt{nd}}\log^{2}(dmT)$ (51) $$

To proceed in simplifying the equation above, we need the following lemma which can be straight-forwardly proved by induction.

###### Lemma 6 .

Let $v_{k},\beta,\gamma_{k}\in\mathbb{R}$. If $v_{0}=0$ and for every $t<T$:

$$ $\displaystyle v_{t+1}\leq\beta\sum_{k=0}^{t}v_{k}+\sum_{k=0}^{t}\gamma_{k}$ (52) $$

Then, it can be checked that this results in:

$$ $\displaystyle v_{T}\leq\sum_{k=0}^{T-1}(\beta+1)^{k}\gamma_{T-k-1}.$ (53) $$

Therefore, continuing from Eq. 51:

$$ $\displaystyle z^{t}$ $\displaystyle\leq\sum_{\tau=0}^{t-1}(2+\frac{\sqrt{d}\log^{2}(dmT)}{\sqrt{n}}) ^{\tau}\cdot\frac{2^{t-\tau-1}+\sqrt{d}\log(d)}{\sqrt{nd}}\log^{2}(dmT)$ $\displaystyle\leq\frac{t\log^{2}(dmT)}{\sqrt{nd}}(2+\frac{\sqrt{d}\log^{2}(dmT )}{\sqrt{n}})^{t-1}+\frac{(2+\frac{\sqrt{d}\log^{2}(dmT)}{\sqrt{n}})^{t}\log^{ 3}(dmT)}{\sqrt{n}(1+\frac{\sqrt{d}\log^{2}(dmT)}{\sqrt{n}})}$ $$

Assume $t\leq\log(d)$ and $n\geq d\cdot\log^{2c}(d)$ for some constant $c$. Then,

$$ $\displaystyle z^{t}$ $\displaystyle\leq\frac{d}{d\log^{c-3}(d)}+\frac{2d}{\sqrt{d}\log^{c-3}(d)}$ $\displaystyle\leq\frac{3\sqrt{d}}{\log^{c-3}(d)}$ $$

where in the above we used the fact that for $t\leq T\leq\log(d)$ and any $m=O(\mathrm{poly}(d))$:

$$ $\displaystyle(2+\frac{\sqrt{d}\log^{2}(dmT)}{\sqrt{n}})^{t-1}=(2+\frac{1}{\log ^{c-2}(d)})^{t-1}\leq(2+\frac{1}{\log(d)})^{t-1}\leq d$ $$

Therefore

$$ $\displaystyle\max_{k}|v^{t}(k)|=O(\frac{\sqrt{d}}{\log^{c-3}(d)})$ $$

We next show that $\|v^{t}(3:d)\|$ is much smaller. To see this, we replace our derived bound in the third equation above. define $\beta^{t}:=\|v^{t}(3:d)\|$ and note that $\beta^{t}\leq\sqrt{d}|v^{t}(k)|$ for every $k\in[3,d]$ and thus:

$$ $\displaystyle\beta^{t}$ $\displaystyle\leq\sum_{\tau=0}^{t-1}\frac{\|v^{\tau}\|\log^{2}(dmT)\sqrt{d}}{ \sqrt{n}}+\sqrt{d}\sum_{\tau=0}^{t-1}\gamma_{\tau}$ $\displaystyle\leq\sum_{\tau=0}^{t-1}\frac{\|v^{\tau}\|}{\log^{c-2}(d)}+\sqrt{d }\sum_{\tau=0}^{t-1}\gamma_{\tau}$ $\displaystyle\leq\sum_{\tau=0}^{t-1}\frac{\sqrt{({\beta^{\tau}})^{2}+\frac{d}{ \log^{2c-6}(d)}}}{\log^{c-2}(d)}+\sqrt{d}\sum_{\tau=0}^{t-1}\gamma_{\tau}$ $\displaystyle\leq\sum_{\tau=0}^{t-1}\frac{{\beta^{\tau}}+\frac{\sqrt{d}}{\log^ {c-3}(d)}}{\log^{c-2}(d)}+\sum_{\tau=0}^{t-1}\frac{2^{\tau}+\sqrt{d}\log(d)}{ \sqrt{n}}\log^{2}(dmT)$ $$

Using the lemma for the recursive summation (Lemma 6) again:

$$ $\displaystyle\beta^{t}$ $\displaystyle\leq\sum_{\tau=0}^{t-1}(1+\frac{1}{\log^{c-2}(d)})^{\tau}(\frac{ \sqrt{d}}{\log^{2c-5}(d)}+\frac{2^{t-\tau-1}+\sqrt{d}\log(d)}{\sqrt{n}}\log^{2 }(dmT))$ $\displaystyle\leq(1+\frac{1}{\log^{c-2}(d)})^{t}\frac{\sqrt{d}}{\log^{c-3}(d)} +\frac{2^{t-1}t}{\sqrt{d}\log^{c-2}(d)}+(1+\frac{1}{\log^{c-2}(d)})^{t}\log^{2 }(dmT)$ $$

where in the above we used the conditions on $m,T$. By using them again and recalling $t\leq T\leq\log(d)$, the above simplifies into:

$$ $\displaystyle\beta^{t}$ $\displaystyle\leq\frac{2\sqrt{d}}{\log^{c-3}(d)}+\frac{\sqrt{d}}{\log^{c-3}(d) }+2\log^{2}(dmT)$ $\displaystyle=O(\frac{\sqrt{d}}{\log^{c-3}(d)})$ $$

recall that $\mathbb{E}_{x}[w_{i}^{t}]$ had the following form

$$ $\displaystyle\mathbb{E}_{x}[w_{i}^{t}(1)]$ $\displaystyle=\left(\sum_{r\in[t],r:even}\binom{t}{r}(\frac{\eta a_{i}}{m})^{r }\right)w_{i}^{0}(1)+\left(\sum_{r\in[t],r:odd}\binom{t}{r}(\frac{\eta a_{i}}{ m})^{r}\right)w_{i}^{0}(2)$ $\displaystyle=\frac{1}{2}((1+a_{i})^{t}+(1-a_{i})^{t})w_{i}^{0}(1)+\frac{1}{2} ((1+a_{i})^{t}-(1-a_{i})^{t})w_{i}^{0}(2)$ $\displaystyle\mathbb{E}_{x}[w_{i}^{t}(2)]$ $\displaystyle=\left(\sum_{r\in[t],r:even}\binom{t}{r}(\frac{\eta a_{i}}{m})^{r }\right)w_{i}^{0}(2)+\left(\sum_{r\in[t],r:odd}\binom{t}{r}(\frac{\eta a_{i}}{ m})^{r}\right)w_{i}^{0}(1)$ $\displaystyle=\frac{1}{2}((1+a_{i})^{t}+(1-a_{i})^{t})w_{i}^{0}(2)+\frac{1}{2} ((1+a_{i})^{t}-(1-a_{i})^{t})w_{i}^{0}(1)$ $\displaystyle\mathbb{E}_{x}[w_{i}^{t}(k)]$ $\displaystyle=w_{i}^{0}(k),k>2.$ $$

where the last line is derived by choosing w.l.o.g. $\eta=m:$

$$ $\displaystyle\mathbb{E}_{x}[w_{i}^{t}(1)]$ $\displaystyle=2^{t-1}(w_{i}^{0}(1)+a_{i}w_{i}^{0}(2)),$ $\displaystyle\mathbb{E}_{x}[w_{i}^{t}(2)]$ $\displaystyle=2^{t-1}(w_{i}^{0}(2)+a_{i}w_{i}^{0}(1)),$ $\displaystyle\mathbb{E}_{x}[w_{i}^{t}(k)]$ $\displaystyle=w_{i}^{0}(k),\,\,k>2.$ $$

therefore,

$$ $\displaystyle|\mathbb{E}_{x}[w_{i}^{t}(1)]|$ $\displaystyle=\Theta(\frac{2^{t}}{\sqrt{d}}),$ $\displaystyle|\mathbb{E}_{x}[w_{i}^{t}(2)]|$ $\displaystyle=\Theta(\frac{2^{t}}{\sqrt{d}}),$ $\displaystyle|\mathbb{E}_{x}[w_{i}^{t}(k)]|$ $\displaystyle=\Theta(\frac{1}{\sqrt{d}}).$ $$

Select $t=\log(d)$. Then, w.h.p. it holds that

$$ $\displaystyle|\mathbb{E}_{x}[w_{i}^{t}(1)]|$ $\displaystyle=\Theta(\sqrt{d}),$ $\displaystyle|\mathbb{E}_{x}[w_{i}^{t}(2)]|$ $\displaystyle=\Theta(\sqrt{d}),$ $\displaystyle|\mathbb{E}_{x}[w_{i}^{t}(k)]|$ $\displaystyle=\Theta(\frac{1}{\sqrt{d}}).$ $$

The calculations above indicate that the signal strength is already larger than the noise after $\log(d)$ SGD steps. We are now ready to compute $y\Phi(w^{t},x)$. Assume for simplicity and without loss of generality that $x(1)=x(2)=y=1$ as other cases lead to the same result. Recall $v_{i}^{t}(k):=\mathbb{E}_{x}[w_{i}^{t}(k)]-w_{i}^{t}(k)$. Then,

$$ $\displaystyle y\Phi(w^{t},x)$ $\displaystyle=\frac{1}{m}\sum_{i=1}^{m}a_{i}\left(w_{i}^{t}(1)+w_{i}^{t}(2)+ \sum_{k=3}^{d}w_{i}^{t}(k)x(k)\right)^{2}$ $\displaystyle=\frac{1}{m}\sum_{i=1}^{m}a_{i}\left(d(w_{i}^{0}(1)+w_{i}^{0}(2)) (a_{i}+1)-v_{i}^{t}(1)-v_{i}^{t}(2)+\sum_{k\geq 3}(w_{i}^{0}(k)-v_{i}^{t}(k))x (k)\right)^{2}$ $\displaystyle=\frac{1}{m}\sum_{i\in P}\left(2d(w_{i}^{0}(1)+w_{i}^{0}(2))-v_{i }^{t}(1)-v_{i}^{t}(2)+\sum_{k\geq 3}(w_{i}^{0}(k)-v_{i}^{t}(k))x(k)\right)^{2}$ $\displaystyle\hskip 36.135pt+\frac{1}{m}\sum_{i\in N}\left(v_{i}^{t}(1)+v_{i}^ {t}(2)+\sum_{k\geq 3}(v_{i}^{t}(k)-w_{i}^{0}(k))x(k)\right)^{2},$ $$

where $P$ and $N$ denote the set of neurons for which $a_{i}=1$ and $a_{i}=-1$, respectively.

Note that the last term in each summation can be simplified as follows:

$$ $\sum_{k\geq 3}(w_{i}^{0}(k)+v_{i}^{t}(k))x(k)=O(\|v_{i}^{t}\|)=O\left(\beta^{t }\right)=O\left(\frac{\sqrt{d}}{\log^{c-3}(d)}\right).$ $$

Concretely, by Hoeffding w.p. $1-C(p_{1}-p_{2}-p_{x}-p_{w})$ over data sampling and initialization :

$$ $\displaystyle\Pr_{x}\left(\sum_{k\geq 3}\left(w_{i}^{0}(k)+v_{i}^{t}(k)\right) x(k)>\frac{\sqrt{d}}{\log^{\frac{c-4}{2}}(d)}\right)\leq e^{-\log(d)}=\frac{1} {d}.$ (54) $$

Define $\zeta_{i}^{t}:=v_{i}^{t}(1)+v_{i}^{t}(2)+\sum_{k\geq 3}(v_{i}^{0}(k)-w_{i}^{t}
(k))x(k)$. Then,

$$ $\displaystyle y\Phi(w^{t},x)$ $\displaystyle=\frac{1}{m}\sum_{i\in P}\left(2d(w_{i}^{0}(1)+w_{i}^{0}(2))- \zeta_{i}^{t}\right)^{2}+\frac{1}{m}\sum_{i\in N}\left(\zeta_{i}^{t}\right)^{2}$ $\displaystyle=\frac{1}{m}\sum_{i\in P}4d^{2}(w_{i}^{0}(1)+w_{i}^{0}(2))^{2}-4d \zeta_{i}^{t}(w_{i}^{0}(1)+w_{i}^{0}(2))+\frac{1}{m}\sum_{i\in[m]}(\zeta_{i}^{ t})^{2}$ $\displaystyle\geq\frac{1}{m}\sum_{i\in P}4d^{2}(w_{i}^{0}(1)+w_{i}^{0}(2))^{2} -4d|\zeta_{i}^{t}|\cdot|w_{i}^{0}(1)+w_{i}^{0}(2)|.$ (55) $$

Moreover, we had w.p. at least $1-C(p_{1}+p_{2}+p_{x}+p_{w})$ that $|v_{i}^{t}(1)|+|v_{i}^{t}(2)|<\frac{C\sqrt{d}}{\log^{c-3}(d)}$ for some absolute constant $C.$ Adding these two together, we deduce that w.p. $1-o_{d}(1)$ it holds uniformly over $i,t$,

$$ $\displaystyle|\zeta_{i}^{t}|<\frac{C\sqrt{d}}{\log^{\frac{c-4}{2}}(d)}.$ $$

It remains to obtain lower- and upper-bounds on $(w_{i}^{0}(1)+w_{i}^{0}(2))^{2}$ as required by Eq. 55. To proceed, note that $\frac{d}{2}(w_{i}^{0}(1)+w_{i}^{0}(2))^{2}\sim\chi^{2}(1).$ Therefore with standard concentration bounds for $\chi^{2}(1)$ and $\chi^{2}(m)$ distributions (e.g., []) we find for any $u\geq 0$,

$$ $\displaystyle\Pr\left(\frac{d}{2}(w_{i}^{0}(1)+w_{i}^{0}(2))^{2}-1\geq 2\sqrt{ u}+2u\right)\leq e^{-u}$ $\displaystyle\Pr\left(m-\frac{d}{2}\sum_{i=1}^{m}(w_{i}^{0}(1)+w_{i}^{0}(2))^{ 2}\geq 2\sqrt{mu}\right)\leq e^{-u}$ $$

By selecting $u=\log^{2}(d)$ in the first inequality above, we find w.p. at most $e^{\log(m)-\log^{2}(d)}$ it holds uniformly for all $i\leq m$:

$$ $\displaystyle(w_{i}^{0}(1)+w_{i}^{0}(2))^{2}\geq\frac{2}{d}+\frac{4}{d}\sqrt{ \log^{2}(d)}+\frac{4}{d}\log^{2}(d),$ $$

and by noting that $|P|=|N|=m/2$ for some $m\geq 2$ and selecting $u=m/16$ in the second inequality we have with probability at least $1-e^{-\frac{m}{16}}$,

$$ $\displaystyle\frac{1}{m}\sum_{i\in P}(w_{i}^{0}(1)+w_{i}^{0}(2))^{2}\geq(1- \frac{1}{\sqrt{2}})\frac{1}{d}.$ $$

Hence w.p. $1-e^{\log(m)-\log^{2}(d)}-e^{-\frac{m}{16}}-C(e^{-\log^{2}(d)}+e^{\log(dmT)-2
\log^{2}(dmT)}+e^{\log(dT)-2\log^{2}(dT)}+e^{-\frac{1}{2}\log^{2}(dm)+\log(dm)})$ over data sampling and initialization,

$$ $\displaystyle\Pr_{x,y}\left(y\Phi(w^{t},x)\geq d-\frac{4Cd}{\log^{\frac{c-6}{2 }}(d)}\right)\geq 1-1/d.$ (56) $$

For the theorem’s choice of $T=\log(d),m=O(\mathrm{poly}(d))$ it is deduced that $e^{\log(m)-\log^{2}(d)}+C(e^{-\log^{2}(d)}+e^{\log(dmT)-2\log^{2}(dmT)}+e^{
\log(dT)-2\log^{2}(dT)}+e^{-\frac{1}{2}\log^{2}(dm)+\log(dm)})=o_{d}(1)$. Moreover, for the theorem’s choice $c=7$ and large enough $d$ it holds that $d-4Cd/\log^{\frac{c-6}{2}(d)}\geq 0.$ This leads to the following lower bound on test accuracy which holds with probability at least $1-2e^{-\frac{m}{16}}-e^{\log(m)-\log^{2}(d)}-o_{d}(1)$ over the randomness in data sampling and initialization,

$$ $\displaystyle\Pr_{x,y}\left(y\Phi(w^{t},x)\geq 0\right)\geq 1-\frac{1}{d},$ $$

This completes the proof of Theorem 2.4.

###### Lemma 6 .

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 8
presentation: 3
questions: 1) Do you think that the bounds holding for ReLu networks could be improved to match the bounds of your paper, or is the smoothness assumption necessary to obtain the improvement?

2) Could you clarify the main novel techniques or approaches used in the proofs of your main results, particularly in relation to previous work in this area?
soundness: 3
strengths: -The paper is clearly written and well organized.

-The paper provides novel convergence and generalization bounds for neural networks with smooth activations, which are significantly better than prior bounds derived for ReLU networks (the test loss bound obtained is width independent). This is important for understanding how generalization of neural networks with smooth activations compares to generalization of ReLU neural networks.

-The result showing efficient learning for the XOR problem with linear sample complexity is a notable contribution.
summary: This paper obtains convergence and generalization bounds for gradient descent training of neural networks with smooth activation functions. An algorithmic stability analysis exploiting the Hessian properties in the gradient descent path allows to bound the generalization gap with the cumulative optimization error.  Ultimately, the bound on the test loss depends on the Euclidean distance with initialization and the Lipschitz constant of the network at initialization. Experiments with a 2-hidden layer network on MNIST and FashionMNIST show that the theoretical bounds provide accurate approximations for the test loss and generalization gap.  Under an NTK separability condition, the results of the paper (with smooth activations) lead to a better bound on the test loss then previous results with ReLU activation functions ($\sqrt{\frac{m}{n}}$ is replaced with $\frac{1}{n}$). Furthermore, on the specific problem of a d-dimensional XOR distribution, using quadratic activations with a large step size, only log(d) iterations (and constant width) are needed to obtain perfect test accuracy with n = O(d) samples. An experiment also confirms the logarithmic dependency on data dimension.
weaknesses: -Experiments on deep networks are done on a small neural network. Also, the expectations over S are estimated with one realization of S.

-The method to bound the spectral norm of the Hessian based on the distance from initialization is not new, limiting the novelty of the approach.

-The paper does not clarify whether the improved bounds are fundamentally due to the smoothness of the activation functions or if they result from limitations in the proof techniques used in previous bounds for ReLU networks, which may have led to weaker results in the non-smooth case.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. In Eq. (5), consider the special case of $n=1$, where we have only one sample. Suppose we initialize $w_0$ at $w^\star$ and suppose $w^\star$ achieves zero training error. Then by Eq. (5), the generalization error will be $0$. This means that even in the case where we have only one trainig sample, every neural network parameter that achieves zero training error will generalize well, which is counter-intuitive. Am I missing something here?

2. In Eq. (6), plugging in $\hat{F} \leq \frac{4\rho^{*2}}{\eta T}$, should we have an additional $\log T$ on the right-hand side due to a summation of $\frac{1}{t}$ for $t = 1,\dots, T-1$?

3. In the cited paper Lei and Ying (2020), I did not found Lemma B.3, which is used in the proof of the generalization error. Is the paper referring to a different lemma? Also, it seems that the leave-one-out error in Lei and Ying is defined by replacing the $i$th sample with another random sample, not deleting the $i$th sample.
soundness: 3
strengths: 1. The paper obtains a tigher bound on the generalization error compared with previous work. In particular, while previous work usually achieves a $O(\frac{1}{\sqrt{n}})$ generalization error, this paper achieves a $O(\frac{1}{n})$ error.

2. In the noiseless data scheme, the paper derives favorable convergence guarantee under only a constant over-parameterization requirement.

3. The paper has a nice extension of their theoretical work from the noiseless data scheme to the noisy data scheme as well as the XOR classification task.
summary: This paper studies the training convergence guarantee and the generalization error of deep MLPs with quadratic activation when applied to binary classification tasks. The paper considers the parameterization of the neural network in the NTK regime, and shows a tigher generalization bound together with a smaller over-parameterization requirement compared with previous work. In addition to the analysis of the noiseless separable data scheme, the paper also considers a noisy data scheme and a XOR classification task. In particular, the paper shows that for the XOR classification, using a large step size and effectively escape the NTK regime and learns useful features. The paper also provided experimental results to verify its theoretical conclusion.
weaknesses: 1. In Theorem 2.1, the over-parameterization requirement scales with $\\Omega(\\rho^{*6L+4})$. Later in the discussion, the paper stated $\\rho = O(\\log n / \\gamma)$, where $\\gamma$ is the margin width. When we treat $L$, the network depth, not as a constant, the over-parameterization requirement can be as large as $\\Omega(\\frac{(\\log n)^{6L+4}}{\\gamma^{6L+4}})$, which might be unfavorable.

2. The setting studied in the paper are not clearly explained, and some notations are not introduced before use. For instance, it is not clear what is $\mathcal{S}$ is in Eq.(5), and what it means to take expectation over it. Moreover, in Section 2.1.2, it is not clear how the noise in the data are generated.

3. In Assumption 1, the separability condition is defined over the tangent feature at initialization. Because the separability is defined in this random feature model, it would naturally imply an implicit lower bound on the over-parameterization: for instance, the more parameters we have, the more random features we will have, and a higher probability the data will be separable when mapped to the feature space.

4. Some theoretical results seems confusing; see Questions below.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1.	Table in page 2 is mis-labeled as Table 2. 
2.	The authors claim that leveraging the Hessian structure in the gradient descent path is the key for their theoretical improvements over the past literature. But I feel the mathematical intuition is not well stated in the main paper. When it is compared with the previous literature (for instance Chen et al, 2021), how does considering the hessian structure benefit the author to prove Theorem 2.1? 
3.	In least-squares regression, the Neural Tangent Kernel (NTK) benefits from the fact that gradient descent keeps weight updates close to their initialization, which linearizes the dynamics in parameter space and makes the analysis more tractable. I suspect this approach will similarly apply to classification tasks if the authors use the NTK framework. Given this, the role of initialization mentioned in the main paper isn't particularly surprising. Could the authors elaborate on the unique insights that arise when studying classification problems, which might not be evident when analyzing regression problems, specifically in terms of the role of initialization?
4.     Somehow, I felt Subsection 2.2 is disconnected to the remaining part of the paper, as the problem settings seem to be different from the previous parts. I suspect this is because of technical difficulties authors might have encountered, but it would be great if authors can make connections between this Subsection and previous parts. For example, how are the particular techniques from earlier sections linked to the XOR distribution analysis in 2.2?
5.	In Theorem 2.4, the step size is equal to the width of the network. Can it be arbitrarily large? Are there any restrictions on the step size?
soundness: 3
strengths: Overall, the contribution of this paper is great filling the gap between practice and theory.
summary: This paper studies the data-dependent convergence and generalization of gradient methods in neural networks with smooth activations. The authors provide a tighter, algorithm-dependent bound on excess risk for deep networks trained with logistic loss, improving on Rademacher complexity-based bounds. These results emphasize the role of initialization, particularly in small-width networks and margin-separable data.
weaknesses: Please see the questions below

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: None.
soundness: 3
strengths: This work presents several improved results from prior works. In particular, the authors are able to show an improved bound on the network width by requiring $\lvert w^\star - w_0 \rvert \lesssim m^{O(1/L)}$, which implies that the NTK regime can happen under a far more relaxed conditions than what previous work suggested. Such results can be important and interesting. The authors then utilize this improvement to derive improved bounds for several classical NTK settings like NTK separable data and noisy data.
summary: This paper studies convergence and generalization for deep neural networks with smooth activation trained by gradient descent. This paper shows several improved results on test error compared with previous works for the following scenarios: (1) noiseless data with separable NTK feature, and (2) noisy data and (3) XOR distribution.
weaknesses: Although the results in this work are impressive, a more detailed description on how the authors achieved this improvement is needed. In particular, the authors mentioned in line 84, 85 that the reason for the improvement in the network's width is by exploring the objective's Hessian structure in the gradient descent path. This sounds very interesting. However, I didn't find any high level discussion/description on this in the main text. I would like the authors to explain how they get this improvement. How does your analysis compare with the prior work that also utilized the Hessian structure? What is the novelty in your analysis such that you can achieve this improvement?

Meta-review: The work presents new advances on generalization bounds for learning neural networks with smooth activations. The approach is based on algorithmic stability, avoiding the challenges with Rademacher complexity based bounds, and also applicable in a small width regime, effectively extending the scope of NTK style analysis. The paper presents precise results under data separability with margin $\gamma$, with sharper results compared to past work. The paper also presents results on the noisy setting and also studies the special case of the XOR distribution. The reviewers overall appreciated the advances reported in the paper, including the new techniques and the implications of the results. Some concerns were brought up which have been mostly addressed by the authors during author discussion.

Decision: Accept (Poster)