Maximum Likelihood With Bias-Corrected Calibration is Hard-To-Beat at Label Shift Adaptation 08 Jul 2020

Accepted to ICML 2020
Authors: Amr Alexandari*, Anshul Kundaje†, Avanti Shrikumar*†
(*co-first authors, †co-corresponding authors)

Introduction

Imagine we train a classifier to predict whether or not a person has a disease based on observed symptoms, and the classifier predicts reliably when deployed in the clinic. Suppose that there is a sudden surge in cases of the disease. During such an outbreak, the probability of persons having the disease given that they show symptoms rises, but the symptoms generated by the disease do not change. How can we adapt the classifier to cope with the difference in the baseline prevalence of the disease?

The problem of distribution shift in its most general form is intractable. However, there are reasonable assumptions that can help us solve this problem. If the disease generates similar symptoms in patients regardless of its spread in the community, then this is known as the label shift or prior probability shift (Amos, 2008), and it corresponds to anti-causal learning (i.e. predicting the cause \(y\) from its effects \(\boldsymbol{x}\)). Anti-causal learning is appropriate for diagnosing diseases given observations of symptoms because diseases cause symptoms. In this example, a pandemic would induce a shift in \(p(y)\); there would be a surge in cases. However, the process generating \(\boldsymbol{x}\) given \(y\) is fixed, i.e. \(p_s(\boldsymbol{x} \mid y) = p_t(\boldsymbol{x} \mid y)\); the symptoms generated by the disease do not change. An automated diagnostic system should cope with the difference in the baseline prevalence of the disease and adjust its predictions accordingly. In our paper, we propose using a maximum likelihood procedure with appropriate calibration for doing that, and we show that our approach is hard to beat at adapting to label shifts.

Bias Corrected Temperature Scaling

Calibration has a long history in the machine learning literature (DeGroot and Fienberg, 1983; Platt, 1999; Zadrozny and Elkan; 2002; Niculescu-Mizil and Caruana, 2005; Kuleshov and Liang, 2015; Naeini et al., 2015; Kuleshov and Ermon, 2016). In the context of modern neural networks, Guo et al. (2017) showed that Temperature Scaling (TS), a single-parameter variant of Platt Scaling (Platt, 1999), was effective at reducing miscalibration. Let \(z(\boldsymbol{x_k})\) be a function that returns the original logit vector. With temperature scaling, we have:

\[p(y_i\mid \boldsymbol{x_k})=\frac{e^{z(\boldsymbol{x_k})_i / T}}{\sum_j e^{z(\boldsymbol{x_k})_j /T}}\]

The parameter \(T\) is optimized with respect to the Negative Log Likelihood on a held-out portion of the training set, such as the validation set.

Guo et al. (2017) compared TS to an approach defined as Vector Scaling (VS), where a different scaling parameter was used for each class along with class-specfic bias parameters, namely:

\[p(y_i\mid \boldsymbol{x_k})=\frac{e^{z(\boldsymbol{x_k})_i W_i+b_i}}{\sum_j e^{z(\boldsymbol{x_k})_j W_j+b_j}}\]

The authors found that vector scaling had a tendency to perform slightly worse than TS as measured by a metric known as the Expected Calibration Error (Naeini et al., 2015).

As shown in the figure, we often found that TS alone resulted in systematically biased estimates of \(p(y_i \mid \boldsymbol{x_k} )\), while VS, a generalization of TS that contains both class-specific bias terms and class-specific scaling terms, did not exhibit as much systematic bias. Intrigued by this observation, we investigated the performance of two intermediaries between Temperature Scaling and Vector Scaling. The first, which we refer to as No Bias Vector Scaling (NBVS), is equivalent to vector scaling but with all the class-specific bias parameters fixed at zero. The second, which we refer to as Bias-Corrected Temperature Scaling (BCTS), is equivalent TS Scaling but with the addition of the class-specific bias terms from VS. BCTS is defined as follows:

\[p(y_i\mid \boldsymbol{x_k})=\frac{e^{z(\boldsymbol{x_k})_i / T +b_i}}{\sum_j e^{z(\boldsymbol{x_k})_j / T +b_j}}\]

As with TS and VS, the parameters are optimized to minimize the Negative Log Likelihood (NLL) on the validation set. Note that in the case of binary classification, the parameterization of BCTS reduces to Platt Scaling (Platt, 1999). Thus, BCTS can be viewed as a multi-class generalization of Platt scaling.

Another Look at Maximum Likelihood

Saerens et al. 2002 proposed an Expectation Maximization (EM) algorithm that estimates \(p_t(y)\) but assumes access to a classifier that outputs the true source distribution conditional probabilities \(p_s(y \mid \boldsymbol{x})\). The algorithm has the following update steps:

\[p_t(y_i)^{(0)}=p_s(y_i)\] \[p_t(y_i \mid \boldsymbol{x_k} )^{(r)}=\frac{\frac{p_t(y_i)^{(r)}}{p_s(y_i)}p_s(y_i \mid \boldsymbol{x_k} )}{\sum_{j=1}^{n}\frac{p_t(y_j)^{(r)}}{p_s(y_j)}p_s(y_j \mid \boldsymbol{x_k} )}\] \[p_t(y_i)^{(r+1)}=\frac{1}{N}\sum_{k=1}^{N} p_t(y_i \mid \boldsymbol{x_k} )^{(r)}\]

where \(p_s(y_i)\) is our estimate of the prior probability of observing class \(i\) on the training set, \(p_t(y_i)^{(r)}\) is the estimate in EM step \(r\) of the prior probability of observing class \(i\) on the testing set, \(p_s(y_i \mid \boldsymbol{x_k} )\) is the conditional probability of observing class \(i\) given features \(\boldsymbol{x_k}\) on the training set, \(p_t(y_i \mid \boldsymbol{x_k} )^{(r)}\) is the conditional probability in EM step \(r\) of observing class \(i\) given features \(\boldsymbol{x_k}\) on the testing set, and \(N\) is the number of examples in the testing set.

Since there is no need to estimate \(p(\boldsymbol{x} \mid y)\) in any step of the EM procedure, the algorithm can scale to high-dimensional datasets. Unfortunately, estimates of \(p(y \mid \boldsymbol{x})\) derived from modern neural networks are often poorly calibrated (Guo et al., 2017), and the lack of calibration can decrease the effectiveness of EM. For this reason, comparisons against the EM algorithm have been absent in the label shift adaptation literature.

Recently, Black Box Shift Estimation (BBSE) (Lipton et al., 2018) and a variant called Regularized Learning Label Shift (RLLS) (Azizzadenesheli et al., 2019): leverage (possibly uncalibrated) predictions off-the-shelf classifiers to estimate the shift. Both of these moment-matching estimators require model retraining with importance weights which can be challenging at large scales.

In our paper, we revisit maximum likelihood. We show that in combination with good calibration, a maximum likelihood procedure outperforms all other methods empirically and achieves state of the art results. Let’s examine the maximum likelihood objective. Let \(\omega_i\) denote membership in class \(i\). We seek target-domain priors \(p_t(\omega_i)\) that maximize the log-likelihood:

\[l(\boldsymbol{X}; p_t(\omega_i)) = 􏰀\sum_{k} log \sum_{i} p_t(\boldsymbol{x_k}\mid\omega_i) q(\omega_i)\]

which we show is equivalent to

\[l(\boldsymbol{X}; p_t(\omega_i)) = 􏰀\sum_{k} log(p_s(\boldsymbol{x_k})) + log \sum_{i}\frac{p_s(\omega_i\mid\boldsymbol{x_k})}{p_s(\omega_i)}p_t{\omega_i}\]

The paper contains proof that the maximum likelihood objective above is concave, so EM converges to the global maximum. Furthermore, any optimization algorithm could be used, so we are not restricted to EM.

Our approach is as follows:

• Given a model that outputs predicted probabilities, calibrate the predictions of the model on a held-out validation using an appropriately strong calibration algorithm. We observed that BCTS and VS work well.
• Average the calibrated predictions over this held-out validation set to obtain the estimated source-domain class priors. This is a principled strategy for estimating source-domain priors improves robustness to poor calibration.
• Given samples from the target domain, use the estimated source-domain class priors and the calibrated predictor to optimize the concave maximum likelihood objective mentioned above w.r.t. the estimated target-domain class proportions.
• After estimating target-domain class proportions, compute the adapted predictions for the target domain similar to the E-step of EM.

In our experiments, the proposed approach results in state of the art performance.

This work demonstrates that the maximum likelihood with appropriate calibration is a formidable and efficient baseline for label shift adaptation. Follow up work by Garg et al. that studies why well-calibrated maximum likelihood is effective indepedently verified our findings: “Across all shifts, MLLS (with BCTS-calibrated classifiers) uniformly dominates BBSE, RLLS, …”

Summary & Additional Resources

In summary:

• The real world is dynamic and real world data distributions are capricious. Enabling automated systems to be robust to the ever-changing shifts present in real world data remains a critical goal of machine learning systems, and is known as "domain adaptation".
• Label shift adaptation arises in settings such as medical diagnosis, and deals with shifts in the distribution of labels while the process generating the properties and manifestations of the label is fixed.
• The popular calibration approach of Temperature Scaling (TS) does not tend to achieve the best results for adaptation to label shift, possibly owing to large systematic biases in the calibrated probabilities. Best results are obtained with variants of TS containing class-specific bias parameters that can correct for systematic bias.
• Maximum likelihood algorithms, such as EM, achieve state-of-the-art results when used with an appropriate type of calibration. Although moment matching algorithms benefit from calibration, they did not tend to outperform maximum likelihood when the probabilities were well-calibrated.

A number of useful resources:

• Paper: https://arxiv.org/abs/1901.06852
• Notebooks reproducing experiments: https://github.com/kundajelab/labelshiftexperiments
• Follow up work by Garg et al. that establishes sufficient conditions for consistency of the maximum likelihood approach and derives finite-sample error bounds: https://arxiv.org/abs/2003.07554