On Invariance Penalties for Risk Minimization

1 INTRODUCTION . Under the learning paradigm of Empirical Risk Minimization ( ERM ) ( Vapnik , 1992 ) , data is assumed to consist of independent and identically distributed ( iid ) samples from an underlying generating distribution . As the data generating distribution is often unknown in practice , ERM seeks predictors with minimal average training error ( i.e. , empirical risk ) over the training set . However , shuffling and treating data as iid risks possibly losing important information about the underlying conditions of the data generating process . Despite becoming a ubiquitous paradigm in machine learning , a growing body of literature ( Arjovsky et al. , 2019 ; Teney et al. , 2020 ) has revealed that ERM and the the common practice of shuffling data inadvertently results in capturing all correlations found in the training data , whether spurious or causal , and produces models that fail to generalize to test data . The potential variation of experimental conditions that can exist at training time and during deployment in real-world applications , manifests in discrepancies between training and testing distributions . This , in turn , highlights the need for machine learning algorithms to generalize out-of-distribution ( OoD ) . Shuffling and treating data as iid risks possibly losing important information about the underlying conditions of the data generating process . As will be demonstrated in this work , partitioning training data into environments , e.g. , based on the conditions under which data is generated , can exploit these differences to enhance generalization . One promising approach based on this observation is that of Arjovsky et al . ( 2019 ) , in which the principle of Invariant Risk Minimization ( IRM ) is introduced . The objective of IRM is to find a predictor that is invariant across all training environments ( see Definition 1 and Equation 1 ) . Because of the conceptually appealing nature of IRM and its potential to address the OoD-generalization problem , there is a stream of literature scrutinizing various facets of the original framework , e.g. , extensions to other settings including online learning ( Javed et al. , 2020 ) and treatment effect estimation ( Shi et al. , 2020 ) , fairness ( Adragna et al. , 2020 ) , introducing game-theoretic interpretations ( Ahuja et al. , 2020 ) , and raising concerns on the drawbacks and limitations of current IRM implementations ( Rosenfeld et al. , 2021 ; Kamath et al. , 2021 ) . For an indepth overview of the broader generalization literature , we refer the interested reader to ( Arjovsky , 2020 ) and the references therein , and for an empirical evaluation of the performance of a number of the state-of-the-art methods on various test cases , we refer the reader to ( Gulrajani & Lopez-Paz , 2020 ) . In this paper , we introduce two practical implementations of the IRM principal to increase its appeal and applicability in real-world settings . First , we propose an invariance penalty that is directly related to risk . More precisely , we show that the risk in each environment under an arbitrary classifier equals to the risk under the optimal classifier for that environment plus the newly proposed invariance penalty between the said classifier and the optimal one . Second , we build on these initial findings to provide practitioners who currently use the original IRMv1 implementation with an adaptive rule by which to choose the penalty coefficient appropriately . In doing so , we characterize the difference between our proposed invariance penalty and the one proposed by Arjovsky et al . ( 2019 ) in terms of the eigenvalues of the Gram matrix of the data representation . This eigenstructure plays a significant role in the failure of invariance penalties including the one proposed by Arjovsky et al . ( 2019 ) . In addition to providing practitioners with two valid approaches for implementing IRM , this work serves to illustrate the importance of the eigenstructure of the Gram matrix of the data representation for IRM . In particular , we revisit the counterexample of Rosenfeld et al . ( 2021 ) where the invariance penalty of Arjovsky et al . ( 2019 ) can be made arbitrarily small for a non-invariant representation . We show that the Gram matrix is ill-conditioned in such cases . We then provide a practical solution to alleviate such behavior , in particular for the case where the representation is parameterized by a neural network . Moreover , we show that the proposed framework finds an invariant predictor for the setting in which the data is generated according to a linear Structural Equation Model ( SEM ) when provided a sufficient number of training environments under a mild non-degeneracy condition , which is similar in nature to the ones considered in ( Arjovsky et al. , 2019 ; Rosenfeld et al. , 2021 ) . Finally , we evaluate our method on various test cases including InvarianceUnitTests ( Aubin et al. , 2020 ) and HealthyGutTests . InvarianceUnitTests is a test bed with three synthetically generated data sets capturing different structures of spurious correlations . HealthyGutTests is a curated collection of biomedical data sets based on a prior meta-analysis of various microbiome studies ( Gupta et al. , 2020 ) in which the relationships between human gut microbiome composition and various disease phenotypes can be investigated . The remainder of the paper is organized as follows . In Section 2 , we formally define the notion of invariant prediction , the invariant risk minimization principle , and its relaxation proposed by Arjovsky et al . ( 2019 ) . In Sections 3 and 4 , we introduce our more practical implementation and the rationale for its design . In Section 5 , we evaluate the efficacy of our proposed model and compare it with other variations of IRM over a series of experiments . We conclude the paper in Section 6 . All mathematical proofs are presented in the Appendix . 2 BACKGROUND : INVARIANT PREDICTION . In this paper , we consider data ( Xe , Y e ) collected from multiple training environments e ∈ Etr where the distribution of ( Xi , Y i ) and ( Xj , Y j ) may be different for i ̸= j with i , j ∈ Etr . We denote by Re the risk under environment e. That is , for predictor f : X → Y , and loss function ℓ : Y × Y → R , the risk under environment e is defined as Re ( f ) = EXe , Y e [ ℓ ( f ( Xe ) , Y e ) ] . 2.1 INVARIANT RISK MINIMIZATION . Arjovsky et al . ( 2019 ) define the notion of invariant predictors under a multi-environment setting as follows . Definition 1 ( Invariant Predictor ) . A data representation φ : X → H is said to elicit an invariant predictor w ◦ φ across environments E if there exists a classifier w : H → Y , which is optimal for all environments , i.e. , w ∈ argminw̃ : H→Y Re ( w̃ ◦ φ ) for all e ∈ E . To find such invariant predictors , Arjovsky et al . ( 2019 ) introduce the notion of the Invariant Risk Minimization ( IRM ) principle : min φ : X→H w : H→Y ∑ e∈Etr Re ( w ◦ φ ) subject to w ∈ argmin w̃ : H→Y Re ( w̃ ◦ φ ) , ∀e ∈ Etr . ( 1 ) As this bi-leveled optimization problem is rather intractable , Arjovsky et al . ( 2019 ) propose a practical implementation of IRM by relaxing the invariance constraint ( which itself requires solving an optimization problem ) to an invariance penalty . We review its derivation in what follows . 2.2 IRMV1 : A RELAXATION OF IRM . In order to provide an implementation of IRM , Arjovsky et al . ( 2019 ) restrict the classifier w to linear functions , i.e. , min φ : X→H w∈Rdφ ∑ e∈Etr Re ( w⊤φ ) subject to w ∈ argmin w̃∈Rdφ Re ( w̃⊤φ ) , ∀e ∈ Etr . ( 2 ) To motivate their proposed penalty , Arjovsky et al . ( 2019 ) first consider the squared loss , i.e. , ℓ ( f ( x ) , y ) = ∥f ( x ) − y∥2 where ∥ · ∥ denotes the Euclidean norm . Define matrix Ie ( φ ) as Ie ( φ ) : = EXe [ φ ( Xe ) φ ( Xe ) ⊤ ] . ( 3 ) Assuming that Ie ( φ ) is full rank for a fixed φ , its respective optimal classifier is unique , i.e. , argminw̃∈RdφRe ( w̃⊤φ ) = w⋆e ( φ ) where w⋆e ( φ ) : = Ie ( φ ) −1EXe , Y e [ φ ( Xe ) Y e ] . ( 4 ) Hence , in this setting , the invariant constraint of IRM in equation 2 can be simplified to w = w⋆e ( φ ) for all e ∈ Etr . A natural relaxation of the constraint w−w⋆e ( φ ) = 0 to a penalty is ∥w−w⋆e ( φ ) ∥2.1 However , Arjovsky et al . ( 2019 ) show that this penalty may not capture invariance by constructing an example for which ∥w − w⋆e ( φ ) ∥2 is not well-behaved ( see Section 4.3 for more details ) . They argue that undoing the matrix inversion Ie ( φ ) −1 , which appears in the computation of w⋆e ( φ ) could improve the behavior of the invariance penalty . That is , considering ∥Ie ( φ ) ( w − w⋆e ( φ ) ) ∥2 as the invariance penalty . Moreover , for the squared loss , one can show that ∥Ie ( φ ) ( w − w⋆e ( φ ) ) ∥ 2 = ( 1/4 ) ∥∥∇wRe ( w⊤φ ) ∥∥2 . ( 5 ) Hence , Arjovsky et al . ( 2019 ) propose the following invariance penalty ρIRMv1e ( φ , w ) : = ∥∥∇wRe ( w⊤φ ) ∥∥2 . ( 6 ) Using the penalty equation 6 , the relaxation of IRM is given by min φ , w ∑ e∈Etr Re ( w ⊤φ ) + λρIRMv1e ( φ , w ) , ( 7 ) where λ ≥ 0 is the penalty coefficient . Notice that for a given w and φ , the predictor w ◦ φ can be expressed using different classifiers and data representations , i.e. , w◦φ = w̃◦φ̃where w̃ = w◦ψ−1 and φ̃ = ψ ◦ φ for some invertible mapping ψ : H → H. Hence , in principle , it is possible to fix w without loss of generality . By relying on this observation , Arjovsky et al . ( 2019 ) fix the classifier as a scalar w = 1 , and , thus , search for an invariant data representation of the form φ ∈ R1×dx . Their relaxation of IRM , which they refer to by IRMv1 is given by min φ ∑ e∈Etr Re ( φ ) + λρ IRMv1 e ( φ , 1.0 ) . ( IRMv1 ) Although equation 5 only holds for squared loss , Arjovsky et al . ( 2019 , Theorem 4 ) justify the choice of ∥∇w|w=1.0Re ( w⊤φ ) ∥2 as an invariance penalty for other loss functions , e.g. , cross-entropy loss . More precisely , let Φ be the matrix that parameterizes the data representation . They show that for all convex differentiable loss functions , ( w⊤Φ ) ⊤∇wR ( w⊤Φ ) = 0 if and only if w is optimal for all environments . Although the penalty ρIRMv1e ( φ , 1.0 ) seems to be an appropriate invariance penalty , which is easily implementable , its effectiveness to capture invariance has been brought into question . In particular , Rosenfeld et al . ( 2021 ) construct a non-invariant data representation such that the penalty ρIRMv1e ( φ , 1.0 ) is arbitrarily small . In what follows , we propose an alternative invariance penalty that is directly comparable to risk . Hence , if the invariance penalty of a classifier is small , then so is the difference between its risk and the risk under the optimal classifier . 1In the presence of known exogenous ( environment dependent ) variables , one can utilize anchor regression Rothenhäusler et al . ( 2021 ) that is conceptually related to the invariance penalization . Instead of an invariance penalty , anchor regression relies on the projection onto the span of the said variables .

Summary: In this paper, the authors propose an alternative Invariant Risk Minimization (IRM) penalty for the particular case of a linear output layer trained with a mean-squared error loss. In this particular case, the original formulation can be understood as penalizing the squared Euclidean norm of the output layer's gradient, that is,  $$\rho_{e}^{\mathrm{IRMv1}}(\Phi, w) := \vert\vert \Phi_{e}^{T}(\Phi_{e} w - y_{e}) \vert\vert_{2}^{2} = \vert\vert \Phi_{e}^{T}\Phi_{e}(w -w_{e}^{*})\vert\vert_{2}^{2},$$  where $\Phi_{e} \in \mathbb{R}^{n_{e} \times d}$  is a matrix with the (learnt) data representation for environment $e$, $y_{e}$ the corresponding targets and $w_{e}^{*} = (\Phi_{e}^{T}\Phi_{e})^{-1}\Phi_{e}^{T}y_{e}$ the optimal, unregularized linear regressor. Witth this in mind, this paper proposes to use instead:  $$\rho_{e}^{\mathrm{IRMv2}}(\Phi, w) := \vert\vert (\Phi_{e}^{T}\Phi_{e})^{1/2}(w -w_{e}^{*})\vert\vert_{2}^{2},$$  which interestingly can be trivially shown to lead to an optimal classifier $w^{*}$ identical to that of the ERM formulation. Crucially, however, the authors propose to explicitly backpropagate through the dependence of this classifier on the data representations $\Phi_{e}$ when minimizing a cost function that incorporates $\rho_{e}^{\mathrm{IRMv2}}(\Phi, w)$.  The authors prove that, under similar non-degeneracy conditions as those assumed when analyzing other IRM penalties, $\rho_{e}^{\mathrm{IRMv2}}(\Phi, w)$ will either find an invariant data representation (or a trivial classifier) for a linear SEM as in Rosenfeld et al. 2021. Furthermore, the authors argue that there might be a connection between ill-conditioning of the Gram matrix $\Phi_{e}^{T}\Phi_{e}$ and failure modes of IRM penalties and propose a heuristic rule to adapt the penalty weight in the original IRM formulation.