On Calibration and Out-of-Domain Generalization

1 Introduction . Machine learning models have recently displayed impressive success in a plethora of fields [ 19 , 9 , 41 ] . However , as models are typically only trained and tested on in-domain ( ID ) data , they often fail to generalize to out-of-domain ( OOD ) data [ 23 ] . The problem is especially pressing when deploying machine learning models in the wild , where they are required to perform well under conditions that were not observed during training . For instance , a medical diagnosis system trained on patient data from a few hospitals could fail when deployed in a new hospital . Many methods have been proposed to improve the OOD generalization of machine learning models . Specifically , there is rapidly growing interest in learning models that display certain invariance properties under distribution shifts and do not rely on spurious correlations in the training data [ 34 , 17 , 1 ] . While highlighting the need for learning robust models , so far these attempts have limited success scaling to realistic high-dimensional data , and in learning truly invariant representations [ 37 , 11 , 20 ] . In this paper , we argue that an alternative and relatively simple approach for learning invariant representations could be achieved through model calibration across multiple domains . Calibration asserts that the probabilities of outcomes predicted by a model match their true probabilities . Our claim is that simultaneous calibration over several domains can be used as an observable indicator ⇤Equal contribution 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . for favorable performance on unseen domains . For example , if we take all patients for whom a classifier outputs a probability of 0.9 for being ill , and in one hospital the true probability of illness in these patients is 0.85 while in the other it is 0.95 , then we may suspect the classifier relies on spurious correlations . Intuitively , the features which lead the classifier to predict a probability of 0.9 imply different results under different experimental conditions , suggesting that their correlation with the label is potentially unstable . Conversely , if the true probabilities in both hospitals match the classifier ’ s output , it may be a sign of its robustness . Our contributions are as follows : We prove that in Gaussian-linear models , under a general-position condition , being concurrently calibrated across a sufficient number of domains guarantees a model has no spurious correlations . We then introduce three methods for encouraging multi-domain calibration in practice . These are , in ascending order of complexity : ( i ) model selection by a multi-domain calibration score , ( ii ) robust isotonic regression as a post-processing tool , and ( iii ) directly optimizing deep nets with a multi-domain calibration objective , based on the method introduced by Kumar et al . [ 26 ] . We show that multi-domain calibration achieves the correct invariant classifier in a learning scenario presented by Kamath et al . [ 20 ] , unlike the objective proposed in Invariant Risk Minimization [ 1 ] . Finally , we demonstrate that the proposed approaches lead to significant performance gains on the WILDS benchmark datasets [ 23 ] , and also succeed on the colored MNIST dataset [ 21 ] . 2 Calibration and Invariant Classifiers . 2.1 Problem Setting . Consider observable features X , a label Y and an environment ( or domain ) E with sample spaces X , Y , E accordingly . We mostly focus on regression and binary classification , therefore Y = R or Y = { 0 , 1 } . To lighten notation , our definitions will be given for the binary classification setting and we will point out adjustments to regression where necessary . There is no explicit limitation on |E| , but we assume that training data that has been collected from a finite subset of the possible environments Etrain ⇢ E . The number of training environments is denoted by k , and Etrain = { ei } ki=1 ⇢ E , so that our training data is sampled from a distribution P [ X , Y | E = ei ] 8i 2 [ k ] . Our goal is to learn models that will generalize to new , unseen environments in E . Ideally , we would like to learn a classifier that is optimal for all environments E . Unfortunately , we only observe data from the limited set Etrain and even if this set is extremely large , the Bayes optimal classifiers on each environment do not necessarily coincide . Following other recent work [ 34 , 17 , 1 ] we therefore aim for a different goal – learning classifiers whose per-instance output will be stable across environments E , as we explain below . We assume the data generating process for E , X , Y follows the causal graph in Figure 1 . 2 We differentiate between causal and anti-causal components of X , and further differentiate between the anti-causal variables which are affected or unaffected by E , denoted as Xac-spurious and Xac-non-spurious , respectively . As an illustrative example , consider again predicting illness across different hospitals . When predicting lung cancer , Y , from patient health records , Xcausal could be features like smoking . Xac-non-spurious are symptoms of Y such as infections that appear in chest X-rays , while Xac-spurious can be marks that technicians put on X-rays as in [ 51 ] . Smoking habits may vary across hospital populations , as might X-ray markings ; but the influence of smoking on cancer and the manifestation of cancer in an X-ray do not vary by hospital . We do not assume to know how to partition X into Xcausal , Xac-spurious , Xac-non-spurious . The main assumptions made in the causal graph in Fig . 1 are that there are no hidden variables , and that there is no edge directly from environment E to the label Y . Such an arrow would imply the conditional distribution of Y given X can be arbitrarily different in an unseen environment E , compared to those present in the training set . Note that for simplicity we do not include arrows from Xcausal to Xac-spurious and Xac-non-spurious but they may be included as well . 2See Appendix A.3 for a brief introduction to causal graphs . We will say a representation ( X ) contains a spurious correlation with respect to the environments E and label Y , if Y ? 6 E | ( X ) ; this motivates our naming of Xac-spurious and Xac-non-spurious in Fig . 1 , as Y ? 6 E | Xac-spurious but Y ? E | Xac-non-spurious . Similar observations have been made by [ 17 , 1 ] . Having a spurious correlation implies that the relation between ( X ) and Y depends on the environment – it is not transferable nor stable across environments . In this work we will simply consider the output f ( X ) of a classifier f : X ! [ 0 , 1 ] as a representation . The crux of this paper is the observation that having E [ Y | f ( X ) , E = e ] = f ( X ) for every value of E , i.e . f being a calibrated classifier across all environments , is equivalent up-to a simple transformation to having Y ? E | f ( X ) , and thus to f having no spurious correlations with respect to E. We prove this assertion in section 2.2 , and as a demonstration of this principle we prove ( section 3 ) that linear models which are calibrated across a diverse set of environments E are guaranteed to discard Xac-spurious as viable features for prediction . 2.2 Invariance and Calibration on Multiple Domains . We define calibration , along with a straightforward generalization to the multiple environment setting . Definition 1 . Let P [ X , Y ] be a joint distribution over the features and label , and f : X ! [ 0 , 1 ] a classifier . Then f ( x ) is calibrated w.r.t to P if for all ↵ 2 [ 0 , 1 ] in the range of f , EP [ Y | f ( X ) = ↵ ] = ↵ . In the multiple environments setting , f ( x ) is calibrated on Etrain if for all ei 2 Etrain and ↵ in the range of f restricted to ei , E [ Y | f ( X ) = ↵ , E = ei ] = ↵ . For regression problems , we consider regressors that output estimates for the mean and variance of Y , and say they are calibrated if they match the true values similarly to the definition above . The precise definition can be found in the supplementary material . We now tie the notion of calibration on multiple environments with OOD generalization , starting with its correspondence with our definition of spurious correlations . Recall that a representation ( X ) does not contain spurious correlations if Y ? E | ( X ) . Treating the output f ( X ) of a classifier as a representation of the data , and considering classifiers satisfying the above conditional independence with respect to training environments , we arrive at a definition of an invariant classifier . Definition 2 . Let f : X ! [ 0 , 1 ] . f is an invariant classifier w.r.t Etrain if for all ↵ 2 [ 0 , 1 ] and environments ei , ej 2 Etrain , where ↵ is in the range of f restricted to each of them : E [ Y | f ( X ) = ↵ , E = ei ] = E [ Y | f ( X ) = ↵ , E = ej ] . ( 1 ) Lemma 1 gives the correspondence between invariant classifiers and classifiers calibrated on multiple environments . The proof is in Section A.1 of the supplementary material . Lemma 1 . If a binary classifier f is invariant w.r.t Etrain , then there exists some g : R ! [ 0 , 1 ] such that ( i ) g f is calibrated on all training environments , and ( ii ) the mean squared error of g f on each environment does not exceed that of f . On the other hand , if a classifier is calibrated on all training environments it is also invariant w.r.t Etrain . Now , we can note how the above notion of invariance relates to that of Invariant Risk Minimization [ 1 ] , where invariance of a representation : X ! H is linked to a shared classifier w⇤ : H ! [ 0 , 1 ] , w⇤ being optimal on all environments w.r.t a loss l : [ 0 , 1 ] ⇥Y ! R 0 . Under the representation ( X ) = f ( X ) , and the cross-entropy or squared losses it turns out that the original IRM definition coincides with Equation ( 1 ) 3 . Hence we aim for a similar notion of conditional independence , yet we approach it from the point-of-view of calibration . In Section 5 we will see that taking this approach leads to different methods that are highly effective in achieving and assessing invariance . We further note that the original IRM objective was deemed too difficult to optimize by the original IRM authors , leading them to propose an alternative called IRMv1 . This alternative however does not capture the full set of required invariances , as shown by [ 20 ] , whereas we show in section 6.1 that multi-domain calibration does indeed capture the required invariances . Having established the connection between calibration on multiple environments and invariance , there are several interesting questions and points to consider : Calibration and sharpness . Calibration alone is not enough to guarantee that a classifier performs well ; on a single environment , always predicting E [ Y ] will give a perfectly calibrated classifier . 3See Observation 2 in [ 20 ] for a proof . Hence , multi-domain calibration should be combined with some sort of guarantee on accuracy . In the calibration literature , this is often referred to as sharpness . To this end , in Section 5 we will propose regularizing models during training or fine-tuning with Calibration Loss Over Environments ( CLOvE ) . Combining this regularizer with standard empirical loss functions helps balance between sharpness and multi-domain calibration . Even without training a new model , we will propose methods for model selection and post-processing that are very easy to apply and help improve multi-domain calibration without a significant effect on the sharpness of the models . Generalization and dependence on Xac-spurious . Suppose that f ( X ) is calibrated on Etrain . Under what conditions does this imply it is calibrated on E ? It is easy to show that calibration on several environments entails calibration on any distribution which can be expressed as a linear combination of the distributions underlying said environments . However , can we go beyond that ? Given a general set E we would like to know what conditions and how many training environments are required for calibration to generalize . We also wish to understand when does calibration over a finite set of training environments indeed guarantee that a classifier is free of spurious correlations . We now turn to answer these questions in the setting of linear-Gaussian models .

Summary: This paper draws an important connection between out-of-domain generalization and calibrated models. Specifically, it is argued that calibrated models are related to invariant representations w.r.t spuriously correlated features, which is necessary for generalization across multiple domains. The paper proves this statement in the case of a multivariate linear-Gaussian model. The authors empirically find a correlation between the Expected Calibration Error (ECE) and OOD performance on Colored MNNIST. Based on these insights, the authors propose three methods for achieving multi-domain calibration of trained models in practice: Approach #1 suggests doing model selection based on the Expected Calibration Error; Approach #2 is intended to augment the trained model by post-processing calibration methods and Approach #3 introduces a calibration objective at train time. In multiple experiments across 5 datasets, reported results suggest that their method outperforms the baselines algorithms without calibration on OOD test environments.