Stability & Generalisation of Gradient Descent for Shallow Neural Networks without the Neural Tangent Kernel
Year: 2021

Abstract:
We revisit on-average algorithmic stability of Gradient Descent (GD) for training overparameterised shallow neural networks and prove new generalisation and excess risk bounds without the Neural Tangent Kernel (NTK) or PolyakŁojasiewicz (PL) assumptions. In particular, we show oracle type bounds which reveal that the generalisation and excess risk of GD is controlled by an interpolating network with the shortest GD path from initialisation (in a sense, an interpolating network with the smallest relative norm). While this was known for kernelised interpolants, our proof applies directly to networks trained by GD without intermediate kernelisation. At the same time, by relaxing oracle inequalities developed here we recover existing NTK-based risk bounds in a straightforward way, which demonstrates that our analysis is tighter. Finally, unlike most of the NTK-based analyses we focus on regression with label noise and show that GD with early stopping is consistent.

heading: 1 Introduction
text: In a canonical statistical learning problem the learner is given a tuple of independently sampled training examples S = (z1, . . . , zn), where each example zi = (xi, yi) consists of an input xi and label yi jointly distributed according to some unknown probability measure P . In the following we assume that inputs belong to an Euclidean ball Bd2(Cx) of radius Cx and labels belong to [−Cy, Cy]. Based on training examples the goal of the learner is to select parameters W from some parameter spaceW in order to minimise the statistical risk
L(W) def= 1 2
∫ (fW(x)− y)2 dP
where fW is a predictor parameterised by W. The best possible predictor in this setting is the regression function f?, which is defined as f?(x) = ∫ y dPY |X=x, while the minimum possible risk
is equal to the noise-rate of the problem, which is given by σ2 = ∫ Z(f
?(x)− y)2 dP . In this paper we will focus on a shallow neural network predictor that takes the form
fW(x) def = m∑ k=1 ukφ (〈Wk,x〉) ( x ∈ Rd, W ∈ Rd×m ) defined with respect to some activation function φ : R → R, fixed output layer u ∈ {±1/√m}m, and a tunable (possibly randomised) hidden layer W. In particular, we will consider fWT , where the
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
hidden layer is obtained by minimising an empirical proxy of L called the empirical risk LS(W) def = (2n)−1 ∑n i=1(fW(xi)− yi)2 by running a Gradient Descent (GD) procedure: For t ∈ [T ] steps with initial parameters W0 and a step size η > 0, we have iterates Wt+1 = Wt − η∇Lt(Wt) where ∇Lt(W) is the first order derivative of Lt(W). Understanding the behaviour of the statistical risk for neural networks has been a long-standing topic of interest in the statistical learning theory [Anthony and Bartlett, 1999]. The standard approach to this problem is based on uniform bounds on the generalisation gap
Gen(WT ) def = L(WT )− LS(WT ) ,
which, given a parameter space W , involves controlling the gap for the worst possible choice of W ∈ W under some unknown data distribution. The theory then typically leads to the capacity-based (Rademacher complexity, VC-dimension, or metric-entropy based) bounds which hold with high probability (w.h.p.) over S [Bartlett and Mendelson, 2002, Golowich et al., 2018, Neyshabur et al., 2018]: 1
sup W∈W
| Gen(W)| . √
capacity(W) n . (1)
Thus, if one could simultaneously control the empirical risk and the capacity of the class of neural networks, one could control the statistical risk. Unfortunately, controlling the empirical risk LT turns out to be a challenging part here since it is non-convex, and thus, it is not clear whether GD can minimise it up to a desired precision.
This issue has attracted considerable attention in recent years with numerous works [Du et al., 2018, Lee et al., 2019, Allen-Zhu et al., 2019, Oymak and Soltanolkotabi, 2020] demonstrating that overparameterised shallow networks (in a sense m & poly(n)) trained on subgaussian inputs converge to global minima exponentially fast, namely, LS(WT ) . (1− η · d/n)T . Loosely speaking, these proofs are based on the idea that a sufficiently overparameterised network trained by GD initialised at W0 with Gaussian entries, predicts closely to a solution of a Kernelised Least-Squares (KLS) formulation minimised by GD, where the kernel function called the Neural Tangent Kernel (NTK) [Jacot et al., 2018] is implicitly given by the activation function. This connection explains the observed exponential rate in case of shallow neural networks: For the NTK kernel matrix K the convergence rate of GD is (1 − ηλmin(K))T , where λmin(K) is its smallest eigenvalue, and as it turns out, for subgaussian inputs λmin(K) & d/n [Bartlett et al., 2021]. 2 Naturally, the convergence was exploited to state bounds on the statistical risk: Arora et al. [2019] showed that for noise-free regression (σ2 = 0) when T & 1/(ηλmin(K)),
L(WT ) . √ 〈y, (nK)−1y〉
n . (2)
Clearly, the bound is non-vacuous whenever 〈 y,K−1y 〉 . n2α for some α ∈ [0, 1), and Arora et al. [2019] present several examples of smooth target functions which satisfy this. More generally, for f? which belongs to the Reproducing kernel Hilbert space (RKHS) induced by NTK one has〈 y,K−1y 〉 ≤ ‖f?‖2HNTK [Schölkopf and Smola, 2002]. The norm-based control of the risk is standard in the literature on kernels, and thus, one might wonder to which extent neural networks are kernel methods in disguise? At the same time, some experimental and theoretical evidence [Bai and Lee, 2019, Seleznova and Kutyniok, 2020, Suzuki and Akiyama, 2021] suggest that connection to kernels might be good only at explaining the behaviour of very wide networks, much more overparameterised than those used in practice. Therefore, an interesting possibility is to develop alternative ways to analyse generalisation in neural networks: Is there a more straightforward kernelfree optimisation-based perspective?
In this paper we take a step in this direction and explore a kernel-free approach, which at the same time avoids worst-case type uniform bounds such as Eq. (1). In particular, we focus on the notion of the algorithmic stability: If an algorithm is insensitive to replacement (or removal) of an observation
1Throughout this paper, we use f . g to say that there exists a universal constant c > 0 and some k ∈ N such that f ≤ cg logk(g) holds uniformly over all arguments.
2Which is a tightest known bound, Oymak and Soltanolkotabi [2020] prove a looser bound without distributional assumption on the inputs.
in a training tuple, then it must have a small generalisation gap. Thus, a natural question is whether GD is sufficiently stable when training overparameterised neural networks.
The stability of GD (and its stochastic counterpart) when minimising convex and non-convex smooth objective functions was first explored by Hardt et al. [2016]. Specifically, for a time-dependent choice of a step size ηt = 1/t and a problem-dependent constant α ∈ (0, 1) they show that
E [ Gen(WT ) ∣∣ W0,u] . ln(T )n−α . Unfortunately, when combined with the NTK-based convergence rate of the empirical risk we have a vacuous bound since LS(WT ) . 1. 3 This is because the stability is enforced through a quickly decaying step size rather than by exploiting a finer structure of the loss, and turns out to be insufficient to guarantee the convergence of the empirical risk. That said, several works Charles and Papailiopoulos [2018], Lei and Ying [2021] have proved stability bounds exploiting an additional structure in mildly non-convex losses. Specifically, they studied the stability of GD minimising a gradient-dominated empirical risk, meaning that for all W in some neighbourhood W and a problem-dependent quantity µ, it is assumed that LS(W)−minW′ LS(W′) ≤ ‖∇LS(W)‖2/(2µ). Having iterates of GD withinW and assuming that the gradient is ρ-Lipschitz, this allows to show
E [ Gen(WT ) ∣∣ W0,u] . ρ µ · 1 n .
As it turns out, the condition is satisfied for the iterates of GD training overparameterised networks [Du et al., 2018] with high probability over W0. The key quantity that controls the bound is ρ/µ . 1/λmin(K) which can be interpreted as a condition number of the NTK matrix. However, it is known that for subgaussian inputs the condition number behaves as n/d which renders the bound vacuous [Bartlett et al., 2021].

heading: 2 Our Contributions
text: In this paper we revisit algorithmic stability of GD for training overparameterised shallow neural networks, and prove new risk bounds without the Neural Tangent Kernel (NTK) or PolyakŁojasiewicz (PL) machinery. In particular, we first show a bound on the generalisation gap and then specialise it to state risk bounds for regression with and without label noise. In the case of learning with noise we demonstrate that GD with a form of early stopping is consistent, meaning that the risk asymptotically converges to the noise rate σ2.
Our analysis brings out a key quantity, which controls all our bounds, the Regularised Empirical Risk Minimisation (R-ERM) Oracle defined as
∆oracleS def = min
W∈Rd×m LS(W) +O
( ‖W −W0‖2F
ηT
) as ηT →∞ , (3)
which means that ∆oracleS is essentially an empirical risk of solution closest to initialisation (an interpolant when for m large enough). We first consider a bound on the generalisation gap.

heading: 2.1 Generalisation Gap
text: For simplicity of presentation in the following we assume that the activation function φ, its first and second derivatives are bounded. Assuming parameterisation m & (ηT )5 we show (Corollary 1) that the expected generalisation gap is bounded as
E [ Gen(WT ) ∣∣ W0,u] ≤ C · ηT n ( 1 + ηT n ) E [ ∆oracleS ∣∣ W0,u] , (4) where C is a constant independent from n, T, η.
3If ηs = 1s we have LS(WT ) . exp(µ ∑T j=1 1 j ) ≈ 1 Tµ , thus, if µ ≈ 1 n we then require T ∼ −n for LS(WT ) . . Plugging this into the Generalisation Error bound we get that log(T )n−α = n1−α log(1/ ) which is vacuous as n grows.
Dependence of m on the total number of steps T might appear strange at first, however things clear out once we pay attention to the scaling of the bound. Setting the step size η to be constant, T = nα, and overparameterisation m & n5α for some free parameter α ∈ (0, 1] we have
E [ Gen(WT ) ∣∣ W0,u] = O( 1 n E [ ‖Ŵ −W0‖2F ∣∣∣ W0,u]) as n→∞ , where Ŵ is chosen as parameters of a minimal-norm interpolating network, in a sense Ŵ ∈ argminW∈Rd×m{‖W −W0‖2F | LS(W) = 0}. Thus, as Eq. (4) suggests, the generalisation gap is controlled by a minimal relative norm of an interpolating network. In comparison, previous work in the NTK setting obtained results of a similar type where in place of Ŵ one has an interpolating minimal-norm solution to an NTK least-squares problem. Specifically, in Eq. (2) the generalisation gap is controlled by 〈 y, (nK)−1y 〉 , which is a squared norm of such solution.
The generalisation gap of course only tells us a partial story, since the behaviour of the empirical risk is unknown. Next, we take care of this and present bounds excess risk bounds.

heading: 2.2 Risk Bound without Label Noise
text: We first present a bound on the statistical risk which does not use the usual NTK arguments to control the empirical risk. For now assume that σ2 = 0 meaning that there is no label noise and randomness is only in the inputs.
NTK-free Risk Bound. In Corollary 2 we show that the risk is bounded as E [L(WT ) |W0,u] ≤ (
1 + C · ηT n
( 1 + ηT
n
)) E [ ∆oracleS ∣∣ W0,u] . (5) Note that this bound looks very similar compared to the bound on generalisation gap. The difference lies in a “1+” term which accounts for the fact that L(WT ) ≤ ∆oracleS as we show in Lemma 2. As before we let the step size be constant, set T = nα, and let overparameterisation be m & n5α for some α ∈ (0, 1]. Then our risk bound implies that
E [L(WT ) |W0,u] = O ( 1 nα E [ ‖Ŵ −W0‖2F ∣∣∣ W0,u]) as n→∞ , which comes as a simple consequence of bounding ∆oracleS while choosing Ŵ to be parameters of a minimal relative norm interpolating network as before. Recall that yi = f?(xi): The bound suggests that target functions f? are only learnable if n−α E[‖Ŵ −W0‖2F |W0,u] → 0 as n → ∞ for a fixed α and input distribution. A natural question is whether a class of such functions is non-empty. Indeed, the NTK theory suggests that it is not and in the following we will recover NTK-based results by relaxing our oracle inequality.
Similarly as in Section 2.1 we choose T = nα and so we require m & n5α, which trades off overparameterisation and convergence rate through the choice of α ∈ (0, 1]. To the best of our knowledge this is the first result of this type, where one can achieve a smaller overparameterisation compared to the NTK literature at the expense of a slower convergence rate of the risk. Finally, unlike the NTK setting we do not require randomisation of the initialisation and the risk bound we obtain holds for any W0. 4
Comparison to NTK-based Risk Bound. Our generalisation bound can be naturally relaxed to obtain NTK-based empirical risk convergence rates. In this case we observe that our analysis is general enough to recover results of Arora et al. [2019] (up to a difference that our bounds hold in expectation rather than in high probability).
By looking at Eq. (5), our task is in controlling ∆oracleS : by its definition in Eq. (3) we can see that it can be bounded by the regularised empirical risk of any model, and we choose a MoorePenrose pseudo-inverse solution to the least-squares problem supplied with the NTK feature map x 7→ (∇WfW(x))(W0). Here, entries of W0 are drawn from N (0, ν2init) for some appropriately
4In the NTK literature randomisation of W0 is required to guarantee that λmin(K) > 0 [Du et al., 2018].
chosen ν2init, and entries of the outer layer u are distributed according to the uniform distribution over {±1/√m}, independently from each other and the data. It is not hard to see that aforementioned pseudo-inverse solution Wpinv will have zero empirical risk, and so we are left with controlling ‖Wpinv −W0‖2F as can be seen from the definition of ∆oracleS . Note that it is straightforward to do since Wpinv has an analytic form. In Theorem 2 we carry out these steps and show that with high probability over initialisation (W0,u),
∆oracleS = ÕP ( 1
n
〈 y, (nK)−1y 〉) as n→∞ ,
which combined with Eq. (5) recovers the result of Arora et al. [2019]. Note that we had to relax the oracle inequality, which suggests that the risk bound we give here is tighter compared to the NTKbased bound. Similarly, Arora et al. [2019] demonstrated that ‖WT−W0‖2F = ÕP (〈 y, (nK)−1y 〉) holds w.h.p. over initialisation and ReLU activation function, however their proof requires a much more involved “coupling” argument where iterates (Wt)Tt=1 are shown to be close to the GD iterates minimising a KLS problem with NTK matrix.

heading: 2.3 Risk Bound with Label Noise and Consistency
text: So far we considered regression with random inputs, but without label noise. In this section we use our bounds to show that GD with early stopping is consistent in the regression with label noise. Here labels are generated as yi = f?(xi) + εi where zero-mean random variables (εi)ni=1 are i.i.d., almost surely bounded,5 and E[ε21] = σ 2.
For now we will use the same settings of parameters as in Section 2.2: Recall that η is constant, T = nα,m & n5α for a free parameter α ∈ (0, 1]. In addition assume a well-specified scenario which means that m is large enough such that for some subset of parametersW we achieve L(W?) = σ2 for some parameters W?. Employing our risk bound of Eq. (5), we relax R-ERM oracle as
E[∆oracleS |W0,u] . σ2 + n−α‖W? −W0‖2F , where W? ∈ argminW∈Rd×m L(W) .
Then we immediately have an α-dependent risk bound E [L(WT ) |W0,u] ≤ ( 1 + C
n1−α
)( σ2 +O ( ‖W? −W0‖2F
nα
)) as n→∞ . (6)
It is important to note that for any tuning α ∈ (0, 1), as long as W? is constant in n (which is the case in the parametric setting), we have E [L(WT ) |W0,u]→ σ2 as n→∞, and so we achieve consistency. Adaptivity to the noise achieved here is due to early stopping: Indeed, recall that we take T = nα steps and having smaller α mitigates the effect of the noise as can be seen from Eq. (6). This should come as no surprise as early stopping is well-known to have a regularising effect in GD methods [Yao et al., 2007]. The current literature on the NTK deals with this through a kernelisation perspective by introducing an explicit L2 regularisation [Hu et al., 2021], while risk bound of [Arora et al., 2019] designed for a noise-free setting would be vacuous in this case.

heading: 2.4 Future Directions and Limitations
text: We presented generalisation and risk bounds for shallow nets controlled by the Regularised Empirical Risk Minimisation oracle ∆oracleS . By straightforward relaxation of ∆ oracle S , we showed that the risk can be controlled by the minimal (relative) norm of an interpolating network which is tighter than the NTK-based bounds. There are several interesting venues which can be explored based on our results. For example, by assuming a specific form of a target function f? (e.g., a “teacher” network), one possibility is to characterise the minimal norm. This would allow us to understand better which such target function are learnable by GD.
One limitation of our analysis is smoothness of the activation function and its boundedness. While boundedness can be easily dealt with by localising the smoothness analysis (using a Taylor approximation around each GD iterate) and randomising (W0,u), extending our analysis to non-smooth activations (such as ReLU φ(x) = max(x, 0)) appears to be non-trivial because our stability analysis
5We assume boundedness of labels throughout our analysis to ensure that LS(W0) is bounded almost surely for simplicity. This can be relaxed by introducing randomised initialisation, e.g. as in [Du et al., 2018].
crucially relies on the control of the smallest eigenvalue of the Hessian. Another peculiarity of our analysis is a time-dependent parameterisation requirement m & (ηT )5. Making parameterisation n-dependent introduces an implicit early-stopping condition, which helped us to deal with label noise. On the other hand, such requirement might be limiting when dealing with noiseless interpolating learning where one could require to have ηT →∞ while keeping n finite.
Notation
In the following denote `(w, (x, y)) def= 12 (fW(x)−y) 2. Operator norm is denoted as ‖·‖op, while L2 norm is denoted by ‖ · ‖2 or ‖ · ‖F . We use notation (a∨ b) def = max {a, b} and (a∧ b) def= min {a, b} throughout the paper. Let (W(i)t )t be the iterates of GD obtained from the data set with a resampled data point:
S(i) def = (z1, . . . , zi−1, z̃i, zi+1, . . . , zn)
where z̃i is an independent copy of zi. Moreover, denote a remove-one version of S by
S\i def = (z1, . . . , zi−1, zi+1, . . . , zn) .

heading: 3 Main Result and Proof Sketch
text: This section presents both the main results as well as a proof sketch. Precisely, Section 3.1 presents the main result and Section 3.2 the proof sketch.

heading: 3.1 Main Results
text: We formally make the following assumption regarding the regularity of the activation function. Assumption 1 (Activation). The activation φ(u) is continuous and twice differentiable with constant Bφ, Bφ′ , Bφ′′ ≥ 0 bounding |φ(u)| ≤ Bφ, |φ′(u)| ≤ Bφ′ and |φ′′(u)| ≤ Bφ′′ for any u ∈ R.
This is satisfied for sigmoid as well as hyperbolic tangent activations. Assumption 2 (Inputs, labels, and the loss function). For constants Cx, Cy, C0 > 0, inputs belong to Bd2(Cx), labels belong to [−Cy, Cy], and loss is uniformly bounded by C0 almost surely.
A consequence of the above is an essentially constant smoothness of the loss and the fact that the Hessian scales with 1/ √ m (see Appendix A for the proof). These facts will play a key role in our stability analysis. 6
Lemma 1 (Smoothness and curvature). Fix W,W̃ ∈ Rd×m. Consider Assumption 1, Assumption 2, and assume that LS(W̃) ≤ C20 . Then, for any S,
λmax(∇2LS(W)) ≤ ρ where ρ def = C2x ( B2φ′ +Bφ′′Bφ +
Bφ′′Cy√ m
) ,
min α∈[0,1]
λmin(∇2LS(W̃ + α(W − W̃))) ≥ − Bφ′′
( Bφ′Cx + C0 ) √ m · (1 ∨ ‖W − W̃‖F ) .
Given these assumptions we are ready to present a bound on the Generalisation Error of the gradient descent iterates. Theorem 1 (Generalisation Error). Consider Assumptions 1 and 2. Fix t > 0. If η ≤ 1/(2ρ) and
m ≥ 144(ηt)2C4xC20B2φ′′ ( 4Bφ′Cx √ ηt+ √ 2 )2
(7)
then E [ Gen(Wt+1) ∣∣ W0,u] ≤ b( η n + η2t n2 ) t∑ j=0 E [LS(Wj) |W0,u]
where b = 16e3C 3 2 x B2φ′(1 + C 3 2 x B2φ′).
6We use notation (a ∨ b) def= max {a, b} and (a ∧ b) def= min {a, b} throughout the paper.
Proof. For full proof see Appendix C with sketch proof in Section 3.2.
Theorem 1 then provides a formal upper bound on the Generalisation Gap of a shallow neural network trained with GD. Specifically, provided the network is sufficiently wide the generalisation gap can be controlled through both: the gradient descent trajectory evaluated at the Empirical Risk∑T t=0 E[LS(Wt) |W0,u], as well as step size and number of iterations in the multiplicative factor. Similar bounds for the test performance of gradient descent have been given by Lei and Ying [2020], although in our case the risk is non-convex, and thus, a more delicate bound on the Optimisation Gap is required. This is summarised within the lemma, which is presented in an “oracle” inequality manner. Lemma 2 (Optimisation Error). Consider Assumptions 1 and 2. Fix t > 0. If η ≤ 1/(2ρ), then
1
t t∑ j=0 LS(Wj) ≤ min W∈Rd×m { LS(W) + ‖W −W0‖2F ηt + b̃‖W −W0‖3F√ m } + b̃C0 · (ηt) 3 2 √ m
where b̃ = C2xBφ′′ (Bφ′Cx + C0).
Proof. Full proof is given in Appendix B with sketch proof in Section 3.2.
By combining Theorem 1 and Lemma 2 we get the following where ∆oracleS is defined in Eq. (3): Corollary 1. Assume the same as in Theorem 1 and Lemma 2. Then,
E [ Gen(WT ) ∣∣ W0,u] ≤ C · ηT n ( 1 + ηT n ) E [ ∆oracleS ∣∣ W0,u] where ∆oracleS is defined in Eq. (3) and C is a constant independent from n, T, η.
The above combined with Lemma 2, and the fact that LS(WT ) = mint∈[T ] LS(Wt) gives us: Corollary 2. Assume the same as in Theorem 1 and Lemma 2. Then,
E [L(WT ) |W0,u] ≤ (
1 + C · ηT n
( 1 + ηT
n
)) E [ ∆oracleS ∣∣ W0,u] . Proof. The proof is given in Appendix E.
Finally, ∆oracleS is controlled by the norm of the NTK solution, which establishes connection to the NTK-based risk bounds: Theorem 2 (Connection between ∆oracleS and NTK). Consider Assumption 1 and that ηT = n. Moreover, assume that entries of W0 are i.i.d., that λmin(K) & 1/n, and assume that u ∼ unif ({±1/ √ m})m independently from all sources of randomness. Then, with probability least 1− δ for δ ∈ (0, 1), over (W0,u),
∆oracleS = ÕP ( 1
n
〈 y, (nK)−1y 〉) as n→∞ .
Proof. The proof is given in Appendix D.

heading: 3.2 Proof Sketch of Theorem 1 and Lemma 2
text: Throughout this sketch let expectation be understood as E[·] = E[· |W0,u]. For brevity we will also vectorise parameter matrices, so that W ∈ Rdm and thus ‖ · ‖2 = ‖ · ‖F . Let us begin with the bound on the generalisation gap, for which we use the notion of algorithmic stability [Bousquet and Elisseeff, 2002]. With W(i)t denoting a GD iterate with the dataset with the resampled data point S(i), the Generalisation Gap can be rewritten [Shalev-Shwartz and Ben-David, 2014, Chap. 13]
Gen = E[L(WT )− LS(WT )] = 1
n n∑ i=1 E[`(WT , z̃i)− `(W(i)T , z̃i)] .
The equality suggests that if trained parameters do not vary much when a data point is resampled, that is WT ≈ W(i)T , then the Generalisation Gap will be small. Indeed, prior work [Hardt et al., 2016, Kuzborskij and Lampert, 2018] has been dedicated to proving uniform bounds on L2 norm maxi∈[n] ‖WT −W (i) T ‖2. In our case, we consider the more refined squared L2 expected stability and consider a bound for smooth losses similar to that of Lei and Ying [2020]: Lemma 3. Consider Assumptions 1 and 2. Then,
Gen . √ E [LS(WT )] √√√√ 1 n n∑ i=1 E [ ‖WT −W(i)T ‖2op ] + 1 n n∑ i=1 E [ ‖WT −W(i)T ‖ 2 op ] .
Proof. The proof is given in Appendix C.1.
To this end, we bound ‖WT −W(i)T ‖22 recursively by using the definition of the gradient iterates and applying Young’s inequality (a+ b)2 ≤ (1 + 1t )a 2 + (1 + t)b2 to get that for any i ∈ [n]:
‖Wt+1 −W(i)t+1‖22 (8)
= ‖Wt −W(i)t − η(∇LS(Wt)−∇LS(i)(W (i) t ))‖22
= ‖Wt −W(i)t − η(∇LS\i(Wt)−∇LS\i(W (i) t )) +
η n (∇`(Wt, zi)−∇`(W(i)t , z̃i)‖22
≤ (1 + 1 t ) ‖Wt −W(i)t − η
( ∇LS\i(Wt)−∇LS\i(W (i) t ) ) ‖22︸ ︷︷ ︸
Expansiveness of the Gradient Update
+ η2(1 + t)
n2
( ‖∇`(Wt, zi)‖22 + ‖∇`(W (i) t , z̃i)‖22 ) . (9)
The analysis of the expansiveness of the gradient update then completes the recursive relationship [Hardt et al., 2016, Kuzborskij and Lampert, 2018, Lei and Ying, 2020, Richards and Rabbat, 2021]. Clearly, a trivial handling of the term (using smoothness) would yield a bound (1 + ηρ)2‖Wt −W(i)t ‖22, which leads to an exponential blow-up when unrolling the recursion. So, we must ensure that the expansivity coefficient is no greater than one. While the classic analysis of Hardt et al. [2016] controls this by having a polynomially decaying learning rate schedule, here we leverage on observation of Richards and Rabbat [2021] that the negative Eigenvalues of the loss’s Hessian∇2L
S \i (·) ∈ Rdm×dm can control the expansiveness. This boils down to Lemma 1:
λmin (∫ 1 0 ∇2LS\i(W (i) t + α(Wt −W (i) t )) dα ) & −1 ∨ ‖Wt −W (i) t ‖2√ m & − √ ηt m
where ‖Wt −W(i)t ‖2 . √ ηt comes from smoothness and standard analysis of GD (see Lemma 6).
Note that the eigenvalue scales 1/ √ m because Hessian has a block-diagonal structure.
Controlling expansiveness of GD updates. Abbreviate ∆ = ∇LS\i(Wt) −∇LS\i(W (i) t ) and ∇2 = ∫ 1 0 ∇2LS\i(W (i) t + α(Wt −W (i) t )) dα. Now we “open” the squared norm
‖Wt −W(i)t − η∆‖22 = ‖Wt −W (i) t ‖22 + η2 ‖∆‖ 2 2 − 2η 〈 Wt −W(i)t ,∆ 〉 and use Taylor’s theorem for gradients to have
η2 ‖∆‖22 − 2η 〈 Wt −W(i)t ,∆ 〉 = 〈 ∆, ( η2∇2 − 2ηI ) (Wt −W(i)t ) 〉 = 〈 Wt −W(i)t , η∇2 ( η∇2 − 2I ) (Wt −W(i)t )
〉 . ( η √ ηt
m + η2
ηt
m
) ‖Wt −W(i)t ‖22
where we noted that (η∇2 − 2I) has only negative eigenvalues due to assumption η ≤ 1/(2ρ).
By rearranging this inequality one can verify that the gradient operator is approximately cocoercive [Hardt et al., 2016]. This result, crucial to our proof, is supported by some empirical evidence: In Fig. 1 we show a synthetic experiment where we train a shallow neural network with sigmoid activation. Our handling of expansiveness eventually gives
‖Wt −W(i)t − η∆‖22 .
( 1 + η √ ηt
m
) ‖Wt −W(i)t ‖22 .
Plugging this back into Eq. (8)-(9) we unroll the recursion:
‖Wt+1 −W(i)t+1‖22 . η2t
n2 t∑ j=0 ( 1 + 1 t )t( 1 + η √ ηt m )t ( ‖∇`(Wj , zi)‖22 + ‖∇`(W (i) j , z̃i)‖ 2 2 ) .
The above suggest that to prevent exponential blow-up (and ensure that the expansiveness remains under control), it’s enough to have η √
ηt m ≤ 1/t, which is an assumption of the theorem, enforcing
relationship between ηt and m.
On-average stability bound. Now we are ready to go back to Lemma 3 and complete the proof of the generalisation bound. Since we have a squared loss we have ‖∇`(Wj , zi)‖22 . `(Wj , zi), and thus, when summing over i ∈ [n] we recover the Empirical risk 1n ∑n i=1 ‖∇`(Wj , zi)‖22 . LS(Wj). Finally, noting that E[`(Wj , zi)] = E[`(W (i) j , z̃i)] we then arrive at the following bound on the expected squared L2 stability
1
n n∑ i=1 E[‖Wt+1 −W(i)t+1‖22] . η2t n2 t∑ j=0 E[LS(Wj)] .
Note that the Generalisation Gap is directly controlled by the Optimisation Performance along the path of the iterates, matching the shape of bounds by Lei and Ying [2020]. Using smoothness of the loss as well as that LS(WT ) ≤ 1T ∑T j=0 LS(Wj), the risk is bounded when m & (ηt)3 as
E[L(WT )] . ( 1 + ηT
n
( 1 + ηT
n
)) 1
T T∑ j=0 E[LS(Wj)] .
Optimisation error. We must now bound the Optimisation Error averaged across the iterates. Following standard optimisation arguments for gradient descent on a smooth objective we then get
LS(Wt+1) ≤ LS(Wt)− η
2 ‖∇LS(Wt)‖22
At this point convexity would usually be applied to upper bound LS(Wt) ≤ LS(Ŵ) − 〈∇LS(Wt),Ŵ −Wt〉 where Ŵ is a minimiser of LS(·). The analysis differs from this approach in two respects. Firstly, we do not consider a minimiser of LS(·) but a general point W, which is then optimised over at the end. Secondly, the objective is not convex, and therefore, we leverage that the Hessian negative Eigenvalues are on the order of − 1√
m to arrive at the upper bound (Lemma 1):
LS(Wt) . LS(W)− 〈∇LS(Wt),W −Wt〉︸ ︷︷ ︸ Convex Component + 1√ m ‖W −Wt‖32︸ ︷︷ ︸
Negative Hessian Eigenvalues
After this point, the standard optimisation analysis for gradient descent can be performed i.e. −〈∇LS(Wt),W −Wt〉 − η2‖LS(Wt)‖ 2 2 = 1 η ( ‖W −Wt‖22 − ‖W −Wt+1‖22 ) , which then yields to a telescoping sum over t = 0, 1, . . . , T − 1. This leads to the upper bound for W ∈ Rm×d
1
T T∑ j=0 E[LS(Wj)] . LS(W) + ‖W −W0‖22 ηT + 1√ m · 1 T T∑ j=0 ‖W −Wj‖32
where we plug this into the generalisation bound and take a minimum over W ∈ Rmd.
1+e〈W?,x〉
where W? is once sampled from N (0, I10) at the beginning of the training. The initialisation W0 is sampled from N (0, I10) and each experiment is performed 10 times (which is reflected by standard deviation). The shallow network is then trained by GD with η = 1, T = 104, and sigmoid activation.

heading: 4 Additional Related Literature
text: Stability of Gradient Descent. Algorithmic stability [Bousquet and Elisseeff, 2002] has been used to investigate the generalisation performance of GD in a number of works [Hardt et al., 2016, Kuzborskij and Lampert, 2018, Chen et al., 2018, Lei and Ying, 2020, Richards and Rabbat, 2021]. While near optimal bounds are achieved for convex losses, the non-convex case aligning with this work is more challenging. Within [Hardt et al., 2016, Kuzborskij and Lampert, 2018] in particular a restrictive 1/t step size is required if t iterations of gradient descent are performed. This was partially alleviated within [Richards and Rabbat, 2021] which demonstrated that the magnitude of Hessian’s negative eigenvalues can be leveraged to allow for much larger step sizes to be taken. In this work we also utilise the magnitude of Hessian’s negative Eigenvalue, although a more delicate analysis is required as both: the gradient is potentially unbounded in our case; and the Hessian’s negative eigenvalues need to be bounded when evaluated across multiple points. We note the scaling of the Hessian network has also previously been observed within [Liu et al., 2020a,b], with empirical investigations into the magnitude of the Hessian’s negative Eigenvalues conducted within [Sagun et al., 2016, Yuan et al., 2019]. Stability-based generalisation bounds were also shown for certain types of non-convex functions, such Polyak-Łojasiewicz functions [Charles and Papailiopoulos, 2018, Lei and Ying, 2021] and functions with strict saddles among stationary points [Gonen and Shalev-Shwartz, 2017]. In line with our results, Rangamani et al. [2020] shows that for interpolating kernel predictors, minimizing the norm of the ERM solution minimizes stability.
Several works have also recently demonstrated consistency of GD with early stopping for training shallow neural networks in the presence of label noise. The concurrent work [Ji et al., 2021] also showed that shallow neural networks trained with gradient descent are consistent, although their approach distinctly different to ours e.g. leveraging structure of logistic loss as well as connections between shallow neural networks and random feature models / NTK. Earlier, Li et al. [2020] showed consistency under a certain gaussian-mixture type parametric classification model. In a more general nonparametric setting, Kuzborskij and Szepesvári [2021] showed that early-stopped GD is consistent when learning Lipschitz regression functions.
Certain notions of stability, such the uniform stability (taking sup over the data rather than expectation) allows to prove high-probability risk bounds (h.p.) [Bousquet and Elisseeff, 2002] which are known to be optimal up to log-terms [Feldman and Vondrak, 2019, Bousquet et al., 2020]. In this work we show bounds in expectation, or equivalently, control the first moment of a generalisation gap. To prove a h.p. bounds while enjoying benefits of a stability in expectation, one would need to control higher order moments [Maurer, 2017, Abou-Moustafa and Szepesvári, 2019], however often this is done through a higher-order uniform stability: Since our proof closely relies on the stability in expectation, it is not clear whether our results can be trivially stated with high probability.

heading: Acknowledgements
text: D.R. is supported by the EPSRC and MRC through the OxWaSP CDT programme (EP/L016710/1), and the London Mathematical Society ECF-1920-61.

Review summaries:
This paper presents a bound on the expected risk of a two-layer neural network with smooth activation, where the first layer is trained by gradient descent. The bound is in terms of the smallest L2 distance of the weights to initialization among interpolating solutions. When choosing the total number of iterations to be $n^\alpha$ ($\alpha \in(0,1]$), and the width to be bigger than $n^{5\alpha}$, the bound essentially scales like $n^{-\alpha}$ times the squared L2 distance. The analysis is based on stability, and relies on controlling the smallest eigenvalue of the Hessian of the loss.

This paper derives bounds for generalization gap of two-layer neural networks using algorithmic stability. The bounds covers noisy and noiseless cases, and depend on the distance of the interpolation solution which is closest to the initialization. With properly chosen number of iterations and network width (with respect to the number of data n), the bounds decay as n increases.   The bounds requires the width of the networks to be bigger than a polynomial of the number of iterations (specifically, number of iterations times the step size). Hence, it works out of the NTK regime. Through, the bounds can also yield reasonable results in the NTK regime. 

The authors propose to apply stability-based bound on the overparameterized neural networks and reach non-vacuous bound with early stopping. The authors show that the excess risk can be bounded by an interpolating network with the shortest GD path from initialization. The proof can be extended to NTK regimes and brings some interesting conclusions, e.g., GD with early stopping is consistent. Overall, I think this paper is interesting, and I tend to accept this paper.