Generalization and Stability of Interpolating Neural Networks with
  Minimal Width
We investigate the generalization and optimization properties of shallow
neural-network classifiers trained by gradient descent in the interpolating
regime. Specifically, in a realizable scenario where model weights can achieve
arbitrarily small training error $\epsilon$ and their distance from
initialization is $g(\epsilon)$, we demonstrate that gradient descent with $n$
training data achieves training error $O(g(1/T)^2 /T)$ and generalization error
$O(g(1/T)^2 /n)$ at iteration $T$, provided there are at least
$m=\Omega(g(1/T)^4)$ hidden neurons. We then show that our realizable setting
encompasses a special case where data are separable by the model's neural
tangent kernel. For this and logistic-loss minimization, we prove the training
loss decays at a rate of $\tilde O(1/ T)$ given polylogarithmic number of
neurons $m=\Omega(\log^4 (T))$. Moreover, with $m=\Omega(\log^{4} (n))$ neurons
and $T\approx n$ iterations, we bound the test loss by $\tilde{O}(1/n)$. Our
results differ from existing generalization outcomes using the
algorithmic-stability framework, which necessitate polynomial width and yield
suboptimal generalization rates. Central to our analysis is the use of a new
self-bounded weak-convexity property, which leads to a generalized local
quasi-convexity property for sufficiently parameterized neural-network
classifiers. Eventually, despite the objective's non-convexity, this leads to
convergence and generalization-gap bounds that resemble those found in the
convex setting of linear logistic regression.
Hossein Taheri, Christos Thrampoulidis
# Generalization and Stability of Interpolating Neural Networks with Minimal Width

###### Abstract

We investigate the generalization and optimization properties of shallow neural-network classifiers trained by gradient descent in the interpolating regime. Specifically, in a realizable scenario where model weights can achieve arbitrarily small training error \(\epsilon\) and their distance from initialization is \(g(\epsilon)\), we demonstrate that gradient descent with \(n\) training data achieves training error \(O(g(1/T)^{2}/T)\) and generalization error \(O(g(1/T)^{2}/n)\) at iteration \(T\), provided there are at least \(m=\Omega(g(1/T)^{4})\) hidden neurons. We then show that our realizable setting encompasses a special case where data are separable by the model's neural tangent kernel. For this and logistic-loss minimization, we prove the training loss decays at a rate of \(\tilde{O}(1/T)\) given polylogarithmic number of neurons \(m=\Omega(\log^{4}(T))\). Moreover, with \(m=\Omega(\log^{4}(n))\) neurons and \(T\approx n\) iterations, we bound the test loss by \(\tilde{O}(1/n)\). Our results differ from existing generalization outcomes using the algorithmic-stability framework, which necessitate polynomial width and yield suboptimal generalization rates. Central to our analysis is the use of a new self-bounded weak-convexity property, which leads to a generalized local quasi-convexity property for sufficiently parameterized neural-network classifiers. Eventually, despite the objective's non-convexity, this leads to convergence and generalization-gap bounds that resemble those found in the convex setting of linear logistic regression.

## 1 Introduction

Neural networks have remarkable expressive capabilities and can memorize a complete dataset even with mild overparameterization. In practice, using gradient descent (GD) on neural networks with logistic or cross-entropy loss can result in the objective reaching zero training error and close to zero training loss. Zero training error, often referred to as "interpolating" the data, indicates perfect classification of the dataset. Despite their strong memorization ability, these networks also exhibit remarkable generalization capabilities to new data. This has motivated a surge of studies in recent years exploring the optimization and generalization properties of first-order gradient methods in overparameterized neural networks, with a specific focus in the so-called Neural Tangent Kernel (NTK) regime. In the NTK regime, the model operates as the first-order approximation of the network at a sufficiently large initialization or at the large-width limit (Jacot et al., 2018; Chizat et al., 2019). Prior works on this topic mostly focused on quadratic-loss minimization and their optimization/generalization guarantees required network widths that increased polynomially with the sample size \(n\). This, however, is not in line with practical experience. Improved results were obtained more recently by (Ji and Telgarsky, 2020; Chen et al., 2020) who have investigated the optimization and generalization of ReLU neural networks with logistic loss, which is more suitable for classification tasks. Assuming that the NTK with respect to the model can interpolate the data (i.e. separate them with positive margin \(\gamma\)), they showed through a Rademacher complexity analysis that GD on neural networks with polylogarithmic width can achieve generalization guarantees that decrease with the sample size \(n\) at a rate of \(\tilde{O}(\frac{1}{\sqrt{n}})\).
In this paper, we provide rate-optimal optimization and generalization analyses of GD for shallow neural networks of minimal width assuming that the model itself can interpolate the data. We focus on two-layer networks with smooth activations that can almost surely separate \(n\) training samples from the data distribution. Concretely, we consider a realizability condition where data and initialization are such that model weights can achieve arbitrarily small training error \(\varepsilon\) while their distance from initialization is \(g(\varepsilon)\) for some function \(g:\mathbb{R}_{+}\rightarrow\mathbb{R}_{+}\). Under this condition, we demonstrate generalization guarantees of order \(O(\frac{g(\frac{1}{T})^{2}}{n}).\) More generally, for any iteration \(T\) of GD and assuming network width \(m=\Omega(g(\frac{1}{T})^{4})\), we obtain an expected test-loss rate \(O(\frac{g(\frac{1}{T})^{2}}{T}+\frac{g(\frac{1}{T})^{2}}{n})\). Additional to the generalization bounds, we provide optimization guarantees under the same setting by showing that the training loss approaches zero at rate \(O(\frac{g(\frac{1}{T})^{2}}{T})\). We note that these results are derived without NTK-type analyses. For demonstration and also for connection to prior works on neural-tangent data models, we specialize our generalization and optimization results to the class of NTK-separable data. We show this is possible because the NTK-data separability assumption implies our realizability condition holds. Thus, for logistic-loss minimization on NTK-separable data, we show that the expected test loss of GD is \(\tilde{O}(\frac{1}{T}+\frac{1}{n})\) provided polylogarithmic number of neurons \(m=\Omega(\log^{4}(T))\). This further suggests that a network of width \(m=\Omega(\log^{4}(n))\), attains expected test loss \(\tilde{O}(\frac{1}{n})\) after \(T\approx n\) iterations.

In contrast to prior optimization and generalization analyses that often depend on the NTK framework, which requires the first-order approximation of the model, we build on the algorithmic stability approach (Bousquet and Elisseeff, 2002) for shallow neural-network models of finite width. Although the stability analysis has been utilized in previous studies to derive generalization bounds for (stochastic) gradient descent in various models, most results that are rate-optimal heavily rely on the convexity assumption. Specifically, the stability-analysis framework has been successful in achieving optimal generalization bounds for convex objectives in (Lei and Ying, 2020; Bassily et al., 2020; Schliserman and Koren, 2022). On the other hand, previous studies on non-convex objectives either resulted in suboptimal bounds or relied on assumptions that are not in line with the actual practices of neural network training. For instance, (Hardt et al., 2016) derived a generalization bound of \(O(\frac{T^{\beta e/(\beta e+1)}}{n})\) for general \(\beta\)-smooth and non-convex objectives, but this required a time-decaying step-size \(\eta_{t}\leq c/t\), which can degrade the training performance. More recently, (Richards and Rabbat, 2021) explored the use of the stability approach specifically for logistic-loss minimization of a two-layer network. By refining the model-stability analysis framework introduced by (Lei and Ying, 2020), they derived generalization-error bounds provided the hidden width increases polynomially with the sample size. In comparison, our analysis leads to improved generalization and optimization rates and under standard separability conditions such as NTK-separability, only requires a polylogarithmic width for both global convergence and generalization.

### Notation

We denote \([n]:=\{1,2,\cdots,n\}\). We use the standard notation \(O(\cdot),\Omega(\cdot)\) and use \(\tilde{O}(\cdot),\tilde{\Omega}(\cdot)\) to hide polylogarithmic factors. Occasionally we use \(\lesssim\) to hide numerical constants. The Gradient and Hessian of a function \(\Phi:\mathbb{R}^{d_{1}\times d_{2}}\rightarrow\mathbb{R}\) with respect to the \(i\)th input (\(i=1,2\)) are denoted by \(\nabla_{i}\Phi\) and \(\nabla_{i}^{2}\Phi\), respectively. All logarithms are in base \(e\). We use \(\left\|\cdot\right\|\) for the \(\ell_{2}\) norm of vectors and the operator norm of matrices. We denote \([w_{1},w_{2}]:=\{w\;:\;w=\alpha w_{1}+(1-\alpha)w_{2},\alpha\in[0,1]\}\) the line segment between \(w_{1},w_{2}\in\mathbb{R}^{d^{\prime}}\).

## 2 Problem Setup

Given \(n\) i.i.d. samples \((x_{i},y_{i})\sim\mathcal{D},i\in[n]\) from data distribution \(\mathcal{D}\), we study unconstrained empirical risk minimization with objective \(\widehat{F}:\mathbb{R}^{d^{\prime}}\rightarrow\mathbb{R}\):

\[\min_{w\in\mathbb{R}^{d^{\prime}}}\Big{\{}\widehat{F}(w):=\frac{1}{n}\sum_{i=1 }^{n}\widehat{F}_{i}(w)=\frac{1}{n}\sum_{i=1}^{n}f\left(y_{i}\Phi\left(w,x_{i} \right)\right)\Big{\}}. \tag{1}\]
This serves as a proxy for minimizing the _test loss_\(F:\mathbb{R}^{d^{\prime}}\to\mathbb{R}\):

\[F(w):=\mathbb{E}_{(x,y)\sim\mathcal{D}}\left[f\left(y\Phi(w,x)\right)\right]. \tag{2}\]

We introduce our assumptions on the data \((x,y)\), the model \(\Phi(\cdot,x)\), and the loss function \(f(\cdot)\), below. We start by imposing the following mild assumption on the data distribution.

**Assumption 1** (Bounded features).: _Assume any \((x,y)\sim\mathcal{D}\) has almost surely bounded features, i.e. \(\|x\|\leq R\), and binary label \(y\in\{\pm 1\}\)._

The model \(\Phi:\mathbb{R}^{d^{\prime}}\times\mathbb{R}^{d}\to\mathbb{R}\) is parameterized by trainable weights \(w\in\mathbb{R}^{d^{\prime}}\) and takes input \(x\in\mathbb{R}^{d}\). For our main results, we assume \(\Phi\) is a one-hidden layer neural-net of \(m\) neurons, i.e.

\[\Phi(w,x):=\frac{1}{\sqrt{m}}\sum_{j=1}^{m}a_{j}\,\sigma(\langle w_{j},x\rangle), \tag{3}\]

where \(\sigma:\mathbb{R}\to\mathbb{R}\) is the activation function, \(w_{j}\in\mathbb{R}^{d}\) denotes the weight vector of the \(j\)th hidden neuron and \(\frac{a_{j}}{\sqrt{m}},j\in[m]\) are the second-layer weights. For the second layer weights, we assume that they are fixed during training taking values \(a_{j}\in\{\pm 1\}\). We assume that for half of second layer weights we have \(a_{j}=1\) and for the other half \(a_{j}=-1\). On the other hand, all the first-layer weights are updated during training. Thus, the total number of trainable parameters is \(d^{\prime}=md\) and we denote \(w=[w_{1};w_{2};\ldots;w_{m}]\in\mathbb{R}^{d^{\prime}}\) the vector of trainable weights. Throughout, we make the following assumptions on the activation function.

**Assumption 2** (Lipschitz and smooth activation).: _The activation function \(\sigma:\mathbb{R}\to\mathbb{R}\) satisfies the following for non-negative constants \(\ell,L\):_

\[|\sigma^{\prime}(u)|\leq\ell,\;\;\;|\sigma^{\prime\prime}(u)|\leq L,\qquad \forall u\in\mathbb{R}.\]

We note that the smoothness assumption which is required by our framework excludes the use of ReLU. Examples of activation functions that satisfy the smoothness condition include Softplus \(\sigma(u)=\log(1+e^{u})\), Gaussian error linear unit (GELU) \(\sigma(u)=\frac{1}{2}u(1+\operatorname{erf}(\frac{u}{\sqrt{2}}))\), and Hyperbolic-Tangent where \(\sigma(u)=\frac{e^{u}-e^{-u}}{e^{u}+e^{-u}}\). On the other hand, Lipschitz assumption is rather mild, since it is possible to restrict the parameter space to a bounded domain.

Next, we discuss conditions on the loss function. Of primal interest is the commonly used logistic loss function \(f(u)=\log(1+e^{-u})\). However, our results hold for a broader class of convex, non-negative and monotonically decreasing functions (\(\lim_{u\to\infty}f(u)=0\)) that satisfy the following:

**Assumption 3** (Lipschitz and smooth loss).: _The convex loss function \(f:\mathbb{R}\to\mathbb{R}_{+}\) satisfies for all \(u\in\mathbb{R}\)_

_3.A:_ _Lipschitzness:_ \(|f^{\prime}(u)|\leq G_{f}\)_._

_3.B:_ _Smoothness:_ \(f^{\prime\prime}(u)\leq L_{f}\)_._

**Assumption 4** (Self-bounded loss).: _The convex loss function \(f:\mathbb{R}\to\mathbb{R}_{+}\) is self-bounded with some constant \(\beta_{f}>0\), i.e., \(|f^{\prime}(u)|\leq\beta_{f}f(u),\forall u\in\mathbb{R}\)._

The self-boundedness Assumption 4 is the key property of the loss that drives our analysis and justifies the polylogarithmic width requirement, as will become evident. Note that the logistic loss naturally satisfies Assumptions 3.A and 3.B (with \(G_{f}=1,L_{f}=1/4\)), as well as, Assumption 4 with \(\beta_{f}=1\). Other interesting examples of loss functions satisfying those assumptions include polynomial losses, with the tail behavior \(f(u)=1/u^{\beta}\) for \(\beta>0\), which we discuss in Remark 2. To lighten the notation and without loss of generality, we set \(G_{f}=L_{f}=\beta_{f}=1\) for the rest of the paper. We remark that our training-loss results also hold for the exponential loss \(e^{-u}\). The exponential loss is self-bounded and while it is not Lipschitz or smooth it satisfies a second-order self-bounded property \(f^{\prime\prime}(u)\leq f(u)\), which we can leverage instead; see Appendix A for details.
Main Results

We present bounds on the train loss and generalization gap of gradient-descent (GD) under the setting of Section 2. Formally, GD with step-size \(\eta>0\) optimizes (1) by performing the following updates starting from an initialization \(w_{0}\):

\[\forall t\geq 0\,:\ w_{t+1}=w_{t}-\eta\nabla\widehat{F}(w_{t}).\]

### Key properties

The key challenge in both the optimization and generalization analysis is the non-convexity of \(f(y\Phi(\cdot,x))\), and consequently of the train loss \(\widehat{F}(\cdot)\). Despite non-convexity, we derive bounds analogous to the convex setting, e.g. corresponding bounds on linear logistic regression in (Ji and Telgarsky, 2018; Shamir, 2021; Schliserman and Koren, 2022). We show this is possible provided the loss satisfies the following key property, which we call _self-bounded weak convexity_.

**Definition 1** (Self-bounded weak convexity).: _We say a function \(\widehat{F}:\mathbb{R}^{d^{\prime}}\to\mathbb{R}\) is self-bounded weakly convex if there exists constant \(\kappa>0\) such that for all \(w\),_

\[\lambda_{\min}\left(\nabla^{2}\widehat{F}(w)\right)\geq-\kappa\, \widehat{F}(w)\,. \tag{4}\]

Recall a function \(G:\mathbb{R}^{d^{\prime}}\to\mathbb{R}\) is weakly convex if \(\exists\kappa\geq 0\) such that uniformly over all \(w\in\mathbb{R}^{d^{\prime}}\), \(\lambda_{\min}\left(\nabla^{2}G(w)\right)\geq-\kappa.\) If \(\kappa=0\), the function is convex. Instead, property (4) lower bounds the curvature by \(-\kappa\,G(w)\) that changes proportionally with the function value \(G(w)\). We explain below how this is exploited in our setting.

To begin with, the following lemma shows that property (4) holds for the train loss under the setting of Section 2: training of a two-layer net with smooth activation and self-bounded loss. The lemma also shows that the gradient of the train loss is self bounded. Those two properties together summarize the key ingredients for which our analysis applies.

**Lemma 1** (Key self-boundedness properties).: _Consider the setup of Section 2 and let Assumptions 1-2 hold. Further assume the loss is self-bounded as per Assumption 4. Then, the objective satisfies the following self-boundedness properties for its Gradient and Hessian:_

1. _Self-bounded gradient:_ \(\left\|\nabla\widehat{F}_{i}(w)\right\|\leq\ell R\,\widehat{F}_{i}(w),\ \ \forall i\in[n]\)_._
2. _Self-bounded weak convexity:_ \(\lambda_{\min}\left(\nabla^{2}\widehat{F}(w)\right)\geq-\frac{LR^{2}}{\sqrt{ m}}\widehat{F}(w)\)_._

Both of these properties follow from the self-boundedness of the convex loss \(f\) combined with Lipshitz and smoothness of \(\sigma\). The self-boundedness of the gradient is used for generalization analysis and in particular in obtaining the model stability bound. The self-bounded weak convexity plays an even more critical role for our optimization and generalization results. In particular, the wider the network the closer the loss to having convex-like properties. Moreover, the "self-bounded" feature of this property provides another mechanism that favors convex-like optimization properties of the loss. To see this, consider the minimum Hessian eigenvalue \(\lambda_{\min}(\nabla^{2}\widehat{F}(w_{t}))\) at gradient descent iterates \(\{w_{t}\}_{t\geq 1}\): As training progresses, the train loss \(\widehat{F}(w_{t})\) decreases, and thanks to the self-bounded weak convexity property, the gap to convexity also decreases. We elaborate on the role of self-bounded weak convexity in our proofs in Section 5.

### Training loss

We begin with a general bound on the training loss and the parameter's norm, which is also required for our generalization analysis.
**Theorem 2** (Training loss - General bound).: _Suppose Assumptions 1-4 hold. Fix any training horizon \(T\geq 0\) and any step-size \(\eta\leq 1/L_{\widehat{F}}\) where \(L_{\widehat{F}}\) is the objective's smoothness parameter. Assume any \(w\in\mathbb{R}^{d^{\prime}}\) and hidden-layer width \(m\) such that \(\|w-w_{0}\|^{2}\geq\max\{\eta T\widehat{F}(w),\eta\widehat{F}(w_{0})\}\) and \(m\geq 18^{2}L^{2}R^{4}\|w-w_{0}\|^{4}\). Then, the training loss and the parameters' norm satisfy_

\[\widehat{F}(w_{T}) \;\leq\;\frac{1}{T}\sum_{t=1}^{T}\widehat{F}(w_{t})\;\leq\;2 \widehat{F}(w)+\frac{5\|w-w_{0}\|^{2}}{2\eta T}, \tag{5}\] \[\forall t\in[T] :\;\;\|w_{t}-w_{0}\|\;\leq\;4\|w-w_{0}\|.\]

A few remarks are in place regarding the theorem. First, Eq. (5) upper bounds the running average (also known as regret) of train loss for iterations \(1,\ldots,T\) by the value, at an arbitrarily chosen point \(w\), of a ridge-regularized objective with regularization parameter inversely proportional to \(\eta T\). Because of smoothness and Lipschitz Assumption 3 of \(f\), it turns out that the training objective is \(L_{\widehat{F}}\)-smooth. Hence, by the descent lemma of GD for smooth functions, the same upper bound holds in Eq. (5) for the value of the loss at time \(T\), as well. Moreover, the theorem provides a uniform upper bound of the norm of all GD iterates in terms of \(\|w-w_{0}\|\). Notably, and despite the non-convexity in our setting, our bounds are same up to constants to analogous bounds for logistic linear regression in (Shamir, 2021; Schliserman and Koren, 2022). As discussed in Sec. 3.1 this is possible thanks to the self-bounded weak convexity property.

The condition \(m\gtrsim\|w-w_{0}\|^{4}\) on the norm of the weights controls the maximum deviations of weights \(w\) from initialization (with respect to network width) required for our results to guarantee arbitrarily small train loss. Specifically, to get the most out of Theorem 2 we need to choose appropriate \(w\) that satisfies both the condition \(m\gtrsim\|w-w_{0}\|^{4}\) and keeps the associated ridge-regularized loss \(\widehat{F}(w)+\|w-w_{0}\|^{2}/(\eta T)\) small. This combined requirement is formalized in the neural-net realizability Assumption 5 below. As we will discuss later in Section 4, this assumption translates into an assumption on the underlying data distribution that ultimately enables the application of Theorem 2 to achieve vanishing training error.

**Assumption 5** (Nn-Realizability).: _There exists a decreasing function \(g:\mathbb{R}_{+}\rightarrow\mathbb{R}_{+}\) which measures the norm of deviations from initialization of models that achieve arbitrarily small training error._

_Formally, for almost surely all \(n\) training samples and for any sufficiently small \(\varepsilon>0\) there exists \(w^{(\varepsilon)}\in\mathbb{R}^{d^{\prime}}\) such that_

\[\widehat{F}(w^{(\varepsilon)})\leq\varepsilon,\;\;\;\text{and} \;\;\;g(\varepsilon)=\left\|w^{(\varepsilon)}-w_{0}\right\|.\]

Since Assumption 5 holds for arbitrarily small \(\varepsilon\), it guarantees that the model has enough capacity to interpolate the data, i.e., attain train error that is arbitrarily small (\(\varepsilon\)). Additionally, this is accomplished for model weights whose distance from initialization is managed by the function \(g(\varepsilon)\). By using these model weights to select \(w\) in Theorem 2 we obtain train loss bounds for interpolating models.

**Theorem 3** (Training loss under interpolation).: _Let Assumptions 1-5 hold. Let \(\eta\leq\min\{\frac{1}{L_{\widehat{F}}},g(1)^{2},\frac{g(1)^{2}}{\widehat{F}( w_{0})}\}\) and assume the width satisfies \(m\geq 18^{2}L^{2}R^{4}\,g(\frac{1}{T})^{4}\) for a fixed training horizon \(T\). Then,_

\[\widehat{F}(w_{T})\;\leq\frac{2}{T}+\frac{5\,g(\frac{1}{T})^{2}} {2\eta T}, \tag{6}\] \[\forall t\in[T]\;:\;\;\|w_{t}-w_{0}\|\;\leq\;4\,g(\frac{1}{T}).\]

To interpret the theorem's conclusions suppose that the function \(g(\cdot)\) of Assumption 5 is at most logarithmic; i.e., \(g(\frac{1}{T})=O(\log(T))\). Then, Theorem 3 implies that \(m=\Omega(\log^{4}(T))\) neurons suffice to achieve train loss \(\tilde{O}(\frac{1}{T})\) while GD iterates at all iterations satisfy \(\|w_{t}-w_{0}\|=O(\log(T))\). In Section 4 (see also Remark 1), we will give examples of data separability conditions that guarantee the desired logarithmic growth of \(g(\cdot)\) for logistic loss minimization, which in turn imply the favorable convergence guarantees described above. Under the same conditions we will show that the step-size requirement simplifies to
\(\min\{3,1/L_{\widehat{F}}\}\) (see Corollary 6.1). Finally, we remark that Theorem 3 provides sufficient parameterization conditions under which GD with \(T=\tilde{\Omega}(n)\) iterations finds weights \(w_{T}\) that yield an interpolating classifier and thus, achieve zero training error. To see this, assume logistic loss and observe setting \(T\gtrsim n\) in Eq. (6) gives \(\widehat{F}(w_{T})\leq\log(2)/n\). This in turn implies that every sample loss satisfies \(\widehat{F}_{i}(w_{T})\leq\log(2)\), equivalently \(y_{i}=\operatorname{sign}\left(\Phi(w_{T},x_{i})\right)\).

### Generalization

Our main result below bounds the generalization gap of GD for training two-layer nets with self-bounded loss functions. We remark that all expectations that appear below are over the training set.

**Theorem 4** (Generalization gap - General bound).: _Suppose Assumptions 1-4 hold. Fix any time horizon \(T\geq 1\) and any step size \(\eta\leq 1/L_{\widehat{F}}\) where \(L_{\widehat{F}}\) is the objective's smoothness parameter. Let any \(w\in\mathbb{R}^{d^{\prime}}\) such that \(\|w-w_{0}\|^{2}\geq\max\{\eta T\,\widehat{F}(w),\eta\widehat{F}(w_{0})\}.\) Suppose hidden-layer width \(m\) satisfies \(m\geq 64^{2}L^{2}R^{4}\|w-w_{0}\|^{4}\). Then, the generalization gap of GD at iteration \(T\) is bounded as_

\[\mathbb{E}\Big{[}F(w_{T})-\widehat{F}(w_{T})\Big{]}\leq\frac{8\ell^{2}R^{2}} {n}\,\mathbb{E}\left[\eta T\,\widehat{F}(w)+2\|w-w_{0}\|^{2}\right].\]

A few remarks regarding the theorem are in place. The theorem's assumptions are similar to those in Theorem 2, which bounds the training loss. The condition \(\|w-w_{0}\|^{2}\geq\max\{\eta T\widehat{F}(w),\eta\widehat{F}(w_{0})\}\) needs to hold almost surely over the training data, which is non-restrictive, as in later applications of the theorem, the choice of \(w\) arises from Assumption 5. The condition \(m\geq 64^{2}L^{2}R^{4}\|w-w_{0}\|^{4}\) on the width of the network, is also the same as that of Theorem 2 but with a larger constant. This means that the last-iterate train loss bound from Theorem 2 (Eq. (5)) holds under the setting of Theorem 4. Hence, it applies to the expected train loss \(\mathbb{E}[\widehat{F}(w_{T})]\) and, combined with the generalization-gap bound, yields a bound on the expected test loss \(\mathbb{E}[F(w_{T})]\).

To optimize the bound, a proper \(w\) must be selected by minimizing the population version of a ridge-regularized training objective. In interpolation settings, the procedure for selecting \(w\) follows the same guidelines as in Assumption 5 and in a similar style as obtaining Theorem 3.

**Theorem 5** (Generalization gap under interpolation).: _Let Assumptions 1-5 hold. Fix \(T\geq 1\) and let \(m\geq 64^{2}L^{2}R^{4}\,g(\frac{1}{T})^{4}\). Then, for any \(\eta\leq\min\{\frac{1}{L_{F}},g(1)^{2},\frac{g(1)^{2}}{\widehat{F}(w_{0})}\}\) the expected generalization gap at iteration \(T\) satisfies_

\[\mathbb{E}\Big{[}F(w_{T})-\widehat{F}(w_{T})\Big{]}\leq\frac{24\ell^{2}R^{2} \,g(\frac{1}{T})^{2}}{n}\,. \tag{7}\]

Note the width condition is similar in order to that of Theorem 3. Thus, provided \(g(\frac{1}{T})\lesssim\log(T)\) (see Remark 1 and Section 4 for examples), we have generalization gap of order \(\tilde{O}(\frac{1}{n})\) with \(m=\Omega(\log^{4}(T))\) neurons. Combined with the training loss guarantees from Theorem 3, we have test loss rate \(\tilde{O}(\frac{1}{T}+\frac{1}{n})\). This further implies that with \(m\approx\log^{4}(n)\) neurons and \(T=n\) iterations, the test loss reaches the optimal rate of \(\tilde{O}(\frac{1}{n})\). On the other hand, previous stability-based generalization bounds (e.g., [Richards and Rabbat, 2021]) required polynomial width \(m\gtrsim T^{2}\) and eventually obtained sub-optimal generalization rates of order \(O(\frac{T}{n})\). We further discuss the technical novelties resulting in these improvements in Section 5.

**Remark 1** (Example: Linearly-separable data).: _Consider logistic-loss minimization, \(\tanh\) activation \(\sigma(u)=\frac{e^{u}-e^{-u}}{e^{u}+e^{-u}}\) and data distribution that is linearly separable with margin \(\gamma\), i.e., for almost surely all \(n\) samples there exists unit-norm vector \(v^{\star}\in\mathbb{R}^{d}\) such that \(\forall i\in[n]:y_{i}\langle v^{\star},x_{i}\rangle\geq\gamma\). We initialize the weights to zero, i.e. \(w_{0}=0\) and show that the realizability Assumption 5 naturally holds in this setting. To see this, for any fixed \(\varepsilon>0,\) set \(\alpha=\frac{2(\log(1/\varepsilon))}{\gamma\sqrt{m}}\) and assume \(m\geq 4\log^{2}(1/\varepsilon)\). With this choice, select weights \(w_{j}^{(\varepsilon)}:=\alpha v^{\star},a_{j}=\frac{1}{\sqrt{m}}\) for \(j\in[1,\cdots,\frac{n}{2}]\) and \(w_{j}^{(\varepsilon)}:=-\alpha v^{\star},a_{j}=\frac{-1}{\sqrt{m}}\) for \(j\in\{\frac{m}{2}+1,\cdots,m\}\). Then, the
model output for any sample \((x_{i},y_{i})\) satisfies_

\[y_{i}\Phi(w^{(\varepsilon)},x_{i})=\frac{y_{i}\sqrt{m}}{2}\left(\sigma(\alpha \langle v^{\star},x_{i}\rangle)-\sigma(-\alpha\langle v^{\star},x_{i}\rangle) \right)=y_{i}\sqrt{m}\sigma(\alpha\langle v^{\star},x_{i}\rangle)\geq\sqrt{m} \sigma(\alpha\gamma)\geq\frac{\sqrt{m}}{2}\alpha\gamma=\log(\frac{1}{ \varepsilon})\]

_where the first equality uses the fact that \(tanh\) is odd, the first inequality follows by the increasing nature of \(tanh\) and data separability, and the last inequality follows since \(\alpha\gamma\leq 1\) and \(\sigma(u)\geq u/2\) for all \(u\in[0,1]\). Thus, the loss satisfies \(\widehat{F}(w^{(\varepsilon)})\leq\varepsilon\) since for the logistic function \(\log(1+e^{u})\leq e^{u}.\) Moreover, our choice of \(\alpha\) implies \(g(\varepsilon)=\|w^{(\varepsilon)}-w_{0}\|=\|w^{(\varepsilon)}\|=\alpha\sqrt {m}=2\log(1/\varepsilon)/\gamma\). To conclude, the NN-Realizability Assumption 5 holds with \(g(\varepsilon)=2\log(1/\varepsilon)/\gamma\) and thus applying Theorems 3, 5 shows that with \(m=\Omega(\log^{4}(T))\) neurons, the training loss and generalization gap are bounded by \(\tilde{O}(\frac{1}{\gamma^{2}T})\) and \(\tilde{O}(\frac{1}{\gamma^{2}n})\), respectively. We note that the same conclusion as above holds for other smooth activations such as Softmax or GELU._

## 4 On Realizability of NTK-Separable Data

In this section, we interpret our results for NTK-separable data by showing that our realizability condition holds for this class. We recall the definition of NTK-separability below (Nitanda et al., 2019; Chen et al., 2020; Cao and Gu, 2020).

**Assumption 6** (Separability by NTK).: _For almost surely all \(n\) training samples from the data distribution there exists \(w^{\star}\in\mathbb{R}^{d}\) and \(\gamma>0\) such that \(\|w^{\star}\|=1\) and for all \(i\in[n]\),_

\[y_{i}\Big{\langle}\nabla_{1}\Phi(w_{0},x_{i}),w^{\star}\Big{\rangle}\geq\gamma. \tag{8}\]

We also assume a bound on the model's output at initialization. Similar assumptions, but for the value of the loss, also appear in prior works that study generalization using the algorithmic stability framework (Richards and Kuzborskij, 2021; Lei et al., 2022).

**Assumption 7** (Initialization bound).: _There exists parameter \(C\) such that \(\forall i\in[n]:|\Phi(w_{0},x_{i})|\leq C,\) for almost surely all \(n\) training samples from the data distribution_

The next proposition relates the NTK-separability assumption to our realizability assumption. The proofs for this section are given in Appendix C.

**Proposition 6** (Realizability of NTK-separable data).: _Let Assumptions 1-2,6-7 hold. Assume \(f(\cdot)\) to be the logistic loss. Fix \(\varepsilon>0\) and let \(m\geq\frac{L^{2}R^{4}}{4\gamma^{4}C^{2}}(2C+\log(1/\varepsilon))^{4}\). Then the realizability Assumption 5 holds with \(g(\varepsilon)=\frac{1}{\gamma}(2C+\log(1/\varepsilon))\). In other words, there exists \(w^{(\varepsilon)}\) such that_

\[\widehat{F}(w^{(\varepsilon)})\leq\varepsilon,\ \ \text{and}\ \ \ \left\|w^{( \varepsilon)}-w_{0}\right\|=\frac{1}{\gamma}\left(2C+\log(1/\varepsilon) \right). \tag{9}\]

Having established realizability, the following is an immediate corollary of the general results presented in the last section.

**Corollary 6.1** (Results under NTK-separability).: _Let Assumptions 1-2,6-7 hold and assume logistic loss. Suppose \(m\geq\frac{64^{2}L^{2}R^{4}}{\gamma^{4}}(2C+\log(T))^{4}\) for a fixed training horizon \(T\). Then for any \(\eta\leq\min\{3,\frac{1}{L_{F}}\}\), the training loss and generalization gap are bounded as follows:_

\[\widehat{F}(w_{T})\leq\frac{5(2C+\log(T))^{2}}{\gamma^{2}\eta T},\] \[\mathbb{E}\left[F(w_{T})-\widehat{F}(w_{T})\right]\leq\frac{24 \ell^{2}R^{2}}{\gamma^{2}n}(2C+\log(T))^{2}.\]

A few remarks are in place regarding the corollary. By Corollary 6.1, we can conclude that the expected generalization rate of GD on logistic loss and NTK-separable data as per Assumption 6 is \(\tilde{O}(\frac{1}{n})\) provided
width \(m=\Omega(\log^{4}(T))\). Moreover, the expected training loss is \(\mathbb{E}[\tilde{F}(w_{T})]=\tilde{O}(\frac{1}{T})\). Thus, the expected test loss after \(T\) steps is \(\tilde{O}(\frac{1}{T}+\frac{1}{n})\). In particular for \(T=\Omega(n)\), the expected test loss becomes \(\tilde{O}(\frac{1}{n})\). This rate is optimal with respect to sample size and only requires polylogarithmic hidden width with respect to \(n\), specifically, \(m=\Omega(\log^{4}(n))\). Notably, it represents an improvement over prior stability results, e.g., (Richards and Rabbat, 2021) which required polynomial width and yielded suboptimal generalization rates of order \(O(T/n)\). It is worth noting that the test loss bound's dependence on the margin, particularly the \(\frac{1}{\gamma^{2}n}\)-rate obtained in our analysis, bears similarity to the corresponding results in the convex setting of linearly separable data recently established in (Shamir, 2021; Schliserman and Koren, 2022). Additionally, our results improve upon corresponding bounds for neural networks obtained via Rademacher complexity analysis (Ji and Telgarsky, 2020; Chen et al., 2020) which yield generalization rates \(\tilde{O}(\frac{1}{\sqrt{n}})\). Moreover, these works have a \(\gamma^{-8}\) dependence on margin for the minimum network width, whereas in Corollary 6.1 this is reduced to \(\gamma^{-4}\). We also note that in general, both \(\gamma\) and \(C\) may depend on the data distribution, the data dimension, or the nature of initialization. This is demonstrated in the next section where we apply the corollary above to the noisy XOR data distribution and Gaussian initialization.

**Remark 2** (Benefits of exponential tail).: _We have stated Corollary 6.1 for the logistic loss, which has an exponential tail behavior. For general self-bounded loss functions and by following the same steps, we can show a bound on generalization gap of order \(O(\frac{1}{n}(f^{-1}(\frac{1}{T}))^{2})\) provided \(m=\Omega((f^{-1}(\frac{1}{T}))^{4})\). Hence, the tail behavior of \(f\) controls both the generalization gap and minimum width requirement. In particular, under Assumption 6, polynomial losses with tail behavior \(f(u)\sim 1/u^{\beta}\) result in generalization gap \(O(T^{2/\beta}/n)\) for \(m=\Omega(T^{4/\beta})\). Thus, increasing the rate of decay \(\beta\) for the loss, improves both bounds on generalization and width. This suggests the benefits of self-bounded fast-decaying losses such as exponentially-tailed loss functions for which the dependence on \(T\) is indeed only logarithmic._

### Example: Noisy XOR data

Next, we specialize the results of the last section to the noisy XOR data distribution (Wei et al., 2019) and derive the corresponding margin and test-loss bounds. Consider the following \(2^{d}\) points,

\[x_{i}=(x_{i}^{1},x_{i}^{2},\cdots,x_{i}^{d})\in\{(1,0),(0,1),(-1,0),(0,-1)\} \times\{-1,1\}^{d-2},\]

where \(\times\) denotes the Cartesian product and the labels are determined as \(y_{i}=-1\) if \(x_{i}^{1}=0\) and \(y_{i}=1\) if \(x_{i}^{1}=\pm 1\). Moreover, consider normalization \(\overline{x}_{i}=\frac{1}{\sqrt{d-1}}x_{i}\) so that \(R=1.\) The noisy XOR data distribution is the uniform distribution over the set with elements \((\overline{x}_{i},y_{i})\). For this dataset and Gaussian initialization, (Ji and Telgarsky, 2020) have shown for ReLU activation that the NTK-separability assumption holds with margin \(\gamma=\Omega(1/d)\). In the next result, we compute the margin for activation functions that are convex, Lipshitz and locally strongly convex.

**Proposition 7** (Margin).: _Consider the noisy XOR data \((\overline{x}_{i},y_{i})\in\mathbb{R}^{d}\times\{\pm 1\}\). Assume the activation function is convex, \(\ell\)-Lipschitz and \(\mu\)-strongly convex in the interval \([-2,2]\) for some \(\mu>0\), i.e., \(\min_{t\in[-2,2]}\sigma^{\prime\prime}(t)\geq\mu\). Moreover, assume Gaussian initialization \(w_{0}\in\mathbb{R}^{d^{\prime}}\) with entries iid \(N(0,1)\). If \(m\geq\frac{80^{2}d^{3}\ell^{2}}{2\mu^{2}}\log(2/\delta)\), then with probability at least \(1-\delta\) over the initialization, the NTK-separability Assumption 6 is satisfied with margin \(\gamma=\frac{\mu}{80d}\)._

An interesting example of an activation function that satisfies the mentioned assumptions is the Softplus activation where \(\sigma(u)=\log(1+e^{u})\). This activation function has \(\mu=0.1\) and \(\ell=1\), and it is also smooth with \(L=1/4\). Therefore, the results on generalization and training loss presented in Corollary 6.1 hold for it. For noisy XOR data, Proposition 7 shows the margin in Assumption 6 is \(\gamma\gtrsim 1/d\). Additionally, for standard Gaussian initialization we have by Lemma C.5 that with high-probability the initialization bound in Assumption 7 satisfies \(C\lesssim\sqrt{d}\). Putting these together, and applying Corollary 6.1 shows that GD with \(n\) training samples reaches test loss rate \(\tilde{O}(\frac{d^{3}}{n})\) after \(T\approx n\) iterations and given \(m=\tilde{\Omega}(d^{6})\) neurons. It is worth noting that the number of training samples can be exponentially large with respect to \(d\). In this case the minimum width requirement is only polylogarithmic in \(n\).
Proof Sketches

We discuss here high-level proof ideas for both optimization and generalization bounds of Theorems 2 and 4. Formal proofs are deferred to Appendices A and B.

### Training loss

As already discussed in Section 3.1, the key insight we use to obtain bounds that are analogous to results for optimizing convex objectives, is to exploit the self-bounded weak convexity property of the objective in Eq. (4). Thanks to this property, the Hessian minimum eigenvalue \(\lambda_{\min}(\nabla^{2}\widehat{F}(w_{t}))\) becomes less negative at the same rate at which the train loss \(\widehat{F}(w_{t})\) decreases.

The technical challenge at formalizing this intuition arises as follows. Controlling the rate at which \(\widehat{F}(w_{t})\) converges to \(\widehat{F}(w)\) for the theorem's \(w\) requires controlling the Hessian at _all_ intermediate points \(w_{\alpha t}:=\alpha w_{t}+(1-\alpha)w,\alpha\in[0,1]\) between \(w\) and GD iterates \(w_{t}\). This is due to Taylor's theorem used to relate \(\widehat{F}(w_{t})\) to the target value \(\widehat{F}(w)\) as follows:

\[\widehat{F}(w)\geq\widehat{F}(w_{t})+\left\langle\nabla\widehat{F}(w_{t}),w-w _{t}\right\rangle+\frac{1}{2}\,\lambda_{\min}\left(\nabla^{2}\widehat{F}(w_{ \alpha t})\right)\left\|w-w_{t}\right\|^{2}.\]

Thus from self-bounded weak convexity, to control the last term above we need to control \(\widehat{F}(w_{\alpha t})\) for any intermediate point \(w_{\alpha t}\) along the GD trajectory. This is made possible by establishing the following generalized local quasi-convexity property.

**Proposition 8** (Generalized Local Quasi-Convexity).: _Suppose \(\widehat{F}:\mathbb{R}^{d^{\prime}}\to\mathbb{R}\) satisfies the self-bounded weak convexity property in Eq. (4) with parameter \(\kappa\). Let \(w_{1},w_{2}\in\mathbb{R}^{d^{\prime}}\) be two arbitrary points with distance \(\|w_{1}-w_{2}\|\leq D<\sqrt{2/\kappa}\). Set \(\tau:=\left(1-\kappa D^{2}/2\right)^{-1}\). Then,_

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

Recall that quasi-convex functions satisfy Eq. (10) with \(\tau=1\) and \(D\) can be unboundedly large. The Proposition 8 indicates that our neural-net objective function is approximately quasi-convex (since \(\tau>1\)) and this property holds locally, i.e. provided that \(w_{1},w_{2}\) are sufficiently close.

Applying (10) for \(w_{1}=w_{t},w_{2}=w\) allows controlling \(\widehat{F}(w_{\alpha t})\) in terms of the train loss \(\widehat{F}(w_{t})\) and the target loss \(\widehat{F}(w)\). The only additional requirement in Proposition 8 for this to hold is that

\[1/\kappa\propto\sqrt{m}\gtrsim\|w_{t}-w\|^{2}. \tag{11}\]

This condition exactly determines the required neural-net width. Formally, we have the following.

**Corollary 8.1** (GLQC of sufficiently wide neural nets).: _Let Assumptions 1,2, 4 hold. Fix arbitrary \(w_{1},w_{2}\in\mathbb{R}^{d^{\prime}}\), any constant \(\lambda>1\), and \(m\) large enough such that \(\sqrt{m}\geq\lambda\frac{LR^{2}}{2}\|w_{1}-w_{2}\|^{2}\). Then,_

\[\max_{v\in[w_{1},w_{2}]}\widehat{F}(v)\leq\left(1-1/\lambda\right)^{-1}\cdot \max\{\widehat{F}(w_{1}),\widehat{F}(w_{2})\}. \tag{12}\]

To conclude, using Corollary 8.1, we can show the regret bound in Eq. (5) provided (by (11)) that \(\sqrt{m}\gtrsim\|w_{t}-w\|^{2}\) is true for all \(t\in[T].\) To make the width requirement independent of \(w_{t}\), we then use a recursive argument to prove that \(\|w_{t}-w\|\leq 3\|w-w_{0}\|\). These things put together, lead to the parameter bound \(\|w_{t}-w_{0}\|\leq 4\|w-w_{0}\|\) and the width requirement \(\sqrt{m}\gtrsim\|w-w_{0}\|^{2}\) in the theorem's statement. We note that the GLQC property is also crucially required for the generalization analysis which we discuss next.
### Generalization gap

We bound the generalization gap using stability analysis (Bousquet and Elisseeff, 2002; Hardt et al., 2016). In particular, we use (Lei and Ying, 2020, Thm. 2) that relates the generalization gap to the "on average model stability". Formally, let \(w_{t}^{\neg i}\) denote the \(t\)-th iteration of GD on the leave-one-out loss \(\widehat{F}^{\neg i}(w):=\frac{1}{n}\sum_{j\neq i}\widehat{F}_{j}(w)\). As before, \(w_{t}\) denotes the GD output on full-batch loss \(\widehat{F}\). We will use the fact (see Corollary D.2.1) that \(f(y\Phi(\cdot,x))\) is \(G_{\widehat{F}}\)-Lipschitz with \(G_{\widehat{F}}=\ell R\) under Assumptions 2 and 3.A. Then, using (Lei and Ying, 2020, Thm. 2(a)) (cf. Lemma B.3) it holds that

\[\mathbb{E}\Big{[}F(w_{T})-\widehat{F}(w_{T})\Big{]}\leq 2G_{\widehat{F}}\ \mathbb{E}\Big{[}\frac{1}{n}\sum_{i=1}^{n}\lVert w_{T}-w_{T}^{\neg i}\rVert \Big{]}. \tag{13}\]

In order to bound the on-average model-stability term on the right-hand side above we need to control the degree of expansiveness of GD. Recall that for convex objectives GD is non-expansive (e.g. (Hardt et al., 2016)), that is \(\lVert(w-\eta\nabla\widehat{F}(w))-(w^{\prime}-\eta\nabla\widehat{F}(w^{ \prime}))\rVert\leq\lVert w-w^{\prime}\rVert\) for any \(w,w^{\prime}\). For the non-convex objective in our setting, the lemma below establishes a generalized non-expansiveness property via leveraging the structure of the objective's Hessian for the two-layer net.

**Lemma 9** (GD-Expansiveness).: _Let Assumptions 1 and 2 hold. For any \(w,w^{\prime}\in\mathbb{R}^{d^{\prime}}\), any step-size \(\eta>0\), and \(w_{\alpha}:=\alpha w+(1-\alpha)w^{\prime}\) it holds for \(H(w):=\eta\frac{LR^{2}}{\sqrt{m}}\widehat{F}^{\prime}(w)+\max\Big{\{}1,\eta \ell^{2}R^{2}\widehat{F}^{\prime\prime}(w)\Big{\}}\) that_

\[\Big{\lVert}\Big{(}w-\eta\nabla\widehat{F}(w)\Big{)}-\Big{(}w^{\prime}-\eta \nabla\widehat{F}(w^{\prime})\Big{)}\Big{\rVert}\leq\max_{\alpha\in[0,1]}H(w _{\alpha})\ \lVert w-w^{\prime}\rVert\,,\]

_where we define \(\widehat{F}^{\prime}(w):=\frac{1}{n}\sum_{i=1}^{n}\lvert f^{\prime}(y_{i}\Phi (w,x_{i}))\rvert\) and \(\widehat{F}^{\prime\prime}(w):=\frac{1}{n}\sum_{i=1}^{n}f^{\prime\prime}(y_{i }\Phi(w,x_{i}))\)._

This lemma can be further simplified for the class of self-bounded loss functions. Specifically, using \(\lvert f^{\prime}(u)\rvert\leq f(u)\) and \(f^{\prime\prime}(u)\leq 1\) from Assumptions 4 and 3.B, we immediately deduce the following.

**Corollary 9.1** (Expansiveness for self-bounded losses).: _In the setting of Lemma 9, further assume the loss satisfies Assumptions 3.B and 4. Provided \(\eta\leq 1/(\ell^{2}R^{2})\), it holds for all \(w,w^{\prime}\in\mathbb{R}^{d^{\prime}}\) that_

\[\Big{\lVert}\left(w-\eta\nabla\widehat{F}(w)\right)-\left(w^{\prime}-\eta \nabla\widehat{F}(w^{\prime})\right)\Big{\rVert}\leq\left(1+\eta\frac{LR^{2} }{\sqrt{m}}\max_{\alpha\in[0,1]}\widehat{F}(w_{\alpha})\right)\Big{\lVert}w- w^{\prime}\Big{\rVert}\,. \tag{14}\]

In Eq. (14) the expansiveness is weaker than in a convex scenario, where the coefficient would be \(1\) instead of \(1+\frac{\eta LR^{2}}{\sqrt{m}}\max_{\alpha\in[0,1]}\widehat{F}(w_{\alpha})\). However, for self-bounded losses (i.e. \(\lvert f^{\prime}(u)\rvert\leq f(u)\)) the "gap to convexity" \(\frac{LR^{2}}{\sqrt{m}}\max_{\alpha\in[0,1]}\widehat{F}(w_{\alpha})\) in Corollary 9.1 is better than the gap from Lemma 9 for \(1\)-Lipschitz losses (i.e. \(\lvert f^{\prime}(u)\rvert\leq 1\)), which would be \(\frac{\eta LR^{2}}{\sqrt{m}}\). Indeed, after unrolling the GD iterates, the latter eventually leads to polynomial width requirements (Richards and Rabbat, 2021).

Instead, to obtain a polylogarithmic width, we use the expansiveness bound in Eq. (14) for self-bounded losses together with the generalized-local quasi-convexity property in Corollary 8.1 as follows. From Corollary 8.1, if \(m\) is large enough such that

\[\sqrt{m}\geq LR^{2}\lVert w_{t}-w_{t}^{\neg i}\rVert^{2},\ \ \ \ \forall t\in[T],\ \forall i\in[n],\]

then Eq. (12) holds on the GD path. This further simplifies the result of Corollary 9.1 applied for \(w=w_{t},w^{\prime}=w_{t}^{\neg i}\) into

\[\Big{\lVert}(w_{t}-\eta\nabla\widehat{F}^{\neg i}(w_{t}))-(w_{t}^{\neg i}- \eta\nabla\widehat{F}^{\neg i}(w_{t}^{\neg i}))\Big{\rVert}\leq\widetilde{H }_{t}^{\,i}\left\lVert w_{t}-w_{t}^{\neg i}\right\rVert,\]

where \(\widetilde{H}_{t}^{\,i}:=1+\frac{2\eta LR^{2}}{\sqrt{m}}\max\{\widehat{F}^{ \neg i}(w_{t}),\widehat{F}^{\neg i}(w_{t}^{\neg i})\}.\) Now from the optimization analyses in Sec. 5.1, we know intuitively that \(\widehat{F}^{\neg i}(w_{t})\leq\widehat{F}(w_{t})\) decays at rate \(\tilde{O}(1/t)\); thus, so does \(\widehat{F}^{\neg i}(w_{t}^{\neg i})\). Therefore, for all \(i\in[n]\)
the expansivity coefficient \(\widehat{H}_{t}^{i}\) in the above display is decaying to \(1\) as GD progresses.

To formalize all these and connect them to the model-stability term in (13), note using triangle inequality and the Gradient Self-boundedness property of Lemma 1 that

\[\left\|w_{t}-w_{t}^{\neg i}\right\|\leq\left\|(w_{t}-\eta\nabla\widehat{F}^{ \neg i}(w_{t}))-(w_{t}^{\neg i}-\eta\nabla\widehat{F}^{\neg i}(w_{t}^{\neg i})) \right\|+\frac{\eta\ell R}{n}\widehat{F}_{i}(w_{t})\,.\]

Unrolling this display over \(t\in[T]\), averaging over \(i\in[n]\), and using our expansiveness bound above we show in Appendix B the following bound for the model stability term

\[\frac{1}{n}\sum_{i=1}^{n}\left\|w_{T}-w_{T}^{\neg i}\right\|\leq\frac{\eta \ell Re^{\beta}}{n}\sum_{t=0}^{T-1}\widehat{F}(w_{t})\,, \tag{15}\]

where \(\beta\lesssim\left(\sum_{t=1}^{T}\widehat{F}(w_{t})+\sum_{t=1}^{T}\widehat{F} ^{\neg i}(w_{t}^{\neg i})\right)/\sqrt{m}\,.\) But, we know from training-loss bounds in Theorem 2 that \(\sum_{t=1}^{T}\widehat{F}(w_{t})\lesssim\|w-w_{0}\|^{2}\) (and similar for \(\sum_{t=1}^{T}\widehat{F}^{\neg i}(w_{t}^{\neg i})\)). Thus, \(\beta\lesssim\|w-w_{0}\|^{2}/\sqrt{m}\). At this point, the theorem's conditions guarantees \(\sqrt{m}\gtrsim\|w-w_{0}\|^{2}\), so that \(\beta=O(1)\). Plugging back in (15) we conclude with the following stability bound: \(\frac{1}{n}\sum_{i=1}^{n}\|w_{T}-w_{T}^{\neg i}\|\lesssim\sum_{t=0}^{T} \widehat{F}(w_{t})/n.\) Applying the train-loss bounds of Theorem 2 once more completes the proof.

## 6 Prior Works

The theoretical study of generalization properties of neural networks (NN) is more than two decades old (Bartlett, 1996; Bartlett et al., 1998). Recently, there has been an increased interest in understanding and improving generalization of SGD/GD on over-parameterized neural networks, e.g. (Allen-Zhu et al., 2019; Oymak and Soltanolkotabi, 2020; Javanmard et al., 2020; Richards and Rabbat, 2021). These results however typically require very large width where \(m=\text{poly}(n)\). We discuss most-closely related-works below.

**Quadratic loss.** For quadratic loss, (Li and Liang, 2018; Soltanolkotabi et al., 2018; Allen-Zhu et al., 2019; Oymak and Soltanolkotabi, 2020; Liu et al., 2022) showed that sufficiently over-parameterized neural networks of polynomial width satisfy a local Polyak-Lojasiewicz (PL) condition \(\|\nabla\widetilde{F}(w)\|^{2}\!\!\geq 2\mu(\widetilde{F}(w)-\widetilde{F}^{ \star})\), where \(\mu\) is at least the smallest eigenvalue of the neural tangent kernel matrix. The PL property in this case implies that the training loss converges linearly with the rate \(\widehat{F}(w_{t})=O((1-\eta\mu)^{t})\) if the GD iterates remain in the PL region. Moreover, (Charles and Papailiopoulos, 2018; Lei and Ying, 2020), have used the PL condition to further characterize stability properties of corresponding non-convex models. Notably, (Lei and Ying, 2020) derived order-optimal rates \(O(\frac{1}{\mu n})\) for the generalization loss. However these rates only apply to quadratic loss. Models trained with logistic or exponential loss on separable data do _not_ satisfy the PL condition even for simple interpolating linear models. Aside from the PL condition-related results, but again for quadratic loss, (Oymak et al., 2019) showed under specific assumptions on the data translating to low-rank NTK, that logarithmic width is sufficient to obtain classification error of order \(O(n^{-1/4})\). In general, they achieve error rate \(O(n^{-1/2})\), but for \(m=\tilde{\Omega}(n^{2})\).

**Logistic-loss minimization with linear models.** Logistic-loss minimization is more appropriate for classification and rate-optimal generalization bounds for GD have been obtained recently in the linear setting, where the training objective is convex. In particular, for linear logistic regression on data that are linearly separable with margin \(\gamma>0\), (Shamir, 2021) proved a finite-time test-error bound \(O(\frac{\log^{2}T}{\gamma^{2}T}+\frac{\log^{2}T}{\gamma^{2}n})\). Ignoring log factors, this is order-optimal with the sample size \(n\) and training horizon \(T\). Their proof uses exponential-decaying properties of the logistic loss to control the norm of gradient iterates, which it cleverly combines with Markov's inequality to bound the fraction of well-separated datapoints at any iteration. This in turn translates to a test-error bound by standard margin-based generalization bounds. More recently, (Schliserman and Koren, 2022) used algorithmic-stability analysis proving same rates (up to log factors) for the test loss. Their results hold for general convex, smooth, self-bounded and decreasing objectives under a realizability assumption suited for convex objectives (analogous to Assumption 5). Specifically,
this includes linear logistic regression with linearly separable data. Here, we show that analogous rates on the test loss hold true for more complicated nonconvex settings where data are separable by shallow neural networks.

**Stability of GD in NN.** State-of-the-art generalization bounds on shallow neural networks via the stability-analysis framework have appeared very recently in (Richards and Rabbat, 2021; Richards and Kuzborskij, 2021; Lei et al., 2022). For Lipschitz losses, (Richards and Rabbat, 2021) shows that the empirical risk is weakly convex with a weak-convexity parameter that improves as the neural-network width \(m\) increases. Leveraging this observation, they establish stability bounds for GD iterates at time \(T\) provided sufficient parameterization \(m=\tilde{\Omega}(T^{2})\). Since the logistic loss is Lipschitz, these bounds also apply to our setting. Nevertheless, our work improves upon (Richards and Rabbat, 2021) in that: (i) we require significantly smaller width, poly-logarithmic rather than polynomial, and (ii) we show \(\tilde{O}(1/n)\) test loss bounds in the realizable setting, while their bounds are \(O(T/n)\). Central to our improvements is a largely refined analysis of the curvature of the loss via identifying and proving a generalized quasi-convexity property for neural networks of polylogarithmic width trained with self-bounded losses (see Section 5 for details). Our results also improve upon the other two works (Richards and Kuzborskij, 2021; Lei et al., 2022), which both require polynomial widths. However, we note that these results are not directly comparable since (Richards and Kuzborskij, 2021; Lei et al., 2022) focus on quadratic-loss minimization. See also Appendix E.

**Uniform convergence in NN.** Uniform bounds on the generalization loss have been derived in literature via Rademacher complexity analysis (Bartlett and Mendelson, 2002); see for example (Neyshabur et al., 2015; Arora et al., 2019; Golowich et al., 2020; Vardi et al., 2022; Frei et al., 2022) for a few results in this direction. These works typically obtain the bounds of order \(O(\frac{\mathcal{R}}{\sqrt{n}})\), where \(\mathcal{R}\) depends on the Rademacher complexity of the hypothesis space. Recent works by (Ji and Telgarsky, 2020; Chen et al., 2020) also utilized Rademacher complexity analysis to obtain test loss rates of \(O(1/\sqrt{n})\) under an NTK separability assumption (see also (Nitanda et al., 2019)) with polylogarithmic width requirement for shallow and deep networks, respectively. Instead, while maintaining minimal width requirements, we obtain test-loss rates \(\tilde{O}(1/n)\), which are order-optimal. Our approach, which is based on algorithmic-stability, is also different and uncovers new properties of the optimization landscape, including a generalized local quasi-convexity property. On the other hand, the analysis of (Ji and Telgarsky, 2020; Chen et al., 2020) applies to ReLU activation and bounds the test loss with high-probability over the sampling of the training set. Instead, we require smooth activations similar to other studies such as (Oymak et al., 2019; Chatterji et al., 2021; Bai and Lee, 2020; Nitanda et al., 2019; Richards and Rabbat, 2021; Richards and Kuzborskij, 2021; Lei et al., 2022) and we bound the test loss in expectation over the training set. Finally, we also note that data-specific generalization bounds for two-layer nets have also appeared recently in (Cao et al., 2022; Frei et al., 2022). However, those results require that data are nearly-orthogonal.

**Convergence/implicit bias of GD.** Convergence and implicit bias of GD for logistic/exponential loss functions on linear models and neural networks have been investigated in (Ji and Telgarsky, 2018; Soudry et al., 2018; Nacson et al., 2019; Lyu and Li, 2020; Chizat and Bach, 2020; Chatterji et al., 2021). In particular, (Lyu and Li, 2020; Ji and Telgarsky, 2020) have shown for homogeneous neural-networks that GD converges in direction to a max-margin solution. While certainly powerful, this implicit-bias convergence characterization becomes relevant only when the number \(T\) of GD iterations is exponentially large. Instead, our convergence bounds apply for finite \(T\) (on the order of sample size), thus are more practically relevant. Moreover, their results assume a GD iterate \(t_{0}\) such that \(\widehat{F}(w_{t_{0}})\leq\log(2)/n\). Similar assumption appears in (Chatterji et al., 2021), which require initialization \(\widehat{F}(w_{0})\leq 1/n^{1+C}\) for constant \(C>0\). Our approach is entirely different: we prove that sufficient parameterization benefits the loss curvature and suffices for GD steps to find an interpolating model and attain near-zero training loss, provided data satisfy an appropriate realizability condition.
Conclusions

In this paper we study smooth shallow neural networks trained with self-bounded loss functions, such as logistic loss. Under interpolation, we provide minimal sufficient parameterization conditions to achieve rate-optimal generalization and optimization bounds. These bounds improve upon prior results which require substantially large over-parameterization or obtain sub-optimal generalization rates. Specifically, we significantly improve previous stability-based analyses in terms of both relaxing the parameterization requirements and obtaining improved rates. Although our focus was on binary classification with shallow networks, our approach can be extended to multi-class settings and deep networks, which will be explored in future studies. Extending our results to the stochastic case by analyzing SGD is another important future direction. Moreover, while our current treatment relies on smoothness of the activation function to exploit properties of the curvature of the training objective, we aim to examine the potential of our results to extend to non-smooth activations. Finally, our generalization analysis bounds the expectation of the test loss (over data sampling) and it is an important future direction extending these guarantees to a high-probability setting.
