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 work studies random features regression on a setting where data $(x_{i}, y_{i})\in\mathbb{R}^{n_{0}+1}$, $i=1, \cdots,n$ is generated from a noisy linear model $y_{i}=\beta^{\top}x_{i}+\epsilon_{i}$ with correlated covariates $x_{i}\sim\mathcal{N}(0, \Sigma)$, planted weights $\beta\sim\mathcal{N}(0,\Sigma_{\beta})$ and independent Gaussian noise $\epsilon_{i}\sim\mathcal{N}(0,\sigma_{\epsilon}^2)$. Estimation is given by $\hat{y}(x) = \hat{\beta}^{\top}\sigma\left(Wx\right)$ with $\sigma$ a real-valued activation function and $W\in\mathbb{R}^{n_{1}\times n_{0}}$ a Gaussian matrix with iid $\mathcal{N}(0, 1/n_{0})$ entries, and $\hat{\beta}\in\mathbb{R}^{n_{1}}$ minimising the ridge regression risk on the features $f = \sigma\left(Wx\right)\in\mathbb{R}^{n_{1}}$.  Its main theoretical contribution is to provide an asymptotic characterisation of the test MSE, the bias and the variance on the proportional regime $m,n_{0}, n_{1}\to\infty$ with fixed ratios $\phi=n_{1}/m, \psi=n_{0}/n_{1}$, as a function of the asymptotic spectral statistics of $(\Sigma, \Sigma_{\beta})$, the ratios $(\phi,\psi)$, the $\ell_2$ penalty strength $\gamma>0$, the noise variance $\sigma_{\epsilon}^2$ and the constants $\rho = \mathbb{E}\left[z\sigma(z)\right]$, $\eta = \rm{Var}\left[\sigma(z)\right]$, $z\sim\mathcal{N}(0,1)$, characterising the activation $\sigma$.   It provides numerical experiments confirming the validity of their formulas for finite but large dimension, and discusses the implications of the anisotropies $(\Sigma, \Sigma_{\beta})$  on the bias, variance and test error in different regimes. In particular, it shows that alignment between $\beta$ and the leading eigenvectors of $\Sigma$ reduces the test MSE, and that anisotropy can induce double-descent like behaviour in the learning curves.