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 work proposes a new bound for the excess risk based on a decomposition into a bias and variance term. The authors use a stability-based analysis to bound the bias excess risk, and uniform convergence to bound the variance excess risk. They illustrate their framework in two contexts: overparametrised linear regression with SGD and low-rank diagonal matrix recovery with gradient flow.