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

1 Introduction . 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 ] . 2 Our Contributions . 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 . 2.1 Generalisation Gap . 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 .

Summary: 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.