Robust Calibration with Multi-domain Temperature Scaling
Abstract
Uncertainty quantiﬁcation is essential for the reliable deployment of machine learning models to high-stakes application domains.
Uncertainty quantiﬁcation is all the more challenging when training distribution and test distribution are different, even if the distribution shifts are mild.
Despite the ubiquity of distribution shifts in real-world applications, existing uncertainty quantiﬁcation approaches mainly study the in-distribution setting where the train and test distributions are the same.
In this paper, we develop a systematic calibration model to handle distribution shifts by leveraging data from multiple domains.
Our proposed method—multi-domain temperature scaling—uses the heterogeneity in the domains to improve calibration robustness under distribution shift.
Through experiments on three benchmark data sets, we ﬁnd our proposed method outperforms existing methods as measured on both in-distribution and out-of-distribution test sets.
Introduction
To make learning systems reliable and fault-tolerant, predictions must be accompanied by uncertainty estimates.
A signiﬁcant challenge to accurately codifying uncertainty is the distribution shift that typically arises over the course of a system’s deployment [Quiñonero-Candela et al., 2008].
For example suppose health providers from 20 different hospitals employ a model to make diagnostic predictions from fMRI data.
The distributions across hospitals could be quite different as a result of differing patient populations, machine conditions and so on.
In such a setting, it is critical to provide uncertainty quantiﬁcation that is valid for every hospital—not just on average across all hospitals.
Going even further, our uncertainty quantiﬁcation should be informative when a new 21st hospital goes online, even if the distribution shifts from those already encountered.
As another example, a centralized model is trained on training data from existing clients in federated learning.
It is important for the central server to provide uncertainty quantiﬁcation for every client.
Similar to the fMRI example, the centralized model should still produce valid uncertainty quantiﬁcation for unseen new clients.
Another example is applying animal recognition models on images in wildlife monitoring, where one set of camera traps corresponds to one domain, and the model will be deployed under distribution shift, i.e., new camera traps.
In this work, we study calibration in the multi-domain setting.
We ﬁnd that by requiring accurate calibration across all observed domains, our method provides more accurate uncertainty quantiﬁcation on unseen domains.
Calibration is a core topic in learning [Platt et al., 1999; Naeini et al., 2015; Gal and Ghahramani, 2016; Lakshminarayanan et al., 2017; Guo et al., 2017; Bates et al., 2021], but most techniques are targeted at settings with no distribution shift.
To see this, we consider a simple experiment on the [Figure] ImageNet-C [Hendrycks and Dietterich, 2019] dataset, which consists of 76 domains.
Here, each domain corresponds to one type of data corruption applied with a certain severity.
We apply the temperature scaling technique [Guo et al., 2017] on the pooled data from all domains.
In Figure 1(a) and 1(b), we display the reliability diagrams for the pooled data and for one individual domain.
We ﬁnd that even under a relatively mild distribution shift—i.e., subpopulation shift from the mixture of all domains to the single domain—temperature scaling does not produce calibrated conﬁdence estimates on the stand-alone domain.
This behavior is pervasive; in Figure 1(c), we see that the calibration on individual domains is much worse than the the reliability diagram from the pooled data would suggest.
To address this issue, we develop a new algorithm, multi-domain temperature scaling , that leverages multi-domain structure in the data.
Our algorithm takes a base model and learns a calibration function that maps each input to a different temperature parameter that is used for adjusting conﬁdence in the base model.
Empirically, we ﬁnd our algorithm signiﬁcantly outperforms temperature scaling on three real-world multi-domain datasets.
In particular, in contrast to temperature scaling, our proposed algorithm is able to provide well-calibrated conﬁdence on each domain.
Moreover, our algorithm largely improves robustness of calibration under distribution shifts.
This is expected, because if the calibration method performs well on every domain, it is likely to have learned some structure that generalizes to unseen domains.
Theoretically, we analyze the multi-domain calibration problem in the regression setting, providing guidance about the conditions under which robust calibration is possible.
Contributions.
The main contributions of our work are as follows: Algorithmically, we develop a new calibration method that generalizes the widely used temperature scaling concept from singledomain to multi-domain.
The proposed new method exploits multi-domain structure in the data distribution, which enables model calibration on every domain.
We conduct detailed experiments on three real-world multi-domain datasets and demonstrate that our method signiﬁcantly outperforms existing calibration methods on both in-distribution domains and unseen out-of-distribution domains .
Theoretically, we study multi-domain calibration in the regression setting and develop a theoretical understanding of robust calibration in this setting.
Related Work
Calibration methods.
There is a large literature on calibrating the well-trained machine learning models, including histogram binning [Zadrozny and Elkan, 2001], isotonic regression [Zadrozny and Elkan, 2002], conformal prediction Vovk et al. [2005], Platt scaling [Platt et al., 1999], and temperature scaling [Guo et al., 2017].
These calibration methods apply a validation set and postprocess the model outputs.
As shown in Guo et al.
[2017], temperature scaling, a simple method that uses a single (temperature) parameter for rescaling the logits, performs surprisingly well on calibrating conﬁdences for deep neural networks.
We focus on this approach in our work.
More broadly, there has been much recent work develop methods to improve calibration for deep learning models, including augmentation-based training [Thulasidasan et al., 2019; Hendrycks et al., 2019b], calibration for neural machine translation [Kumar and Sarawagi, 2019], neural stochastic differential equation [Kong et al., 2020], self-supervised learning [Hendrycks et al., 2019a], ensembling [Lakshminarayanan et al., 2017], and Bayesian neural networks [Gal and Ghahramani, 2016; Gal et al., 2017], as well as statistical guarantees for calibration with black-box models Angelopoulos et al.
Calibration under distribution shifts.
Ovadia et al.
[2019] conduct an empirical study on model calibration under distribution shifts and ﬁnd that models are much less calibrated under distribution shifts.
Minderer et al.
[2021] revisit calibration of recent state-of-the-art image classiﬁcation models under distribution shifts and study the relationship between calibration and accuracy.
Wald et al. [2021] study model calibration and out-of-distribution generalization.
Other works consider providing uncertainty estimates under structured distribution shifts, such as covariate shift [Tibshirani et al., 2019; Park et al., 2021], label shift [Podkopaev and Ramdas, 2021], and f -divergence balls [Cauchois et al., 2020].
Another line of work studies calibration in the domain adaptation setting [Wang et al., 2020; Park et al., 2020], which require unlabeled samples from the target domain.
Problem setup
Notation.
We denote the input space and the label set by X ✓ R d and Y = { 1 , . . .
, J } .
We let [ x ] i denote the i -th element of vector x .
We use P ( X ) to denote the marginal feature distribution on input space X , P ( Y | X ) to denote the conditional distribution, and P ( X, Y ) to denote the joint distribution.
For the multiple domains scenario, we let P k ( X ) and P k
( Y | X ) denote the feature distribution and conditional distribution for the k -th domain.
We let f : X !
R J denote the base model, e.g., a deep neural network, where J is the total number of classes.
We assume f returns an (unnormalized) vector of logits.
Throughout the paper, the base model is trained with training data and will not be modiﬁed.
The class prediction of model f on input x 2 X is denoted by ˆ y =
argmax
j 2 { 1 ,...,J } [ f ( x ; ✓ )] j .
We use 1 {·} to represent the indicator function.
We use h ( · ; f,  ) :
X ! [0 , 1] to denote a calibration map (parameterized by  ) that takes an input x 2 X and returns a conﬁdence score—this is a post-processing of the base model f .
We let ˆ ⇡ = h ( x ; f,  ) 2 [0 , 1] denote the conﬁdence estimate for sample x when using model f .
For instance, if we have 100 predictions { ˆ y 1 , . . .
, ˆ y 100 } with conﬁdence ˆ ⇡ 1 = · · · = ˆ ⇡ 100 = 0 . 7 , then the accuracy of f is expected to be 70% on thesesamples (if the conﬁdence estimate is well calibrated).
Data from the domains P 1 , . .
, P K are used for learning the calibration models, and we call the in-distribution (InD) domains.
We use e
P to denote the unseen out-of-distribution (OOD) domain which is not used for calibrating the base model.
Our goal is to learn a calibration map h that is well calibrated on the OOD domain e
P .
To do this, we will learn a calibration map that does well on all InD domains simultaneously.
To measure calibration, we ﬁrst review the deﬁnition of approximate expected calibration error.
Deﬁnition 2.1 (ECE) .
For a set of samples D = { ( x i , y i ) } ni =1 with ( x i , y i ) i .
i .
d .
⇠ P ( X, Y ) , the (empirical) expected calibration error (ECE) with M bins evaluated on D is deﬁned as ECE [Equation] B
[Equation]
[Equation]
[Equation]
and B m , acc( B m ) , conf( B m ) are deﬁned as
[Equation]
where ˆ ⇡ i and ˆ y i are the conﬁdence and predicted label of sample x i .
The empirical ECE deﬁned in Eq.
(1) approximates the expected calibration error (ECE) E [ | p  P (ˆ y =
y
| ˆ ⇡ = p ) | ] with bin size equal to M Naeini et al. [2015]; Guo et al. [2017]; see Lee et al.
[2022] for statistical results about about the empirical ECE as an estimator.
The perfect calibrated map corresponds to the case when P (ˆ y =
y
| ˆ ⇡ = p ) =
p holds for all p 2 [0 , 1] .
Multi-domain calibration.
Although the standard ECE measurement in Eq.
(1) provides informative evaluations for various calibration methods in the single-domain scenario, it does not provide ﬁnegrained evaluations when the dataset consists of multiple domains, P 1 , . .
, P K .
It is possible that the
ECE evaluated on the pooled data D pool K
= D 1 [ · · · [ D K is small while the ECE evaluated on one of the domains is large.
For example, as shown in Figure 1(c), there may exist a domain, k 2 [ K ] , such that the ECE evaluated on domain k is much higher than the ECE evaluated on the pooled dataset, i.e., ECE( D k )  ECE( D pool K ) .
In the fMRI application mentioned in Section 1, producing well-calibrated conﬁdence on data from every hospital is a more desirable property compared to only being calibrated on the pooled data from all hospitals.
Therefore, it is natural to consider the ECE evaluated on every domain, which we refer to as “per-domain ECE.”
Next, we introduce the notion of
Multi-domain ECE to formalize per-domain calibration.
Deﬁnition 2.2 (Multi-domain ECE) .
For a dataset D pool K = D 1 [ · · · [ D K consisting of samples from K domains, where D k = { ( x i,k , y i,k ) } n k i
=
and ( x i,k , y i,k ) i .
i .
d .
⇠ P k
( X, Y ) , the (empirical) multi-domain expected calibration error (Multi-domain ECE) with M bins evaluated on D pool K is deﬁned as MDECE ( D pool K )
= 1 K
P
[Equation]
ECE [Equation] D k ) .
Remark 2.3. In Deﬁnition 2.2, we weight each domain equally to balance across domains, which could better reﬂect how the calibration method performs on each individual domain. Furthermore, in our experiments, we also visualize the ECE measured on each domain to provide additional information on model performance on every domain. Compared with the standard ECE evaluated on the pooled dataset, multi-domain ECE provides information about per-domain model calibration.
In the multi-domain setting, we aim to learn a calibration map ˆ h that can produce calibrated conﬁdence estimates on every InD domain.
Intuitively, if the unseen OOD domain e
D is similar to one or multiple InD domains, ˆ h can still provide reliable conﬁdence estimates on the new domain.
We formally study the connection between “well-calibrated on each InD domain” and “robust calibration on the OOD domain” in Section 5.
Temperature scaling.
Next, we review a simple and effective calibration method, named temperature scaling (TS) [Platt et al., 1999; Guo et al., 2017], that is widely used in single-domain model calibration.
Temperature scaling applies a single parameter T > 0 and produces the conﬁdence prediction for the base model f as
[Equation]
where [Softmax( z )]
j = exp([ z ] j ) /
P
Ji =1 exp([ z ] i ) .
The parameter T is the so-called temperature , with larger temperature yielding more diffuse probability estimates.
To learn the temperature parameter T from dataset D
= { ( x i , y i ) } ni =1 , Guo et al.
[2017] propose to ﬁnd T by solving the following convex optimization problem,
[Equation]
which optimizes the temperature parameter such that the negative log likelihood is minimized.
We use
TS-Alg to denote the temperature scaling learning algorithm; given inputs dataset D and base model f , TS-Alg outputs the learned temeperature parameter by solving Eq.
e.g., ˆ T = TS-Alg ( D , f ) .
Multi-domain temperature scaling
We propose our algorithm—multi-domain temperature scaling—that aims to improve the calibration on each domain.
One key observation is that if we apply temperature scaling to each domain separately, then TS is able to produce calibrated conﬁdence on every domain.
Therefore, the question becomes how to “aggregate” these temperature scaling models and learn one calibration model, denoted by ˆ h , that has similar performance to the k -th calibration model ˆ h k evaluated on domain k for every k 2 [ K ] .
At a high level, we propose to learn a calibration model that maps samples from the input space X to the temperature space R + .
To start with, we learn the temperature parameter ˆ T k for the base model on every domain k by applying temperature scaling on D k .
Next, we apply the base deep model to compute feature embeddings of samples from different domains, 1 and label feature embeddings from the k -th domain with ˆ T k .
In particular, we construct K new datasets, ˆ D 1 , . . .
, ˆ D K , where each dataset contains feature embeddings and temperature labels from one domain, i.e., ˆ D k
= { ( ( x i,k ) , ˆ T k ) } n k i =1 .
Finally, we apply linear regression on these labeled datasets.
In detail, our algorithm is as follows:
Learn temperature scaling model for each domain.
For every domain k , we learn temperature ˆ T k by applying temperature scaling on validation data D k
= { ( x i,k , y i,k ) } n k i =1 from k -th domain, i.e., ˆ T k = TS-Alg ( D k , f ) and TS-Alg denotes the TS algorithm.
Learn linear regression of temperatures.
Extract the feature embeddings of the base deep model f on each domain.
Use ( x i,k ) 2 R p to denote the feature embedding of the i -th sample from k -th domain.
Then we learn ˆ ✓ by solving the following optimization problem,
[Equation]
Predict temperature on unseen test samples.
Given an unseen test sample
[Equation]
we ﬁrst compute the predicted temperature
e
T using the learned linear model [Equation]
h
[Equation]
x [Equation] , ˆ ✓ i .
Then we output the conﬁdence estimate for sample [Equation]
x as
[Equation]
We denote our proposed method by MD-TS (
M ult- D omain T emperature S caling).
A presentation of the algorithm in pseudocode can be found in Algorithm 1, Appendix A.
We pause to consider the basic concept in more detail.
The goal of our proposed algorithm is to predict the best temperature for samples from different several domains.
In an ideal setting where the learned linear model ˆ ✓ results in good calibration on every InD domain, we can expect that ˆ ✓ will continue to yield good calibration on the OOD domain e
P when e
P is close to one or several InD domains.
For example, e
P will work well if e
P is a mixture of the K domains, i.e., e
P =
P
Kk =1 ↵ k P k and ↵ 2  K  1 .
Regarding the algorithmic design, linear regression is one of the simplest models for solving the regression problem.
It is computationally fast to learn such linear models as well as make predictions on new samples, making it attractive.
We test alternative, more ﬂexible, regression algorithms in Section 4 but do not observe signiﬁcant gains over linear regression.
To illustrate how our proposed algorithm MD-TS performs differently from standard TS, we return to the
ImageNet-C dataset.
We compare the predicted temperature of our algorithm on new samples from domain k with the temperature that results from running TS on domain k alone.
The results are summarized in
Figure 2, where each circle corresponds to the mean predicted temperature on one InD domain.
For each domain, we also visualize the standard deviation of the predicted temperatures for samples from that domain (the horizontal bar around each point).
We ﬁnd that our algorithm predicts the temperature quite well.
Note that it does not have access to the domain index information of the fresh samples.
By contrast, TS always uses the same temperature, regardless of the input point.
Experiments
In this section, we present experimental results evaluating our proposed method, demonstrating its effectiveness on both in-distribution and out-of-distribution calibration.
We focus on three realworld datasets, including ImageNet-C [Hendrycks and Dietterich, 2019]—a widely used robustness [Table] image {+dataset, WILDS-RxRx
+}
[Koh et al., 2021] image of cells (by microscopy) et al., 2020] landmark recognition dataset in federated learning.
Additional experimental results and implementation details can be found in Appendix C.
Our code is available at https://github. com/yaodongyu/MDTS .
Datasets.
We evaluate different calibration methods on three datasets, ImageNet-C, WILDS-RxRx1, and GLDv2.
ImageNet-C contains 15 types of common corruptions where each corruption includes ﬁve severity levels.
Each corruption with one severity is one domain, and there are 76 domains in total (including the standard ImageNet validation dataset).
We partition the 76 domains into disjoint in-distribution domains and out-of-distribution by severity level or corruption type.
WILDS-RxRx1 is a domain generalization dataset, and we treat each experimental domain as one domain.
We adopt the default val/test split in Koh et al.
[2021]: use the four validation domains as in-distribution domains and the 14 test domains as the out-of-distribution domains.
We also provide experimental results of other random splits in Appendix C. For GLDv2, each client corresponds to one domain, and there are 823 domains in total.
We randomly select 500 domains for training the model, and then use the remaining 323 domains for evaluation denoted by validation domains.
We further screen the validation domains by removing the domains with less than 300 data points.
There are 44 domains after screening, and we use 30 domains as in-distribution domains and the remaining 14 domains as out-of-distribution domains.
For all datasets, we randomly sample half of the data from in-distribution domains for calibrating models and use the remaining samples for InD ECE evaluation.
We use all the samples from OOD domains for ECE evaluation.
Models and training setup.
We consider multiple network architectures for evaluation, including
ResNet-50 [He et al., 2016], ResNext-50 [Xie et al., 2017], DenseNet-121 [Huang et al., 2017], BiTM-50 [Kolesnikov et al., 2020], Efﬁcientnet-b1 [Tan and Le, 2019], ViT-Small, and ViT-Base [Dosovitskiy et al., 2020].
To evaluate on ImageNet-C, we directly evaluate models that are pre-trained on ImageNet [Deng et al., 2009].
For WILDS-RxRx1 and GLDv2, we use the ImageNet pre-trained models as initialization and apply SGD optimizer to training the models on training datasets.
Evaluation metrics.
We use the Expected Calibration Error (ECE) as the main evaluation metric.
We set the bin size as 100 for ImageNet-C, and set bin size as 20 for WILDS-RxRx1 and GLDv2.
We evaluate ECE on both InD domains and OOD domains.
Speciﬁcally, we evaluate the ECE of each InD/OOD domain.
Meanwhile, we also evaluate the ECE of the pooled InD/OOD domains, i.e., the ECE evaluated on all samples from InD/OOD domains.
We use unseen samples from the InD domain to measure the per-domain ECE.
We also measure the averaged per-domain ECE results (i.e., per-domain ECE averaged across domains).
4.1 Main results
We summarize the ECE results of different methods on three datasets in Table 1 and Figure 3. We use
TS to denote temperature scaling [Guo et al., 2017], and use MSP to denote applying the maximum softmax probability [Hendrycks and Gimpel, [Figure]
In Table 1, we use the ImageNet validation dataset and ImageNet-C datasets with severity level s 2 { 1 , 5 } as the InD domains and use the remaining datasets as OOD domains.
We present the averaged per-domain ECE results in Table 1, and visualize the ECE of each domain in Figure 3.
As shown in Table 1 and Figure 3(a)-3(c), we ﬁnd that our proposed approach achieves much better InD per-domain calibration compared with baselines.
Also, TS does not signiﬁcantly improve over MSP on ImageNet-C InD domains in Table 1, but our proposed method largely improve the ECE compared with MSP and TS.
For instance, the ECE results of MSP and TS on Efﬁcientnet-b1 are 6.93 and 6.99, and our method achieves 3.84.
Intuitively, when there are a diverse set of domains in the calibration dataset, a single temperature cannot provide well-calibrated conﬁdences.
In contrast, our proposed method is able to produce much better InD conﬁdence estimates by leveraging the multi-domain structure of the data.
Next we study the performance of different methods on out-of-distribution domains.
From Table 1, we ﬁnd that MD-TS achieves the best performance on OOD domains arcoss all the settings.
On
ImageNet-C with BiT-M-R50, MD-TS improves the ECE from 6.54 (MSP) to 4.05, while the performance of TS is similar to MSP.
Moreover, MD-TS signiﬁcantly outperforms MSP and TS on
WILDS-RxRx1, where MD-TS improves over TS by around 5.00 measured in ECE.
Figure 3(d)-3(f) display the per-domain ECE performance on out-of-distribution domains.
MD-TS improves over TS on more than half of the domains in all three datasets.
For the remaining domains, MD-TS performs slightly worse than TS.
Furthermore, on those domains that TS performs poorly (ECE > 8), MD-TS largely improves over TS by large margins.
Further comparisons in Appendix C.6 show that these improvements continue to hold when relative to two other calibration techniques: MC dropout [Gal and Ghahramani, 2016] and deep ensembles [Lakshminarayanan et al., 2017].
4.2 Predicting generalization
Suppose a model can produce calibrated conﬁdences on unseen samples, in which case we could leverage the calibrated conﬁdence to predict the model performance.
Speciﬁcally, based on the deﬁnition of ECE in Eq.
(1), when the model is well-calibrated, the average of the calibrated conﬁdence is close to the average accuracy, i.e., Conf ( D ) ⇡
Acc ( D ) .
Meanwhile, predicting model [Table] [Figure] performance accurately is an essential ingredient in developing reliable machine learning systems, especially under distributional shifts [Guillory et al., 2021].
As shown in Table 1, we ﬁnd that our proposed method produces well-calibrated conﬁdence values on both InD and OOD domains.
We now measure its performance on predicting model performance and compare with existing methods.
We measure the performance using mean absolute error (MAE), MAE =
(1 /K ) ·
[Equation]
Acc [Equation] k ) | where S k is the dataset from the k -th domain.
We show the predicting model accuracy results in Table 2.
MD-TS signiﬁcantly improves over existing methods on predicting model performance across all three datasets.
For example, on ImageNet-C, calibrated conﬁdence of MD-TS produces fairly accurate predictions on both InD and OOD domains (less than 2% measured in MAE), which largely outperforms MSP and TS.
In Figure 4, we compared the prediction performance of TS and MD-TS on every OOD domain.
We ﬁnd that MD-TS achieves better prediction performance compared to TS on most of the domains.
Refer to Appendix C.1 for more results in which other architectures are tested.
4.3 MD-TS ablations
To learn a calibration model that performs well per-domain, we apply linear regression on feature representations  ( x k ) such that h  ( x k ) , ✓ i ⇡ ˆ T k , where x k is from domain k and ˆ T k is the temperature parameter for domain k .
We investigate other methods for learning the map from feature representations to temperatures in a regression framework.
Speciﬁcally, beside the ordinary least squares (OLS) used in Algorithm 1, we consider ridge regression (Ridge), robust regression with
Huber loss (Huber), kernel ridge regression (KRR), and K -nearest neighbors regression (KNN).
The implementations are mainly based on scikit-learn [Pedregosa et al., 2011].
We use grid search (on InD domains) to select hyperparameters for Ridge, Huber, KRR, and KNN.
We summarize the comparative results for different regression algorithms in Table 3.
Compared to
OLS, other regression algorithms do not achieve signiﬁcant improvement.
Speciﬁcally, KRR achieves slightly better performance on OOD domains, while other algorithms have similar performance compared to OLS.
Moreover, there are no hyperparameter in OLS, which makes it more practical in [Table] real-world problems.
Meanwhile, the results suggest that our proposed MD-TS is stable to the choice of speciﬁc regression algorithms.
Theoretical analysis
In this section, we provide theoretical analysis to support our understanding of our proposed algorithm in the presence of distribution shifts.
We use h ? k
( · ) = h ( · ; f,  ? k ) :
X ! [0 , 1] to denote the best calibration map for the base model f on the k -th domain; this map minimizes the expected calibration error (ECE) E [ | p  P (ˆ y = y
| ˆ ⇡ = p ) | ] over distribution P k .
We also call h ?
k a hypothesis in the hypothesis class H .
Next, given the ﬁxed base model f , we aim to learn ˆ h ( · ) =
h ( · ; f, ˆ  ) such that " (ˆ h, P k,X ) =
E X ⇠ P k,X [ | h ?
k ( X )  ˆ h ( X ) | ] is small for every domain k , where " k (ˆ h ) denotes the risk of ˆ h w.r.t.
the the best calibration map h ?
k under domain P k .
In addition, we are interested in generalizing to new domains: suppose there is an unseen OOD domain
[Equation] and its marginal feature distribution is different from existing domains, i.e.,
[-P X 
 = P k,X for k 2 [ K ] .-][Equation]
Our goal is to understand the conditions under which ˆ h can have similar calibration on OOD domains as the InD domains.
For example, if the OOD domain is similar to the mixture distribution of InD domains, we would expect ˆ h performs similarly on InD and OOD domains.
Previous work [Krueger et al., 2021] consider a similar multiple training domains setting in the context of out-of-distribution generalization, where they show that minimizing the differences in training risks (w.r.t.
different domains) can lead to good OOD performance.
To quantify the distance between two distributions, we ﬁrst introduce the H -divergence [Ben-David et al., 2010] to measure the distance between two distributions:
Deﬁnition 5.1 ( H -divergence) .
Given an input space X and two probability distributions P X and P {+
+} X on X , let H be a hypothesis class on X , and denote by A the collection of subsets of X which are the support of hypothesis h [Equation]
H , i.e., A H = { h  1 (1) | h 2 H } .
The distance between P X and P 0 X is deﬁned as
[Equation]
The H -divergence reduces to the standard total variation (TV) distance when H contains all measurable functions on X , which implies that the H -divergence is upper bounded by the TV-distance, i.e., d H ( P X , P 0 X )  d TV ( P X , P 0 X ) .
On the other hand, when the hypothesis class H has a ﬁnite VC dimension or pseudo-dimension, the H -divergence can be estimated using ﬁnite samples from P X and P 0 X [Ben-David et al., 2010].
Next, we deﬁne the mixture distribution of the K in-distribution domains P ↵ K,X on input space X as follows:
[Equation]
where
K
[Equation]
and ↵
[Equation]
Given multiple domains { P 1 , . . .
, P K } , we can optimize the combination parameters ↵ such that
P ↵ K,X minimizes the H -divergence between P ↵ K,X and e
P X .
More speciﬁcally, we deﬁne ˆ ↵ as
[Equation]
where h ? :
= argmin h 2 H { " ( h, P ↵ K,X )+ " ( h, e
P X ) } and ¯ H is deﬁned as ¯ H :
= { sign( | h ( x )  h 0 ( x ) |  t ) | h, h 0 2 H , 0  t  1 } .
We now give an upper bound on the risk on the unseen OOD domain.
This result follows very closely those of Blitzer et al.
[2007]; Zhao et al.
[2018], instantiated in our calibration setup.
Details can be found in Appendix D.
Theorem 5.2.
Let H be a hypothesis class that contains functions h : X ! [0 , 1] with pseudodimension Pdim( H ) =
d . Let { D k,X } Kk =1 denote the empirical distributions generated from { P k,X } Kk =1 , where D k,X contains n i.i.d.
samples from the marginal feature distribution P k,X of domain k .
Then for  2 (0 , 1) , with probability at least 1   , we have
[Equation]
where ˆ ↵ and  ( P ˆ ↵ K,X , e
P X ) are deﬁned in Eq.
(3) , e
P X denotes the marginal distribution of the
OOD domain, Pdim( H ) is the pseudo-dimension of the hypothesis class H , and ˆ " (ˆ h, D k,X ) is the empirical risk of the hypothesis ˆ h on D k,X . Remark 5.3.
As shown in Theorem 5.2, even if the OOD domain is very different from the indistribution domains, the Eq. (4) still implies that we could decrease the risk upper bound on the OOD domain if we perform multi-domain calibration. More speciﬁcally, if we could achieve good calibration performance on each individual domain by using multi-domain calibration (which that the ﬁrst term in the RHS of Eq.
(4) is small), then the term 12 d ¯ H
( P ˆ ↵ K,X , e
P X )
+  ( P ˆ ↵ K,X , e
P X ) is always smaller or equal to 12 d ¯ H ( P 0 , e
P X )
+  ( P 0 , e
P X ) , where P 0 is the pooled distribution or any individual domain distribution.
Remark 5.4.
As suggested by Eq.
(4) of Theorem 5.2, larger risks on in-distribution domains will lead to a larger upper bound for the risk evaluated on the OOD domain.
On the other hand, as shown in Figure 1, a universal temperature is not sufﬁcient to achieve good calibration performance on each individual in-distribution domain.
Therefore, even in the mixture of in-distribution domain setting, a universal temperature is suboptimal and applying multi-domain temperature scaling could be better than using a universal temperature.
This result means that if we can learn a hypothesis ˆ h that achieves small empirical risk ˆ " (ˆ h, D k,X ) on every domain, then ˆ h is able to achieve good performance on the OOD domain if distribution of the OOD domain is similar to the mixture distribution of InD domains measured by H -divergence.
In this case, if the learned calibration map ˆ h is well-calibrated on every domain P k , then ˆ h is likely to provide calibrated conﬁdence for the OOD domain e
P .
Recall from Section 4, we proposed an algorithm that performs well across InD domains.
The upper bound in Eq.
(4) provides insight into understanding why this algorithm is effective.
Discussion
We have developed an algorithm for robust calibration that exploits multi-domain structure in datasets.
Experiments on real-world domains indicate that multi-domain calibration is an effective way to improve the robustness of calibration under distribution shifts.
Our proposed algorithm still needs validation domains to achieve strong calibration performance on OOD domains, one interesting direction for future work would be to extend our algorithm to a scenario where no domain information is available.
We hope the multi-domain calibration perspective in this paper can motivate further work to close the gap between in-distribution and out-of-distribution calibration. Acknowledgments and Disclosure of Funding
We would like to thank the anonymous reviewers for their constructive suggestions and comments.
Yaodong Yu acknowledges support from the joint Simons Foundation-NSF DMS grant #2031899.
Stephen Bates was supported by the Foundations of Data Science Institute and the Simons Institute.
Yi Ma acknowledges support from ONR grants N00014-20-1-2002 and N00014-22-1-2102 and the joint Simons Foundation-NSF DMS grant #2031899.
Michael Jordan acknowledges support from the
ONR Mathematical Data Science program.