Kernelized Heterogeneous Risk Minimization
Year: 2022

Abstract:
The ability to generalize under distributional shifts is essential to reliable machine learning, while models optimized with empirical risk minimization usually fail on non-i.i.d testing data. Recently, invariant learning methods for out-of-distribution (OOD) generalization propose to find causally invariant relationships with multienvironments. However, modern datasets are frequently multi-sourced without explicit source labels, rendering many invariant learning methods inapplicable. In this paper, we propose Kernelized Heterogeneous Risk Minimization (KerHRM) algorithm, which achieves both the latent heterogeneity exploration and invariant learning in kernel space, and then gives feedback to the original neural network by appointing invariant gradient direction. We theoretically justify our algorithm and empirically validate the effectiveness of our algorithm with extensive experiments.

heading: 1 Introduction
text: Traditional machine learning algorithms which optimize the empirical risk often suffer from poor generalization performance under distributional shifts caused by latent heterogeneity or selection biases that widely exist in real-world data[9, 20]. How to guarantee a machine learning algorithm with good generalization ability on data drawn out-of-distribution is of paramount significance, especially in high-stake applications such as financial analysis, criminal justice and medical diagnosis, etc.[13, 18], which is known as the out-of-distribution(OOD) generalization problem[1].
To ensure the OOD generalization ability, invariant learning methods assume the existence of the causally invariant correlations and exploit them through given environments, which makes their performances heavily dependent on the quality of environments. Further, the requirements for the environment labels are too strict to meet with, since real-world datasets are frequently assembled by merging data from multiple sources without explicit source labels. Recently, several works[5, 15] to relax such restrictions have been proposed. Creager et al.[5] directly infer the environments according to a given biased model first and then performs invariant learning. But the two stages cannot be jointly optimized and the quality of inferred environments depends heavily on the pre-provided biased model. Further, for complicated data, using invariant representation for environment inference is harmful, since the environment-specific features are gradually discarded, causing the extinction of latent heterogeneity and rendering data from different latent environments undistinguishable. Liu et al.[15] design a mechanism where two interactive modules for environment inference and invariant learning respectively can promote each other. However, it can only deal with scenarios where invariant and variant features are decomposed on raw feature level, and will break down when the decomposition can only be performed in representation space(e.g., image data). ∗Equal Contributions †Corresponding Author
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
This paper focuses on the integration of latent heterogeneity exploration and invariant learning on representation level. In order to incorporate representation learning with theoretical guarantees, we introduce Neural Tangent Kernel(NTK[10]) into our algorithm. According to NTK theory[10], training the neural network is equivalent to linear regression using Neural Tangent Features(NTF), which converts non-linear neural networks into linear regression in NTF space and makes the integration possible. Based on this, our Kernelized Heterogeneous Risk Minimization (KerHRM) algorithm is proposed, which synchronously optimizes the latent heterogeneity exploration module Mc and invariance learning moduleMp in NTF space. Specifically, we propose our novel Invariant Gradient Descent(IGD) forMp, which performs invariant learning in NTF space and then feeds back to neural networks with appointed invariant gradient direction. ForMc, we construct an orthogonal heterogeneity-aware kernel to capture the environment-specific features and to further accelerate the heterogeneity exploration. Theoretically, we demonstrate our heterogeneity exploration algorithm for Mc with rate-distortion theory and justify the orthogonality property of the built kernel, which jointly can illustrate the mutual promotion between the two modules. Empirically, experiments on both synthetic and real-world data validate the superiority of KerHRM in terms of good out-of-distribution generalization performance.

heading: 2 Preliminaries
text: Following [1, 3], we consider dataD = {De}e∈supp(Etr) with different sources dataDe = {Xe, Y e} collected from multiple training environments Etr. Here environment labels are unavailable as in most of the real applications. Etr is a random variable on indices of training environments and P e is the distribution of data and label in environment e. The goal of this work is to find a predictor f(·) : X → Y with good out-of-distribution generalization performance, which is formalized as:
arg min f max e∈supp(E)
L(f |e) (1)
where L(f |e) = Ee[`(Xe, Y e)] represents the risk of predictor f on environment e, and `(·, ·) : Y × Y → R+ the loss function. Note that E is the random variable on indices of all possible environments such that supp(Etr) ⊂ supp(E). Usually, for all e ∈ supp(E) \ supp(Etr), the data and label distribution P e(X,Y ) can be quite different from that of training environments Etr. Therefore, the problem in equation 1 is referred to as Out-of-Distribution (OOD) Generalization problem [1]. Since it is impossible to characterize the latent environments E without any prior knowledge or structural assumptions, the invariance assumption is proposed for invariant learning: Assumption 2.1. There exists random variable Ψ∗S(X) such that the following properties hold: a. Invariance property: for all e, e′ ∈ supp(E), we have P e(Y |Ψ∗S(X)) = P e ′ (Y |Ψ∗S(X)) holds. b. Sufficiency property: Y = f(Ψ∗S) + , ⊥ X .
This assumption indicates invariance and sufficiency for predicting the target Y using Ψ∗S , which is known as invariant representations with stable relationships with Y across E . To acquire such Ψ∗S , a branch of works[4, 11, 15] proposes to find the maximal invariant predictor(MIP) of an invariance set, which are defined as follows: Definition 2.1. The invariance set I with respect to E is defined as:
IE = {ΨS(X) : Y ⊥ E|ΨS(X)} = {ΨS(X) : H[Y |ΨS(X)] = H[Y |ΨS(X), E ]} (2)
where H[·] is the Shannon entropy of a random variable. The corresponding maximal invariant predictor (MIP) of IE is defined as S = arg maxΦ∈IE I(Y ; Φ), where I(·; ·) measures Shannon mutual information between two random variables.
Firstly, we propose that using the maximal invariant predictor S of IE can guarantee OOD optimality in Theorem 2.1. The formal statement is similar to [15] and can be found in Appendix A.3. Theorem 2.1. (Optimality Guarantee, informal) For predictor Φ∗(X) satisfying Assumption 2.1, Ψ∗S is the maximal invariant predictor with respect to E and the solution to OOD problem in equation 1 is EY [Y |Ψ∗S ] = arg minf supe∈supp(E) E[L(f)|e].
However, recent works[4, 11] on finding MIP solutions rely on the availability of data from multiple training environments Etr, which is hard to meet with in practice. Further, their validity is highly determined by the given Etr. Since IE ⊆ IEtr , the invariance regularized by Etr is often too large and the learned MIP may contain variant components and fails to generalize well. Based on this,
Heterogeneous Risk Minimization(HRM[15]) proposes to generate environments Etr with minimal |IEtr | and to conduct invariant prediction with learned Etr. However, the proposed HRM can only deal with simple scenarios where X = [Ψ∗S ,Ψ ∗ V ] T on raw feature level (Ψ∗S are invariant features and Ψ∗V variant ones), and will break down where X = h(Ψ ∗ S ,Ψ ∗ V ) (h(·, ·) is an unknown transformation function), since the decomposition can only be performed in representation space. In this work, we focus on the integration of latent heterogeneity exploration and invariant learning in general scenarios where invariant features are latent in X , which can be easily fulfilled in real applications.
Problem 1. (Problem Setting) Assume thatX = h(Ψ∗S ,Ψ∗V ) ∈ Rd, where Ψ∗S satisfies Assumption 2.1, h(·) is an unknown transformation function and Ψ∗S ⊥ Ψ∗V (following functional representation lemma[7]), given heterogeneous dataset D = {De}e∈supp(Elatent) without environment labels, the task is to generate environments Elearn with minimal |IElearn | and meanwhile learn invariant models.

heading: 3 Method
text: Remark. Following the analysis in section 2, to generate environments Elearn with minimal |IElearn | is equivalent to generate environments with as varying P (Y |Ψ∗V ) as possible, so as to exclude variant parts Ψ∗V from the invariant set IElearn .
In spite of such insight, the latent Ψ∗S ,Ψ ∗ V make it impossible to directly generate Elearn. In this work, we propose our Kernelized Heterogeneous Risk Minimization (KerHRM) algorithm with two interactive modules, the frontendMc for heterogeneity exploration and backendMp for invariant prediction. Specifically, given pooled data, the algorithm starts with the heterogeneity exploration moduleMc with a learned heterogeneity-aware kernel κc to generate Elearn. The learned environments are used byMp to produce invariant direction θinv in Neural Tangent Feature(NTF) space that captures the invariant components ΨS , and then θinv is used to guide the gradient descent of neural networks. After that, we update the kernel κc to orthogonalize with the invariant direction θinv so as to better capture the variant components ΨV and realize the mutual promotion between Mc andMp iteratively. The whole framework is jointly optimized, so that the mutual promotion between heterogeneity exploration and invariant learning can be fully leveraged. For smoothness we
Algorithm 1 Kernelized Heterogeneous Risk Minimization (KerHRM) Algorithm Input: Heterogeneous dataset D = {De = (Xe, Y e)}e∈Etr Initialization: MLP model fw(·) with initialized w0, Neural Tangent Feature Φ(X) = ∇wfw0(X)(fixed in following), clustering kernel initialized as κ(0)c (x1, x2) = xT1 x2 for t = 1 to T do
1. Generate E(t)learn with clustering kernel κ (t−1) c : E(t)learn =Mc((Φ(X), Y ), κ (t−1) c ) 2. Learn invariant model parameters θ(t)inv with E (t) learn in NTF space: θ (t) inv =Mp(E (t) learn) 3. Feedback to Neural Network fw(·) with θ(t)inv: w (t) inv = arg minw L(w;X,Y ) + Reg(w, θ (t) inv) 4. Update the clustering kernel κ(t)c with θ (t) inv: κ (t) c ← Orthogonal Transform(κ(t−1)c , θ(t)inv)
end for
begin with the invariant prediction step to illustrate our algorithm, and the flow of whole algorithm is shown in figure 1.
3.1 Mp: Invariant Gradient Descent with Elearn (Step 1)
For our invariant learning moduleMp, we propose Invariant Gradient Descent (IGD) algorithm. Taking the learned environments Elearn as input, our IGD firstly performs invariant learning in Neural Tangent Feature (NTF[10]) space to obtain the invariant direction θinv, and then guides the whole neural network fw(·) with θinv to learn the invariant model(neural network)’s parameters winv . Neural Tangent Feature Space The NTK theory[10] shows that training the neural network is equivalent to linear regression using non-linear NTFs φ(x), as in equation 4. For each data point x ∈ Rd, where d is the feature dimension, the corresponding feature is given by φ(x) = ∇wfw(x) ∈ Rp, where p is the number of neural network’s parameters. Firstly, we would like to dissect the feature components within φ(x) by decomposing the invariant and variant components hidden in φ(x). Therefore, we propose to perform Singular Value Decomposition (SVD) on the NTF matrix:
Φ(X) T︸ ︷︷ ︸
Rn×p
≈ U︸︷︷︸ Rn×k · S︸︷︷︸ Rk×k · V T︸︷︷︸ Rk×p
where p n ≥ k (3)
Intuitively, in equation 3, each row V Tj,: of V T represents the j-th feature component of Rp and we
take k such different feature components with the top k largest singular values to represent the data, and the rationality of low rank decomposition is guaranteed theoretically[21, 16] and empirically[2]. Since SVD ensures every feature components orthogonal, the neural tangent feature of the i-th data point can be decomposed into φ(xi)T ≈ ∑k j=1 Ui,j ·Sj,j ·V Tj,:, where Ui,j ·Sj,j denotes the strength of the j-th feature component in the i-th data. However, since neural networks have millions of parameters, the high dimension prevents us from learning directly on high dimensional NTFs Φ(X). Therefore, we rewrite the initial formulation of linear regression into:
fw(X) ≈ fw0(X) + Φ(X) T (w − w0) ≈ fw0(X) + USV T (w − w0) (4) = fw0(X) + Ψ(X) ( V T (w − w0) ) = fw0(X) + Ψ(X)θ (5)
where we let θ = V T (w−w0) ∈ Rk which reflects how the model parameter w utilizes the k feature components. Since V T is orthogonal, fitting w − w0 with features Φ(X) is equivalent to fitting θ using reduced NTFs Ψ(X). In this way, we convert the original high-dimensional regression problem into the low-dimensional one in equation 5, since in wide neural networks, we have p n ≥ k. Invariant Learning with Reduced NTFs Ψ(X) We could perform invariant learning on reduced NTFs Ψ(X) in linear space. In this work, we adopt the invariant regularizer proposed in [11] to learn θ = V T (w − w0) due to its optimality guarantees, and the objective function is:
θinv = arg min θ ∑ e∈Elearn Le(θ; Ψ, Y ) + α ·VarElearn (∇θL e) (6)
Guide Neural Network with invariant direction θinv With the learned θinv, it remains to feed back to the neural network’s parameters w. Since for neural networks with millions of parameters whose p ≈ 108, it is difficult to directly obtain w as w = w0 + V θinv. Therefore, we design a loss function to approximate the projection (w−w0//V θinv). Note that fw(X) = fw0(X)+USV T (w− w0) = fw0(X) + USθinv , we have
S−1UT (fw(X)− fw0(X)) = θinv (7)
Therefore, we can ensure the updated parameters w satisfy that S−1UT (fw(X)− fw0(X)) ∈ Rk is parallel to θinv , which leads to the following loss function:
winv = arg min w L(w;X,Y ) + λ
( 1− | 〈 θinv, S −1UT (fw(X)− fw0(X)) 〉 |
‖θinv‖‖S−1UT (fw(X)− fw0(X))‖
) (8)
where L(w;X,Y ) is the empirical prediction loss over training data and the second term is to force the invariance property of the neural network.

heading: 3.2 Variant Component Decomposition with θinv (Step 2)
text: The core of our KerHRM is the mutual promotion of the heterogeneity exploration moduleMc and the invariant learning moduleMp. From our insight, we should leverage the variant components ΨV
to exploit the latent heterogeneity. Therefore, with the better invariant direction θinv learned byMp that captures the invariant components in data, it remains to capture better variant components ΨV so as to further accelerate the heterogeneity exploration procedure, for which we design a clustering kernel κc on the reduce NTF space of Rk with the help of θinv learned in section 3.1. Recall the NTF decomposition in equation 3, the initial similarity of two data points xi and xj can be decomposed as:
κ(0)c (xi, xj) = φ(xi) Tφ(xj) = 〈UiS,UjS〉 (9)
With the invariant direction θ(t)inv learned byMp in iteration t, we can wipe out the invariant components used by θ(t)inv via
Ψ (t+1) V (xi)← UiS − 〈 UiS, θ (t) inv 〉 θ (t) inv/‖θ (t) inv‖ 2 (10)
which gives a new heterogeneity-aware kernel that better captures the variant components Ψ∗V as κ (t+1) c (xi, xj) = Ψ (t+1) V (xi) TΨ (t+1) V (xj).
3.3 Mc: Heterogeneity exploration with κc (Step 3)
Mc takes one heterogeneous dataset as input, and outputs a learned multi-environment partition Elearn for invariant prediction module Mp, and we implement it as a clustering algorithm with kernel regression given the heterogeneity-aware κc(xi, xj) = ΨV (xi)TΨV (xj) that captures the variant components in data. Following the analysis above, only the variant components Ψ∗V should be leveraged to identify the latent heterogeneity, and therefore we use the kernel κc as well as ΨV (X) learned in section 3.2 to capture the different relationship between Ψ∗V and Y , for which we use P (Y |ΨV ) as the clustering centre. Specifically, we assume the j-th cluster centre PΘj (Y |ΨV (X)) to be a Gaussian around f(Θj ; ΨV (X)) as:
hj(ΨV (X), Y ) = PΘj (Y |ΨV (X)) = ( √ 2πσ)−1 exp(−(Y − f(Θj ; ΨV (X)))2/2σ2) (11)
For the given N = ∑ e∈supp(Elatent) |D
e| data points D = {ψV (xi), yi}Ni=1, the empirical distribution can be modeled as P̂N = 1N ∑N i=1 δψV (xi),yi . Under this setting, we propose one convex clustering algorithm, which aims at finding a mixture distribution in distribution set Q defined as:
Q = {Q = ∑ j∈[K] qjhj(ΨV (X), Y ),q ∈ ∆K} (12)
to fit the empirical data best. Therefore, the original objective function and the simplified one are:
min Q∈Q DKL(P̂N‖Q)⇔ min Θ,q Lc = − 1N ∑ i∈[N ] log ∑ j∈[K] qjhj(ψV (xi), yi)  (13) Note that our clustering algorithm differs from others since the cluster centres are learned models parameterized with Θ. As for optimization, we use EM algorithm to optimize the centre parameters Θ and the mixture weights q iteratively. Specifically, when optimizing the cluster centre model f(Θj ; ·), we use kernel regression with κc(·, ·) to avoid computing ΨV (X) and allow large k. For generating the learned environments Elearn, we assign i-th point to j-th cluster with probability Pi,j = qjhj(ψV (xi), yi)/ ∑ l∈[K] qlhl(ψV (xi), yi).

heading: 4 Theoretical Analysis
text: In this section, we provide theoretical justifications of the mutual promotion betweenMc andMp. Since our algorithm does not violate the theoretical analysis in [11] and [10] which proves that better Elearn fromMc benefits the MIP learned byMp, to finish the mutual promotion, we only need to justify that better θinv fromMp benefits the learning of Elearn inMc. 1. Using ΨV ∗ benefits the clustering. Firstly, we introduce Lemma 4.1 from [15] to show that using Ψ∗V benefits the clustering in terms of larger between-cluster distance. Lemma 4.1. For ei, ej ∈ supp(Elatent), assume thatX satisfying Assumption 2.1, then under reasonable assumption([15]), we have DKL(P ei(Y |X)‖P ej (Y |X)) ≤ DKL(P ei(Y |Ψ∗V )‖P ej (Y |Ψ∗V )).
Then similar to [14], we use the rate-distortion theory to demonstrate why larger DKL between cluster centres benefits our convex clustering as well as the quality of Elearn.
Theorem 4.1. (Rate-Distortion) For the proposed convex clustering algorithm, we have:
min Q∈Q DKL(P̂N ||Q) = min Θ I(I; J) + (1/2σ2)EI,J [d(ψV (xi), yi,Θj)] + Const (14)
where rij = P (j|ψV (xi), yi) is a discrete random variable over the space {1, 2, ..., N} × {1, 2, ...,K} which denotes the probability of i-th data point belonging to j-th cluster, I, J are the marginal distribution of random variable rij respectively, d(ψV (xi), yi,Θj) = (fΘj (ψV (xi))−yi)2 and I(·; ·) the Shannon mutual information. Note that the optimal r can be obtained by the optimal Θ and therefore we only minimize the r.h.s with respect to Θ.
Actually d models the conditional distribution P (Y |ΨV ). If in the underlying distribution of the empirical data P (Y |ΨV ) differs a lot between different clusters, the optimizer will put more efforts in optimizing EI,J [d(ψV (xi), yi,Θj)] to avoid inducing larger error, resulting in smaller efforts put on optimization of I(I; J) and a relatively larger I(I; J). This means data sample points I have a larger mutual information with cluster index J , thus the clustering is prone to be more accurate.
2. Orthogonality Property: Better θinv for better ΨV . Firstly, we prove the orthogonality property between θinv(equation 6) and parameters Θ of clustering centres fΘj (·).
Theorem 4.2. (Orthogonality Property) Denote the data matrix of j-th environment Xj and ΨjV = ΨV (X j), then for each Θj(j ∈ [K]) = ((ΨjV )TΨ j V ) −1(ΨjV )
TY j , we have Span(Θ) ⊆ Ker(θinv) and Span(ΨjV ) ⊆ Ker(θinv), where Span denotes the column space and Ker the null space.
Theorem 4.2 justifies that the parameter space for clustering model fΘ(·) as well as the space of learned variant components ΨV is orthogonal to the invariant direction θinv, which indicates that better invariant direction θinv regulates better variant components ΨV and therefore better heterogeneity. Taking (1) and (2) together, we conclude that better results(θinv) ofMp promotes the latent heterogeneity exploration inMc because of larger between-cluster distance. Finally, we use a linear but general setting for further clarification.
Example. Assume that data points from environments e ∈ E are generated as follows:
X = Y (Ψ∗S + βeΨ ∗ V ) +N (0,Σ) ∈ Rd (15)
where Y = ±1 with equal probability, the coefficient βe varies across environment e, Ψ∗S ∈ Rd is the invariant feature and following functional representation lemma [7] Ψ∗V is the variant feature with Ψ∗V ⊥ Ψ∗S ∈ Rd and its relationship with the target Y relies on the environment-specific βe. Remark. In example 4, whenMp achieves optimal, we have θinv = Ψ∗S , which is the mid vertical hyperplane of the two Gaussian distribution. Then following equation 10, we have ΨV = X − (XT θinv)θinv = Y βeΨ ∗ V , which directly shows that in the next iteration,Mc uses solely variant components Ψ∗V in X to learn environments Elearn with diverse P (Y |X) = P (Y |Ψ∗V ), which by lemma 4.1 and theorem 4.1 gives the best clustering results.

heading: 5 Experiments
text: In this section, we validate the effectiveness of our method on synthetic data and real-world data.
Baselines We compare our proposed KerHRM with the following methods:
• Empirical Risk Minimization(ERM): minθ EPtr [`(θ;X,Y )] • Distributionally Robust Optimization(DRO [6]): minθ supQ∈Df (Q,Ptr)≤ρ EQ[`(θ;X,Y )]
• Environment Inference for Invariant Learning(EIIL [5]):
min Φ max u ∑ e∈E 1 Ne ∑ i ui(e)`(w Φ(xi), yi)+λ‖∇w|w=1.0 1 Ne ∑ i ui(e)`(w Φ(xi), yi)‖2 (16)
• Heterogeneous Risk Minimization(HRM [15])
• Invariant Risk Minimization(IRM [1]) with environment Etr labels:
min Φ ∑ e∈Etr Le + λ‖∇w|w=1.0Le(w Φ)‖2 (17)
We choose one typical method[6] of DRO as DRO is another main branch of methods for OOD generalization problem of the same setting with us (no environment labels). And HRM and EIIL are another methods for inferring environments for invariant learning without environment labels. We choose IRM as another baseline for its fame in invariant learning, but note that IRM is based on multiple training environments and we provide Etr labels for it, while the others do not need. Further, for ablation study, we run KerHRM for only one iteration without the feedback loop and denote it as Static KerHRM(KerHRMs). For all experiments, we use a two-layer MLP with 1024 hidden units.
Evaluation Metrics To evaluate the prediction performance, for task with only one testing environment, we simply use the prediction accuracy of the testing environment. While for tasks with multiple environments, we introduce Mean_Error defined as Mean_Error = 1|Etest| ∑ e∈Etest L e, Std_Error
defined as Std_Error = √
1 |Etest|−1 ∑ e∈Etest(L
e −Mean_Error)2, which are mean and standard deviation error across Etest. And we use the average mean square error for Le.

heading: 5.1 Synthetic Data
text: Classification with Spurious Correlation Following [19], we induce the spurious correlation between the label Y ∈ {+1,−1} and a spurious attribute A ∈ {+1,−1}. Specifically, each environment is characterized by its bias rate r ∈ (0, 1], where the bias rate r represents that for 100 ∗ r% data, A = Y , and for the other 100 ∗ (1− r)% data, A = −Y . Intuitively, r measures the strength and direction of the spurious correlation between the label Y and spurious attribute A, where larger |r − 0.5| signifies higher spurious correlation between Y and A, and sign(r − 0.5) represents the direction of such spurious correlation, since there is no spurious correlation when r = 0.5. We assume X = H[S, V ]T ∈ R2d, where S ∈ Rd is the invariant feature generated from label Y and V the variant feature generated from spurious attribute A:
S|Y ∼ N (Y 1, σ2sId), V |A ∼ N (A1, σ2vId) (18)
and H ∈ R2d×2d is an random orthogonal matrix to scramble the invariant and variant component, which makes it more practical. Typically, we set σ2v ≥ σ2s to let the model more prone to use spurious V since V is more informative.
In training, we set d = 5 and generate 2000 data points, where 50% points are from environment e1 with r1 = 0.9 and the other from environment e2 with r2. For our method, we set the cluster number K = 2. In testing, we generate 1000 data points from environment e3 with r3 = 0.1 to induce distributional shifts from training. In this experiments, we vary the bias rate r2 of environment e2 and the scrambled matrix H which can be an orthogonal or identity matrix (as done in [1]), and results after 10 runs are reported in Table 1.
From the results, we have the following observations and analysis: ERM suffers from the distributional shifts between training and testing, which yields the worst performance in testing. DRO can only provide slight resistance to distributional shifts, which we think is due to the over-pessimism problem[8]. EIIL achieves the best training performance but also performs poorly in testing. HRM outperforms the above three baselines, but its testing accuracy is just around the random guess(0.50), which is due to the disturbance of the simple raw feature setting in [15]. IRM performs better when the heterogeneity between training environments is large(r2 is small), which verifies our analysis in section 2 that the performance of invariant learning methods highly depends on the quality of the given Etr. Compared to all baselines, our KerHRM performs the best with respect to highest testing accuracy and lowest (Train_Acc− Test_Acc), showing its superiority to IRM and original HRM. Further, we also empirically analyze the sensitivity to the choice of cluster numberK of our KerHRM. We set r2 = 0.80 and test the performance with K = {2, 3, 4, 5} respectively. Results compared with IRM are shown in Table 2. From the results, we can see that the cluster number of our methods does not need to be the ground truth number(ground truth is 2) and our KerHRM is not sensitive to the choice of cluster number K. Intuitively, we only need the learned environments to reflect the variance of relationships between P (Y |Ψ∗V ), but do not require the environments to be ground truth. However, we notice that when K is far away from the proper one, the convergence of clustering algorithm is much slower.
Regression with Selection Bias In this setting, we induce the spurious correlation between the label Y and spurious attributes V through selection bias mechanism, which is similar to that in [12]. We assume X = H[S, V ]T ∈
Rd and Y = f(S) + , where f(·) is a non-linear function such that P (Y |S) remains invariant across environments while P (Y |V ) changes arbitrarily. For simplicity, we select data (xi, yi) with probability P (xi, yi) according to a certain variable Vb ∈ V :
P̂ (xi, yi) = |r|−5∗|yi−sign(r)∗Vb| (19)
where |r| > 1. Intuitively, r eventually controls the strengths and direction of the spurious correlation between Vb and Y (i.e. if r > 0, a data point whose Vb is close to its y is more probably to be selected.). The larger value of |r| means the stronger spurious correlation between Vb and Y , and r > 0 means positive correlation and vice versa. Therefore, here we use r to define different environments.
In training, we generate 1000 points from environment e1 with a predefined r and 100 points from e2 with r = −1.1. In testing, to simulate distributional shifts, we generate data points for 6 environments with r ∈ [−2.9,−2.7, . . . ,−1.9]. We compare our KerHRM with ERM, DRO, EIIL and IRM. We conduct experiments with different settings on r and the scrambled matrix H .
From the results in Table 3, we have the following analysis: ERM, DRO and EIIL performs poor with respect to high average and stability error, which is similar to that in classification experiments(Table 1). The results of HRM are quite different in two scenarios, where Scenario 1 corresponds to the simple raw feature setting(H = I) in [15] but Scenario 2 violates such simple setting with random orthogonal H and greatly harms HRM. Compared to all baselines, our KerHRM achieves lowest average error in 5/6 settings, and its superiority is especially obvious in our more general setting(Scenario 2).
Colored MNIST To further validate our method’s capacity under general settings, we use the colored MNIST dataset, where data X are high-dimensional non-linear transformation from invariant features(digits Y ) and variant features(color C). Following [1], we build a synthetic binary classification task, where each image is colored either red or green in a way that strongly and spuriously correlates with the class label Y . Firstly, a binary label Y is assigned to each images according to its digits: Y = 0 for digits 0∼4 and Y = 1 for digits 5∼9. Secondly, we sample the color id C by flipping Y with probability e and therefore forms environments, where e = 0.1 for the first training environment, e = 0.2 for the second training environments and e = 0.9 for the testing environment. Thirdly, we induce noisy labels by randomly flipping the label Y with probability 0.2.
We randomly sample 2500 images for each environments, and the two training environments are mixed without environment label Etr for ERM, DRO, EIIL, HRMs and HRM, while for IRM, the Etr labels are provided. For IRM, we sample 1000 data from the two training environments respectively and select the hyper-parameters which maximize the minimum accuracy of two validation environments. Note that we have no access to the testing environment while training, therefore we cannot resort to testing data to select the best one, which is more reasonable and different from that in [1]. For the others, since we have no access to E labels, we simply pool the 2000 data points for validation. The results are shown in Table 4, where Perfect Inv. Model represents the oracle
results that can be achieved under this setting. We run each method for 5 times and report the average accuracy, and since the variance of all methods are relatively small, we omit it in the table.
From the results, our KerHRM generalize the HRM to much more complicated data and consistently achieves the best performances. KerHRM even outperforms IRM significantly in an unfair setting where we provide perfect environment labels for IRM, which shows the limitation of manually labeled environments. Further, to best show the mutual promotion betweenMc andMp, we plot the training and testing accuracy as well as the KL-divergence DKL of P (Y |C) between the learned Elearn over iterations in figure 2. From figure 2, we firstly validate the mutual promotion between Mc andMp since DKL and testing accuracy escalate synchronously over iterations. Secondly, figure 2 corresponds to our analysis in section 2 that the performance of invariant learning method is highly correlated to the heterogeneity of Etr, which sheds lights to the importance of how to leverage the intrinsic heterogeneity in training data for invariant learning.

heading: 5.2 Real-world Data
text: In this experiment, we test our method on a real-world regression dataset (Kaggle) of house sales prices from King County, USA3, where the target variable is the transaction price of the house and each sample contains 17 predictive variables, such as the built year, number of bedrooms, and square footage of home, etc. Since it is fairly reasonable to assume the relationships between predictive variables and the target vary along the time (for example, the pricing mode may change along the time), there exist distributional shifts in the price-prediction task with respect to the build year of houses. Specifically, the houses in this dataset were built between 1900 ∼ 2015, and we divide the whole dataset into 6 periods, where each contains a time span of two decades. Notice that the later periods have larger distributional shifts. We train all methods on the first period where built_year ∈ [1900, 1920) and test on the other 5 periods and report the average results over 10 runs in figure 3. For IRM, we further divide the period 1 into two decades for the Etr provided.
3https://www.kaggle.com/c/house-prices-advanced-regression-techniques/data
Analysis The testing errors of ERM and DRO increase sharply across environments, indicating the existence of the distributional shifts between environments. IRM performs better than ERM and DRO, which shows the usefulness of environment labels for OOD generalization and the possibility of learning invariant predictor from multiple environments. The proposed KerHRM outperforms EIIL and HRM, which validates its superiority of heterogeneity exploration. KerHRM even outperforms IRM, which indicates the limitation of manually labeled environments in invariant learning and the necessity of latent heterogeneity exploration.

heading: 6 Limitations
text: Although the proposed KerHRM is a competitive method, it has several limitations. Firstly, since inMc we take the model parameters as cluster centres, the strict convergence guarantee for our clustering algorithmMc is quite hard to analyze. And empirically, we find when the pre-defined cluster number K is far away from the ground-truth, the convergence ofMc will become quite slow. Further, such restriction also affects the analysis of the mutual promotion betweenMc andMp, which we can only empirically provide some verification. Besides, although we incorporate Neural Tangent Kernel to deal with data beyond raw feature level, how to deal with more complicated data still remains unsolved. Also, how to reduce the randomness induced by Neural Tangent Kernel and how to incorporate deep learning with the mutual promotion between the two modules needs further investigation, and we left it for future work.

heading: 7 Conclusion
text: In this paper, we propose the KerHRM algorithm for the OOD generalization problem, which achieves both the latent heterogeneity exploration and invariant prediction. From our theoretical and empirical analysis, we find that the heterogeneity of environments plays a key role in invariant learning, which is consistent with some recent analysis[17] and opens a new line of research for OOD generalization problem. Our code is available at https://github.com/LJSthu/Kernelized-HRM.

heading: Acknowledgements
text: This work was supported National Key R&D Program of China (No. 2018AAA0102004).

Review summaries:
This paper studies the out-of-distribution (OOD) generalization problem. It proposes a Kernel heterogeneity & Invariance Learning method, which jointly learns both the latent heterogeneity exploration and invariant learning in kernel space. This paper provides strong theoretical guarantees for the proposed method. We know the proposed method makes clear improvements over related baselines in OOD scenarios from the empirical results.

This paper proposes a novel optimization framework for Out-Of-Distribution (OOD) generalization problems. OOD generalization problem is one of the most important problems in today's machine learning,  and it is highly non-trivial to investigate it at the optimization framework level. The proposed framework focuses on the integration of latent heterogeneity discovery and invariant learning on representation level, by introducing neural tangent kernel (NTK). Theoretical analysis is also conducted to verify the mutual promotion of these two modules. Experiments show that the proposed method can significantly outperform the commonly used and latest optimization algorithms.