Continual Invariant Risk Minimization

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 . ( 2019 ) and extended by Ahuja et al . ( 2020 ) . 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 . 1 INTRODUCTION . 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 . ( 2018 ) ) and domain adaptation ( e.g. , Johansson et al . ( 2019 ) ) . Both of these research directions , however , have their own intrinsic limitations ( Ahuja et al . ( 2020 ) ) . 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 defined by intervening on one or more variables of the world. ” The same applies to IRMG ( Ahuja et al . ( 2020 ) ) 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 efficiently 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 . 2 BACKGROUND : OFFLINE INVARIANT RISK MINIMIZATION . We consider a multi-environment setting where , given a set of training environments E = { e1 , e2 , · · · , em } , the goal is to find parameters θ that generalize well to unseen ( test ) environments . Each environment e has an associated training data set De and a corresponding risk Re Re ( w ◦ φ ) .= E ( x , y ) ∼De ` e ( ( w ◦ φ ) ( x ) , y ) , ( 1 ) where fθ = w ◦φ is the composition of 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 ( x , y ) ∼∪e∈EDe ` ( fθ ( x ) , y ) . ( 2 ) ERM has strong theoretical foundations in the case of iid data ( Vapnik ( 1992 ) ) but can fail dramatically when test environments differ significantly 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 classifier w given φ is the same across all training environments . This leads to the following multiple bi-level optimization problem min φ∈Hφ , w∈Hw ∑ e∈E Re ( w ◦ φ ) s.t . w ∈ arg min we∈Hw Re ( we ◦ φ ) , ∀e ∈ E , ( 3 ) where Hφ , Hw are the hypothesis sets for , respectively , feature extractors and classifiers . Unfortunately , 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 a representation φ such that the classifier is “ approximately locally optimal ” ( Arjovsky et al . ( 2019 ) ) min φ∈Hφ ∑ e∈E Re ( φ ) + λ||∇w|w=1.0Re ( wφ ) ||2 , ∀e ∈ E , ( 4 ) 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 classifiers 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 . Specifically , IRMG optimizes the following objective : min φ∈Hφ ∑ e∈E Re ( w̄ ◦ φ ) s.t . we = arg min w∈Hw Re ( 1 |E| ( w + w−e ) ◦ φ ) , ∀e ∈ E , ( 5 ) where w̄ = 1|E| ∑ e∈E we is the average and w−e = ∑ e′∈E , e′ 6=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 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 difficult 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 justified by the property of the Kullback–Leibler ( KL ) divergence that promotes invariant distributions when used in sequential learning ( as shown in Theorem 3 ) . 3.1 VARIATIONAL CONTINUAL LEARNING . Following prior work in continual learning ( Nguyen et al . ( 2018 ) ) , let Dt be the training data from the t-th environment et , let Dt1 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 ) p ( θ|Dt1 ) ∝ p ( θ|Dt−11 ) p ( Dt|θ ) , ( 6 ) that is , once we have the posterior distribution p ( θ|Dt−11 ) at time t − 1 , we can obtain , by applying Bayes rule , the posterior p ( θ|Dt1 ) at time t up to a normalization constant . This is achieved by multiplying the previous posterior with the current data likelihood p ( Dt|θ ) . The posterior distribution is in general not tractable and we use an approximation . With the variational approximation , p ( θ|Dt1 ) ≈ qt ( θ ) , 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 qt ( θ ) = arg min q ( θ ) E ( x , y ) ∼DtEθ∼q ( θ ) { ` ( y , fθ ( x ) ) } +DKL ( q ( θ ) ||qt−1 ( θ ) ) , ( 7 ) from the variational distribution at step qt−1 ( θ ) , with fθ = w ◦ φ , a function with parameters θ . 3.2 EQUIVALENT FORMULATION OF IRM AS A BILEVEL OPTIMIZATION PROBLEM ( 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 first introduce the following new equivalent definition of IRM ( equation 3 ) . Definition 1 ( Bilevel IRM ( BIRM ) ) . Let Hφ be a set of feature extractors and let Hw be the set of possible classifiers . 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 min φ∈Hφ , w∈Hw ∑ e∈E Re ( w ◦ φ ) ( 8a ) s.t . ∇wRe ( w ◦ φ ) = 0 , ∀e ∈ E. ( 8b ) 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 efficient solution methods and simplifies the conditions of IRM . It 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 4 . To solve the BIRM problem , we propose to use the Alternating Direction Method of Multipliers ( ADMM ) ( Boyd et al . ( 2011 ) ) . 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 .

Summary: In 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.