C ONTINUAL I NVARIANT R ISK M INIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Empirical risk minimization can lead to poor generalization behaviour on unseen environments if the learned model does not capture invariant feature representations.
Invariant risk minimization (IRM) is a recent proposal for discovering environment-invariant representations.
It was introduced by Arjovsky et al.
and extended by Ahuja et al.
The assumption of IRM is that all environments 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 sequentially.
We show that existing approaches, including those designed for continual 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 using multiple datasets and with multiple sequential environments that the proposed methods outperforms or is competitive with prior approaches.
I NTRODUCTION
Empirical risk minimization (ERM) is the predominant principle for designing machine learning models.
In numerous application domains, however, the test data distribution can differ from the training data distribution.
For instance, at test time, the same task might be observed in a different environment.
Neural networks trained by minimizing ERM objectives over the training distribution tend to generalize poorly in these situations.
Improving generalization of learning systems has become a major research topic in recent years, with many different threads of research including, but not limited to, robust optimization (e.g., Hoffman et al.
and domain adaptation (e.g., Johansson et al.
Both of these research directions, however, have their own intrinsic limitations (Ahuja et al.
Recently, there have been proposals of approaches that learn environment-invariant 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 et al.
(2019)), which pioneered this idea, introduces a new optimization loss function to identify non-spurious causal feature-target interactions.
Invariant risk minimization games (IRMG, Ahuja et al.
(2020)) expands on IRM from a 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 learning 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 a 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.
(2020), “IRM Arjovsky et al.
(2019) requires sampling data from multiple environments simultaneously for computing a regularization term pertinent to its learning objective, where different environments are deﬁned by intervening on one or more variables of the world.”
The same applies to IRMG (Ahuja et al.
To address the problem of learning environment-invariant ML models in sequential environements, we make the following contributions:
• We expand both IRM and IRMG under a Bayesian variational framework and develop novel objectives (for the discovery of invariant models) in two scenarios: (1) the standard multienvironment scenario where the learner receives training data from all environments at the same time; and (2) the scenario where data from each environment arrives in a sequential manner.
• We demonstrate that the resulting bilevel problem objectives have an alternative formulation, which allows us to compute a solution efﬁciently using the alternating direction method of multipliers (ADMM).
• We compare our method to ERM, IRM, IRMG, and various continual learning methods (EWC, GEM, MER, VCL) on a diverse set of tasks, demonstrating comparable or superior performance in most situations.
B ACKGROUND : O FFLINE I NVARIANT R ISK M INIMIZATION
We consider a multi-environment setting where, given a set of training environments E = { e 1 , e 2 , · · · , e m } , the goal is to ﬁnd parameters θ that generalize well to unseen (test) environments.
Each environment e has an associated training data set D e and a corresponding risk R e
[Equation]
(1) where f θ = w ◦ φ is the composition of a feature extraction function φ and a classiﬁer (or regression function) w .
Empirical Risk Minimization (ERM) minimizes the average loss across all training examples, regardless of environment:
[Equation]
ERM has strong theoretical foundations in the case of iid data (Vapnik (1992)) but can fail dramatically when test environments differ signiﬁcantly from training environments.
To remove spurious features from the model, Invariant Risk Minimization (IRM, Arjovsky et al.
(2019)) instead aims to capture invariant representations φ such that the optimal classiﬁer w given φ is the same across all training environments.
This leads to the following multiple bi-level optimization problem
[Equation]
where H φ , H w are the hypothesis sets for, respectively, feature extractors and classiﬁers.
Unfortunately, solving the IRM bi-level programming problem directly is difﬁcult since solving the outer problem requires solving multiple dependent minimization problems jointly.
We can, however, relax
IRM to IRMv1 by ﬁxing a scalar classiﬁer and learning a representation φ such that the classiﬁer is “approximately locally optimal” (Arjovsky et al.
[Equation]
where w is a scalar evaluated in 1 and λ controls the strength of the penalty term on gradients on w .
Alternatively, the recently proposed Invariant Risk Minimization Games (IRMG) (Ahuja et al.
(2020)) proposes to learn an ensemble of classiﬁers with each environment controlling one component of the ensemble.
Intuitively, the environments play a game where each environment’s action is to decide its contribution to the ensemble aiming to minimize its risk.
Speciﬁcally, IRMG optimizes the following objective:
[Equation]
where ¯ w = 1 | E | (cid:80) e ∈ E w e is the average and w − e
=
(cid:80) e (cid:48)
[Table]
C ONTINUAL IRM BY A PPROXIMATE B AYESIAN I NFERENCE
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 to combine principles from IRM and continual learning.
Experience replay, that is, memorizing examples of past environments and reusing them later, could be possible in some scenarios but it is often difﬁcult to estimate a-priori the extend of replay necessary to achieve satisfactory generalization 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 justiﬁed by the property of the Kullback–Leibler (KL) divergence that promotes invariant distributions when used in sequential learning (as shown in Theorem 3).
3. V ARIATIONAL C ONTINUAL L EARNING
Following prior work in continual learning (Nguyen et al.
, let D t be the training data from the t -th environment e t , let D t 1 be the cumulative data up to the t -th environment, and let θ 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 proofs are provided in the supplementary material)
[Equation]
that is, once we have the posterior distribution p ( θ | D t − 1) at time t − 1 , we can obtain, by applying Bayes rule, the posterior p ( θ | D t 1 ) at time t up to a normalization constant.
This is achieved by multiplying the previous posterior with the current data likelihood p ( D t | θ ) .
The posterior distribution is in general not tractable and we use an approximation.
With the variational approximation, p ( θ | D t 1 ) ≈ q t ( θ ) , it is thus possible to propagate the variational distribution from one environment to the next.
From Corollary 14 (in the supplementary material) we can write the continual variational
Bayesian inference objective as
[Equation]
from the variational distribution at step q t − 1 ( θ ) , with f θ = w ◦ φ , a function with parameters θ .
3. E QUIVALENT FORMULATION OF IRM AS A B ILEVEL O PTIMIZATION P ROBLEM (BIRM)
In order to extend the IRM principle of Equation 3 using the principle of approximate Bayesian inference, by applying Lemma 5 (in supplementary material), we ﬁrst introduce the following new equivalent deﬁnition of IRM (equation 3).
Deﬁnition 1 (Bilevel IRM (BIRM)) .
Let H φ be a set of feature extractors and let
H w be the set of possible classiﬁers.
An invariant predictor w ◦ φ on a set of environments E is said to satisfy the Invariant Risk Minimization (IRM) property, if it is the solution to the following bi-level Invariant Risk Minimization (BIRM) problem
[Equation]
This formulation results from substituting the minimization conditions in the constraint set of the original IRM formulation with the Karush–Kuhn–Tucker (KKT) optimality conditions.
This new formulation allows us to introduce efﬁcient solution methods and simpliﬁes the conditions of IRM.
It also justiﬁes the IRMv1 model; indeed, when the classiﬁer is a scalar value and the equality constraint is included in the optimization cost function, we obtain Equation 4.
To solve the BIRM problem, we propose to use the Alternating Direction Method of Multipliers (ADMM) (Boyd et al.
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 in the supplementary material.
3. B ILEVEL V ARIATIONAL IRM
At this point, we cannot yet directly extend the IRM principle using variational inference.
That is because if we observe all environments at the same time, the prior of the single environment is data independent.
Therefore, we substitute q t − 1 ( θ ) from Equation 7 with priors p φ ( θ ) and p w ( ω ) , where θ and ω are now the parameters of the two functions φ and w .
We also substitute q t ( θ ) with the variational distributions q φ ( θ ) and q w ( ω ) .
Deﬁnition 2 (Bilevel Variational IRM (BVIRM)) .
Let P φ be a family of distributions over feature extractors, and let P w be a family of distributions over classiﬁers.
A variational invariant predictor on a set of environments E is said to satisfy Bilevel Variational Invariant Risk Minimization (BVIRM) if it is the solution to the following problem:
[Equation]
with
[Equation]
and
[Equation]
and where p φ and p w are the priors of the two distributions.
β is a hyper-parameter balancing the
ERM and closeness to the prior.
Deﬁnition 2 extends Deﬁnition 1 with the objective of Eq.7, where the parameters φ and w are substituted by their distributions q φ and q w .
The gradient of the cost in the inner problem is taken with respect to the distribution q w .
When we parameterize q φ with θ and q w with ω , the gradient is evaluated with respect to these parameters 1 , since the condition implies that the solution is locally optimal.
If Q ( p, q ) is convex in the ﬁrst argument, then the solution is globally optimal.
This deﬁnition extends the IRM principle to the case where we use approximate Bayes inference, shaping the variational distributions q w and q φ , to be, in expectation, invariant and optimal across environments.
3. T HE BVIRM ADMM A LGORITHM
As noted for the BIRM deﬁnition, the solution of the variational BVIRM formulation can be obtained by using ADMM (Boyd et al.
While in general there are no convergence results of ADMM methods for this problem, for local minima, under proper conditions 2 , the stochastic version of ADMM converges with rate O (1 / √ t ) for convex functions and O (log t/t ) for strongly convex functions (Ouyang et al.
We are now in the position to write the BVIRM-ADMM formulation of the BVIRM problem.
ADMM is deﬁned by the update Eq.11, where we denote with the apexes − and + the value of any variable before and after the update.
Moreover, we abbreviate as follows
[Equation]
with
[Equation]
Here, φ is ﬁxed and θ is updated in an external loop or given (e.g. the identity function).
In the experiment we use stochastic Gradient Descent (SGD) to update both the model parameters w e and the feature extractor parameters φ .
The result follows by applying Lemma 11 in the supplementary material and substituting x i ← (cid:16) w e φ (cid:17) , f i ( x i ) ← Q eφ (
ω e , θ )
and g i ( x i ) ← ∇ q ( ω ) Q ew ( ω + e , θ ) .
We provide a pseudo-code implementation leveraging Equation 11 as Algorithm 1.
One of the advantages of the ADMM formulation of BVIRM of Eq.11, is that it can be computed in parallel, where only Eq.11b requires synchronization among environments, while the other steps can be computed independently.
Algorithm 1: w, φ ←
BVIRMADMM( E, R e ) ADMM version of the
Bilevel Variational IRM Algorithm
Result: w ◦ φ : feature extraction and classiﬁer for the environment E // Randomly initialize the variables 1 ω, ω e , u e , v e , θ ← Init() ; //
Outer loop (on θ ) and
Inner loop (on ω ) 2 while not converged do // Update φ using SGD θ =
SGD θ ( (cid:80) e ∈ E Q eφ ( q w , q φ )) ; for k = 1 , . .
, K do for e ∈ E do ω
e
=
SGD ω e L ρ ( ω e , u e , ω, v e ) ; ω = 1 / | E | (cid:80) e ( ω e + u e ) ; u e =
u e + ( ω e − ω ) ; v e = v e + ∇ ω Q e ( ω e , θ ) ; end end 12 end
Algorithm 2: w, φ ←
C-BVIRMADMM( E, R e ) ADMM version of the
Bilevel Variational IRM Algorithm
Result: w ω ◦ φ − θ : feature extraction and classiﬁer for the environment E
Randomly initialize the variables 1 ω, ω e , u e , v e , θ ← Init() ; 2 ¯ ω = 0 ; 3 for e ∈ E do for k = 1 , . .
, K do θ =
SGD θ ( Q eφ ( q w , q φ )) ; while not converged do // Update ω using SGD and ADMM ω
e = SGD ω e L ρ ( ω e , u e , ω, v e ) ; ω = 1 / 2( ω e + u
e
+ ¯ ω ) ; u e =
u e + ( ω e − ω ) ; v e = v e + ∇ ω Q e
( ω e , θ ) ; end end ¯ ω = ω e ; 14 end
3.4. T HE C ONTINUAL BVIRM ADMM A LGORITHM
In presence of sequential environments, the priors for the new environment are given by the previous environment’s distributions q − φ and q − w , this is obtained by comparing the BVIRM deﬁnition in Eqs.
(9) with the continual Bayesian learning Equation (7).
In Equation 10 we thus now have Q eφ ( q w , q φ ) =
E w ∼ q w ,φ ∼ q φ R e ( w ◦ φ ) + βD KL ( q φ || q − φ )+ βD KL ( q w || q − w ) and Q ew ( q w , q φ ) =
E w ∼ q w ,φ ∼ q φ R e ( w ◦ φ ) + βD KL ( q w || q − w ) Algorithm 2 presents an example implementation of
ADMM 3 applied to the continual BVIRM formulation.
3. I NFORMATION -T HEORETIC INTERPRETATION OF C-BVIRM
[Equation]
p
[Equation]
p t ,
[Equation]
The KL divergence provides an additional motivation for the methods we propose.
Indeed, for causal discovery Peters et al.
(2015) suggests a discovery mechanism for causal variables as the intersection of the invariant conditional distributions across environments subject to interventions.
The KL divergence is asymmetric and only components present in the ﬁrst 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 sequentially.
This can be made more explicit by the property of the information projection Cover (1999)
Theorem 3 (Information Projection) .
If P and Q are two families of distributions with partially overlapping support, ∅ ⊂ supp( P ) (cid:84) supp( Q ) , and q ∈ Q , then
[Equation]
has support in the intersection for the support of P and q , or supp( p ∗ )
⊆ supp( P ) (cid:84) supp( q ) .
Therefore, if we have a sequence of sets of distributions of models from intervention environments and we compute the projection in sequence, the ﬁnal projected distribution has support on the intersection of all previous distribution families, or supp( P t ) =
(cid:84) ti =1 supp
( P i ) (see Figure 1) , since at each step p i +1 = arg min p ∈ P
i +1 D KL
( p || p i ) .
R ELATED W ORK
Generalization Domain adaptation (Ben-David et al., 2007; Johansson et al., 2019) aims to learn invariant features or components φ ( x ) that have similar P ( φ ( x )) on different (but related) domains by explicitly minimizing a distribution discrepancy measure, such as the Maximum Mean Discrepancy (MMD) (Gretton et al., 2012) or the Correlation Alignment (CORAL) (Sun & Saenko, 2016).
The above condition, however, is not sufﬁcient to guarantee successful generalization to unseen domains, even when the class-conditional distributions of all covariates changes between source and target domains (Gong et al., 2016;
Zhao et al., 2019).
Robust optimization (Hoffman et al., 2018;
Lee & Raginsky, 2018), on the other hand, minimizes the worst performance over a set of possible environments E , that is, max e ∈ E R e ( θ ) .
This approach usually poses strong constraint on the closeness between training and test distributions (Bagnell, 2005) which is often violated in practical settings (Arjovsky et al., 2019; Ahuja et al., 2020).
Incorporating the machinery of causality into learning models is a recent trend for improving generalization.
(Bengio et al., 2019) argued that causal models can adapt to sparse distributional changes quickly and proposed a meta-learning objective that optimizes for fast adaptation.
IRM, on the other hand, presents an optimization-based formulation to ﬁnd non-spurious actual causal factors 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 et al., 2020), which models the expected value of target y given each feature as a Markov decision process (MDP) and identiﬁes the spurious feature
x i if E [ y | x i ] 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 as 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 is not divided into environments, Environment Inference for Invariant Learning (EIIL) classiﬁcation 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.
(2017); De Lange et al.
(2019) addresses the problem of learning one classiﬁer that performs well across multiple tasks given in 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 a single task that is observed in different environments.
E XPERIMENTAL E VALUATION 5. D ATASETS AND E XPERIMENT S ETUP
Colored MNIST Figure 2 (left) shows a sample of train (upper) and test (lower) samples.
In each training environment, the task is to classify whether the digit is, respectively, even or odd.
As in prior work, we add noise to the preliminary label by randomly ﬂipping it with a probability
[Table]
{+of 0 . 
+}
The color of the image is deﬁned by the variable z , which is the noisy label ﬂipped with probability p c ∈ [0 . 1 , 0 . 2] .
The color of the digit is green if z is even and red if z is odd.
Each train environment contains 30 , 000 images of size 28 × 28 pixels, while the test environment contains 10 , 000 images where the probability p c = 0 .
The color of the digit ( b01 ) or the background ( b11 ) is thus generated from the label but depends on the environment.
Figure 3 depicts 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.
[Table]
relationships of colored
MNIST. Colored FashionMNIST, KMNIST, and EMNIST Figure 2 (right) shows the Fashion-MNIST dataset, where the variable z deﬁnes the background color.
Again, we add noise to the preliminary label ( y = 0 for “t-shirt”, “pullover”, “coat”, “shirt”, “bag” and y = 1 for “trouser”, “dress”, “sandal”, “sneaker”, “ankle boots”) by ﬂipping it with 25 percent probability to construct the ﬁnal label.
Besides, we also consider
Kuzushiji-MNIST dataset Clanuwat et al.
(2018) 4 and the EMNIST Letters dataset Cohen et al.
The former includes 10 symbols of Hiragana, whereas the latter contains 26 letters in the modern English alphabet.
For EMNIST, there are 62 , 400 training samples per environment and 20 , 300 test samples.
We set y = 0 for letters ‘a’, ‘c’, ‘e’, ‘g’, ‘i’, ‘k’, ‘m’, ‘o’, ‘q’, ‘s’, ‘u’, ‘v’, ‘y’ and y = 1 for remaining ones.
Reference Methods.
We compare with a set of popular reference methods in order to show the advantage of the variational Bayesian framework in learning invariant models in the sequential environment setup.
For completeness, we also evaluate the performances of four reference continual learning methods.
These include Elastic Weight Consolidation (EWC, Kirkpatrick et al.
, Gradient Episodic Memory (GEM, (Lopez-Paz & Ranzato, 2017)) 6 , Meta-Experience Replay (MER,
Riemer et al.
(2018)) 7 , and Variational Continual Learning (VCL, Swaroop et al.
; Nguyen et al.
ERM is the classical empirical risk minimization method; we always use the crossentropy loss.
IRMv1 enforces the gradient of the model with respect to a 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 to the previous task, where the relevance is measured by Fisher Information (FI); GEM uses episodic memory and computes the updates such that accuracy on previous tasks is not reduced, using gradients stored from previous tasks; MER uses an efﬁcient replay memory and employs the meta-learning gradient update to obtain a smooth adaptation among tasks; VCL and Variational Continual
Learining with coreset VCLC apply variational inference to continual learning.
C-VIRMv1 and C-VIRMG refer to, respectively, our proposed variational extensions of IRMv1 and IRGM in sequential environments.
C-BVIRM is the implementation with ADMM.
All hyper-paramter optimization strategies and simulation conﬁgurations are discussed in detail in the supplementary material. [Table]
5. R ESULTS
Table 1 lists the training and test accuracy on the MNIST dataset with the color correction b01 (see
Figure 2 left).
Since we introduced label noise by randomly ﬂipping 25 percent of the given labels, a hypothetical optimal classiﬁer would be able to achieve an accuracy of 75% in both training and test environments.
ERM, IRMv1, and IRMG perform poorly in the setup where environments are 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 features for the classiﬁcation problems at hand, here: the colors of the digits (red ∼ odd; green ∼ even), they perform poorly (much worse than a random baseline) when the spurious feature properties are inverted (green ∼ odd; red ∼ even).
In contrast, our variational extensions to both IRM and IRMG achieve a classiﬁcation accuracy higher than 45% on the test data.
This implies that our model is not relying exclusively on spurious correlations present in the color of digits.
By comparing the performance between C-VIRMv1 and C-BVIRM, we conclude that (1) our proposed bilevel invariant risk minimization framework (i.e., the BIRM in Deﬁnition 1) is an effective alternative to the original formulation Arjovsky et al.
; and (2) ADMM is effective in solving the BIRM optimization problem and has the potential to improve the generalization performance.
In addition, one can observe that the KL divergence term in VCL and our framework signiﬁcantly improves the test accuracy with respect to the baseline counterparts.
This result further justiﬁes our motivation of using a 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 sequential environments (d) (central rows) different datasets, and (c) (lower rows) the two given color correlation schemes.
We can observe that there is a general trend in the results.
IRMG and IRM, with an accuracy of less than 10% , are not able to 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 to learn more robust features and classiﬁcation models.
The higher variance of the accuracy is caused by the stochastic nature of the variational Bayesian formulation.
5. E NVIRONMENT I NFERENCE FOR C ONTINUAL INVARIANT LEARNING
In practical applications, the environmental labels are usually unavailable, which means that it is difﬁcult or impossible to manually partitioning the training set into “domains” or “environments”. [Table] [Table]
[Table]
In order to generalize our continual invariant learning models to an environment-agnostic setting, we leverage the recently proposed Environment Inference for Invariant Learning (EIIL) by Creager et al.
(2020) to automatically infer environment partitions from observational training data, and integrate
EIIL into our continual invariant learning models.
We take our proposed C-VIRMv1 as an example.
According to Table 3, it is easy to observe that inferring environments directly from observational data (using EIIL) has the potential to improve (continual) invariant learning relative to using the hand-crafted environments.
Moreover, C-VIRMv 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 4 further suggests that the generalization accuracy improves for both IRMv1 and C-VIRMv as the number of training samples increases.
Again, we observed that, when combined with EIIL,
C-VIRMv1 always outperforms IRMv1.
C ONCLUSIONS
We aim to broaden the applicability of IRM to settings where environments are observed sequentially.
We show that reference approaches fail in this scenario.
We introduce a 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 a signiﬁcant improvement in generalization capabilities
A S UPPLEMENTARY M ATERIAL A. V ARIATIONAL I NVARIANT R ISK M INIMIZATION GAMES
We now consider the IRMG objective and extend it with the variational Bayesian inference.
If we observe all environment at the same time, the prior of the single environment is data independent.
From Equation 7, we thus substitute q t − 1 ( θ ) with a priors p φ ( θ ) and q w ( ω ) , where θ and ω are now the parameters of the two functions φ and w .
While we substitute q t ( θ ) , with the variational distributions q φ ( θ ) and q w ( ω ) .
The outer problem is now
[Equation]
where ¯ w =
[Equation]
|
[Equation]
w
[Equation]
, w e ∼ q w e ( w ) is the average classiﬁer and w − e =
[Equation]
, w e (cid:48) ∼ q w e (cid:48) ( w ) is the complement classiﬁer.
In the reformulation of the IRMG model, we weight the distance of the varional distribution to the prior with β .
We notice how the difference of the variational formulation of IRMG differs on the presence of the mean on the distribution of the function over the variational distributions and the KL term.
We can now ﬁnally extend IRMG when the environments are observed sequentially.
Combining the deﬁnition of IRMG Eqs.
(5) with the continual bayesian learning Equation (7), we obtain the variational objective of IRMG in sequential environment case.
[Equation]
We can similarly extend the deﬁnition of IRMv1 when all environments are seen at the same time and sequentially.
A. M EAN F IELD P ARAMETRIZATION AND REPARAMETRIZATION TRICK
When we want to implement Equation 11 and Equation 12 and the different variation, we use the mean ﬁeld approximation and the reparametrization trick Kingma & Welling (2013).
In this case the density function of our model is parameterized by θ and ω and constraints becomes ∇ q ( ω ) Q ew ( ω + e , θ ) =
→ ∇ ω Q ew ( ω, θ ) = 0 .
If we then parametrize µ ( ω µ ) and σ ( ω σ ) the mean and standard deviation and model the distribution as q ω ( w ) = µ ( ω µ ) + (cid:15)σ ( ω σ ) , with (cid:15) ∼ N (0 , 1)
We now want to compute the gradient (in the following we ignore the dependence on the φ and its parameters)
[Equation]
The second term is
[Equation]
with
[Equation]
where we assume σ p , σ q to be diagonal, in this way the previous equation can be evaluated elementwise and where the D KL ( q || p ) is deﬁned as
[Equation]
The ﬁrst term is evaluated by Monte Carlo sampling
[Equation]
with
[Equation]
and w i ∼ N (0 , 1) .
Also in this case
[Equation]
A. T HE BIRM-ADMM A LGORITHM
We observe that to solve BIRM we can use Lemma 11 and write the following algorithm
[Equation]
where
[Equation]
We denote .
+ , .
− the values of the variable after and before the update.
In order to implement the method we use the SGD to update the model w e and in a outer loop updating for φ .
Algorithm 3: w, φ ← BIRM-ADMM( E, R e ) ADMM version of the Bilevel IRM Algorithm
Result: w ◦ φ : feature extraction and classiﬁer for the environment E // Randomly initialize the variables
[Equation]
e , u e , v e , φ ← Init() ; // Outer (on φ ) and Inner loop (on w ) 2 while not converged do // Update φ using stochastic gradient descent(SGD)
φ = SGD φ ( (cid:80)
e R e ( w ◦ φ )) ; for k = 1 , . . .
, K do for e ∈ E do w e = SGD w e L ρ ( w e , u e , w, v e ) ; w = 1 / | E | (cid:80) e ( w e + u e ) ; u e =
u e + ( w e − w ) ; v e = v e + ∇ w R e ( w e ◦ φ ) ; end end 12 end
A. V ARIATIONAL I NVARIANT R ISK M INIMIZATION
Deﬁnition 4 (VIRM) .
Give a set of distribution over the mapping P φ and a distribution over the set of classiﬁer P w , a variational invariant predictor on a set of environments E is said to satisfy
the
Variational Invariant Risk Minimization (VIRM) if it is the solution of the following problem
[Equation]
where
[Equation]
and p φ , p w are the priors of the two distributions.
A. B ILEVEL A LTERNATIVE F ORMULATION
We state here a general result on solving Bilevel Optimization Problems
Lemma 5 (Bilevel Reformulation) .
[Equation]
then we can solve the equivalent problem
[Equation]
Proof of Lemma 5 .
Lemma 5 follows by applying the Karush-Kuhn-Tucker conditions (Chapter
Boyd et al.
(2004)) to Eq.18, where the Lagrangian function is L ( x, y, u ) =
f ( x, y )+ u T g ( x, y ) .
Lemma 6 (Equivalence of Deﬁnition 1) .
Deﬁnition 1 is equivalent to Eq.
3, the Invariant Risk
Minimization.
Proof of Lemma 6 .
The result follows by apply Lemma 5 to Eq.3.
Lemma 7 (Deﬁnition 2) .
Deﬁnition 2 is the extension of Eq.
8, the Bilevel Invariant Risk Minimization, when the function is described by the distributions of their variable φ and w .
Proof of Lemma 7 .
The result follows by inspecting Eq. 8.
The equation requires the minimisation of the aggregated loss function, which is now, from Eq.7:
[Equation]
where we have separated the two contributions in φ and w , and used genetic prior distributions p φ and p w .
This is by the additive property of KL divergence:
[Equation]
since we model the two distributions independently, i.e. q φ,w = q φ q w and q φ,w = p φ p w .
Since the classiﬁers’ losses shall be minimal for all environments, this condition is substituted by requiring the gradient with respect to q w to be zero, ∀ e .
The gradient w.r.t.
q w of the second term of Eq.20 is zero.
A. T HEOREM 3 AND IRM CONNECTION A.6. S EQUENTIAL I NFORMATION P ROJECTION
In Theorem.3, we show that the Information Projection (IP) shrinks the support of the output distribution.
Lemma 8.
If we have a sequence of families of distributions P i .
Let p 1 ∈ P 1 and
[Equation]
then
[Equation]
Proof of Lemma 8.
We have that ∀ i, supp p i ⊆ supp( P i ) (cid:84) supp( P i − 1 ) , where the ﬁrst condition follows from p i ∈ P i in the minimization and the second from Theorem 3.
The results follows by iterating the property.
A.6. IRM AND I NFORMATION P ROJECTION
[Figure]
We show now two ways to state the connection of the
IRM principle and the sequential IP.
Let q − be the distribution of the previous environment and R ( q ) the loss function of the current environment, where q denotes the distribution of the network parameters.
Let
[Equation]
be the optimal distribution for the current environment.
We can then consider the Taylor expansion of the parameters distribution around the optimal distribution as
[Equation]
we can compute the new distribution as
[Equation]
and then
[Equation]
Or alternatively
[Equation] [Equation]
[Equation]
and then
[Equation]
Or more simply
[Equation]
This last equation, shows how the new distribution is the intersection of the optimal distribution at the previous step q − and the current optimal distribution q ∗ .
Fig 8 shows visually, how the new distribution is the result of projecting into two distributions q ∗ and q − .
A. O UT OF DISTRIBUTION G ENERALIZATION
The question arises if the property of generalization to out of distributions given by Theorem 9 in Arjovsky et al. (2019) also holds for BIRM and BVIRM.
Lemma 9.
If φ and w are linear functions and w ◦ φ =
Φ T w is a solution of Eq.8 it then satisﬁes
[Equation]
Proof.
Lemma 9 follows from the fact that
[Equation]
For BIRM thus Theorem 9 of Arjovsky et al.
(2019) applies directly.
A similar results holds for the the BVIRM model
Lemma 10.
If φ ∼ p φ and w ∼ p φ are linear functions and w ◦ φ = Φ T w is a solution of Eq.9, with β = 0 , it then satisﬁes
[Equation]
where ¯Φ =
E Φ ∼ p
φ [Φ] and ¯ w =
E w ∼ p w [ w ] are the mean values.
Proof.
Lemma 10 follows from the fact that
[Equation]
[Equation]
We now take the Fr ´ echet directional derivative in the η direction that is the limit of
[Equation]
which is obtained when we differentiate the distribution q w → q w + (cid:15)η .
Since
[Equation]
[Equation]
we can factorize for the direction η and obtain
[Equation]
We can now derive the Lemma by requiring δ q w R e ( w ◦ φ ) =
Theorem 9 of Arjovsky et al.
(2019) now holds when φ has rank r > 0 in expectation with respect to the invariant distribution q φ , i.e. E φ ∼ q φ rank(Φ)
= r . A. G ENERALIZED ADMM The following generalization of ADMM holds:
Lemma 11 (GADMM) .
Suppose we want to minimized
[Equation]
27) we can equivalently solve the following problem
[Equation]
28) using the following update role (scaled ADMM)
[Equation]
where the augmented Lagrangian
[Equation]
and x − i = { x j , j (cid:54) = i } is the set of all other variable, expect
the i -th.
A. C ONTINUAL V ARIATIONAL I NFERENCE
Following Nguyen et al.
we can state the following lemma.
Lemma 12 (Variational Continual Learning) .
Suppose we have a sequence of datasets
D i , i = 1 , . . .
, t drown i.i.d, then the variational estimation of the distribution q t at step t is given as projection on KL divergence
[Equation]
with Z t = (cid:82) q t − 1 ( θ ) p ( D t | θ ) dθ the normalization factor, which does not depends on q .
Proof of Lemma 12.
Let denote D t 1 = (cid:83) ti =1 D
i , from i.i.d.
p ( D t 1 ) = (cid:81) ti =1 p
( D i ) .
We are interested to maximase the a posteriori probability of the paramters give the data p ( θ | D t 1 )
[Equation]
since
[Equation]
thus
[Equation]
We now use a probability distribution which approximates the distribution at step t −
[Equation]
t
[Equation]
[Equation]
1) when then want to approximate at time t
[Equation]
This can be obtain by minimizing the KL divergence of the variational distribution q t and the distribution induced by the previous step approximation, thus
[Equation]
Lemma 13 (VCLv2) .
The minimization of the VCL deﬁned in Lemma 12, is equivalent to solve the following minimization
[Equation]
with N t i.i.d.
samples
[Equation]
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 a general loss function we can substitute the reconstruction probability with the loss function associated with a neural network parametrized by θ
[Equation]
Proof of Lemma 13.
The Lamma follows from the deﬁnition of the KL diveregnce
D KL
[Equation]
The last term does not depend on q .
Thus the result follows.
If we substitute the log of the posterior probability with a speciﬁc loss function we obtain the following Corollary.
Corollary 14 (Continual Variational Bayesian Inference) .
Given a loss function (cid:96) ( y, ˆ y ) , the variational continual learning is formulated as
[Equation]
with f θ = ( w ◦ φ ) θ
A. P ROOFS
Proof of Theorem 3.
Let ﬁrst ﬁrst recall that D KL ( p || q )
= (cid:82) p ( x ) ln p ( x ) q ( x ) dx .
If q ( x ) =
then p ( x ) =
otherwise the distance is inﬁnite.
Second if p ( x ) = 0 , then the contribution of q ( x ) is not considered since the integral is taken of the support of p , thus, since the intersection is not null and p is the result of an optimization, the support of p is the intersection of the support of q and the support of P .
A. D ATASETS 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 6 we show the dataset as generated from Ahuja et al.
In Figure 7 we show the
EMINST (letter) and KMNIST dataset.
[Figure]
[Figure] [Figure]
A. H YPER - PARAMETER S EARCH AND E XPERIMENTAL S ETUP
We performed hyper-parameter search around the suggested values from the original works and the values selected based on the best performance on the test environment.
To implement a complete comparison we used for training 1 (cid:48) 000 samples randomly drawn from each environment.
All methods were trained on the same data, using random seed reset.
We trained all method with 100 epochs on a batch size of 256 .
• IRM : γ = 91257 , threshold = 1/2 epochs, learning rate 2 .
5 e −• IRMG : warm start= 300 , termination accuracy 0 .
6 , learning rate 2 .
5 e − 4 , dropout probability 75% , weight decays = .•
ERM : learning rate 1 e − 3 , dropout probability 75% , weight decays = .•
MER :
memory size 100 ( 10% of the samples), learning rate 1 e − 3 , replay batch size = 5 , β = .
03 , γ = 1 .•
GEM : memory size 100 ( 10% of the samples), learning rate 1 e −• EWC : memory size 100 ( 10% of the samples), learning rate 1 e − 3 , regularization 0 .•
VCL,VCLC : learning rate = 5 e − 3 , corset size 100 ( 10% of the samples), • C-BVIRM, C-VIRMG, C-VIRMv1 : weight decays = .
00125 , β = 1 . , number evaluations 5 , ρ 0 = ρ 1 = 10 , step threshold =1/2 epochs, δρ = 100 , learning rate 1 e −
The neural network architecture is composed of 2 non-linear Exponential Linear Unit (ELU) activated Full Connected layers of size 100 , followed by a 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 a limited memory of samples from previous environments.
A. H YPER - PARAMETERS OF ENVIRONMENT INFERENCE FOR CONTINUAL INVARIANT LEARNING
We list below the values of hyper-parameters in EIIL for continual invariant learning: 2 layers,
390 hidden neurons, 501 epochs, l2 regularizer weight
learning rate
numbber of runs
, number
of EIIL
iterations
(cid:48) 000 , number
Monte Carlo evaluations
: 3 , penalty anneal iterations: 190 , penalty weight: 191257 .
18613115902 , prior weight: 1 e −
A.13. S YNTHETIC D ATASET
The Synthetic Dataset is described in (Arjovsky et al.
for testing IRM
and it is deﬁned by a
Structural Causal Model (Pearl (2009)), where a variable y ∈ R N is generated by x 1 ∈ R N , while x 2 ∈ R N is generated by y .
The observed variable is x = ( x 1 , x 2 ) .
The structural equations are
[Equation]
[Equation]
[Equation]
(34) with σ 1 ﬁxed and σ e dependent on the environment.
We compared with ERM, IRM (Arjovsky et al.
, IPC (Invariant Prediction), which is the method proposed in Peters et al. (2015), and EIIL (
[Table]