Out-of-Distribution Generalization in Kernel Regression
Year: 2022

Abstract:
In real word applications, the data generating process for training a machine learning model often differs from what the model encounters in the test stage. Understanding how and whether machine learning models generalize under such distributional shifts remains a theoretical challenge. Here, we study generalization in kernel regression when the training and test distributions are different using the replica method from statistical physics. We derive an analytical formula for the out-of-distribution generalization error applicable to any kernel and real datasets. We identify an overlap matrix that quantifies the mismatch between distributions for a given kernel as a key determinant of generalization performance under distribution shift. Using our analytical expressions we elucidate various generalization phenomena including possible improvement in generalization when there is a mismatch. We develop procedures for optimizing training and test distributions for a given data budget to find best and worst case generalizations under the shift. We present applications of our theory to real and synthetic datasets and for many kernels. We compare results of our theory applied to Neural Tangent Kernel with simulations of wide networks and show agreement. We analyze linear regression in further depth.

heading: 1 Introduction
text: 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).

heading: 2 OOD Generalization Error for Kernel Regression from the Replica Method
text: 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.

heading: 2.1 Problem setup
text: 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.

heading: 2.2 Overview of the calculation
text: 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.

heading: 2.3 Main Result
text: 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.

heading: 3 Applications to Real Datasets
text: First, we test our theory on real datasets and demonstrate its broad applicability. To do that, we define Dirac training and test probability measures for a fixed dataset D of size M with some point mass on each of the data points p(x) = P xµ2D p µ (x xµ) and p̃(x) = P xµ2D p̃ µ (x xµ). With this measure, the kernel eigenvalue problem in Eq.(4) becomes Kdiag(p) = ⇤, where diag(p) denotes the diagonal matrix of probability masses of each data point. In this setup the number of
eigenfunctions, or eigenvectors, are equal to the size of entire dataset D. Once the eigenvalues ⇤ and the eigenvectors have been identified, we compute the target function coefficients by projecting the target data yc onto these eigenvectors ac = >diag(p)yc for each target c = 1, . . . , C. Once all of these ingredients are obtained, theoretical learning curves can be computed using Proposition 1. This procedure is similar to the one used in [15, 16].

heading: 3.1 Shift in test distribution may help or hurt generalization
text: To analyze the influence of distribution shift on generalization error, we first study the problem of a fixed training distribution and a shifting test distribution. While one may naively expect that a shift in the test distribution would always result in worse performance, we demonstrate that mismatch can be either beneficial or detrimental to generalization. A similar point was made in [4].
Since our generalization error formula is a differentiable function with respect to the training measure, we numerically optimize Eg with respect to {p̃µ} using gradient descent to find beneficial test measures where the error is smaller than for the case where training and test measures are matched. Similarly, we perform gradient ascent to find detrimental test measures, which give higher test error than the matched case.
The result of this procedure is shown in Figure 1. We perform gradient descent (labeled beneficial) and gradient ascent (labeled detrimental) on M = 1000 MNIST digits 8’s and 9’s, where target outputs are binary { 1, 1} corresponding to two digits. We use a depth 3 ReLU fully-connected neural network with 2000 hidden units at each layer and its associated Neural Tangent Kernel (NTK) which are computed using the NeuralTangents API [29]. The probability mass assigned to the first 50 points are provided in Figure 1(a). We see that points given high mass for the beneficial test distribution are given low probability mass for the detrimental test distribution and vice versa. We plot the measure change matrix O 0 which quantifies the change in the generalization error due to distribution shift. Recall that in the absence of noise, Eg is of the form Eg = Ematchedg + v>O 0v for
Algorithm 1: Optimizing Training Measure at sample size P Result: GET_LOSS(z 2 RM , K 2 RM⇥M , y 2 RM⇥C , 2 R+)
Compute normalized distribution p = softmax(z); Diagonalize on Train Measure Kdiag(p) = ⇤ with >diag(p) = I ; Get Target Function Weights a = >diag(p)y; Get Overlap Matrix O = 1M > ;
Solve Implicit Equation  = ODE-INT h ̇ = +  P k ⇤kk ⇤kkP+  i
; return Eg(, P,⇤,O) (see Proposition 1);
; Result: OPT_MEASURE(K 2 RM⇥M , y 2 RM , P , T ,⌘, )
Initialize z = 0 2 RM , t = 0; Diagonalize Kernel on Uniform Measure 1MK = ̃⇤̃̃
>; while t < T do
z = z ⌘ GRAD[GET_LOSS(z,K,y, )]; t = t+ 1;
end return softmax(z) ;
a vector v which depends on the task, kernel, training measure and sample size P . Depending on the eigenvalues and eigenvectors of O 0, and the vector v, the quadratic form v>O 0v may be positive or negative, resulting in improvement or detriment in the generalization error. In Figure 1(b), we see that the O 0 matrix for the beneficial test measure appears roughly opposite that of the detrimental measure. This is intuitive since it would imply the change in the generalization error to be opposite in sign for beneficial and detrimental measures v> O 0 v = v>O 0v. The spectrum of the matrix, shown in Figure 1(c), reveals that it possesses both positive and negative eigenvalues. We plot the theoretical and experimental learning curves in Figure 1(d). As promised, the beneficial (blue) measure has lower generalization error while the detrimental measure (orange) has higher generalization error than the matched case (black). Experiments show excellent match to our theory. In Section SI.2, we present another application of this procedure to adversarial attacks during testing.

heading: 3.2 Optimized Training Measure for Digit Classification with Wide Neural Networks
text: Often in real life tasks test distribution is fixed while one can alter the training distribution for more efficient learning. Therefore, we next study how different training distributions affect the generalization performance of kernel regression when the test distribution is fixed. We provide pseudocode in Algorithm 1 which describes a procedure to optimize the expected generalization error with respect to a training distribution of P data points. This optimal training distribution has been named the dual distribution to p̃(x) [4]. All of the operations in the computation of Eg support automatic differentiation with respect to logits z 2 RM that define the training distribution through softmax(z), allowing us to perform gradient descent on the training measure [30].
As an example, we fix the test distribution to be uniform MNIST digits for 8’s and 9’s and optimize over the probability mass of each unique digit to minimize Eg, again using a depth 3 ReLU neural tangent kernel and binary targets { 1, 1}. Our results indicate that it is possible to reduce generalization when the training distribution is chosen differently than the uniform test distribution (See Section SI.2 for extended discussion). In Figure 2, we optimize Eg for P = 30 training samples to extract the optimal training distribution on NTK. We observe that the high mass training digits are closer to the center of their class, Figure 2(e). However we also find that optimizing the distribution for different values of P may cause worse generalization beyond the P used to optimize the distribution Figure 2(g, h).
We also test our results on neural networks, exploiting a correspondence between kernel regression with the Neural Tangent Kernel, and training infinitely wide neural networks [12]. Figure 2(i) shows that our theory matches experiments with ReLU networks of modest width 2000.

heading: 4 Linear Regression: An Analytically Solvable Model
text: Next we study linear regression to demonstrate various interesting OOD generalization phenomena. Consider a linear target function f̄(x) = >x where x, 2 RD and a linear kernel K(x,x0) = 1 Dx
>x0 (in this model M = D). Data are sampled from zero-mean Gaussians with arbitrary covariance matrices C and C̃ for training and test distributions, respectively.
In this case, the kernel eigenvalue equation can be solved exactly. Denoting the eigenvalues and eigenvectors of the covariance matrix of the training distribution as CU = U⌃ where ⌃⇢ = 2⇢ ⇢ is the diagonal eigenvalue matrix and U is an orthogonal matrix with columns being eigenvectors of
C, the integral eigenvalue problem becomes ⌘⇢ ⇢(x) = hK(x, .), ⇢(.)ip(x) = ⇢ D u > ⇢ x. Therefore, the normalized eigenfunctions are ⇢(x) = u>⇢ x/ ⇢ and the eigenvalues are ⌘⇢ = 2⇢/D. The overlap matrix is
O⇢ = h ⇢(.), (.)ip̃(x) = ⌃ 1/2U>C̃U⌃ 1/2. (11)
Finally computing the target weights as a = ⌃1/2U> , we obtain the generalization error Eq.(1):
Eg = E 0,p(x) g + 0 1 "
2 + (D)2 >(PC+ DI) 1 ✓ C̃ 1
0 1 C ◆ (PC+ DI) 1 ,
= P TrC2(PC+ DI) 2, 0 = P Tr C̃C(PC+ DI) 2,
 = + TrC(PC+ DI) 1 (12) and
E0,p(x)g = 1 " 2 +
(D)2 1 >C(PC+ DI) 2 . (13)
As a consistency check, we note that the generalization error is minimized (Eg = 0) when C̃ = 0 corresponding to a Dirac measure at the origin. This makes sense since target function at origin is 0 and the estimator on the test distribution is also 0.
Next, we consider a diagonal covariance matrix C̃ = ̃2I for test distribution and C = diag ( 2, . . . 2| {z }
Mr , 0, . . . 0| {z } D Mr ) (14)
for training distribution to demonstrate how training distribution affects generalization in a simplified setting. The integer Mr corresponds to the rank of the training measure’s covariance which is also the number of non-zero kernel eigenvalues for the linear kernel. Furthermore we take the target function to have power only in the first N features ( ⇢>N = 0) where we adopt to the normalizationPN
⇢=1 2 ⇢ = 1. Thus, the target does not depend on the D N remaining dimensions x⇢>N and we
study how compressing some directions in training distribution influences generalization.
In this case, the self-consistent equation for  becomes exactly solvable. The generalization error reduces to:
Eg = ̃ 2
✓ "2
2 ↵ (0 + ↵)2 ↵ + 02 (0 + ↵)2 ↵
MrX
⇢=1
2⇢ + NX
⇢=Mr+1
2⇢
◆ , (15)
where 0 =  2Mr/D = 1 2
⇥ (1 + ̃ ↵) + q (1 + ̃+ ↵
2 4↵], ̃ = 2Mr/D and ↵ = P/Mr. This result matches that of [16] analyzing in-distribution generalization error when Mr = N = D and 2 = ̃2. We identify ̃ as an effective regularization parameter as it assumes the role plays for in-distribution generalization (compare to Eq. 7 of [16]).
First, we note that the generalization linearly scales with the variance of the test distribution. This is due to the linearity of the target and the estimator; any mismatch between them is amplified when the test points are further away from the origin.
Next, when the dimension of training measure is smaller than the dimension of the target function, Mr < N , there is an irreducible error due to the fact that kernel machine cannot fit the N Mr dimensions which are not expressed in the training data (Figure 3a). However in certain cases choosing Mr < N may help generalization if the data budget is limited to P ⇠ Mr since the learning happens faster compared to larger Mr. As illustrated in Figure 3a, if the target power is distributed such that PN ⇢=N 0+1 2 ⇢ ⌧ PN 0 ⇢=1 2 ⇢ , choosing Mr = N 0 can be a better strategy despite the irreducible error due to unexplored target power. In Figure 3a, we picked N = 60 and N 0 = 40 which is the number of directions target places most of its power on. If the data budget was limited to P < N , choosing Mr = N 0 (orange curve in Figure 3a) performs better than Mr = N (green curve) although for P N the Mr < N curve has irreduible error while the Mr = N does not. Reducing Mr below N 0 (blue curve in Figure 3a) on the other hand does not provide much benefit at small P .
Next, we observe that 0 is a monotonically decreasing function of ↵ and shows a sharp decrease near ↵ ⇡ 1 + ̃ which also implies a sharp change in generalization error. In fact when ̃ = 0, 0 becomes zero at ↵ = 1 and its first derivative diverges implying either a divergence in generalization error due to label noise or vanishing of Eg in the absence of noise (Section SI.3). When the labels are noisy, this non-monotonic behavior ↵ = 1 in Eg (called sample-wise double-descent) signals the over-fitting of the labels beyond which the kernel machine is able to average over noise and converge to the true target function [24, 31, 32, 33, 34, 35, 36, 37, 38, 39, 16]. Hence reducing Mr means that this transition occurs earlier in the learning curve (P ⇠ (1 + ̃)Mr) which implies that less training samples are necessary to estimate the target function well. Intuitively, sampling from an effectively higher dimensional space increases the necessary amount of training data to fit the modes in all directions. In Figure 3(b), we demonstrate this prediction for a linear target and obtain perfect agreement with experiment.
Finally, we show that increasing effective regularization ̃, which is controlled by the training distribution ( 2, Mr), leads to the suppression of double-descent peak, hence avoiding possible non-monotonicities. However large ̃ leads to slower learning creating a trade-off. Results from previous studies [36, 16] have shown the existence of an optimal ridge parameter in linear regression minimizing the generalization error for all P . Minimizing Eq.(15) with respect to , we find that the optimal ridge parameter is given by:
⇤ = Mr D "2. (16)
In Figure 3(b), we show that choosing optimal ridge parameters indeed mitigates the double descent and results in the best generalization performance.

heading: 5 Further Results
text: In SI, we present other analytically solvable models and further analysis on distribution mismatch for real data applications.
• In Section SI.2, we analyze the gradient descent/ascent procedures performed in Section 3 and provide further experiments to motivate how our theory can be used to find mismatched train/test distributions which improve generalization. We also apply our theory to adversarial attacks during testing.
• In Section SI.3, we study a more general linear regression model with diagonal overlap matrix where train/test distributions, target function and number of directions kernel expresses vary. When the target has out-of-RKHS components, we show how distribution shifts may help in generalization.
• In Section SI.4, we apply our theory to rotation invariant kernels such as the Gaussian RBF kernel and NTK [12] acting on spherical data in high-dimensions. We examine how a mismatched sphere radius affects the generalization error in the limit P,D ! 1 similar to the case solved in [15, 16].
• In Section SI.5, we study interpolation versus extrapolation in linear regression and regression in Fourier space by studying rectangular distributions with different ranges for one-dimensional inputs.

heading: 6 Discussion
text: Interest in kernel methods has recently surged due to their profound connections to deep neural networks [12, 13, 40, 41], and the replica method of statistical physics proved useful in studying its generalization properties, inductive biases and non-monotonic learning curves [15, 16, 28, 42, 43, 44, 37, 45]. Along the lines of these works, we analyzed generalization performance of kernel regression under distribution mismatch and obtained analytical expressions which fit experiments perfectly including those with wide neural networks. We demonstrated that our formula can be used to optimize over training distribution on MNIST which revealed that certain digits are better to sample more often than others. We considered several analytically solvable models, particularly linear regression for which we demonstrated how dimensional reduction of the training distribution can be helpful in generalization, including shifting of a double-descent peak. Our theory brings many insights about how kernel regression generalizes to unseen distributions.
In this work, we focused on generalization error for specific mismatched training and test distributions, as in [4, 46]. One could, instead, consider a set of possible test distributions that the model could encounter, and assess out-of-distribution generalization error based on the worst performance on that set. This distributional robustness approach [47] has been the focus of most previous studies [19, 48, 49, 50] where one can train a model by minimizing a more general empirical risk on a subset of all test distributions using robust optimization techniques [51]. More recent approaches aim to learn invariant relationships across all distributions [21, 22, 20]. It will be interesting to see if our approach can be generalized to ensembles of training and test distributions.
While our theory accurately describes experiments, it has limitations. First, using it to optimize for a training distribution requires the knowledge of exact test distribution which could not be available at the time of training. Second, the theory requires an eigendecomposition of the kernel which is computationally costly for large datasets. This problem, however, can potentially be solved by using stochastic methods to compute the kernel regression solution [52]. Third, our theory uses the replica theory [14], which is not fully rigorous. Fourth, if to be used to describe neural networks, our theory’s applicability is limited to their kernel regime.

Review summaries:
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.

This paper studies the generalization capabilities in kernel regression in the particular case where the distribution of training and testing is different.  They quantify the mismatch between distributions for a given kernel and show that its a key determinant of generalization performance.  They also propose how to optimize training and test distributions for a given data budget to find the best and worst case generalizations under a shift. The authors investigate how in real data this hypothesis can be confirmed and show how learning rate with training set size changes with training distribution and cases where choosing it different than test distribution helps. Their approach provides analytical tool for having more precise predictions.

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.