Overparameterization Improves Robustness to Covariate Shift in High Dimensions

1 Introduction . Theoretical justification for almost all machine learning methods relies upon the equality of the distributions from which the training and test data are drawn . Nevertheless , in many real-world applications , this equality is violated—naturally-occurring distribution shift between the training data and the data encountered during deployment is the rule , not the exception [ 31 ] . Even non-adversarial changes in distributions can uncover the surprising fragility of modern machine learning models [ 55 , 56 , 43 , 26 , 12 , 50 ] . Such shifts are distinct from adversarial examples , which require explicit poisoning attacks [ 22 ] ; rather , they can result from mild corruptions , ranging from changes of camera angle or blur [ 26 ] , to subtle , unintended changes in data acquisition procedures [ 56 ] . Moreover , this fragility limits the application of deep learning in certain safety-critical areas [ 31 ] . Empirical studies of distribution shift have observed several intriguing phenomena , including linear trends between model performance on shifted and unshifted test distributions [ 56 , 26 , 31 ] , dramatic degradation in calibration [ 50 ] , and surprising spurious inductive biases [ 12 ] . Theoretical understanding of why such patterns occur across a variety of real-world domains is scant . Even basic questions such as what makes a certain distribution shift likely to hurt ( or help ) a model ’ s performance , ⇤Equal contribution . †Work performed while the author was at Google . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . and by how much , are not understood . One reason that these phenomena have eluded theoretical understanding is that there is often a strong coupling between model and distribution , implying that the effect of a given shift can not usually be understood in a model-agnostic way . Another reason that satisfactory explanations have remained lacking is that the go-to formalism for studying generalization in classical models , namely uniform convergence theory ( see e.g . [ 66 ] ) , may be insufficient to explain the behavior of modern deep learning methods ( even in the absence of distribution shift ) [ 47 , 68 ] . Indeed , classical measures of model complexity , such as various norms of the parameters , have been found to lead to ambiguous conclusions [ 49 ] . In this paper , we follow a different approach : instead of focusing on worst-case bounds for generic distributions , we study average-case behavior for narrowly specified distributions . While this change in perspective sacrifices generality , it allows us to derive more precise predictions , which we believe are necessary to fully capture the relevant phenomenology . We study a specific type of distribution shift called covariate shift , in which the distributions of the training and test covariates differ , while the conditional distribution of the labels given the covariates remains fixed . Using random matrix theory , we perform an asymptotically exact computation of the generalization error of random feature regression under covariate shift . The random feature model provides a useful testbed to ( 1 ) investigate the interplay between various factors such as model complexity , label noise , bias , variance , and covariate shift ; ( 2 ) rigorously define a model-agnostic notion for the strength of covariate shift ; and ( 3 ) provide a theoretical explanation for the linear relationships recently observed between in-distribution and out-of-distribution generalization performance [ 55 , 56 , 27 ] . 1.1 Contributions . Our primary contributions are to : 1 . Provide a model-agnostic partial order over covariate shifts that is sufficient to determine when a shift will increase or decrease the test error in random feature regression ( see Def . 4.1 ) ; 2 . Compute the test error , bias , and variance of random feature regression for general multivariate Gaussian covariates under covariate shift in the high-dimensional limit ( see Sec . 5.1 ) ; 3 . Prove that overparameterization enhances robustness to covariate shift , and that the error , bias , and variance are nonincreasing functions of the number of excess parameters ( see Sec . 5.3 ) ; 4 . Deduce an exact linear relationship between in-distribution and out-of-distribution generalization performance , offering an explanation for this surprising recent empirical observation ( see Sec . 5.4 ) . 1.2 Related work . There is extensive literature on the empirical analysis of distribution shift in all of its myriad forms , ranging from domain adaptation [ 62 , 20 , 8 , 70 , 69 , 37 , 36 ] to defenses against adversarial attacks [ 39 , 60 ] to distributionally robust optimization [ 59 , 15 , 16 ] , among many others . Interestingly , for naturally occurring distribution shifts [ 31 , 26 ] , standard robustness interventions provide little protection [ 56 , 63 ] . Indeed , empirical risk minimization on clean , unshifted training data often performs better on out-of-distribution benchmarks than more sophisticated methods [ 31 ] . One of the most striking observations in the context of natural distribution shifts is that model robustness improves with the classifier ’ s accuracy [ 56 , 63 , 26 , 43 ] . For example , if a classifier ’ s accuracy increases by 1.0 % on the unshifted CIFAR-10 test set , this tends to increase its accuracy by 1.7 % on the CIFAR-10.1 dataset ( a dataset with natural distribution shift ) [ 56 ] . Moreover , such linear trends between the unshifted and shifted measures of error have now been observed in several contexts [ 56 , 63 , 43 , 40 , 44 ] . The number of theoretical works studying the impact of distribution shift on generalization is far smaller . One pioneering work provides VC-dimension-based error bounds for classification that are augmented by a discrepancy measure between source and target domains [ 10 ] , while another demonstrates a similar class of uniform convergence-based results in the setting of kernel regression [ 11 ] . Recent work shows that learning domain-invariant features is insufficient to guarantee generalization when the class-conditional distributions of features may shift [ 69 ] . When the source domain gradually shifts toward the target domain , non-vacuous margin-based bounds for self-training can be established [ 32 ] . In [ 40 ] , assumptions based on model similarity are used to help explain why classifiers exhibit linear trends between their accuracies on shifted and unshifted test sets [ 56 , 43 ] . Our technical tools build on a series of works that have studied the exact high-dimensional limit of the test error for a growing class of model families and data distributions . In the context of linear models , recent work analyzes ridge regression for general covariances and a general non-isotropic source condition on the parameters which generate the targets [ 57 ] , extending earlier work studying minimum-norm interpolated least squares and ridge regression in the random design setting [ 9 , 14 , 24 ] . The non-isotropy of source parameter effectively induces a shift on the bias term of this model , but the phenomenon is distinct from the covariate shifts we study here . Beyond linear regression , random feature models provide a rich but tractable class of models to gain further insight into generalization phenomena [ 2 , 3 , 41 , 35 ] . These methods are of particular interest because of their connection to neural networks , with the number of random features corresponding to the network width ( or model complexity ) [ 48 , 33 , 29 ] , and because they serve as a practical method for data analysis in their own right [ 54 , 61 ] . In this context , a precise characterization of the gaps between uniform convergence and the ( asymptotic ) exact test error as a function of the sample size and number of random features can be derived [ 68 ] . Since this paper ’ s publication , we released follow-up work considering unequal scales in the training and test distributions and optimal regularization [ 64 ] . From the technical perspective , our analytic techniques build upon these works and a series of recent results stemming from the literature on random matrix theory and free probability [ 53 , 52 , 1 , 2 , 38 , 51 , 19 , 45 ] . 2 Preliminaries . 2.1 Problem setup and notation . As in prior work studying random feature regression [ 24 , 41 , 2 , 1 ] , we compute the test error in the high-dimensional , proportional asymptotics where the dataset size m , input feature dimension n0 , and hidden layer size n1 all tend to infinity at the same rate , with : = n0/m and : = n0/n1 held fixed . We refer to / as the overparameterization ratio , which is the limit of n1/m and characterizes the normalized complexity of the ( random ) feature model . Interestingly , in this high-dimensional limit , the conditional distribution of a linear labeling function is asymptotically equivalent to a wide class of nonlinear teacher functions ( see [ 41 , 2 ] for more details ) . With this in mind , we consider the task of learning an unknown function from m i.i.d . samples ( xi , yi ) 2 Rn0 ⇥ R for i 2 { 1 , . . . , m } , where the covariates are Gaussian , xi ⇠ N ( 0 , ⌃ ) with positive definite covariance matrix ⌃ , and the labels are generated by a linear function parameterized by 2 Rn0 , drawn from N ( 0 , In0 ) . In particular y ( xi ) = > xi/ p n0 + ✏i , ( 1 ) where ✏i ⇠ N ( 0 , 2✏ ) is additive label noise on the training points . We study the class of prediction models defined by kernel ridge regression using unstructured random feature maps [ 54 ] . The random features are given by a single-layer , fully-connected neural network with random weights . Given a set of training data X = [ x1 , . . . , xm ] and a prospective test point x , the random features embeddings of the training and test data are given by F : = ( WX/ p n0 ) and f : = ( Wx/ p n0 ) , ( 2 ) for a random weight matrix W 2 Rn1⇥n0 with i.i.d . standard Gaussian entries and an activation function : R ! R applied elementwise . The induced kernel is K ( x1 , x2 ) : = 1 n1 ( Wx1/ p n0 ) > ( Wx2/ p n0 ) , ( 3 ) and the model ’ s predictions are given by ŷ ( x ) = Y K 1Kx , where Y : = [ y ( x1 ) , . . . , y ( xm ) ] , K : = K ( X , X ) + Im , Kx : = K ( X , x ) , and 0 is a ridge regularization constant1 . Owing to the implicit regularization effect of the nonlinear feature maps [ 7 ] , in low noise settings the optimal value of can sometimes be negative [ 30 ] . For simplicity , we nevertheless make the standard assumption that 0 , though we emphasize our techniques readily accommodate negative values . 1We overload the definition of K to include the additive regularization whenever no arguments are present . Also , if = 0 and K is not full-rank , K 1 should be understood as the Moore–Penrose pseudoinverse . Our central object of study is the expected test loss for a datapoint x ⇠ N ( 0 , ⌃⇤ ) where ⌃⇤ may be different from the training covariance ⌃ . The test error ( without label noise on the test point ) is E⌃⇤ = E [ ( > x/ p n0 Y K 1Kx ) 2 ] ( 4 ) = Ex , [ ( E [ ŷ ( x ) ] y ( x ) ) 2 ] | { z } B⌃⇤ +Ex , [ V [ ŷ ( x ) ] ] | { z } V⌃⇤ , ( 5 ) where the inner expectations defining the bias and variance are computed over W , X , and Y . We decompose the training and test covariance matrices into eigenbases as ⌃ = P n0 i=1 ivivi > and ⌃⇤ =P n0 i=1 ⇤ i v⇤ i v⇤ i > , where the eigenvalues are in nondecreasing magnitude , i.e . 1 2 . . . n0 and ⇤1 ⇤2 . . . ⇤n0 . We define the overlap coefficients ri : = v > i ⌃⇤vi = n0X j=1 ( v⇤ j · vi ) 2 ⇤j ( 6 ) to measure the alignment of ⌃⇤ with the ith eigendirection of ⌃ . In particular , ri is the induced norm of vi with respect to ⌃⇤ . We use t̄r to denote the dimension-normalized trace : for a matrix A 2 Rn⇥n , t̄r ( A ) = 1 n tr ( A ) . We use kAk1 and kAkF to denote the operator norm and Frobenius norm of matrix A respectively . Finally , we use x to denote the Dirac delta function centered at x .

Summary: The paper studies the test error, as well as the bias and variance, of random feature regression in the asymptotic regime, in a setting with covariate shift — i.e., when the covariance matrix of the features is different at testing time than at training time. The paper defines a partial ordering over covariance shifts, and it proves that this corresponds to increased or decreased performance compared to the case of no shift. It also obtains a few results that relate these notions to overparameterization.