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: This paper presents a theoretical study of covariate shift under the setting of kernel ridge regression in the asymptotic regime. They charaterize the limiting test error, bias, and variance in this setting. They relates these with a quantity named hardness. They claim that their findings agree with several interesting empirical phenomena.