Pith. sign in

REVIEW 3 major objections 4 minor 13 references

Robust Invariant Representation Learning by Distribution Extrapolation

T0 review · 3 major / 4 minor · reviewed 2026-08-07 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read This paper claims that extrapolating a per-sample squared-gradient version of the IRMv1 penalty across synthetic environments improves out-of-distribution accuracy and calibration compared with standard IRM variants.

desk verdict A clean REx-style extension with a load-bearing flaw: in the paper's own SEM, the stated objective prefers the zero predictor over the invariant one, so the reported gains cannot come from Eqs. (6)-(9). read the letter →

arxiv 2505.16126 v2 pith:VK65K5MH submitted 2025-05-22 cs.LG

Kotaro Yoshida , Konstantinos Slavakis This is my paper
classification cs.LG
keywords invariantriskminimizationout-of-distributiongeneralizationdistributionextrapolationper-samplegradientpenaltycalibrationover-parameterizationIRMv1
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

IRMv1, the standard practical approximation of invariant risk minimization, penalizes the squared gradient of the expected loss, and this paper argues that this penalty can be small for models that rely on spurious features whenever the training environments are similar or the model is over-parameterized. To make the penalty robust, the paper replaces it with the expectation of squared per-sample gradients, which is linear in the data distribution and therefore can be extrapolated to pseudo-unseen environments in the style of risk extrapolation. Two concrete penalties result: a max-affine extrapolation ($\mathrm{mm}$-IRMv1) and a variance-penalized alternative ($\mathrm{v}$-IRMv1), both of which can be dropped into existing IRM variants. On the paper's experiments—SEMs, CMNIST, CFMNIST, PACS, and VLCS—these penalties consistently improve accuracy and calibration over the base variants. A sympathetic reader would take the contribution to be a modular fix for a known failure mode of IRMv1, not a new theory of invariance.

What carries the argument

The load-bearing object is the per-sample squared-gradient penalty $J_{\mathrm{IRM},e}(\pi,\Phi)$ of Eq. (6), which replaces $\|\nabla_\pi R_e(\pi\Phi)\|^2$ and is linear in the environment distribution, enabling extrapolation. Lemma 4.1 (Jensen's inequality) shows $|\nabla_\pi R_e|^2\le J_{\mathrm{IRM},e}$, so the new penalty is a majorant of the original. The extrapolation machinery is the same affine-combination trick as risk extrapolation: allowing negative coefficients on environment risks creates pseudo-environments outside the convex hull of training distributions. This yields the closed-form penalty $C_{\mathrm{mm}}(\Phi)=(1-\alpha_{\min}|E_{\mathrm{train}}|)\max_e J_{\mathrm{IRMv1},e}(\Phi)+\alpha_{\min}\sum_e J_{\mathrm{IRMv1},e}(\Phi)$ and the variance variant $C_v(\Phi)=\gamma\,\mathrm{Var}(\{J_{\mathrm{IRMv1},e}\})+\sum_e J_{\mathrm{IRMv1},e}$. Theorem 3.1 supplies the motivation: if the summed ERM loss is below $\delta$, then the original penalty is bounded by $2L_\Phi\delta$, which can be small even for spurious-feature extractors under similar environments.

What would settle it

On the SEM of Eq. (10) with training environments $\{0.2,0.6\}$, train $\mathrm{mm}$-IRMv1 and $\mathrm{v}$-IRMv1 and record the learned $(\hat{w}_{\mathrm{inv}},\hat{w}_{\mathrm{spu}})$; if either method converges to $(\mathbf{0}_d,\mathbf{0}_d)$—the zero predictor, for which $J=0$—rather than to $(\mathbf{1}_d,\mathbf{0}_d)$, the claim that the extrapolated penalty enforces invariance is falsified.

Watch

Extended reading notes

Core claim

The paper's central claim is that IRMv1's gradient penalty is the wrong object to penalize: because it is the squared norm of the gradient of an expectation, it can vanish for spurious-feature solutions when training environments are insufficiently diverse or when the model has spare capacity. The proposed remedy is to replace it by $J_{\mathrm{IRM},e}(\pi,\Phi)=\mathbb{E}_{(x,y)\sim P_e}\{|\nabla_\pi\ell(\pi\Phi(x),y)|^2\}$, the per-sample squared-gradient risk, which is linear in the data distribution. Lemma 4.1 shows this quantity upper-bounds the original penalty, and extrapolating it over affine combinations of environments with negative coefficients gives $C_{\mathrm{mm}}$; adding a variance term gives $C_v$. The paper reports that on the SEM of Arjovsky et al. and on four vision benchmarks, these penalties reduce causal and non-causal error and improve accuracy, ECE, and ACE compared with IRMv1, BIRM, and BLO, with the variance variant more stable under over-parameterization.

Load-bearing premise

The load-bearing premise is that minimizing the extrapolated per-sample squared-gradient penalty still forces the model to learn invariant features; the paper proves only that the new penalty upper-bounds the old one, not that its minimizer satisfies the original IRM constraint, and in the paper's own SEM the zero predictor has zero penalty but no invariant features.

Editorial extensions

If this is right

  • IRMv1, BIRM, and BLO can be upgraded by swapping in either $C_v$ or $C_{\mathrm{mm}}$ without changing their training loops, yielding better average OOD accuracy and calibration.
  • The variance-based penalty $C_v$ is the safer default in over-parameterized vision models, while the max-based $C_{\mathrm{mm}}$ excels in low-dimensional SEM settings.
  • The observed failure mode—IRMv1's training penalty collapsing to near zero while test accuracy and calibration stay poor—is mitigated by both proposed penalties.
  • Calibration improves alongside accuracy, consistent with the paper's use of ECE and ACE as a second evaluation axis.
  • The same extrapolation recipe can be layered onto other IRM variants whose penalties are expectations of per-sample quantities.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • The paper leaves implicit that any penalty expressible as an expectation over per-sample losses—not just IRMv1's—can be extrapolated this way; a natural test is applying $C_v$ to game-theoretic or information-bottleneck IRM variants.
  • Because the extrapolated penalty stays larger than the raw IRMv1 penalty during training, monitoring $J_{\mathrm{IRMv1},e}$ on held-out environments might serve as a cheap early-stopping or model-selection signal, an application the paper does not explore.
  • The superiority of $\mathrm{mm}$ in low-dimensional settings and $\mathrm{v}$ in over-parameterized ones suggests a smoothness-versus-diversity tradeoff; one could interpolate between the two penalties with a single hyperparameter rather than choosing one form.
  • A cautionary consequence of Lemma 4.1 alone: a degenerate predictor that makes $J=0$ (e.g., a constant feature map) is not ruled out by the majorant property, so the method's success in the reported experiments is an empirical, not a proven, guarantee.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 4 minor

Summary. This paper proposes two modifications to the IRMv1 penalty for invariant risk minimization. The proposal replaces the squared norm of the gradient of the expected loss with the expectation of the squared per-sample gradient norm, J_IRM,e(π,Φ) = E[|∇πℓ(πΦ(x),y)|²], and then applies risk-extrapolation ideas (REx) to this penalty, yielding two objectives called mm-IRMv1 and v-IRMv1 (Eqs. (7) and (9)). The authors prove that J_IRM,e majorizes the original IRMv1 penalty (Lemma 4.1) and present experiments on synthetic SEMs and four vision datasets, claiming consistent improvements in accuracy and calibration over IRMv1, BIRM, and BLOC-IRM.

Significance. If the proposed method were correct, it would provide a drop-in replacement penalty for IRM-family methods with a few percent gains in accuracy and calibration. The paper includes a reproducibility link, and the closed-form of the mm penalty is derived carefully in Appendix C. However, the theoretical justification is incomplete: Lemma 4.1 only establishes an upper bound and does not show that minimizing J_e enforces the IRM lower-level constraint. The numerical evidence is weakened by test-domain validation tuning and by several within-noise differences.

major comments (3)
  1. [Section 4.2, Eq. (6); Section 5.1] The proposed surrogate J_IRM,e is not a safe replacement for the IRMv1 penalty. For the paper's own SEM (Eq. (10)) with d=5 and E_train={0.2,2}, the invariant feature extractor Φ(x)=1^T x_inv gives J_e=4d e² (0.8 for e=0.2 and 80 for e=2), whereas the trivial extractor Φ≡0 gives J_e=0. Including the ERM term, the total objective for the trivial solution is Σ(d e² + 1)=22.2, while the invariant solution gives 2 + λ·80.8 for v-IRMv1 (and similarly for mm-IRMv1 with α_min ≤ -1), which is larger for every λ in the reported grid {1,10}. Thus the proposed objective prefers a constant predictor to the invariant one. Table 1 nevertheless reports low causal errors for v-IRMv1 and mm-IRMv1; this is only possible if the implementation optimizes a different objective than Eqs. (6)–(9) or if the optimizer is trapped away from the objective's minimizer. Lemma 4.1 alone does not bridge the gap: it proves only that |∇πRe|² ≤ J_IRM,e, and a large J_IRM,e does not imply a large original penalty, while a zero J_IRM,e does not imply a non-degenerate invariant predictor.
  2. [Section D.2] For the correlation-shift datasets (CMNIST, CFMNIST), the hyperparameters α_min and γ are selected using the 'test-domain validation set' strategy (Gulrajani & Lopez-Paz, 2020). This leaks test-domain information into model selection, which is known to overestimate OOD performance. Since the same tuning is not applied to the baselines' hyperparameters, the comparison is not fair, and the claimed gains (Table 2) cannot be attributed to the proposed penalty.
  3. [Tables 2–4, Section 5.2.2] The claim that the proposed variants 'consistently outperform' the baselines is not supported by the reported standard deviations. For example, IRMv1 v shows no improvement on VLCS (58.4±2.1 vs 58.4±0.5) and is within noise on PACS (75.9±3.8 vs 75.5±1.4) and CFMNIST (74.8±0.5 vs 74.3±1.1); v-BIRM worsens average ECE (16.6 vs 16.3) and v-IRMv1 worsens ECE on CMNIST (10.5 vs 10.4) and CFMNIST (17.9 vs 17.8). The 95% confidence intervals overlap for most individual comparisons, so the reported improvements are not statistically robust.
minor comments (4)
  1. [Tables 2–4] The column headers 'P ACS' and 'A vg.' should be 'PACS' and 'Avg.'.
  2. [Section 2.5] The notation Re(π·Φ) is used both as an expectation functional and as a function argument; a consistent notation would improve readability.
  3. [References] The reference to Quiñonero-Candela et al. (2022) is a reprint of a 2008 monograph; please cite the original publication year.
  4. [Section 5.1.1] In the sentence 'x_inv_e, x_spu_e and y_e are d×1vector-valuedandscalar-valuedrealizations', spaces are missing between words.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: proposed penalties and experiments are self-contained, with only non-load-bearing self-citations.

full rationale

The derivation chain is not circular. JIRM,e in Eq. (6) is defined as an expectation of squared per-sample gradients; Lemma 4.1 is proved from Jensen's inequality and does not assume the target result. Equations (7) and (9) are algebraic/LP consequences of that definition, and the closed form in (7) is derived from KKT conditions, not fitted. The experimental claims are not predictions from the fitted loss by construction: causal/non-causal errors are measured on learned SEM parameters, and vision accuracies/calibrations are independent benchmarks. The self-citations (Yoshida & Naganuma 2024; Naganuma et al. 2025) appear only in the related-work calibration discussion and are not used to justify the extrapolation step or to exclude alternative methods. The use of the test-domain validation set for correlation-shift model selection in D.2 is transparently reported; it is a benchmark-selection protocol, not a step in which a fitted parameter is renamed as a prediction. One legitimate soundness concern—Lemma 4.1 only proves J_e majorizes the IRMv1 penalty, so surrogate minimizers need not match the original penalty's—is a correctness risk, not a circularity, because nothing in the proof assumes the proposed method works.

Assumptions & free parameters 3 free parameters · 5 assumptions · 0 invented entities

The paper adds no new physical entities. Its central claim depends on standard smoothness assumptions and on the domain assumption that extrapolation with negative weights is useful; the three hyperparameters αmin, γ and λ are tuned per dataset and are free parameters.

free parameters (3)
  • alpha_min = SEM: -1, -5, -10; vision: -0.1 to -1
    User-defined lower bound in Eq. (8); controls extrapolation strength; tuned by grid search (Appendices D.1, D.2).
  • gamma = SEM: 1, 10, 100; vision: 0.1 to 1
    Variance weight in Eq. (9); tuned by grid search.
  • lambda = SEM: 1, 10; vision: 1e6 (correlation shifts), 1 (diversity shifts)
    Regularization weight in (mm-IRMv1) and (v-IRMv1); set by grid search or protocol.
assumptions (5)
  • standard math Assumption 3.1: for each Φ, ∇πRe(π·Φ) is LΦ-Lipschitz in π
    Used in Theorem 3.1 via the descent lemma; not verified for the neural networks in the vision experiments.
  • standard math Re(f) >= 0 for all f
    Stated in Section 2.1; needed for the contradiction in Appendix A.
  • standard math Expectation and gradient can be interchanged (dominated convergence)
    Invoked in Appendix B for Lemma 4.1; no integrability check is given.
  • domain assumption Negative mixture weights define useful pseudo-environments
    The extrapolation set A in Eq. (8) includes α_e < 0, which does not correspond to a probability distribution; the paper assumes optimizing over these pseudo-environments improves OOD generalization, following Krueger et al. (2021).
  • standard math The dummy classifier can be fixed at π=1 without loss
    Scaling ambiguity argument in Section 2.5; valid when the lower-level minimizer is nonzero.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Robust Invariant Representation Learning by Distribution Extrapolation." pith.science (2026). https://pith.science/paper/VK65K5MH

@misc{pith2026250516126,
  author       = {Pith},
  title        = {Pith review of: Robust Invariant Representation Learning by Distribution Extrapolation},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/VK65K5MH}},
  note         = {Machine review of arXiv:2505.16126}
}
read the original abstract

Invariant risk minimization (IRM) aims to enable out-of-distribution (OOD) generalization in deep learning by learning invariant representations. As IRM poses an inherently challenging bi-level optimization problem, most existing approaches -- including IRMv1 -- adopt penalty-based single-level approximations. However, empirical studies consistently show that these methods often fail to outperform well-tuned empirical risk minimization (ERM), highlighting the need for more robust IRM implementations. This work theoretically identifies a key limitation common to many IRM variants: their penalty terms are highly sensitive to limited environment diversity and over-parameterization, resulting in performance degradation. To address this issue, a novel extrapolation-based framework is proposed that enhances environmental diversity by augmenting the IRM penalty through synthetic distributional shifts. Extensive experiments -- ranging from synthetic setups to realistic, over-parameterized scenarios -- demonstrate that the proposed method consistently outperforms state-of-the-art IRM variants, validating its effectiveness and robustness.

Figures

Figures reproduced from arXiv: 2505.16126 by the authors.

Figure 1
Figure 1. The figure illustrates the relationship between the IRMv1 penalty values in the training environ [PITH_FULL_IMAGE:figures/full_fig_p010_1.png] view at source ↗
Figure 2
Figure 2. The relationship between the IRMv1 penalty values in the training environment and the corre [PITH_FULL_IMAGE:figures/full_fig_p018_2.png] view at source ↗
Figure 3
Figure 3. The relationship between the IRM penalty values in the training environment and the corresponding [PITH_FULL_IMAGE:figures/full_fig_p019_3.png] view at source ↗
Figures from the paper (1 more)
Figure 4
Figure 4. Figure 4: The relationship between the IRMv1 penalty values in the training environment and the corre [PITH_FULL_IMAGE:figures/full_fig_p019_4.png]

Discussion (0). Sign in to comment.

Reference graph

Works this paper leans on

13 extracted references · 4 canonical work pages

  1. [1]

    Invariant risk minimization games.arXiv preprint arXiv:2002.04692,

    K Ahuja, K Shanmugam, KR Varshney, and A Dhurandhar. Invariant risk minimization games.arXiv preprint arXiv:2002.04692,

    arXiv 2002

  2. [3]

    Invariant risk minimization.arXiv preprint arXiv:1907.02893,

    Martin Arjovsky, L’eon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization.arXiv preprint arXiv:1907.02893,

    arXiv 1907

  3. [9]

    Antonio Torralba and Alexei A Efros

    URL https://arxiv.org/abs/1911.08731. Antonio Torralba and Alexei A Efros. Unbiased look at dataset bias. InCVPR 2011, pp. 1521–1528. IEEE,

    arXiv 1911

  4. [11]

    URLhttps://proceedings.mlr.press/v162/ zhou22e.html. 14 Appendices A Proof of Theorem 3.1 First, by the definition ofFδ, ∑ e∈Etrain Re(π· Φ)≤δ, ∀(π, Φ)∈F δ, and the initial hypothesis thatRe(·)≥ 0, it can be easily concluded that Re(π· Φ)≤δ, ∀(π, Φ)∈F δ,∀e∈Etrain. Assume, for contradiction to the claim of the theorem, that there exist anˆe∈Etrain and (ˆπ,...

    work page 2004

  5. [12]

    This completes the proof. B Proof of Lemma 4.1 Notice that∀(π, Φ), |∇πRe(π· Φ)|2 =|∇π E(x,y)∼Pe{ℓ(π· Φ(x),y )}| 2 =| E(x,y)∼Pe{∇πℓ(π· Φ(x),y )}| 2 ≤ E(x,y)∼Pe { |∇πℓ(π· Φ(x),y )|2} =JIRM,e(π, Φ), where the interchange between the expectation and the partial gradient in the second equality follows from the dominated convergence theorem, while the inequalit...

    work page 2020

  6. [224]

    test-domain validation set

    across 5 classes, drawn from four different sources. For the experimental setup, the Caltech101, LabelMe, and VOC2007 domains are used for training, while SUN09 serves as the test domain. The Adam optimizer is used with a learning rate of 5× 10−4 across all datasets. For correlation shift experiments, models are trained using full-batch optimization for50...

    work page 2022

  7. [1991]

    Out-of- distribution generalization via risk extrapolation (rex).arXiv preprint arXiv:2003.00688,

    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).arXiv preprint arXiv:2003.00688,

    arXiv 2003

  8. [2013]

    Pareto invariant risk minimization: Towards mitigating the optimization dilemma in out-of-distribution generalization.arXiv preprint arXiv:2206.07766,

    Y Chen, K Zhou, Y Bian, B Xie, B Wu, Y Zhang, K Ma, H Yang, P Zhao, B Han, et al. Pareto invariant risk minimization: Towards mitigating the optimization dilemma in out-of-distribution generalization.arXiv preprint arXiv:2206.07766,

    arXiv

Show all 13 references
  1. [2016]

    In search of lost domain generalization

    12 Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. arXiv preprint arXiv:2007.01434,

    2007 arXiv

  2. [2017]

    Lutz Prechelt

    URL https://arxiv.org/abs/1712.04621. Lutz Prechelt. Early stopping-but when? In Neural Networks: Tricks of the trade, pp. 55–69. Springer,

    arXiv

  3. [2020]

    Invariance principle meetsinformationbottleneckforout-of-distributiongeneralization

    K Ahuja, E Caballero, D Zhang, JC Gagnon-Audet, Y Bengio, I Mitliagkas, and I Rish. Invariance principle meetsinformationbottleneckforout-of-distributiongeneralization. arXiv preprint arXiv:2106.06607,

    arXiv

  4. [2022]

    doi: 10.1016/j.engappai.2022.105151

    ISSN 0952-1976. doi: 10.1016/j.engappai.2022.105151. URL http://dx.doi.org/10.1016/j.engappai.2022.105151. Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarial training of...

    1976

  5. [2024]

    URL https://openreview.net/forum?id=9YqacugDER

    ISSN 2835-8856. URL https://openreview.net/forum?id=9YqacugDER. Yihua Zhang, Pranay Sharma, Parikshit Ram, Mingyi Hong, Kush Varshney, and Sijia Liu. What is missing in IRM training and evaluation? challenges and solutions.arXiv preprint arXiv:2303.02343,

    arXiv

Pith tools

Reviewed August 7, 2026 · model on record in the stance chip above.