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: The paper applies the so-called "replica method" to tackle the generalization issue for kernel regression where there is a distribution shift between training data and testing data. It presents the derivation of kernel regression by using an alternative perspective on RKHS functions. Empirical studies on real datasets including MNIST have been shown. Discussions with the neural tangent kernel are presented.  Despite the interesting idea, the clarity of the paper can be improved to better appreciate the ideas and presentations. Moreover, the significance and impact of the work is also unclear, making it difficult to appreciate the advancement from existing results.