Integrated Latent Heterogeneity and Invariance Learning in Kernel Space

1 Introduction . 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 can not 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 . 2 Preliminaries . 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 . 3 Method . 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 .

Summary: 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.