Continual Invariant Risk Minimization
Authors: Francesco Alesiani
Shujian Yu
Mathias Niepert
Venue:  (2021)
Keywords: Supervised Learning
Causal Learning
Invariant Risk Minimization
Continual Learning

Abstract: Empirical risk minimization can lead to poor generalization behaviour on unseen environments if the learned model does not capture invariant feature represen- tations. Invariant risk minimization (IRM) is a recent proposal for discovering environment-invariant representations. It was introduced by Arjovsky et al. (2019) and extended by Ahuja et al. (2020). The assumption of IRM is that all environ- ments are available to the learning system at the same time. With this work, we generalize the concept of IRM to scenarios where environments are observed se- quentially. We show that existing approaches, including those designed for contin- ual learning, fail to identify the invariant features and models across sequentially presented environments. We extend IRM under a variational Bayesian and bilevel framework, creating a general approach to continual invariant risk minimization. We also describe a strategy to solve the optimization problems using a variant of the alternating direction method of multiplier (ADMM). We show empirically us- ing multiple datasets and with multiple sequential environments that the proposed methods outperforms or is competitive with prior approaches.

Under review as a conference paper at ICLR 2021
CONTINUAL INVARIANT RISK MINIMIZATION
Anonymous authors Paper under double-blind review
ABSTRACT
Empirical risk minimization can lead to poor generalization behaviour on unseen environments if the learned model does not capture invariant feature represen- tations. Invariant risk minimization (IRM) is a recent proposal for discovering environment innoommmenttiiviiantt invariant representations. It was introduced by Arrovsky al. (2019) and extended by Ahuja et al.] (020). The assumption of IRM is that all environ- ments are available the learning system the same time. With this work, we generalize the concept of IRM to scenarios where environments are observed se- quentially. We show that existing approaches, including those designed for contin- ual learning, fail to identify the invariant features and models across sequentially presented environments. We extend IRM under a variational Bayesian and bilevel framework, creating a general approach to continual invariant risk minimization. We also describe strategy to solve the optimization problems using a variant of the alternating direction method of multiplier (ADMM). We show empirically us- ing multiple datasets and with multiple sequential environments that the proposed methods outperforms or is competitive with prior approaches.
INTRODUCTION
Empirical risk minimization (ERM) is the predominant principle for designing machine learning models. In numerous application domains, however, the test data distribution differ from the training data distribution. For instance, at test time, the same task might be observed in a differ- ent environment. Neural networks trained by minimizing ERM objectives over the training distri- bution tend to generalize poorly in these situations. Improving generalization of learning systems has become major research topic in recent years, with many different threads of research includ- ing, but not limited to, robust optimization (e.g., Hoffman et al. (2018)) and domain adaptation (e.g., Johansson al. )2019)) Both of these research directions, however, have their own intrin- sic limitations (Ahuja et al. (2020)). Recently, there have been proposals of approaches that learn anvironment-invarion representations. The motivating idea is that the behavior of a model being invariant across environments makes it more likely that the model has captured a causal relationship between features and prediction targets. This in turn should lead to a better generalization behavior. Invariant risk minimization (IRM, Arjovsky al. (019)), which pioneered this idea, introduces new optimization loss function to identify non-spurious causal feature-targe interactions. Invari- ant risk minimization games (IRMG, Ahuja et al. (2020)) expands on IRM from game-theoretic perspective.
The assumption of IRM and its extensions, however, is that all environments are available to the learning system at the same time, which is unrealistic in numerous applications. A learning agent experiences environments often sequentially and not concurrently. For instance, in a federated learn- ing scenario with patient medical records. each hospital's ' (environment) data might be used to train a shared machine learning model which receives the data from these environments in sequential manner. The model might then be applied to data from an additional hospital (environment) that was not available at training time. Unfortunately, both IRM and IRMG are incompatible with such a continual learning setup in which the learner receives training data from environments presented in a sequential manner. As already noted by Javed et al. (020)). "IRM Arjovsky et al. (2019) requires sampling data from multiple environments simultancously for computing regularization term per- tinent to its learning objective, where different environments are defined by intervening on one or more variables of the world." The same applies to IRMG (Ahuja et al. (2020)

---

Under review as a conference paper at ICLR 2021
To address the problem of learning environment invarriat ML models in sequential environements, we make the following contributions:
We expand both IRM and IRMG under Bayesian variational framework and develop novel objectives (for the discovery of invariant models) in two scenarios: (I) the standard multi- environment scenario where the learner receives training data from all environments at the same time; and (2) the scenario where data from each environment arrives a sequential manner.
We demonstrate that the resulting bilevel problem objectives have an alternative formulation, which allows us to compute a solution efficiently using the alternating direction method of multi- pliers (ADMM).
We compare our method to ERM. IRM, IRMG, and various continual learning methods (EWC, GEM, MER, VCL.) on diverse set of tasks, demonstrating comparable or superior performance in most situations.
2 BACKGROUND: OFFLINE INVARIANT RISK MINIMIZATION
We consider a multi-environment setting where, given a set of training environments E' {21, 62. Jm the goal is to find parameters that generalize well unseen (test) environments. Each environment has associated training data set De and corresponding risk R
Rº (w Φ) E( lo(ww ()),,
where fo = w o ø is the composition a feature extraction function @ and a classifier (or regression function) w. Empirical Risk Minimization (ERM) minimizes the average loss across all training examples, regardless of environment:
RERM E(eea)))UEEED {(fo(r).
ERM has strong theoretical foundations in the case of iid data (apnik 1992) but can fail dramat- ically when test environments differ significantly from training environments. To remove spurious features from the model, Invariant Risk Minimization (IRM, [Arjvvsky et al.) (2019) instead aims to capture invariant representations & such that the optimal classifier w given @ is the same across all training environments. This leads the following multiple bi-level optimization problem
min &E E
(w Φ) cEE
S.t. w E arg min 12 (w, Φ), Ve E E, W.EH
(3)
where Hş, Hw are the hypothesis sets for, respectively. feature extractors and classifiers. Unfortu- nately, solving the IRM bi-level programming problem directly is difficult since solving the outer problem requires solving multiple dependent minimization problems jointly. We can, however, relax IRM to IRMv1 by fixing a scalar classifier and learning representation & such that the classifier is "approximately locally optimal" (Arjovsky et al. (019))
min HE
[e( + \ w/ w| (wo) (w)) Ve E F.
(4)
BEE
where is scalar evaluated in and À controls the strength of the penalty term on gradients on w. Alternatively, the recently proposed Invariant Risk Minimization Games (IRMG) (Ahuja et al. 20200) proposes to learn an ensemble of classifiers with each enviromment controlling one compo- nent of the ensemble. Intuitively, the environments play game where each environment's 's action is to decide its contribution to the ensemble aiming to minimize its risk. Specifically, IRMG optimizes the following objective:
min Σ Re Φ dEH,
s.t. W. = arg min Rd wEH
(W Ye E E,
(5)
E
cEE
where to E
EeEE is the average and w = Ed'E,e We the complement classifier.
3 CONTINUAL IRM BY APPROXIMATE BAYESIAN INFERENCE
Both IRM and IRMG assume the availability of training data from all environments at the same time, which is impractical and unrealistic in numerous applications. A natural approach would be
2

---

Under review as a conference paper at ICLR 2021
to combine principles from IRM and continual learning. Experience replay, that is, memorizing amples of past environments and reusing them later, could be possible in some scenarios but it is often difficult to estimate a-priori the extend of replay necessary achieve satisfactory general- ization capabilities. Here, we propose to adopt a probabilistic approach, exploiting the propagation of the model distribution over environments using Bayes' rule. We integrate both IRM and IRMG with stochastic models, introducing their variational counterparts that admit a continual extension. In addition, our approach is justified by the property of the Kullback-Leibler (KL) divergence that promotes invariant distributions when used in sequential learning (as shown in Theorem
3.1 VARIATIONAL CONTINUAL LEARNING
Following prior work in continual learning (guuyen et al. (2018)) ) let D, be the training data from the t-th environment et, let Dí be the cumulative data up the t-th environment, and let 0 be the parameters of the feature extractor. When each environment is given in a sequential manner, we can use Bayes' rule and we have (all proof' are provided in the supplementary material)
p(0 DJ) x ()(,))))(()1))),,
that is, once we have the posterior distribution p(0 at time t - 1, we can obtain, by apply- ing Bayes rule, the posterior p(0 Di) at time up to normalization constant. This is achieved by multiplying the previous posterior with the current data likelihood p( D, (0). The posterior distribu- tion in general not tractable and we use an approximation. With the variational approximation, Pi) R ge (0), it is thus possible to propagate the variational distribution from one environment to the next. From Corollary [4n the supplementary material) we can write the continual variational
Bayesian inference objective
go arg min ((,,)) Eema(0) (4y fo(x))) + Dra,(g())la- (())
from the variational distribution step (0), with fo w o. function with parameters 8.
3.2 EQUIVALENT FORMULATION OF IRM BILEVEL OPTIMIZATION PROBLEM (BIRM)
In order to extend the IRM principle of Equation S using the principle of approximate Bayesian inference, by applying Lemma (in supplementary material), we first introduce the following new equivalent definition of IRM (equation 3).
Definition (Bilevel IRM (BIRM)). Let Hg be a set of feature extractors and let Ho be the of possible classifiers. An invariant predictor
ó set of en- vironments is said to satisfy the Invariant Risk Minimization (IRM) property, if it the solution the sollowing bi-level Invariant Risk Minimization (BIRM) problem
min wE H. ER(
(8a)
S.f. =
(8b)
This formulation results from substituting the minimization conditions in the constraint set of the original IRM formulation with the Karush-Kuhn- (KKT) optimality conditions. This new formulation allows introduce efficient solution methods and simplifies the conditions of IRM. also justifies the IRMv1 model; indeed, when the classifier is a scalar value and the equality constraint is included in the optimization cost function, we obtain Equation To solve the BIRM problem, we propose to use the Alternating Direction Method of Multipliers (ADMM) (Boyd et al. (20TT). ADMM is an alternate optimization procedure that improves convergence and exploits the decomposability of the objective function and constraints. Details of the BIRM-ADMM algorithm are presented the supplementary material.
3.3 BILEVEL VARIATIONAL IRM
At this point, cannot yet directly extend the IRM principle using variational inference. That is because we observe all environments at the same time, the prior of the single environment is data independent. Therefore, we substitute (f) from Equation with priors Po( and pw (w), where and w are now the parameters of the two functions ó and w. We also substitute (.) with the variational distributions qo and (w))
3

---

Under review as a conference paper at ICLR 2021
Definition 2 (Bilevel Variational IRM (BVIRM)). Let Pg be family of distributions over feature extractors. and let Pw be a family of distributions over classifiers. A vari- ational invariant predictor a set of environments E is said to satisfy Bilevel Varia- tional Invariant Risk Minimization (BVIRM) if it is the solution to the following problem:
min PoEPs qu cEE
Eg 96)
(9a)
s.t. V qu Qw (gw 9o) = 0, E,
(9b)
with
Qa(qus 96) = Burgg
R o 6) 3D x (qwllp))
(10a)
orgg
and
Qu (Aw: 96) = Burrqu
R (w 3DLL (4 (pwlpw))
(10b)
drugs
and where Pe and Pwo are the priors of the two distributions. B is hyper-purameter balancing the ERM and closeness to the prior.
Definition extends Definition [ with the objective of Eq 7 where the parameters & and wo are substituted by their distributions 4s and qu- The gradient of the cost in the inner prohlem is taken with respect to the distribution lw- When we parameterize Is with 9 and qu with w, the gradient evaluated with respect these parameter since the condition implies that the solution locally optimal. If (p, q) is convex in the first argument, then the solution is globally optimal. This defini- tion extends the IRM principle to the case where we use approximate Bayes inference, shaping the variational distributions qu and go. to be, in expectation, invariant and optimal across environments.
3.4 THE BVIRM ADMM ALGORITHM
As noted for the BIRM definition, the solution of the variational BVIRM formulation can be ob- tained by using ADMM (Boyd et al. (2011)). While in general there are no convergence results of ADMM methods for this problem, for local minima, under proper conditions F the stochastic version of ADMM converges with rate O(I/Vt) for convex functions and Ollog t/t) for strongly convex functions (Ouyang et al. 2)13)). We are now in the position to write the BVÍRM-ADMM formulation of the BVIRM problem. ADMM defined by the update [q where we denote with the apexes and the value of any variable before and after the update. Moreover, we abbreviate as follows ((w,, 0) = ((a(w),((0)).
, w = arg min Lp (Wes uE W v. E,
(1la)
We
w+
1/|E| Σ (we + ue)
(11b)
ut
ue + (w.t wt)
(11c)
ve = vf + Va(w) Qi (wt (utt
(11d)
with
PI Lp(wer We: w. ve) (wee 0)+ + Po |lie - w 1el1 2
VV Vi (wr o の) + vell?. (12)
2
Here, & is fixed and 0 is updated in an externa loop or given (e.g. the identity function). In the experiment we use stochastic Gradient Descent (SGD) to update both the model parameters we and the feature extractor parameters é. The result follows by applying Lemmal
I in the supplementary
material and substituting x; + (a.))) <- 0% (we, 0) and (x;r;)
V Vatw) (w 0). We
provide a pseudo-code implementation leveraging Equation 11 as Algorithm [ One of the advan- tages of the ADMM formulation of BVIRM of [q[ is that it can be computed in parallel, where only Eq requires synchronization among environments, while the other steps can be computed independently.
Implementation detail using the mean field parameterization and reparametrization trick is provided in the Supplementary Material These conditions specific bounds on the magnitude and variance of the (sub-)gradients of the stochastic function (uyang al. (013)). We used ELU e C in the experiments.
4

---

Under review as a conference paper at ICLR 2021
Algorithm 1:
W.
BVIRM-
Algorithm 2: w.
C-BVIRM-
ADMM(E, R ADMM version of the Bilevel Variational IRM Algorithm
ADMM(E, R ADMM version of the
Bilevel Variational IRM Algorithm
Result: w o o: feature extraction and classifier for the environment E
Result: ww o ob - : feature extraction and
classifier for the environment E // Randomly initialize the variables
// Randomly initialize the variables
w. we: Me, de, 0 + Init(:
w, Ve,, 0 t Init(;
// Outer loop (on 8) and Inner loop (on w)
2 =0;
for E E do
while not converged do using SGD
4
for 1,..., K do
// Update
0 SGD a (q:(w:: 96));
0 = SGDD((ees (gw {o)): for k = 1,.... K do for e E do
while not converged do Update w using SGD and ADMM
We = SGD w. Lp (We, we we)
We SGD w, Lp(we, ue,w, ve):
w 1/| Ee (We ue)
w 1/2(we + Ue W);
Ue = + (we w):
ue (wo w):
Ve vee + V Qt (Wes 0):
10
Uc Ve + V (Wee 0)
10
end
end
end
end
end
13
W
We
14
end
3.4.1 THE CONTINUAL BVIRM ADMM ALGORITHM
In presence of sequential environments, the priors for the new environment are given by the pre- vious environment's distributions 4s and tw this is obtained by comparing the BVIRM defini- tion Eqs.
( with the continual Bayesian learning Equation 77. In Equation TO we thus now have O (Bur, Io Em-ag.sr-g Re (w Φ) + 3DKI 1 ) +.DDII.( and Qn (gw, (o) = Ewrnqw drogg 12d (w Φ) + 3DKL (qwellul (w)) Algorithm presents an example implementation of ADMM applied the continual BVIRM formulation.
3.5 INFORMATION- IAPORMATION-THEORETTT THEORETIC INTERPRETATION OF C-BVIRM
The KL divergence provides an additional motivation for the meth- ods we propose. Indeed, for causal discovery Peters et al. (2015) suggests a discovery mechanism for causal variables as the intersec- tion of the invariant conditional distributions across environments subject to interventions. The KL divergence is asymmetric and only components present in the first argument distribution are evaluated. This implies that by using the KL divergence we can compute the intersection of the distributions, even when these are observed se- quentially. This can be made more explicit by the property of the information projection Cover (1999)
Pr
Theorem 3 (Information Projection). If and Q are two fam- ilies of distributions with partially overlapping support, @ (P) SQpp(Q), and q E Q, then
Figure 1: Sequential pro- jection of distributions
Pi....Ph Pt. where Pi+1 arg minpee... DKL(pllp)
p¹ arg min (rl(())))) pEP
has support in the intersection for the support of P and q, or supp(p") C supp(P) supp(())
3 In Algorithm the ADMM update equation is implemented from line 6 to line 9, while in Algorithm 2, from line to line 10.
5

---

Under review as a conference paper at ICLR 2021
6
8 9
0 0473773 2
Figure 2: The two color models (on the left b01, on the right b1l) 1) for the train (upper row) and test (lower row) of the MNIST (left) and FashionMNIST (right) datasets.
Therefore, if we have a sequence sets of distributions of models from intervention environments and we compute the projection in sequence, the final projected distribution has support on the inter- section of all previous distribution families, or supp Pt) N supp( P:) (see Figure 1] since at each step Pi+1 = arg minpeP.4. P.+1 DKL (pllp.).
4 RELATED WORK
et al. 2007 Johansson et al. 2019) aims to learn invariant features components (xx) that have similar P(o(+)) different (but related) domains by explicitly minimizing a distribution discrepancy measure, such as the Maximum Mean Discrep- ancy (MMD) (Gretton al. 2012) the Correlation Alignment (CORAL) (Sun & Saenko 6 The above condition, however, is not sufficient to guarantee successful generalization to unseen do- mains, even when the class-conditional distributions of all covariates changes between source and larget domains (Gong et al. 2016 Zhao et al. 2019). Robust optimization (Hoffmann et al. 2018 Lee Raginsky (018), on the other hand, minimizes the worst performance over set of possi- ble environments E, that is, max.eg (0). This approach usually poses strong constraint on the closeness between training and test distributions (Bagnell (2055) which often violated in practical settings (Arjovsky et al. 2019 Ahuja et al. 2020).
Generalization Domain adaptation
Incorporating the machinery of causality into learning models a recent trend for improving gener- alization. (Bengio al. 2019) argued that causal models can adapt to sparse distributional changes quickly and proposed meta-learning objective that optimizes for fast adaptation. IRM, on the other hand, presents an optimization-bast based formulation to find non-spurious actual causal faccorr to target y. Extensions of IRM include IRMG and the Risk Extrapolation (REx) (Krueger et al.) 2020). Our work's motivation is similar to that of online causal learning (Javed al. (2020), which models the expected value of target given each feature as Markov decision process (MDP) and identifies the spurious feature x; if BBy [r:]] is not consistent to temporally distant parts of the MDP. The learning is implemented with a gating model and behaves as a feature selection mechanism and, therefore, can be seen learning the support of the invariant model. The proposed solution, however, is only applicable to binary features and assumes that the aspect of the spurious variables is known (e.g. the color). It also requires careful hyper-parameter tuning. In the cases where data not divided into environments, Environment Inference for Invariant Learning (EIIL) classification method (Creager) (et al. (2020)) aims at splitting the samples into environments. This method proves to be effective also when the environment label is present.
Continual Learning Kirkpatrick et al. (017): :: De Lange et al. (0119) addresses the problem of learning one classifier that performs well across multiple tasks given a sequential manner. The focus is on the avoidance of catastrophic forgetting. With our work, we shift the focus of continual learning to the study of single task that observed in different environments.
EXPERIMENTAL EVALUATION
5.1 DATASETS AND EXPERIMENT SETUP
Colored MNIST Figure 2 (left) shows sample of train (upper) and test (lower) samples. In each training environment, the task to classify whether the digit is. respectively, even or odd. As in prior work, we add noise the preliminary label by randomly flipping it with probability

---

for "trouser", Under review as a conference paper at ICLR 2021
Table I: Mean accuracy (N 5) on train and environments when training on consecutive environments on MNIST and the b01 color correlation.
C-BVIRM
C-VIRMG
C-VIRVI
ERM
EWC
GEM
IRMG
IRMv1
MER
VCL
VCLC
train
71.3
69.1
51.4
86.4
87.4
87.4
86.3
85,3
87.3
89.3
89.3
(4.2)
(2.8)
(3.4)
(1.2)
(2.7)
(2.7)
(1.2)
(0.8)
(1.7)
(0.6)
(0.6)
Test
29.6
27.9
46.0
12.7
15.7
15.7
12.8
9.9
14.8
24.9
24.9
(3.3)
(8.5)
(2.1)
(2.7)
(4.2)
(4.2)
(2.6)
(0.2)
(3.8)
(1.9)
(1.9)
of 0.25. The color of the image is defined by the variable 2, which the noisy label flipped with probability Pe [0.1, 0.2]. The color of the digit is green if is even and red if z is odd. Each train environment contains 30, 000 images of size 28 x 28 pixels, while the test environment contains 10, 000 images where the probability P. 0.9. The color of the digit (b01) the background (b11) is thus generated from the label but depends on the environment. Figure depicct the causal graph (the hammer indicating the effect of the intervention) of the environment. The variable "Color" is inverted when moving from the training to test environment.
Colored FNShionMNIST, KMNIST, and EMNIST Figure 2 (right) shows the Fashion-MNIST dataset, where the variable = defines the back- ground color. Again, we add noise to the preliminary label (y 0 for "t-shirt", "pullover", "coat", "shirt", "bag" and "dress", "sandal", "sneaker", "ankle boots") by flipping it with 25 per- cent probability to construct the final label. Besides, we also consider Przzushjj-MMNiSS dataset Clanuwat et 2018 and the EMNIST Let- ters dataset Cohen et al. (2017 The former includes 10 symbols of Hi- ragana, whereas the latter contains 26 letters in the modern English alpha- bet. For EMNIST, there are 62, 400 training samples per environment and 20, 300 test samples. We set y = for letters 'a', 'c', 'c', 'g', 'i', K. 'm', "o", 'q", 's', 'u', 'v", 'y" and y 1 for remaining ones.
Color
Digit
Image
(ataset(Cohen
Figure 3: Causal re- lationships of colored MNIST.
Reference Methods. We compare with a set of popular reference methods in order to show the ad- vantage of the variational Bayesian framework in learning invariant models in the sequential environ- ment setup. For completeness, we also evaluate the performances of four reference continual learn- ing methods. These include Elastic Weight Consolidation (EWC, Kirkpatrick al. (2017). Gra-
dient Episodic Memory (GEM, (Lopez-Paz & Ranzato 2017 6 Meta-Experience Replay (MER, Riemer et al. 2018 and Variational Continual Learning (VCL. Swaroop al. 2019); Nguyen et al. (2018) 8 ERM is the classical empirical risk minimization method; we always use the cross- entropy loss. IRMv1 enforces the gradient of the model with respect to scalar to be zero. IRMG models the problem as a game among environments, where each environment learns a separate model. EWC imposes a regularization cost on the parameters that are relevant the previous task, where the relevance is measured by Fisher Information (FI); GEM uses episodic memory and com- putes the updates such that accuracy on previous tasks not reduced, using gradients stored from previous tasks; MER uses an efficient replay memory and employs the meta-learning gradient up- date to obtain a smooth adaptation among tasks; VCL and Variational Continual Learining with coreset VCLC apply variational inference continual learning. C-VIRMv1 and C-VIRMG refer to, respectively, our proposed variational extensions of IRMv1 and IRGM in sequential environ- ments. C-BVIRM is the implementation with ADMM.
All hyper-paramter optimization strategies and simulation configurations are discussed in detail in the supplementary material.
https://gituub..oo/rris cond/m/ooooiiisssccniiiisss
wttp://tt/ppoooccccondddeee/ennn services/emnist -arrrcsss/emmitt-tttssstt https//ggthubbccomm/mtttteeeneee

---

Under review as a conference paper at ICLR 2021
Table 2:
Mean accuracy (over runs) and standard deviation at time for (n) 6, 10 environ- ments, (d) across datasets, and (c) for the two color correlations (b01,b11).
C-BVIRM
C-VIRMG
C-VIRMV
ERM
EWC
GEAA
IRMG
IRMvI
MER
VCL.
VCLC
2
29.6
27.9
46.0
12.7
15.7
15.7
12.8
9.9
14.8
24.9
24.9
(3.3)
(8.5)
(2.1)
(2.7)
(4.2)
(4.2)
(2.6)
(0.2)
(3.8)
(1.9)
(1.9)
28.8
27.5
47.1
11.1
15.6
15.6
12.2
9.7
15.3
15.4
15,4
(4.1)
(2.6)
(3.0)
(2.9)
(5.7)
(5.7)
(2.8)
(0.3)
(4.4)
(1.1)
(1.1)
10
21.8
25.2
31.0
10.2
17.7
17.7
12.4
10.2
15.5
10.8
10,8
(2.4)
(4.5)
(7.1)
(0.2)
(5.5)
(5.5)
(2.1)
(0.2)
(4.1)
(0.2)
(0.2)
MNIST
29.6
27.1
46.8
12.7
15.7
15.7
12.8
9.9
14.8
24.9
24.9
(3.3)
(7.6)
(2.6)
(2.7)
(4.2)
(4.2)
(2.6)
(0.2)
(3.8)
(1.9)
(1.9)
Fa-MNIST
26.7
10.7
15.4
15.3
10.8
9.9
13.2
24.9
24.9
(4.3)
(8.7)
(3.6)
(1.5)
(5.1)
(5.4)
(1.4)
(0.2)
(2.8)
(2.0)
(2.0)
KMNIST
32.8
24.2
46.5
12.0
14.0
14.0
12.1
9.9
15.6
24.9
24.9
(4.6)
(6.0)
(1.9)
(2.1)
(3.5)
(3.5)
(2.4)
(0.2)
(4.1)
(2.0)
(2.0)
EMNIST
32.1
25.0
45.9
10.8
14.8
10.0
12.6
24,9
(4.6)
(7.5)
(2.5)
(1.0)
(3.6)
(3.7)
(1.2)
(0.2)
(2.3)
(2.0)
(2.0)
b01
29.6
27.1
46.8
14.9
18.8
18.8
14.6
9.8
18.0
24.9
24.9
(3.3)
(7.6)
(2.6)
(0.8)
(1.3)
(1.3)
(0.8)
(0.1)
(1.0)
(2.0)
(2.0)
b11
38.8
23.9
43.3
9.9
12.5
12.5
9.8
9.9
11.5
24.9
24.9
(4.2)
(6.9)
(4.5)
(0.2)
(3.6)
(3.6)
(0.1)
(0.2)
(2.0)
(2.0)
(2.0)
5.2 RESULTS
Table 1 lists the training and test accuracy on the MNIST dataset with the color correction b01 (see Figure left). Since we introduced label noise by randomly flipping 25 percent the given la- bels, hypothetical optimal classifier would be able achieve accuracy of 75% in both training and test environments. ERM, IRMv1, and IRMG perform poorly in the setup where environments given sequentially. Similarly, reference continual learning methods also fail to learn invariant representation in the new environment. As these models are learning to mainly use spurious fea- tures for the classification problems at hand, here: the colors of the digits (red~-odd; green~evenn, they perform poorly (much worse than a random baseline) when the spurious feature properties are both IRM and IRMG
inverted (green-odd; red-even). In contrast, achieve a classification accuracy higher than 45% on the test data. This implies that our model is relying exclusively on spurious correlations present in the color of digits. By comparing the performance between C-VIRMyl and C-BVIRM, we conclude that (1) our proposed bilevel invari- risk minimization framework (i.e., the BIRM in Definition 1) is an effective alternative to the original formulation AArjovsky et al. (2019); and (2) ADMM is effective in solving the BIRM op- addition, timization problem and has the potential to improve the generalization performance. one can observe that the KL divergence term in VCL and our framework significantly improves the test accuracy with respect to the baseline counterparts. This result further justifies our motivation of using variational Bayesian framework for the problem of continual invariant risk minimization. Table 2 lists the accuracy on the test environment for: (n) (upper rows) an increasing number of se- quential environments (d) (central rows) different datasets, and (c) (lower rows) the given color correlation schemes. We can observe that there general trend in the results. IRMG and IRM, with an accuracy of less than 10%, are not able learn invariant models. Similarily, the continual learning reference methods (MER, EWC, MER, VCL, VCLC) also fail with a test accuracy of under 25%. The proposed methods on the other hand provide mechanism learn more robust features and classification models. The higher variance of the accuracy is caused by the stochastic nature of the variational Bayesian formulation.
variational extensions
5.3 ENVIRONMENT INFERENCE FOR CONTINUAL INVARIANT LEARNING
In practical applications. the environmental labels are usually unavailable, which means that it is difficult impossible manually partitioning the training set into "domains" or "environments".
8

---

Under review as a conference paper at ICLR 2021
Table 3: Mean accuracy (over 10 runs) on train and test environments when training off-line on environments on Colored-MNIST, with the EIlL. (pc1 = 0.2, pc2 0.1, 50'000 samples)
IRMv1 Train
C-VIRMv1
Test
Train
Test
No EIIL
70.73 (1.16)
67.48 (1.96)
70.99 (0.90)
66. 60 (2.66)
EIIL
73.78 (0.61)
67.96 (3.01)
75.29 (0.53)
68.40 (1.11)
Table 4: Mean accuracy (over runs) on train and test environments when training on environment on Colored-MNIST, with and without EIll.. (pe) 0.1)
Without Environment Inference
With Environment Inference (EIIL)
IRMv1
C-VIRMvI
IRMvl
C-VIRMv1
N,
Train
Test
Train
Test
Train
Test
Train
Test
1'000
93.7 (0.7)
13.5 (1.7)
94.1 (1.1)
13.7 (1.5)
95.5 (0.3)
12.7 (2.0)
96.0 (0.4)
16.5 (6.0)
2'000
91.5 (0.4)
(0.9)
91.1 (0.7)
11.9 11.9
92.6 (0.4)
27.8 (2.8)
93.3 (0.7)
(3.4)
5'000
90.2 (0.4)
(1.1)
90.1 (0.4)
(0.7)
91.6 (0.4)
91.6 (4.8)
30,6 (3.2)
10'000
89.9 (0.3)
10.1 (0.5)
(0.2)
(0.1)
85.3
43.9)
83.7 (1.2)
50.4 50.3)
20'000
90.0 (0.2)
10.1 (0.2)
90.1 (0.2)
10.1 (0.0)
77.2 (1.2)
(2.2)
77.9 (1.1)
57.6 (2.0)
50'000
90.1 (0.1)
9.7 (0.4)
90.0 (0.1)
10.0
73.9 (0.5)
67.2
74.0 (0.5)
67.3 (1.0)
In order to generalize our continual invariant learning models to an environment--gnnosiic setting, we leverage the recently proposed Environment Inference for Invariant Learning (EIIL) by Creager al. 2020) to automatically infer environment partitions from observational training data, and integrate EIIL into continual invariant learning models.
We take our proposed C-VIRMy1 as an example. According to Table 3 it is easy to observe that inferring environments directly from observational data (using EIIL)) has the potential improve (continual) invariant learning relative to using the hand-crafted environments. Moreover, C- C-VIRMv1 with EIIL improves both training and test accuracy, compared with IRMv1 with EIIL. In fact, this environment partition strategy also enables invariant learning with only one environmental data. Table f further suggests that the generalization accuracy improves for both IRMvl and C-VIRMy] as the number of training samples increases. Again, we observed that, when combined with EIIL, C- C-VIRMv1 always outperforms IRMv1.
CONCLUSIONS
We aim broaden the applicability of IRM settings where environments are observed sequen- tially. We show that reference approaches fail in this scenario. We introduce variational Bayesian approach for the estimation of the invariant models and a solution based on ADMM. We evaluate the proposed approach with reference models, including those from continual learning, and show significant improvement in generalization capabilities.
REFERENCES
Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk min- imization games. arXiv preprint (XXi:2002204469) 2020.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv: 1907 2019.
J Andrew Bagnell. Robust supervised learning. In Proceedings of the national conference on artifi- cial intelligence, volume 20, pp. 714, 2005.
Shai Ben- Denidavid, John Blitzer, Koby Crammer, and Fernando Pereira. Analysis representations for domain adaptation. In Advances in neural information processing systems, pp. 137-144, 2007.

---

Under review as a conference paper at ICLR 2021
Yoshua Bengio, Tristan Deleu, Nasim Rahaman, Rosemary Ke, Sébastien Lachapelle, Olexa Bila- niuk, Anirudh Goyal, and Christopher Pal. meta-transfer objective for learning to disentangle causal mechanisms. arXiv preprint arXiv:1901.10912, 2019.
Stephen Boyd, Stephen P Boyd, and Lieven Vandenberghe. Convex optimization. Cambridge uni- versity press, 2004.
Stephen Boyd, Neal Parikh, and Eric Chu. Distributed optimization and statistical learning via the alternating direction method of multipliers. Now Publishers Inc, 2011.
Tarin Clanuwat, Mikel Bober-Irizar, Asanobu Kitamoto, Alex Lamb, Kazuaki Yamamoto, and David Ha. Deep learning for classical japanese literature. arXiv preprint arXiv:1812.01718. 2018.
Gregory Cohen, Saeed Afshar, Jonathan Tapson, and Andre Van Schaik. Emnist: Extending mnist to handwritten letters. 2017 International Joint Conference on Neural Networks (LJCNN), pp- IEEE, 2017.
Thomas M Cover. Elements of information theory. John Wiley & Sons, 1999.
Elliot Creager, Jörn-Henrik Jacobsen, and Richard Zemel. Enviromment inference for invariant learn- ing. In ICML Workshop on Uncertainty and Robustness, 2020.
Matthias De Lange. Rahaf Aljundi, Mare Masana. Sarah Parisot, Xu Jia, Ales Leonardis, Gregory Slabaugh, and Tinne Tuytelaars. Continual learning: A comparative study on how to defy forget- ting in classification tasks. arXiv preprint arXiv: (00::999900088 009908833,....00.0 2019.
Mingming Gong, Kun Zhang, Tongliang Liu, Dacheng Tao, Clark Glymour, and Bernhard Schölkopf. Domain adaptation with conditional transferable components. In International con- ference on machine learning, pp. 9939-2844 2016.
Arthur Gretton, Karsten M Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alexander Smola.
A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723 2012.
Judy Hoffman, Mehryar Mohri, and Ningshan Zhang. Algorithms and theory for multiple-source adaptation. In Advances in Neural Information Processing Systems, pp. 8246-8256, 2018.
Khurram Javed, Martha White, and Yoshua Bengio. Learning causal models online. arXiv preprini (0l::20066.74611, 2020.
Fredrik D Johansson, David Sontag, and Rajesh Ranganath. Support and invertibility in domain- invariant representations. arXiv preprint
2019.
Diederik P Kingma and Max Welling. 1.Xiv:1312.6114,00 2013.
Auto-encoding variational bayes.
arXiv preprint
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska- Grabska-Barwinska Barwinska, et al. Overcom- ing catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 114(13):3511-322 2017.
David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Remi Le Priol, and Aaron Courville. anttof-distribution generalization via risk extrapolation (rex). arXiv preprint S00:::2003000688 2020.
Jaeho Lee and Maxim Raginsky. Minimax statistical learning with wasserstein distances. In Ad- vances Neural Information Processing Systems, pp. 2687-2696, 2018.
David Lopez-Paz and Marc' Aurelio Ranzato. Gradient episodic memory for continual learning. In Advances in Neural Information Processing Systems, pp. 6467-6476, 2017. 2017.
Cuong V. Nguyen, Yingzhen Li, Thang D. Bui, and Richard E. Tumne. Variational continual learn- ing. In International Conference on Learning Representations, 2018.
Hua Ouyang, Niao He, Long Tran, and Alexander Gray. Stochastic alternating direction method of multipliers. In International Conference on Machine Learning, pp. 80-88, 2013.
10

---

Under review as a conference paper at ICLR 2021
Judea Pearl. Causality. Cambridge university press, 2009.
Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference using invariant predic- tion: identification and confidence intervals. arXiv preprint arXiv: 1e0i:::551..33222 2015.
Matthew Riemer, Ignacio Cases, Robert Ajemian, Miao Liu, Irina Rish, Yuhai Tu, and Gerald Tesauro. Learning learn without forgetting by maximizing transfer and minimizing interfer- ence. arXiv preprint arXiv: /9/0,, 2018.
Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In European conference on computer vision, pp. 443-450. Springer, 2016.
Siddharth Swaroop, Cuong V. Nguyen, Thang D. Bui, and Richard E. Turner. Improving and understanding variational continual learning. arXiv:1905.008 [es, staf], 2019. URL http:
Warxivvorg/ab//1955.020999
Vladimir Vapnik. Principles of risk minimization for learning theory. In Advances in neural infor- mation processing systems, pp. 831-838, 1992.
Han Zhao, Remi Tachet des Combes, Kun Zhang, and Geoffrey J Gordon. On learning invariant representation for domain adaptation. arXiv preprint arXiv:1901.0944 2019.

---

Under review as a conference paper at ICLR 2021
A SUPPLEMENTARY MATERIAL
A.1 I VARIATIONAL INVARIANT RISK MINIMIZATION GAMES
We now consider the IRMG objective and extend it with the variational Bayesian inference. we observe all environment at the same time, the prior of the single environment is data independent. From Equation 7 we thus substitute (r- with a priors (0) and qw,(w,, where and w are now the parameters of the two functions and w. While we substitute qe (0), with the variational distributions (f) and (w). The outer problem is now
min 4s
Bd Rd (w o @) + BD&L llpp)
(13a)
s.t.
Qw.
arg min Ew-quu que R qwe
(w + tw-) o + (D, E E (13b)
qw (w) is the average classifier and w-e EeeEr We' Wel ~ qu (w) is the complement classifier. In the reformulation of the IRMG model, we weight the distance of the varional distribution to the prior with 3. We notice how the difference of the variational formulation of IRMG differs the presence of the mean the distri- bution of the function over the variational distributions and the KL term.
where w
1 EeEEr Wes We We
[E]
We can now finally extend IRMG when the environments observed sequentially. Combining the definition of IRMG Eqs. 5 with the continual bayesian learning Equation 7 we obtain the variational objective of IRMG sequential environment case.
min 9e
E (ly,, v x)} + 3 (uol(uullu
(14a)
s.t.
w (w w)-1),, w ~ yoo (w)) wi-1 ~ Tw (w)
(14b)
Iu arg min 1~ (e(y 1 w1-1) a o) + DD 1l
(14c)
4-
We can similarly extend the definition of IRMv1 when all environments are seen at the same time and sequentially.
A.2 MEAN FIELD PARAMETRIZATION AND REPARAMETRIZATION TRICK
When want implement Equation II and Equation and the different variation, the mean field approximation and the reparametrization trick Kingma & Welling (2013). In this case the density function of our model is parameterized by 0 and w and constraints becomes Qu (4 0) = 0 + Q 0) = 0. we then parametrize j(ww)) and (ww)) the mean and standard deviation and model the distribution as qu (w) = ,سww) + <a((w) with e ~ NN 1) We now want to compute the gradient (in the following we ignore the dependence the ø and its parameters)
VwQ(w)) V. Vu R(w o) (V.
The second term is
V DKL = (. V. Dr
with
Valt = 1, Vwo=
V D. p) = -0 M4)
V. DKL (glp)
(iiag(;))~1 -1 diag(op)
where we assume pp og to be diagonal, in this way the previous equation can be evaluated element- wise and where the DKL (gll)) is defined as
EΣ Σ
= In
n + to Ex Eng} + (up 11) IT (pp mg)
12

---

Under review as a conference paper at ICLR 2021
The first term is evaluated by Monte Carlo sampling
V Wul R(w) V.
R(we)
with
- =(w) + (
and wi (0, 1). Also in this case
N
1 N
R(ws)= = V R(w, Voll +
Σ R(w;Vuwe
A.3 THE BIRM-ADMM ALGORITHM
We observe that to solve BIRM we can use Lemma I1 and write the following algorithm
w = arg min
Lp We, ue w 22 ),
(15a)
We
w 1/ Σ (e + ue)
(15b)
o ug + (w+ - -
(15c)
ve ue + R. Rf (wt ())
(15d)
where
Lp(Wer Uos w. U6)
R R (we o Φ)
PU
W.
+ uol1 + 2 I Re (w, o + vell (16)
denote .t.- the values of the variable after and before the update. In order to implement the method the SGD to update the model we and in outer loop updating for ø.
Algorithm 3: w. ø + BIRM-ADMMEL R*) ADMM version of the Bilevel IRM Algorithm
Result: feature extraction and classifier for the environment E // Randomly initialize the variables
w, we, Ue, Ve, Ø +- Init();
// Outer (on Φ) and Inner loop (on w) while not converged do
// Update ø using stochastic gradient descent (SGD) (w
== SGD,
for k = 1,... K do for e E E
We SGD w. L L (We, Mey W, Ve);
w = 1/| De [w, + up)
u. Ue + [e w)
Ve = Ue + V m
R (. o Φ)
10
end
11
end
12
end
A.4 VARIATIONAL INVARIANT RISK MINIMIZATION
Definition (VIRM). Give set of distribution over the mapping P,, and distribution over the set of classifier Pa., variational inyariant predictor set of environments E. is said satisfy the
13

---

Under review as a conference paper at ICLR 2021
Variational Invariant Risk Minimization (VIRM) it is the solution of the following problem
min g Eu
Σ cEE
Volaw, (allw,yo) 4o)
(17a)
s.t.
tu arg min Q (qw: (6), oLEPu
(17b)
where
2 (hes (ho, 96) Ewwyn
( Φ) + (gs (iD
(17c)
Q (9w, 96) = Bwriqu
R (w o) + 3DL (gw lpw)
(17d)
grugo
and ps. Pw are the priors of the two distributions.
A.5 BILEVEL ALTERNATIVE FORMULATION
We state here a general result on solving Bilevel Optimization Problems Lemma (Bilevel Reformulation).
min
(().)))( <
(18a)
s.t.
g(z) E arg min f(a, ((a,))
(18b)
then we can solve the equivalent problem
min
F(x,y) ((r,,((())) So,
(19a)
2,y.M
V,L(a, u) 0,
(19b)
u
(19c)
g(r,y) $0
(19d)
u=g(z,v)= =0 f(x,)) + g(
(19e)
L(,,y,u)
(19f)
Proof of Lemma
Lemma
follows by applying the Karush-Kuu where the Lagrangian function is L(x, y, u) [(x, + (flggu/g, u)- D
conditions (Chapter
Boyd et al. 2004) to Eq 18
Lemma 6 (Equivalence of Definition
D
Definition
is equivalent to Eq. 3 the Invariant Risk
Minimization.
Proof of Lemma 0 The result follows by apply Lemma 5 to Eq
Lemma (Definition D. Definition the extension of Eq. S the Bilevel Invariant Risk Minimiza- tion, when the function is described by the distributions of their variable and w.
Proof of Lemma The result follows by inspecting Eq. of the aggregated loss function, which is now, from Eq
The equation requires the minimisation
Qg go) = Eming (w o Φ) +
+ (Drl((uullp))))
(20)
druqa
where we have separated the two contributions øi and w, and used genetic prior distributions and pw- This is by the additive property of KI. divergence:
(go P&P.)) DKL(qsllpo) + DkL(gullpw),
(21)
since we model the two distributions independently, i.e. Jo.w Lo and fa.w PaPa. Since the classifiers' losses shall be minimal for all environments, this condition is substituted by requiring the gradient with respect qu to be zero, Ve. The gradient w.r.t. qu of the second term of Eq 20 zero.
14

---

Under review as a conference paper at ICLR 2021
A.6
THEOREM AND IRM CONNECTION
A.6.1
SEQUENTIAL INFORMATION PROJECTION
In Theorem 3
we show that the Information Projection (IP) shrinks the support of the output distri-
bution.
Lemma 8. we have sequence of families of distributions P. Let Pt E Pi and
Pi+1 arg min pEPr11
(p.(p.P)) (p.P:)
then
dudp it C supp(P3)
We have that Vi, supp p. C supp(P:) N supp(P:-1). where the first condition follows from p; E P; in the minimization and the second from Theorem 3 The results follows by iterating the property.
Proof of Lemma g
A.6.2 IRM AND INFORMATION PROJECTION
We show now two ways to state the connection of the IRM principle and the sequential IP. Let g be the dis- tribution of the previous environment and R(g) the loss function of the current environment, where q denotes the distribution the network parameters. Let
= arg min R(g)
be the optimal distribution for the current environment. We can then consider the Taylor expansion of the param- eters distribution around the optimal distribution as
Figure 4: Sequential IRM projec- tion of distributions, where qt arg min (pll + (xI.(lll"))
R(g) = R(q") Aq R(q")
we can compute the new distribution as
Aq"
arg min DKL + Aqlla Ag
s.t.
Aq VgR(g") Se
and then
p+ = q" + Aq
Or alternatively
p".g"
arg min DKT. (plla + (x1l(1llg) P:o
s.t.
VgR(g)=
and then
at p*
Or more simply
q arg mig DKL (pllp (KL(pllg")
(22)
This last equation, shows how the distribution is the intersection of the optimal distribution at the previous step q and the current optimal distribution q" Fig shows visually, how the new distribution is the result of projecting into two distributions q" and
15

---

Under review as a conference paper at ICLR 2021
A.7 OUT OF DISTRIBUTION GENERALIZATION
The question arises if the property of generalization to out of distributions given by Theorem 9 in Arjovsky al. 2019 also holds for BIRM and BVRRM.
Lemma 9. If o and are linear functions and w t w is solution of Eg 8 it then satisfies
DExx
[x" 4" w DEx=,y [
Proof. Lemma follows from the fact that
V (w 6) = 4l x v"w -e-x [X
(23)
(24)
For BIRM thus Theorem 9 of Arjossyy al. (2019) applies directly. A similar results holds for the the BVIRM model
Lemma 10. If à ~ Ps and w ~ Ps are linear functions and w o ò = 6² w is a solution of Eq
with 3 = 0, it then satisfies
Barigg
[xxx [x" x") TExeN. [X-Xe]]
where क -Epppppl ( and Pe [w]] are the mean values.
Proof. Lemma 10
follows from the fact that
qn Qi (hwr 4o 18=0
E Bngnngu drong
R Φ)
(25)
(26)
We now take the Fréchet directional derivative in the n direction that the limit
mi Ewaggur R d) lim (Ewrrg drags
R (((( en) x) Ew
R (w 6))
drwgg
which is obtained when we differentiate the distribution qui - 6p Bewng grugs we can factorize for the direction and obtain
Ou
et). Since
R (w o ②) Engeogu druuge
27 x
XeXe Th w 21/² + Txxx,r [X
8q Rd ( Φ) = 2 gu 4
X T -EE--EEx-- [X-yeT]
grugo
We can now derive the Lemma by requiring 6y., RR (w Φ) =0
Theorem 9 of Arjovsky et al. (019) now holds when & has rank r > 0 in expectation with respect to the invariant distribution 9a. I.e. Egngg rank(1) (
A.8 GENERALIZED ADMM
The following generalization of ADMM holds: Lemma 11 (GADMM). Suppose we want to minimized
min =
(27)
we can equivalently solve the the following problem
min
-((o)aa=2 -
(28)
ZC
using the following update role (scaled ADMM)
16

---

Under review as a conference paper at ICLR 2021
xt arg min Lep 11, ,z u/ ), E
(29a)
zt 1/N Σ (re + u.))
(29b)
uf = uj + (x - =)
(29c)
o vi + g((x))
(29d)
where the augmented Lagrangian
Lolwi, Ui, vi)
E.t..) Pa Σ |1: uill² 2=ll2+
P1 2
Σ
(1a(())) + usll2 (30)
and x-x = {2 I # il is the set of all other variable, expect the i-th.
A.9 CONTINUAL VARIATIONAL INFERENCE
Following Nguyen al. 2018 we can state the following lemma.
Lemma 12 (Variational Continual Learning). Suppose we have a sequence of datasets D,, = 1,...y t drown i.idd then the variational estimation of the distribution qu at step t is given as projec- tion on KL divergence
le (0) arg min DKL.
1 7, V (00((()))(
(((0)
((0)
with Z, = 1 d8 the normalization factor, which does not depends on
Proof of Lemma 12
Let denote Dj = U D;, from i.dd = pDD)). We are interested to maximase the a posteriori probability of the paramters give the data p(8) Di)
p(0)D)) p(D,) ((DD)
since
p(0, D1)
((D))) p(0)p(D)10)
P(0) II p(D(0)
110)0(1)111))00001)))
p(0, "(0,D:-))(DD))
"10101=")((00"))000000000
thus
D(0
P(D)) "DD"
P(D(-1)
P(D)) P(00
We now use probability distribution which approximates the distribution at step @t (0) ~ 111
when then want to approximate at time t
(e) ~ (0))
17

---

Under review as a conference paper at ICLR 2021
This can be obtain by minimizing the KL divergence of the variational distribution q and the distri- bution induced by the previous step approximation, thus
[[()]] Z, (0)
q: ()) = arg min DKL 9(0)
Lemma 13 (VCLv2). The minimization of the VCL defined in Lemmal following minimization
is equivalent to solve the
(i()) = arg max (g (lo DkL(a(0))la-(0)
9(0)
with N2 i..dd. samples
N ) flogp(v((0,aa))
1
Ep g (0) (logp(D,(0)
N₂
Where the second term can be computed in closed form for known distribution as for example with the Gaussian distributions, whereas the expectation can be approximated by Monte Carlo sampling. For general loss function we can substitute the reconstruction probability with the loss function associated with neural network parametrized by 0 logp(u! 00, x))
Elyl (w d){(z())
N.
1
Es- (log p(D,)))>
o0 (w )){(1))))
N
Proof of Lemma 13
The Lamma follows from the definition of the KL diveregnce
DKL
(()]]
Zt 9 (0-1(0)(()1))) = (In - In qt 1 (0) - In + In Z,)
(In - In gt- E.g In p( De + E,g In Z.
= (kalq(0)]la-1(00) En + In Z.
The last does not depend on q. Thus the result follows.
If we substitute the log of the posterior probability with a specific loss function we obtain the fol- lowing Corollary.
Corollary 14 (Continual Variational Bayesian Inference). Given a loss function Ely, 1). the varia- tional continual learning formulated as
(0) arg min ao
flly, (xL((()))))-1(0())
(31)
((0)
with fo (w Φ)
A.10 PROOFS
p(x) In da If g(x) = 0 then p(x) otherwise the distance infinite. Second p(x) = O, then the contribution of q(z) is not considered since the integral is taken of the support of p. thus, since the intersection is not null and p the result an optimization, support of p is the intersection of the support g and support of P.
Proof of Theorem Let first first recall that (kI((1)1g)
A.11 DATASETS AND COLOR CORRECTION
We here visualize few of the dataset and color correlations. Figure 5 shows Fashion-MNIST and the b11 color correlation. In the test environment the background color of each class is inverted. In Figure we show the dataset as generated from Ahuja et al. (2020). In Figure 7 we show the EMINST (letter) and KMNIST dataset.
18

---

Under review as a conference paper at ICLR 2021
(a)
Figure 5: Fashion MNIST dataset training (a) and testing (b) environments; the color is inverted based on the b11 color correlation scheme, where the background color depends on the class of the image. In the test environment the dependency is inverted.
3
9
3
2
7
4
6
7
1
3
5
6
7
0
5 1
4
2 6
2 2
8
2
0
0
2
9
(a)
(b)
Figure 6: MNIST dataset (a) and Fashion MNIST (b) environments defined in Ahuja et (2020)
K
4
Y
K
SSC
5
e
y
+
20
W
V
S
9
B
w
(a)
(b)
Figure 7: Examples of the EMNIST dataset (a) and of the KMNIST (b).
19

---

Under review as a conference paper at ICLR 2021
A.12 PAPER-PARAMETER SEARCH AND EXPERIMENTAL SETUP
We performed hyper-parameter search around the suggested values from the original works and the values selected based on the best performance the test environment. To implement complete comparison we used for training I 1'000 samples randomly drawn from each environment. All meth- ods were trained on the same data, using random seed reset. We trained all method with 100 epochs on a batch size of 256.
IRM: 7 917 threshold = 1/2 epochs, learning rate 2.5e
IRMG: warm start=300, termination accuracy 0.6, learning rate 2.5e dropout probabil- ity 75%, weight decays = =001
ERM: learning rate 1e 3 dropout probability 75%, weight decays = .00125
MER: memory size 100 (10% of the samples), learning rate 1e-3, replay batch size =5,
3 =
.03,
1.0
GEM: memory size 100 (10% of the samples), learning rate 1e-
EWC: memory size 100 (10% of the samples), learning rate 1e-3, regularization 0.1
VCL,VCLC; learning rate = 5e corset size 100 (10% of the samples),
C-BVIRM, C-VIRMG, C-VIRMyl: weight decays .00125, 3 1., number evaluations 5, Po P1 10, step threshold =1/2 epochs, 6p 100, learning rate 1e-3
The neural network architecture is composed of 2 non-linear Exponential Linear Unit (ELU) acti- vated Full Connected layers of size 100, followed by linear full connected layer. Each layer with dropout. Dropout is not present in VCL/VCLC since not implemented in the original work. Training loss is the Cross Entropy. We tested also with the feature extraction layer separated, but with no advantage, since the test set-up only consist of one task.
The IRMv1, IRMG and ERM methods, similarly to the other methods, are trained sequentially as data from each new environment arrives. The Continual Learning methods are allowed to have limited memory of samples from previous environments.
A.13 EYPER-PARAMETER OF ENVIRONMENT INFERENCE FOR CONTINUAL INVARIANT LEARNING
We list below the values of hyper-parameten EIIL for continual invariant learning:
2 layers,
390 hidden neurons,
501 epochs,
12 regularizer weight: 0.00110794565
learning rate: 0.000459853365655466611.
numbber of runs: 10,
number EIIL iterations: 10'000,
number Monte Carlo evaluations: 3,
penalty anneal iterations: 190,
penalty weight:
prior weight: 1e-6
A.13.1 SYNTHETIC DATASET
The Synthetic Dataset is described in (Arjovsky et (2019) for testing IRM and it is defined by Structural Causal Model (Pear) 2009)), where a variable y E RN is generated by xj E RR,, while F2 E R. is generated by y. The observed variable is z (21,22). The structural equations are
F 61, 6 N(0,
(32)
y = 21 21 by (O,
(33)
21 y c 6 1()
(34)
with di fixed and a, dependent on the environment. We compared with ERM. IRM (Arjovsky et al.
2019 ), IPC (Invariant Prediction), which is the method proposed in Peters et al. 2015). 5) and EIIl.
20

---

Under review as a conference paper at ICLR 2021
Table 5:
Mean accuracy (over runs) on Synthetic Dataset (Arjovsky 2019 Creager et al.
2020 ). BIRM refers to our bilevel objective Eq. (8) optimized with ADMM.
Causal MSE
Noncausal MSE
ERM
0.827 0.016
0.824 0.015
ICP
1.000 1 0.000
0.756 0.423
IRM
0.280 1 0.006
0.290 + 0.009
BIRM
0.183
= 0.005
0.184 0.002
EIIL(IM))
0.180 0.026
0.188 0.033
Creager et al. (2020)) ). We use a similar set up of Creager et al. (020), with N = 4. The invariant model is given by (1, 0).
21


Reviews:
[{"review_id": "r5tbGSrYuSw", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "# Post-discussion update \n\nThe authors have significantly updated the paper during the discussion period. I have seen the changes, but they unfortunately do not substantiate the claim that the proposed methods are learning anything meaningful. \n\nIn the new results in table 1, the train accuracy now is better than random, but the test accuracy on the unseen environments is worse than a model that makes uniformly random predictions. This implies that the proposed method is not learning to ignore the spurious features at all. It might seem from table 1 that C-VIRMv1 is performing well --- it has 46% accuracy on the test environment after all. However, C-VIRMv1 also has the worse performance on the train set (50% accuracy). In-fact, the test or train performance alone does not mean anything in this benchmark. The goal of the benchmark is to do well on the train set while also generalizing to an unseen environment. Doing well on test set by doing poorly on the training set is not progress. \n\nThe new results in table 3 are equally troubling. The authors claim that the results in table 3 show that C-VIRM with EIIL is performing better than IRM with EIIL, but the data does not support the claim. Both IRM and C-VIRM are performing similarly; all the results are with-in error margins of each others, and the table can not be used to make any claims. \n\nI would encourage the authors to do a more systematic study of the proposed method. Investigate if the proposed methods are learning to ignore the spurious correlation at all (Table 1 suggests they are not) and only make claims that are supported by the data. In its current form, I cannot recommend this paper for acceptance. \n\n# Initial review\n\nI'm going to keep the initial review short to point out a fundamental issue with the paper. Until this issue is addressed, I see little point in discussing the paper at length. Hopefully the authors can address the issue in the discussion period and I will update my review to include other aspects of the paper. \n## Issue No 1 ##\n The proposed solution is arguably not learning at all. Both MNIST and FashionMNIST, as used in this paper and earlier IRM papers, are binary classification benchmarks (Five classes have label 0 and remaining five have label 1). A completely randomly initialized neural network gets ~50% accuracy on a binary classification task. This is not just speculation, I have run experiments using a randomly initialized network on these benchmarks to confirm that they do indeed get ~50% +- 2% accuracy on both train and test set. Coincidently, both C-BVIRM and C-VIRMG also get ~50% accuracy on the train/test set. This would imply that the model learned by C-BVIRM and C-VIRMG is indistinguishable from a model that does no learning. Here are some other baselines that will do as well as C-BVIRM: \n\n1. Set learning rate of ERM to be infinitesimally small so that no learning happens in the given number of steps (set learning rate to $e^{-100}$ for instance). \n\n2. Make completely random changes in the weight instead of updating the weights using gradients. \n\n3. Make no changes to the network. \n\nClearly 1, 2 and 3 are bad learning algorithms (the first and third are not learning at all whereas the second is making arbitrary updates), yet all three of them would perform as well as the two algorithms proposed in the paper. \n\nDoes this mean the two methods proposed in the paper are bad? Not necessarily. The BIRM formulation of the problem seems technically correct and the authors are tackling and important and interesting problem. However, the poor performance of C-BVIRM does imply that the authors need to look at their experiment results in more detail and do some sanity checks. \n\nA sanity check for checking if C-BVIRM is learning anything at all is to test it on a benchmark for which both train and test environments have a constant $p_c$ (say 0.1). A good learner should get 90% accuracy on both train and test environment in such a benchmark (Both IRM and ERM would get 90% accuracy). My guess is that C-BVIRM and C-VIRMG would get 50% accuracy on train/test when trained using the same parameters as used to run the experiments in the paper.  \n\n## Issue No 2 ##\nPlease see the private comment for the second issue. \n\nUntil issue 1 is addressed, discussing the technical details of the paper wouldn't make a difference and I can not recommend an acceptance. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "OhARvWKdaw", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nIn this work, the authors consider the problem of continual learning with distribution shifts.  The work extends the recent work invariant risk minimization (IRM) from Arjovsky et al. to a continual learning setup. IRM was designed as an offline learning framework. In this work, the authors consider the setting where the different domains arrive sequentially. The authors propose a Bayesian extension of the IRM framework that allows sequential updates as the environments arrive. They provide a justification for how the KL divergence term helps in shrinking the support continually to arrive at an approximately invariant support. Several experiments were carried out on colored MNIST and its variants to show how the proposed scheme is better than existing continual learning methods and existing IRM frameworks. \n\nPros: \nI like the problem the authors consider. It is indeed important to understand how to extend IRM type frameworks in settings such as continual learning. I also liked the ADMM based approach taken by the authors as this approach gives another way of learning IRM based predictors in addition to existing approaches. \n\nCons:\n\nThere are several problems with the paper. I discuss these problems below in different subsections. If these major issues are addressed in the rebuttal, I would be open to changing my score. \n\na) Information Projection (Theorem 3)\nIn the work, the authors show in Theorem 3 as to how when the model is trained sequentially the support of the distribution shrinks. \n\nLet us consider optimization in 2 time steps\n\nt=1, q_{t} = min_{q_{t}\\in P} KL (q_t||q_{t-1}) --> support of q_1 = P \\cap support of q_0 \n\nt=2, q_{t} = min_{q_{t}\\in P} KL (q_t||q_{t-1}) --> support of q_2 = P \\cap support of q_1 = P \\cap P \\cap support of q_0 = P \\cap support of q_0 \n\nIt is not clear from just the KL term why support would strictly shrink. The reason why the authors approach works is because the first term based on the IRM loss allows to update and shrink the support and the KL divergence term ensures that there is no need to expand the existing support. Define the first term in the loss associated with IRMv1 as Q_{IRMV1}(q_t)\n\nt=1, q_{t} = min_{q_{t}\\in P} Q_{IRMv1}(q_t) + KL (q_t||q_{t-1}) --> support of q_1 \\subset P \\cap support of q_0 \n\nt=2, q_{t} = min_{q_{t}\\in P} Q_{IRMv1}(q_t)  + KL (q_t||q_{t-1}) --> support of q_2 \\subset P \\cap support of q_1 \n\nIt is the combination of the IRM term with KL term that ensures shrinkage and KL term on its own only ensures support does not expand. A proof or an illustration of what I state would have been nice as the current intuition from the authors is incomplete and to some extent a bit misleading.\n\nb) Why penalize the IRM constraints as well?  \nThe authors arrived at a variational formulation in 9a and 9b with new objective functions defined in 10 a and 10b. It is unclear why it was chosen to impose KL penalty in the constraint (10b) as well. If not a theoretical justification, an experiment in the supplement illustrating why these choices were made would be nice. At least some good intuition needs to be provided. \n\nc) About the ADMM approach: \nI am happy to see the authors tried the ADMM approach. However, one thing that is unclear is it does not necessarily serve the purpose of solving the problem authors want. From the experiments it seems a continual extension of IRMv1 itself suffices. Is there any other reason why ADMM approach was used. Also, if ADMM approach was proposed as another way to solve standard                                                                                                               IRM in the offline setting, then some experiments explaining any advantage would be useful. At least some offline comparisons would make a better case for using ADMM. \n\nd) Why not empirically compare with Javed et al.?\n\nAs the authors correctly identify there is a recent work from Javed et al., which solves the same problem but makes more assumptions about the data.  I understand that the work uses one hot encoding of color etc. Despite that you can allow your model to also have access to the same data format and then see if your model has any more advantage to offer over and above what was done in Javed et al.\n\ne) Comparisons with offline IRM based methods (IRMv1 Arjovsky et al., IRMG Ahuja et al.) do not seem to have been carried out in a fair manner.\nThe authors have not provided a very clear description of how IRM based methods were trained. It seems to me that the authors are running offline IRM based methods in the following manner.  They provide the data from the first environment and then keep the offline IRM based models fixed and not update them when the data from the next environment arrives. If they indeed update the offline IRM based methods, then there is no reason why the methods perform so poorly.  What was the reason for not updating these models? The reason I am suspecting is that everytime retraining the entire model on the entire data is computationally expensive. If that is the concern, then authors should have shown a comparison of run-times. The authors should have run experiments allowing all methods the same run-times and then shown that IRM based methods are still not able to perform better. \nFor instance if the time to train continual IRM based model proposed by authors for the two environments is a total of 20 seconds. \nSuppose the total time to train a standard IRMv1 is say a total of 15 seconds for the first environment. There is still a five second window for the IRMv1 model to be updated.  If run-time comparisons are non-trivial, then at least there should be a budget on the total number of gradient computations that is fixed across methods. \n\n\nf) Compare with Creager et al. \n\nThe reason IRM based offline methods did not perform well is because they were given data from one environment only.  Recently, in Creager et al., it was shown how one can learn to split the data to create new environments and then train an IRM on those environments.  The authors should compare with the Creager et al. as follows. When the data from first environment comes in you can use the approach introduced in the paper http://www.gatsby.ucl.ac.uk/~balaji/udl2020/accepted-papers/UDL2020-paper-045.pdf to learn environments (code available at https://github.com/ecreager/eiil). Therefore, by doing this you can split the data in the first environment into smaller environments and use IRM. This would lead to a lot of performance improvement for the offline IRM based methods.\n\n\ng) Current experiments do not reflect the true potential of the method proposed by the authors: \n\nThe current experiments do not do a good job of reflecting a continual learning setup. Simple modifications on existing IRM based methods can outperform the method proposed by the authors. However, I believe the method proposed by the authors is nice and much more experimentation is needed. I make some suggestions on how to improve the experiments. \n\nIt seems that in the current experiments after 2 environments there is no more gain from using any method, in fact more environments seem to hurt. I would encourage the authors to do an experiment where adding more environments actually removes spurious correlations more and more and helps.  To this end, there are three suggestions I would like to make:\n\ni) First let us understand is why is two environment sufficing in colored MNIST dataset.  The reason is that the dimension of the spurious factor, i.e., color is small (only two colors). You can increase the number of colors, then more environments will be needed to decorrelate. \n\nii) If  adding colors and doing comparisons using raw images is hard, then you can try using the one hot encoded colored MNIST from Javed et al. and add more colors to it. \n\niii) Another suggestion would be to try the regression experiments from Arjovsky. In Arjovsky, the authors use two environments only. However, the two environments differ a lot 0.2 and 2.0 variance. If you introduce a sequence of environments from 0.2 to 2.0, say 0.2, 0.4, 0.6, ...2.0., then it is likely that more environments will have some benefit.  \n\nQuality: The solution proposed seems promising. There are problems with theoretical justification. Experiments need a lot of work as I stated above.  \n\nSignificance: The problem considered in the paper is interesting. The experiments carried out do not do a justice to the problem being considered. \n\nClarity: The writing of the paper can be improved. The experiments need to be described much more clearly and use the space in the supplement to give all the details. Also, the steps of ADMM shoould have been better explained.  \n\n\n\n\nReferences\n\nJaved et al. \"Learning causal models online\"\n\nCreager et al. \"Environment Inference for Invariant Learning\"\n\n\n \n\n\n\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "EfNvGph3hyX", "reviewer": "AnonReviewer3", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper extends the idea of invariant risk minimization (IRM) initially introduced by Arjovsky et al. (2019) to the setting of continual learning in which environments are observed sequentially rather than concurrently. This extension is implemented under a variational Bayesian and bilevel framework and the optimization is solved using a variant of the alternating direction method of multiplier (ADMM). The authors demonstrate the superiority of the proposed methods on variants of Colored MNIST. \n\nPros:\n+ The proposed method is a natural and practical extension of IRM and IRMG, which I believe is more applicable to many real-world scenarios. \n+ The resulting bilevel problem can be efficiently addressed using the alternating direction method of multipliers.\n\nConcerns:\n- My main concern is about the generalization theory in this continual setting. In both IRM and IRMG, the authors provide the conditions under which the OOD generalization can be guaranteed, as stated in Assumption 8 and Theorem 9 of Arjovsky et al. (2019) and in Assumption 2 and Theorem 2 of Ahuja et al. (2020). I did not see any similar theoretical guarantees in this paper. Without them, it is hard to judge whether or not  the proposed method really generalizes out-of-distribution in a sequential manner. \n- I do not think the contents in Definition 1&2 are just definitions. Instead, they should be treated as propositions, or lemma or so on, because they should be supported by proofs which seem not to be provided yet. \n- The section of Continual IRM by Approximate Bayesian Inference is not friendly with readers unfamiliar with ADMM, which makes that part hard to go through. \n- No section index\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Ef8gE5-xPiW", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "## Summary\nThe paper proposes a generalization of the invariant risk minimization objective to the continual learning setting, where environments are observed sequentially. An ADMM strategy for the solution of the resulting bilevel problem is proposed. In extensive experiments on smaller MNIST-like datasets, the method is shown to perform favorable to recent approaches for continual learning.\n\n## Explanation of Rating\nThe main strength of the paper is that it attempts to tackle an important and open problem using a reasonably principled approach. The work is well-written, and the methods performs well in practice and is evaluated against a large set of competitors. I lean towards accepting this submission. The weaknesses of the paper are the scalability of the approach (see comment #1) and the lack of theoretical guarantees for the ad hoc ADMM scheme (see comment #2). \n\n## Detailed Comments \n1. An evaluation of the method for larger architectures (e.g. a convolutional network) would make the approach more convincing. At the moment, I get the impression that there are issues with the scalability of the approach to larger data sets and models. \n2. The ADMM formulation seems rather ad-hoc, and since the loss is not convex, it is unclear whether the scheme is numerically stable / convergent. The statement about convergence rates in the strongly convex / convex setting are a bit puzzling, as they do not apply to the problem at hand. What is the main advantage of the ADMM approach over, say, solving the bilevel problem using implicit differentiation? \n\n\n\n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors address the problem of learning environment-invariant representations in the case where environments are observed sequentially.
This is done by using a variational Bayesian and bilevel framework.

The paper is borderline, with two reviewers (R2 and R3) favoring slightly acceptance and two reviewrs (R4 and R1) favoring rejection.

R4 points out that the current experiments do not do a good job of reflecting a continual learning setup and that simple modifications on existing IRM based methods could outperform the method proposed by the authors. The authors are encouraged to take into account the reviewer's
suggestions to improve the paper.

R1 argued initially that the proposed solution is not learning at all since it has errors very close to random guessing. While the authors have improved their method in the revision, the results are still close to random guessing, which questions the practical usefulness of the proposed approach. Also, in the revision, the authors managed to obtain better results when their method is combined with Environment Inference for Invariant Learning (EIIL), but these results are secondary and not the main part of the paper.

The authors should improve the work taking into account the reviewrs' comments.
Decision: Reject