Weighted Risk Invariance: Domain Generalization under Invariant Feature Shift
Authors: Gina Wong
Joshua Gleason
Rama Chellappa
Yoav Wald
Anqi Liu
Venue: Accepted by TMLR ()

Abstract: Learning models whose predictions are invariant under multiple environments is a promising approach for out-of-distribution generalization. Such models are trained to extract features $X_{\text{inv}}$ where the conditional distribution $Y \mid X_{\text{inv}}$ of the label given the extracted features does not change across environments. Invariant models are also supposed to generalize to shifts in the marginal distribution $p(X_{\text{inv}})$ of the extracted features $X_{\text{inv}}$, a type of shift we call an invariant covariate shift. However, we show that proposed methods for learning invariant models underperform under invariant covariate shift, either failing to learn invariant models---even for data generated from simple and well-studied linear-Gaussian models---or having poor finite-sample performance. To alleviate these problems, we propose weighted risk invariance (WRI). Our framework is based on imposing invariance of the loss across environments subject to appropriate reweightings of the training examples. We show that WRI provably learns invariant models, i.e. discards spurious correlations, in linear-Gaussian settings. We propose a practical algorithm to implement WRI by learning the density $p(X_{\text{inv}})$ and the model parameters simultaneously, and we demonstrate empirically that WRI outperforms previous invariant learning methods under invariant covariate shift.

Published in Transactions on Machine Learning Research (07/2024)
Weighted Risk Invariance: Domain Generalization under Invariant Feature Shift
Gina Wong
guong150jhu. edu
Deportment of Computer Science Johns Hopkins Universily
Joshua Gleason
gleason Gound. edu
Department of Electrical and Computer Engineering University of Moryland, College Park
Rama Chellappa
rchella4 @jhu edu
Department of Electrical and Computer Engineering Johns Hopkins University
Yoav Wald
your. wald edu
Center for Data Science New York Universily
Anqi Liu
abiu. cs @iuuu. edu
Department of Computer Science Johns Hopkins University
Reviewed OpenReview: pttps://openrevieee net/ forum? id= WyPKLWPPSr
Abstract
Learning models whose predictions are invariant under multiple environments is a promising approach for cittof-distribution generalization. Such models are trained to extract features Xinv whore the conditional distribution Y Xinv the label given the extracted fcatures does not change across environments. Invariant models are also supposed to generalize to shifts in the marginal distribution p(Xinv) of the extracted features Xiav, type of shift call an invariant covoriate shift. However, show that proposed methods for learning invariant models underperform under invariant covariate shift, either failing to learn invariant
even for data generated from simple and well-studied linear-Gaussian models or having poor finite-sample performance. To alleviate these problems, we propose weighted risk invariance (WRI). Our framework is based on imposing invariance the loss across environments subject to appropriate reweightings of the training examples. We show that WRI provably learns invariant models, i.c. discards spurious correlations, linear- Gaussian settings. We propose practical algorithm implement WRI by learning the density p(Xinv) and the model parameters simultaneously, and we demonstrate empirically that WRT outperforms previous invariant learning methods under invariant covariate shift.
1 Introduction
Although traditional machine learning methods can be incredibly effective, many are based on the i. assumption that the training and test. data are independent and identically distributed. As a result, these methods are often brittle to distribution shift, failing to generalize training performance to out-of distribbtttt (OOD) data. (oorallaa Efros 2011 Beery et al. 2018 Hendrycks & Dietterich 2019 Goirhos et al. 2020). Distribution shifts abound in the real world: as general examples, we encounter them when we collect test
data under different conditions than we collect the training data (Adini et al. 1997 Huang et. al. 2006

---

Published in Transactions on Machine Learning Research (07/2024)
Data cistributions
Environment 1 (train)
Environment (train)
Environment 3 (test)
pi(z)
pa(z)
Pa(e)
p(v|V - 0)
p6r|V 13
-5-
-5
-5
0
Zicn
Jiee
Zion
ERM
IRM
VREx
WRI
1.0
5
0
0-
(1)
0.4
-5
-5 0
5
0
Figure 1: Top row: Gaussian setup exhibiting both heteroskedasticcty and invariant covariate shift. Note that same across environments, while p(y]= changes across environments,
for invariant features Fine: (ine) following causal model. Bottom row: visualize the performance of different algorithms under this setup, and find the WRI objective recovers a more invariant predictor than the other algorithms. Specifically, the WRI predictor bases its predictions more on the invariant features Zasv and less the spurious features Zspu. so it learns decision boundary that is vertical in this case. The full parameters for this simulation can be found in Appendix
when we train on synthetic data and evaluate on real data (Chen et al. 2020 Beery et al. 2020 Learning to generalize under distribution shift, then, is a critical step loward practical deployment.
To achieve generalization, a typical setup is to use multiple training environments with the hope that, if the model learns what is invariant across the training environments, it can leverage that invariance on unseen test environment as well. An extensive line of research therefore focuses learning statistical relationships that are invariant across training environments (Ganin et al. 2016 Tzeng et al. 2021). Of these, many works focus on learning invariant relationship between the data X and labels Y. as such relationships are posited to result from invariant causal mechanisms (Pearl 1995 Schölkopf et al. 2012 Peters et al.] 2016 Rooas-Carulla et al.] 2018). Applying this idea to representation learning, several works have found that learning features ((X) such that the label conditional distribution pXY) is invariant across training environments is an elfective method for improving generalization to unobserved test
2017 Li et al.
environments (Arjovsky et al. 2019 Krueger al. 2021 Wald al. 2021 Eastwood et al. 2022).
Predictors that rely conditionally invariant features be of OOD generalization (Koyama & Yamaguchi 2020). However, demonstrate that existing methods for finding these predictors can struggle or even fail when the conditionally invariant features undergo covariate shift, a type of shift we call invoriant covariate shiftt. We show that invariant covariate shift introduces two types of challenges: first, it leads risk invariance objectives (such as Krueger et al. (202T)) to incorrectly identify the conditionally invariant features in a heteroskedastic sctting Jand second, it adversely impacts the sample complexity (Shimodaira) (200)) of recovering an optimal invariant predictor (in the style of Arjovsky
optimal solution to the problem
et al. (2019)). For illustration of these failure cases, the learned predictors for VREx and IRM
respectively in Figure 1
In this work, we demonstrate how to mitigate the effects of invariant covariate shift by accounting for the underlying structure the data in our learning process. We propose Weighted Risk Invariance (WRI), where we enforce invariance of a model's reweighted losses across training environments. We show that under appropriate choices of weights, models satisfying WRI provably achieve OOD generalization by learning to reject spurious correlations, under common model for learning under spurious correlations (Arjovsky)
That is, where classification is more difficult some instances than heres..
2

---

Published in Transactions on Machine Learning Research (07 (024))
al. 2019 Rosenfeld et al. 2020 Wald et al. 2021p. To learn models that satisfy WRI in practice, we introduce method that alternates between estimation of reweightings (corresponding to densities in our implementation) and optimization of prediction loss penalized by an empirical estimate of WRI. We show that solving for WRI recovers invariant predictors, even in the difficult setting of hetcroskedasticity and invariant covariate shift, where popular previously proposed invariant learning principles fail (Figure 1 Finally, we show that under invariant covariate shift, the WRI solution outperforms widely used basclines.
In summary, our main contributions are:
We introduce an invariant learning method that uses weighted risk invariance to tackle invariant covariate shift. We formally show that enforcing weighted risk invariance recovers an invariant predictor under a linear causal setting (Prop. Thm.
We propose WRI, an optimization problem for enforcing weighted risk invariance, and provide an algo- rithm to solve WRI in practice via alternating minimization ($3.4).
We verify the efficacy our approach with experiments on simulated and real-world data. We find that, under invariant covariate shift, our method outperforms widely used basclines on benchmarks like ColOTedMNIST and the DomainBed datasels. Moreover, our learned densities report when conditionally invariant features in training, making them useful for downstream tasks like OOD detection ($4
2 Preliminaries
2.1 Domain generalization
the problem of learning some invariant relationship from multiple training domains /enirronments Err = (ei,...,ex] that we assume to hold across the set all possible environments we may encounter E. Each training environment E e consists of a dataset De {(xf, -= and we assume that data pairs (xf,yf) are sampled i.dd from distributions Pe(X,, Y²) with x E X and y y
Domain generalization was first posed (Blanchard et al. 2011 Muandet et al. 2013)
For loss function l P y I, we define the statistical risk of predictor f: X 5 over environment
RT) Ex re) [((f(x)),
(1)
and its empirical realization
De Σ ()(())))))
(2)
Empirical Risk Minimization (ERM) (aapnik 1991 aims to minimize the average loss over all of the training samples
RERM(F)
E
R().)
(3)
Exr BEE.
Unfortunately, ERM fails to capture distribution shifts across training environments (Arjovsky et al. 2019 Krueger et al. 2021) and so can fail catastrophically depending on how those distribution shifts extend the test data. (Previous works point to cases where ERM successfully generalizes, but the reason is not well understood (Vedantam al.) 2021 Gulrajani & Lopez-Paz, 2020), and its success does not hold across all types of distribution shifts (Ye al. 2022 Wiles et al 2021 In our work, therefore aim to learn predictor that captures the underlying causal mechanisms instead, under the assumption that such mechanisms are invariant across the training and test environments.
2.2 Our causal model
We assume that input X contains both iinsaa/iivaiantt and spurious components such that there exists a deterministic mechanism to predict the label from invariant. components Xine, while the relationship
When the context clear, we write the data-generating random vector X, output random variable and their realizations without the environment superscript.
3

---

Published in Transactions on Machine Learning Research (07/202) 2024)
E
Xnnv
Xspa
X
Figure 2: Causal graph depicting our data data-generating process. The environment E is dashed to emphasize it takes on unobserved values test time.
belween spurious components Xspn and label Y depends on the environment. More explicitly, say that Y I E Xspu. The observed features X are function of the invariant and spurious components, or X = X which we assume is injective so that there exists g (-) that recovers Xapu: almost surely. We represent this data-generating process with the causal graph in Figure 2 From the graph, we observe that PY Xine) = Pa(Y Xiw) for any e, E &, i.e. is fixed across environments. On the other hand, the distribution (Xppu Y) is not fixed across various We summarize these assumptions below.
and we assume it may change arbitrarily.
Assumption (Causal Prediction Coupling). A causal prediction coupling set of distributions P {pele: where each Pa P is distribution over real vectors X (Xinv: Xepu) and target Y that decomposes according to
Pe (x = (P X (X Y).
We call Xinv and Xspu the invariant and spurious components accordingly. The distribution p(Y Xinv) is fixed across all members of P.
Intuitively, we seek robust predictor f where (f) does not grow excessively high for any Pe P. Let us focus on the square and logistic losses, so an optimal model with respect to Pe predicts Y according to its conditional expectation [Y X]. Focusing on couplings where Pe(X YY) can undergo large changes, we observe that guarantee on the worst case loss can only be obtained when f(x) does not depend on Xspu- This can be shown formally for the case of linear regression, by showing that any model f(x) = x where Wepu o 0 can incur an arbitrarily high loss under some Pe P (Arjovsky et al. 2019 Rosenfeld et al. 2020). Thus, we wish to learn predictor that satisfies f(X) ((Xinv) inv) almost everywhere.
Restricted to such predictors, we observe two things. First, there is a single model that is simultaneously optimal for all P, and this model the predictor we would like to learn. Second, the worst-case loss this predictor is controlled by the loss of the most difficult instances terms of Xinvs since Pe(Xin)) can shift mass to these instances where the predictor incurs the highest loss. Following these arguments, we arrive definitions of a spurious-Iree predictor and an invariant predictor.
Definition 1. measurable function X y a spurious-free predictor with respect to a causal prediction coupling P if f(X) f(Xinv) almost everywhere. Considering the square logisitic losses, we call the optimal spurious- spurious-free predictor f*(x) =B]y Xinv Xin]] the invariant predictor for short.
Invariant predictors, or representations, are usually defined models that satisfy Ep [Y [(X)] Ep [Y] [(X) for any e, E g' (Arjovsky et al. 2019 3 For the settings we study in this paper, the
invariant predictor as defined in Detmiiion satisties this property.
2.3 Invariant learning under covariate shift
Causal predictors assume there exist features Xinv = ((X) such that the mechanism that generates the target variables Y from Φ(X) is invariant under distribution shift; in other words, the conditional distribution p(Y)((X)) is invariant across environments (Schölkopf et al. 2012: Peters et al. 2016). This is similar to covariate shift, where the conditional distribution p(Y X) is assumed to be the same between training
and test while the training and test distributions (X) and (X) differ. However, while the covariato
Demanding equality in distribution and not just expectation is also common.
4

---

Published in Transactions on Machine Learning Research (07/2024)
shift assumption refers to distribution shift in the covariates X, study the case of covariate shift in the conditionally invariant features Xinv, which we call invoriant covariate shift. For simplicity, we will also call conditionally invariant features simply invariant when the context is clear, using terminology consistent with previous works (Liu et al., 2021a Ahuja et al. 2021
As a motivating example, consider the case where we have patient demographics X and diagnostic infor- mation Y from multiple hospitals. Some demographic features, like whether a patient smokes or not, may have direct influence diagnosis like lung cancer. We can model these features invariant features Xine- Under invariant covariate shift, although the distribution of smokers p(Xine) can differ from hospital to hospital, the relationship ine) remains the same. Unlike regular covariate shift, which applics to the entire covariate, the assumption here is more targeted: under regular covariate shift, we assume that the conditional distribution given all the patient demographics pX) is the same across hospitals. In contrast, the assumption that only mo)) is the same across hospitals is more realistic.
Several popular works search for proxy forms of invariance in an attempl to recover the invariant features from complex, high-dimensional data. Methods along this line include Risk Extrapolation (REx) (Krueger et al. 2021) and Invariant Risk Minimization (IRM) (Arjovsky et al. 2019). Unfortunately, these methods are adversely impacted by invariant covariate shift, as we demonstrate in $3.3 We will see there are two types of challenges, namely: (1) the failure to learn an invariant predictor, even given infinite samples- this is where methods like REx fail, even though they have favorable performance when there is no invariant covariate shift; and (2) finite sample performance, where methods like IRM underperform when there is invariant covariate shift, even though they are guaranteed to learn an invariant predictor as the amount of data grows to infinity.
A classical approach to counteract the effects of covariate shift is importance weighting (Shimodaira) 2000). where the training loss is reweighted with a ratio of test to training densities so that training better reflects the test distribution. Generalizing this concept, Martin et al. (2023) finds there are many choices of weights that can help aligu the training and test distributions; that is, there exist many choices of weights a(x) and ((x) such that
(X.Y) [x(x){[f(x) EE. (x))
(4)
Note that this equation reduces to importance weighting under density ratio weighting, or when B(x) = (x))x/pp and axx) = 1. Under this more general framework, density weighting, where ((x) = (x)
and (xx) = (x), achieves a similar effect.
Our work studies the question of how lo adapt reweighting to the case of invariant learning. Unfortunately, a naive choice of weights, where the covariates of one environment are reweighted to match the distribution of the other (i.e. importance weighting), will not provide the desired result in the scenario of invariant learning [ In what follows, we show that for reweighting to be effective in this setting, it needs to explicitly account for the density of Xinv-
3 Dware generalization
3.1 Invariant feature density weighting
Consider choosing some set of weights (xxin)), (Xim)) such that a(xine Pe, (Xivv) = (xinv)) Pay (Xim)) Martin et al. 2023). This defines a pseudo-distrriuttion g(Xivv) & o(Xinv)Pe, (Xivv) over the invariant features. It turns out that for spurious-fres models fx), even when we reweigh the entire joint distribution Pe, (X, Y) instead of the marginal distribution over Xnnv alone, the expected loss will still be equal to that over g(X (up to a normalizing constant). This also holds when we do the rewcighting on another environment, which leads to the following result proved in the appendix.
Proposition Let Assumption ! hold over a set of environments E C If predictor f is spurious- free over E, then for every pair of enwironments ei, eg E E, their weighted risks are equal if their respective weighting functions a and 3 obey a(xinw) (Xinw) 3(xinn))Pe (Xiw)).
*This occurs because invariant model the original distribution may become not-invariant under the weighted data, where distributions of covariates are matched across environments. It. is easy to see this with simple examples.
5

---

Published in Transactions on Machine Learning Research (07/2024)
This result shows that this form of weighted risk invariance is signature of spurious-free prediction, and it is easy to verify that the optimal invariant predictor also satisfies this invariance. We formalize this notion of invariance with the following definition.
Definition We define weighted risk invariance between two environments es and ey to mean that
(X [a
)((((x), v) (((())))(((())))))))))))))))) Byc, (x",,")) [8(x (((x),))
where a(Xinv Pe, (Xinv) = ((x Pe, (Xiv).. Weighted risk invariance over a set of more than two envi- ronments means that weighted risk invariance holds for all pairwise combinations of environments in that set. When one of the weighting functions is density ratio of the invariant. features, i.e. when either a(Xanv) Pez (Xe (Xmv) and B(Xmv) = 1; or a(Xinv))
and ((Xinv) Pe. (Xmv)/Pe, (Xime), say this is density ratio weighting. When the weighting functions are invariant feature densities, i.e. when (xXiv) = Pe, (Xiv) and ((Xinv) = Pe, (Xmv), we say this is density weighting.
However, weighted risk invariance alone not enough to guarantee that a predictor is spurious-free As example, the case where a(Xme) ((x 0 allows any predictor to satisfy weighted risk invariance in trivial manner, without excluding spurious features. To avoid these degenerate cases, we assume that the reweighted environments are in general position Then, in previous work principled methods for learning invariant models (Arjovsky et al. 2019 Krueger et al. 2021 Wald et al., 2021], , we prove that predictors with weighted risk invariance discard spurions features Xapu for OOD generalization under general position conditions. We show this for the case of linear regression, where we have the following data-generating process parameterized by Wine Ty: and P4, E3 for training environment E [k]:
Y wil Xin + EE, E~ N(
(5)
= =
77 ~ N,),
Theorem 1. Consider a regression problem following the data generating process of equation 5. Let & be a set of environeents with |8] > dapa that satisfies general position. A linear regression model f(x) = w x with w [Wans: W ssya] that salisfies weighted risk invoriance w.r.l the squared loss must also salisfy W spn For density ratio weighting, general position holds with probability 1.
This result shows that in a linear regression setting, solving for weighted risk invariance allows us to learn a prodictor that discards spurious correlations to generalize across environments. The general position condition excludes cases such as (( = ((xmn) = 0 that reduce to a degenerate set of equations. We further show that density ratio weighting achieves this non-degeneracy with probability 1, and we conjecture that this holds for other choices of weighting functions.
3.2 Finding an effective weighting function
Each weighting function introduces different challenges the learning process. Density ratio weighting has the typical importance weighting limitation, where the support of the distribution in the numerator needs to be contained in the support of the distribution in the denominator. In addition, density ratio weights tend to have a high maximum and high variance, and as a result, lead to a high sample complexity another known issuc of importance weighting (Cortes et al. 2010 @ Conversely, while the number of woights//constrannss for the general form of weighted risk invariance combinatorial weighting only requires weights and constraints that are linear in the number of environments. (Intuitively, this is because we set a reference environment, weighted by a 1, and weight all other environments by a density ratio that aligns them with the reference environment: we discuss this in more detail in our proof of Theorem
the number of environments, density ratio
In practice, we find this benefit is not computationally significant when the number of environments is sufficiently small. Instead, we propose to use the invariant feature densities as the weighting functions. That is, for two environments e and Ej₁ we use a weighted risk
Reies (f) Ep,, Exs, (x), [pe, (X )((()) v)]
5See Appendix A for the full definition of general position.
68ee Appendix for more discussion.

---

Published in Transactions on Machine Learning Research (07 /024)
(4)
0.3
(b)
Spuriousness of linear classifee
AND
Pulla
0.2
ERM
IBM
VREX
-2
-5.0
0.0
Xinv
Number of samples
Figure 3: (a) We demonstrate the finite sample behavior of ERM, IRM, VREx, and WRI methods in the case where the distribution shift is large. We sample data from the distributions shown in (b). Even as the number of samples increases, ERM, IRM and VREx methods continue to select spurious classifiers while the WRI method quickly converges
invariant classifier.
so that weighted risk invariance means (f) Rf) There are several practical benefits to this approach. For onc, density weighting docs not impose roquirements on the support of the environments, and allows for invariant learning over all regions of invariant feature support overlap And although density weighting uses combinatorial number of constraints, the number of weighting functions we need to learn is still equal to the number of environments. As another benefit, learning the invariant feature densities gives us signal indicating the distribution shift in invariant features across domains that is useful for downstream tasks. We therefore choose the weighting functions to be the invariant feature densities, and more precisely define weighted risk invariance as risk invariance with density weighting.
To enforce this weighted risk invariance, we propose to minimize the objective
(w) R Σ -
(6)
cEEer
eveeEEEr
The first term enforces ERM to ensure good average performance across the training environments. The second term enforces weighted risk invariance, so call it the WRI penalty term. Hyperparameter À controls the weight of the WRI penalty.
3.3 Comparison to other invariant learning methods
Comparison to REx Risk invariance across environments is popular and practically effective method for achieving invariant prediction (Xie et al., 2020 Krueger et al. 2021 Liu et al. 2021b Eastwood et al. 2022) that was formalized by REx (Kueeer el. al. 2021). Specifically, the VREx objective is a widely used approach of achieving risk invariance, where the variance between risks is penalized in order to enforce equality risks:
RVRE=(f) (ERM(F) + AVRE ( (f)())
(7)
The variance between risks is equal to the average mean squared error between risks, up to constant. Note that the WRI penalty in equation reduces to the VREx variance penalty under uniform weighting.
Under invariant covariate shift, risk invariance objectives can be misaligned with the goal of learning a predictor that only depends on Xin- To illustrate this, consider a problem with heteroskodastic label noise (where some instances more difficult classify than others (Krueger et al., 2021)). Under this setting, invariant covariate shift results in case where one training environment has more difficult examples than the other. Then invariant predictor would obtain different losses both environments, meaning it will not satisfy risk invariance. Figure demonstrates this case, and shows that risk invariance penalty (VREx) does not learn the invariant classifier. In contrast, weighted risk invariance holds when there both heteroskedasticity and invariant covariate shift.
SSpportt overlap necessary for invariant learning to meaningfully occur; this fundamental limitation of invariant learning general (Ahuja 22021), and violating this does not misalign the weighted risk objective. non-overlapping regions, the weighted risk would simply zero, causing the ERM objective take over.
7

---

Published in Transactions on Machine Learning Research (07 (224))
(a)
(b)
1.0
IRM pecalty
Ponelity tarm (hormalized)
WXI penalty
0.6
0.2
(b)
(c)
0.00
0.75
1.00
1.25
1.50
2.00
Strength of invariant covortate shift (distanoe betwoen invartant means)
0.0
5.0
Figure 4: The WRI penalty is more sensitive to spurious classifiers than the IRM penalty under invariant covariate shift. (a) We start with two data distributions (pı and p2) and increase the distance between their means from to 2 in the invariant direction, essentially creating invariant covariate shift: ranging in degree from the diagram in (b) to the diagram in (c). We then search through the set of challenging, slightly spurious classifiers with near-optimal ERM loss, so that we must rely on the invariant penalty to reject these classifiers. We plot how the minimum IRM and WRI penalties these regions change as covariate shift increases.
Comparison to IRM IRM searches for the optimal predictor on latent space that is invariant across
environments. This work generated many follow-up extensions (Ahuja et al. 2021 Lu et al. 2021], as well as some critical exploration of its failure cases (Rosenfeld et al. 2020 Ahuja et al. 2020 Kamath et al.
2021 Guo et al., 2021). The original objective a bi-level optimization problem that non-convex and
challenging to solve, its relaxation IRMv1 is more often used in practice:
RJRas(f) RERM(() + AFRM Σ
R NI
(8)
Note that IRM uses the average gradient across the dataset to enforce invariance. This approach inherently introduces a sample complexity issuc: when invariance violations are sparse or localized within a large dataset, the average gradient only indirectly accounts for these anomalies, requiring more samples to recognize the violations. Empirically, we observe that even for toy example, the sample complexity for IRM is high under invariant covariate shift. We demonstrate this in Figure where WRI converges to an invariant predictor while IRM does not. (Note that REx does not converge an invariant predictor either- but we do not expect it to converge to an invariant predictor in this case, regardless of the amount of training data.)
In Figure 4 we demonstrate how the average loss-based WRI penalty is more sensitive to invariance violations than the average gradient- gradient-based IRM penalty. To create this figure, we start with a model that uses spurious features and observe how the IRM and WRI penalties for the model change as we introduce covariate shift the invariant features. Ideally, the penaltics should be consistently nonzero (regardless of the degree of shift) the model is not invariant. We that the IRM penalty changes quite significantly the invariant covariate shift increases, while the change for WRI is less pronounced.
3.4 Implementation
Evaluating the WRT penalty requires weighting the loss by functions of the invariant features. However, we do not have access the invariant featurcs-other otherwise the invariant learning problem would already be solved. In Appendix C we empirically show that for a simple case of our regression setting where we can compute exact densities, optimizing for WRI the joint features (rather than the invariant features) still leads to learning invariant feature weighting functions. Unfortunately, density estimation on higl-dimensional data is notoriously difficult, and we often cannot compute exact density values in practice.
To tackle more realistic settings where we do not have exact density values, our approach must optimize for weighted risk invariance while also learning the fcature densitics. We therefore propose to use alternating minimization, an effective strategy for nonconvex optimization problems involving multiple interdependent terms. We propose to alternate between WRI-regularized prediction model and environment-specific
8

---

Published in Transactions on Machine Learning Research (07/2024)
density models. In step, we keep the prediction model fixed and update density model for each environment; in the alternating step, we keep all density models fixed and train the prediction model.
We do not offer any convergence guarantees for this approach. Even so, our experiments support that alternating minimization can recover an invariant predictor and density estimates effectively. Most of the experiments in our main body focus on using WRT to learn an invariant predictor. However, in Appendix D.3, we show that running alternating minimization on our regression setting recovers close approximation of the ground truth densities (up to normalizing constant). We also find that the learned densities are useful for predicting whether an imput is part of the training invariant feature distributions or OOD in $4.1
Practical objective We want to avoid learning a representation where there is no overlap between environments (corresponding to zero weights for all data points), so we introduce a negative log penally term to encourage high entropy and a large support for each environment. We observe empirically that this term encourages overlap between environments. When our density model matches the truc distribution lamily of the features, the negalive log term encourages a maximun likelihood estimate of the feature distribution parameters.
To make the dependence of the weighted risk on the invariant feature densities explicit, we write the as- sociated densities d's as a parameter of the weighted risk R d As in 22. we denote the learned representation as (xx so that f f= $ for some predictor g. The full loss function is
C(F, ( = Σ ( + A Σ (R (f:d') Re (f; + > eEEir
Ep.(x) I- log d
(*(x))
(9)
cEErr
eife E
As in equation 0 the loss comprised of the ERM and WRI penalties, with an additional log penalty term to regularize the learned representations. For full implementation details, refer to Appendix D
4 Experiments
We evaluate our WRI implementation on synthetic and real-world datasets with distribution shifts, particu- larly focusing on cases of covariate shift in the invariant features. In all of the datasets, we seleet our model based on validation set from the test environment, and we report the test set accuracy averaged over 5 random seeds with standard errors. Training validation selection assumes that the training and test data are drawn from similar distributions, which is not the case we aim to be robust to (Gulrajani & Lopez-Paze 2020); test validation selection is more useful for demonstrating OOD capabilities (Ruan et al.] 2021).
Our major basclines ERM (apnnk 1991. IRM (Arjovsky et al., 2019), and VREx (Krueger et al., 2021 IRM and VREx are two other sausallymootvvated works that also search for conditional invariance as a signature of an underlying causal structure. Because WRI shares a similar theoretical grounding, we find it particularly important to compare our empirical performance with these works. Appendix L includes comparisons with other non-causal basclines, as well as additional experiments and details.
4.1 ColoredMNIST
ColoredMNIST (CMNIST) is a dataset proposed by Arjovsky et al. (2019) as a binary classification extension of MNIST. In this dataset, digit shapes are the invariant features and digit colors are the spurious features. The digit colors are more closely correlated to the labels than the digit shapes are, but the correlation between color and label is reversed in the test environment. The design of the dataset allows invariant predictors, which base their predictions on the digit shape, to outperform predictors that use spurious color information. We create a heteroskedasticc variant on this dataset, HCMNIST, where we vary the label flip probability with the digit. We also create HCMNIST-CS. as version of HCMNIST with invariant covariate shift, by enforcing different distributions of digits in each environment. For a predictor to perform well on the test set of these datasets, it must learn to predict based on the invariant features, even under heteroskedasttic noise (present in both HCMNIST and HCMNIST-CS) and invariant covariate shift (HCMNIST-SS) Finally, we create ideal versions of these datasets where we simplify the image data to two-dimensional features of digit value and color. (More details how these datasets were generated can be found in A Appendix E.3
9

---

Published in Transactions on Machine Learning Research (07/2024)
Table 1: Ideal HCMNIST digit and color penalty compari- son (xle-3).
Table 2: Accuracy on HCMNIST without covariate shift and with covariate shift.
Algorithm
Dataset
Digit
Color
Min
Algorithm
HCMNIST
HCMNIST-CS
Avg
HCMNIST
0.0
4.7
Digit
ERM
53.8 2.3
51.1
52.4
WRI
HCMNIST-CS
0.0
0.8
Digit
IRM
67.2 + 2.7
61.3 - 4.7
64.3
VREx
67.4 + 2.3
58.2 = 2.3
62.8
HCMNIST
0.0
5.4
Digit
VREx
HCMNIST-CS
1.1
0.4
Color X
WRI
75.1 + 2.9
74.7 + 3.8
74.9
Table 3: CMNIST OOD detection
Algorithm
1pr0/pr=209% 20%
1pr 40%
(pr0fpr=60%
1pr0fpr=80%
AUROC
ERM
45.6 1 2.8
64.7 11.7
82.211.2
0.539.10.018
IRM
29.8 3.2
46.8 + 2.8
66.4 +1,9
84.8+0.9
0.565+0.017
VRFx
25.3
2.5
43.5 + 2.3
61.6 +1.5
80.8+
0.525+0.014
WRI
36.7 6.7
54.1 1 5.8
68.1 14.55
8..0:1.5
0.59510.000
We use the ideal datasets to evaluate the WRI and VREx penalties with a predictor that uses only the digit value (i.e. an invariant. predictor) and a predictor that uses only the color (i.e. a spurious predictor), and we report the results in Table [ As anticipated, WRI registers zero penalty for the invariant predictor in both the HCMNIST and HCMNIST-CS scenarios. In contrast, VREx registers zero penalty for the invariant predictor HCMNIST but non-zero penalty HCMNIST-CS. favoring the spurious predictor under heteroskodasticity and invariant covariate shift.
We compare our method with all baselines on the image (non-ideal) data in Table 2 These results demon- strate that WRI consistently performs well in both the heteroskedastic and the heteroskedastic and distri- bution shift scttings, while VREx and IRM show significant degradation under invariant covariate shift.
OOD detection performance of our learned densities To Lest the estimated invariant density from WRI, compare its utility an OOD detector to the model confidences from ERM, IRM, and VREx. The evaluation is conducted on modified CMNIST test split with mirrored digits. Since the estimated invariant density should be function of the shape, we expect the invalid digits to have lower density estimates. We compute AUROC values for each experiment and report the results in Table 3 We observe that at all computed false positive rates, the true positive rate higher for our method; this means that our learned density values are better for detecting OOD digits than the prediction confidences from other methods. For more details, see Appendix E.4
4.2 Real-world datasets from DomainBed
We evaluate our method on 5 real-world datasets that are part of the DomainBed suite, namely VLCS (Fang et al. 2013), PACS (Li et al. 2017). OfficeHome (Venkateswara et al.) 2017), TerraIncognita (Beery et al. 2018), and DomainNet (Peng et al. 2019). We run on FRM-trained features for computational efficiency, as these should still contain spurious information as well as the invariant information necessary generalize OOD (Rosenfeld al. 2022). In order achieve a fair comparison, we evaluate all methods starting from the same set of pretrained features. We report the average performance across environments on each dataset in Table 4
While the WRI predictor achieves higher accuracy than the baselines, we find that the performances are all similar. Methods like VREx that rely on assumptions like homoskedasticity may still see high accuracy on datascts that fulfill those assumptions. We note that, addition to generalization accuracy, our method has the additional benefit of reporting when predictions are based on invariant features that have low probability in training. Additional information and results DemainBed are provided in Appendix E.5
10

---

Published in Transactions on Machine Learning Research (07/2024)
Table DomainBed results on feature data
Algorithm
VLCS
PACS
OfficeHome
TerraIncognita
DomainNet
Avg
ERM
76.5 + 0.2
84.7 + 0.1
64.5 + 0.1
51.2 + 0.2
33.5 + 0.1
62.0
TRM
76.7 + 0.3
84.7 + 0.3
63.8 + 0.6
52.8 + 0.3
22.7 + 2.8
60.1
VREx
76.7 + 0.2
84.8 + 0.2
64.6 + 0.2
52.2 + 0.3
26.6 + 2.1
61.0
WRI
77.0 + 0.1
85.2 + 0.1
64.5 i 0.2
52.7 = 0.3
32.8 + 0.0
62.5
5 Related works
Invariant learning Much of the causal motivation for invariant learning stems from Schölkopf et al. (011
who connected the invariant mechanism assumption to invariance in conditional distributions across multiple environments. Since then, invariant learning works have found that the invariant mechanism assumption results in additional forms of invariance, including invariance the optimal predictor AAjovssy et al. (2019) and risk invariance under homoskodasticity (Krueger et al.] 2021). Uncovcring the underlying invariant mechanisms through explicit causal discovery (Peters et al., 2016) can be difficult to scale up to complex, high nigh-dimensional data, significant branch of invariant learning instead leverages these additional forms of invariance to recover a predictor based on invariant mechanisms. While these methods do not provide all
the guarantees typically offered by methods in causal discovery, they are faster and simpler in comparison, and can still provide theoretical guarantees for certain causal structures.
Learning thcory provides another framework for understanding invariant learning. The seminal work of Ben-David et al. (2010) shows that the risk on test environment is upper bounded by the error on the training environment(s), the total variation between the marginal distributions of training and test, and the difference in labeling functions between training and test. While the first term is minimized in ERM, the second term motivates learning marginally invariant features (X) such that Pe (X))) is invariant
across e (Pan et al. 2010 Baktashmotlagh et al. 2013 Ganin et al. 2016: Tzeng et al. 2017 Long et al.
2018 Zhao et al.] 2018), and the third term motivates learning conditionally invariant features ((X) such
that Pe (e({((((X)) is invariant across e (Muandet et al.j 2013 Koyama & Yamaguchi 2020). Observing
the importance of both approaches, several works even attempt to learn features that are both marginally
and conditionally invariant (Zhang et al. 2013 Long et al., 2015; Li et al. 2021 Guo et al. 2021 Our
work follows a similar vein: we focus learning conditionally invariant features, under the assumption that these align with the underlying causal features, but also reweight the conditionally invariant features to
mitigate any shift in their marginal distributions.
Importance weighting methods Starting from the carly and influential method of importance weighting Shimodaira 2000), many methods for dealing with covariate shift use some form of weighting approach (Martin et al. (2023). Weighting methods work by weighting the loss with density ratio of test features to training fcatures (Shimodaira 2000; Huang et al.] 2006), or vice versa (Lu & Zicbart 2014). However, these methods assume that the support of the features in the numerator is contained in the support of the features in the denominator. Further, they can perform poorly when the density ratio is large (i.e. in areas where the features in the denominator are much more rare); these large values lead to high variance and low effective sample size in the density ratio estimator (see Appendix B) For this reason, we propose density weighting rather than density ratio weighting, as density weighting is not. limited by the support of any feature distríbution and leads to better performance practice, as established by Martin et al. (2023 Alternatively, we can also overcome the density ratio support limitation with a generalized form of importance weighting (Fang et al. 2024
Density estimates for downstream tasks Density estimates are useful for variety of downstream lasks. They serve as an intuilive measure of epistemic uncertainty (Hällermeier & Waegernan 2021 in our work, the invariant feature density estimates approximate how much reliable (spurious-freet evidence we have in the training data. As a result, our density estimates are useful for detecting when samples are OOD, along the line of previous/comei concurrent work (Breunig et al.) 2000 CCharpentirr et al., 22000: Peng et al., 2024] In general, density estimates are useful for building trustworthy systems, and they can be leveraged for both
11

---

Published in Transactions on Machine Learning Research (07/2024)
fairness (Cho et al. 2020
Balunović et al. 2021 and calibration (Kuleshov & Deshpande 2022
Deshpande
& Kuleshov 2023
Xiong et al. 2023).
6 Discussion
Limitations of our method The invariant feature weighting functions in WRI allow its solutions to be robust to the case of heteroskedasticity and invariant covariate shift. However, these weighting functions are constrained by the invariant feature densities of each environment; we do not have access to these densities in practice, and can only recover estimate of that density through alternating minimization. Future work could focus better understanding the optimization process (in the style of Chen et al. (2022 2024)), with the goal of recovering the invariant feature densities with more rigorous guaranlees. We believe that an accurate invariant feature density estimate would be useful for reporting the confidence of predictions during deployment. In general, the question of how to better quantify the uncertainties of domain adaptation/generea /ee/greezzzaaiinnn models remains an open line of research.
Limitations of invariant prediction Invariant prediction methods assume there exists some general- izable relationship across training environments, but there are cases where it is impossible to extract the relationship. Notably for linear classification, Ahuja et al. 2021 proves that generalization is only possible if the support of invariant features in the test environment. is subset of the union of the supports in the training environments. For this reason, it is important to report when we are given data points that are rare in the invariant feature distribution; then, users know when a model is extrapolating outside of the training domain.
Broader impacts Compared to other invariant learning methods, our WRI method offers the additional benefit of providing density estimates that allow level of OOD detection, which can be useful in safety- critical situations. In general, we believe these estimates provide useful measure of epistemic uncertainty that is not usually present in domain gencralization methods.
7 Conclusion
Our work demonstrates the utility of weighted risk invariance for achicving invariant learning, even in the difficult case of heteroskedasttciity and invariant covariate shift. We proved that weighted risk invariance is signature of spurious- free prediction, and we further proved that enforcing weighted risk invariance allows to recover spurious-free prodictor under general causal setting. We proposed an efficient and useful method of weighting by the invariant feature densities, WRI, and we introduced an algorithm to practically enforce WRT by simultaneously learning the model parameters and the invariant feature densities. In our experiments, we demonstrated that the WRI predictor outperforms popular baselines under invariant covariate shift. Finally, we showed that our learned invariant feature density reports when invariant features are outside of the training distribution, a useful signal for the predictor's trustworthiness.
8 Reproducibility Statement
The code for generating the figures and empirical results can be found at httpp:://githb..commm weighted_ braeeedd__rr__-eeeee invariance/
Acknowledgements
GW and RC were partially supported by a MURI from the Army Research Office under Grant No. W9 17-1-0304. This is part of the collaboration between US DOD, UK MOD, and UK Engineering and Physical Rescarch Council (EPSRC) under the Multidisciplinary University Research Initiative. AL was partially supported by JHU Discovery Award, Amazon-JHU AI2AI Award, and a seed grant from JHU Institute of Assured Autonomy.
12

---

Published in Transactions on Machine Learning Research (07 (204))
References
Yael Adini, Yael Moses, and Shimon Ullman. Face recognition: The problem of compensating for changes in illumination direction. IEEE Transactions on pattern analysis and machine intelligence, 19(7):721 732, 1997.
Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk minimization games. In International Conference on Machine Learning, pp. 145-155. PMLR, 2020.
Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, Jean-Christophe Gagnon-Audel, Yoshua Bengio, loan- nis Mitliagkas, and Irina Rish. Invariance principle meets information bottleneck for out-of-distribution generalization. Advances in Neural Information Processing Systems, 1133488-4554 2021.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint wrXic:1907.08899 2019.
Mahsa Baktashmotlagh, Mehrtash T Harandi, Brian C Lovell, and Mathieu Salzmann. Unsupervised do- main adaptation by domain invariant projection. In Proceedings of the IEEE international conference compuler vision, pp. 769 776, 2013.
Mislav Balunović, Anian Ruoss, and Martin Vechev. Fair normalizing flows. arXiv preprint arXiv:2106.05977 2021.
Sara Beery, Grant Van Horn, and Pictro Perona. Recognition in terra incognita. In Proceedings of the European conference computer vision (ECCV), pp. 456 473, 2018.
Sara Beery, Yang Liu, Dan Morris, Jim Piavis, Ashish Kapoor, Neel Joshi, Markus Meister, and Pietro Perona. Synthetic examples improve generalization for rare classes. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. 863 863-873, 2020.
Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman Vaughan. A theory of learning from different domains. Machine learning, 19:151-175, 2010.
Gilles Blanchard, Gyemin Lee, and Clayton Scott. Generalizing from several related classification tasks to a new unlabeled sample. Advances in neural information processing systems, 24, 2011.
Markus M Breunig, Hans-Peter Kriegel, Raymond T Ng, and Jörg Sander. Lof: identifying density-basse local outliers. In Proceedings of the 2000 ACM SIGMOD international conference on Management of data, pp. 93 1841 2000.
Junbum Cha, Kyungjac Lee, Sungrae Park, and Sanghyuk Chun. Domain generalization by mutual- information regularization with pre-trained models. In European Conference on Computer Vision, pp. 457. Springer, 2022.
Bertrand Charpentier, Daniel Zügner, and Stephan Günnemann. Posterior network: Uncertainty estimation without ood samples via density-based pseudo-counts. Advances in Neural Informalion Processing Systems, 33:1356-133 11656,,37,,..1 2020.
Wuyang Chen, Zhiding Yu, Zhangyang Wang, and Animashree Anandkumar. Automated synthetic-to-real generalization. In International Conference on Machine Learning, pp. 1746-1756. PMLR, 2020.
Yongqiang Chen, Kaiwen Zhou, Yatao Bian, Binghui Xie, Bingzhe Wu, Yonggang Zhang, Kaili Ma, Han Yang, Peilin Zhao, Bo Han, et al. Pareto invariant risk minimization: Towards mitigating the optimization dilemma in out-of-distribution generalization. arXiv preprint arXiv:220600770 2022.
Yongqiang Chen, Wei Huang, Kaiwen Zhou, Yatao Bian, Bo Han, and James Cheng. Understanding and im- proving feature learning for out-of-distribution generalization. Advances in Neural Information Processing Systems, 36, 2024.
Jacwoong Cho, Gyeongjo Hwang, and Changho Suh. A fair classifier using kernel density estimation. Ad- vances in neural information processing systems, 33:15088-150888 2020.
13

---

Published in Transactions on Machine Learning Research (07/2024)
Corinna Cortes, Yishay Mansour, and Mehryar Mohri. Learning bounds for importance weighting. Advances in neural informalion processing systems, 23, 2010.
Shachi Deshpande and Volodymyr Kuleshov. Calibrated propensity scores for causal effect estimation. arXiv preprint ar Xrx::2306. 00382, 2023.
Cian Eastwood, Alexander Robey, Shashank Singh, Julius Von Kügelgen, Hamed Hassani, George J Pappas, and Beruhard Schölkopf. Probable domain generalization via quantile risk minimization. orXiv preprinl arXiv:2207.0900 2022.
Chen Fang, Ye Xu, and Daniel N Rockmore. Unbiased metric learning: On the utilization of multiple datasets and web images for softening bias. In Proceedings of the IEEE International Conference on Computer Vision, pp. 1657 1664, 2013.
Tongtong Fang, Nan Lu, Gang Niu, and Masashi Sugiyama. Generalizing importance weighting to a universal solver for distribution shift problems. Advances in Neural Information Processing Systems, 36, 2024.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversan training of neural networks. The journal of machine learning research, 17(1):2066-2330, 2016.
Robert Gcirhos, Henn-k Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. Shorteut learning in deep neural networks. Nature Machine Inlelli- gence, 2(11):665-677 2020.
Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. arXi::200700494, 2020.
ar preprint
Ruocheng Guo, Pengchuan Zhang, Hao Liu, and Emre Kiciman. Out-of-distribution prediction with invariant risk minimization: The limitation and an effective fix. arXiv preprint arXiv:2101.07732, 2021.
prediccno
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness common corruptions and perturbations. arXie preprint arXiv:1903. 12261, 2019.
Jiayuan Huang, Arthur Gretton, Karsten Borgwardt, Bernhard Schölkopf, and Alex Smola. Correcting sample selection bias by unlabeled data. Advances in newral information processing systems, 19, 2006.
Eyke Hüllermeier and Willem Wacgeman. Alcatoric and epistemic uncertainty in machine learning: An introduction to concepts and methods. Machine Learning, 110(3):457 506, 2021.
Pritish Kamath, Akilesh Tangella, Danica Sutherland, and Nathan Srebro. Docs invariant risk minimization capture invariance? In International Conference on Artificial Intelligence and Statistics, pp. 0069-4077.
PMLR, 2021.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
Masanori Koyama and Shoichiro Yamaguchi. Out-of-distrbbatio generalization with maximal invariant predictor. 2020.
David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Ont-of-distributioo generalization via risk extrapolation (rex). In International Conference on Machine Learning, pp. 5815 5826. PMLR, 2021.
Volodymyr Kuleshov and Shachi Deshpande. Calibrated and sharp uncertainties in deep learning via density estimation. In International Conference on Machine Learning, pp. 11683 11693. PMLR, 2022.
Bo Li, Yezhen Wang, Shanghang Zhang, Dongsheng Li, Kurt Keutzer, Trevor Darrell, and Han Zhao. Learning invariant representations and risks for semi-supervised domain adaptation. In Proceedings of the IFEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 1104-1113, 2021.
14

---

Published in Transactions on Machine Learning Research (07 (07/2024)
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain gener- alization. In Proceedings of the IEEE international conference computer vision, pp. 1542-55500 2017.
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy Hospedales. Learning to generalize: Mela-learning for domain generalization. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
Anqi Liu and Brian Ziebart. Robust classification under sample selection bias. Advances in neural information processing systems, 27, 2014.
Jiashuo Liu, Zheyuan Hu, Peng Cui, Bo Li, and Zheyan Shen. Heterogeneous risk minimization. In Inter- national Conference on Machine Learning, pp. 6804 6814. PMLR, 2021a.
Jiashuo Liu, Zheyan Shen, Pong Cui, Linjum Zhou, Kun Kuang, Bo Li, and Yishi Lin. Stable adversarial learning under distributional shifts. In Proceedings of the AAAT Conference on Artificial Intelligence, volume 35, pp. 8662 8670, 2021b.
Mingsheng Long, Yue Cao, Jianmin Wang, and Michael Jordan. Learning transferable features with deep adaptation networks. In International conference on machine learning, pp- 97-105. PMLR, 2015.
Mingsheng Long, Zhangjie Cao, Jianmin Wang, and Michael I Jordan. Conditional adversarial domain adaptation. Advances in newral information processing systems, 31, 2018.
Chaochao Lu, Yuhuai Wu, Jose Miguel Hernández-Loben and Bernhard Schölkopf. Nonlinear invariant risk minimization: A causal approach. arXiv preprinl arxiv:2102. 12355, 2021.
Jose Ignacio Segovia Martin, Santiago Mazuelas, and Anqi Liu. Double- for covariate shift adap- tation. In International Conference on Machine Learning, pp. 00439-30457, PMLR, 2023.
Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In International conference on machine learning, pp. 10-18. PMLR, 2013.
Sinno Jialin Pan, Tvor W Tsang, James Kwok, and Qiang Yang. Domain adaptation via transfer component analysis. IEEE transactions on neural networks, 22(2):199 210, 2010.
Judea Pearl. Causal diagrams for empirical research. Biometrika, 82(4):669 688, 1995.
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer,
R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12:2825 2830, 2011.
Bo Peng, Yadan Luo, Yonggang Zhang, Yixuan Li, and Zhen Fang. Conjnorm: Tractable density estimation for out-of-distribution donection. In The Twelfth International Conference on Learning Representations, 2024. URL /tttps://operrevie:: wttpeiiwwwwwnneeeeeeeeeeee net wew.eeorrorid--llii2ee.i...
Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In Proceedings of the IEEE/CVE international conference computer vision, pp. 1406,141,,.. 2019.
Jonas Peters, Peter Büihlmann, and Nicolai Meinshansen. Causal inference by using invariant prediction: identification and confidence intervals. Journal of the Royal Statistical Sociely: Series B (Statistical Methodology), 78(5):947 1012, 2016.
Mateo Rojas-Carulla, Bernhard Schölkopf, Richard Turner, and Jonas Peters. Invariant models for causal transfer lcarning. The Journal of Machine Learning Research, 1911:1300-1132, 2018.
Elan Rosenfeld, Pradcep Ravikumar, and Andrej Risteski. The risks of invariant risk minimization. arXiv preprint XrXio:2010.051: 057661, 2020.
Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. Ddmain-adjusted regression or: Erm may already learn features sufficient. for out-of-disttiiutiio generalization. ar Xío preprint (xXin:2202.06855 2022.
15

---

Published in Transactions on Machine Learning Research (07/2024)
Yangjun Ruan, Yann Dubois, and Chris J Maddison. Optimal representations for covariate shift. arXiv preprint or Xrvi:2200..00 100057, 2021.
Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang, Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint arXiv:1911.089911 2019.
Bernhard Schölkopf, Dominik Janzing, Jonas Peters, Eleni Sgouritsa, Kun Zhang, and Joris Mooij. On causal and anticausal learning. arXic preprinl 12Xi::1066.64 1222006667,,, 71, 2012.
Hidctoshi Shimodaira. Improving predictive inference under covariate shift by weighting the log-likelihood function. Journal of slatistioal planning and inference, 90(2):227 244, 2000.
Baochen Sun and Kate Sacnko. Deep coral: Correlation alignment for deep domain adaptation. In Computer Vision ECCV 2016 Workshops: Amsterdam, The Netherlands, October 8-10 and 15-16, 2016, Proceedings, Part III 14, pp. 443 450. Springer, 2016.
Antonio Torralba and Alexei Efros. Unbiased look at dataset bias. In CVPR 2011, pp. 1521 1528. IEEE, 2011.
Eric Tzeng, Judy Hoffman, Kate Saenko, and Trevor Darrell. Adversarial discriminative domain adaptation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7167 7176, 2017.
Vladimir Vapnik. Principles of risk minimization for learning theory. Advances in newral information processing systems, 4, 1991.
Ramakrishna Vedantam, David Lopez-Paz, and David J Schwab. An empirical investigation of domain generalization with empirical risk minimizers. Advances in Neural Information Processing Systems, 34: 28131-28143, 2021.
Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. In Proccedings of the IEEE conference on computer vision and poltern recognition, pp. 5018 5027, 2017.
Yoav Wald, Amir Feder, Daniel Greenfeld, and Uri Shalit. On calibration and out-of-domain generalization. Advances in neural information processing systems, 34:2215 2227, 2021.
Olivia Wiles, Sven Gowal, Florian Stimberg, Sylvestre Albise-Rebuff, Ira Ktena, Taylan Cemgil, et al. A fino-grained analysis on distribution shift. arXiv preprint 2rXiv:2110.1100 11328, 2021.
Xi Wu, Yang Guo, Jiefeng Chen, Yingyu Liang, Somesh Jha, and Prasad Chalasani. Representation bayesian risk decompositions and multi-source domain adaptation. arXiv preprint ar arXiv:2004. 103900 2020.
Chuanlong Xie, Haotian Ye, Fei Chen, Yue Liu, Rui Sun, and Zhenguo Li. Risk variance penalization. arXiv preprint Xiv:20006 07544, 2020.
Miao Xiong, Ailin Deng, Pang Wei Koh, Jiaying Wu, Shen Li, Jianqing Xu, and Bryan Hooi. Proximity- informed calibration for deep neural networks. arXiv preprint arXiv:2306.0450 2023.
Nanyang Ye, Kaican Li, Haoyue Bai, Rumpeng Yu, Lanqing Hong, Fengwei Zhou, Zhenguo Li, and Jun Zhu. Ood-bench: Quantifying and understanding two dimensions of out-o--distiibution generalization. In Proceedings of the IEEE/CVF Conference on Compuler Vision and Pallern Recognition, pp- 7947-7958, 2022.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint arXiv: 09412, 2017.
Kun Zhang, Bernhard Schölkopf, Krikamol Muandet, and Zhikun Wang. Domain adaptation under target and conditional shift. In International conference on machine learning, pp. $19-827. PMLR, 2013.
Han Zhao, Shanghang Zhang, Guanhang Wu, José MF Moura, Joao P Costcira, and Geoffrey Gordon. Adversarial multiple source domain adaptation. Advances in neural information processing systems, 31, 2018.
16

---

Published in Transactions on Machine Learning Research (07 (07/2024)
A Proofs
Proposition Let Assumption 1 hold over a set of environments & S P. If predictor f is spurious- free over &, then for every pair of environments ei, E E, their weighted risks are equal if their respective weighting functions a and B obey (e (Ximx) ((xin)) (Xtre).
Proof. The weighted risk for environment is given by
Ie Jx Pi (Xinv, Xapus v) (X((xinv Xspu), a(x) d dxspu dinne dy-
If f and are invariant, then can transform this integral follows.
J Xun Xepu
J. S pi (Xivvs Xapu: v) e(f Xapu); 1) (xXinv) d d'sp d Xinv dy
S, / ((I(xinv) a (xxin)) f
pi (Xinv, Xspu: 11 d Xapu Xinw dy
Xin
Kapa
I, [ ((f(Ximv)))) (xxin)) pp (Xave, v) d Xiuv dy
Xiw
[ I ((f(xiw)),)) (xxxm)) pr pi(yl Xinw) Pı (Xamx) d Xiav dy.
A.1
Xun
Similarly, if B is invariant, then the weighted risk for environment given by
J/ (([(Xinvss)) y) ((X P2(1) Xinv) p2(Xinn) Xinv dy.
(A.2)
Because py( Xine) the same for all e, then equation A.1 and equation A.2 equal if (xxinv))) (Xinv) (Xxnw) P2
To show that weighted risk invariance leads to OOD generalization, must first make assumption the diversity of our training environments. We give definition for general position in similar style to previous works (Arjovsky et al. 2019 Wald et al. 2021
Definition A.I. We given set of k environments such that ( > 2d where dapu is the dimension of spurious features. Each environment, has weighted covariance End and weighted correlation p²d with another environment j. We define these
IN Ex,(X [xx] (Xiv)] 3 = 2 Ep,(X)))) [x your (Xiv)]]
We say that this set of environments is in general position if for any scalar y E IR and all nonzero x E Rdppa
dim(span (( E +x n² }(SEN))) = dapa-
(A.3)
With this definition, we proceed to show that for a set of environments in general position, linear regression model that satisfies weighted risk invariance must be spurious-free
Theorem 1. Consider a regression problem following the data generoting process of equation 5. Let & be a
set of environments with [3| > dapn. that satisfies general position. A linear regression model f(x) T x with w [Wew: W Wayul that satisfies weighted risk invoriance w.r.t the squared loss must also satisfy W Wapu For densily ratio weighting, general position holds with probability 1.
17

---

Published in Transactions on Machine Learning Research (07/2024)
Proof. Assume we have weighted risk invariance between environments 1 and 2, so
J / Kepur
[Xiuv: Xspu] y)² 3)2 .p (Xivs Xspu: y) a (x d dxspu Xuw dy
Ainv
Ja / (w" [Ximv Xapull 8)² P2 (Xinv Xapu: v) ((xmmv) dxspu dxmnv dy.
(A.4)
Xunx Kapa
When we expand the square, the left hand side becomes
TT (1o Jaan Ja
[ Xspu Knnv Xspul' P1 (Xinvi Xspus y) a(xinv) d Xspu d d'iuv dy
2 So. fo
[Xiav, Xapuly Pı (Xinv Xapu 8) a (X dxppu d d dy
Nua Kopa
J Jenn Jo 5² -P- (Xinv, Xapu: 11 al (xxiv)) d Xapu d Xine dy.
Xun Sxau
We define
11 |, Nun Eipu
/
1 [Xiav Xapul [xvn, Xapul" (Xinv, Xapu y) a(xin) d d'pppu d d dy
f
(Xivys Xspul X Xspul Pr (Niuvr Xspu) a (xxin)) d Xapu d Xiav
Xiav Napa
1. 1/-
[Xinv Xepuly Pi (Xinv, Xspu: y) a(Xmn)) Xapu d d dy
y Maav
Napa
cl.2 Is / Xin Xipu
ppi (Xinvs Xspus y) af (Xin)) d Xapa d Xinv dy
S [ y² (P y) a(xiv) d Xinv dy,
with analogous definitions for 2,1 1221, and o2.1, Then equation A.4 can be written more succinctly
w y12 w w 2 + ot2 = w 2 w w 2211 + 2211
(A.5)
We decompose the quadratic term 212 into invariant, spurious, and mixed components:
f
12
[Xinvi Xspu (xivs Xspul' T Pi (Xinvs Xspu) a a(xinv) d Xapu d Xiny
Nine Xapu
Xinv X Xzpu X
Xinv xnn pr (Xinv. Xspu) at a(Xinv) d dxspu diav
Xier Keger
Xspu x Xppu
=fo Jxmn
Xinv XJnv
Pi (Xinv) ar (x(xivv) d Xinv
+
Pi (Xinv, Xsnu) orl (((Xinv) Xepu d Xine
Xspu T
Xuvv
Xapu
Xinv xpu 0
Pi (Xinv Xapu) a(xinv) Xxpu d d
Xapu.
Xepu X
1.2
X-Inv
07
0
1,2 miiz
0
0
0
1 Spp
+
18

---

Published in Transactions on Machine Learning Research (07 (07/2024)
We break down the third term E Emir in more detail, following the data-generating process for Xapu
211,2 miz
0
Xinv X
pi (Xinv, Xspu) a (xxiv)) d Xepu d Xini
Kepu
Xspu Xinv
Jk
Naa [ (Winne) Xinv Mi X tep, x +1x
0
Xinv Hu x Wine +Xin ET + Xinv 17
...P((s) pl (T) pi (Xinv) a(xinv) dn de d Xine
Xinv 1o x Winv (X(au))) al a(xim)) Xiav
Xan (W T Xinv Sa xhav
0
Xinv 14g
[uix
E (P) Pı (Xinv) a a(xinv) de d Xinv
Li x
Nua
Xinv 17
pi())) (X((um)) axxmn) dn d Xivay.
!1) xI
0
Both and y) are mean-zero random variables that are independent Xinv. This means the second and third terms go to zero, leaving us with only the first term. We further decompose the first term, now explicitly labeling the dimensions of the 0 submatrices for clarity.
Idine
Odmo x1 pii
1 1.2
Idacx O1 Oxxdire
Oduev X dapu
O Oupa diav
14
where
1.2 -ine
Odnu xdin
Xinv (Winv) Xinv
(Xinv) c((Xiv)) Xinv
Xivv Winv Xinv
01x1
Nin
O xdui
v222
(vl.2)
0 01x1
with v1.2 Xiav (War)) T Xinv being a din-dimennional column vector.
We also separate the linear term 2 into invariant and spurious components:
|, Xivv Mapu
|
[/ [Xivv Xspuly Pi (Xinvs Xepu 3) ai (xxin)) ) d d'sppu Xinv dy
12 =2
=2 ( / (Xinv;
Oly
pi (Xinv, Xepu, y) a (Xinx) d Xapa Xinv dy
Ainv
Xopu
+ 2
[0, Xspuly
Pi (Xiuv, Xepus v) a (Niv)) d d'iny d'spp dy
Xepm
Xnn
/
=2
(Xinv; Oly pı (vl Xine) pr (Xime) a axXinv) d Xinv dy
Anv
/ [0, Xyply / Xmv
2 yy Nopu
Pr (Xiuss Xepus y) a (Xiinv)) dl d'vsv xxpp dy
= 1
1.2 0] + [0, 1.2)-
Finally, we break down the constant term c2.2.
S Xuv
d.2
y Pi (Xiuv v) a(xivv) d Xinv dy
/ 1 y² Pi (vl xmn) ) pı (Xinv) (xXnv) dxinv dy.
19

---

Published in Transactions on Machine Learning Research (07 (07/2024)
Given the condition on the weighting functions ((X = (X )p it is clear that E x2.1 v1.2 = v2,1 1,2
nnv Pinv and ct.2 = c2.1. With this information, we define y == 2 whhx v - 1 and simplily
equation A.5 to
21,
Wspa +
1 Wspu w E2,1 spu spu Wspu (1( 2.1 T Wspu
(A.6)
spu "sp
Since equation A.6 holds for each pairwise combination of environments, then
apu (Zi
Wspu +
(app) Wspu Vi,j [k].
(A.7)
We define the (9) dapu matrix
w T (21.2
Σ2.1
p T
apu
epu
w [I1.3 spu
Σ3.1
1.3
(Eah-1 Ek-1,k)
((upp 1.
1spu T (
w spu
The environments are in general position, so matrix M is full rank for any nonzero Wspu- That means that there is no nonzero vector that solves
M x = 0.
(A.8)
If there is no nonzero solution to equation A.8 then there is no nonzero wspu that solves equation A.7 equation A.5 implies that Wspu = 0.
Thus,
We now examine the case of density ratio weighting, For a pair of environments and j' with the weighting constraint a ((Xav))e, (Xav) = 8 (Xin) (Xin)) we let a(xiv) = 1 so that 3(xmm) = Per (Xim))/ee. (Xim). The weighting function a no longer depends on an enviromment-specffe distribution; instead, it is the same for all pairs of environments that include environment i'. This weighting allows us give an additional result based
smaller set of equations centered around environment i!
Under density ratio weighting, covariance L is identical for all environment indices i, is mean n This means that if we generalize equation A.6 to all environment pairs with environment 1, the left hand side of the equation identical. Then for some scalar / E TR,
T I w spu spu Wapu (H(pp)) Wepu =t Vi [k] 1.
(A.9)
This system of equations can be written as a k x dspu matrix
wwT. w spu - + (M2p)) ) T
spu 2spu + 7 2.1 (Mepu)
M'
w spu
Supu + (1 T
The steps by which the density ratio weighting simplifies matrix M to M' should illustrate how our general position definition also be stated more simply under this weighting, in line with the definitions given by similar works. Then, ftrriaht application of Theorem 10 of Arjovsky et al. (2019) or Lemma $4 of Wald et al. (2021) gives our final result: that the subset of environments in our setting that are outside of general position has measure zero.
20

---

Published in Transactions on Machine Learning Research (07/2024)
B On the sample complexity of different weighting functions
Our work differs from other works on weighting to mitigate distribution shift. We are the first to propose weighting by functions of the invariant features; this allows us to use weighted risks to recover spurious-free predictor under shift, whereas most forms of weighting only focus on mitigating shift, without any invariant learning mechanisms. Still, work on the generalization bounds of importance weighting (to mitigate shift) can be readily extended our case. In this section, we use the arguments of Cortes et al. (2010) show that convergence guarantees
improved by selecting weights with lower variance and lower maximum value. This result explains some of the benefits of using density weights over density ratio weights.
We fix f H, where H denotes the hypothesis set under consideration. For some weighting functions a and 3 such that (x) (x) = Pea) (x), we define the weighted risk over two environments e1 and to be
R(F) = E Pe, (X.Y) (a(x) 8)] P., (x,x) ((x))
(B.10)
The empirical weighted risk is defined to be
(.) [De
Σ
a(x)e([(x);)) v), and
x.yEDe
(B.11)
R₂(1) = iDeg x.yED
((x))([(x))))
Following the proof of Theorem 1 from Cortes et al. (2010), we assumee that {(f(x),)) v) E [0, 1]. we de- fine a random variable Z. ((X)e(((X)))) - R()) ) where X, Y are drawn from Pet Then, (²() =
pp,, (x,Y) [x² R(f)2. Applying Bernstein's inequality yields
-Peel exp 202(Z) 2x 2s/a33)
Pr(R(() -
(B.12)
where Mo max{1, ). Setting å equal to the bound from Eq. B.12 and solving for E, it follows that with probability at least 1 - the following bound holds for any > 0:
R(() - Rn(
Ma log R
M2 log 2 202 log 1
+
3/Deal
9 2
[Dey
(B.13)
2M4 log 1/
20 (Za) log
$ 3/D.
+
Decl
Following the same steps, get bound for - If we combine these events with the union bound, then the probability of both of these bounds holding is at least 1-26. 28.
Note that for density ratio weights, the maximum values M, and Mg can be large (e.g. when a sample is much more rare in one environment than another); this also leads to the variances o² (Za) and (Z being large. In contrast, density weights tend have lower maximum value and lower variance.
We emphasize that this section does not provide a sample complexity bound for method; it should be treated exploratory comparison of different weighting functions only, with the caveat that our result relies on several assumptions that would not hold in practice. We fix the hypothesis f, but it would typically be learned from training data. We also assume the optimal weighting functions are known, when these would need to be learned from the data/ trained model as well.
21

---

Published in Transactions on Machine Learning Research (07 (07/2024)
C Optimization of WRI for the 2D Gaussian case
In the main body of the paper, show that for set of invariant weighting functions, weighted risk invariance is necessary and sufficient condition for spurious-free prediction. This demonstrates the importance of correcting for invariant covariate shift, and extends the importance weighting method to the invariant learning setting. However, we encounter a chicken-and-egg problem in the implementation of our method; specifically, that we weight by the density of the invariant features, without having access to the invariant features in practice. At best, we only have access to the density the joint features, both spurious and invariant.
In this section, we show that for 2D case of our regression setting in equation optimizing for WRI leads to learn model that discards the spurious features, even when we slart with the densily of the joint features. In particular, we see from Figure C.1 that optimizing for risk weighted by the joint feature density eventually leads to the weights on the spurious features going to zero, leaving only the weights on the invariant features.
1.0
(initial weights spurious)
WRI (initial weights equal)
ERM (initial weights all spurious)
0.8
ERM (initial weights equal)
Spurious weight (normalized)
0.6
0.4
0.2
0.0
2000
4000
6000
8000
10000
Training epochs
Figure C.1: Given data of the form [ine Xapal following setting in equation 5 train model with two- parameter feature layer, where one weight scales the invariant features and one weight scales the spurious features. For both the WRI and ERM methods, weights and with weights on the spurious features only. We find that optimizing WRI always leads to the spurious weights going zero, whereas optimizing ERM converges nonzero spurious weights.
do initializations where initialize the model with equal feature
C.1 Experiment details
The data generation process for this experiment follows equation 5 where each data point x [Xivvs X This data then used train two-layer linear model, where the feature layer has two weight parameters: one weight for the invariant features, and one for the spurious features. The model layers do not have any activation functions, they are straightforward linear transformations.
To decouple our investigation from the difficulties associated with density estimation, we assume that Xinv is Gaussian-distrrbuttee.. this way, we have a closed form solution for the density of the features. When we optimize for WRI, we therefore weight by the exact feature density.
We use this setup to optimize for both the WRI and ERM methods. To start, we initialize the model with either equal feature weights (where the invariant and spurious feature weights equal), with the
22

---

Published in Transactions on Machine Learning Research (07/2024)
invariant feature weight (so all the weight is the spurious features). We find that, regardless of initialization, the WRI-trained model learns to reject the spurious weight. This means that as the model trains, we learn to weight by an invariant feature density and predict from the invariant features only. The ERM-trained model, on the other hand, never reaches a nonzero spurious weight.
C.2 Discussion
The results of Figure practical implementation, even though we only have access to the input data and not the invariant features. However, please note that we only show this for specific case of our regression setting, where the invariant and spurious features are casily separable and the invariant features Xnnv follow Gaussian distribution. If we were run this experiment on data from an unknown distribution, the results would be less clear: our estimate feature density, inexact operation in weighting functions would depend on how correctly high dimensional cases.
c are promising evidence that invariant feature reweighting can be realized in a
23

---

Published in Transactions on Machine Learning Research (07 2024)
Algorithm D.1: WRI with model-based density
Parameters:
n: total number of optimization steps
density update frequency
ng: number of density update steps
initialize random model weights for and de for all e E Etr
initialize Adam (Kingma & Ba 2017 optimizer for f initialize Adam optimizer for d for -...n do
sample featurized minibatch from all environments compute L equation 9 on minibatch and step optimizer if i-1 is multiple of then repeat n. times
sample featurized minibatch from all environments compute C on minibatch and step de optimizer
D Implementation details
When the optimal invariant predictor was first introduced with IRM, it was first molivated as a constrained optimization problem that searched over all spurious-free predictors for the predictor with the lowest training loss. The practical objective IRMv1 was then introducod, where the constraint was approximated by gradient penalty. Many follow-up works took similar path, where penalty is used as some form of invariance constraint. Our penalty is weighted risk invariance; according to our theory, models with zero weights on spurious features will satisfy the WRI constraint. Therefore, a learning rule that finds the most accurate classifier satisfying WRI will learn an optimal hypothesis that relies on the invariant features alone.
We implement this learning rule with alternating minimization; our specific algorithm is described in Al- gorithm [.I] and a graphical representation of our network architecture is shown in Figure D.2 This is complex optimization procedure, and we do not derive guarantees its convergence invariant repre- sentations in this work. We simply motivate this implementation with the knowledge that spurious-fiee predictor would minimize the WRI penalty, which we optimize for when the penalty cocfficient in equation g is sufficiently large. We test that optimizing WRI on the observable joint feature density still focuses model weights on the invariant features in Appendix(C] In the main body of the paper, we empirically show that our objective effective at recovering a generalizable predictor, as well as density estimates with a meaningful signal for OOD detection. Later in this section (Appendix D.3). we also test the quality of our density estimates for 2D case of regression setting in equation 5
D.1 Optimizing WRI leads to invariant features
To provide additional intuition on why optimizing for WRI leads to learning invariant features, let us consider trying to reweight a representation ((X) that is not invariant. Specifically, assune that for two environments, Pi (Y) ((X)) # Pz(Y (()).. In what follows, we demonstrate why WRI would reject such a representation. We focus on density ratio reweighting for simplicity, so that p = p1 (X)) P2 ((X)).
To satisly invariance under this reweighting, must have 5(Y +(X)) = (Y ((X)). Clearly this does not hold, since p(Y ((X)) PY *(X)), and this lack of invariance can be identified from data under mild assumptions on overlap. Therefore, enforcing this type of weighted invariance constraint (i.c. matching of the conditionals Pı (Y &(X)) and the density ratio weighted distribution p(Y (()))) would rule out a non-invariant representation (XX). Conversely, an invariant representation would naturally satisfy this constraint.
The gap between this constraint and weighted risk invariance lies in the fact that we only enforce invariance of the risk, not the full conditionals (we do this because imposing risk invariance is more tractable in practice).
24

---

Published in Transactions on Machine Learning Research (07 (07/202) /204)
Inputxx x
Data Representation
(x)
Predictor g
Weightfunction d ... dk
Predictorf
Σ (R Rejei(fidd))) + B j QLEBEE
==
Re (f) +
E p log de (Φ(x))
eEEt
EEE
Figure D.2: Graphic overview our network architecture. The red lines indicate back pack-propagation paths for components that impact the first two terms. The blue lines indicate the paths for the last two terms.
While the two are not equivalent, there are cases where they coincide. For example, Theorems 1 and 2 and the experiments of Krueger et al. (2021) support the claim that risk invariance learns an invariant predictor.
A subtle but important point is that these theorems presuppose the existence of risk invariant predictor. Under the invariant covariate shift we consider, some invariant fcatures must be discarded to obtain a risk invariant predictor (e.g. in the 2D Gaussian case, risk invariant predictor would have to discard the causal feature completely). With weighted risk invariance, this issue is alleviated. By using (for example) density ratio reweighting, the optimal invariant predictor becomes risk invariant again.
D.2 Practical strategies
We observe that the magnitude of the WRT regularization term from equation g varies significantly in practice, making it difficult to choose good value for X. This is true even when the true invariant density is known. We employ two strategics to deal with this. First, we constrain the density model to predict values within a pre-defined range by applying sigmoid activation on the final prediction with constant scale and shift factor. Additionally, we divide the WRI term by the average negative lieliio,, which also helps to decouple the empirical risk from the WRI regularization term.
We allow different optimization parameters for the prediction model and density estimation models. Both op- timizers have different learning rate, weight decay, batch size, and À penalty. (The range of hyperparameter values are provided with the DomainBed details in Appendix E.5
D.3 Density estimation quality
Optimizing WRI with exact feature density values successfully recovers an invariant predictor, as we show in Appendix C Yet, density values are often estimated with some inaccuracy in practice. In this scction, we investigate the accuracy of our alternating minimization method in estimating feature densities. We use the Gaussian case of our regression setting in equation 5 detailed in Appendix C Combining this with
25

---

Published in Transactions on Machine Learning Research (07/2024)
linear model for feature learning provides closed-form solution for the feature density that can compare against our estimated densities at each alternating step
Learning the parameters of density model When we know the distribution family we are estimating the densitics of, we should set up our density estimation model so that we are learning the parameters of the same distribution. In this case, the log penalty pulls us toward a maximum likelihood estimate of the parameters that would generate the observed feature densities.
Learning unconstrained density estimation model When make assumptions about the distribution family, we opt to use multilayer perceptron (MLP) to model our unnormalized density models, with a sigmoid activation function to limit the minimum and maximum output. As we specify in equa- tion g we fit our density models with the sum of the WRI penalty and a negative log penalty. The WRI penalty attempts to fulfill the weighted risk constraint (where learned weighting/ weighting/fuucttons fuuccions a and & satisfy pe, (Xin)) = Pe, (Xme)). Since the total mass not fixed with this model, the log penalty primarily acts as regularization to disincentivize low estimates.
After the density models trained and the unnormalized density estimates computed, scale the den- sity estimates by a constant of proportionality C to obtain proportional density estimates, where C is the scalar that minimizes the mean squared error (MSE) between the estimated and exact densities. We then normalize the MSE between the proportional density estimates and the exact densities, so that an error of one corresponds to a density estimator that predicts a constant (which would reduce the WRI penalty to the VREx penalty). We show how the normalized/relativ MSE of the proportional density estimates goes down train in Figure D.3
2.00
Ch error
C error
1.75
Proportional density relative MSE
constant density
1.50
1.25
1.00
0.75
0.50
0.25
0.00
1000
2000
3000
4000
5000
Training epochs
Figure D.3: Under alternating minimization with an unconstrained density estimation model, the relative MSE of the proportional density estimates goes down over time.
Where are density estimates expected to be accurate? To understand how the accuracy the den- sity estimates is impacted by the support of the environments, consider a data model with two environments: el and e2. The density estimate d is evaluated only on samples from enviroument ei, and similarly the density estimate d is evaluated only samples from environment 62 For this reason, expect the WRI penalty to incentivize accurate density estimates for d where the ground truth density Pe, is high, and for a where Pe, is high. In regions where the density Pe, is small, we expeet the learning of d to be influenced by other factors, including the log penalty and the lendency loward learning smooth functions (due to other regularization, the continuous nature of commonly used activation functions, etc.) The combination of these factors results in the extrapolation of the density estimates outside the region of support overlap, as shown in Figure D.4 This extrapolation does not significantly impact the WRI penalty, as at least one probability density will be low in this region, the penalty will be low well.
26

---

Published in Transactions on Machine Learning Research (07 (07/2024) /024)
0.25
C-del
C-gez
Pst
Pat
0.15
Density
0.10
0.05
-6
-2
Feature ((x)
Figure D.4: We train unconstrained density estimation model for 3000 epochs (to convergence) and compare the proportional density learned for each feature to the ground truth density of the invariant features. We see that for given environment e,, the learned density accurately approximates the shape of the ground truth density where it overlaps with the other environment Cj.
E
Extensions on experiments and additional experimental details
E.1 Toy dataset hyperparameters (Figure 1 details)
define toy dataset consisting of the three environments described here. Environ- ments 1 and 2 are used for training and enviromment 3 is used as test. We let Y = = 1 + > 0), where e is a standard normal random variable and is the indicator function. The following data distributions are used:
To produce Figure !
X lav ~ 2(0,22),
inv ~ N(0,, (t) X² Y (1,22))
X²
x3
inv ~ N(0,33),
X1
~ N(1, ( spm
X spu
= 0 ~ NN-1,11)
spu
x1 Y -V=~~((-L) (1)3). spu
x2 Y N(-1,2 and spu
= N(1,
spu
We sample 10 points from each environment. The predictor a simple lincar model. The ERM, IRM, VREx, and WRI objectives are optimized using scikit-learn (Pedregosa el al. 2011). We use a penally weight of 10a for IRM, a penalty weight of 1 for VREx, and a penalty weight of 1500 for WRL We note that VREx converges to an absurd solution when using larger weight, despite generally using much larger weight in the DomainBed implementation. The penalty weights were roughly optimized by eye to achieve the best qualitative results (regarding invariance).
E.2 Experiments multi-dimensional synthetic dataset
Data generated following structural equation model ulation based on the linear causal model shown in Figure 2 The spurious and invariant distributions are drawn from normal distributions Xiny Nxr,," and X Y ~ N(n"", Eea) and X is the concatenation of Xinv and Xspu- To simulate the classification scenario, we define the class label to be Y = arg max, w Xiav tep where Ey N.0,
We construct multi-class classification sim-
Rather than specifying all of the distribution parameters manually, we sample them randomly according
4º (O, Lo J ~ NN, s
e RandCov(1 Atav. 1+ Amw).
W US
(E.14)
E RandCoo( (1) ))
o = og
27

---

Published in Transactions on Machine Learning Research (07 (07/2022 2024)
(b)
0.6
0.6
0.5
Test Acc
Test Ace
Tost Acc
0.4
0.5
0.5
0.3
0.4
0.2
0.0
0.05
0.15
0.2
0.25
90
87
85
80
70
Ainu
parity
ERM
IRM
VREx
WRI
Figure E.5: Our simulated classification setting from E.2 where we measure the test accuracy under different shifts in the invariant features. (a) We vary the average distance between invariant distributions. (b) We vary the covariance matrices between invariant distributions. (c) With dinv 0.15, Ain 0.5, we vary the correlation of spurious features between the training environments with test having the opposite correlation. In all cases, the WRI predictor outperforms other methods on average.
where Sa.... is the din-ddimensional hypersphere, and RandCov(a, b) is a covariance malrix-geeeralinn random variable that selects the square root of the eigenvalues i.id. from U(a,b) b).
This formulation allows to specify the data simulation using four scalar values: Jnv, Ssuu. Dav: and By. The din and spu terms define the variability of the means between different invariant and spurious distributions respectively. For example, when Oinv is zero, Xinv will have the same mean value across all enviromments. Ainv is a value between 0 and that defines how much the invariant covariance matrices vary between environments, and 0y specifies the noise in the label generation process.
Figure E.5 shows how we compare to other baselines when we run on data with 20 dimensions (dmv 10), with training environments and 1 test. We see that (a) with isotropic covariance, we outperform other baselines when the invariant distribution means overlap and when we increase the distance between the means. Further, (b) with identical means, we outperform other basclines when the invariant distributions all have isotropic covariance and when the covariance matrices vary between distributions. Both cases represent the covariate shift in the invariant features, which our predictor is designed to be more robust to. On average, our improvement over the baselines is more significant when we have isotropic covariance for all environments and increase the distance between the environment mcans; we believe this because the isotropic covariance allows for a clearer (more rolle)) case of covariate shift.
Testing dependence on controlled spurious correlation We run modified version of our simulation that tests the mislcading nature of the spurious features. Similar the ColoredMNIST approach of defining the train and test environments to have opposite color-label correlation, randomly select a vector for each class and ensure that its correlation with the mean of the spurious distributions is positive for each training environment and negative for the test environment. This is accomplished by first sampling the data as in equation E.14 flipping correlation if necessary, then applying a transform to ensure that the means are in a hypercone with opening angle Operity. As 0parity decreases from 90% the spurious data increasingly predicts Y consistently across training environments, making it more challenging to disentangle the spurious and invariant features. In Figure E.5(c), we see that although decreasing parity degrades the performance of all methods, WRI still performs better than the other baselines.
Additional discussion of simulation hyperparameters Here, we further discuss the hyperparameters
Cinv, Any, oy, and aspu: and how they impact the data generation.
inv non-negative scalar that encapsulates the distance between invariant data for each environment. When this value is zero, the invariant distributions are all centered on the origin; when this value is large enough, the average distance between invariant distributions increases.
28

---

Published in Transactions on Machine Learning Research (07/2024)
Ainv is scalar in [0 that encapsulates how much the covariance matrices of the invariant distributions differ between environments. When this value is zero, all environments have identity covariance. Intuitively, this value is similar to dnv in the sense that larger values will produce invariant distributions that overlap less between environments.
By is the only directly specified parameter of the data model and it controls the amount of label noise present. When set zero, the label is deterministic function of Xivvi when set very high, Y becomes more difficult to predict from the invariant data.
Osu is non-negative scalar that encapsulates the distance between spurious data distributions. As this value gets larger, the spurious distributions are less likely to overlap both between classes and between environments. This has two primary effects. First, within a single environment, less overlap will make spurious data more predictive of the label. However, because the locations of the distributions change between environments, this also means that the spurious nature of Xspu should be more apparent during training.
One caveat to our method of data generation is that controlling for uniform number of classes is difficult since the class label is a function of Xinv For this reason, we sample a new set of data model parameters when one class is overrepresented in any environment by a factor of 1.5 times the expected uniform representation (e.g- for 5 classes, we resample when more than 30% of data from any environment have the same label). This also ensures that any models with accuracy above 1.5 times uniform are making non-trivial predictions.
Controlling the spurious correlation For the experiment reported in Figure E.5 (c), we explicitly manipulate the data model to ensure higher level of spurious correlation via the Oparity hyperparameter. This produces a scenario similar to ColoredMNIST, where the correlation between the label and spurious data (color) is relatively consistent between training environments (+90% and +80%) but has the opposite relationship in test (-90%). This means that an algorithm that learns to rely on color will be heavily penalized at test time. In this section, we describe the mechanism behind Qpriity in more detail.
We begin by sampling distribution parameters for all the invariant and conditional spurious datasets (the same as in prior experiments). Next, we select a random unit vector Cy ~ US for each class label. Our goal is to ensure that the spurious mcan n correlates positively with cy for cach of the training environments and negatively with the test environment. In other words, the training and test environments have opposite parity w.r.t. the hyperplane. We achieve this by negating 4 when necessary, i.e.
MM sign(c sigu(cTpuoo))0 yoy sign (ogu(cTn::)) otherwise
if Eer
(E.15)
While this ensures the desired correlation, the relationship is not very strong due to the fact that random vectors in high-dimensional space tend to be nearly orthogonal. Therefore, we introduce 0parity.
For each class, we transform Se' again to ensure it is inside the hypercone with axis c, and opening angle Oparity. This is accomplished by rotating 1 towards the axis so that the angle to the axis is scaled by a factor of @parsty/909 /900. The identity transform is, therefore, synonymous with Sperity = 90%, and as Oparity decreases towards 0 the spurious data becomes more consistently predictive of the label.
Numerical Results
For experiments on the synthetic datasets, we integrate the DomainBed implementa- tions of IRM and VREx. All experiments use 20 dimensional data (10 invariant dimensions and 10 spurious dimensions). They also all use 5 environments (1 test and 4 training) and 5 class labels. The numerical
values visualized in the Figure| E.5 bar graphs are shown in Tables E.1
E.2 and E.3
E.3 Heteroskedastic CMNIST
ColoredMNIST (CMNIST) a dataset proposed by Arjovsky et al. (2019) as a binary classification extension of MNIST where the shapes of the digits are invariant features and the colors of the digits are spurious features that are more tightly correlated to the label than the shapes are, with the correlation between the color and the label being reversed in the test environment. Specifically, digits 0-4 and 5 9 are classes 0 and 1
29

---

Published in Transactions on Machine Learning Research (07/2022 /204)
Table E.1: Sweep diev (Ainv = 0,0y = 0.5, 1). Plotted in Figure E.5 (a)
Fivv 0.15
Algorithm
0.0
0.05
0.1
0.2
0.25
0.3
ERM
51.3 = 1.6
52.5 + 1.1
53.3
+ 0.9
51.5 + 2.0
50.3
+ 1.3
53.7 + 1.4
54.8 + 2.1
IRM
52.4 = 2.7
51.9
+ 2.3
52.2
- 2.7
51.2 = 3.2
49.2
+ 1.8
51.5 = 3.2
54.3 - 3.3
VREx
52.5 = 1.0
51.1
+ 1.5
54.0
+ 1.0
52.4 + 1.4
50.5
+ 1.3
53.6 = 1.3
56.8 + 1.1
WRI
55.9 2.5
55.0
1.3
55.9
2.3
54.4 + 2.3
52.1 3.1
54.1 1.8
56.1 + 2.6
Table E.2: Sweep Dinv (amn
0, Fy = 0.5, epu = 1). Plotted Figure
E.5 (b)
Algorithm
0.0
0.2
0.4
0.6
0.8
ERM
51.3 + 1.6
51.4 + 1.5
51.1
+ 1.2
51.1
+ 1.3
51.7 - 1.2
IRM
52.4 + 2.7
51.3
+ 2.4
51.5
- 2.1
50.6
+ 1.7
51.3 - 1.8
VREx
52.5
1.0
51.6
+ 0.3
51.2
- 0.9
52.0
1.5
52.5 = 1.3
WRI
55.9
+ 2.5
52.3 + 1.4
52.9 = 1.6
53.3 + 0.7
53.8 = 2.5
Table E.3: Sweep Qparity (oiv 0.15, A..
0.5, ov O.5, Japa
1). Plotted in Figure E.5
i parity 85
Algorithm
90%
87
80
70
ERM
44.8 1.2
43.0 + 1.4
40.5 - 1.1
37.5
0.7
30.5 = 0.9
IRM
43.8
3.6
41.9
3.9
40.9
- 4.2
37.0
4.0
29.9
3.0
VREx
45.2 + 2.0
42.9
1.9
41.4
= 1.6
37.2
+ 2.7
30.5 1.3
WRI
46.9 + 1.5
45.4 + 1.9
43.9 + 2.6
39.5 + 3.2
30.9 - 2.2
30

---

Published in Transactions on Machine Learning Research (07 (204))
3
6
6
2
5
9
8
7
d
P
f
Figure E.6: Examples of CMNIST in-distribution (top) and out-of-distribution (bottom) used for OOD detection. The out-of-disibbution samples are generated by flipping a subset of digits horizontally or vertically.
respectively; after injecting 25% label noise to all environments, the red digits in the first two environments have an 80% and 90% chance of belonging to class 0, while the red digits in the third environment have 10% chance. The design of the dataset allows invariant predictors (that only base their predictions on the shape of the digits) to outperform predictors that use spurious color information, ideally achieving 75% accuracy on all environments.
Hetcroskedastic CMNIST with Covariate Shift Variants of CMNIST have been proposed, with the aim of incorporating additional forms of distribution shift (Ahuja et al. 2020; Wu et al. 2020 Krueger et al. 2021). To demonstrate the efficacy of our method under heteroskedastic covariate shift, we construct a heteroskedastic variant of CMNIST, which we call HCMNIST, where label flips occur for digits 0, 1, 5, and 6 with probability 5% and label flips occur for other digits with probability 25%. Additionally, we generate a variant of HCMNIST with covariate shift, HCMNIST-CS, where we redistribute data among environments. We place 65% of the digits 0, 1, 5, and 6 into the first training environment and 5% into the second enviromment (with the remaining 30% in test). The remaining digits are distributed so that all enviromments have the same number of samples.
The empirical results on the HCMNIST with and without covariate shift are shown in Table 2 These experiments use the practical version of the WRI algorithm and are computed using the same evaluation strategy as DomainBed. The results demonstrate that the practical WRI implementation performs well in both the heterogencous case with and without covariate shift. Conversely, both IRM and VREx have significant degradation when covariate shift is introduced.
We also create an idealized version of these datasets that simplifies the images into two-dimensional features consisting of the digit value and the color. We evaluate the WRI penalty and VREx penalty using three optimal predictors: one that operates on only the digit value (invariant), one that operates on the color value (spurious), and one that operates on both. The results of this experiment are presented Table 1 As expected, we find that both WRI and VREx have zero penalty on HCMNIST. However, in the presence of covariate shift, the VREx penalty for the invariant predictor is greater than both the penalty for the spurious classifier and the penalty for the mixed classifier. This demonstrates simple and concrete where VREx does not select for an invariant classifier, suggesting that the degradation in VREx performance under shift in Table 2 can be attributed to true failure of VREx to recover the invariant predictor.
E.4 Out-of-distribution Detection with CMNIST
To evaluate the out-of-distribution (OOD) detection performance CMNIST, we first train model the two training environments of (traditional) CMNIST, then assess the model's ability to classify sample digits as in-domain or out-of-domain on modified version of the CMNIST test environment. The modified test environment is a mix of unaltered digits for the "in-domain" samples and flipped digits for the "out-of- domain' samples. Specifically, we horizontally flip digits 3, 4, 7, and 9, and vertically flip digits 4, 6, 7, and 9. In this way, we create CMNIST samples that would not exist in the training distribution but still appear plausible. FigureE.6) provides examples of these in-domain and out-of-domain samples. Note that we only flip digits that are unlikely to be confused with real digits after flipping.
31

---

Published in Transactions on Machine Learning Research (07 (274))
In order to measure OOD performance, require an OOD score for each sample that predicts if the sample is in-distribution. Fach OOD score is then compared against its in-distribution label to compute Receiver Operating Characteristic (ROC) curve. This curve captures the trade-off between the true positive rate and false positive rate at various confidence thresholds. Lacking an explicit density estimate for ERM, IRM, and VREx, we instead use the maximum class prediction score for the OOD score. We then compare the effectiveness of WRI's density estimate as an OOD score against ERM, IRM, and VREx. Finally, we calculate the area under the ROC curve (AUROC) an aggregate metric. The results, presented in Table indicate that the WRI density estimates more accurately detect OOD digits than the prediction confidence scores from the other methods.
E.5
DomainBed experimental details and results
E.5.1
Featurizer
We use a ResNet50 model to train featurizers using the default ERM DomainBed parameters for each dataset. We train one featurizer corresponding to each test enviromment with both 32 and 64 dimensional features. Each featurizer is trained on the training environments for a fixed number of steps. Once trained, the features are pre-computed and all experiments are then performed on these features. The dimensionality of the features used is selected randomly as a hyperparameter. For cach experiment, we use single hidden-layer MLP with ReLU activation which operates directly on ResNet features.
E.5.2 Hyperparameters
Since we are training on pre-extracted features, running for 5000 steps is unnecessary. Instead, we reduce the total number of steps per experiment to 500 and increase the default learning rate to 1e-2. Because we have fewer steps than the original (non-featurized) DomainBed implementation, we reduce the annealing iterations on IRM and VREx by a factor of 10 and limit the maximum anncaling steps to 500. Otherwise, VREx and IRM would nearly always finish training before annealing has finished. Additionally, we define + to be an identity function,
found this to produce the most consistent results.
We specified the following DomainBed hyperparameters and selected them according to the following dis- tributions. Note that in Bed, the default hyperparameters are used during the first hyperparameter seed, and random hyperparameters are selected for subsequent seeds. For any hyperparameters not listed here, we use the defaults provided by DomainBed.
Learning-rate default: 1e-2, random: log uniform over [5c-3, 1e-1].
IRM-anneal iters default 50, random: log-uniform over [1, 500]. Number of steps before penalty weight is increased.
VREx-anneal iters - default 50, random: loguurformm over [1, 500]. Number of steps before penalty weight increasod.
Featurizer dimensions default: 64, random: discrete uniform from (32,64), The dimensionality of the pretrained features used.
WRI-A default: 1, random: log-uniform over [le-1, 5e1]. The penalty weight used when computing in the predictor optimization step.
WRI-annealing- default: 0, random: discrete uniform from {0, 10). Number of steps before WRI
regularization term is included in loss function.
WRI-density update freq (w) default: 1, random: discrete uniform from {1, 2,24488, Number of predictor optimization steps between density optimization steps.
WRI-density learning rate- default: 20-2, random: log-uniform over [1o-2, 5c-2]. Learning rate used for the density estimate optimizer.
32

---

Published in Transactions on Machine Learning Research (07/2024)
WRI-density weight decay default: 1e-5, random: log- ungguniform over [1e-6, 1c-2]. The weight decay used for the density estimate optimizer.
WRI-density batch size default: 256, random: discrete uniform from {128, 256). The batch size used when optimizing for densily estimates.
WRI-density X default: 5, random:
log-uniform from [5 50). The penalty weight used when computing in the density optimization step.
WRI-density default: 2e2, random: log-uniform from [5e-2.5]. The negative log penalty weight used when computing L in the density optimization step.
WRI-density steps (na) default: 4, random: discrete uniform from (4,16,32) 32). The number of optimization steps taken each time the density estimators updated.
WRI-min density default: 0.05, random: uniform from [0.01, The minimum density imposed via scaled shifted sigmoid activation on density estimator models.
WRI-max density- default: 1, random: uniform from [0.4, 2J. The maxnimum density imposed via scaled shifted sigmoid activation on density estimator models.
E.5.3
DomainBed with additional baselines
We run DomainBed experiments on additional baselines to place our work in larger context, expanding Table to Table E.4 Specifically, we also compare to GroupDRO (Sgaaaa et al. 2019 Mixup (Zhang et al. 2017 MLDG Li et al. 2018), and CORAL (Sun & Saenko 2016). These are the (current) top- performing methods that are implemented and tested in the Domain Bed test suite that can be used with an ERM-trained fcaturizer- as that is sufficient learn the invariant features nocessary train an OOD predictor [Rosenfeld al.) 2022), but also significantly speeds up training. We did not include methods that are not implemented in DomainBed, but it is worth mentioning examples like MIRO (Cha et al. 2022) that demonstrate stato-of-the-art performance are different lines of work; they are not causally motivated, but can sometimes have better generalization accuracy than methods with a causal basis.
image data. Note that all of the aforementioned methods
Table E.4: DomainBed results feature data, with additional non-causal baselines
Algorithm
VLCS
PACS
OfficeHome
Terralncognita
DomainNet
Avg
ERM
76.5 + 0.2
84.7 + 0.1
64.5 + 0.1
51.2 +: 0.2
33.5 + 0.1
62.0
IRM
76.7 = 0.3
84.7 + 0.3
63.8 = 0.6
52.8 + 0.3
22.7 2.8
60.1
GroupDRO
77.0 = 0.2
84.8 + 0.1
65.1 = 0.1
51.3 + 0.2
30.9 0.1
61.9
Mixup
76.6 = 0.2
85.0 0.1
66.1 = 0.0
52.3 - 0.7
33.1 0.1
62.6
MLDG
75.8 + 0.2
82.3 + 0.2
64.8 + 0.2
48.3 + 0.2
33.0 0.1
60.8
CORAL
76.5 - 0.2
85.1 1 0.2
65.0 + 0.1
51.8 - 0.4
33.5 0.1
62.5
VREx
76.7 + 0.2
84.8 + 0.2
64.6 + 0.2
52.2 + 0.3
26.6 + 2.1
61.0
WRI
77.0 0.1
85.2 0.1
64.5 + 0.2
52.7
0.3
32.8 0.0
62.5
E.5.4 Individual dataset results
This section contains the individual Domainßed results on VLCS, PACS, OfficeHome,, Terralncognita, and Net The Average column from each dataset table is reported in Table E.4 For all other dataset tables, the column labels indicates which environment was held out for test.
VLCS
33

---

Published in Transactions on Machine Learning Research (07 (07/2024)
A Agoorithm
C
S
V
Avg
ERM
97.2 + 0.1
66.0 = 0.6
68.9 + 0.5
73.8 + 0.5
76.5
IRM
97.0 + 0.0
66.3 = 0.5
69.1 0.5
74.2 + 0.2
76.7
GroupDRO
96.9 + 0.2
67.1 + 0.4
09.3 0.2
74.6 + 0.3
77.0
Mixup
97.3
66.7 0.4
7 0.5
76.6
MLDG
97.2 0.1
63.9 0.6
0.3
13.7=0.3
75.8
CORAL
96.9 + 0.2
66.3 - 0.2
+
73.9 +.3
76.5
VREx
97.2 - 0.1
66.9 = 0.0
6
0.4
76.7
WRI
97.1 + 0.2
67.0 = 0.2
69.6 0.6
74.3 + 0.2
77.0
PACS
Algorithm
A
C
P
S
Avg
ERM
80.5 + 0.3
81.5 + 0.2
96.6 + 0.1
0.1
80.2 0.1
84.7
IRM
80.8 + 0.9
82.0 0.3
96.3
79.9
0.3
84.7
GroupDRO
81.4 = 0.3
81.2 0.3
8022 9.2
84.8
Mixup
80.9 + 0.3
81.7 0.3
96.6
80.8
85.0
MLDG
80.0 = 0.5
81.1 = 0.4
95.7 0.2
72.2
82.3
CORAL
81.4 + 0.2
81.8 + 0.4
96.8 0.1
80.5 + 0.2
85.1
VREx
80.6 + 0.5
81.8 = 0.7
96.8
- 0.1
80.0 = 0.2
84.8
WRI
81.2 + 0.3
8233 = 0.1
96.5 + 0.2
0.2
80.7 0.1
85.2
OfficeHome
Algorithm
A
P
R
Avg
ERM
58.9 + 0.5
50.2 + 0.3
74.4 + 0.2
74.5 + 0.4
64.5
IRM
57.9 + 0.6
49.8 # 0.5
73.4 0.8
74.1
63.8
GroupDRO
59.7 0.3
50.7 0.2
74.8 0.1
75.4
65.1
Mixup
61.1
52.2 0.1
75.6 0.1
7
66.1
MLDG
59.1 0.4
50.7 2 0.2
74.8 4 0.1
74.8 0.3
64.8
CORAL
59.3 0.2
50.6
74.9 0.2
75.2
0.2
65.0
VREx
58.8 0.5
50.4 = 0.3
74.4 0.2
74.8 + 0.4
64.6
WRI
58.9 + 0.6
49.8 = 0.2
74.7 + 0.1
74.7 + 0.3
64.5
TerraIncognita
Algorithm
L100
L38
L43
L46
Avg
ERM
50.3 + 0.4
50.0 + 0.6
57.8 0.3
46.6 + 0.4
51.2
IRM
53.6 + 0.7
53.2 2 0.8
57.7 0.2
46.6 - 0.8
52.8
GroupDRO
50.4 + 0.5
50.1 0.3
+7.4
47.3 + 0.3
51.3
Mixup
53.1 1.8
52.4 0.9
57.1
46.6 1.6.03 0.3
52.3
MLDG
44.7 0.6
49.3 2 0.3
57.2 0.2
41.8 - 0.3
48.3
CORAL
51.0 0.5
51.5 1.0
57.8 0.2
47.0 0.5
51.8
VREx
52.5 + 1.1
51.2 - 0.6
57.9 7.9
47.2 0.4
52.2
WRI
51.7 + 0.5
55.0 + 0.6
57.2 0.5
47.1 + 0.3
52.7
Nomain
34

---

Published in Transactions on Machine Learning Research (07/2024)
Algorithm
clip
info
paint
quick
real
sketch
Avg
ERM
47.9 + 0.2
16.3 + 0.1
40.6 + 0.2
9.6 + 0.2
46.7 + 0.2
40.2 + 0.1
33.5
IRM
31.9 + 4.4
11.4 + 1.3
28.9 + 3.0
6.5 = 0.7
30.8 + 3.9
26.9 + 3.7
22.7
GroupDRO
43.8 + 0.3
15.8 + 0.2
37.5 + 0.3
8.8 + 0.1
42.6 0.2
37.1 + 0.3
30.9
Mixup
47.2 + 0.1
16.3
0.2
40.2 + 0.2
9.1
45.5 0.3
39.6 - 0.2
33.1
MLDG
46.9 0.2
16.2
- 0.1
40.0
9.2 0.1
46.2 0.1
39.4 = 0.1
33.0
CORAL
47.9 0.3
16.5
0.1
40.5 0.2
0.1
46.6 +
39.8
0.2
33.5
VREx
37.1
3.2
14.0
1 0.9
31.9
7.5 0.7
37.2 - 2.5
32.0 2.5
26.6
WRI
46.7
+ 0.2
15.9 + 0.1
39.8 0.1
9.3 + 0.1
46.0 + 0.2
39.4
0.1
32.8
35


Reviews:
[{"review_id": "mnZFWX5GHx", "reviewer": "Reviewer_RbPh", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths:\n* The writing is generally quite good. The paper provides a lot of extra information in the appendix, and a good amount of effort has gone into relating the what has been proposed in this paper with previous work.\n* The experiments are focused on understanding rather than trying to beat benchmarks, and as a result the paper is quite interesting to read.\n\n## Weaknesses:\n* It's difficult to see how the specific causal structure assumed in this work maps onto real-world situations. The motivation and scope of the paper would be easier to understand if the reader was given an example problem where this structure is present. I would also be interested in an intuitive example to compare and contrast the relationship of invariant covariate shift with regular covariate shift.\n* It's not clear why we should expect optimising the proposed objectives in either of Equation 6 or Equation 9 to result in accurate feature densities. The derivation only explains how one can obtain robust models if one already has accurate densities, not the other way around. The experimental results (Table 3) indicate that the learned densities, while more useful that other methods, are still of low quality.\n* Appendix B does not provide sample complexity results; in order to apply Bernstein's inequality in this context one must assume that the model is independent of the training data, which is not true when the data is used to train the model. Frameworks based on, e.g., covering numbers have been developed to get around this. The union bound should also be used to combine the two applications of Bernstein inequality used to obtain B.14, since these two separate bounds only hold with high probability. Assuming this is fixed, it should also be made explicit that this convergence rate does not necessarily apply to the proposed algorithm, since this bound assumes optimal weighting functions are known. Moreover, the conclusion that one should select densities with lower maximum values and lower variances does not make sense; surely the ideal densities are the true densities of the underlying data distributions? Also, if the densities are probability densities, reducing the maximum would require increasing the variance in order for the density to remain normalised.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This submission proposes a method for learning invariant representations in the presence of covariate shift. A causal data generation process is proposed and an idealised objective is derived from this. A heuristic agorithm for minimising this objective given only a finite sample is prposed, and experimentally validated on some domain generalisation problems. There are also some additional experiments that provide some corroboration to other claims made along the way.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* Provide an example problem where the causal structure assumed in this work is present.\n* Provide an intuitive example to compare and contrast the relationship of invariant covariate shift with regular covariate shift.\n* Explain why we should expect optimising the proposed objectives in either of Equation 6 or Equation 9 to result in accurate feature densities. Alternatively, be upfront that there are no guarantees and quantify experimentally the extent to which alternating optimisation on the heuristic density estimation objective does find good approximations of the (un-normalised) densities.\n* Fix the sample complexity results in Appendix B, or remove them.\n* State more clearly some of the assumptions about how different environments must be related in order for this method to work. E.g., should all environments have the same support, as is typically required for density ratio methods?", "broader_impact_concerns": "The work is mainly theoretical in general, but the topic (robust generalisation) is more likely to lead to a positive broader impact than negative, in my view. The authors include some discussion on broader impact concerns in their paper.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "NjpNN0GhB8", "reviewer": "Reviewer_d8hB", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. This paper is well written with clear structures. On the theoretical side, it is explicit about the assumptions, such as causal prediction coupling and spurious-free predictor.\n2. It compares with VREx and IRM and discuss the similarity and difference.\n3. Based on weighted invariance this paper proposes a new loss function and compares the algorithm on variant benchmarks such as Colored MNIST and DomainBed.\n\nWeaknesses:\n1. I still have some questions regarding the assumptions. As mentioned in the paper, most common assumptions in DG are covariant shifts, label shifts, and invariant predictors. However, this paper is based on a new assumption called causal prediction coupling. Why is this assumption valid and practical? Or could this be the reason why it doesn't perform significantly better in practice? Also, the linear causal setting seems limited as well. How does WRI perform in non-linear settings? \n2. Figure 1 seems confusing. Is this figure to demonstrate that WRI is better than other algorithms, and in which sense? \n3. Although Prop. 1 seems intriguing for a new principle of domain invariance, it might be a weak condition. As the authors mentioned, a trivial solution can be $\\alpha = \\beta = 0$. Is there a way to say that the WRI condition is tight or sufficient as well (in general)? Or could the authors share some intuition?\n4. In order to enforce WRI, in eq. (6) the authors need to minimize all pairs of environments. This quadratic scaling makes it infeasible for a large number of envs. Could the authors show some comparison of the computation time with a large number of environments?\n5. Sec 3.3 only compares with REx and IRM in some special examples, but fails to prove the advantage in general.\n6. In Sec 3.4, a log density loss was added ($\\log d^e$). Could the authors explain why it is necessary or theoretically justify it? Some convergence proof might be helpful.\n7. The main experiments only compare with IRM, ERM and VREx. If we look at the appendix, WRI is in fact not better than simpler algorithms like CORAL. This minor improvement might show that the underlying assumption (causal prediction coupling) does not hold in general.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a new algorithm for domain generalization, called weighted risk invariance (WRI). The objective is to minimize some type of variance after reweighting with the environment distributions.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See above.", "broader_impact_concerns": "The broader impact is properly discussed in Sec 6.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "SMXmcWUtnX", "reviewer": "Reviewer_svP3", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "(+) The problem setting is novel, and is presented clearly.\n\n(+) The discussion and comparison with previous prevalent methods such as REx and IRM are comprehensive.\n\n(+) Empirical studies demonstrate certain advantages of WRI;\n\n(-) Despite a novel problem setting, it's unclear whether it's realistic and feasible;\n\n(-) There is no guarantee for the solution;\n\n(-) Some important related haven't been compared or discussed;", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper focuses on domain generalization under the covariate shifts of invariant features. The authors argue that it needs to also consider the density ratios of the underlying invariant features in order to generalize better under the considered distribution shifts. Then, they propose a new framework called weighted risk invariance (WRI), that learns environment weights based on invariant features, and optimizes for invariance under the weighted environment risks. They discuss the difference of WRI with IRM and REx, as well as its implementation. They demonstrate the effectiveness of WRI through experiments on ColoredMNIST and DomainBed.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Despite a novel problem setting, it's unclear:\n- whether the assumptions about the invariant features are realistic. It'd be better to provide some illustrative examples, case studies as well as demonstrations in the benchmarks. For example, do the invariant features really shift in some datasets from DomainBed?\n- whether the problem is feasible. It's already been argued that the weights for covariate shifts are not unique. Given that, is there really a solution for the proposed problem setting? Or why can previous approaches for covariate shifts/concept shifts/subpopulation shifts/invariant learning not solve it? Can we solve it without any additional assumptions?\n\n2. There is no guarantee for the solution. Meanwhile, critical details such as how to find the weights (i.e., calculating $R^{e_i,e_j}(f)$) have not been well elaborated in the main text. Since a proper estimation of the weights depends on the identification of the invariant features, \nwould not the problem be a \"the chicken-and-egg conundrum\"?\n\n3. Some important related haven't been compared or discussed. For example, [1,2,3,4] also discuss the drawbacks of the invariant learning framework and propose new solutions. [5] generalize the importance weighting to broader distribution shifts.\n\n**References**\n[1] Does Invariant Risk Minimization Capture Invariance? AISTATS'21. \n\n[2] Fishr: Invariant Gradient Variances for Out-of-Distribution Generalization, ICML'22.\n\n[3] Pareto Invariant Risk Minimization, ICLR'23.\n\n[4] Understanding and Improving Feature Learning for Out-of-Distribution Generalization, NeurIPS'23.\n\n[5] Generalizing Importance Weighting to A Universal Solver for Distribution Shift Problems, NeurIPS'23.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work combined two directions in domain generalization: covariate shift and invariant risk minimization, by allowing the invariant features to go through a (marginal) distribution shift among different domains (environments). This new problem setting, albeit a bit straightforward, could be interesting in some applications. The proposed algorithm is also a combination of existing ideas in covariate shift (density weighting) and invariant risk minimization (variance regularization). The chicken-egg problem raised by reviewer svP3 was addressed by alternating minimization, which, albeit lack of formal guarantee, seems to be a reasonable workaround and shows some promise in the authors' experiments. Overall the reviewers agreed that the problem setting is interesting. There are some doubts on the comparison against stronger baselines such as CORAL and GroupDRO, but I think the authors have done reasonable justification of their methodology through experiments on both synthetic examples and real datasets. 

Some minor suggestions: 

(a). For Eq (9), instead of saying the last log term encourages high entropy (why is this a good thing?), why not just point out we are performing maximum likelihood density estimation for each environment (in the feature space)? It would be nice to comment on the role of each term in Eq (9). For instance, without the first ERM term, the result could be trivial? 

(b). In comparison to VREx (7), it would be nice to point out the proposed method in Eq (6) is exactly Var(R^{e_i}(f)), so the main difference simply lies in the weighting in the risk R^{e_i}(f): In particular, if one uses uniform weighting we recover VREx. The authors seem to know this: the reply to Reviewer d8hB (On quadratic scaling) perhaps is also worth incorporating into the final version. 

(c). "In Appendix C, we show that ... WRI on the joint features (rather than the invariant features) still leads to learning invariant feature weighting functions." Appendix C is an experiment on a toy dataset, so I would suggest changing "show" to "empirically show" in this claim.
Decision: Accept (poster)