Weighted Risk Invariance for Density-Aware Domain Generalization

Weaknesses (cont'd)

4. Before addressing this concern, we would first like to clarify some of our results on the DomainBed test suite. DomainBed comprises 5 real world datasets: VLCS, PACS, OfficeHome, TerraIncognita, and DomainNet. These datasets all include a variety of image styles and natural shifts. Generalizability is not only measured by performance on individual datasets, but also on the datasets on average. Over all the datasets, WRI performs the best on average, and also outperforms ERM on all datasets but OfficeHome (where it has the same average accuracy) and DomainNet. In all cases but DomainNet, WRI is also the clear first by outperforming the other methods by more than a standard deviation, or else it is within a standard deviation of the best method. Except for ERM on DomainNet, this cannot be said for any other method/dataset combination.

   That being said, we do recognize the reviewer’s concerns about the similar performance levels. We would like to emphasize that in addition to achieving competitive prediction accuracy, our method is unique in that it also trains density estimators for invariant features, and these allow users to get additional information at test time. We demonstrate empirically that this density is useful for OOD detection, but it should be useful for detecting rare/minority examples in general. This is an additional signal that other invariant prediction methods do not provide.

We thank the reviewer for their thoughtful engagement. We are glad that our approach seems both interesting and novel, and we are grateful for the reviewer’s recognition of the work that went into the theoretical justification. We hope that by answering the reviewer’s concerns, we can better clarify our technical contributions.

Weaknesses

1. We wish to clarify two points in this regard:

   a. Our theory shows that under the model we consider, WRI learns a model that only relies on $X_{inv}$. This is the required condition for OOD generalization in this case, and is the goal of our method and implementation.

   Our goal is not causal discovery, and hence we do not wish to prove that some representation obtained by the model corresponds explicitly to $X_{inv}$. This is in line with many other works on learning invariant representations, e.g. IRM and many more follow-up works [1-3]. This is different from works such ICP [4], that indeed recover causal features and therefore perform causal discovery.

   b. An important distinction between performing invariant representation learning and performing causal discovery, is that generally invariant features are not guaranteed to be causes of $Y$. This is also a distinction made in prior work, e.g. in Wald et al. [5] (Fig. 1, Fig. 2a and corresponding analysis). Our method can also be applied in the data generating processes studied there, where $X_{inv}$ is an anti-causal feature that is a stable predictor under distribution shifts.

   We added a clarification about this distinction between causal discovery and invariant prediction to our discussion section, with thanks for the valuable input.

2. We thank the reviewer for their insight on this important point. We corrected our definition of an invariant predictor: now, we more precisely define a predictor that only uses invariant features as spurious-free, and we clarified the connection to the standard definition of invariant predictor as mentioned by the reviewer. We also modified our statement of Proposition 1 accordingly.

3. We are confident in the accuracy of the claim and proof, and are open to answering any questions about the proof. We see how the claim may seem less intuitive due to the point regarding the meaning of the term invariance, which we discussed earlier in the response. However, we believe that using the term "spurious free", as we do in the updated paper, makes the claim much more intuitive (we emphasize that the claim is still mathematically correct according to the definitions we used in the earlier version as well, it is just the name of the mathematical property that we updated here).

   To explain why the claim holds intuitively, please note that if we write down the expected risk of a classifier that is only a function of the invariant features, $f(X_{inv})$, then the spurious features can be marginalized out and the risk is just the expected risk under the marginal distribution of $X_{inv}$ and $Y$. Because reweighting is done only on the invariant features, then the marginal distribution over $X_{inv}$ does not shift between environments and then the risk of any classifier that only depends on $X_{inv}$ is fixed across them. We hope that this clarifies the proof.

   The intuition behind the weighted risk in Definition 2 is to adjust samples so that environments with different distributions of invariant features can be compared. This is a similar concept to importance reweighting [6], but with adjustments to favor invariant learning. The weighted risk between two environments is asymmetric, and we use the phrase “weighed risk invariance” to describe when the risk of environment A, weighted by the density from environment B, is equal to the risk of environment B weighted by the density from environment A.

   We can verify that weighted risk invariance holds using the predictor and density functions we learn after training with eq 9. For an environment A, the predictor is used to compute the loss at each point, and the density function trained for another environment B is used to weight the losses. We then take the average over all the weighted losses at each point to approximate the risk of environment A weighted by the density from environment B. We flip the environments to compute the risk of environment B weighted by the density from A, and these quantities should be equal.

[1] Jiang and Veitch (2022). Invariant and Transportable Representations.

[2] Mitrovic et al (2021). Representation Learning via Invariant Causal Mechanisms

[3] Rojas-Carulla et al (2018). Invariant Models for Causal Transfer Learning

[4] Peters et al (2015). Causal inference using invariant prediction

[5] Wald et al (2021). On calibration and out-of-domain generalization

[6] Shimodaira (2000). Improving predictive inference under covariate shift.

Thank you for taking the time to read and review our submission. We would be happy to answer any remaining questions before the discussion period ends tomorrow.

We thank the reviewer for their thoughtful engagement, and we are glad that the reviewer finds our approach theoretically sound and interesting. However, we feel there may be a misunderstanding of some of our contributions.

Weaknesses

1. To answer this comment, we wish to lay out the following points:

   a. Heteroskedasticity in our setting is exactly the same as defined in REx. Under heteroskedasticity, different values of $X$ have different levels of noise. We see that perhaps our use of the phrase "heteroskedastic covariate shift" (in paragraph 4 of the Introduction and the Conclusion) may be misleading in that sense, and we changed that phrasing to clarify that our meaning is exactly the same as in REx.

   b. Covariate shift is a completely distinct notion from heteroskedasticity. However, our point is that under heteroskedasticity, covariate shift leads the invariant classifier to have different risks across environments. Thus, REx will not learn a correct classifier.

   c. The above point is recognized in the REx paper, where the authors give the following explanation:

   REx…performs poorly in the domain-heteroskedastic case….Intuitively, this is because the irreducible error varies across domains in these tasks, meaning that the risk will be larger on some domains than others.

   We offer this explanation via example in “Comparison to REx”. Following your suggestion, we added details to make the example more explicit.

   d. We would appreciate further clarification on what you mean by varying noise variance for $Y$ across environments. If $Y$'s variance changes across environments, then $p_e(Y \mid X_{inv})$ is not consistent across environments, hence the overwhelming majority of invariant learning algorithms should not learn a classifier proportional to $p_e(Y \mid X_{inv})$. Note that this setting is different than heteroskedastic noise, where noise levels may change across $X$ but are fixed across environments (i.e. $p_e(Y \mid X_{inv})$ is the same for all $e$, but $Y$ may have varying noise levels across values of $X_{inv}$).

2. In “Comparison to IRM”, we show that as the degree of covariate shift between two environments increases, then the IRM penalty becomes less sensitive to invariance violations, resulting in increased sample complexity. The sample complexity issues of IRM are already extensively discussed in previous work, often where the entire paper is focused on analyzing the failure cases of IRM [1-2]. The focus of our paper is on the potential benefit of invariance weighting in invariant learning, so another proof of the sample complexity of IRM would be out of scope.

   We thank the reviewer for the recommendation for Figure 4, but we suggest that adding REx to Figure 4 may be misleading. Depending on how difficult examples are distributed across the two environments, the REx penalty may be misaligned with recovering an invariant predictor (for reasons discussed above). Thus, the magnitude of the REx penalty may not even be relevant to its sensitivity to invariance violations.

3. The reason we employ alternating minimization is precisely due to the difficulty of identifying the invariant features. Because we do not have access to these, we cannot learn their density by directly optimizing equation 9. Instead, we rely on the observation that if $d(X) = p_e(X_{inv})$, then an invariant predictor would have zero WRI penalty, and if we have invariant features, then we can perform density estimation on those features. With the density regularization term, we find that we do not recover any other trivial solutions in practice.

4. Causally-motivated methods have the benefit of interpretability and sound theoretical grounding. As a trade off, they can be difficult to scale to complex, high dimensional data, and therefore cannot always achieve the generalization accuracy of non-causal methods. We do reference and compare to several of these other methods to place our work in context (in the section of the Appendix you mention). Overall, we believe that pushing forward the state of the art in causally-motivated methods is, in itself, a significant contribution.

5. To address this, we created Appendix section C.4 to include a more detailed explanation of the modified CMNIST test split; this should better clarify how we create examples that are OOD in the invariant feature space. We thank the reviewer for the suggestion.

6. We thank the reviewer for the references, which we have added to the discussion section. We do want to emphasize that these works are about general covariate shift, and do not discuss the particular setting we are dealing with here, which is covariate shift over latent invariant features.

[1] Ahuja et al (2021). Empirical or Invariant Risk Minimization?

[2] Rosenfeld et al (2021). Risks of Invariant Risk Minimization

Thank you for taking the time to read and review our submission. We would be happy to answer any remaining questions before the discussion period ends tomorrow.

Regarding how IRMv1 and other invariant learning methods recover $P(Y \mid X_{inv})$. In the specific case of IRMv1, they define invariance as $\mathbb{E}\_{p_e}[Y \mid \phi(X)] = \mathbb{E}\_{p_{e'}}[Y \mid \phi(X)]$ for any two environments $p_e$ and $p_{e'}$, instead of the full conditional probabilities ($p_e[Y \mid \phi(X)] = p_{e'}[Y \mid \phi(X)]$). Indeed in their model the conditional mean is invariant, yet as you suggested, the conditional distributions aren't. Since minimizing the squared loss only involves finding the conditional mean, their analysis works out in that case and it is guaranteed to learn a predictor that is free of spurious correlations.

Hence, they just use a weaker form of conditional independence, whereas subsequent work discusses many forms of such conditional "independencies", showing cases where these are appropriate. For instance, [1] also match the variance instead of the mean, while [2] and many others seek to match the full conditional probabilities (i.e. having $p_e[Y \mid \phi(X)] = p_{e'}[Y \mid \phi(X)]$). Invariance of the loss, which is what we study, is discussed in many other papers that are cited in our work, and conditions are drawn for cases where this is a sufficient condition to discard spurious correlations. In our model, like many other works such as ICP, REx, and others cited in our paper, we do assume that the noise level is constant across environments, and that $p(Y \mid X_{inv})$ is fixed across environments. We should also note that we choose to focus on risk invariance since it is a popular type of invariance to study, and the theory and main ideas in our paper can likely be applied with other invariance principles.

Regarding the discussion on how our objective and others recover invariant features. We agree that this can be better clarified, and we have modified our implementation discussion (Appendix B) to give a more detailed discussion on this. The easiest way to explain this is perhaps starting from a constrained optimization problem, where weighted risk invariance is enforced as a constraint, and the objective is to minimize the training loss. Then the theoretical statement says that models with zero weights on spurious features will be the ones satisfying the constraint. Optimizing the loss chooses the best performing one among those, which under the conditions of the model is just the optimal invariant predictor (corresponding to the conditional mean when we consider the squared loss or logistic loss for classification---note that trivial solutions that satisfy the constraint but do not minimize empirical loss will not be selected). Then our objective is simply a penalized version of this constrained problem. This largely follows a similar path to works like IRM and many of its follow-ups (and more generally works that approximate constrained problems with a penalty), hence we chose to focus on the properties of WRI in our submission as this type of development is not a contribution of our work. However, we see how in the process we lost some of the accessibility of our arguments to audiences that are not familiar with exactly this type of argumentation. Hence we have added the aforementioned discussion to our appendix and we will emphasize it in the main paper. Thanks very much for your engagement and your help in improving our paper.

[1] Bellot, Alexis, and Mihaela van der Schaar. (2020) "Accounting for unobserved confounding in domain generalization." arXiv preprint arXiv:2007.10653.

[2] Puli, Aahlad Manas, et al. (2021) "Out-of-distribution Generalization in the Presence of Nuisance-Induced Spurious Correlations." International Conference on Learning Representations.

We thank the reviewer for their thoughtful engagement. We are glad that the reviewer finds the proposal well formulated, novel, and non-trivial.

Weaknesses

1. We combined the introduction and comparison to VREx and IRM in our latest version. We appreciate the suggestion to improve the readability of our paper.

2. The use of density estimates is a key benefit of our work. Following your suggestion, we added Appendix C.4 to include more technical detail. We now spend more time discussing how we construct our OOD samples, as well as how our density estimates are used as an OOD score (in line with current OOD literature [1]).

3. We agree with the reviewer that theory for the nonlinear case would be ideal. Our current theory revolves around systems of quadratic equations having unique solutions, and a nonlinear generalization would result in a different system of nonlinear equations that seems quite difficult to analyze and is beyond the scope of this paper. We wish to emphasize that existing literature only provides theoretical justification for either the linear case [2-4] or for highly specific nonlinear distributions under additional assumptions [5]. Many works in the area also do not provide theoretical justification at all. Therefore, we believe that a generalized theory for nonlinear models would be a truly seminal contribution which deserves its own work. In this work, our focus is on the effects of covariate shift in the invariant features, which we study both empirically and theoretically within the framework that is currently common in the literature on invariant learning.

   In context with current literature, we believe our formulation is still a significant contribution. We first provide theoretical and empirical evidence for the linear model, then empirical evidence for the nonlinear model where some version of our causal structure holds (CMNIST), and finally empirical evidence on DomainBed, a benchmark of highly unstructured datasets with unknown shifts. In this way, we hope to form a coherent narrative from our theory to our experiments.

4. The practical objective aligns closely with the WRI objective: both use an ERM term (to encourage good performance across environments) and a WRI term (to encourage that we learn to predict on only the invariant features). The main difference between the two is in access to the density of the invariant features, so the practical objective additionally includes density as a learnable parameter and a regularization term for density estimation. Following your concern, we updated the paper to make the similarity/difference more clear (below equation 9).

[1] Liu et al (2021). Energy-based OOD Detection

[2] Peters et al (2015). Causal inference using invariant prediction

[3] Arjovsky et al (2019). Invariant Risk Minimization

[4] Krueger et al (2021). OOD Generalization via Risk Extrapolation

[5] Dong and Ma (2023). First Steps Toward Understanding the Extrapolation of Nonlinear Models to Unseen Domains

Thank you for taking the time to read and review our submission. We would be happy to answer any remaining questions before the discussion period ends tomorrow.

We thank the reviewer for their engagement and support of our paper. We tried to study invariance in a more general setting, and we are glad that the reviewer finds our formulation novel and the results interesting.

Weaknesses

1. Our main intention with the real life datasets was to demonstrate that our method does scale up to large image data, and is robust enough to perform well on a variety of image styles and shifts. However, current benchmarks like DomainBed are complex and difficult to parse in terms of what shift is occurring. Therefore, most proposed methods do not show a significant improvement on these benchmarks, and it is even unclear whether improvement on these specific benchmarks is the best way to evaluate generalization methods [1].

   Our goal is to raise an issue that we believe is relevant in any future application of invariant learning, and even though there are some more steps and further research required in order to make invariant learning relevant in practice, we hold that research into the common pitfalls of invariant learning methods is useful down the line as the literature forms more effective methods. Our application of reweighting to invariant learning may be incorporated as a component on top of future improvements in invariant learning methods.

2. Many real life scenarios exhibit heteroskedasticity and/or shift in the invariant features, settings our method aims to be robust to. As an example, we can imagine having data from various hospitals where demographics change. If we think of smoking as a causal feature for predicting cancer, and the percentage of smokers changes across environments, then this is a clear case of covariate shift in invariant features. Some types of cancer being more difficult to diagnose than others is a case of heteroskedasticity.

Questions

1. We make no assumption on the support of the invariant features. We do want to point out that if there were no overlapping support in the invariant features at all, then there is no way to guarantee that our method, or any invariant prediction method, recovers an invariant classifier [2]; this is a fundamental limit of the invariant prediction problem. When there is some overlapping support, then invariant learning is possible.

   The practical objective has been proposed over the original formulation because the original formulation requires the density of the invariant features, and we do not usually have access to these in practice. Instead, we propose to learn both the invariant feature density and an invariant predictor via alternating minimization. We see that this results in our learning useful invariant feature densities (through our OOD experiments) in addition to our invariant predictor.

2. Good question. We believe whether our method should be used depends on the dataset and use case. To elaborate:

   • WRI should be considered when you are looking for a causally-motivated method. Causally-motivated methods can be seen as more interpretable and theoretically sound, although they can also be difficult to scale to high-dimensional data. If this is not important to your use case, we can recommend non-causal methods with strong empirical performance like MIRO [3] and Fishr [4].
   • In our paper, we compare to VREx as another causally-motivated baseline. When we have a setting with heteroskedasticity and shift in the invariant features (see our response to Weakness 2), then VREx will fail to recover an invariant predictor. Thus, we recommend WRI if you suspect your data to be more heterogeneous. Our performance should be similar to VREx otherwise—in this case, we would recommend just using VREx directly.
   • Learning the density estimates of invariant features is unique to our method. In an experiment, we demonstrated that this is useful for OOD detection. More generally, this is also useful for informing a user when they are encountering test samples that are rare/minority samples in training. Thus, we recommend our method if you want to learn an invariant predictor that also includes density information at test time—e.g. for fairness or bias mitigation reasons.

[1]. Eastwood et al (2022). Probable Domain Generalization

[2]. Ahuja et al (2021). Invariance Principle Meets Information Bottleneck for OOD Generalization

[3]. Cha et al (2022). Domain generalization by mutual-information regularization

[4]. Rame at al (2022). Fishr: Invariant gradient variances for OOD generalization

Suggestions

1. Some of our experimental results are unconventional to present as a plot, e.g. CMNIST and DomainBed test results are always shown as tables, and these experiments do not have a continuously varying aspect that would benefit from being presented as a plot. However, we fully agree with the spirit of the suggestion; in trying to improve the readability of our results, we identified Table 2 as being difficult to read/interpret. Thus, we reworked Table 2 so that it matches the general format of the other tables (methods on the left, overall results on the right) and added some slight visuals. We believe this significantly improves its presentation and we thank the reviewer for the suggestion. Please feel free to raise additional concerns.

2. We appreciate the suggestion. We have created Appendix C.4 to visualize examples from the proposed CMNIST dataset and to include additional technical detail.

Thank you for taking the time to read and review our submission. We would be happy to answer any remaining questions before the discussion period ends tomorrow.