Anisotropic Random Feature Regression in High Dimensions

1 INTRODUCTION . While deep learning races ahead with its impressive list of practical successes , it has often done so by building larger models and training on ever larger datasets . For example , recent large language models leverage billions ( Brown et al. , 2020 ; Raffel et al. , 2019 ) or even trillions ( Fedus et al. , 2021 ) of parameters . As the state-of-the-art pushes the computational frontier , model development , experimentation , and exploration become increasingly costly , both in terms of the computational resources and person-hours . To facilitate the development of improved models , it is becoming ever more necessary to develop theoretical models that can help guide model development in these high-dimensional scenarios . Indeed , our theoretical understanding is so limited that even linear regression , the most basic machine learning method of all , is still yielding surprises in high dimensions ( Hastie et al. , 2019 ) . While linear regression can capture many salient features of high-dimensional data , it is not a reasonable model for high-dimensional models . In particular , there is no natural data-independent notion of model size : the number of parameters in the model is intimately tied ( usually equal ) to the dimensionality of the datapoints . This limitation means that linear regression may not accurately capture the effect of overparameterization , which appears to be a crucial ingredient in modern deep learning models ( Zhang et al. , 2016 ) . Beyond linear regression , the next simplest model is arguably the random feature model of Rahimi et al . ( 2007 ) , which turns out to be tractable in high dimensions while also allowing for a natural data-independent lever to vary model size – the number of random features . Random feature models have also gained independent interest through their relationship to neural kernels and Gaussian processes ( Neal , 1996 ; Lee et al. , 2017 ; Novak et al. , 2018 ; Lee et al. , 2019 ; Jacot et al. , 2018 ) . While recent works have examined the asymptotic performance of random feature models , most have done so in the restricted setting of isotropic covariates and target functions ( Mei & Montanari , 2019 ; d ’ Ascoli et al. , 2020 ; Adlam & Pennington , 2020a ; b ) . Based on the conjectured Gaussian equivalence principle of Goldt et al . ( 2020b ) , a few recent works have begun relaxing these simplifying assumptions in an effort to describe more natural data distributions ( Loureiro et al. , 2021 ; d ’ Ascoli et al. , 2021 ) . In this work , we study high-dimensional random feature regression for anisotropic data and target functions , setting the results of d ’ Ascoli et al . ( 2021 ) on a more rigorous footing and revealing a host of novel phenomena related to the alignment of the data and the target weights . 1.1 CONTRIBUTIONS . Our primary contributions are to : 1 . Derive theoretical predictions for the test error , bias , and variance of random feature regression for anisotropic Gaussian covariates under general linear target functions in the high-dimensional asymptotic limit ( Section 3.1 ) ; 2 . Prove that overparameterization reduces the total test error , bias , and variance , even in the presence of anisotropy ( Section 3.2 ) ; 3 . Define a partial order on the space of weight-data alignments and prove that generalization performance improves in response to stronger alignment ( Definition 2.1 , Section 3.3 ) ; 4 . Demonstrate that anisotropy can induce sample-wise multiple descent , but prove that parameter-wise multiple descent is limited to double descent ( Section 3.4 ) ; 5 . Provide a theoretical explanation for steep cliffs observed in parameter-wise learning curves that persist under optimal regularization , and relate them to phase transitions in the spectrum of the feature kernel matrix ( Section 3.4 ) . 1.2 RELATED WORK . Our analysis builds on a series of works that have studied the exact high-dimensional asymptotics of the test error for a growing class of model families and data distributions . Early work in this direction examined minimum-norm interpolated least squares and ridge regression in the high-dimensional random design setting ( Belkin et al. , 2019 ; Dobriban et al. , 2018 ; Hastie et al. , 2019 ) , finding surprising phenomena such as double descent and a negative optimal ridge constant ( Kobak et al. , 2020 ) . These analyses were generalized in ( Richards et al. , 2021 ; Wu & Xu , 2020 ; Hastie et al. , 2019 ) to allow for anisotropic covariance matrices and weights that generate the targets . This anisotropic configuration was studied in further detail in ( Mel & Ganguli , 2021 ) , uncovering nuanced behavior such as sample-wise multiple descent and a sequence of phase transitions in the performance and optimal regularization . While linear regression can provide valuable insight into generalization performance in high dimensions , it is not sufficiently complex to faithfully reproduce the features of non-linear models , such as the effects of overparameterization . Random feature models , on the other hand , can capture many such features while still maintaining analytical tractability . Recent results have developed a detailed understanding of isotropic random feature regression in the high-dimensional asymptotic limit , providing a precise characterization of peaks in the test error and the benefits of overparameterization ( Mei & Montanari , 2019 ; d ’ Ascoli et al. , 2020 ; Adlam & Pennington , 2020a ) . The origin of these peaks was explained by means of a fine-grained decomposition of the bias and variance in ( Adlam & Pennington , 2020b ; Lin & Dobriban , 2020 ) . While most prior work examining the high-dimensional asymptotics of random feature models have leveraged the isotropy of the covariate distribution , this assumption was recently relaxed in ( Liao et al. , 2020 ) for random Fourier feature models , in ( Anonymous , 2021 ) in the study of covariate shift , and in a series of works based on the Gaussian equivalence principle of Goldt et al . ( 2020a ; b ) ( Gerace et al. , 2020 ; Loureiro et al. , 2021 ) . Though most of these works did not pursue a detailed investigation of the anisotropy itself , while finalizing this manuscript we became aware of concurrent work ( d ’ Ascoli et al. , 2021 ) which does focus on the interplay of the anisotropic input features and target functions , as we do here . Whereas d ’ Ascoli et al . ( 2021 ) leverage the Gaussian equivalence principle and the replica method from statistical physics , we use a completely different set of techniques stemming from the literature on random matrix theory and operator-valued free probability ( Pennington & Worah , 2018 ; 2019 ; Adlam et al. , 2019 ; Adlam & Pennington , 2020a ; Louart et al. , 2018 ; Péché et al. , 2019 ; Far et al. , 2006 ; Mingo & Speicher , 2017 ) , which can be made completely rigorous . d ’ Ascoli et al . ( 2021 ) investigates the role of anisotropy not just under squared loss , as we do here , but also under logistic loss , and obtains good numerical support for the Gaussian equivalence conjecture in that setting . In contrast , our narrower and deeper focus on squared loss allows us not only to uncover a host of novel phenomena , including multiple descent , steep error cliffs , and spectral phase changes , but also to rigorously prove several propositions about their existence and properties . We provide a more in-depth discussion of these and other related works in App . A . Our work relies heavily on the concept of alignment ( see Def . 2.1 ) between the data covariance and the target vector coefficients . A completely analogous definition appears in Anonymous ( 2021 ) in the context of covariate shift . Owing to the similarity of these definitions as well as the underlying formulas , many of the techniques and analyses of Anonymous ( 2021 ) can be applied in our context , leading to a kind of mirroring of results . As the underlying implications and interpretations are otherwise quite different , we refrain from repeatedly commenting on these recurring connections . 2 PRELIMINARIES . 2.1 PROBLEM SETUP AND NOTATION . We study random feature regression ( Rahimi et al. , 2007 ) as a model for learning an unknown linear function of the data , y ( xi ) = > xi/ p n0 + ✏i , ( 1 ) from m independent samples ( xi , yi ) 2 Rn0 ⇥ R , i = 1 , . . . , m , where the covariates are Gaussian , xi ⇠ N ( 0 , ⌃ ) , and where ✏i ⇠ N ( 0 , 2✏ ) is additive noise ( that we take to be present on the training samples only ) . We focus on the high-dimensional proportional asymptotics in which the input feature dimension n0 , the hidden layer size n1 , and the number of samples m all tend to infinity at the same rate , with : = n0/m and : = n0/n1 held constant . The overparameterization ratio / = n1m is a measure of the normalized complexity of the random feature model . Interestingly , under this model in our high-dimensional setup , any nonlinear component of the signal to be learned behaves like additive noise , so the linear function in Eq . ( 1 ) does not sacrifice generality ( Mei & Montanari , 2019 ; Adlam & Pennington , 2020a ) . Following many works on high-dimensional regression ( see e.g . Dobriban et al . ( 2018 ) ) , we assume the coefficient vector is random , ⇠ N ( 0 , ⌃ ) , though allowing for deterministic would be a straightforward extension and would simply involve replacing ⌃ ! > throughout . Denoting the training dataset as X = [ x1 , . . . , xm ] and the test point as x , the random features are , F : = ( WX/ p n0 ) and f : = ( Wx/ p n0 ) , ( 2 ) where W 2 Rn1⇥n0 is a random weight matrix with i.i.d . standard Gaussian entries , and ( · ) : R ! R is an activation function that is applied elementwise . These random features define the kernel K ( x1 , x2 ) : = 1 n1 ( Wx1/ p n0 ) > ( Wx2/ p n0 ) , ( 3 ) which can be used to compute the model ’ s predictions on a point x as ŷ ( x ) = Y K 1Kx . ( 4 ) Here we have defined the training labels as Y : = [ y ( x1 ) , . . . , y ( xm ) ] , and the regularized kernel matrix as K : = K ( X , X ) + Im , as well as Kx : = K ( X , x ) for 0 . The training error is given by , E⌃train = E E [ ( y ( X ) ŷ ( X ) ) 2 ] = E [ ( > X/ p n0 Y K 1K ( X , X ) ) 2 ] , ( 5 ) and the test error ( without label noise on the test point ) can be written as E⌃ = E ExE [ ( y ( x ) ŷ ( x ) ) 2 ] 2✏ = ExE [ ( > x/ p n0 Y K 1Kx ) 2 ] , ( 6 ) where the outermost expectation over has been suppressed since the quantity concentrates sharply around its mean and the inner expectation is computed over all the randomness from training , i.e . W , X , and ✏ . The test error can be decomposed into its bias and variance components as E⌃ = ExE [ ( E [ ŷ ( x ) ] y ( x ) ) 2 ] | { z } B⌃ +Ex [ V [ ŷ ( x ) ] ] | { z } V⌃ . ( 7 ) The bias and variance are computed with respect to all the randomness from training , i.e . W , X , ✏ , which differs from common practice in the statistics literature , but which is necessary to obtain a decomposition that is intuitive and unambiguous when the predictive function relies on randomness from multiple sources ( Adlam & Pennington , 2020b ; Lin & Dobriban , 2020 ) . As noted by ( Hastie et al. , 2019 ; Wu & Xu , 2020 ; Mel & Ganguli , 2021 ) , the test error of linear regression depends on the geometry of ( ⌃ , ⌃ ) ( or ( ⌃ , ) in the case of nonrandom ) , and we will find the same to be true for random feature regression . As such , we decompose the training and covariance matrices into eigenbases as ⌃ = Pn0 i=1 ivivi > and ⌃ = Pn0 i=1 i v i v i > , where the eigenvalues are in nondecreasing magnitude , i.e . 1 2 . . . n0 and 1 2 . . . n0 . Following ( Anonymous , 2021 ) , we then define the overlap coefficients as qi : = v > i ⌃ vi = n0X j=1 ( v j · vi ) 2 j , ( 8 ) to measure the alignment of ⌃ with the ith eigendirection of ⌃ . Note that for deterministic this definition simply gives the square of the component of in the ith eigendirection , i.e . qi = ( v > i ) 2 .

Summary: This paper studies nonlinear random feature regression of a linear response in the *anisotropic* setting where both the covariates and parameter have arbitrary covariance matrix. The training error and test error are computed explicitly as the dimension, sample size, and number of features jointly tend to infinity. The proof method relies on linearization and a technique from free probability involving the construction of linear pencils. The asymptotic formulas are used to theoretically describe (i) the beneficial effect of overparameterization on the bias and variance components of the test error and (ii) the number of critical points of the test error as a function of overparameterization. This paper graphs the asymptotic formulas for some specific models to demonstrate parameter-wise and sample-wise multiple descent phenomena.