Towards Understanding Variants of Invariant Risk Minimization through the Lens of Calibration
Authors: Kotaro Yoshida
Hiroki Naganuma
Venue: Accepted by TMLR ()

Abstract: Machine learning models traditionally assume that training and test data are independently and identically distributed. However, in real-world applications, the test distribution often differs from training. This problem, known as out-of-distribution (OOD) generalization, challenges conventional models. Invariant Risk Minimization (IRM) emerges as a solution that aims to identify invariant features across different environments to enhance OOD robustness. However, IRM's complexity, particularly its bi-level optimization, has led to the development of various approximate methods. Our study investigates these approximate IRM techniques, using the consistency and variance of calibration across environments as metrics to measure the invariance aimed for by IRM. Calibration, which measures the reliability of model prediction, serves as an indicator of whether models effectively capture environment-invariant features by showing how uniformly over-confident the model remains across varied environments. Through a comparative analysis of datasets with distributional shifts, we observe that Information Bottleneck-based IRM achieves consistent calibration across different environments. This observation suggests that information compression techniques, such as IB, are potentially effective in achieving model invariance. Furthermore, our empirical evidence indicates that models exhibiting consistent calibration across environments are also well-calibrated. This demonstrates that invariance and cross-environment calibration are empirically equivalent. Additionally, we underscore the necessity for a systematic approach to evaluating OOD generalization. This approach should move beyond traditional metrics, such as accuracy and F1 scores, which fail to account for the model’s degree of over-confidence, and instead focus on the nuanced interplay between accuracy, calibration, and model invariance.

Published in Transactions on Machine Learning Research (06/2024)
Towards Understanding Variants of Invariant Risk Minimization through the Lens of Calibration
Yoshida Kotaro", Hiroki Naganuma2.3
yoshida. k. 6.6.mm titech.ac.jp, nagonuma. hin quebec,
Tokyo Institute of Technology, Mila Quebec Al Institute,
3 Université de Montréal,
Reviewed on OpenReview: https: openreview. net/ forum? id= 9QqauuDER
Abstract
Machine learning models traditionally assume that training and test data are indepen- dently and identically distributed. However, in real-world applications, the test distribution often differs from training. This problem, known as out-of-distribution (OOD) general- ization, challenges conventional models. Invariant Risk Minimization (IRM) emerges as a solution that aims to identify invariant features across different enviromments enhance OOD robustness. However, IRM's complexity, particularly its bi-level optimization, has led to the development of various approximate methods. Our study investigates these approxi- mate IRM techniques, using the consistency and variance of calibration across environments as metrics to measure the invariance aimed for by TRM. Calibration, which measures the reliability of model prediction, serves as an indicator of whether models effectively cap- ture irrironment-invaaiann features by showing how uniformly over-confident the model remains across varied environments. Through comparative analysis of datasets with dis- tributional shifts, observe that Information Bottleneck-based IRM achieves consistent calibration across different euvironments. This observation suggests that information com- pression techniques, such as Information Bottleneck are potentially effective in achieving model invariance. Furthermore, our empirical evidence indicates that models exhibiting consistent calibration across environments are also well-calibrated. This demonstrates that invariance and cross-envirommet calibration are empirically equivalent. Additionally, we underscore the necessity for a systematic approach to evaluating OOD generalization. This approach should move beyond traditional metrics, such accuracy and F1 scores, which fail to account for the model' degree of over-confidence,, and instead focus on the nuanced interplay between accuracy, calibration, and model invariance. Our code is available at https://githes http//ttt..ccco/totooo////////// Variants, Calibration "arrianta_Calibaation
1 Introduction
In the evolving landscape of machine learning, the importance of out-of-distribution (OOD) gencralization (Vapnik, 1991) and calibration (Guo al., 2017) cannot be understated, especially in real-world applications and critical scenarios. Traditional machine learning approaches, grounded in the assumption that data are independently and identically distributed (IID), often struggle to cope with OOD scenarios. Additionally, notable trend in contemporary machine learning models is their tendency towards overconfidence, which undermines the reliability of their confidence immatioo This issue is primarily recognized as a calibration challenge. Responding to these limitations, there has been a surge in research focusing on methodolo- gies like Invariant Risk Minimization (IRM) (Arjovsky et al., 2020). IRM designed to identify constant features across different environments, facilitating more robust generalization in the face of environmental "equal contributions

---

Published in Transactions on Machine Learning Research (06/2024)
variations. However, the computational demands of IRM have led to the exploration of more feasible, ap- proximate methods. Despite the development of several IRM variants, their inability to consistently surpass the performance of finely-tuned Empirical Risk Minimization (ERM) models has been observed (Gulrajani & Lopez-Paz, 2020; Zhang et al., 2023a). This performance gap raiscs critical questions about the inherent deficiencies of these IRM variants. As we showed Figure 1, even by leveraging IRM variants, there is a trade-off between in-distribution (TD) and OOD accuracy. This implies that models trained by IRM variants failed to obtain invariant features since they sacrificed ID accuracy to achieve better OOD accuracy. The invariance aimed for by IRM can, from another perspective, be conceptualized the invariance of calibra- tion degree, meaning that the model remains equally over-confident across all environments (see Section 2.3 for details). Our study intersects with this insight by evaluating the invariance of calibration through the variance of calibration metrics across environments, examining the extent. to which IRM variants mimic the original model in terms invariance. We hypothesize that an effective implementation IRM should result in consistent calibration across all environments, indicated by minimal variance in calibration. The results reveal significant disparities in how calibration performance behaves across different environments depending the IRM methodology used, even when OOD accuracy levels similar. This suggests that relying solely OOD accuracy for method development may be inadequate.
Our key contributions are as follows:
We study discrepancies between IRM for- mulation and its implemented variants us- ing the calibration metrics as a key metric to measure how much the model learned en- vironmental invariant features.
80
50
50
20
20-
We observed that Information Bottleneck- based IRM (IB-IRM) improves the calibra- tion performance across the environments (Figure and Figure 2).
1000
2000
3000
4000
5000
60
80
100
STEP
ID ACC
(a)
(b)
IRMv1
OOD INMv1
ERM
I8Mv1
We demonstrate that the trend of IB-IRM aligns more closely with the original objec- tives of IRM since its ECE variants across all environments are the smallest (Table 2).
BIRM
IB-IMM
IRM Game
PAIR
Figure 1: A graph observed in CMNIST showing the trade-off between and OOD accuracy in approximation methods of IRM. Figure (a) visualizes the accuracy of the IRMv1 model ID and OOD, plotted against the training steps the horizontal axis. The two solid lines ap- proximately symmetrie around an accuracy of 50%. Fig- (b) visualizes the ID accuracy (horizontal axis) against the OOD accuracy (vertical axis) for typical approxima- tion methods of TRM. There clear trend that the accuracy improves OOD, the ID accuracy decreases.
2 Background
2.1 OOD Generalization
In traditional deep neural network scenarios, it is assumed that the learning and test environments are
IID. The most commonly used learning algorit hi in this assumption has been ERM (Vapnik, 1991), which has been successful various scenarios. Specifically, ERM aims to minimize the sum of the risk of the model f X + y in the set of training environments Etrain-
Σ R$(
min f.
Eevuin
where (f) is represented as Ex (X((X)) YY)], which is the risk in a environment e.
However, in the real world, the assumption of IID does not always hold, and distribution shifts between training and testing environments can occur (Quinonero-Candel et al., 2008). When the model relies on features that are only effective in the training environment, known as shortcut features, can decrease the testing environment (Geirhos et al., 2020). This issue is commonly known OOD generalization and has become urgent problem recently. Therefore, it is necessary to train robust models
performance
2

---

Published in Transactions on Machine Learning Research (06/2024)
that are not dependent specific environments and can perform consistently across all environments. model that is robust across the environment, is defined as follows.
A
[[X* = E|y 15(x)) Ve, E Eall
(2)
where Eall represents the set of all possible environments. Equation (2) indicates that, for a robust model, the conditional probability of the label given the model's prediction should be equal belween any two envi- ronments.
Since it is difficult to access all real-world enviromments during training, various methods have been proposed to realize invariant models satislying (eq- (2)) using limited learning environments.
2.2 Calibration
The alignment of model's predictive confidence with its actual accuracy extremely important in practical scenarios such as medical image diagnosis and autonomous driving. For example, if there are 100 instances where model predicts with 80% confidence, then 80 of those predictions should be correct. However, in recent deep neural networks, issues such as increased model capacity have led to problems with miscalibration, and many methods have been proposed to address this calibration issue (Guo et al., 2017). The correct calibration in all environments is defined for any predictive probability a of f as follows:
B|X*)((X)) = a] = Ve Eall
Equation (3) implics that the model's conditional prodictive probability (in other words, is called confi- dence) is always consistent with its actual accuracy.
Expected Calibration Error
The quality of uncertainty calibration is often quantified as ECE (Nacini et al., 2015) and is commonly used. It is defined as follows:
M
Bm N
ECF
Σ
Acc( Bm) Conf((Bu)
(4)
Acc(Bm) and (Bm)(Bm) are defined as follows
Acc( Bm)
Σ Bm BE
19o yb), Conf(Bmm)
Σ Pb
Bm
Bm
where the predictive probabilities ranging from 0 to 1 are divided into M parts, Acc( (e((B)) is the prediction accuracy within the bin Bn and Conf(Bm) is the average of the prediction probability Po (confidence level) for each data in the bin Bm-
Adaptive Calibration Error
Nixon et al. (2019) claimed that ECE is not perfect for measuring model calibration performance in some cases. Specifically, ECE calculates the discrepancy between predicted probabilities and actual accuracy for fixed bins based predicted probabilities. If the model's predictions exhibit sharpness, this result in data imbalance within each bin. This trend of sharpness has been particularly noted in recent largo-scale neural networks. The inbalance in the number of data points within bins implies a bias in the quality of calibration estimation, potentially leading to inaccurate assessments. To address the aforementioned issue, the Adaptive Calibration Error (ACE) has been proposed (Nixon al., 2019). Unlike ECE, ACE divides the bins such that each bin contains the same number of samples, rather than being based on predicted probabilities. The formulation as follows:
3

---

Published in Transactions on Machine Learning Research (06/2024)
K ACE
Σ AAc((mmx) Con (Bm,k)
(5)
KM
where K denotes the number of classes. The predicted probabilities corresponding to class for each data point sorted and divided equally into M bins, with the m-th bin denoted Bm, In other words, each bin contains [N/M] samples.
However, it should be noted that, in our scope of experiments, there no significant difference between FCE and ACE results.
Negative Log-Lite LoggLiiei
Additionally, the Negative Log-Tikeelhood (NLL) is commonly used as indirect metric for measuring calibration performance (Minderer et al., 2021). The NLL is expressed by the following equation, and the closer the model'ss confidence in the correct label is to 1, the more favorable it is.
NLL Σ log p(yn 0)
(6)
Therefore, even if the accuracy is 100%, improving the model's confidence still enhance the NLL. If the model is overfitting on the NLL in the training domain and becoming overconfident, it is more likely to make incorrect predictions with overconfidence in the test domain. This leads to a higher NLL, indicating poor calibration performance (Mukhoti et al., 2020).
2.3 OOD Generalization and Calibration
Some empirical study (Nagamuma & Hataya, 2023) reveals the correlation between OOD generalization and calibration. Theoretically, eq. (3) can be a sufficient condition for eq- (2), demonstrating that calibration across all environments leads to OOD generalization. Wald et al. (2021) capitalize on this observation and propose CLOvE, a method that calibrates models across multiple environments, thereby enabling general- ization to unseen data distributions.
Why do use calibration to evaluate IRM variants?
Since eq. (3) is sufficient condition for eq. (2) and not a necessary one, it is possible for model to achieve OOD while making overconfident predictions. Therefore, calibration performance does not necessarily per- fectly explain O0D performance.
However, comparing the degree of calibration across different environments allows for measuring a model's invariance. The invariance condition defined by the equality of the conditional probability E[Va[[(X.a) across all environments, which can be restated the consistency of a model's miscalibration across all environments. Thus, evaluating the variance of model's calibration metric across environments provides means to assess its invariance. For example, if a model is consistently over-confident. by 0.1 across all environments, meaning that [EX = a] = min(1.0, 0.1) Ve E Eall holds, then the model satisfies invariance.
How do we approximate invariance condition by using ECE?
Each term of invariance condition eq. (2), can be approximated by the conditional accuracy, E[(a|f(X..] Ace(B...) Then, Equation (2) could be approximated as Acc(Be m) Acc(B" Ve,,''' Eau- And this equation be easily transformed into Con m. Ace(B" m) Conf(B m) Ace( m Me,, Eau- Given that recent models are empirically overconfident. (Minderer et al., 2021), we assume Conf( > Acc( ) so the equation Gap(B = Conf(Br) fore, Equation (2) can be approximated as ECF,E = ECF.VV. Vevee E Eaut and the consistency of ECE across environments can be used to approximate the model's invariance.
Ax) [Acc((B)) CBn) holds true. There-
Furthermore, our empirical results indicate that cases of equivalent calibration failure between ID and OOD rare. When calibration is consistent. across domains, it is typically in well-calibrated cases. Hence,
4

---

Published in Transactions on Machine Learning Research (06/2024)
empirical evidence suggests an equivalence between the invariance condition and calibration across all envi- ronments.
In this paper, we utilize this fact that comparing calibration performance across different environments will provide a better assessment of OOD generalization capabilities and evaluate the approximation methods of IRM, introduced Section
3
Preliminary
3.1 Domain Invariant Representation
One approach to making models more robust in OOD scenarios is through representation learning. Rep- resentation learning a method that teaches the model to learn effective data representations from raw data for better performance in predictions. Traditionally, features of data were designed by human, but representation learning reduces this necessity and can also compress the dimensions of input data, enabling faster computations.
In the context of OOD scenarios, there is application for reprosentation lcarning, specifically aimed extracting envirommenntinnvarien features from raw data and reducing environmenl-speccfic spurious features to achieve more robust predictions. For example, in the task of classifying animals from images, if the model learns to associate parts of an animal with extrancous elements, it may lead to high performance within the training domain but fail to make robust predictions when the background distribution shifts in the test domain. However, if the model discards spurious features (e-g., background) from the raw data, can consistently make robust predictions across all environments. One particularly well-known method that applics representation lcarning to OOD generalization scenarios is Domain-Adversart Neural Networks (DANN) (Ganin et al., 2016). This approach uses adversarial training to adjust the learning process that it is impossible to discern from the extracted features which environment the input comes from, thereby aiming to extract invariant features.
3.2 Invariant Learning
In the context of applying representation learning OOD generalization, the method we focus on in this paper is IRM (Arjovsky et al., 2020). IRM is proposed with the aim of extracting enviromment-invarian features from input data to enable consistent predictions across any environment. IRM is defined with the purpose of realizing (eq. (2)) as follows.
min Σ R' o b)
(7)
4:x 7 oE Feruin wiH H
subject to w E arg min G:H
(a o D). Ve E Etrain
Here, X' represents the input space of each data point, H denotes the extracted invariant feature space, and y is the output space of the model. IRM defines the function as f= w o, where maps the input data to H. and based on these extracted invariant features, the final prediction made by w. IRM aims achieve consistent prodictions as a form of OOD gencralization.
The formulation in eq. (7) represents a bi-level optimization problem, which is challenging to solve exactly. As result, various implementable approximation methods have been proposed. Many variants add various constraints to the IRM objective and approximate the bi-level optimization problem as singlc-level problem to make more tractable (Arjovsky el. al., 2020; Ahuja et al., 2022; Chen el. al., 2022; Ahuja et al., 2020; Lin et al., 2022). Additionally, Zhang et al. (2023b) propose an improved variant IRM Game, another IRM variant, that retains the bi-level optimization framework during training. However, despite numerous contributions the development of approximation methods for IRM, method that generally surpasses well- tuned FRM in terms of OOD generalization accuracy has not yet been proposed. Therefore, it is essential
5

---

Published in Transactions on Machine Learning Research (06/2024)
to understand why these approximation methods have not succeeded and what approaches might be more suitable for addressing the approximation problem.
While the original IRM formulation poses implementation challenges, an alternative method has also been proposed to address this issue. IRM aims to learn invariant features by keeping [[X*] constant across environments. However, Huh Baidya (2023) introduced complementary notion of invariance called MRI, which seeks conserve the label-conditioned feature expectation E [[ across environments, in contrast the original IRM formulation.
4
Related Works
Methods
Featurizer
Predictor
Linearity of w
Information Bottleneck
Game theory
Ensemble
Bayesian inference
IRM
w
X
X
IRMv1
④
1.0
X
X
X
IB-IRM
w ④
1.0
/
/
X
X
X
PAIR
w ④
1.0
V
X
x
IRM Game
Wene
/
X
BIRM
w. +
1.0
X
X
X
Table Approximation methods of TRM
4.1 IRMv1
IRMv1 (Arjovsky et al., 2020) constrains w in eq. (7) as a linear mapping and allows the transformation '' (w. -)) w/ 1.0. The following one-variable optimization problem can relaxed and computed.
min -
Σ R*(D) + A V R $6) 1|2
(8)
BE Pirain
4.2 Information Bottleneck based (RM(IB-IRM)
While [RMv1 has been successful with certain OOD generalization datasets such CMNIST, it has been noted that if not enough invariant features are included in each environment in the training data, robust prediction is lost. Therefore, IB-IRM (Ahuja et al., 2022) aims to improve IRMv1 by compressing the entropy of the feature extractor Φ' using the information bottleneck method, thereby preventing excessive reliance on features that adversely affoct the invariant prediction. The variance of the fcatures extracted from cach data by &' is added as a regularization term.
A ll W -4) 12 + >Var(
(9)
4.3 Pareto IRM(PAIR)
In general, there is trado-off between ERM and OOD generalization, and PAIR (Chen et al., 2022) focuses on the need lo properly manage that trade-off when generalizing a model out of distribution. To achieve more robust models, PAIR is designed to find their Pareto optimal solutions in terms multi-objective optimization. In fact, it implemented as a multi-objective optimal problem for ERM, IRMv1 penalties, and VREx (Krueger et al., 2021) penalties.
Enc 20
"Wnss Birainl

---

Published in Transactions on Machine Learning Research (06/2024)
OOD ECE
60
OOD ACE
60
NLN
000
40
40
20
20
0
0
20
40
60
80
100
0
20
40
60
80
100
ID ECE
ID ACE
ID NLL
(a) ECE va ECE
(b) ACE vs ACE
(c) NLL vs NLL
ERM
IRMv1
BIRM ID OOD
IB-IRM
IRM Game
PAIR
Figure ECE, ACE, and NLL evaluation the CMNIST. The X-axis indicates the evaluation of ID sets, and the Y-axis indicates the evaluation of OOD sets. The red solid line represents the case when the metric values are equal for both ID and OOD, indicating that the model has achieved the same level of calibration performance in both domains. IRMv1, IB-IRM, BIRM, and PAIR are distributed relatively close to the red solid line, indicating environments. Those data points that show better OOD calibration
that they achieve consistent calibration performance for each IRM variant tend to be closer to the red line.
min (CERM, Cinnsos, CIRMO,,, LVnEx) 4'-0-
(10)
where
LERM
Σ
LIRMA =ll Vo (w 12, LynES = Var((P((T') )oe Eerann))) (D))),
eEErrain
4.4
IRM Game
IRM Game (Ahuja et al., 2020) aims to achieve Nash equilibrium regarding inference accuracy by introducing game-theoretie framework to the IRM between each environment. Each environment in the training data is assigned its own predictor w", each of which is trained to be optimal in each environment, and the final Wens is the ensemble of all predictors Beran EeeF w. By learning environment-specifee predictors, the limitation linearity made in IRMv1 is eliminated, and implementation closer to eq. (7) is aimed at.
min x: x 7l Wens R R.S eEBirai
R (Wens b)
(11)
s.t. w E arg
min W
Rd
we o ob 3.
Ve
Etroón
where Wh
w + E
Ee'fe Wa!
Etraia
4.5 Bayesian (RM(BIRM)
In BIRM (Lin et al., 2022), IRMv1 may overfit the training environment the training data are insufficient, introduced Bayesian estimation approach there. Noting that the posterior probability p(w (xx)), Y*)

---

Published in Transactions on Machine Learning Research (06/2024)
is invariant in all environments given the data mapped by & and the correct label if the model is able to acquire invariant features, the penalty of IRMv1 is modified. By using Bayesian estimation instead of point estimation, they aim to improve the generalization performance of the model by taking into account the uncertainty of the data.
Table presents an overview of TRM and its approximation methods.
5
Experiments
Although the approximation methods in the previous chapter have achieved some success in terms of OOD accuracy on OOD datasets (Arjovsky et al., 2020; Ahuja et al., 2022; Chen et al., 2022; Ahuja et al., 2020; Lin et al., 2022), it not known how invariant features they actually acquire and how robust they are unknown environments. Therefore, we evaluated them by comparing their calibration performance with ECE, ACE, and NLL. Since aim to investigate OOD performance through invariance condition eq. (2), we do not use some metrics that do not incorporate the model's conditional probability, such as the Fl score.
5.1 Setting
We compared ECE of each model on four different OOD-datasets, Specifically, the datasets used were Colored MNIST (CMNIST) (Arjovsky et al., 2020), Rotated MNIST (RMNIST) (Ghifary et al., 2015), PACS (Li et al., 2017), and VLCS (Fang et al., 2013), sourced from the DomainBed benchmark (Gulrajani & Lopez-Paz, 2020). These datasets can be distinguished based on the type of distribution shift they represent (Ye et al.,
2022). The distribution shift in CMNIST is known as the correlation shift, where the conditional probability of labels given spurious features changes across environments. The other three datasets Mataaets-RMMISTT. PACS, and VL.CS are classified under diversity shift, where the prior probability of spurious features changes across environments.
For optimization, Adam (Kingma & Ba, 2015) used consistently across all models, and the tuning of learning rates and hyperparameters for each method was conducted in accordance with their respective papers.
5.2
Correlation Shift
5.2.1 Dataset
CMNIST is a dataset based on MNIST, containing 70,000 dimension examples (1, 28, 28). It is adapted for a binary classification task where digits less than are labeled as class 0, and those or above as class 1. Spurious correlations are added by assigning colors (red or green) to the digits, with the proportion of the two colors varying in each class depending on the environment. Denoting the proportion of green in class as environment e, then the training environments are set e [10%, 20%], and the test environment [90%].
5.2.2 Results
Figure 2 presents scatter plot showing the relationship between evaluation on TD and that on OOD for multiple IRM variants, cach trained with multiple hyperparameters. The red solid line represents the case when the metric values equal for both ID and OOD, indicating that the model has achieved the same level of calibration performance in both domains. IRMv1 (green), IB-IRM (light blue), BIRM (red), and PAIR (yellow) are distributed near the red solid line. Additionally, they exhibit better calibration performance in both domains. On the other hand, ERM (blue) and IRM Game (purple) are distributed above the red solid line, showing worse calibration performance in OOD. This suggests that they may be overfitting to the training domain, resulting in overconfidence in the test domain.
In some cases, IRM Game shows poor calibration performance but still appears near the red solid line. However, it does not exhibit stable distribution. This instability is likely due to the presence of separate
8

---

Published in Transactions on Machine Learning Research (06/2024)
50
60
45
10"
50
.000 ECE
ACE M
000 NLL
35
40
10
25
20
15
10
1000
5000
1000
2000
1000
2000
3000
STEP
(a) OOD Accuracy
(b) OOD ECE
(c) OOD ACE
(d) OOD NLL
À in IRMv1
10
100
1000
10000
Figure Comparison of the model's OOD calibration performance in CMNIST, with various values of X in the formulation of IRMv1 eq. (8). It was observed that increasing A, thereby increasing the regularization penalty in IRMv1, results in improved calibration performance for the model in the OOD environment.
classifiers w fitted for each environment, which may lead to overfitting within each environment, making it difficult to balance across them.
Empirically, cases, where calibration performance is consistent between ID and OOD, are limited to those where both domains exhibit better calibration performance. Specifically, as shown in Figure 2c, instances distributed along the red solid line are confined to points with low NLL in both domains. This trend is also observed ECE (Figure 2a) and ACE (Figure 2b), suggesting that, empirically, cross-cnviromment calibration is a necessary condition for invariance.
Moreover, Figure 3 visualizes how the penalty introduced in IRMv1 for the purpose of OOD generalization actually affects the calibration performance. It specifically compares different values of X in IRMv1's formu- lation eq. (8), demonstrating that as the impact of IRMv1's penalty increases, the calibration performance improves. It was also found that the information bottleneck penalty in IB-IRM is effective in improving calibration. Detailed results are presented in Figure 7 and Figure 8 in Appendix B.2.
5.3
Diversity Shift
5.3.1
Datasets
The details for RMNIST, PACS, and VLCS are follows:
RMNIST: RMNIST is datasct based on MNIST, containing 70,000 examples of the same dimension (1, 28, 28) and includes 10 classes. The environments are distinguished by the rotation angle of the digits, denoted as e. The training environments are set at e = [15", , 45, 1 60%, 60%,75% 75%], while the test environment
is at e = [0]].
PACS: PACS is a dataset containing 9,991 examples with dimensions (3, 224, 224) and includes 7 classes. The environments are distinguished by four different styles of images. Specifically, the training environments [Art].
are set [Cartoons, Photos, Sketches), and the test environment is set
VLCS:
VLCS is a dataset containing 10,729 examples with dimensions (3, 224, 224) and includes classes. There are four types of environments. The training environments are set as e [Cattech101 LabelMe, SU N09), and the test enviroument is set as e = [VOC2007].
5.3.2 Results
Considering that a model's OOD accuracy affects its calibration (Ovadia et al., 2019), we compared the calibration performance with OOD accuracy for a fairer comparison of the calibration metrics. A threshold set for validation accuracy in the training environment, and upon reaching this threshold, early stopping

---

Published in Transactions on Machine Learning Research (06/2024)
30
35
14
25
12
25
10
20
OOD ECE
ECE
20
15
10
4
54
0
0
4
6
8
10
5
10
15
20
5
10
15
ID ECE
ID
ID ECE
(a) RMNIST
(b) PACS
(c) VLCS
ERM
IRMv1
BIRM
IB-IRM
IRM Game
PAIR
OOD
Figure Comparison of the relationship between ECE in the training environment (horizontal axis) and the test environment (vertical axis). The red solid line represents the case where the ECE is equal in both enviromments. It observed that IB-IRM (in light blue) is distributed the red solid line, indicating lendency not to overfit to the training environment compared to other methods.
of training was implemented to create conditions of comparable OOD accuracy. The threshold for carly stopping was selected using grid search to minimize the variance in test accuracy across different hods
Table shows the variance in calibration performance across all environments for each method on differ- ent datascts. Despite similar OOD accuracy, the calibration performance behavior varies across mothods. Notably, IB-IRM consistently achieves lower variance in most cases, indicating that it maintains similar cal- ibration across all environments and achieves invariance. In Figure 4, it also evident that IB-IRM achieves similar calibration performance for both ID and OOD. Additional ACE and NLL results are presented in Appendix B.1.
Similar to the Correlation Shift, was observed that when calibration performance is consistent across both domains, there is a strong tendency to achieve good calibration performance. Detailed explanations along with more experimental results are provided in Appendix B.1.
6 Information Bottleneck in Calibration
In our experiments, we observe that IRM with information bottleneck is effective calibration, as shown Figure 2 and Figure 4. Therefore, in this section, we provide the results the ablation study for IB-IRM.
We investigated the impact of the penalty term related the information bottleneck in IB-IRM on calibra- tion. Figure visualizes various evaluation metrics in CMNIST when altering the strength of the information bottleneck in IB- IRM, that is, the magnitude of in eq. (9). As shown in the top row of Figure 5, which visualizes this based on OOD calibration and OOD accuracy, we confirm the clear trend that strengthening the regularization o improves both accuracy and ECE. The bottom row of Figure 5 shows the relationship between the calibration metrics in ID and OOD. Increasing 7 aligns both calibration metrics towards the red line, representing that this constitutes a successful calibration across environments. These findings suggest that not only is the IRMvl's regularization term aimed at OOD sccnarios, but the information bottlenock also further promotes the learning of invariant. features.
As an intuitive explanation of the information bottleneck approach's 's benefit for ECE is as follows:
The information bottleneck explicitly regularizes entropy.
2. This leads to the prevention of overfitting to one-hot encoded labels.
3. As result, the model trained by information bottleneck well calibrated
10

---

Published in Transactions on Machine Learning Research (06/2024)
ERM
IRMv1
IB-IRM
BIRM
IRM Game
PAIR
RMNIST
88.3+0.1
89.7+2.3
85.4+13
87.6+11
87.8+14
82.7+14
PACS
73.8
76.9219
75.5
67.221.7
72.7=1.8
76.0+1.
OOD ACC
VLCS
70.0
67.7+2.0
69.5+0.5
70.2+2.0
69.5=007
67.9+0.2
Avg.
77.411.
78.112.1
76.8
75.011.
76.7.1.3
75.5111
RMNIST
2.27+0.97
3.61+1.73
0.17 0..03
1.85+0.49
4.37+1.79
11.17+3.04
PACS
19.86+17.38
23.62=1.77
22.10+6.39
45.78+11.00
16.73=6.05
6.19=7.49
ECE Variance
VLCS
+7.97
65.99
5.80
12.57
6.49
33.50
Avg.
13.35
31.07=1.11
9.3623.28
20.1626.00
9.22=27.59
16.95210.0
RMNIST
2.31
3.68
0.18
1.85a0.48
4.37
11.65:22.6
PACS
21.15+18.69 25.33+1442
20.2447.27
47.31+
18.94=6.86
6.29
ACE Variance
VLCS
20.2012.23 67.65
6.29
12.29
6.98
33.90
Avg.
14.55=7.26
32.22=3.09
8.90=3.65
20, 20.48
10.10=3.26
17.28
RMNIST
0.01=0.00 0.02+0.01
0.02+0.00
0.02+0000
0.03
0,05
PACS
0.15++..1 0.12+0.04
0.09+0.01
0.23+0006
0.12+0.01
0.05=0.01
NLL Variance
VLCS
0.12+0.01
0.32+00.01
0. 0.10=0.00
0.11+0.01
0.10#a.01
0. 0.1900.03
Avg.
0.0910.05
0.15.0.02
0.07 10.01
0.12
0.08
0.10...02
Table 2: Comparison of calibration performance consistency across environments for RMNIST, PACS, and VLCCS. This analysis presents the variance in ECE, ACE, and NLL across environments, where smaller variances indicate more consistent calibration across environments. Each model was trained with hyperpa- rameter luning using three different seeds, and the average variance across environments was calculated. The results show that IB-IRM achieved the smallest average variance, indicating that consistently maintains calibration across different environments.
It is known that the reason for the degradation calibration performance in recent large-scale neural networks due to the overconfidence caused by their complexity (Guo et al., 2017). By keeping entropy low and maintaining the model relatively simple, overconfidence is prevented, which explains the superior calibration performance of the information bottleneck.
As technical explanation, we describe the validity of our results based on the theoretical principles of the information bottleneck. In OOD scenarios, the objective of the information bottleneck is to minimize the mutual information ((x$(()) between the input X and the intermediate representation ((X) while preserving the mutual information ((#(X); between F(X) and the output Y. Intuitively, this means learning (XX to roduce as much information from X as possible while retaining as much information about Y. When ((X) is a deterministic transformation of X, the entropy (((()) can be used minimize I(X; (XX) (Kirsch et al., 2021). In the context of IRM, which aims to learn invariant features, keeping H()(X)) small encourages + to capture invariant feat Xive from X and discard spurious features X eup It should be noted that Ahuja et al. (2022) thcoretically demonstrated, using realistic SEMs, that minimizing H()(X)) in IRM leads to discarding spurious features and grasping invariant features.
Our results empirically support the effectiveness of applying the information bottleneck method to invariant feature learning.
7 Discussion and Conclusion
In this paper, we focus on IRM, which is challenging to implement due to the difficulties of bi-level optimiza- tion. We conducted comparative evaluations using ECE, ACE, and NLL, computationally feasible metrics for calibration, understand its approxination methods better and measure how much the model learned environmental invariant features.
11

---

Published in Transactions on Machine Learning Research (06/2024)
10
104
10*
55
55
10a
10a
10'
IBRM
00D ACC
V IBRM
00D ACC
V IBRM
4.
45
45
ODD
10a
10
10c
940
40
440
10a
10a
10
35
35
30
10a
30
10a
30
10
40
20
40
60
10a
OOD ECE
OOD ACE
OOD NLL
(a)
(b)
(c)
10
10a
10a
35
10a
10a
10a
OOD 1C
25
V iBIRM
0DD ACE
30
in IBiRM
OOD NLL
in 18/RM
20
10a
102
10*
10
25
101
20
104
101
5
15
000
0
10a
10
109
10a
0
10
15
20
10
15
20
25
10"
ID ECE
ID ACE
ID NLL
(d)
(e)
(f)
ID OOD
Figure 5: Impact of the information bottleneck on calibration with IB-IRM on CMNIST task. The top row displays the relationship between calibration metrics and Accuracy in the OOD context, investigated by varying the coefficient y of the information bottleneck penalty in the IB-IRM formulation (eq. (9)). As the value of increases, both metrics show improvement. The bottom visualizes the values of calibration metrics in both ID and OOD. Altering the value of y shows that the larger the value, the more the data aligns with the red line, which indicates equality between the two calibration performances. With sufficiently large 7, the points almost perfectly distributed along the red line, indicating successful calibration across multiple environments.
First, our empirical results suggest that calibration performance is equivalent across environments only when calibration performance relatively high. This suggests that calibration across environments not only a sufficient condition for invariance, but also necessary condition empirically, making it clear that calibration is extremely effective for learning invariant features.
Tt is also suggested that IB-IRM, which applies the information bottleneck method to IRM, achieves more consistent and higher calibration under dataset shift, indicating that it performs invariant feature learning better than other variants.
Another finding from our empirical results that regularization penalties introduced for OOD generalization effectively improve calibration. However, more practical datascts (under diversity shift), we observe a trade-off between OOD calibration performance and accuracy. Although IB-IRM mitigates this tradeoff compared to IRMv1 and other variants; carefully balancing these two metrics is important in aiming for OOD generalization.
As a future challenge, it has become clear that the degree of invariant feature learning cannot be fully measured by accuracy on OOD data alone, necessitating the establishment systematic generalization metrics and measurement methods. Furthermore, a better understanding of the relationship between OOD generalization and uncertainty calibration, as well as the development of novel approximation methods for IRM based this understanding, is anticipated. In particular, gaining a deeper understanding of the impact
12

---

Published in Transactions on Machine Learning Research (06/2024)
of the information bottleneck on OOD generalization and calibration will be key to addressing the issues of OOD generalization and calibration. A more comprehensive grasp of how the information bottleneck principle affects the model's ability to generalize to unseen environments and maintain well-calibrated predictive uncertaintics will be crucial in developing more robust and reliable models.
Acknowledgments
We acknowledge the generous allocation of computational resources from the TSUBAME3.0 supercomputer facilitated by the Tokyo Institute of Technology. This assistance came through the TSUBAME Encourage- ment Program for Young/ / Female Users, whose support was instrumental this rescarch.
13

---

Published in Transactions on Machine Learning Research (06/2024)
References
K Ahuja, K Shanmugam, KR Varshney, and A Dhurandhar. Invariant risk minimization games. arXiv preprint Xrx::22002.0 04692, 2020.
K Ahuja, E Caballero, D Zhang, JC Gagnon-Audet, Y Bengio, I Mitliagkas, and Rish. Invariance principle meets information bottleneck for out-of-distribution generalization. arv preprint arXiv:2166.0600 2022.
Martin Arjovsky, L'eon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk minimization. arv
preprinl XrXio:1907.0285 028933, 2020.
Y Chen, Zhou, Y Bian, Xie, B Wu, Y Zhang, K Ma, Yang, P Zhao, B Han, et Pareto invariant risk minimization: Towards mitigating the optimization dilemma outtof-distriibution generalization. arXiv preprint ar Xiv:22006. 07766, 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.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Donain-adversarial training of neural networks. The journal of machine learning research, 17(1):2096-200) 2016.
Robert Geirhos, J'orn-Henrik Jacobsen, Clandio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. Shorteut learning in deep neural networks. Nature Machine Intelli- gence, 2(11):665-67) 673, 2020.
Muhammad Ghifary, W Bastiaan Kleijn, Mengjie Zhang, and David Balduzzi. Domain generalization for object recognition with multi-task autoencoders. In Proceedings of the IEEE international conference on computer vision, pp. 2551 2559, 2015.
Ishaan Guirajani and David Lopez-Paz. In search of lost domain generalization. arXiv preprint arXiv:2007.0148/ 2020.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In International Conference on Machine Learning, pp. 1321-1330, 2017.
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models, 2021.
Dongsung Huh and Avinash Baidya. The missing invariance principle found the reciprocal twin of invariant risk minimization, 2023.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representations, 2015.
Andreas Kirsch, Clare Tyle, and Yarin Gal. Unpacking information bottlenecks: Unifying information- theoretic objectives in deep learning, 2021.
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. 006888 2021.
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. 5542-5550, 2017.
Y Lin, Н Dong, Wang, and T Zhang. Bayesian invariant risk minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognàtion, pp. 16021-16030, 2022.
Matthias Minderer, Josip Djolonga, Rob Romijnders, Frances Hubis, Xiaolna Zhai, Neil Houlsby, Dustin Tran, and Mario Lucic. Revisiting the calibration of modern neural networks. Advances in Neural Infor- mation Processing Systems, 34:15682- -61166 2021.
14

---

Published in Transactions on Machine Learning Research (06/2024)
Jishnu Mukhoti, Viveka Kulharia, Amartya Sanyal, Stuart Golodetz, Philip Torr, and Puncet Dokania. Calibrating deep neural networks using focal loss. Advances in Newral Information Processing Systems, 33:15288 15299, 2020.
Mahdi Pakdaman Naeini, Gregory Cooper, and Milos Hauskrecht. Obtaining well calibrated probabilities using bayesian binning. In Proceedings of the AAAI conference artificial intelligence, 2015.
Hiroki Naganuma and Ryuichiro Hataya, An empirical investigation of pre-trained model selection for out- of-distribution generalization and calibration. arXiv preprint orXiv:2307.08187 2023.
Jeremy Nixon, Michael W Dusenberry, Linchuan Zhang, Ghassen Jerfel, and Dustin Tran. Measuring
calibration in deep learning. In CVPR workshops, 2019.
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, David Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snoek. Can you trust your model's uncertainly? evaluating predictive uncertainty under dataset shift. In Advances in neural information processing systems, volume 32, 2019.
Joaquin Oaadell,--Candela, Masashi Sugiyama, Anton Schwaighofer, and Neil D Lawrence. Dataset shift in machine learning. Mit Press, 2008
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748-8766 8763. PMLR, 2021.
Yangjun Ruan, Yann Dubois, and Chris J Maddison. Optimal representations for covariate shift. arXiv preprint arXi::2201.000 00005,, 2021.
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.
Vladimir Vapnik. Principles of risk minimization for learning theory. Advances in newral information processing systems, 4, 1991.
Yoav Wald, Amir Feder, Daniel Greenfeld, and Uri Shalit. On calibration and out-of-domain generalization. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.). Advonces in Neural Information Processing Systems, 2021. URL
Nanyang Ye et al. Ood-bench: Quantifying and understanding two dimensions of out-of-distribution gen- eralization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
Yihua Zhang, Pranay Sharma, Parikshit Ram, Mingyi Hong, Kush Varshncy, and Sijia Liu. What is missing in irm training and evaluation? challenges and solutions. arXiv preprinl 230 0200000.2333.. 2023a.
Yihua Zhang, Pranay Sharma, Parikshit Ram, Mingyi Hong, Kush Varshney, and Sijia Liu. What is missing in irm training and evaluation? challenges and solutions, 2023b.
15

---

Published in Transactions on Machine Learning Research (06/2024)
Appendix
A Details of experiments
A.1
Implementation of IRM Variants
This section provides detailed description of each method used in the experiments.
A.1.1
IRMv1
The learning process conducted using the loss calculated based on the formulation shown in 8. The hyperparameter À is used to adjust the impact of the IRMv1 penalty.
A.1.2
IBIRM
The learning process was performed using the loss calculated based on the formulation shown in 9, The hy- perparameters À and are used to adjust the influence of the IRMv1 penalty and the information bottlenock penally, respectively.
A.1.3 IRM Game
IRM Game has two variants: F-IRM Game and V-IRM Game. The difference between them lies whether the featurizer & fixed or not. F-IRM Game fixes & as an identity matrix, while V-IRM Game considers variable Φ. For each datasct, the variant that achieved higher performance in terms of accuracy was selected.
A.1.4 PAIR
PAIR employs weighted average of ERM loss, IRMv1 loss, and VREx loss the final loss for learning, When calculating the weights, a hyperparameter called freference is used to adjust the scaling of each loss.
A.1.5 BIRM
In BIRM, when calculating the IRMv1 loss, Monte Carlo sampling is used to sample N times from the posterior distribution of the parameters, and the average of these samples used to compute the loss. This approach takes into account the uncertainty the parameters when calculating the loss. The standard deviation of the Gaussian noise added to the parameters during sampling is treated as a hyperparameter denoted birm sd. In all experiments, N was set to
A.2 Datasets
A.2.1 Domainbed
We conducted experiments using Colored MNIST (CMNIST) (Arjovsky et al., 2020), Rotated MNIST (RM- NIST) (Ghifary et al., 2015), PACS (Li et al., 2017), and VLCS (Fang et al., 2013) datasets from DomainBed (Gulrajani & Lopez-Paz, 2020). We split each dataset into training and validation sets, with 80% used for training and the remaining 20% for validation. The environment partitions were follows:
CMNIST: Frrain
110%, 20%), Frest [90%]
RMNIST: Errain [15,, 109, 45%, 60", 75%, Erest [0']
PACS: Etrain
Phoolo, Painting, Sketch), Erest [Arl]
VLCS: Etrain
1aaltech10, MU, Fress [VOC2007]
16

---

Published in Transactions on Machine Learning Research (06/2024)
A.2.2 Hyperparameters
In the experiments, we set the batch size to 256 for CMNIST, 128 for RMNIST, and 16 for PACS and VLCS. Grid search was performed on the learning rate for all experiments, with values of [le-4, 5e-4, le-3, 5e-3). For the hyperparameters specific to each approximation method, grid search was conducted as shown in Table 3.
Parameter
A
[1. lel, 1e2, 1e3, 1e4]
[1. lel, 1e2, 1e3, 1ed]
preference
2,
hirm_ sd
[e-2, 1e-1, 2e-1]
Table 3: Hyperparameters for each method
Model selection was performed using an oracle based on accuracy in the test environment for Correlation shift, while for Diversity shift, was based on validation accuracy in the training environment. Correlation shift is known to be a relatively challenging task P[Y[X, where X represents the input data and Y represents the corresponding ground truth labels. As shown in Figure 1(b), there exists a trade-off between accuracy in the training and test environments for all methods, necessitating oracle model selection. In contrast, for Diversity shift, there is correlation between ID and OOD accuracy across all methods shown in Figure 10, allowing for model selection based on performance in the training environment.
al., 2022) as it involves a shift in the conditional probability
B Additional results
B.1
Ablation Study Diversity Shift
First, Table 4, similar to Table 2, compares the calibration performance of each method under the same early stopping settings. While Table 2 compares the variance in calibration performance across environ- ments, Table 4 provides straightforward comparison of the calibration metrics. Although the methods achieve similar OOD accuracy, their calibration performance varies, with IB-IRM achieving consistently better average calibration performance.
Next, Figure 4 visualizes the relationship between ID and OOD calibration performance for each dataset. The red solid line represents cases where calibration performance is equal in both domains, indicating successful cross-environment calibration. Notably, IB-IRM (light blue) is distributed relatively close to the red solid line in RMNIST and VLCS, while PAIR (yellow) is close to the red line in PACS. This suggests that these methods achieve consistent calibration across environments.
Combined with the results in Table 2, it can be concluded that IB-IRM consistently achieves high calibration performance across environments on average.
Additionally, from the plots in Figure 4, a trend can be observed where lower calibration performance tends to be closer to the red solid line. This implies that achieving consistent calibration across environments, in other words, achieving invariance, requires higher calibration performance.
B.2 Impact of penalties of IBIRM
Figure 3 in CMNIST compares the OOD ECE of the model when varying the regularization penalty À in IRMv1. Additionally, we compare the OOD ECE of the model when varying the coefficients of the two regularization penalties in IBIRM Equation (9) as shown in Figure and Figure 8, Similar Figure 3, the regularization penalty increases, the OOD ECE decreases.
17

---

Published in Transactions on Machine Learning Research (06/2024)
ERM
IRMv1
IB-IRM
BIRM
IRM Game
PAIR
RMNIST
88.3+0.1
89.7-23
85.4+1.
87.6+L1
87.8+14
82.7-1.
PACS
73.8
76.9=1.9
75.5208
67.211.7
72.7=18
76.0=1.6
OOD ACC
VLCS
70.0=0.8
67.7=200
69.5+0.5
70.2+2.
69.5+0.7
67.9=022
Avg.
17.4
78.1 12.1
16.8
75.011.8
76.7113
75.5_1.1
RMNIST
5.01
5.74+1.42
2.40=0.64
4.45+0.42
6.71+1.00
9.55++..00
PACS
13.05=5.21
12.97=2.15
11.98+0.13
18.61+1800
12.16=1.52
7.03=3.32
OOD ECE
VLCC
12.55
22.79
7.97 2222
9.94
8.78
17.79-2.00
Avg.
10.20=2.56
13.83=1.17
7.45#1.12
11.00#1.51
9.22+1.28
11. 11.45=2.86
RMNIST
1.
5.65
2.26=0.59
4.4020.12
6.5611.20
9.49 20.
PACS
12.99=5.22
12.91=2.09
11.86+0.41
18.5
12.06=1611
6.74+3.38
OOD ACE
VLCS
12.48 11.37
22.71 22.7110.88 L0.82
8.0711..16
9.88
8.72
17.79
Avg.
10. 14.155
13. 13.7611.15
7.40=1.05
10.94+151
9.11+1333
11.34=2.77
RMNIST
0.41
0.43
0.48+0.03
0.42=0001
0.54+0.11
0.6622.01
PACS
1.07+0.36
1.00+0.19
0.87+0.03
1.34+0.11
1.06+0.07
0.75+0.05
OOD NLL
VLCS
0.97=0.08
1.63+0.08
0.88=0.04
0.90=0.08
0.89+0.01
1.25+0.11
Avg. 0.82
1.02
0.74 10.03
0.89
0.83
0.89
Table 4: Comparison of OOD accuracy and the calibration metrics across RMNIST, PACS, and VLCS datasets. Following hyperparameter tuning for each method, measurements were taken across three different random seeds, presenting the mean and variance for cach metric. It was observed that on average, IB-IRM exhibits a lower the calibration metrics while maintaining comparable inference performance.
B.3 Feature visualization study
To investigate how differences in ECE impact the model's ability to learn invariant features, we conducted feature visualization study. Using toy CNN, visualized the latent features of IRMv1 and IB-IRM RMNIST. The visualizations were created by applying t-SNE (Van der Maaten & Hinton, 2008) to the outputs of the final convolutional layer of the CNN. The results are shown in Figure 9.
First, the table presents the variance in ECE across the six environments of RMNIST, indicating that IB- IRM achieves more stable calibration. Next, consider the four figures. The top row visualizes the feature distributions of each model in the ID setting, while the bottom row shows the OOD setting. The color of each point represents the class, as RMNIST involves 10-class classification.
In truly successful invariant feature learning scenario, the distributions should be similar in both the top and bottom rows. IB-IRM (right column) shows stable class feature distributions in both ID and OOD. In contrast, IRMv1, with its higher ECE variance, displays variations in these distributions, indicating that the features it extracts differ depending on the environment. This study demonstrates that stable calibration performance across environments reflocts true invariance.
B.4 Relationship between ID and O0D accuracy in diversity shift
As shown in Figure 1(b), in the presence of correlation shift, including the CMNIST dataset, a trade-off exists bet ween ID and OOD accuracy for all methods. However, as illustrated Figure 10, for diversity shift, positive correlation between ID and OOD accuracy was observed across all methods.
18

---

Published in Transactions on Machine Learning Research (06/2024)
30
14
25
12
25
10
20
OOD ACE
20
15
10
0
0
4
6
8
10
10
15
20
2.5
5.0
7.5
10.0
12.5
15.0
ID
ID
ID ACE
(a) RMNIST
(b) PACS
(c) VLCS
3.0
2.5
1.0
2.5
OOD NLL
ODD NLL
1.5
0.6
ODD
1.5
1.0
0.2
0.5
0.5
0.0
0.0
0.0
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.0
0.5
1.0
1.5
2.0
0.4
0.6
0.8
1.0
1.2
1.4
NLL
NLL
ID NLL
(d) RMNIST
(e) PACS
(f) VLCS
ERM
IRMv1
BIRM
IB-IRM
IRM Game
PAIR
ID OOD
Figure 6: Comparison of the relationship between the calibration metrics in the training enviromment (borizontal axis) and the test environment (vertical axis). The red solid line represents the case where the calibration metrics is equal in both environments. It was observed that IB-IRM (in light blue) is distributed near the red solid line, indicating tendency not overfit the training environment compared to other methods.
C Additional discussion and potential future directions
The experimental results demonstrated that IBIRM, which employs an information bottleneck, excels in terms of model calibration. This can be attributed to the fact that applying the information bottleneck restricts the representational power the featurizer & (X), limiting its complexity and enabling it to discard spurious features, thereby reducing dependence on them. This approach of constraining the complexity of data representations has been applicd in various contexts, and there are related studics the field of Large Language Models (LLMs) well as follows.
Ruan et al. (2021) demonstrated that applying the information bottleneck method to the vasion-language model CLIP (Radford et al., 2021) improved its OOD generalization performance. Furthermore, Hu et al. (2021) proposed a new fine-tuning method called Low-Rank Adaptation (LoRA). This method involves attaching a very low-rank linear layer as an adapter to a pre-trained model and fine-tuning only that part for a specific task, achieving performance comparable to full fine-tuning with significantly lower memory requirements. This approach can be interpreted as compressing task-specific information into low-rank representation during learning, suggesting connection with the information bottleneck method.
19

---

Published in Transactions on Machine Learning Research (06/2024)
50
000 ACC
50
000 ECE
000 NLL
30
3.0
10a
30
20
10-
10-
109
10
1000
1000
2000
3000
(a) OOD Accuracy
(b) OOD ECE
(c) OOD ACE
(d) OOD NLL
in IBIRM
10
100
1000
10000
Figure Comparison of the model's OOD calibration performance in CMNIST, with various values of y and À fixed at a sufficiently large value (10000) in the formulation of IBIRM 9. It was observed that as y increases, i.e., as the regularization penalty from the information bottleneck method becomes more significant, the calibration performance of the model in OOD scenarios improves.
60
Mww
55
50
10"
ACE
000 ECE
ACE
000 NLL
30-
000
35
30
10
20-
25
10P
10
1000
2000
1000
4000
5000
0
1000
2000
3000
4000
5000
0
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
STEP
STEP
STEP
STEP
(a) OOD Accuracy
(b) OOD ECE
(c) OOD ACE
(d) OOD NLL
À in IBIRM
1
10
100
1000
10000
Figure 8: Comparison the model's OOD calibration performance in CMNIST, with various values of and y fixed at sufficiently large value (10000) in the formulation of IBIRM 9. It observed that À increases, Le., the regularization penalty from IRMv1 in IB-IRM becomes more significant, the calibration performance of the model in OOD scenarios improves.
Based on the aforementioned research, future studies are anticipated to investigate the impact of learning through information compression on OOD generalization and calibration, as well as to develop novel met that simultaneously address these two issues.
20

---

Published in Transactions on Machine Learning Research (06/2024)
100
100
50
50
550
100
75
25
100
100
75
75
100
(a) ID IRMv1(ACC:08). ECE:00008)
(b) ID IB-IRM(ACC:98) ECE:0.01)
75
0
-25
-2
-50
-75
-75
-75
-50
-25
25
-75
-50
-25
0
25
75
(c) OOD IRCM1(ACC:92.55 ECE:0.02)
(d) IB-IRMAACC:92. ECE:0.02)
IRMv1 [B-IRM
ECE Variance
2.1e-5
1.0e-5
Figure Visualization the latent feature spaces for IRMv1 and IB-IRM on RMNIST. top row shows visualizations, and the bottom row shows OOD visualizations. The color of each point represents the class (RMNIST involves 10-class classification). When invariant features are properly acquired, the latent features should remain consistent across environments. IB-IRM, which achieves more stable ECE variance across environments, demonstrates stable feature distributions in both ID and OOD. In contrast, IRMv1, which exhibits worse ECE variance, shows variations in feature distributions between ID and OOD, suggesting it may be acquiring different features depending on the environment.
21

---

Published in Transactions on Machine Learning Research (06/2024)
90
80
96
75
80
92
70
000 ACC
75
90
88
00D
65
65
86
84
55
80
ACC
ID ACC
ACC
(a) RMNIST
(b) PACS
(c) VLCS
ERM
IRMv1
BIRM
IB-IRM
IRM Game
PAIR
Figure 10: This figures illustrate the relationship between ID and OOD accuracy each dataset of Diversity shift. It is evident that there exists a positive correlation between accuracy on the ID data and accuracy on the OOD data.
22


Reviews:
[{"review_id": "WU8dQydMP6", "reviewer": "Reviewer_Loo1", "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\nThe paper is an empirical investigation paper, so its value lies in its findings and observations. Here are some findings I think are interesting to the community of invariant feature learning:\n\n1. Reveals that IRM variants with strong regularization typically achieve lower ECE.\n\n2. Observes that Information Bottleneck-based IRM (IB-IRM) lowers ECE across environments, aligning more closely with IRM's original objectives while relatively maintaining accuracy. Discusses the theoretical basis for why the information bottleneck approach is effective for invariant feature learning.\n\n3. Identifies a trade-off between ECE and accuracy for IRMv1, underscoring the compromises needed for robust generalization and calibration. Provides ablation studies that offer useful insights into the impact of different regularization terms in IB-IRM.\n\n## Weaknesses\nI think the major weakness is that the paper only uses ECE as a metric to evaluate variants of IRM, which makes the paper not deep enough. I understand this paper is an empirical investigation paper, and it indeed has some interesting findings. However, the paper could be strengthened by using more evaluation metrics other than accuracy and ECE to provide a more comprehensive picture, e.g., F1 score, AUC, Adaptive Calibration Error, Maximum Calibration Error, Negative log likelihood. Further, I think some feature visualization could also benefit the paper, as it could provide more insights into the learned features.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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 investigates variants of Invariant Risk Minimization (IRM) using Expected Calibration Error (ECE) as a key metric to measure how well the model learns domain-invariant features. The authors conduct comparative evaluations of these IRM variants on datasets with distributional shifts. They conclude that the degree of invariant feature learning cannot be fully captured by OOD accuracy alone, necessitating systematic generalization metrics. The paper also highlights the need to better understand the relationship between OOD generalization and uncertainty calibration to develop improved IRM variants.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. [Critical] Discuss additional evaluation metrics that could be used to assess the IRM variants more comprehensively. Consider metrics such as F1 score, AUC, Adaptive Calibration Error, Maximum Calibration Error, and Negative log-likelihood. To provide new insights, apply some of these metrics to at least one of the datasets used in the study. If certain metrics are deemed not useful for this particular investigation, discuss the reasons behind this decision.\n\n\n2. [Suggested] Include feature visualizations to provide a deeper understanding of the learned features captured by the different IRM variants. One approach to consider is the visualization technique used in the paper \"Enhancing Compositional Generalization via Compositional Feature Alignment\" (https://arxiv.org/abs/2402.02851), which employs similar datasets for domain-invariant learning. By incorporating such visualizations, readers can gain insights into how these methods capture invariant features and how they differ from each other.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ZHkRxUO5z4", "reviewer": "Reviewer_GgdA", "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- This paper shows many intriguing findings for invariant learning, many IRM variants are extensively studied.  \n- The paper is well-organized.  \n- The empirical analysis is intuitive and inspiring.\n\nWeaknesses:\n- Lack of motivation during demonstrating the idea. Why using ECE as a measure and how can it act as an important measure for invariant learning performance should be more clearly explained. Moreover, why it is important to compare various IRM variants? The reason behind some methods achieving good ECE and accuracy results are not intuitively explained.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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 studies invariant risk minimization (IRM) by using Expected Calibration Error (ECE) as a matric to measure the learning of invariant features. The authors studies different variants of IRM and show which is observed to achieve both low ECE and high accuracy. Through empirical evaluation, the authors show the IB-IRM can have low ECE and maintaining original accuracy. Moreover, a tradeoff is identified to achieve robust generalization and meanwhile accurate calibration. The finding shows potential for the development of invariant learning by enhancing the prediction confidence without sacrificing accuracy.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please see weaknesses for detail.", "broader_impact_concerns": "No ethical issues.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "EHiqw76N8t", "reviewer": "Reviewer_vARo", "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": "Strength:\n* Out-of-distribution experiments are conducted on various distribution shifts.\n* The paper carefully studies the effects of regularization on the accuracy and calibration of IRM methods on OOD data.  \n\nWeakness:\n* **Confusing motivation**: The paper's main contribution is using the confidence calibration metric, expected calibration error (ECE), to measure the OOD generalization performance of invariant feature learning. It first establishes that good calibration is sufficient for OOD generalization by invoking the definition of calibration and OOD generalization. However, the experiments show that calibration does not always correlate with OOD data accuracy. Intrinsically, calibration is a measure of alignment between the uncertainty and predictions of a model. A model can be very robust on OOD data (high accuracy) while over-confident or under-confident (poor calibration). Does this mean this model is not robust? Conversely, a model can perform very poorly on OOD data (low accuracy) while being perfectly calibrated. As stated in the paper, good calibration is only a *sufficient* but not a *necessary* condition for good OOD generalization. This brings it to my question: why should we consider calibration as an important metric for invariant feature learning? since it cannot faithfully reflect a model's robustness and generalization on OOD data as the experiments suggested. \n\n* **ECE is not the best metric for measuring confidence calibration**: Expected calibration error has many drawbacks. Several prior works focusing on measuring calibration have proposed improved metrics such as Adaptive calibration error (AECE) [1] and test-based calibration error (TCE) [2].\n\n[1] Nixon, J., Dusenberry, M. W., Zhang, L., Jerfel, G., & Tran, D. (2019). Measuring Calibration in Deep Learning. In CVPR workshops (Vol. 2, No. 7).\n\n[2] Matsubara, T., Tax, N., Mudd, R., & Guy, I. (2023). TCE: A Test-Based Approach to Measuring Calibration Error. In: Proceedings of the Conference on Uncertainty in Artificial Intelligence, PMLR", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "The paper studies the interaction between confidence calibration and accuracy on out-of-distribution (OOD) data for a family of invariant risk minimization (IRM) methods. Specifically, it notices a trade-off between the calibration and accuracy of OOD data. The main takeaway is that accuracy alone cannot measure invariant feature learning. Future research should also take into account calibration.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* The main takeaway is unsatisfying in its current form. Specifically, it failed to convince me that calibration measures the quality of learned invariant features.  \n\n* It would be great if the authors could consider including other calibration metrics in the experiments.", "broader_impact_concerns": "No major concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper investigates the efficacy of approximate Invariant Risk Minimization techniques, as measured by calibration metrics measured across a variety of distribution shifts. The empirical results suggest that Information Bottleneck-based IRM is consistently well-calibrated, suggesting that it effectively captures environment-invariant features. The reviewers found the conclusions intriguing and paper and investigation well-organized. Some reviewers requested additional metrics, visualizations, and justifications, which were supplied in the author response and revision. The reviewers and I are now all in agreement that the paper meets the TMLR standards and should be published.
Decision: Accept (poster)