Invariance Principle Meets Information Bottleneck for Out-of-Distribution Generalization

1 Introduction . Recent years have witnessed an explosion of examples showing deep learning models are prone to exploiting shortcuts ( spurious features ) ( Geirhos et al. , 2020 ; Pezeshki et al. , 2020 ) which make them fail to generalize out-of-distribution ( OOD ) . In Beery et al . ( 2018 ) , a convolutional neural network was trained to classify camels from cows ; however , it was found that the model relied on the background color ( e.g. , green pastures for cows ) and not on the properties of the animals ( e.g. , shape ) . These examples become very concerning when they occur in real-life applications ( e.g. , COVID-19 detection ( DeGrave et al. , 2020 ) ) . To address these out-of-distribution generalization failures , invariant risk minimization ( Arjovsky et al. , 2019 ) and several other works were proposed ( Ahuja et al. , 2020 ; Pezeshki et al. , 2020 ; Krueger et al. , 2020 ; Robey et al. , 2021 ; Zhang et al. , 2021 ) . The invariance principle from causality ( Peters et al. , 2015 ; Pearl , 1995 ) is at the heart of these works . The principle distinguishes predictors that only rely on the causes of the label from those that do not . The optimal predictor that only focuses on the causes is invariant and min-max optimal ( Rojas-Carulla et al. , 2018 ; Koyama and Yamaguchi , 2020 ; Ahuja et al. , 2021 ) under many distribution shifts but the same is not true for other predictors . ∗Equal contribution . †Mila - Quebec AI Institute , Université de Montréal . Correspondence to : kartik.ahuja @ mila.quebec . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Our contributions . Despite the promising theory , invariance principle-based approaches fail in settings ( Aubin et al. , 2021 ) where invariant features capture all information about the label contained in the input . A particular example is image classification ( e.g. , cow vs. camel ) ( Beery et al. , 2018 ) where the label is a deterministic function of the invariant features ( e.g. , shape of the animal ) , and does not depend on the spurious features ( e.g. , background ) . To understand such failures , we revisit the fundamental assumptions in linear regression tasks , where invariance-based approaches were shown to provably generalize OOD . We show that , in contrast to the linear regression tasks , OOD generalization is significantly harder for linear classification tasks ; we need much stronger restrictions in the form of support overlap assumptions3 on the distribution shifts , or otherwise it is not possible to guarantee OOD generalization under interventions on variables other than the target class . We then proceed to show that , even under the right assumptions on distribution shifts , the invariance principle is insufficient . However , we establish that information bottleneck ( IB ) constraints ( Tishby et al. , 2000 ) , together with the invariance principle , provably works in both settings – when invariant features completely capture the information about the label and also when they do not . ( Table 1 summarizes our theoretical results presented later ) . We propose an approach that combines both these principles and demonstrate its effectiveness on linear unit tests ( Aubin et al. , 2021 ) and on different real datasets . 2 OOD generalization and invariance : background & failures . Background . We consider a supervised training data D gathered from a set of training environments Etr : D = { De } e∈Etr , where De = { xei , yei } n e i=1 is the dataset from environment e ∈ Etr and ne is the number of instances in environment e. xei ∈ Rd and yei ∈ Y ⊆ Rk correspond to the input feature value and the label for ith instance respectively . Each ( xei , y e i ) is an i.i.d . draw from Pe , where Pe is the joint distribution of the input feature and the label in environment e. Let X e be the support of the input feature values in the environment e. The goal of OOD generalization is to use training data D to construct a predictor f : Rd → Rk that performs well across many unseen environments in Eall , where Eall ⊃ Etr . Define the risk of f in environment e as Re ( f ) = E [ ` ( f ( Xe ) , Y e ) ] , where for example ` can be 0-1 loss , logistic loss , square loss , ( Xe , Y e ) ∼ Pe , and the expectation E is w.r.t . Pe . Formally stated , our goal is to use the data from training environments Etr to find f : Rd → Y to minimize min f max e∈Eall Re ( f ) . ( 1 ) So far we did not state any restrictions on Eall . Consider binary classification : without any restrictions on Eall , no method can reduce the above objective ( ` is 0-1 loss ) to below one . Suppose a method outputs f∗ ; if ∃ e ∈ Eall \ Etr with labels based on 1− f∗ , then it achieves an error of one . Some assumptions on Eall are thus necessary . Consider how Eall is restricted using invariance for linear regressions ( Arjovsky et al. , 2019 ) . Assumption 1 . Linear regression structural equation model ( SEM ) . In each e ∈ Eall Y e ← w∗inv · Zeinv + e , Zeinv ⊥ e , E [ e ] = 0 , E [ | e|2 ] ≤ σ2sup Xe ← S ( Zeinv , Zespu ) ( 2 ) where w∗inv ∈ Rm , Zeinv ∈ Rm , Zspu ∈ Ro , S ∈ Rd× ( m+o ) , S is invertible ( m+ o = d ) . We focus on invertible S but several results extend to non-invertible S as well ( see Appendix ) . 3Support is the region where the probability density for continuous random variables ( probability mass function for discrete random variables ) is positive . Support overlap refers to the setting where train and test distribution maybe different but share the same support . We formally define this later in Assumption 5 . Assumption 1 states how Y e andXe are generated from latent invariant features Zeinv 4 , latent spurious features Zespu and noise e. The relationship between label and invariant features is invariant , i.e. , w∗inv is fixed across all environments . However , the distributions of Z e inv , Z e spu , and e are allowed to change arbitrarily across all the environments . Suppose S is identity . If we regress only on the invariant features Zeinv , then the optimal solution is w ∗ inv , which is independent of the environment , and the error it achieves is bounded above by the variance of e ( σ2sup ) . If we regress on the entire Ze and the optimal predictor places a non-zero weight on Zespu ( e.g. , Z e spu ← Y e + ζe ) , then this predictor fails to solve equation ( 1 ) ( ∃ e ∈ Eall , Zespu →∞ , error→∞ , see Appendix for details ) . Also , not only regressing on Zeinv is better than on Z e , it can be shown that it is optimal , i.e. , it solves equation ( 1 ) under Assumption 1 and achieves a value of σ2sup for the objective in equation ( 1 ) . Invariant predictor . Define a linear representation map Φ : Rr×d ( that transforms Xe as Φ ( Xe ) ) and define a linear classifier w : Rk×r ( that operates on the representation w · Φ ( Xe ) ) . We want to search for representations Φ such that E [ Y e|Φ ( Xe ) ] is invariant ( in Assumption 1 if Φ ( Xe ) = Zeinv , then E [ Y e|Φ ( Xe ) ] is invariant ) . We say that a data representation Φ elicits an invariant predictor w ·Φ across the set of training environments Etr if there is a predictor w that simultaneously achieves the minimum risk , i.e. , w ∈ arg minw̃ Re ( w̃ · Φ ) , ∀e ∈ Etr . The main objective of IRM is stated as min w∈Rk×r , Φ∈Rr×d 1 |Etr| ∑ e∈Etr Re ( w · Φ ) s.t . w ∈ arg min w̃∈Rk×r Re ( w̃ · Φ ) , ∀e ∈ Etr . ( 3 ) Observe that if we drop the constraints in the above which search only over invariant predictors , then we get the standard empirical risk minimization ( ERM ) ( Vapnik , 1992 ) ( assuming all the training environments occur with equal probability ) . In all our theorems , we use 0-1 loss for binary classification Y = { 0 , 1 } and square loss for regression Y = R. For binary classification , the output of the predictor is given as I ( w ·Φ ( Xe ) ) , where I ( · ) is the indicator function that takes 1 if the input is ≥ 0 and 0 otherwise , and the risk is Re ( w ·Φ ) = E [ |I ( w ·Φ ( Xe ) ) −Y e| ] . For regression , the output of the predictor is w ·Φ ( Xe ) and the corresponding risk is Re ( w ·Φ ) = E [ ( w ·Φ ( Xe ) − Y e ) 2 ] . We now present the main OOD generalization result from Arjovsky et al . ( 2019 ) for linear regressions . Theorem 1 . ( Informal ) If Assumption 1 is satisfied , Rank [ Φ ] > 0 , |Etr| > 2d , and Etr lie in a linear general position ( a mild condition on the data in Etr , defined in the Appendix ) , then each solution to equation ( 3 ) achieves OOD generalization ( solves equation ( 1 ) , @ e ∈ Eall with risk > σ2sup ) . Despite the above guarantees , IRM has been shown to fail in several cases including linear SEMs in ( Aubin et al. , 2021 ) . We take a closer look at these failures next . Understanding the failures : fully informative invariant features vs. partially informative invariant features ( FIIF vs. PIIF ) . We define properties salient to the datasets/SEMs used in the OOD generalization literature . Each e ∈ Eall , the distribution ( Xe , Y e ) ∼ Pe satisfies the following properties . a ) ∃ a map Φ∗ ( linear or not ) , which we call an invariant feature map , such that E [ Y e ∣∣Φ∗ ( Xe ) ] is the same for all e ∈ Eall and Y e 6⊥ Φ∗ ( Xe ) . These conditions ensure Φ∗ maps to features that have a finite predictive power and have the same optimal predictor across Eall . For the SEM in Assumption 1 , Φ∗ maps to Zeinv . b ) ∃ a map Ψ∗ ( linear or not ) , which we call spurious feature map , such that E [ Y e ∣∣Ψ∗ ( Xe ) ] is not the same for all e ∈ Eall and Y e 6⊥ Ψ∗ ( Xe ) for some environments . Ψ∗ often creates a hindrance in learning predictors that only rely on Φ∗ . Note that Ψ∗ should not be a transformation of some Φ∗ . For the SEM in Assumption 1 , suppose Zespu is anti-causally related to Y e , then Ψ∗ maps to Zespu ( See Appendix for an example ) . In the colored MNIST ( CMNIST ) dataset ( Arjovsky et al. , 2019 ) , the digits are colored in such a way that in the training domain , color is highly predictive of the digit label but this correlation being spurious breaks down at test time . Suppose the invariant feature map Φ∗ extracts the uncolored digit and the spurious feature map Ψ∗ extracts the background color . Ahuja et al . ( 2021 ) studied two variations of the colored MNIST dataset , which differed in the way final labels are generated from original MNIST labels ( corrupted with noise or not ) . They showed that the IRM exhibits good OOD generalization ( 50 % improvement over ERM ) in anti-causal-CMNIST ( AC-CMNIST , original data from Arjovsky et al . ( 2019 ) ) but is no different from ERM and fails in covariate shift-CMNIST ( CSCMNIST ) . In AC-CMNIST , the invariant features Φ∗ ( Xe ) ( uncolored digit ) are partially informative about the label , i.e. , Y 6⊥ Xe|Φ∗ ( Xe ) , and color contains information about label not contained 4In many examples in the literature , invariant features are causal , but not always ( Rosenfeld et al. , 2021 ) . in the uncolored digit . On the other hand in CS-CMNIST , invariant features are fully informative about the label , i.e. , Y ⊥ Xe|Φ∗ ( Xe ) , i.e. , they contains all the information about the label that is contained in input Xe . Most human labelled datasets have fully informative invariant features ; the labels ( digit value ) only depend on the invariant features ( uncolored digit ) and spurious features ( color of the digit ) do not affect the label . 5 In the rare case , when the humans are asked to label images in which the object being labelled itself is blurred , humans can rely on spurious features such as the background making such a data representative of PIIF setting . In Table 2 , we divide the different datasets used in the literature based on informativeness of the invariant features . We observe that when the invariant features are fully informative , both IRM and ERM fail but only in classification tasks and not in regression tasks ( Ahuja et al. , 2021 ) ; this is consistent with the linear regression result in Theorem 1 , where IRM succeeds regardless of whether Y e ⊥ Xe|Zeinv holds or not . Motivated by this observation , we take a closer look at the classification tasks where invariant features are fully informative .

Summary: This work gives an indepth analysis to the sufficient conditions on the generalization of invariant risk minimization (IRM) in classification setting. Authors show that in order to guarantee good generalization, the support overlap condition of invariant features is required. Further, authors also show and validated that if the support overlap condition does not satisfied for spurious features, the combination of IRM objective and information bottleneck regularization (i.e., minimzing the mutual information $I(X; \Phi(X))$) is a reliable remedy.