Towards Understanding Generalization via Decomposing Excess Risk Dynamics

1 INTRODUCTION . Generalization is one of the essential mysteries uncovered in modern machine learning ( Neyshabur et al. , 2014 ; Zhang et al. , 2016 ; Kawaguchi et al. , 2017 ) , measuring how the trained model performs on unseen data . One of the most popular approaches to generalization is uniform convergence ( Mohri et al. , 2018 ) , which takes supremum over parameter space to decouple the dependency between the training set and the trained model . However , Nagarajan & Kolter ( 2019 ) pointed out that uniform convergence itself might not be powerful enough to explain generalization , because the uniform bound can still be vacuous under overparameterized linear regimes even if we reduce the parameter space to the minimum . One alternative solution beyond uniform convergence is to analyze the generalization dynamics , which measures the generalization gap during the training dynamics . Stability-based bound is among the most popular techniques in generalization dynamics analysis ( Lei & Ying , 2020 ) , which is derived from algorithmic stability ( Bousquet & Elisseeff , 2002 ) . Fortunately , one can derive nonvacuous bounds under general convex regimes using stability frameworks ( Hardt et al. , 2016 ) . However , stability is still far from explaining the remarkable generalization abilities of neural networks , mainly due to two obstructions . Firstly , stability-based bound depends heavily on the gradient norm in non-convex regimes ( Li et al. , 2019 ) , which is typically large at the beginning phase in training neural networks . Secondly , stability-based bound usually does not work well under general nonconvex regimes ( Hardt et al. , 2016 ; Charles & Papailiopoulos , 2018 ; Zhou et al. , 2018b ) but neural networks are usually highly non-convex . The aforementioned two obstructions mainly stem from the coarse-grained analysis of the signal and noise . As Zhang et al . ( 2016 ) argued , neural networks converge fast when fitting signal but converge relatively slowly when fitting noise1 , indicating that the training dynamics over signal and noise are significantly different . Consequently , on the one hand , the fast convergence of signal-related training contributes to a large gradient norm at the beginning phase ( see Figure 1a ) , resulting in 1In this paper , we refer the signal to the clean data without the output noise , and the noise to the output noise . See Section 2 for the formal definitions . poor stability . On the other hand , the training on signal forces the trained parameter away from the initialization , making the whole training path highly non-convex ( see Figure 1b ) . The above two phenomena inspire us to decompose the training dynamics into noise and signal component and only apply the stability-based analysis over the noise component . To demonstrate that such decomposition generally holds in practice , we conduct several experiments of neural networks on both synthetic dataset and real-world dataset ( see Figure 2 for more details ) . Based on the above discussion , we improve the stability-based analysis by proposing a decomposition framework on excess risk dynamics2 , where we handle the noise and signal components separately via bias-variance decomposition . In detail , we decompose the excess risk into variance excess risk ( VER ) and bias excess risk ( BER ) , where VER measures how the model fits noise and BER measures how the model fits signal . Under the decomposition , we apply the stability-based techniques to VER and apply uniform convergence to BER inspired by Negrea et al . ( 2020 ) . The decomposition framework accords with the theoretical and experimental evidence surprisingly well , providing that it outperforms stability-based bounds in both linear ( overparameterized linear regression ) and non-linear ( diagonal matrix recovery ) regimes . We summarize our contributions as follows : • We propose a new framework aiming at improving the traditional stability-based bounds , which is a novel approach to generalization dynamics analysis focusing on decomposing the excess risk dynamics into variance component and bias component . Starting from the overparameterized linear regression regimes , we show how to deploy the decomposition framework in practice , and the proposed framework outperforms the stability-based bounds . • We theoretically analyze the excess risk decomposition beyond linear regimes . As a case study , we derive a generalization bound under diagonal matrix recovery regimes . To our best knowledge , this is the first work to analyze the generalization performance of diagonal matrix recovery . • We conduct several experiments on both synthetic datasets and real-world datasets ( MINIST , CIFAR-10 ) to validate the utility of the decomposition framework , indicating that the framework provides interesting insights into the generalization community . 2We decompose the excess risk , which is closely related to generalization , purely due to technical reasons . The excess risk dynamics tracks the excess risk during the training process . 1.1 RELATED WORK . Stability-based Generalization . The stability-based researches can be roughly split into two branches . One branch is about how algorithmic stability leads to generalization ( Feldman & Vondrak , 2018 ; 2019 ; Bousquet et al. , 2020 ) . Another branch focus on how to calculate the stability parameter for specific problems , e.g. , Hardt et al . ( 2016 ) prove a generalization bound scales linearly with time in convex regimes . Furthermore , researchers try to apply stability techniques into more general settings , e.g. , non-smooth loss ( Lei & Ying , 2020 ; Bassily et al. , 2020 ) , noisy gradient descent in non-convex regimes ( Mou et al. , 2018 ; Li et al. , 2019 ) and stochastic gradient descent in non-convex regimes ( Zhou et al. , 2018b ; Charles & Papailiopoulos , 2018 ; Zhang et al. , 2021 ) . In this paper , we mainly focus on applying the decomposition framework to improve the stability-based bound . Uniform Convergence is widely used in generalization analysis . For bounded losses , the generalization gap is tightly bounded by its Rademacher Complexity ( Koltchinskii & Panchenko , 2000 ; Koltchinskii , 2001 ; Koltchinskii et al. , 2006 ) . Furthermore , we reach faster rates under realizable assumption ( Srebro et al. , 2010 ) . A line of work focuses on uniform convergence under neural network regimes , which is usually related to the parameter norm ( Bartlett et al. , 2017 ; Wei & Ma , 2019 ) . However , as Nagarajan & Kolter ( 2019 ) pointed out , uniform convergence may be unable to explain generalization . Therefore , more techniques are explored to go beyond uniform convergence . Other Approaches to Generalization . There are some other approaches to generalization , including PAC-Bayes ( Neyshabur et al. , 2017a ; Dziugaite & Roy , 2017 ; Neyshabur et al. , 2017b ; Dziugaite & Roy , 2018 ; Zhou et al. , 2018a ; Yang et al. , 2019 ) , information-based bound ( Russo & Zou , 2016 ; Xu & Raginsky , 2017 ; Banerjee & Montúfar , 2021 ; Haghifam et al. , 2020 ; Steinke & Zakynthinou , 2020 ) , and compression-based bound ( Arora et al. , 2018 ; Allen-Zhu et al. , 2018 ; Arora et al. , 2019 ) . Bias-Variance Decomposition . Bias-variance decomposition plays an important role in statistical analysis ( Lehmann & Casella , 2006 ; Casella & Berger , 2021 ; Geman et al. , 1992 ) . Generally , high bias indicates that the model has poor predicting ability on average and high variance indicates that the model performs unstably . Bias-variance decomposition is widely used in machine learning analysis , e.g. , adversarial training ( Yu et al. , 2021 ) , double descent ( Adlam & Pennington , 2020 ) , uncertainty ( Hu et al. , 2020 ) . Oymak et al . ( 2019 ) applied bias-variance decomposition on Jacobian of neural networks to explain their different performances on clean and noisy data . This paper considers a slightly different bias-variance decomposition following the analysis of SGD ( Dieuleveut et al. , 2016 ; Jain et al. , 2018 ; Zou et al. , 2021 ) , focusing on the decomposition of the noisy output . Matrix Recovery . Earlier methods for solving matrix recovery problems rely on the convex relaxation techniques for minimum norm solutions ( Recht et al. , 2010 ; Chandrasekaran et al. , 2011 ) . Recently , a branch of works focuses on the matrix factorization techniques with simple local search methods . It has been shown that there is no spurious local minima in the exact regime ( Ge et al. , 2016 ; 2017 ; Zhang et al. , 2019 ) . In the overparameterized regimes , it was first conjectured by Gunasekar et al . ( 2018 ) and then answered by Li et al . ( 2018 ) that gradient descent methods converge to the low-rank solution efficiently . Later , Zhuo et al . ( 2021 ) extends the conclusion to the noisy settings . 2 PRELIMINARY . In this section , we introduce necessary definitions , assumptions , and formally introduce previous techniques in dealing with generalization . Data distribution . Let x ∈ X ⊂ Rp be the input and y ∈ Y ⊂ R be the output , where x , y are generated from a joint distribution ( x , y ) ∼ P . Define Px , Py , and Py|x as the corresponding marginal and conditional distributions , respectively . Given n training samples D ≜ { xi , yi } i∈ [ n ] generated from distribution P , we denote its empirical distribution by Pn . To simplify the notations , we define X ∈ Rn×p as the design matrix and Y ∈ Rn as the response vector . Excess Risk . Given the loss function ℓ ( θ ; x , y ) with parameter θ and sample ( x , y ) , we define the population loss as L ( θ ; P ) ≜ E ( x , y ) ∼P [ ℓ ( θ ; x , y ) ] and its corresponding training loss as L ( θ ; Pn ) ≜ 1n ∑ i ℓ ( θ ; xi , yi ) . Let At denote the optimization algorithm which takes dataset D as input and return the trained parameter θ̂ ( t ) at time t , namely , At ( D ) = θ̂ ( t ) . During the analysis , we focus on the excess risk dynamics EL ( θ̂ ( t ) ; P ) , which measures how the trained parameter θ̂ ( t ) performs on the population loss : EL ( θ̂ ( t ) ; P ) ≜ L ( θ̂ ( t ) ; P ) −min θ L ( θ ; P ) . Although the minimizer can be not unique , we define L ( θ∗ ; Pn ) ≜ minθ L ( θ ; P ) where θ∗ denotes arbitrarily one of the minimizers . We additionally remark that bounding the generalization gap L ( θ̂ ( t ) ; P ) −L ( θ̂ ( t ) ; Pn ) suffices to bound the excess risk under Empirical Risk Minimization ( ERM ) framework by the following split ( see Equation 1 ) . Therefore , we mainly discuss the excess risk following Bartlett et al . ( 2020 ) . EL ( θ̂ ( t ) ; P ) = [ L ( θ̂ ( t ) ; P ) − L ( θ̂ ( t ) ; Pn ) ] ︸ ︷︷ ︸ generalization gap + [ L ( θ̂ ( t ) ; Pn ) − L ( θ∗ ; Pn ) ] ︸ ︷︷ ︸ ≤0 under ERM + [ L ( θ∗ ; Pn ) − L ( θ∗ ; P ) ] ︸ ︷︷ ︸ ≈0 by concentration inequalities . ( 1 ) VER and BER . As discussed in Section 1 , we aim to decompose the excess risk into the variance excess risk ( VER ) and the bias excess risk ( BER ) defined in Definition 1 . The decomposition focuses on the noisy output regimes and we split the output y into the signal component E [ y|x ] and the noise component y − E [ y|x ] . Informally , VER measures how the model performs on pure noise , and BER measures how the model performs on clean data . Definition 1 ( VER and BER ) Given ( x , y ) ∼ P , let ( x , E [ y|x ] ) ∼ Pb denote the signal distribution and ( x , y − E [ y|x ] ) ∼ Pv denote the noise distribution . The variance excess risk ( VER ) EvL ( θ ; P ) and bias excess risk ( BER ) EbL ( θ ; P ) are defined as : EvL ( θ ; P ) ≜ EL ( θ ; Pv ) , EbL ( θ ; P ) ≜ EL ( θ ; Pb ) . To better illustrate VER and BER , we consider three surrogate training dynamics : Standard Training , Variance Training , and Bias Training , corresponding to ER , VER , and BER , respectively . Standard Training : Training process over the noisy data ( X , Y ) from the initialization θ ( 0 ) . We denote the trained parameter at time t by θ̂ ( t ) . Variance Training : Training process over the pure noise ( X , Y − E [ Y |X ] ) from the initialization θ ( 0 ) v . We denote the trained parameter at time t by θ̂ ( t ) v . Bias Training : Training process over the clean data ( X , E [ Y |X ] ) from the initialization θ ( 0 ) b . We denote the trained parameter at time t by θ̂ ( t ) b . When the context is clear , although the trained parameters θ̂ ( t ) , θ̂ ( t ) v , θ̂ ( t ) b are related to the corresponding initialization and the algorithms , we omit the dependency . Besides , we denote θ∗ , θ∗v and θ∗b the optimal parameters which minimize the corresponding population loss L ( θ ; P ) , L ( θ ; Pv ) , and L ( θ ; Pb ) , respectively . Techniques in generalization . We next introduce two techniques in generalization analysis including stability-based techniques in Proposition 1 and uniform convergence in Proposition 2 . These techniques will be revisited in Section 3 . Proposition 1 ( Stability Bound from Feldman & Vondrak ( 2019 ) ) Assume that algorithm At is ϵ-uniformly-stable at time t , meaning that for any two dataset D and D′ with only one different data point , we have sup ( x , y ) EA [ ℓ ( At ( D ) ; x , y ) − ℓ ( At ( D′ ) ; x , y ) ] ≤ ϵ . Then the following inequality holds3 with probability at least 1− δ : |L ( At ( D ) ; P ) − L ( At ( D ) ; Pn ) | = O ( ϵ log ( n ) log ( n/δ ) + √ log ( 1/δ ) √ n ) . Proposition 2 ( Uniform Convergence from Wainwright ( 2019 ) ) Uniform convergence decouples the dependency between the trained parameter and the training set by taking supremum over a parameter space that is independent of training data , namely L ( At ( D ) ; P ) − L ( At ( D ) ; Pn ) ≤ sup θ∈B [ L ( θ ; P ) − L ( θ ; Pn ) ] . where B is independent of dataset D and At ( D ) ∈ B for any time t .

Summary: This paper studies the generalization performance of learning algorithms. The basic idea is to decompose the training dynamics into noise and signal components, and then consider separately the behavior of models trained with noise and signals. The authors then proposed to use the stability approach to tackle the noise components, and the uniform convergence approach to tackle the signal part. The authors consider two specific algorithms: overparameterized linear regression and diagonal matrix recovery. Experimental results are also reported.