SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Charles Marx*
ctmaxx@cs.stanford.edu
Sofian Zalouk*
szalouk@stanford.edu
Stefano Ermon
ermon@cs.stanford.edu
Calibration ensures that probabilistic forecasts meaningfully capture uncertainty by requiring that predicted probabilities align with empirical frequencies. However, many existing calibration methods are specialized for post-hoc recalibration, which can worsen the sharpness of forecasts. Drawing on the insight that calibration can be viewed as a distribution matching task, we introduce kernel-based calibration metrics that unify and generalize popular forms of calibration for both classification and regression. These metrics admit differentiable sample estimates, making it easy to incorporate a calibration objective into empirical risk minimization. Furthermore, we provide intuitive mechanisms to tailor calibration metrics to a decision task, and enforce accurate loss estimation and no regret decisions. Our empirical evaluation demonstrates that employing these metrics as regularizers enhances calibration, sharpness, and decision-making across a range of regression and classification tasks, outperforming methods relying solely on post-hoc recalibration.
Probabilistic forecasts are valuable tools for capturing uncertainty about an outcome of interest. In practice, the exact outcome distribution is often impossible to recover [51], leading to overconfident forecasts [15]. Calibration provides statistical guarantees on the forecasts, ensuring that the predicted probabilities align with the true likelihood of events. For example, consider a weather forecast that assigns an  $80\%$  probability of rain tomorrow. A well-calibrated forecast would imply that, on average, it rains on  $80\%$  of the days with such predictions. By appropriately quantifying uncertainty, calibration empowers decision-makers to efficiently allocate resources and mitigate risks, even when the outcome distribution cannot be recovered exactly.
Many forms of calibration have been proposed in both classification [8, 21, 35] and regression [42, 23, 53] to ensure that uncertainty estimates match true frequencies. From a methodological perspective, the literature on enforcing calibration is fractured; many algorithms to enforce calibration are specialized for a particular form of calibration [39, 29, 52, 10, 17] or are designed exclusively for post-hoc calibration [27, 23]. Although post-hoc methods are effective at enforcing calibration, they often lead to a degradation in the predictive power, i.e. sharpness, of the forecaster [26]. This motivates the need for versatile metrics that can enforce various forms of calibration during training.
In this work, we introduce a unified framework wherein existing notions of calibration are presented as distribution matching constraints. To enforce distribution matching, we use the Maximum Mean Discrepancy (MMD), giving us unbiased estimates of miscalibration that are amenable to gradient-based optimization. We frame these metrics as regularizers and optimize them alongside proper scoring rules. This allows us to enforce calibration while preserving sharpness.
Framing calibration as distribution matching [38] allows us to use the Maximum Mean Discrepancy (MMD) [14] to give unbiased estimates of miscalibration that are amenable to gradient-based optimization. An advantage of this approach is that we express existing calibration measures by varying the kernel in the MMD metric. Moreover, we can express new notions of calibration specific to a downstream decision task. This allows us to prioritize errors that impact decision-making [52]. Our contributions can be summarized as follows:
- We propose a framework based on distribution matching to unify many notions of calibration.
- We provide differentiable training objectives to incentivize a wide variety of popular forms of calibration during empirical risk minimization in supervised learning. We frame these metrics as regularizers that can be optimized alongside proper scoring rules. This allows us to enforce calibration without degrading sharpness.
- We show how to define and enforce new calibration metrics tailored to specific downstream decision problems. When these metrics are successfully optimized, decision makers can accurately estimate the decision loss on unlabeled data.
- We perform an empirical study across 10 tabular prediction tasks (5 regression and 5 classification) and find that using our metrics as regularizers improves calibration, sharpness, and decision-making relative to existing recalibration methods.
**Calibration** Our work builds on the literature on calibrated forecasting [31, 8, 9, 13]. Recently, many forms of calibration have been proposed in classification [5, 25, 24, 32, 35, 33, 10, 17, 28, 36, 27], regression [23, 42, 53, 39, 51], and beyond [47, 22]. Algorithms to enforce calibration are often designed specifically for one form of calibration. For example, significant attention has been devoted to effectively estimating the Expected Calibration Error (ECE) to calibrate confidence levels in binary classification [19, 44, 4, 15, 35]. Our work aims to supplement this work by providing general calibration metrics that can be used as training objectives. These methods can be used in concert with post-hoc calibration, to encourage calibration during model training and correct any remaining miscalibration after training is complete.
Decision Calibration An interesting line of recent work focuses on calibrating predictions for downstream decision problems [39, 52]. Zhao et al. [52] shows that decision calibration can be achieved when the set of possible actions is finite. We show cases in which this result can be extended to infinite action spaces by leveraging a low-dimensional sufficient statistic of the decision problem. Whereas Zhao et al. [52] compute calibration by approximating a worst-case decision problem with a linear classifier, we use the kernel mean embedding form of an MMD objective, making it simple to calibrate with respect to a restricted set of loss functions by adjusting the kernel.
Kernel-Based Calibration Metrics Motivated by issues arising from the histogram binning of the Expected Calibration Error, a recent stream of research uses kernels to define smooth calibration metrics. Zhang et al. [50] and Popordanoska et al. [37] use kernel density estimates to estimate calibration in classification settings. In contrast, our work applies to regression and takes a distribution matching perspective. Widmann et al. [47] and Kumar et al. [26] are the most similar existing works to our own, and also introduce MMD-based calibration metrics. Our work differs from Widmann et al. [46] due to our focus on trainability and optimization, connections to existing forms of calibration, and applications to decision-making. Kumar et al. [26] can be seen as a special case of our method applied to top-label calibration in multiclass classification (see Table 1). Our work extends the literature on MMD-based calibration metrics, showing how to express and optimize 11 existing forms of calibration by varying the kernel in the MMD objective (see Tables 1 and 2).
3 Calibration is Distribution Matching
We consider the regression setting, where, given a feature vector  $x \in \mathcal{X} = \mathbb{R}^d$ , we predict a label  $y \in \mathcal{Y} \subset \mathbb{R}$ . We are given access to  $n$  i.i.d. examples  $(x_1, y_1), \ldots, (x_n, y_n)$  distributed according to some true but unknown cumulative distribution function (cdf)  $P$  over  $\mathcal{X} \times \mathcal{Y}$ . A lower case  $p$  denotes the probability density function (pdf) or probability mass function (pmf), when it exists. We use
Table 1: We express popular forms of calibration in terms of distribution matching, where  $P$  is the true distribution and  $Q$  is the forecast. The forecaster achieves calibration if the forecast variable and target variable are equal in distribution, given the conditioning variable. Here,  ${\delta }_{\ell }^{ * }\left( {Q}_{Y \mid  X}\right)$  is a Bayes optimal action under  ${Q}_{Y \mid  X},{y}_{0}$  is a fixed label in  $\mathcal{V},\alpha  \in  \left\lbrack  {0,1}\right\rbrack$  is a threshold,and  $\phi \left( X\right)$  is a feature mapping. See Appendix A for additional details.
Calibration Objective	Forecast Variable	Target Variable	Conditioning Variable
Quantile Calibration [23]	QY|X(Y)	PY|X(Y)	-
Threshold Calibration [39]	QY|X(Y)	PY|X(Y)	1{QY|X(y0)≤α}
Marginal Calibration [12]	Ŷ	Y	-
Decision Calibration [52]	Ŷ	Y	δ* (QY|X)
Group Calibration [36]	Ŷ	Y	Group(X)
Distribution Calibration [42]	Ŷ	Y	QY|X
Individual Calibration [51]	Ŷ	Y	X
Local Calibration [27]	Ŷ	Y	φ(X)
subscripts for marginal and conditional distributions, so  $P_{Y|x}(\cdot)$  is the conditional cdf for  $Y$  given  $X = x$  and  $p_{Y}(\cdot)$  is the marginal pdf for  $Y$ . When the input  $X$  is random, we write  $P_{Y|X}$  as the cdf-valued random variable that takes value  $P_{Y|x}$  upon the event that  $X = x$ .
Our goal is to learn a forecaster  $Q$  that, given a feature vector  $x$ , forecasts a cdf  $Q_{Y|x}(\cdot)$  over  $\mathcal{V}$ . A forecaster implicitly defines a joint distribution over  $\mathcal{X} \times \mathcal{V}$  by combining the forecasted conditional pdf with the marginal pdf for the features, i.e.  $q(x,y) = p_X(x)q_{Y|x}(y)$ . The forecaster's joint cdf  $Q(x,y)$  is defined by integrating the pdf  $q(x,y)$  over the appropriate region of  $\mathcal{X} \times \mathcal{V}$ . In the same way as  $(X,Y) \sim P$  is a sample from the true distribution, we write a sample from  $Q$  as  $(X,\widehat{Y}) \sim Q$ , where  $X \sim P_{X}$  and  $\widehat{Y}$  is a label drawn from the forecast  $(\widehat{Y}|X = x) \sim Q_{Y|x}$ .
An ideal forecaster perfectly recovers the true probabilities so that  $Q = P$ . However, in practice we often only observe each feature vector  $x$  once, making perfect forecasts unattainable [51]. Instead, calibration enforces that forecasts are accurate on average, leaving us with a tractable task that is strictly weaker than perfect forecasting. Each form of calibration requires the forecasts to respect a particular law of probability. For example, if  $Y$  is binary, distribution calibration states that on examples for which the forecast is  $q_{Y|x}(Y = 1) = c$ , we should observe that  $Y = 1$  with frequency  $c$ . In regression, quantile calibration [23] states that  $y$  should exceed each quantile  $c \in [0,1]$  of the forecast with frequency  $c$ . Calibration can also be enforced within subgroups of the data (i.e. group calibration [36]), in smooth neighborhoods of the input space (i.e. local calibration [27]), and even for a single feature vector (i.e. individual calibration [51]) by using randomization.
Calibration can be expressed as a conditional distribution matching problem, requiring that certain variables associated with the true distribution match their counterparts from the forecasts.
Lemma 3.1 (informal). Quantile calibration is equivalent to distribution matching between  $Q_{Y|X}(Y)$  and  $P_{Y|X}(Y)$ . Distribution calibration is equivalent to distribution matching between  $\widehat{Y}$  and  $Y$  given  $Q_{Y|X}$ . Individual calibration is equivalent to distribution matching between  $\widehat{Y}$  and  $Y$  given  $X$ .
The choice of random variables that we match in distribution determines the form of calibration recovered. We show how to express many popular forms of calibration in terms of distribution matching in Table 1, and include additional details on these correspondences in Appendix A. In Section 4, we describe how to express and estimate these forms of calibration using kernel-based calibration metrics. In Section 5, we show how to tailor these metrics to downstream decision problems. Finally, in Section 6, we apply these metrics as regularizers and explore the empirical effects on calibration, sharpness, and decision-making.
4 Calibration Metrics as Trainable Regularizers
In this section, we introduce a framework to enforce calibration using distribution matching. We quantify calibration using integral probability metrics which admit unbiased estimates from data
and are amenable to gradient-based optimization. This allows us to optimize calibration alongside a standard training objective, achieving calibration without degrading sharpness. Jointly optimizing calibration and sharpness has been shown to outperform only enforcing calibration post-hoc [26].
Integral probability metrics (IPMs) are a flexible set of metrics to quantify the difference between two distributions [30] in terms of a family of witness functions. Since we need to perform distribution matching conditionally, we consider witness functions over both the label  $Y$ , as well as a conditioning variable  $Z$  that assumes the same distribution under  $P$  and  $Q$ .
$$
D (\mathcal {F}, P, Q) = \sup  _ {f \in \mathcal {F}} \left| \mathbb {E} _ {P} [ f (Y, Z) ] - \mathbb {E} _ {Q} [ f (\widehat {Y}, Z) ] \right|. \tag {1}
$$
Different choices for the witness functions  $\mathcal{F}$  and the conditioning variable  $Z$  measure different notions of calibration. Similar to Widmann et al. [47], we focus on the case where  $\mathcal{F}$  is the unit ball of a Reproducing Kernel Hilbert Space (RKHS)  $\mathcal{H}$ , so that the IPM corresponds to the Maximum Mean Discrepancy [14], denoted  $\mathrm{MMD}(\mathcal{F}, P, Q)$ . MMD can either be expressed in terms of the canonical feature map  $\phi(y, z) \in \mathcal{H}$  or the reproducing kernel  $k((y, z), (\widehat{y}, z')) = \langle \phi(y, z), \phi(\widehat{y}, z') \rangle_{\mathcal{H}}$  for  $\mathcal{H}$ :
$$
\begin{array}{l} \operatorname {M M D} ^ {2} (\mathcal {F}, P, Q) = \left\| \mathbb {E} _ {P} [ \phi (Y, Z) ] - \mathbb {E} _ {Q} [ \phi (\widehat {Y}, Z) ] \right\| _ {\mathcal {H}} ^ {2} \tag {2} \\ = \mathbb {E} _ {P} \left[ k \left(\left(Y, Z\right), \left(Y ^ {\prime}, Z ^ {\prime}\right)\right) \right] + \mathbb {E} _ {Q} \left[ k \left(\left(\widehat {Y}, Z\right), \left(\widehat {Y} ^ {\prime}, Z ^ {\prime}\right)\right) \right] - 2 \mathbb {E} _ {P} \mathbb {E} _ {Q} \left[ k \left(\left(Y, Z\right), \left(\widehat {Y} ^ {\prime}, Z ^ {\prime}\right)\right) \right] \\ \end{array}
$$
Given a dataset  $D$  consisting of  $n$  i.i.d. pairs  $(x_{i},y_{i})\sim P$ , we can generate forecasted labels  $\widehat{y}_i\sim Q_{Y|x_i}$ . The conditioning variables  $z_{i}$  are computed from  $(x_{i},y_{i})$  according to the expressions given in Table 1. The plug-in MMD estimator is unbiased and takes the form [14]:
$$
\widehat {\mathrm {M M D}} ^ {2} (\mathcal {F}, D, Q) = \frac {1}{n (n - 1)} \sum_ {i = 1} ^ {n} \sum_ {j \neq i} ^ {n} h _ {i j}, \tag {3}
$$
where  $h_{ij}$  is a one-sample U statistic:
$$
h _ {i j} := k \left(\left(y _ {i}, z _ {i}\right), \left(y _ {j}, z _ {j}\right)\right) + k \left(\left(\widehat {y} _ {i}, z _ {i}\right), \left(\widehat {y} _ {j}, z _ {j}\right)\right) - k \left(\left(y _ {i}, z _ {i}\right), \left(\widehat {y} _ {j}, z _ {j}\right) - k \left(\left(y _ {j}, z _ {j}\right), \left(\widehat {y} _ {i}, z _ {i}\right)\right) \right.
$$
The variance of this estimator can be reduced by marginalizing out the simulated randomness used to sample  $\widehat{y}_i\sim Q_{Y|x_i}$ . This can be done either analytically, or empirically by resampling the forecasted labels multiple times per example. We find that empirically marginalizing out the simulated randomness improves training stability in practice. The full training objective combines a proper scoring rule (we use negative log-likelihood) with an MMD estimate to incentivize calibration. To train, we divide the dataset  $D$  into batches  $D_{b}\subset \{1,\dots ,n\}$  large enough so that the MMD estimate on a batch has low enough variance to provide useful supervision (we find that  $D_{b}$  of size 64 suffices in practice). Given a family of predictive models  $\{Q_{\theta}:\theta \in \Theta \}$  (e.g., a neural network), we optimize the following training objective on each batch:
$$
\min  _ {\theta \in \Theta} \sum_ {i \in D _ {b}} - \log q _ {Y | x _ {i}; \theta} (y _ {i}) + \lambda \cdot \widehat {\mathrm {M M D}} ^ {2} (\mathcal {F}, D _ {b}, Q _ {\theta}) \tag {4}
$$
In order to differentiate  $\widehat{\mathrm{MMD}}^2 (\mathcal{F},D_b,Q_\theta)$  with respect to  $\theta$ , we require that the kernel  $k$  is differentiable and the samples can be expressed as a differentiable function of the model parameters (e.g., using a reparameterization trick). For parameterizations that do not admit exact differentiable sampling, we can apply a differentiable relaxation (e.g., if  $Q_{\theta}$  is a mixture distribution, we can use a Gumbel-Softmax approximation). In principle, it is possible to remove the need for differentiable sampling altogether by optimizing our MMD objective using Monte-Carlo RL techniques (e.g., REINFORCE [48]). However, this is beyond the scope of our work. Note that regularizing for calibration at training time does not give calibration guarantees, unless one can guarantee the objective will be optimized out-of-sample. In this sense, regularization and post-hoc calibration are complementary: regularization mitigates the trade-off between calibration and sharpness while post-hoc methods can provide finite-sample calibration guarantees (e.g., [45, 29]). For this reason, we suggest using regularization and post-hoc calibration together (see Table 3).
On the Choice of the Kernel When the kernel  $k$  is universal, the MMD metric is zero if and only if  $Y$  and  $\widehat{Y}$  are equal in distribution, given  $Z$  [14]. Universality is preserved when we decompose the kernel over  $\mathcal{V} \times \mathcal{Z}$  into a pointwise product of two universal kernels, one over  $\mathcal{V}$  and one over  $\mathcal{Z}$  [43].
Table 2: Popular forms of calibration expressed as distribution matching for a classification problem with classes  $\mathcal{Y} = \{1,\dots ,m\}$ . Recall that  $q_{Y|X}(\cdot)$  is the forecasted pmf. The forecaster is calibrated if the forecast variable and target variable are equal in distribution, given the conditioning variable. For top-label calibration, we use  $Y^{*}:= \arg \max_{y\in \mathcal{Y}} q_{Y|X}(y)$ , which is random due to  $X$ . The marginal calibration condition expresses  $m$  distribution matching constraints, one for each  $y\in \mathcal{Y}$ .
Calibration Objective	Forecast Variable	Target Variable	Conditioning Variable
Canonical Calibration [44, Eq 1]	\(\widehat{Y}\)	Y	\(q_{Y|X}\)
Top-label Calibration [44, Eq 2]	\(\mathbb{1}\{\widehat{Y}=Y^{*}\}\)	\(\mathbb{1}\{Y=Y^{*}\}\)	\(q_{Y|X}(Y^{*})\)
Marginal Calibration [44, Eq 3]	\(\mathbb{1}\{\widehat{Y}=y\}\)	\(\mathbb{1}\{Y=y\}\)	\(q_{Y|X}(y), \forall y\in\mathcal{V}\)
Lemma 4.1. Let  $k$  be the function defined by the pointwise product  $k((y,z),(y',z')) = k_y(y,y')k_z(z,z')$  where  $k_y$  and  $k_z$  are universal kernels over spaces  $\mathcal{V}$  and  $\mathcal{Z}$  respectively. Then, under mild conditions,  $k$  is a universal kernel over  $\mathcal{V} \times \mathcal{Z}$ .
See [43, Theorem 4] for a proof. In our experiments, we explore using a universal RBF kernel over  $\mathcal{V} \times \mathcal{Z}$ , since it satisfies the mild conditions needed for Lemma 4.1. To enforce decision calibration, we also consider a kernel  $k_{Y}$  that is intentionally not universal, and instead only measures differences between  $P$  and  $Q$  that are relevant for decision-making (see Section 5). Furthermore, to measure Quantile Calibration and Threshold Calibration, we must match the distributions of the probability integral transforms  $P_{Y|X}(Y)$  and  $Q_{Y|X}(Y)$  instead of  $Y$  and  $\widehat{Y}$ . Thus, we replace  $k_{Y}$  with a universal kernel over  $[0,1]$ , the codomain of the probability integral transform.
4.1 Calibration Metrics for Classification
Thus far, we have focused on estimating calibration metrics in the regression setting. We now describe how our framework applies to classification.
In the classification setting, given a feature vector  $x \in \mathcal{X} = \mathbb{R}^d$  we predict a label  $y \in \mathcal{Y} = \{1, \dots, m\}$ . As before, we are given access to  $n$  i.i.d. examples  $(x_1, y_1), \ldots, (x_n, y_n)$  distributed according to some unknown distribution  $P$  over  $\mathcal{X} \times \mathcal{Y}$ . Our goal is to learn a forecaster  $Q$  that, given a feature vector  $x$ , forecasts a probability mass function (pmf)  $q_{Y|x}(\cdot)$  over labels  $\mathcal{Y}$ .
We express popular notions of calibration used in classification (i.e., canonical, top-label, marginal) in terms of distribution matching in Table 2. See Appendix A.2 for additional details on calibration and distribution-matching in classification. We use the same unbiased plug-in estimator for MMD shown in Equation 3. In the classification setting, we can analytically marginalize out the randomness originating from the forecast when computing the MMD estimate. After performing this marginalization, the one-sample U-statistic  $h_{ij}$  is given by
$$
h _ {i j} = k \left(\left(y _ {i}, z _ {i}\right), \left(y _ {j}, z _ {j}\right)\right) + \sum_ {y \in \mathcal {Y}} \sum_ {y ^ {\prime} \in \mathcal {Y}} q _ {i} (y) q _ {j} \left(y ^ {\prime}\right) k \left(\left(y, z _ {i}\right), \left(y ^ {\prime}, z _ {j}\right)\right) - 2 \sum_ {y \in \mathcal {Y}} q _ {i} (y) k \left(\left(y, z _ {i}\right), \left(y _ {j}, z _ {j}\right)\right),
$$
where the conditioning variables  $z_{i}$  are computed from  $(x_{i},y_{i})$ , as detailed in Table 2. Here,  $q_{i}(y) = q_{Y|x_{i}}(y)$  is the predicted probability of the label  $Y = y$  given features  $x_{i}$ . When the forecast and target variables depend on variables other than  $Y$  (e.g.,  $Y^{*}$  depends on  $q_{Y|X}$  in top-label calibration), we add those variables as inputs to the kernel. Note that the plug-in MMD estimator with this U-statistic is differentiable in terms of the predicted label probabilities, and that it can be computed in  $O(n^{2}m^{2})$  time.
5 Calibration and Decision-Making
Some of the most practically useful results of calibration are the guarantees provided for downstream decision-making [52]. Universal kernels guarantee distribution matching in principle, but in some cases require many samples. However, if we have information about the decision problem in advance, we can design calibration metrics that are only sensitive to errors relevant to decision-making. In this section, we show how to design such metrics and show that they enable us to (1) accurately evaluate the loss associated with downstream actions and (2) choose actions that are in a sense optimal.
We begin by introducing a general decision problem. For each example, a decision-maker chooses an action  $a \in \mathcal{A}$  and incurs a loss  $\ell(a, y)$  that depends on the outcome. Specifically, given information  $z$  (such as the forecast), a decision policy  $\delta: \mathcal{Z} \to \mathcal{A}$  assigns an action  $\delta(z)$ . The set of all such policies is denoted  $\Delta(\mathcal{Z})$ . The Bayes optimal policy under a label distribution  $P_{Y}$  is denoted  $\delta_{\ell}^{*}(P_{Y}) := \arg \min_{a \in \mathcal{A}} \mathbb{E}_{P}[\ell(a, Y)]$ . Decision calibration [52] provides two important guarantees:
$$
\mathbb {E} _ {Q} [ \ell (a, \widehat {Y}) \mid Z = z ] = \mathbb {E} _ {P} [ \ell (a, Y) \mid Z = z ], \quad \forall a \in \mathcal {A}, z \in \mathcal {Z} \quad (\text {l o s s e s t i m a t i o n})
$$
$$
\mathbb {E} _ {P} \left[ \ell \left(\delta_ {\ell} ^ {*} \left(Q _ {Y | x}\right), Y\right) \mid Z = z \right] \leq \mathbb {E} _ {P} \left[ \ell (\delta (z), Y) \mid Z = z \right], \quad \forall z \in \mathcal {Z}, \delta \in \Delta (\mathcal {Z}) \quad (\text {n o r e g r e t})
$$