Out-of-Distribution Generalization in Kernel Regression

1 Introduction . Machine learning models are trained to accurately predict on previously unseen samples . A central goal of machine learning theory has been to understand this generalization performance [ 1 , 2 ] . While most of the theory in this domain focused on generalizing in-distribution , i.e . when the training examples are sampled from the same distribution as that of the test examples , in real world applications there is often a mismatch between the training and the test distributions [ 3 ] . Such difference , even when small , may lead to large effects in generalization performance [ 4 , 5 , 6 , 7 , 8 ] . In this paper , we provide a theory of out-of-distribution ( OOD ) generalization for kernel regression [ 9 , 10 , 11 ] applicable to any kernel and real datasets . Besides being a popular machine learning method itself , kernel regression is recovered from an infinite width limit of neural networks [ 12 ] , making our results relevant to understanding OOD generalization in deep neural networks . Indeed , it 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . has been argued that understanding generalization in kernel methods is necessary for understanding generalization in deep learning [ 13 ] . Using methods from statistical physics [ 14 ] and building on techniques developed in a recent line of work [ 15 , 16 ] , we obtain an analytical expression for generalization in kernel regression when test data is drawn from a different distribution than the training distribution . Our results apply to average or typical case learning curves and describe numerical experiments well including those on real data . In contrast , most previous works focus on worst-case OOD generalization bounds in the spirit of statistical learning theory [ 17 , 18 , 19 , 20 , 21 , 22 , 23 ] . We show examples of how mismatched training and test distributions affect generalization and in particular , demonstrate that the mismatch may improve test performance . We use our analytical expression to develop a procedure for minimizing generalization error with respect to the training/test distributions and present applications of it . We study OOD generalization with a linear kernel in detail and present examples of how dimensional reduction of the training distribution can be helpful in generalization , including shifting of a double-descent peak [ 24 ] . We present further analyses of various models in Supplementary Information ( SI ) . 2 OOD Generalization Error for Kernel Regression from the Replica Method . We consider kernel regression in the setting where training examples are sampled i.i.d . from a distribution p ( x ) but the test examples are drawn from a different distribution p̃ ( x ) . We derive our main analytical formula for the generalization error here , and illustrate and discuss its implications in the following sections . 2.1 Problem setup . We consider a training set D = { ( xµ , yµ ) } Pµ=1 of size P where the D-dimensional inputs x 2 RD are drawn independently from a training distribution p ( x ) and the noisy labels are generated from a target function yµ = f̄ ( xµ ) + ✏µ where the noise covariance is E [ ✏µ✏⌫ ] = `` 2 µ⌫ . Kernel regression model is trained through minimizing the training error : f⇤ D = argmin f2H 1 2 PX µ=1 f ( xµ ) yµ 2 + hf , fi H , ( 1 ) where subscript D emphasizes the dataset dependence . Here H is a Reproducing Kernel Hilbert Space ( RKHS ) associated with a positive semi-definite kernel K ( x , x0 ) : RD ⇥RD ! R , and h. , .i H is the Hilbert inner product . The generalization error on the test distribution p̃ ( x ) is given by Eg ( D ) =R dx p̃ ( x ) f⇤ D ( x ) f̄ ( x ) 2 . We note that this quantity is a random variable whose value depends on the sampled training dataset . We calculate its average over the distribution of all datasets with sample size P : Eg = ED Z dx p̃ ( x ) f⇤ D ( x ) f̄ ( x ) 2 . ( 2 ) As P grows , we expect fluctuations around this average to fall and Eg ( D ) to concentrate around Eg . We will demonstrate this concentration in simulations . 2.2 Overview of the calculation . We calculate Eg using the replica method from statistical physics of disordered systems [ 14 , 25 ] . Details of calculations are presented in Section SI.1 . Here we give a short overview . Readers may choose to skip this part and proceed to the main result in Section 2.3 . Our goal is to calculate the dataset averaged estimator f⇤ ( x ) ⌘ EDf⇤D ( x ) and its covariance ED ⇥ f⇤ D ( x ) f⇤ ( x ) f⇤ D ( x0 ) f⇤ ( x0 ) ⇤ . From these quantities , we can reconstruct Eg using the bias-variance decomposition of Eg = B + V , where B = R dx p̃ ( x ) f⇤ ( x ) f̄ ( x ) 2 and V = ED ⇥ R dx p̃ ( x ) f⇤ D ( x ) f⇤ ( x ) 2 ⇤ . For this purpose , it will be convenient to work with a basis in the RKHS defined by Mercer ’ s theorem [ 26 ] . One can find a ( possibly infinite dimensional ) complete orthonormal basis { ⇢ } M⇢=1 for L2 ( RD ) with respect to the training probability distribution p ( x ) due to Mercer ’ s theorem [ 26 ] such that K ( x , x0 ) = ( x ) > ⇤ ( x0 ) , Z dx p ( x ) ( x ) ( x ) > = I , ( 3 ) where we defined M ⇥M diagonal eigenvalue matrix ⇤⇢ = ⌘⇢ ⇢ and the column vector ( x ) = 1 ( x ) , 2 ( x ) , . . . , M ( x ) . Eigenvalues and eigenfunctions satisfy the integral eigenvalue equation Z dx0p ( x0 ) K ( x , x0 ) ⇢ ( x 0 ) = ⌘⇢ ⇢ ( x ) . ( 4 ) We note that the kernel might not express all eigenfunctions if the corresponding eigenvalues vanish . Here we assume that all kernel eigenvalues are non-zero for presentation purposes , however in Section SI.1 we consider the full case . We also define a feature map via the column vector ⇢ ( x ) ⌘ ( ⇤1/2 ( x ) ) ⇢ = p ⌘⇢ ⇢ ( x ) so that h ( x ) , ( x ) iH = I . The complete basis ( x ) can be used to decompose the target function : f̄ ( x ) = ā > ( x ) = w̄ > ( x ) , ( 5 ) where w̄ = ( ⇤ 1/2ā ) , and ā and w̄ are vectors of coefficients . A function belongs to the RKHS H if it has finite Hilbert norm hf , fi2 H = a > ⇤ 1a < 1 . With this setup , denoting the estimator as f ( x ) = w > ( x ) and the target function f̄ ( x ) as given in Eq . ( 5 ) , kernel regression problem reduces to minimization of the energy function H ( w ; D ) ⌘ 12 PP µ=1 ( w̄ w ) > ( xµ ) + ✏µ 2 + 12kwk 2 2. with optimal estimator weights w⇤D = argminw H ( w ; D ) . We again emphasize the dependence of the optimal estimator weights w⇤D to the particular choice of training data D. We map this problem to statistical mechanics by defining a Gibbs distribution / e H ( w ; D ) over estimator weights w which concentrates around the kernel regression solution w⇤ D as ! 1 . This can be used to calculate any function O ( w⇤ ; D ) by the following trick : O ( w⇤ ; D ) = lim ! 1 @ @ J logZ [ J ; , D ] J=0 , Z [ J ; , D ] = Z dwe H ( w ; D ) +JO ( w ) , ( 6 ) where Z is the normalizer of the Gibbs distribution , also known as the partition function . Next , we want to compute the average of EDO ( w⇤ ; D ) which requires computing ED logZ . Further , experience from the study of physics of disordered systems suggests that the logarithm of the partition function concentrates around its mean ( is self-averaging ) for large P [ 14 ] , making our theory applicable to the typical case . However , this average is analytically hard to calculate due to the partition function appearing inside the logarithm . To proceed , we resort to the replica method from statistical physics [ 14 ] , which uses the equality ED logZ = limn ! 0 EDZ n 1 n . The method proceeds by calculating the right hand side for integer n , analytically continuing the resulting expression to real valued n , and performing the limit . While non-rigorous , the replica method has been successfully used in the study of the physics of disordered systems [ 14 ] and machine learning theory [ 27 ] . A crucial step in our computation is approximating w > ( x ) as a Gaussian random variable via its first and second moments when averaged over the training distribution p ( x ) . It has been shown that this approximation yields perfect agreement with experiments [ 28 , 15 , 16 ] . Details of our calculation is given in Supplementary Section SI.1 . 2.3 Main Result . The outcome of our statistical mechanics calculation ( Section SI.1 ) which constitutes our main theoretical result is presented in the following proposition . Proposition 1 . Consider the kernel regression problem outlined in Section 2.1 , where the model is trained on p ( x ) and tested on p̃ ( x ) . Consider the Mercer decomposition of the RKHS kernel K ( x , x0 ) = ( x ) > ⇤ ( x0 ) , where we defined M ⇥ M ( M possibly infinite ) diagonal eigenvalue matrix ⇤⇢ = ⌘⇢ ⇢ and the column vector ( x ) = 1 ( x ) , 2 ( x ) , . . . , M ( x ) , withR dx p ( x ) ( x ) ( x ) > = I . Also consider an expansion of the target function in the Mercer basis f̄ ( x ) = ā > ( x ) . The dataset averaged out-of-distribution generalization error is given by : Eg = E 0 , p ( x ) g + 0 1 `` 2 + 2ā > ( P⇤+ I ) 1O 0 ( P⇤+ I ) 1ā , = + Tr ( P + ⇤ 1 ) 1 , = P Tr ( P + ⇤ 1 ) 2 , 0 = P TrO ( P + ⇤ 1 ) 2 , ( 7 ) where must be solved self-consistently , and we defined the M ⇥M overlap matrix O⇢ = Z dx p̃ ( x ) ⇢ ( x ) ( x ) , O 0 = O 1 0 1 I . ( 8 ) Here E0 , p ( x ) g denotes the generalization error when both training and test distributions are matched to p ( x ) and is given by : E0 , p ( x ) g = 1 `` 2 + 2 1 ā > ( P⇤+ I ) 2ā , ( 9 ) which coincides with the findings of [ 15 , 16 ] . Further , the expected estimator is : f⇤ ( x ; P ) = X ⇢ P⌘⇢ P⌘⇢ + ā⇢ ⇢ ( x ) . ( 10 ) Several remarks are in order . Remark 1 . Our result is general in that it applies to any kernel , data distribution and target function . When applied to kernels arising from the infinite width limit of neural networks [ 12 ] , the information about the architecture of the network is in the kernel spectrum⇤ and the target weights ā obtained by projecting the target function onto kernel eigenbasis . Remark 2 . Formally , the replica computation requires a thermodynamic limit in which P ! 1 , where variations in Eg due to the sampling of the training set become negligible . The precise nature of the limit depends on the kernel and the data distribution , and includes scaling of other variables such as D and M with P . We will give examples of such limits in SI . However , we observe in simulations that our formula predicts average learning curves accurately for even as low as a few samples . Remark 3 . We recover the result obtained in [ 15 , 16 ] when the training and test distributions are the same ( O = I which implies Eg = E 0 , p ( x ) g ) . Remark 4 . Mismatched training and test distributions may improve test error . Central to our analysis is the shifted overlap matrix O 0 which may have negative eigenvalues and hence cause better generalization performance when compared to in-distribution generalization . Remark 5 . The estimator f⇤ only depends on the training distribution and our theory predicts that the estimator eventually approaches to the target f⇤ ! f̄ for large enough P which performs perfectly on the training distribution unless there are out-of-RKHS components in target function . The latter case is studied in depth in Section SI.1 . Remark 6 . While stated for a scalar output , our result can be trivially generalized to a vector output by simply adding the error due to each component of the output vector , as described in [ 15 ] . Next we analyze this result by studying various examples .

Summary: This paper studies the out-of-distribution generalization error in kernel ridge regression (KRR) with any choice of regularization parameter $\lambda$ (including the case of interpolation), and as a general functional of the training and test distributions using the replica method in statistical physics. In addition,  -- they provide explicit closed-form expressions for the case of the linear model.  -- using the neural tangent kernel (NTK), they empirically evaluate their asymptotic formulas on modified samples of the MNIST test dataset to provide best-case training distributions, which can sometimes be different from the one that exactly matches.