Invariance Principle Meets Information Bottleneck for Out-of-Distribution Generalization
Authors: Kartik Ahuja
Ethan Caballero
Dinghuai Zhang
Jean-Christophe Gagnon-Audet
Yoshua Bengio
Ioannis Mitliagkas
Irina Rish
Venue: NeurIPS 2021 Spotlight (2021)
Keywords: out-of-distribution generalization
invariance principle
information bottleneck

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 the 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.

Invariance Principle Meets Information Bottleneck for Out-of-Distribution Generalization
Kartik Ahuja
Ethan Caballero*
Dinghuai Zhang"
Jean- SennChhistoopee Gagnon-Audet 1
Yoshua Bengio f
Ioannis Mitliagkas'
Irina Rish
Abstract
The invariance principle from causality 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. contrast the linear regression tasks, we show that for linear classification tasks we need much stronger restrictions on the distribution shifts, or otherwise OOD generalization impossible. Furthermore, even with appropriate restrictions 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 approach that incorporates both of these principles and demonstrate its effectiveness in several experiments.
Introduction
Recent years have witnessed an explosion of examples showing deep learning models are prone to exploiting shorteuts (spurious features) (Geirhos et al., 2020 Pezeshki al.) (2200) which make them fail to generalize ditoo--distribution (OOD). In Beery et al. (2018), convolutional neural network was trained to classify camels from cows; however, it found that the model relied 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 AArjovsky 2019) and several other works were proposed (Ahuja et al. 2020; Pezeshki et al. 2020: Krueger
et al. 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 these works. The principle distinguishes predictors that
only rely the causes of the label from those that do not. The optimal predictor that only focuses
the causes is invariant and min-max optimal (Rojas- Carulla et al. 2018 Koyama and Yamaguchi
2020: Ahuja al. 2021 under many distribution shifts but the same is not true for other predictors.
"Equal contribution.
$ Mila Quebec Institute, Université de Montréal. Correspondence to: kartik mila.quebec.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Our contributions. Despite the promising theory, invariance principle- approaches fail in settings (Aubin et al., (02T) 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-bar 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 assumptions" 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 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 summarizes our theoretical results presented later). We propose approach that combines both these principles and demonstrate its effectiveness on linear unit tests (Aubin et al. (021) and on different real datasets.
approaches were
Task
Invariant fentures capture label info
Support overlap invariant features
Support overlap spurious features
OOD generalization guarantee (E,, - Estr)
ERM
IRM
IB-ERM
IB-IRM
Full/Partial
No
Yes/No
Impossible for algorichm generalize OOD [Thm2]
Linear
Classification
Full
Yes
No
x
[Thm3,4]
Partial
Yes
No
X
x
[Appendix]
Full
Yes
Yes
7
/
7 [Thm3,4]
Partial
Yes
Yes
x
X
Lineur
Full
No
No
Regression
Partial
No
No
x
[Thm4]
Table I: Summary of the new and existing results (Arjovsky et al.) 2019 Rosenfeld et al. 202
IB-ERM (IRM): information bottleneck empirical (invariant) risk minimization ERM (IRM).
2 OOD generalization and invariance: background & failures
Background. We consider supervised training data D gathered from set of training environments E D = DD JeEEL.. where D = {xf,, vf y is the dataset from environment E E,, and n is the number of instances in environment e. x² E Rd and yf E C & correspond to the input feature value and the label for 7th instance respectively. Each (If yf) is an i.i.d. draw from IP, where IP' is the joint distribution of the input feature and the label in environment e. Let X 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 predictor
RRd + R$ that performs well across many unseen environments in Eall, where Eatl C Exr. Define the risk of f in environment e as R²()) = E [[e(((X)) Y+)] where for example& & can be 0-1 loss, logistic loss, square loss, (X", Y²) ~ IP, and the expectation E w.r.. Pe_ Formally stated, our goal is to use the data from training environments Esr find f Rdd -> y to minimize
min max (f). f eEEall
(1)
So far we did not state any restrictions on Eatt. Consider binary classification: without any restrictions on Eall, no method can reduce the above objective (E is 0-1 loss) to below one. Suppose a method outputs f"; if 3 E Eall Err with labels based on 1-f-. then it achieves an error of one. Some assumptions on E.nt are thus necessary. Consider how Eaut is restricted using invariance for linear regressions (Arjovsky et al. (019)).
Assumption 1. Linear regression structural equation model (SEM). In each e E Eall
Y + -w" Vin t Zin Le,, Ele"] 0, E[e < o2
(2)
x + S(Z Kisp))
where Winy E RTT Zi E R", Zspu E R", S (dx(()))) S is invertible (m + o == d). We focus on invertible S but several results extend to non- inniinvertible S as well (see Appendix).
'Support 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

---

Assumption| states how Y and X are generated from latent invariant features Z 4 latent spurious features Zpu and noise c". The relationship between label and invariant features is invariant. i.e., Win is fixed across all environments. However, the distributions of 7 spu: and e are allowed to change arbitrarily across all the environments. Suppose S is identity. If we regress only on the invariant features Zinv. then the optimal solution is w which is independent of the environment, and the error it achieves is bounded above by the variance of E (d"up). If we regress on the entire Z and the optimal predictor places non-zero weight on Zypu (e.g., Zip + Y + ( then this predictor fails to solve equation ①1 (3 e E EaR, Zypu + 20, error + 0o, see Appendix for details). Also, not only regressing Zine is better than on 2", it can be shown that it is optimal, i..., it solves equation TT) under Assumption and achieves a value of Guup for the objective in equation (..
Invariant predictor: Define a linear representation map t R (that transforms X as ((x))
and define a linear classifier w. : RR*x (that operates on the representation w (xxx)). We want to search for representations & such that E|Y"| (X+) is invariant (in Assumption [i +( Z
then [[ is invariant). We say that a data representation @ elicits an invariant predictor w d across the set of training environments Eir if there is predictor that simultancously achieves the minimum risk, i.c., w² E arg min; R (w- 1), Ve E Egg. The main objective of IRM is stated as
min SUIRRK,,@RRPEE
(w - b) s.t. w E arg min eEEer
Σ
R R" D), Ve E Etr. WeRker
|Etrl
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 )' = R. For binary classification, the output of the predictor is given as ((w 4( ())),, where 1(-) is the indicator function that takes 1 if the input is 2 0 and 0 otherwise, and the risk is R (w 4) = [[[[ - -b(X'))x))))))))))))
Y= For regression, the output of the predictor is w (X*)) and the corresponding risk is 125 (w +) E (w x(X*) -Y We now present the main OOD generalization result from (Arjvssyy et al, (2019) for linear regressions.
Theorem 1. (Informal) If Assumption| 7 is satisfied, Rank( ( > O. 11erl > 2d, and Eer lie in a linear general position (a mild condition the data in Exe, defined in the Appendix), then each solution to equation achieves OOD generalization (solves equation L.. A e E Eall with risk > (up))
Despite the above guarantees, IRM has been shown fail in several cases including linear SEMs in (Aubin et al., 202T). We take a closer look at these failures next.
Understanding the failures: fully informative invariant features vs. partially informative in- variant features (FIIF vs. PIIF). We define properties salient to the datasets/SEMSs used in the OOD
generalization literature. Each e E Eau, the distribution (Xx, Y*) Fe satisfies the following proper- ties. a) 3 a map x" (linear or not), which we call an invariant feature map, such that E | g- (X*)] is the same for all e E Eall and Y L x*(x") These conditions ensure & maps to features that have finite predictive power and have the same optimal predictor across Eatl. For the SEM in Assumption I 4- maps to - b) = map di* (linear or not), which we call spurious feature map, such that E ye N (X+)] is not the same for all E Ealt and Y L (x*(x") for some environments. 4 often creates hindrance in learning predictors that only rely on &- Note that 4 should not be a transformation of some x², For the SEM in Assumption I suppose Zipu is anti-causally related to Y then & maps to Zepu (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 T" extracts the uncolored digit and the spurious feature map 3 extracts the background color. Ahuua et al. (202I) 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 causal-CMNIST (AC-CMNIST. original data from Arjovsky et al. (2019) but no different from ERM and fails in covariate shift-CMNIST (CS- CMNIST). In AC- CMNIST, the invariant features 4' (X*) (uncolored digit) are partially informative about the label, i.c.. Y y X 4 (X'). and color contains information about label not contained
In many examples in the literature, invariant features are causal, but not always (Rosenfeld et al. 2021
3

---

Fully informative invariant features (FIIF)
Partially informative invariant features (PITF)
Ve E Eatt, Y 1 _X (x(((()))
e & Eatt Y L __x (X))
Task: classification
Task: classification or regression
Example 2/2S, CS-CMNIST
Example 1/1S, Example 3/3S, AC-CMNIST
SEM in Assumption 2
SEM in Rosenfeld et al. 2021
ERM and IRM fail
ERM fails, IRM succeeds sometimes
Theorem 3,4 (This paper)
Theorem 9, 5.1 AAroossyy et al. 2019 Rosenfeld et al. 2021
Table 2: Categorization of OOD evaluation datasets and SEMs. Example I/1S, 2/2S, 3/3S from (Aubin et al. 2021). AC-CMNIST (Arjovsyy et al. 2019), CS-CMNIST Ahuja et al. 2021).
in the uncolored digit. On the other hand in CS- CS-CMNIST, invariant features are fully informative about the label, i.e., Y I X Ø* (X"), i.e. they contains all the information about the label that is contained in input X 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. In the rare case, when the humans are asked to label images in which the object being labelled itself is blurred, humans can rely 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 al., 202T): ) this is consistent with the linear regression result in Theorem I where IRM succeeds regardless of whether Y T X Z holds or not. Motivated by this observation, we take closer look at the classification tasks where invariant features are fully informative.
3 OOD generalization theory for linear classification tasks
A two-dimensional example with fully informative invariant features. We start with 2D classi- fication example (based on Nagarajan al. (202T). I. which can be understood as a simplified version of the CS CMNIST dataset (Ahuja et al. (021). Example 2/2S of Aubin et al. (2021). where both IRM and ERM fail. The example goes as tollows. In each training environment E Etr
ye 1( (Xim i - 1) where X {0, is Bernoulli
(4)
1 2
Xe + X + W where W E (0, 1} is Bernoulli (1 p¹) with selection bias p"
where Bernoulli(a) takes value with probability and otherwise. Each training environment characterized by the probability yf. Following Assumption T we assume that the labelling function does not change from E,, to Eau, thus the relation between the label and the invariant features does not change. Assume that the distribution of X and X can change arbitrarily. See Figure Ih) for a pictorial representation of this example illustrating the gist of the problem: there are many classifiers with the same error E,r while only the identical to the labelling function ((X 2)
generalizes correctly OOD. Define classifier ((iinv-Piin + spp - 1/2 (Winy + Wspu)). Define set of classifiers S = [Winn: ) s.t. Winv > |Uspul]. Observe that all the classifiers in S achieve zero classification eиror on the training environments. However, only classifiers for which Wspu 0 solve the OOD generalization (eq. (T). 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 l 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 [ (Arjvvsyy 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. (D). [ 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 in our construction.
We use the
'The deterministic labelling case was referred as realizable problems in (Arjovsky 2019
'Please note that this example illustrates certain important facets in very simple fashion: only in this example max- mar-mn classifier can solve the problem but not in general. (Further explanation in the Appendix).

---

T Prothow hortwwhettt
Lobel
Wi
5 Negytivwhet suppors
tnaning support
Spurious feature
Invaiaa predidor
Label
M
00.
9.305t Faluus
Color intensity
(00-5)
11,00
Labelling hyperplane Invoriant
Figure 1: a) 2D classification example illustrating multiple invariant predictors: Most of these predictors rely on spurious features and each of them achieve zero error across all Etr. b) illustration of the impossibility result. If latent invariant features in the training environments are separable, then there are multiple equally good candidates that could have generated the data, and the algorithm cannot distinguish between these.
Assumption 2. Linear classification structural equation model (FIIF). In each E Eall
1 No (Z Zepu);
Y +- 1(wimu Zinu)
+ No, X t- ((Ziv: Zspu),
N ~ "Brnoullll(g))
2'
where winh E Th with |/ = is the labelling hyperplane, Z E IR". 2e E RR, No is binary noise with identical distribution across environments, @ the XOR operator, invertible.
If noise level is zero, then the above SEM covers linearly separable problems. See Figure 2 for the directed acyclic graph (DAG) corresponding to this SEM. From the DAG observe that Y IX Z which implies that the invariant features are fully informative. Contrast this with DAG that follows Assumption shown in Figure 2p). where Y Y X=\Ze iny and thus the invariant features are not fully informative. If Ealt follows the SEM in Assumption 2nd suppose the distribution of Z Ze can change arbitrarily, then it can be shown that only a classifier identical to the labelling function (winy ) can solve the OOD generalization (eq. (D); such classifier achieves error of q (noise level) in all the environments. As a result, if for a classifier we can find E Eall that follows Assumption Senriiion@hhrrrr where the error greater than q, then such classifier does solve equation (). Now we ask what are the minimal conditions on training environments &tr to achieve OOD generalization when Euul follow Assumption 2 To achieve OOD generalization for linear regressions, in Theorem 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 work with similar assumptions for the SEM Assumption 2 and solve the OOD generalization (eq. (I)? We state some assumptions and notations to answer that. Define the support of the invariant (spurious) features Zinv (Zegu) in environment e as Zi (2
Assumption 3. Bounded invariant features. UeEE., Zí,v is bounded set Assumption 4. Bounded spurious features. UsEE,, Zspu is bounded set.
Assumption 5. Invariant feature support overlap. Ve E Ealls Zin Seee,, 2
Assumption 6. Spurious feature support overlap. Ve E Eall; 2ep S UU'EE,, 2
Assumption 5 ( 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 important 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 Zin, separated by the labelling hyperplane - Deñine Inv-Margin min.eu.ee,, zy. sgn (winv 2) (w =). This margin only coincides with the standard margin in support vector machines when the noise level q is 0 (linearly separable) and identity. Inv-Margin > O, then the labelling hyperplane wi separates the support into two halves (see Figurepp)
A set Z is bounded if 3M such that V2 E Z, ||||| $ M.
5

---

Assumption 7. Strictly separable invariant features. Inv-Margin > O.
Next, we show the importance of support overlap for invariant features.
Theorem 2. Impossibillty of guaranteed OOD generalization for linear classification. Suppose each e E Eaul follows Assumption 2 for all the training environments Etr. the latent invariant features are bounded and strictly separable, i.e., Assumption and algorithm fails to solve the OOD generalization (eq. (T), i.e, for the output of every algorithm B e E Eall which the error exceeds the minimum required value q (noise level).
then every deterministic
The proofs to all the theorems are in the Appendix. We provide high-level intuiton as why invariant feature support overlap is crucial to the impossibilityy result. In Figure we show that the support of latent invariant features are strictly separated by the labelling hyperplane winw then can find another valid hyperplane wit that is equally likely to have generated the same data. There is no algorithm that can distinguish between winv and wi As result, if we use data from the region where the hyperplanes disagree (yellow region Figure I)); then the algorithm fails.
Significance of Theorem 2 We showed that without the support overlap assumption the invariant features, OOD generalization is impossible for linear classífication tasks. This is in contrast to linear regression in Theorem I ( (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 generaliza- tion (eq. (T). Each classifier can be written as i X w S i 7spp)) Winv 7 + ". Z If Vupu * 0, then the classifier t 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 7
Sufficiency: If the spurious features satisfy support overlap (Assumption Ianiiionnooolee holds), then both ERM and IRM solve the OOD generalization problem (eq. (T). 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. (.). Also, there exist no such classifiers that rely on featuress
features and also achieve OOD generalization.
Significance of Theorem 3 From the first part, 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. (T)) Interestingly, in this case we can have classifiers that rely on the spurious features and yet solve the OOD generalization (eq. (.). For the 2D case (eq.
this case implies that the entire set S solves the OOD generalization (eq. (.) 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 Xí,v in the set S solve the OOD generalization (eq. (I).
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 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.
4 Information bottleneck principle meets invariance principle
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 T(X; +(X)) is used to measure information compression. If representation ((X) a deterministic transformation of X,
then in principle we can use the entropy of (xX to measure compression (Kirsch et al.) 2020 Let

---

Y
Y
Y
Zin
7:spu
7is
Zin
7
X
X
X
(a) FIIF (this work)
(b) PIF (Arjovsky ct al. 2019)
(c) PIF (Rosenteld ct al.
2021
Figure 2: Comparison of the DAG from Assumption
2 (fully informative invariant features) vs. DAGs
from Rosenfeld al. (2021) Arjovsky al. (2019) (partially informative invariant features).
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 (xx))) when & is the identity is H (yf) + log(2), where H ((r)) is the Shannon entropy in Bernoulli(yf) If selects the invariant/spurious features only, then H ((x))))) = log(2). Among all three choices, the that has the least entropy and also achieves zero error is the representation that focuses 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 t needed? We answer this next.
Why invariance? Consider simple classification SEM. In each e E Eir, Y t- X @b X @ Ne and X + Y @ V*, where all the random variables involved are binary valued. noise No, V are Bernuuli with parameters y (identical across &rr), c (varies across Exr) respectively, If of q. then in &tr predictions based on spu are better than predictions based on X X1Le x22. inv If both X/,e, X22 are uniform Bernoulli, then these features have a higher entropy than Xe In this case,
"ppu"
the information bottleneck would bar using X inv Instead, we want the model to focus on X x2 and not on X Invariance constraints encourage the model focus on X X2,5 example, observe that invariant features are partially informative unlike the 2D case (eq. (4)).
iny. In this
Why invariance and information bottleneck? We have illustrated through simple examples when the information bottleneck is needed but not invariance and vire-versa. We now provide a simple example where both these constraints are needed the same time. This example combines the 2D case (eq. (4)) and the example we highlighted in the paragraph above: Ye < X @ Ne, X1. i + W and X2, <- Y V In this case, the invariance constraint does not allow spu
representations that use X2,e but does not prohibit representations that rely on X1.e. However, spu information bottleneck constraints on top ensure that representations that only use X are used. We now describe an objective that combines both these principles:
spu
min Σ he (w ②)
s.t. |Eerl
Σ REEir
Rs (w ①) S rht E arg min R D), Ve E Eir, (6)
WEREXY
REE...
where he in the above is a lower bounded differential entropy defined below and th 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 certain hidden layer. For both analytical convenience and also because the above setup 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 variable X, (x)) = Ex [log dl x] and dlP x is the Nikooym derivative of Ixx 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 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
Results extend altemate objective with information bottleneck constraints and average risk objective.
7

---

Fully informative invariant features (FIIF), Suppose each e E Eall follows Assumption Assume that the invariant features are strictly separable, bounded, and satisfy support overlap (Assumptions 3 and 7 hold). Also, for each E Etr 7pp + AZ + W where A E Roxm We E Re is continuous, bounded, and zero mean noise. Each solution to IB-IRM (eq. O with l as 0-1 loss, and T 1). and IB-E B--EMM solves the OOD generalization (eq. (1) but ERM and IRM (eq. (e4.33). fail.
Partially informative invariant features (PIIF). Suppose each e Eall follows Assumption lnd B e E Ete such that EEle" spul + 0. If ( > 2d and the set Exr 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, o² < <rth I oy, where oy and o² 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 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 3J This establishes that support overlap of spurious features is not 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, 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 result in Rosenfeld et al. (202I), which focused on linear classification SEMs (DAG shown in Figure ] with partially informative invariant features. Under the assumption of complete support overlap for spurious and invariant features, authors showed IRM succeeds.
4.1 Proposed approach
We take the three terms from the optimization in equation (6) and create a weighted combination as w. w- w. (w-D) (w+h))(+x())))))
Σ R VV w. (w.+)]12+vh+
Σ6 [ (
In the LHS above, the first term corresponds to the risks across environments, the second term approximates invariance constraint (follows the IRMv1 objective (Arjovsky jet al., 2019). ) and the third term is the entropy of the classifier in each environment. In the RHS, h(() is the entropy of & unconditional on the environ- ment (the entropy on the left-hand side is entropy conditional on the environment assuming all the environments are equally likely).
Jotgrll V to
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
I5-ERM
IBM
IRM
lhe
Number of descent steps
Figure 3: Comparing conver-
gence of /wZ. (metric from
|Wagul
(1= [(ee + AIIV
m, 1" (" +)1l + -Var(()):
Nagarajan et al. 2021 for aver-
age selection bias p p-=00 0.9.
On the behavior of gradient descent with and without informa- tion 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 (cq. (I). 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 ct al. 202T). Even though gradient descent can eventually learn the ideal classifier that only relies on the invariant features, training is frustratingly slow was shown by Nagarajan et al. (202I). In the next theorem, we characterize the impact of using IB penalty (Var( (b) in the 2D example (eq.
4 We compare the methods in terms of (t) which was the metric used in Nagarajan et al.
(t)
2021 Wssp(() (1) and (!) are the weights for the spurious feature and the invariant feature at time
of training (assuming training happens with continuous time gradient descent).
8

---

Theorem 5. Impact of IB on learning speed. Suppose each e E Eir follows the 2D case from equation Set À 0. Y > equation 7 to get the IB-ERM objective with & as exponential loss. Mapu(C) (1) I c in time less than
Continuous-time gradient descent on this IB-ERM objective achieves
W
2[I-DE (Wo(.) denotes the principal branch of the Lambert W function), while in the same time the
ratio for ERM (agu(()
2 In
3-22 (12p)2))
Wo( (1 2(1-p)e) where [Eer] Exces, P..
Winu
converges to zero for both methods, but it converges much faster for IB-ERM (for p (0 S 0.001 and ratio for ERM is 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. See Figure 3 for the comparisons.
(ap Win (
0.9,6 0.001, = 0.58, the ratio for IB-ERM is weul()
(2)
(e)
5 Experiments
Methods, datasets & metrics. We compare our approaches information bottleneck based ERM (IB-
ERM) 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 Terra Incognita dataset (Beery et al. 2018), and COCO (Ahmed)
al.] 202T). We follow the same protocol for tuning hyperparameters from Aubin et al. (2021): :
Arjovsky et al. (019) 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,
DS-CMNIST, Terra Incognita, and COCO we use accuracy. For Example I/1S, we use mean square error (MSE). The code for experiments can be found at tattps://giihhae cww.au/aaooomaaaaaaa.aaaa ahaaaaaujaaaaaaa IRMIRM
com/
Summary of results. In Table 3 we provide a comparison of methods for different examples in linear unit tests (Aubin al., 202 1) 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 I/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- 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 g we do not expect IRM and ERM to do well on Example 2/2S CS-CMNIST, Terra Incognita and COCO dataset are fully informative. On these FIFF 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 FIF 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 find that IB -R-MMM is able to extract the benefit of IB penalty. In settings such as COCO dataset, where IB -RM does not perform as well as IB BRMM better hyperparameter tuning strategies should be able to help B--M adapt and put a higher weight IB penalty. Overall, 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).
Theorem
g as these datasets fall in the FIIF category, i.e., the invariant features
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. (021 ))
6 Extensions, limitations, and future work
Extension to non-linear models and multi-class classification. In this work our theoretical analysis focused on linear models. Consider the map X <- S( 7spu) 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 linear model that only relies on the invariant features Zinv to predict the label Y. For
We place Terra Incognita and COCO dataset in the FIF assuming that the humans who labeled the images did not need to rely on unceliable/spurious features such as background to generate the labels.

---

#Envs
ERM
IB-ERM
IRM
IB-IRM
Oracle
Examplel
13.36 + 1.49
12.96 + 1.30
11,15
11.68 0.90
10.42:10.16
Examplels
13.33 1.49
12.92 + 1.30
11.07 0.68
+1.74 1.03
10.45+0.19
Example2
3
0.42 + 0.01
0.00
0.45
0.00 + 0.00
+ +.00
Example2s
3
0.45 1 0.01
0.00
0.451.000
0.06 1: 0.12
0.00 0.00
Example3
3
0.48 1 0.07
0.06
1 0.48
0.48 0.07
0.01 1 0.00
Example3s
0.49 1 0.06
0.49
0.49 1 0.07
0.49 0.07
0.01 1 0.00
Example1
6
33.74 + 60.18
32.03 + 57.05
23,04 = 40.64
25.66 + 45,96
22.21
Examplels
33.62 + 59.80
31.92 + 56.70
22.92 40.60
25.60 45.62
22.13+38..9
Example2
0.37 + 0.06
0.02 0.05
0.46 + 0.01
0.43+0.11
0.00+0.00
Example2s
0.01 0.46 + 0.01
0.02 0.06
0.46 0.01
0.45 + 0.10
0.10
0.00++000
Example3
0.33 + 0.18
0.26 + 0.20
0.14 + 0.18
0.19 + 0.19
0.01+0.00
Example3s
0.36 + 0.19
0.27 + 0.20
0.14 + 0.18
0.19 + 0.19
0.01
Table 3: Comparisons on linear unit tests in terms of mean square error (regression) and classification error (classification). "#Envs" means the number of training environments.
ERM
IB-ERM
IRM
IB-IRM
CS-CMNIST
60.27 + 1.21
71.80 = 0.69
61.49 1.45
71.79 + 0.70
AC-CMNIST
16.84 + 0.82
50.24 = 0.47
66.98 + 1.65
67.67 + 1.78
Terra Incognita
49.80 + 4.40
56.40 = 2.10
54.60
1.30
54.10 + 2.00
COCO
22.70 + 1.04
31.66 2.39
18.47 = 10.20
25.10 + 1.03
Table 4: Classification accuracy percentage on colored MNISTs, Terra Incognita and COCO dataset.
rely the approaches proposed in this work. For part a), need to leverage advancements in the field of non-linear ICA (Khemakhem et al. 2020). The current state-of-the vate-of-thheart to solve part a) requires strong structural assumptions the dependence between all the components of Zinu: Zepu (Lu et al. 202T). Therefore, solving part a) and part b) in conjunction with minimal assumptions forms 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.
part b),
On the choice for IB penalty and IRMv1 penalty. We use the approximation for entropy (in equation 77) described in Kirsch et al. (2020). The approximation (even though upper bound) serves as an effective proxy for the true information bottleneck as shown in the experiments in Kirsch et al. (2020) (e.g., their experiment approximation even in moderately high dimensions, as an example in CS-CMNIST, the dimension of the layer at which bottleneck constraints 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., [021] Kamath et al. 2021 Gulrajani and Lopez-Paz 2021), there has been criticism of different aspects of IRM, e.g., failure of IRMvT penalty in non--onneen linear models, the tuning of IRMv1 penalty, etc. Since we use IRMv1 1 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 2020: Chang et al.) 2020). Exploring their benefits together with information bottleneck is a fruittul future work. Before concluding, we want 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.
Imagenette dataset). Also, our experiments validate this
7 Conclusion
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 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 form of the information bottleneck constraint along with invariance goes a long way in overcoming the failures while retaining the cxisting provable guarantees.
10

---

Acknowledgements
We thank Reyhane Askari Hemmat, Adam Ibrahim, Alexia Selicoeur-Martinenn Divyat Mahajan. Ryan D'Orazio, Nicolas Loizou, Manuela Girotti, and Charles Guille Saclure for the feedback. Kartik Ahuja would also like thank Karthikeyan Shanmugam for discussions pertaining 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-201) 2011)))))- 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.
References
Ahmed, F., Bengio, Y., van Seijen, H., and Courville, A. (2021). Systematic generalisation with group invariant predictions. In International Conference Learning Representations.
Ahuja, K., Shanmugam, K., Varshney, K., and Dhurandhar, A. (2020). Invariant risk minimization games. In International Conference on Machine Learning, pages 145-155. PMLR.
Ahuja, K., Wang, J., Dhurandhar, A., Shanmugam, K., and Varshney, K. R. (2021). Empirical or invariant risk minimization? sample complexity perspective. In International Conference on Learning Representations.
Arjovsky, M., Bottou, L.., Gulrajani, I., and Lopez-Paz, D. (2019). Invariant risk minimization. arXiv preprint arXiv: 1wX.::99070009000 002993...2....
Aubin, B., Stowik, A., Arjovsky, M., Bottou, I.,, and Lopez-Paz, (2021). Linear unit-tests for invariance discovery. arXiv preprint arXiv:2102.10867
Beery, S.. Van Horn, G., and Perona, P. (2018). Recognition in terra incognita. In Proceedings of the European Conference on Computer Vision, pages 456-473.
Chang, S., Zhang, Y., Yu, M., and Jaakkola, T. S. (2020). Invariant rationalization. In International Conference on Machine Learning, 2020.
DeGrave, A. Janizek, D., and Lee, S.-I. (2020). Al for radiographic COVID-19 detection selects shortcuts over signal. medRxiv.
Geirhos, R., Jacobsen, J.-H., Michaelis, C., Zemel, R., Brendel, W., Bethge, M., and Wichmann, F. A. (2020). Shortcut learning deep neural networks. Nature Machine Intelligence, 2(11):665-673.
Gulrajani, and Lopez-Paz, D. (2021). In search of lost domain generalization. In International Conference Learning Representations.
Kamath, P., Tangella, A., Sutherland, J., and Srebro, N. (2021). Does invariant risk minimization capture invariance? arXiv preprint 21X11:2101.01134.
Khemakhem, I., Kingma, D., Monti, R., and Hyvarinen. A. (2020). Variational autoencoders and nonlinear ica: A unifying framework. In International Conference on Artificial Intelligence and Statistics, pages 2207-2217. PMLR.
Kirsch, A., Lyle, C., and Gal, Y. (2020). Unpacking information bottlenecks: Unifying information- theoretic objectives in deep learning. arXiv preprint arxiv:2003.12537
Koyama, M. and Yamaguchi, (2020). Out-of-distributioo generalization with maximal invariant predictor. arXiv preprint 2rXiv:2008.01888

---

Krueger, D., Caballero, E., Jacobsen, J.-H., Zhang, A., Binas. J., Zhang, D., Priol, R. L., and Courville, A. (2020). cat-of-diitribuution generalization via risk extrapolation (rex). arXiv preprint arXiv:2003.00688
Lu, C., Wu, Y., Hernández-Lobato, J. M., and Schölkopf, B. (2021). Nonlinear invariant risk minimization: A causal approach. arXiv preprint arXiv:2102.1233:
Nagarajan, V., Andreassen, A., and Neyshabur, B. (2021). Understanding the failure modes of out-of-distribution generalization. In International Conference on Learning Representations.
Pearl, (1995). Causal diagrams for empirical research. Biometrika, 82(4):669
Peters, J., Bühlmann, P., and Meinshausen, N. (2015). Causal inference using invariant prediction: identification and confidence intervals. arXiv preprint urXiv:/50.0013:
Pezeshki, M., Kaba, S.-O., Bengio, Y., Courville, A., Precup, D., and Lajoie, G. (2020). Gradient starvation: A learning proclivity in neural networks. arXiv preprint arXiv:2011.09466
Robey, A., Pappas, G. J., and Hassani, H. (2021). Model-based domain generalization. arXiv preprint arXiv:2102.1145 14336.
Rojas-Carulla, M., Schölkopf, B., Turner, R., and Peters, J. (2018). Invariant models for causal transfer learning. The Journal Machine Learning Research, 19(1):1309-133 1342.
Rosenfeld, Ravikumar, K., and Risteski, A. (2021). The risks of invariant risk minimization. In International Conference Learning Representations.
Soudry, D., Hoffer, E., Nacson, M. S., Gunasekar, S., and Srebro, N. (2018). The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research, 19(1):8822-28
Tishby, N., Pereira, F. C., and Bialek, W. (2000). The information bottleneck method. arXiv preprint physics/0000077.
Vapnik, V. (1992). Principles of risk minimization for learning theory. In Advances neural information processing systems, pages 8-1-838.
Zhang, D., Ahuja, K., Xu, Y., Wang, Y., and Courville, A. C. (2021). Can subnetwork structure be the key to out-of-distribution generalization? In ICML.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes] See Section 2-5 and the additional details such the proofs in the supplementary material.
(b) Did you describe the limitations of your work? [Yes] See Section 4.1 and Section 6.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See Section A. in the Appendix in the supplementary material.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results....
(a) Did you state the full set of assumptions of all theoretical results? [Yes] See Section
2-4.
(b) Did you include complete proofs of all theoretical results? [Yes] See the Appendix in the Supplementary Material.
3. If you experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main ex- perimental results (either in the supplemental material or as URL.)? [Yes] See htttp://gitte com wom/aaujak/IIB IRM
12

---

(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Section A.2 in the Appendix in the supplementary material.
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [Yes] See Section A.2 in the Appendix in the supplementary material.
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? [Yes] See Section A.2 in the Appendix in the supplementary material.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yss] We
use the codes from following github repositories htttps://githut com/ facebookreeearrch//Dmaaaaaeeddd https://gittub. comaa//cbbccbbbbbessssrcc InvariandRiissiimmiizzzet and https: //t::///ithbb com/faaceoookkeseardd
InvarianceUnii77estt and we have cited the creators in the Section A.2 in the
Appendix in the supplementary material.
(b) Did you mention the license of the assets? [Yes] All the repositories mentioned above use MIT license. We have mentioned this in Section A.2 in the Appendix in the supplementary material.
(c) Did you include any new assets either the supplemental material or as URL? [Yes] We have included code for our experiments in the supplementary material.
(d) Did you discuss whether and how consent obtained from people whose data you're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "hMOWFBNKfC", "reviewer": "Reviewer_s4BJ", "summary": "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. \n\nThe 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.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper is well organized and easy to follow. The theoretical results are presented in a logical order and the conclusions are quite clear. Although all the proofs are in the appendix, authors provide the intuitions by examples and discussed the significance of the theorems. \nWith these theorems, the proposed IB-ERM is well-motivated. Some relaxations are involved from the original formulation in (6) to the final objective in (7). The influence of the IB term to optimization is discussed in Theorem 5. \n\nOverall, this paper is well-written, with comprehensive theoretical analysis to support the proposed IB-ERM for linear classification. \nDetailed comments/questions:\n1.\tAlthough we can construct tasks like CS-CMNIST and the model in assumption 2, is it true that, with high-probability, the oracle in real world is PIIF? Since the conditional independency is a special case of the conditional dependency. Then, what is the advantage of IB-IRM, comparing to other methods designed for PIIF, e.g., (Rosenfeld et al. 2021)?\n2.\tThe proposed IB-IRM in eqn (6) takes IB as the objective and ERM and invariance as constraints. Will theorem 4 still holds if one take IB as an additional constraint in IRM? Can we have more insights to design the relaxed objective later?\n3.\tHow did you tune the coefficient of the two penalty terms at the equation (7) in real experiments? Any ablation study to show the influence of these two penalties in a different situation?\n4.\tIn sec 5, authors make discussion about the results on CMNIST tasks. How about the experimental results on linear unit tests? Why only make comparison on three and six training environments? Do we have any prior knowledge about which tasks are PIIF and which tasks are FIIP? IB-IRM is not the best across all the tasks.\n5.\tMinors: 1) put the DAG into Figure 2? 2) ref the definition of “support overlap invariant/spurious features” in table 1; 3) Is the observation in figure 3 general for different values of \\lambda and \\gamma? 4) In Theorem 3, “some of the ERM and IRM solution..” is not rigorous. “these exist”?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See the detailed questions in the main review, especially question 3 and 4. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "E7ONXjXU3F", "reviewer": "Reviewer_kbVa", "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.  ", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "In general, this work made a few theoretical contributions. It specifies some conditions to guarantee good generalization of IRM in classification, especially the support overlap of invariant (or spurious) features. Some theorems (like Theorem 4) also partially explains why information bottleneck can improve generalization, which is less investigated before. \n\nHowever, I still have some concerns:\n1. Authors combine the objective of IRM and the minimizing of entropy of representation $H(\\Phi(X))$. However, instead of directly optimizing the differential entropy (which is usually hard), authors made a Gaussian assumption and alternatively minimize the variance. I am not so sure how practical is this assumption in real-world data or applications. Because the experiments in this work is just on simple synthetic data or MNIST-like data. \n\nCan you provide more justifications on this approximation or more empricial results on recent domain generalization benchmarks (i.e., the DomainBed [1] or WILDS [2]) in which both IRM and ERM are baselines. \n\n[1] Gulrajani, Ishaan, and David Lopez-Paz. \"In search of lost domain generalization.\" arXiv preprint arXiv:2007.01434 (2020).\n\n[2] Koh, Pang Wei, et al. \"Wilds: A benchmark of in-the-wild distribution shifts.\" International Conference on Machine Learning. PMLR, 2021.\n\n2. Similar to the information bottleneck, I feel the hyperparameter $\\gamma$ (in Eq. (7)) controls a trade-off and plays a significant role to the performance. Here, there is another hyperprameter $\\lambda$ that controls the degree of invariance. Can you give a more indepth analysis on the trade-offs between these terms or how different values of $\\gamma$ or $\\lambda$ influence the performance of the combined objective?\n\n3. I still feel there is a jump from your analysis in Section 3 to the introduction of information bottleneck regularization (i.e., $\\min H(\\Phi(X))$) in Section 4. Just focusing on your illustrative example, when $\\Phi$ only picks invariant features, the entropy is minimum. However, it does not mean when the entropy is minimized, $\\Phi$ exactly picks the invariant features. For example, the suprious features only also gives the minimum entropy. My question is how can you guarantee this regularization alone is sufficient for IB-IRM to only focus on invariant features, rather than other combinations of invariant and suprious features? or did I miss something?\n\nFinally, some minor suggestions:\n1. It would be much better if authors can give some illustrative examples on \"fully\" informative features and \"partially\" informative features. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Authors clearly mentioned limitations and future work directions. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "A8A8pb5bdJD", "reviewer": "Reviewer_bprw", "summary": "In this paper, the authors re-check the invariance principle and show that out-of-distribution generalization is much harder for linear classification tasks, requiring much stronger restrictions in the form of support overlap assumptions on the distribution shifts. They propose an approach combining both information bottleneck constraints and the invariance principle to addressing the OOD generalization in both regression and classification tasks. ", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I pretty like the theoretical analysis presented in this paper, from which we could get very insightful thoughts. \n\nIn this paper, the authors use 0-1 loss for binary classification $\\mathcal{Y}=\\{0, 1\\}$ and square loss for regression $\\mathcal{Y}=\\mathbb{R}$. I am curious if all the results in this paper can be extended to the other losses and to multi-class classification tasks. \n\nAlthough the authors admit in the final section that the proposed method mainly focuses on linear models. I would like to hear more from the authors about what the main challenge is when their theory is extended to nonlinear models.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "two hours"}, "extra_text": {}}, {"review_id": "7D9U1jRUB7I", "reviewer": "Reviewer_YQ8U", "summary": "- The authors first argue that IRM is insufficient to obtain good test performance (only using invariant features) on OOD data at test time when the OOD support is not in the ID support\n- Then they prove for a particular distributional assumption on the spurious features, that (additionally) limiting the differential entropy of the representation can solve this issue and lead to OOD generalizing predictors\n- In some toy experiments, they show that IB-IRM sometimes performs better", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- They argue that the information bottleneck in some situations can help to obtain the invariant predictor when invariance itself cannot (when there’s not enough support overlap)\n- In their (few) experiments on modified colored MNIST datasets, including one where IRM does not improve much compared to ERM (CS-MNIST) and demonstrate that adding the IB regularizer leads to better OOD generalization\n\nWeaknesses:\n\nTheory\n- The insufficiency results for IRM are not surprising, e.g. FIIF problem was already noted in Nagarajan et al\n- Comparison with other works that show limits of IRM such as Srebro - their examples themselves seem to make a similar point as FIIF?\n\nSince the theoretical contributions are insufficient, we now examine the methodological contributions \nMethodology:\n- major criticism: is this supposed to be a game-changer for practice? then it’s absolutely necessary to do OOD generalization experiments on real world datasets (such as WILDS). the current datasets are completely toy (versions of colored MNIST)\n- since variances are huge compared to MSE for examples 1/1s, 3/3s with 6 environments, one can really only compare the performances of the methods using CS-MNIST and AC-MNIST and perhaps example 2/2s which is rather unsatisfying as empirical evidence for a method\n- table 3 suggests that more environments worsens the benefits of adding IB even in an FIIF setting? more discussion/experiments on that would be necessary\n- again regarding the point that the variances are huge for Example 1/1s with 6 vs. 3 environments: a comment on that in the main text would be useful - is that expected and why?\n\nUpdate: Apologies, I actually was not able to understand the problem that I previously had when I had a brief new look at the paper (before even reading your response). I'm guessing I was thinking about some problems that confused my mindset when reading. Now the sufficiency result actually sounded trivial since in the case of spurious correlations AND support overlap all possible environments in E_all would also be spuriously correlated...\n\nThanks for discussion on relation to Kamath et al. \n\nI agree Nagarajan did not provide formal results, but hinted at the problem so this is why I said it is not surprising. \n\nThanks for comment on six-environments to discuss the high variance. That comment should be added to paper.\n\nAdding experiments on Coco addresses my main concern (I still think the methodology is the main contribution). However I would like to see more experimental details before giving the nod + experiments on one more dataset. The lack of real-world experimental evidence was the primal reason for the score and hence I still tend to reject the paper. The number of scenarios where IB-IRM or IB-ERM do better is still rather limited and currently only includes COCO in terms of non-toy dataset (without having been able to see the revision and experimental details) which is in my opinion still quite rather meager for publication at Neurips. How about WILDS? I believe if the authors run a more extensive experimental study this paper can be accepted at the next venue.\n\nI didn't see this point addressed: \"Table 3 suggests that more environments worsens the benefits of adding IB even in an FIIF setting? more discussion/experiments on that would be necessary\"", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Authors do have a section on limitations and motivate future work through that.\nThere's no potential negative societal impact for OOD detection per se ... (for certain applications it could, but this holds for almost all problems that one can think of solving).", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7"}, "extra_text": {}}]

Author rebuttal:

Meta-review: After substantial discussion, including additions of new experimental results from the authors, reviewers have come to a consensus that this paper provides interesting and relevant theoretical results, as well as sufficient empirical evidence supporting the theory, to merit acceptance. We expect that this work will appeal to the ERM, IRM, and Information Bottleneck communities. We particularly think that the exponential speedup of convergence of the IB-based methods will be of theoretical and practical interest. We look forward to reading the updated version of the paper with the additions and revisions the authors have promised.
Decision: Accept (Spotlight)