Paper: Weighted Risk Invariance for Density-Aware Domain Generalization (ICLR-2024)

Reviewer comment (Theory):
The definition of an invariant predictor (in Defn. 1) is not well-defined. Indeed, the invariance of the conditional distribution $p_e(f(X)|X_{inv})$ is not equivalent to the condition $f(X)=f(X_{inv})$. Furthermore, the domain invariant presentation in a general domain generalization should be based on the conditional distribution of the label given the input feature, i.e., $p_e(Y|g(X))$ with $g$ is a presentation mapping, not the predictor $f$.

Author response:
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.

Reviewer comment (Theory):
Though I did not have time to verify the proof of Proposition 1, I am still not convinced the result that the invariance of the predictor $f$ leads to the weighted invariance (Defn. 2). In addition, could you please clarify the real meaning of the definition of the weighted invariance (Defn. 2), also the way to verify that property? I also would like to see the comments from other reviewers on this.

Author response:
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.

Reviewer comment (Evaluation):
The experimental results of the paper are supportive, but not very convincing. In particular, it seems that the proposed WRI can only beat the related baselines (IRM and VREx) on the synthetic dataset (MNIST), while on the real-world dataset Domainbed, their performance is much similar, and WRI's even could not beat the naive ERM approach.

Author response:
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.

Reviewer comment (Theory):
The authors assert that REx is limited to the homoskedastic setting, whereas their method can accommodate the heteroskedastic setting. However, the definitions appear to pertain to disparate scenarios. In the REx literature, homoskedasticity is tied to noise variance discrepancies across different $X$, while heteroskedasticity in this work relates to covariate shifts in invariant features. The rationale behind REx's inability to address the heteroskedastic covariate shift is not lucid. Moreover, the proposed WRI seems incapable of addressing the conventional heteroskedastic scenario, as varying noise variance for $Y$ across environments would render the weighted risk inconsistent across environments.

Author response:
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}$).

Reviewer comment (Theory):
The implementation details are somewhat unclear. The reason behind employing an alternating minimization process, as opposed to direct optimization of Equation (9), is not explicit. Furthermore, ensuring the identification of invariant features via Equation (9) seems challenging. Notably, the final term in Equation (9) gravitates towards $d(x) \rightarrow P_e(x)$, inclining towards a dependency on all features. Furthermore, it is unclear whether the second term of Equation (9) also contains other trivial solutions (like the all zero solution mentioned by the author).

Author response:
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.

Reviewer comment (Experiments):
The omission of critical baselines, such as SWAD [1] and MIRA [2], potentially diminishes the empirical significance of the proposed technique. The authors' assertion in the appendix that non-causally motivated methods can occasionally outperform causally-based methods in domain generalization tasks appears to undermine the essence of leveraging causality-based techniques in this realm.

Author response:
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.

Reviewer comment (Experiments):
In the "OOD detection performance of our learned densities" section, there is an absence of detailed explanation regarding the modified CMNIST test split incorporating mirrored digits.

Author response:
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.

Reviewer comment (Novelty):
The discussion would benefit from a more extensive review of related work addressing covariate-shift generalization, e.g., [3][4].

Author response:
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.

Reviewer comment (Presentation):
The introduction and comparison to (V)REx and IRM are separated, posing some difficulty in reading.

Author response:
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.

Reviewer comment (Theory):
No argument is provided for the non-linear case, even though the experiments on DomainBed does not follow the linear regime.

Author response:
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.

Reviewer comment (Theory):
Lacking analysis on the difference between WRI and the practical objective.

Author response:
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).

Reviewer comment (Evaluation):
While the paper does a good job of motivating the problem and showing results on synthetic setups, my main concern is reg. the experiments on real life datasets: that it is unclear whether a conclusion can be made. When does the proposed method work v/s when does it not, when it comes to real life datasets? Is there any reason behind these observations? What can one infer?

Author response:
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.

Reviewer comment (Writing):
Following up on the previous point, could the authors explain when might such a phenomenon occur in a real life dataset more explicitly? Perhaps a real life intuition would help.

Author response:
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.

Reviewer comment (Theory):
Is there any assumption on the support of the invariant features? Is this why the practical objective has been proposed over the original formulation?

Author response:
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.

Reviewer comment (Evaluation):
Is there any guidance for the reader on when this method should and should not be used, given a dataset?

Author response:
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.

Reviewer comment (Presentation):
It might be a good idea to include a summary of all experimental results in a visual format e.g. plot so that the results are easier to read through.

Author response:
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.

Reviewer comment (Presentation):
It would be nice to visualize the new proposed colored MNIST example via an illustration to understand the setting used in the paper.

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