Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: In this paper, the authors propose an alternative Invariant Risk Minimization (IRM) penalty for the particular case of a linear output layer trained with a mean-squared error loss. In this particular case, the original formulation can be understood as penalizing the squared Euclidean norm of the output layer's gradient, that is,

$$\rho_{e}^{\mathrm{IRMv1}}(\Phi, w) := \vert\vert \Phi_{e}^{T}(\Phi_{e} w - y_{e}) \vert\vert_{2}^{2} = \vert\vert \Phi_{e}^{T}\Phi_{e}(w -w_{e}^{*})\vert\vert_{2}^{2},$$

where $\Phi_{e} \in \mathbb{R}^{n_{e} \times d}$  is a matrix with the (learnt) data representation for environment $e$, $y_{e}$ the corresponding targets and $w_{e}^{*} = (\Phi_{e}^{T}\Phi_{e})^{-1}\Phi_{e}^{T}y_{e}$ the optimal, unregularized linear regressor. Witth this in mind, this paper proposes to use instead:

$$\rho_{e}^{\mathrm{IRMv2}}(\Phi, w) := \vert\vert (\Phi_{e}^{T}\Phi_{e})^{1/2}(w -w_{e}^{*})\vert\vert_{2}^{2},$$

which interestingly can be trivially shown to lead to an optimal classifier $w^{*}$ identical to that of the ERM formulation. Crucially, however, the authors propose to explicitly backpropagate through the dependence of this classifier on the data representations $\Phi_{e}$ when minimizing a cost function that incorporates $\rho_{e}^{\mathrm{IRMv2}}(\Phi, w)$.

The authors prove that, under similar non-degeneracy conditions as those assumed when analyzing other IRM penalties, $\rho_{e}^{\mathrm{IRMv2}}(\Phi, w)$ will either find an invariant data representation (or a trivial classifier) for a linear SEM as in Rosenfeld et al. 2021. Furthermore, the authors argue that there might be a connection between ill-conditioning of the Gram matrix $\Phi_{e}^{T}\Phi_{e}$ and failure modes of IRM penalties and propose a heuristic rule to adapt the penalty weight in the original IRM formulation.
main_review: ## Major points

1. IRM penalties are most often studied alongside representations learnt by deep learning models, which are most often trained using some variant of stochastic gradient descent. Concretely, Arjovsky et al. 2019 proposed an unbiased stochastic estimator of the gradient of the original IRM penalty. However, to the best of my knowledge, the proposed approach as presented would in principle require computing the gradient of the ERM regressor $w^{*}(\Phi) = (\Phi^{T}\Phi)^{-1}\Phi^{T}y$ with respect to $\Phi$. Similarly, the adaptive rule to adjust the penalty weight for IRMv1 presented in Section 3.1. would in principle require estimating the smallest eigenvalue of $\Phi_{e}^{T}\Phi_{e}$, which might be difficult in a stochastic setting.

    I believe it would be ideal if the authors could address this limitations either theoretically (e.g. by proposing unbiased gradient estimators) or empirically (e.g. by proposing heuristic, possibly-biased gradient estimators and demonstrating good empirical performance nevertheless). At the very least, the issue of how to estimate the gradient of $L_{t}(\phi_{\theta_{t}})$ in Line 7 of Algorithm 1 in practice should be described in more detail. If it turns out that the proposed approach can only be used with (non-stochastic) gradient descent, this should at least be explicitly stated in the manuscript to ease the way for future work.

2. The proposed approach can only be used alongside the MSE loss. While this is acknowledged by the authors, which is appreciated, I do believe this to be an important practical limitation, as this loss is rarely used for supervised classification. Moreover, this makes the soundness of the adaptive rule in Section 3.1. questionable whenever IRMv1 is used in conjunction with e.g. the logistic loss. 

3. The adaptive rule in Section 3.1. could be better justified. While the bound in Lemma 3 is to the best of my knowledge sound for the MSE loss, it is not explicitly stated how that motivates Eq. (15) exactly and what the theoretical implications are, if any. This is in my opinion important as it could be argued this adaptive rule is the contribution that has the most practical impact in performance according to the experiments in Section 5.

4. The authors make a point that there might be a link between ill-conditioning of the Gram matrix $\Phi_{e}^{T}\Phi_{e}$ and failure modes of IRM penalties, providing theoretically that some of the most well-known counterexamples to these penalties do exhibit ill-conditioning. However, the empirical results in Figure 2 in the appendix are somewhat conflicting. In particular, it would appear that the "phase transition" in which test accuracy sharply drops precedes the sharp increase in ill-conditioning particularly for IRMv1 and IRMv1A.

5. All in all, the experimental results do not convey a clear benefit of using IRMv2 vs IRMv1. The HealthyGutTests results in Table 2 are, as far as I can tell, well within the margin of error. Most importantly, it would appear IRMv1 (and the IRMv1A variant) dominate IRMv2 in the LinearUnitTests experiments.

## Typos

Proofs in the appendix often refer to the $\mathcal{l} - 2$ ball instead of $\mathcal{l}_{2}$.
summary_of_the_review: Out-of-distribution generalization is one of the most pressing, largely open problems in the field. Concretely, the IRM framework has recently emerged as a promising line of research in this direction, but its characterization is far from conclusive both from a theoretical and a practical standpoint. Thus, I consider the paper's topic timely and of great relevance to the field.

I believe the manuscript  explores an interesting idea and is technically sound, albeit not without what I consider to be important limitations, namely,
+ The proposed approach is only applicable to the MSE loss.
+ It is unclear whether it would be sound to use the proposed penalty in a stochastic gradient-based optimization setting, as it requires backpropagating through the optimal linear regressor, which depends on all data.
+ The experimental results do not provide compelling evidence that the proposed penalty is statistically significantly superior to the original formulation.

All in all, I believe IRM to be a promising yet not fully mature idea with many crucial open questions remaining to be answered. Thus, I consider exploratory ideas such as those presented in this article to be worthwhile despite potential shortcomings that might need to be revisited by future work. Hence, despite the aforementioned limitations, I lean towards supporting acceptance of the manuscript.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper studies the problem of generalizing out of distribution (OOD), i.e., when data is not iid. Motivated by the error decomposition of the least-squares scheme, this paper proposes a novel invariance penalty that is directly related to risk. Based on the newly proposed penalty, this paper introduces another implementation of the IRM, named IRMv2. By characterizing the difference between the proposed penalty and the one in [1], this paper provides an adaptive rule for the choice of penalty parameter in the original IRMv1. Theoretical results guarantee that the proposed scheme can capture invariance when  in the linear case. The conditional number of Gram matrix of the data representation is analyzed.
main_review: This paper is well written but the innovation is insufficient. The form of proposed invariant penalty is similar to the one of invariant penalty proposed in [1] and limited to the least-square loss. Similar theoretical result was obtained in [2] when  in the linear case. Moreover, considering that there have been many alternatives of IRM in recent literature, a detailed explanation should be given to fully illustrate the strength and weakness of the proposed method compared to the existing alternatives.
summary_of_the_review: Specific Comments
1）The similar theoretical result in this paper has been obtained in [2]. In other words, under the same condition, the theoretical result also holds for IRM. It confuses me whether the proposed method is essentially different from IRM.
2）It seems to me that the proposed invariance penalty doesn’t solve the problem of IRM. This paper doesn’t guarantee that the proposed IRMv2 can recover invariant predictor at least in the theoretical scenario where IRM fails.
3）I suggest that a detailed comparison with other alternatives for IRM should be given to better illustrate the strength and weakness of the newly proposed invariant penalty.
4）There are a few places that are not very rigorous. (1) In 5-th line of Algorithm 1,  is computed according to Eq. equation 4. However, the distribution is usually unknown and expectation can’t be computed. (2) In 6-th line of Algorithm1,  is computed according to Eq. equation 12 instead of Eq. equation 11.

[1] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893, 2019.
[2] Elan Rosenfeld, Pradeep Kumar Ravikumar, and Andrej Risteski. The risks of invariant risk minimization. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=BbNIbVPJ-42.
review_title: null

final_score: rating: 3: reject, not good enough
review_content: summary_of_the_paper: The paper proposes a small modification to the IRMv1 objective which leads to an improvement over the original. They propose potential theoretical explanations for the improvement.
main_review: In an earlier version of this review, I wrote that I felt that the objective proposed by this work is very similar to an existing linear regression objective from 2018, known as Anchor Regression [1]. The present paper proposes to consider linear regression *only*, which of course makes for an easier problem, but it also means there is lots of existing work in this area (for which there is insufficient discussion in the related work).

However, after further thought, I've realized that the two are only equal in quite specific circumstances. On the other hand, I have come across an observation that I think demonstrates the extent to which this work does not analyze their objective closely enough.

Consider the objective proposed by the paper:
$$\sum_e R_e(w^T\varphi) + \lambda ||I_e(\varphi)^{1/2} (w-w_e^*)||^2.$$
The first term is simply ERM. The second term, given features $X = \varphi(x)$, can be written $(w-w_e^*)^TX^TX(w-w_e^*)$. Consider what happens when we take the gradient of this objective with respect to our predictor $w$. We arrive at $2X^TX(w-w_e^*)$. Observe that this is precisely the gradient of linear regression, if we are assuming that our target is $Xw_e^*$ instead of the given target $y$. In other words, this is just linear regression where the targets are provided by a proxy classifier, as in two-stage least-squares. This tells us two things:

1. The objective is simply a weighted sum of the risk when predicting the true target $y$ and the "instrumental variable" target $Xw_e^*$ (equivalently, $w_e^{*T}\varphi(x)$). I believe this is nowhere mentioned in the paper and is worth exploring, if only because there doesn't seem to be any real understanding of what is going on with this objective.
2. I think my original impression that this was similar to Anchor Regression now appears pretty close to accurate, since Anchor Regression can also be seen as a trading off between ERM and a two-stage regression with instrumental variables.

Both of these points indicate to me that there is indeed something interesting going on with this objective, but this paper does not adequately explore the implications; nor does it properly discuss the previous related work (this paper takes existing work on logistic regression and applies it to linear regression, but there is a great deal of existing work in the causal literature on linear regression which is not discussed).

Returning to the paper itself, I find the empirical and theoretical results quite weak and I think this submission needs work even if it was a totally new idea. If there were some interesting new technical results to provide justification for this method in a different setting (such as classification) or under a different/more general model, it would certainly be noteworthy. However, I do not believe this to be the case: the setup Lemmas are simple linear algebra, and the primary theorem is almost identical to the related theorem in Rosenfeld et al. (furthermore, it only shows that it has the *same* linear environment complexity as IRM, rather than improving on it as the paper claims). The only new thing here is learning the feature embedder $\varphi$, but this paper provides no additional justification for when/why this should work---that is, all of the arguments for the benefits of this approach are agnostic to the actual quality of $\varphi$.

I also don't buy the claim that this would work in the non-linear setting; the result of Rosenfeld et al. is not specific to the type of loss, just the non-linearity of $\varphi$. The point behind that proof was that given only observations from the training environment, one can construct regions which contain all training points and on which the featurizer is perfect, and then make the featurizer arbitrarily bad elsewhere. In the absolute worst case, imagine a $\varphi = \varphi^*$ on all training points, and $\varphi = 0$ elsewhere. It seems clear to me that whatever claim the authors are making here with regards to the OOD performance of their predictor is not correct or is stated unclearly, because this solution would behave just like $\varphi^*$ at train time and give arbitrary error at test time. Even when observing the full environmental distributions, the fact that as we move away from the distribution mean the measure decays exponentially but the loss only grows quadratically means the same result would hold. I encourage the authors to revisit this section and work to clarify what precise mathematical claim they are making, and what this actually means in English.

Finally, *what is it with bolding empirical results which do not meet the bar for statistically significant improvement?* In Table 2, the confidence intervals of the performances of **all compared methods** overlap. Further, why aren't the standard deviations of the individual test accuracies reported over the five runs? There is zero indication that this method outperforms the ones to which it is being compared.

[1] Anchor regression: heterogeneous data meets causality. Dominik Rothenhäusler, Nicolai Meinshausen, Peter Bühlmann, Jonas Peters.


-------------------- Update after rebuttal --------------------

After discussion with the authors, I believe there are interesting aspects of this algorithm that are promising and worth exploring further. However, I still don't think the proposed method is analyzed or understood well enough, nor does there appear to be any substantial empirical or theoretical improvement here over existing works. I'm leaving my recommendation as is.
summary_of_the_review: The idea behind this paper, while nice, is a repackaging of IV regression, and does not adequately analyze the implications of the objective (nor do they properly contextualize it within the great deal of similar related work in other fields). They provide no new theoretical justification (in fact their main technical result, very similar to another work, simply shows that this algorithm performs on par with the existing IRMv1). Further, their explication of performance in the non-linear setting is incomplete.
review_title: null

final_score: rating: 3: reject, not good enough
review_content: summary_of_the_paper: This paper investigates the failure case of Invariant Risk Minimization and proposes a new invariance penalty for Out-of-Distribution Generalization. The authors provide some insights on how to avoid the potential failure of IRM based on the failure cases. And empirical results also validate the effectiveness of the proposed method.
main_review: a)	Strength: 
i.	This paper analyzes the counter-example of IRM and proposes a new invariance penalty to solve the OOD Generalization problem.
ii.	Practical insights on avoiding the failure case are provided for a better understanding of the counter-examples.
b)	Weaknesses:
i.	The proposed method is restricted to mean squared loss and cannot be extended to more general cases.
ii.	Although the role of I_e is demonstrated in Section 4.3, why the newly-proposed penalty is necessary and good is still unclear. And the derivation of the invariance penalty seems ad-hoc.
iii.	There lack of enough baselines in experiments. The authors only compare with ERM and IRM, while there are many other invariant learning methods, such as Rex[1], EIIL[2], HRM[3], etc.
iv.	Experimental results are not good enough to validate the effectiveness of the proposed method.

[1] D. Krueger, E. Caballero, J.-H. Jacobsen, A. Zhang, J. Binas, D. Zhang, R. Le Priol, and A. Courville, “Out-of-distribution generalization via risk extrapolation (rex),” ICML2021.
[2] E. Creager, J.-H. Jacobsen, and R. Zemel, “Environment inference for invariant learning,”ICML2021.
[3] J. Liu, Z. Hu, P. Cui, B. Li, and Z. Shen, “Heterogeneous risk minimization,”ICML 2021.

summary_of_the_review: This paper analyzes the counter-example of IRM and proposes a new invariance penalty based on that. Although some intuitions are provided, it is unclear why the proposed invariance penalty is necessary. And baselines, as well as the empirical results, are not enough to support the method.
review_title: null

Meta-review: This manuscript was the object of a rich and lengthy discussion. The AC also felt compelled to read the paper in details and discussed it further with the SAC.

The authors did a thorough job at addressing some of the reviewers points. The added results on cross-entropy loss and additional discussion, as well as the points made in "Further Discussion on the Numerical Experiments" are very much appreciated. 

However, significant concerns remain on establishing connections with prior work, including related ideas on invariance from the causality literature, so as to gain deeper understanding of the implications of the proposed objective. We also strongly encourage the authors to further work on strengthening their theoretical analysis to clearly demonstrate the value of the proposed approach.

The proposed formulation is certainly thought provoking and we urge the authors to pursue their work in view of the above comments.
Decision: Reject