Explaining Scaling Laws of Neural Network Generalization

1 SCALING LAWS FOR NEURAL NETWORKS . For a large variety of models and datasets , neural network performance has been empirically observed to scale as a power-law with model size and dataset size ( Hestness et al. , 2017 ; Kaplan et al. , 2020 ; Rosenfeld et al. , 2020b ; Henighan et al. , 2020 ) . These exponents determine how quickly performance improves with more data and larger models . We would like to understand why these power-laws emerge , and what features of the data and models determine the values of the power law exponents . In this work , we present a theoretical framework for understanding scaling laws in trained neural networks . We identify four related scaling regimes with respect to the number of model parameters P and the dataset size D. With respect to each of D , P , there is both a variance-limited regime and a resolution-limited regime . Variance-Limited Regime In the limit of infinite data or an arbitrarily wide model , some aspects of neural network training simplify . Specifically , if we fix one of D , P and study scaling with respect to the other parameter as it becomes arbitrarily large , then the difference between the finite test loss and its limiting value scales as 1/x , i.e . as a power-law with exponent 1 , with x = D or √ P ∝ width in deep networks and x = D or P in linear models . Resolution-Limited Regime In this regime , one of D or P is effectively infinite , and we study scaling as the other parameter increases . In this case , a variety of works have empirically observed power-law scalings 1/xα , typically with 0 < α < 1 for both x = P or D. We derive exponents in this regime precisely in the setting of random feature models ( c.f . next section ) . Empirically , we find that our theoretical predictions for exponents hold in pretrained , fine-tuned models even though these lie outside our theoretical setting . For more general nonlinear models , we propose a refinement of naive bounds into estimates via expansions that hold asymptotically . These rely on the idea that additional data ( in the infinite model-size limit ) or added model parameters ( in the infinite data limit ) are used by the model to carve up the data manifold into smaller components . For smooth manifolds , loss , and network , the test loss will depend on the linear size of a sub-region , while it is the d-dimensional sub-region volume that scales inversely with P or D , giving rise to α ∝ 1/d.1 To test this empirically , we make measurements of the resolution-limited exponents in neural networks and intrinsic dimension of the data manifold , shown in Figure 1b . Explicit Derivation We derive the scaling laws for these four regimes explicitly in the setting of random feature teacher-student models , which also applies to neural networks in the large width limit . This setting allows us to solve for the test error directly in terms of the feature covariance ( kernel ) . The scaling of the test loss then follows from the asymptotic decay of the spectrum of the covariance matrix . For generic continuous kernels on a d-dimensional manifold , we can further relate this to the dimension of the data manifold . Summary of Contributions : . 1 . We propose four scaling regimes for neural networks . The variance-limited and resolutionlimited regimes originate from different mechanisms , which we identify . To our knowledge , 1A visualization of this successively better approximation with dataset size is shown in Figure 1b for models trained to predict data generated by a random fully-connected network . this categorization has not been previously exhibited . We provide empirical support for all four regimes in deep networks on standard datasets . 2 . We derive the variance-limited regime under simple yet general assumptions ( Theorem 1 ) . 3 . We present a hypothesis for resolution-limited scaling through refinement of naive bounds ( Theorems 2 , 3 ) , for general nonlinear models . We empirically test the dependence of the estimates on intrinsic dimension of the data manifold for deep networks on standard datasets ( Figure 1b ) . 4 . In the setting of random feature teacher-student networks , we derive both variance-limited and resolution-limited scaling exponents exactly . In the latter case , we relate this to the spectral decay of kernels . We identify a novel duality that exists between model and dataset size scaling . 5 . We empirically investigate predictions from the random features setting in pretrained , fine-tuned models on standard datasets and find they give excellent agreement . 6 . We study the dependence of the scaling exponent on changes in architecture and data , finding that ( i ) changing the input distribution via switching datasets and ( ii ) the addition of noise have strong effects on the exponent , while ( iii ) changing the target task via superclassing does not . Related Works : There have been a number of recent works demonstrating empirical scaling laws ( Hestness et al. , 2017 ; Kaplan et al. , 2020 ; Rosenfeld et al. , 2020b ; Henighan et al. , 2020 ; Rosenfeld et al. , 2020a ) in deep neural networks , including scaling laws with model size , dataset size , compute , and other observables such as mutual information and pruning . Some precursors ( Ahmad & Tesauro , 1989 ; Cohn & Tesauro , 1991 ) can be found in earlier literature . Recently , scaling laws have also played a significant role in motivating work on the largest models that have yet been developed ( Brown et al. , 2020 ; Fedus et al. , 2021 ) . There has been comparatively little work on theoretical ideas ( Sharma & Kaplan , 2020 ; Bisla et al. , 2021 ) that match and explain empirical findings in generic deep neural networks . In the particular case of large width , deep neural networks behave as random feature models ( Neal , 1994 ; Lee et al. , 2018 ; Matthews et al. , 2018 ; Jacot et al. , 2018 ; Lee et al. , 2019 ; Dyer & Gur-Ari , 2020 ) , and known results on the loss scaling of kernel methods can be applied ( Spigler et al. , 2020 ; Bordelon et al. , 2020 ) . Though not in the original , Bordelon et al . ( 2020 ) analyze resolution-limited dataset size scaling for power-law spectra in later versions . During the completion of this work , Hutter ( 2021 ) presented a specific solvable model of learning exhibiting non-trivial power-law scaling for power-law ( Zipf ) distributed features . This does not directly relate to the setups studied in this work , or present bounds that supersede our results . Concurrent to our work , Bisla et al . ( 2021 ) presented a derivation of the resolution-limited scaling with dataset size , also stemming from nearest neighbor distance scaling on data manifolds . However , they do not discuss requirements on model versus dataset size or how this scaling behavior fits into other asymptotic scaling regimes . In the variance-limited regime , scaling laws in the context of random feature models ( Rahimi & Recht , 2008 ; Hastie et al. , 2019 ; d ’ Ascoli et al. , 2020 ) , deep linear models ( Advani & Saxe , 2017 ; Advani et al. , 2020 ) , one-hidden-layer networks ( Mei & Montanari , 2019 ; Adlam & Pennington , 2020a ; b ) , and wide neural networks treated as Gaussian processes or trained in the NTK regime ( Lee et al. , 2019 ; Dyer & Gur-Ari , 2020 ; Andreassen & Dyer , 2020 ; Geiger et al. , 2020 ) have been studied . In particular , this behavior was used in ( Kaplan et al. , 2020 ) to motivate a particular ansatz for simultaneous scaling with data and model size . The resolution-limited analysis can perhaps be viewed as an attempt to quantify the ideal-world generalization error of Nakkiran et al . ( 2021 ) . This work makes use of classic results connecting the spectrum of a smooth kernel to the geometry it is defined over ( Weyl , 1912 ; Reade , 1983 ; Kühn , 1987 ; Ferreira & Menegatto , 2009 ) and on the scaling of iteratively refined approximations to smooth manifolds ( Stein , 1999 ; Bickel et al. , 2007 ; de Laat , 2011 ) . 2 FOUR SCALING REGIMES . Throughout this work we will be interested in how the average test loss L ( D , P ) depends on the dataset size D and the number of model parameters P . Unless otherwise noted , L denotes the test loss averaged over initialization of the parameters and draws of a size D training set . Some of our results only pertain directly to the scaling with width w ∝ √ P , but we expect many of the intuitions apply more generally . We use the notation αD , αP , and αW to indicate scaling exponents with respect to dataset size , parameter count , and width . All proofs appear in the supplement . 2.1 VARIANCE-LIMITED EXPONENTS . In the limit of large D the outputs of an appropriately trained network approach a limiting form with corrections which scale as D−1 . Similarly , recent work shows that wide networks have a smooth large P limit ( Jacot et al. , 2018 ) , where fluctuations scale as 1/ √ P . If the loss is sufficiently smooth then its value will approach the asymptotic loss with corrections proportional to the variance , ( 1/D or 1/ √ P ) . In Theorem 1 we present sufficient conditions on the loss to ensure this variance dominated scaling . We note , these conditions are satisfied by mean squared error and cross entropy loss , though we conjecture the result holds even more generally . Theorem 1 . Let ` ( f ) be the test loss as a function of network output , ( L = E [ ` ( f ) ] ) , and let fT be the network output after T training steps , thought of as a random variable over weight initialization , draws of the training dataset , and optimization seed . Further let fT be concentrating with E [ ( fT − E [ fT ] ) k ] = O ( ) ∀k ≥ 2 . If ` is a finite degree polynomial , or has bounded second derivative , or is 2-Hölder , then E [ ` ( fT ) ] − ` ( E [ fT ] ) = O ( ) . Dataset scaling Consider a neural network , and its associated training loss Ltrain ( θ ) . For every value of the weights , the training loss , thought of as a random variable over draws of a training set of size D , concentrates around the population loss , with a variance which scales as O ( D−1 ) . If the optimization procedure is sufficiently smooth , the trained weights , network output , and higher moments , will approach their infinite D values , ED [ ( fT − ED [ fT ] ) k ] = O ( D−1 ) . Here , the subscript D on the expectation indicates an average over draws of the training set . This scaling together with Theorem 1 gives the variance limited scaling of loss with dataset size . This concentration result with respect to dataset size has appeared for linear models in Rahimi & Recht ( 2008 ) and for single hidden layer networks with high-dimensional input data in Mei & Montanari ( 2019 ) ; Adlam & Pennington ( 2020a ; b ) . In the supplement we prove this for GD and SGD with polynomial loss as well as present informal arguments more generally . Additionally , we present examples violating the smoothness assumption and exhibiting different scaling . Large Width Scaling We can make a very similar argument in the w → ∞ limit . It has been shown that the predictions from an infinitely wide network , either under Bayesian inference ( Neal , 1994 ; Lee et al. , 2018 ) , or when trained via gradient descent ( Jacot et al. , 2018 ; Lee et al. , 2019 ) approach a limiting distribution at large width equivalent to a linear model . Furthermore , corrections to the infinite width behavior are controlled by the variance of the full model around the linear model predictions . This variance ( and higher moments ) have been shown to scale as 1/w ( Dyer & Gur-Ari , 2020 ; Yaida , 2020 ; Andreassen & Dyer , 2020 ) , Ew [ ( fT − Ew [ fT ] ) k ] = O ( w−1 ) . Theorem 1 then implies the loss will differ from its w =∞ limit by a term proportional to 1/w . We note that there has also been work studying the combined large depth and large width limit , where Hanin & Nica ( 2020 ) found a well-defined infinite size limit with controlled fluctuations . In any such context where the model predictions concentrate , we expect the loss to scale with the variance of the model output . In the case of linear models , studied below , the variance is O ( P−1 ) rather than O ( √ P ) , and we see the associated variance scaling in this case .

Summary: This paper characterized the scaling of generalization error with respect to the number of training samples $D$ and parameters $P$ for certain smooth estimators in two regimes: the variance-limited regime and the resolution-limited regime; this divide is somewhat analogous to the parametric vs. nonparametric estimation. The authors also argued that the large-width and large-sample limit have similar scaling behavior. The findings are supported by empirical evidence.