Invariance Principle Meets Information Bottleneck for Out-of-Distribution Generalization
Year: 2021

Abstract:
The invariance principle from causality is at the heart of notable approaches such as invariant risk minimization (IRM) that seek to address out-of-distribution (OOD) generalization failures. Despite the promising theory, invariance principle-based approaches fail in common classification tasks, where invariant (causal) features capture all the information about the label. Are these failures due to the methods failing to capture the invariance? Or is the invariance principle itself insufficient? To answer these questions, we revisit the fundamental assumptions in linear regression tasks, where invariance-based approaches were shown to provably generalize OOD. In contrast to the linear regression tasks, we show that for linear classification tasks we need much stronger restrictions on the distribution shifts, or otherwise OOD generalization is impossible. Furthermore, even with appropriate restrictions on distribution shifts in place, we show that the invariance principle alone is insufficient. We prove that a form of the information bottleneck constraint along with invariance helps address key failures when invariant features capture all the information about the label and also retains the existing success when they do not. We propose an approach that incorporates both of these principles and demonstrate its effectiveness in several experiments.

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

heading: 2 OOD generalization and invariance: background & failures
text: 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.

heading: 3 OOD generalization theory for linear classification tasks
text: A two-dimensional example with fully informative invariant features. We start with a 2D classification example (based on Nagarajan et al. (2021)), which can be understood as a simplified version of the CS-CMNIST dataset (Ahuja et al., 2021), Example 2/2S of Aubin et al. (2021), where both IRM and ERM fail. The example goes as follows. In each training environment e ∈ Etr
Y e ← I ( Xeinv − 1
2
) , where Xeinv ∈ {0, 1} is Bernoulli (1 2 ) ,
Xespu ← Xeinv ⊕W e, where W e ∈ {0, 1} is Bernoulli ( 1− pe ) with selection bias pe > 1
2 ,
(4)
where Bernoulli(a) takes value 1 with probability a and 0 otherwise. Each training environment is characterized by the probability pe. Following Assumption 1, we assume that the labelling function does not change from Etr to Eall, thus the relation between the label and the invariant features does not change. Assume that the distribution of Xeinv and X e spu can change arbitrarily. See Figure 1a) for a pictorial representation of this example illustrating the gist of the problem: there are many classifiers with the same error on Etr while only the one identical to the labelling function I(Xeinv− 12 ) generalizes correctly OOD. Define a classifier I(winvxinv +wspuxspu − 12 (winv +wspu)). Define a set of classifiers S = {(winv, wspu) s.t. winv > |wspu|}. Observe that all the classifiers in S achieve a zero classification error on the training environments. However, only classifiers for which wspu = 0 solve the OOD generalization (eq. (1)). With Φ as the identity, it can be shown that all the classifiers S form an invariant predictor (satisfy the constraint in equation (3) over all the training environments when ` is the 0-1 loss). Observe that increasing the number of training environments to infinity does not address the problem, unlike with the linear regression result discussed in Theorem 1 (Arjovsky et al., 2019), where it was shown that if the number of environments increases linearly in the dimension of the data, then the solution to IRM also solves the OOD generalization (eq. (1)). 6 We use the above example to construct general SEMs for linear classification when the invariant features are fully informative. We follow the structure of the SEM from Assumption 1 in our construction.
5The deterministic labelling case was referred as realizable problems in (Arjovsky et al., 2019). 6Please note that this example illustrates certain important facets in a very simple fashion; only in this example a max-margin classifier can solve the problem but not in general. (Further explanation in the Appendix).
Assumption 2. Linear classification structural equation model (FIIF). In each e ∈ Eall
Y e ← I ( w∗inv · Zeinv ) ⊕Ne, Ne ∼ Bernoulli(q), q < 1
2 , Ne ⊥ (Zeinv, Zespu), Xe ← S ( Zeinv, Z e spu ) ,
(5)
where w∗inv ∈ Rm with ‖w∗inv‖ = 1 is the labelling hyperplane, Zeinv ∈ Rm, Zespu ∈ Ro, Ne is binary noise with identical distribution across environments, ⊕ is the XOR operator, S is invertible.
If noise level q is zero, then the above SEM covers linearly separable problems. See Figure 2a) for the directed acyclic graph (DAG) corresponding to this SEM. From the DAG observe that Y e ⊥ Xe|Zeinv, which implies that the invariant features are fully informative. Contrast this with a DAG that follows Assumption 1 shown in Figure 2b), where Y e 6⊥ Xe|Zeinv and thus the invariant features are not fully informative. If Eall follows the SEM in Assumption 2 and suppose the distribution of Zeinv, Zespu can change arbitrarily, then it can be shown that only a classifier identical to the labelling function I(w∗inv · Zeinv) can solve the OOD generalization (eq. (1)); such a classifier achieves an error of q (noise level) in all the environments. As a result, if for a classifier we can find e ∈ Eall that follows Assumption 2 where the error is greater than q, then such a classifier does not solve equation (1). Now we ask – what are the minimal conditions on training environments Etr to achieve OOD generalization when Eall follow Assumption 2? To achieve OOD generalization for linear regressions, in Theorem 1, it was required that the number of training environments grows linearly in the dimension of the data. However, there was no restriction on the support of the latent invariant and latent spurious features, and they were allowed to change arbitrarily from train to test (for further discussion on this, see the Appendix). Can we continue to work with similar assumptions for the SEM in Assumption 2 and solve the OOD generalization (eq. (1))? We state some assumptions and notations to answer that. Define the support of the invariant (spurious) features Zeinv (Z e spu) in environment e as Zeinv (Zespu).
Assumption 3. Bounded invariant features. ∪e∈EtrZeinv is a bounded set.7 Assumption 4. Bounded spurious features. ∪e∈EtrZespu is a bounded set.
Assumption 5. Invariant feature support overlap. ∀e ∈ Eall,Zeinv ⊆ ∪e′∈EtrZe ′ inv Assumption 6. Spurious feature support overlap. ∀e ∈ Eall,Zespu ⊆ ∪e′∈EtrZe ′ spu
Assumption 5 (6) states that the support of the invariant (spurious) features for unseen environments is the same as the union of the support over the training environments. It is important to note that support overlap does not imply that the distribution over the invariant features does not change. We now define a margin that measures how much the is training support of invariant features Zeinv separated by the labelling hyperplane w∗inv. Define Inv-Margin = minz∈∪e∈EtrZeinv sgn ( w∗inv · z )( w∗inv · z ) . This margin only coincides with the standard margin in support vector machines when the noise level q is 0 (linearly separable) and S is identity. If Inv-Margin > 0, then the labelling hyperplane w∗inv separates the support into two halves (see Figure 1b)).
7A set Z is bounded if ∃M <∞ such that ∀z ∈ Z, ‖z‖ ≤M .
Assumption 7. Strictly separable invariant features. Inv-Margin > 0.
Next, we show the importance of support overlap for invariant features.
Theorem 2. Impossibility of guaranteed OOD generalization for linear classification. Suppose each e ∈ Eall follows Assumption 2. If for all the training environments Etr, the latent invariant features are bounded and strictly separable, i.e., Assumption 3 and 7 hold, then every deterministic algorithm fails to solve the OOD generalization (eq. (1)), i.e., for the output of every algorithm ∃ e ∈ Eall in which the error exceeds the minimum required value q (noise level).
The proofs to all the theorems are in the Appendix. We provide a high-level intuiton as to why invariant feature support overlap is crucial to the impossibility result. In Figure 1b), we show that if the support of latent invariant features are strictly separated by the labelling hyperplane w∗inv, then we can find another valid hyperplane w+inv that is equally likely to have generated the same data. There is no algorithm that can distinguish between w∗inv and w + inv. As a result, if we use data from the region where the hyperplanes disagree (yellow region Figure 1b)), then the algorithm fails.
Significance of Theorem 2. We showed that without the support overlap assumption on the invariant features, OOD generalization is impossible for linear classification tasks. This is in contrast to linear regression in Theorem 1 (Arjovsky et al., 2019), where even in the absence of the support overlap assumption, guaranteed OOD generalization was possible. Applying the above Theorem 2 to the 2D case (eq. (4)) implies that we cannot assume that the support of invariant latent features can change, or else that case is also impossible to solve.
Next, we ask what further assumptions are minimally needed to be able to solve the OOD generalization (eq. (1)). Each classifier can be written as w̄ ·Xe = w̄ · S(Zeinv, Zespu) = w̃inv · Zeinv + w̃spuZespu. If w̃spu 6= 0, then the classifier w̄ is said to rely on spurious features. Theorem 3. Sufficiency and Insufficiency of ERM and IRM. Suppose each e ∈ Eall follows Assumption 2. Assume that a) the invariant features are strictly separable, bounded, and satisfy support overlap, b) the spurious features are bounded (Assumptions 3-5, 7 hold).
• Sufficiency: If the spurious features satisfy support overlap (Assumption 6 holds), then both ERM and IRM solve the OOD generalization problem (eq. (1)). Also, there exist solutions to ERM and IRM solutions that rely on the spurious features and still achieve OOD generalization.
• Insufficiency: If spurious features do not satisfy support overlap, then both ERM and IRM fail at solving the OOD generalization problem (eq. (1)). Also, there exist no such classifiers that rely on spurious features and also achieve OOD generalization.
Significance of Theorem 3. From the first part, we learn that if the support overlap is satisfied for both the invariant features and the spurious features, then either ERM or IRM can solve the OOD generalization (eq. (1)). Interestingly, in this case we can have classifiers that rely on the spurious features and yet solve the OOD generalization (eq. (1)). For the 2D case (eq. (4)) this case implies that the entire set S solves the OOD generalization (eq. (1)). From the second part, we learn that if support overlap holds for invariant features but not for spurious features, then the ideal OOD optimal predictors rely only on the invariant features. In this case, methods like ERM and IRM continue to rely on spurious features and fail at OOD generalization. For the above 2D case (eq. (4)) this implies that only the predictors that rely only on Xeinv in the set S solve the OOD generalization (eq. (1)). To summarize, we looked at SEMs for classification tasks when invariant features are fully informative, and find that the support overlap assumption over invariant features is necessary. Even in the presence of support overlap for invariant features, we showed that ERM and IRM can easily fail if the support overlap is violated for spurious features. This raises a natural question – Can we even solve the case with the support overlap assumption only on the invariant features? We will now show that the information bottleneck principle can help tackle these cases.

heading: 4 Information bottleneck principle meets invariance principle
text: Why the information bottleneck? The information bottleneck principle prescribes to learn a representation that compresses the input X as much as possible while preserving all the relevant information about the target label Y (Tishby et al., 2000). Mutual information I(X; Φ(X)) is used to measure information compression. If representation Φ(X) is a deterministic transformation of X , then in principle we can use the entropy of Φ(X) to measure compression (Kirsch et al., 2020). Let
us revisit the 2D case (eq. (4)) and apply this principle to it. Following the second part of Theorem 3, where ERM and IRM failed, assume that invariant features satisfy the support overlap assumption, but make no such assumption for the spurious features. Consider three choices for Φ: identity (selects both features), selects invariant feature only, selects spurious feature only. The entropy of H(Φ(Xe)) when Φ is the identity is H(pe) + log(2), where H(pe) is the Shannon entropy in Bernoulli(pe). If Φ selects the invariant/spurious features only, then H(Φ(Xe)) = log(2). Among all three choices, the one that has the least entropy and also achieves zero error is the representation that focuses on the invariant feature. We could find the OOD optimal predictor in this example just by using information bottleneck. Does it mean the invariance principle isn’t needed? We answer this next.
Why invariance? Consider a simple classification SEM. In each e ∈ Etr, Y e ← X1,einv ⊕X 2,e inv ⊕Ne and Xespu ← Y e ⊕ V e, where all the random variables involved are binary valued, noise Ne, V e are Bernoulli with parameters q (identical across Etr), ce (varies across Etr) respectively. If ce < q, then in Etr predictions based on Xespu are better than predictions based on X 1,e inv , X 2,e inv . If both X1,einv , X 2,e inv are uniform Bernoulli, then these features have a higher entropy than X e spu. In this case, the information bottleneck would bar using X1,einv , X 2,e inv . Instead, we want the model to focus on X 1,e inv , X2,einv and not on X e spu. Invariance constraints encourage the model to focus on X 1,e inv , X 2,e inv . In this example, observe that invariant features are partially informative unlike the 2D case (eq. (4)).
Why invariance and information bottleneck? We have illustrated through simple examples when the information bottleneck is needed but not invariance and vice-versa. We now provide a simple example where both these constraints are needed at the same time. This example combines the 2D case (eq. (4)) and the example we highlighted in the paragraph above: Y e ← Xeinv ⊕ Ne, X1,espu ← Xeinv ⊕W e, and X2,espu ← Y e ⊕ V e. In this case, the invariance constraint does not allow representations that use X2,espu but does not prohibit representations that rely on X 1,e spu . However, information bottleneck constraints on top ensure that representations that only use Xeinv are used. We now describe an objective 8 that combines both these principles:
min w,Φ ∑ e∈Etr he ( w ·Φ ) s.t. 1 |Etr| ∑ e∈Etr Re ( w ·Φ ) ≤ rth, w ∈ arg min w̃∈Rk×r Re(w̃ ·Φ),∀e ∈ Etr, (6)
where he in the above is a lower bounded differential entropy defined below and rth is the threshold on the average risk. Typical information bottleneck based optimization in neural networks involves minimization of the entropy of the representation output from a certain hidden layer. For both analytical convenience and also because the above setup is a linear model, we work with the simplest form of bottleneck which directly minimizes the entropy of the output layer. Recall the definition of differential entropy of a random variableX , h(X) = −EX [log dPX ] and dPX is the Radon-Nikodym derivative of PX with respect to Lebesgue measure. Because in general differential entropy has no lower bound, we add a small independent noise term ζ (Kirsch et al., 2020) to the classifier to ensure that the entropy is bounded below. We call the above optimization information bottleneck based invariant risk minimization (IB-IRM). In summary, among all the highly predictive invariant predictors we pick the ones that have the least entropy. If we drop the invariance constraint from the above optimization, we get information bottleneck based empirical risk minimization (IB-ERM). In the above formulation and following result, we assume that Xe are continuous random variables; the results continue to hold for discrete Xe as well (See Appendix for details). Theorem 4. IB-IRM and IB-ERM vs. IRM and ERM
8Results extend to alternate objective with information bottleneck constraints and average risk as objective.
• Fully informative invariant features (FIIF). Suppose each e ∈ Eall follows Assumption 2. Assume that the invariant features are strictly separable, bounded, and satisfy support overlap (Assumptions 3,5 and 7 hold). Also, for each e ∈ Etr Zespu ← AZeinv + W e, where A ∈ Ro×m, W e ∈ Ro is continuous, bounded, and zero mean noise. Each solution to IB-IRM (eq. (6), with ` as 0-1 loss, and rth = q), and IB-ERM solves the OOD generalization (eq. (1)) but ERM and IRM (eq.(3)) fail.
• Partially informative invariant features (PIIF). Suppose each e ∈ Eall follows Assumption 1 and ∃ e ∈ Etr such that E[ eZespu] 6= 0. If |Etr| > 2d and the set Etr lies in a linear general position (a mild condition defined in the Appendix), then each solution to IB-IRM (eq. (6), with ` as square loss, σ2 < r
th ≤ σ2Y , where σ2Y and σ2 are the variance in the label and noise across Etr) and IRM (eq.(3)) solves OOD generalization (eq. (1)) but IB-ERM and ERM fail.
Significance of Theorem 4 and remarks. In the first part (FIIF), IB-ERM and IB-IRM succeed without assuming support overlap for the spurious features, which was crucial for success of ERM and IRM in Theorem 3. This establishes that support overlap of spurious features is not a necessary condition. Observe that when invariant features are fully informative, IB-ERM and IB-IRM succeed, but when invariant features are partially informative IB-IRM and IRM succeed. In real data settings, we do not know if the invariant features are fully or partially informative. Since IB-IRM is the only common winner in both the settings, it would be pragmatic to use it in the absence of domain knowledge about the informativeness of the invariant features. In the paragraph preceding the objective in equation (6), we discussed examples where both the IB and IRM constraints were needed at the same time. In the Appendix, we generalize that example and show that if we change the assumptions in linear classification SEM in Assumption 2 such that the invariant features are partially informative, then we see the joint benefit of IB and IRM constraints. At this point, it is also worth pointing to a result in Rosenfeld et al. (2021), which focused on linear classification SEMs (DAG shown in Figure 2c) with partially informative invariant features. Under the assumption of complete support overlap for spurious and invariant features, authors showed IRM succeeds.

heading: 4.1 Proposed approach
text: We take the three terms from the optimization in equation (6) and create a weighted combination as∑ e ( Re(Φ)+λ‖∇w,w=1.0Re(w·Φ)‖2+νhe(Φ) ) ≤ ∑ e ( Re(Φ)+λ‖∇w,w=1.0Re(w·Φ)‖2+νh(Φ) ) .
In the LHS above, the first term corresponds to the risks across environments, the second term approximates invariance constraint (follows the IRMv1 objective (Arjovsky et al., 2019)), and the third term is the entropy of the classifier in each environment.
gence of |wspu|√
w2spu+w 2 inv
(metric from
Nagarajan et al. (2021)) for average selection bias p = 0.9.
In the RHS, h(Φ) is the entropy of Φ unconditional on the environment (the entropy on the left-hand side is entropy conditional on the environment assuming all the environments are equally likely). Optimizing over differential entropy is not easy, and thus we resort to minimizing an upper bound of it (Kirsch et al., 2020). We use the standard result that among all continuous random variables with the same variance, Gaussian has the maximum differential entropy. Since the entropy of Gaussian increases with its variance, we use the variance of Φ instead of the differential entropy (For further details, see the Appendix). Our final objective is given as∑
e
( Re(Φ) + λ‖∇w,w=1.0Re(w · Φ)‖2 + γVar(Φ) ) . (7)
On the behavior of gradient descent with and without information bottleneck. In the entire discussion so far, we have focused on ensuring that the set of optimal solutions to the desired objective (IB-IRM, IB-ERM, etc.) correspond to the solutions of the OOD generalization problem (eq. (1)). In some simple cases, such as the 2D case (eq. (4)), it can be shown that gradient descent is biased towards selecting the ideal classifier (Soudry et al., 2018; Nagarajan et al., 2021). Even though gradient descent can eventually learn the ideal classifier that only relies on the invariant features, training is frustratingly slow as was shown by Nagarajan et al. (2021). In the next theorem, we characterize the impact of using IB penalty (Var(Φ)) in the 2D example (eq. (4)). We compare the methods in terms of |wspu(t)winv(t) |, which was the metric used in Nagarajan et al. (2021); wspu(t) and winv(t) are the weights for the spurious feature and the invariant feature at time t of training (assuming training happens with continuous time gradient descent).
Theorem 5. Impact of IB on learning speed. Suppose each e ∈ Etr follows the 2D case from equation (4). Set λ = 0, γ > 0 in equation (7) to get the IB-ERM objective with ` as exponential loss. Continuous-time gradient descent on this IB-ERM objective achieves |wspu(t)winv(t) | ≤ in time less than W0( 1 2γ ) 2(1−p) (W0(·) denotes the principal branch of the Lambert W function), while in the same time the
ratio for ERM |wspu(t)winv(t) | ≥ ln( 1+2p 3−2p )/ln
( 1 + W0( 1 2γ ) 2(1−p) ) , where p = 1|Etr| ∑ e∈Etr p e .
|wspu(t)winv(t) | converges to zero for both methods, but it converges much faster for IB-ERM (for p = 0.9, = 0.001, γ = 0.58, the ratio for IB-ERM is |wspu(t)winv(t) | ≤ 0.001 and ratio for ERM is | wspu(t) winv(t)
| ≥ 0.09). In the above theorem, we analyzed the impact of information bottleneck only. The convergence analysis for both the penalties jointly comes with its own challenges, and we hope to explore this in future work. However, we carried out experiments with gradient descent on all the objectives for the 2D example (eq. (4)). See Figure 3 for the comparisons.

heading: 5 Experiments
text: Methods, datasets & metrics. We compare our approaches – information bottleneck based ERM (IBERM) and information bottleneck based IRM (IB-IRM) with ERM and IRM. We also compare with an Oracle model trained on data where spurious features are permuted to remove spurious correlations. We use all the datasets in Table 2, Terra Incognita dataset (Beery et al., 2018), and COCO (Ahmed et al., 2021). We follow the same protocol for tuning hyperparameters from Aubin et al. (2021); Arjovsky et al. (2019) for their respective datasets (see the Appendix for more details). As is reported in literature, for Example 2/2S, Example 3/3S we use classification error and for AC-CMNIST, CS-CMNIST, Terra Incognita, and COCO we use accuracy. For Example 1/1S, we use mean square error (MSE). The code for experiments can be found at https://github.com/ahujak/IB-IRM.
Summary of results. In Table 3, we provide a comparison of methods for different examples in linear unit tests (Aubin et al., 2021) for three and six training environments. In Table 4, we provide a comparison of the methods for different CMNIST datasets, Terra Incognita and COCO dataset. Based on our Theorem 4, we do not expect ERM and IB-ERM to do well on Example 1/1S, Example 3/3S and AC-CMNIST as these datasets fall in the PIIF category, i.e, the invariant features are partially informative. On these examples, we find that IRM and IB-IRM do better than ERM and IB-ERM (for Example 3/3S when there are three environments all methods perform poorly). Based on our Theorem 4, we do not expect IRM and ERM to do well on Example 2/2S, CS-CMNIST, Terra Incognita and COCO dataset,9 as these datasets fall in the FIIF category, i.e., the invariant features are fully informative. On these FIIF examples, we find that IB-ERM always performs well (close to oracle), and in some cases IB-IRM also performs well. Our experiments confirm that IB penalty has a crucial role to play in FIIF settings and IRMv1 penalty has a crucial role to play in PIIF settings (to further this claim, we provide an ablation study in the Appendix). On Example 1/1S, AC-CMNIST, we find that IB-IRM is able to extract the benefit of IRMv1 penalty. On CS-CMNIST and Example 2/2S we find that IB-IRM is able to extract the benefit of IB penalty. In settings such as COCO dataset, where IB-IRM does not perform as well as IB-ERM, better hyperparameter tuning strategies should be able to help IB-IRM adapt and put a higher weight on IB penalty. Overall, we can conclude that IB-ERM improves over ERM (significantly in FIIF and marginally in PIIF settings), and IB-IRM improves over IRM (improves in FIIF settings and retains advantages in PIIF settings).
Remark. As we move from three to six environments, we observe that MSE in Example 1/1S exhibits a larger variance. This is because of the way data is generated, the new environments that are sampled have labels that have a higher noise level (we follow the same procedure as in Aubin et al. (2021)).

heading: 6 Extensions, limitations, and future work
text: Extension to non-linear models and multi-class classification. In this work our theoretical analysis focused on linear models. Consider the map X ← S(Zinv, Zspu) in Assumption 2. Suppose S is non-linear and bijective. We can divide the learning task into two parts a) invert S to obtain Zinv, Zspu and b) learn a linear model that only relies on the invariant features Zinv to predict the label Y . For
9We place Terra Incognita and COCO dataset in the FIIF assuming that the humans who labeled the images did not need to rely on unreliable/spurious features such as background to generate the labels.
part b), we can rely on the approaches proposed in this work. For part a), we need to leverage advancements in the field of non-linear ICA (Khemakhem et al., 2020). The current state-of-the-art to solve part a) requires strong structural assumptions on the dependence between all the components of Zinv, Zspu (Lu et al., 2021). Therefore, solving part a) and part b) in conjunction with minimal assumptions forms an exciting future work. In the entire work, the discussion was focused on binary classification tasks and regression tasks. For multi-class classification settings, we consider natural extension of the SEM in Assumption 2 (See the Appendix) and our main results continue to hold.
On the choice for IB penalty and IRMv1 penalty. We use the approximation for entropy (in equation (7)) described in Kirsch et al. (2020). The approximation (even though an upper bound) serves as an effective proxy for the true information bottleneck as shown in the experiments in Kirsch et al. (2020) (e.g., see their experiment on Imagenette dataset). Also, our experiments validate this approximation even in moderately high dimensions, as an example in CS-CMNIST, the dimension of the layer at which bottleneck constraints are applied is 256. Developing tighter approximations for information bottleneck in high dimensions and analyzing their impact on OOD generalization is an important future work. In recent works (Rosenfeld et al., 2021; Kamath et al., 2021; Gulrajani and Lopez-Paz, 2021), there has been criticism of different aspects of IRM, e.g., failure of IRMv1 penalty in non-linear models, the tuning of IRMv1 penalty, etc. Since we use IRMv1 penalty in our proposed loss, these criticisms apply to our objective as well. Other approximations of invariance have been proposed in the literature (Koyama and Yamaguchi, 2020; Ahuja et al., 2020; Chang et al., 2020). Exploring their benefits together with information bottleneck is a fruitful future work. Before concluding, we want to remark that we have already discussed the closest related works. However, we also provide a detailed discussion of the broader related literature in the Appendix.

heading: 7 Conclusion
text: In this work, we revisited the fundamental assumptions for OOD generalization for settings when invariant features capture all the information about the label. We showed how linear classification tasks are different and need much stronger assumptions than linear regression tasks. We provide a sharp characterization of performance of ERM and IRM under different assumptions on support overlap of invariant and spurious features. We showed that support overlap of invariant features is necessary or otherwise OOD generalization is impossible. However, ERM and IRM seem to fail even in the absence of support overlap of spurious features. We prove that a form of the information bottleneck constraint along with invariance goes a long way in overcoming the failures while retaining the existing provable guarantees.

heading: Acknowledgements
text: We thank Reyhane Askari Hemmat, Adam Ibrahim, Alexia Jolicoeur-Martineau, Divyat Mahajan, Ryan D’Orazio, Nicolas Loizou, Manuela Girotti, and Charles Guille-Escuret for the feedback. Kartik Ahuja would also like to thank Karthikeyan Shanmugam for discussions pertaining to the related works.
Funding disclosure
We would like to thank Samsung Electronics Co., Ldt. for funding this research. Kartik Ahuja acknowledges the support provided by IVADO postdoctoral fellowship funding program. Yoshua Bengio acknowledges the support from CIFAR and IBM. Ioannis Mitliagkas acknowledges support from an NSERC Discovery grant (RGPIN-2019-06512), a Samsung grant, Canada CIFAR AI chair and MSR collaborative research grant. Irina Rish acknowledges the support from Canada CIFAR AI Chair Program and from the Canada Excellence Research Chairs Program. We thank Compute Canada for providing computational resources.

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

This paper investigates the conditions that are needed to provably generalize OOD in the linear classification tasks. Comparing with the linear regression tasks, linear classification tasks need more strong conditions to guarantee the OOD generalization.   The paper first proposes the overlap assumptions on both invariant and spurious feature and analyze the OOD generalization guarantee under these assumptions. The theorem shows that only both two types of features satisfy the support overlap assumption can we expect the IRM to solve the OOD generalization problems. Then the paper study the possibility of solving the problem only with the support overlap assumption on the invariant features. The proposed method uses both information bottleneck and invariance constraints. Theoretical results show that the proposed IB-IRM can successfully solve the problem in both fully and partially informative invariant features situations. The experiment results are consistent with the theoretical results.