SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Qingyao Sun*
Cornell University
qs234@cornell.edu
Sayna Ebrahimi  
Google Cloud AI Research  
saynae@google.com
Kevin Murphy Google DeepMind kpmurphy@google.com
Alexander D'Amour  
Google DeepMind  
alexdamour@google.com
Changes in the data distribution at test time can have deleterious effects on the performance of predictive models  $p(y|x)$ . We consider situations where there are additional meta-data labels (such as group labels), denoted by  $z$ , that can account for such changes in the distribution. In particular, we assume that the prior distribution  $p(y,z)$ , which models the dependence between the class label  $y$  and the "nuisance" factors  $z$ , may change across domains, either due to a change in the correlation between these terms, or a change in one of their marginals. However, we assume that the generative model for features  $p(x|y,z)$  is invariant across domains. We note that this corresponds to an expanded version of the widely used "label shift" assumption, where the labels now also include the nuisance factors  $z$ . Based on this observation, we propose a test-time label shift correction that adapts to changes in the joint distribution  $p(y,z)$  using EM applied to unlabeled samples from the target domain distribution,  $p_t(x)$ . Importantly, we are able to avoid fitting a generative model  $p(x|y,z)$ , and merely need to reweight the outputs of a discriminative model  $p_s(y,z|x)$  trained on the source distribution. We evaluate our method, which we call "Test-Time Label-Shift Adaptation" (TTLSA), on several standard image and text datasets, as well as the CheXpert chest X-ray dataset, and show that it improves performance over methods that target invariance to changes in the distribution, as well as baseline empirical risk minimization methods. Code for reproducing experiments is available at https://github.com/nalzok/test-time-label-shift.
Machine learning systems are known to be sensitive to so-called "spurious correlations" [Geirhos et al., 2020, Wiles et al., 2022, Arjovsky et al., 2019] between irrelevant features of the inputs and the predicted output label. These features and their associated correlations are called "spurious" because they are expected to change across real-world distribution shifts. As a result, models that rely on such spurious correlations often have worse performance when they are deployed on a target domain that is distinct from the source domain on which they were trained [Geirhos et al., 2020, Izmailov et al., 2022]. For example, Jabbour et al. [2020] show that neural networks trained to recognize conditions like pneumonia from chest X-rays can learn to rely on features that are predictive of patient demographics rather than the medical condition itself. When the correlation between demographic
*Work done as a master's student at the University of Chicago.
factors and the condition change (e.g., when the model is deployed on a different patient population), the performance of such models suffers.
To address this issue, recent work has focused on learning predictors that are invariant to changes in spurious correlations across source and target domains, either using data from multiple environments [e.g., Arjovsky et al., 2019], or data where the labels for the nuisance factors are available at training time [Veitch et al., 2021, Makar et al., 2022, Puli et al., 2022, Makino et al., 2022].
However, "spurious" correlations can sometimes provide valuable prior information for examples where the input is ambiguous. Consider the example of calculating the probability that a patient has a particular disease given a positive test. It is well known that the underlying prevalence of the disease (i.e., prior probability it is present) in the patient population is highly informative for making this diagnosis; even if the test is highly sensitive and specific, the patient's probability of having the disease given a positive test may be low if the disease is rare given that patient's background [see, e.g., Bours, 2021]. A similar logic applies in many prediction problems: if we can predict, say, patient demographics from the input, and the prevalence of the target label differs between demographic groups in the current patient population, then the demographic-relevant spurious features provide prior information that supplements the information in the target-relevant features of the input. Thus, if a model could be adjusted to use spurious features optimally in downstream target domains, it could substantially out-perform invariant predictors, as we show in this paper.
Motivated by the above, in this paper we propose a test-time approach for optimally adapting to distribution shifts which arise due to changes in the underlying joint prior between the class labels  $y$  and the nuisance labels  $z$ . We can view these changes as due to a hidden common cause  $u$ , such as the location of a specific hospital. Thus we assume  $p_s(u) \neq p_t(u)$ , where  $p_s$  is the source distribution, and  $p_t$  is the target distribution. Consequently,  $p_i(y,z) = \sum_u p(y,z|u)p_i(u)$  will change across domains  $i$ . However, we assume that the generative model of the features is invariant across domains, so  $p_i(\boldsymbol{x} \mid y,z) = p(\boldsymbol{x} \mid y,z)$ . See Figure 1 for an illustration of our modeling assumptions.
The key observation behind our method is that our assumptions are equivalent to the standard "label shift assumption", except it is defined with respect to an expanded label  $m = (y,z)$ , which we call the meta-label. We call this the "expanded label shift assumption". This lets us use existing label shift techniques, such as Alexandari et al. [2020], Lipton et al. [2018], Garg et al. [2020], to adapt our model using a small sample of unlabeled data  $\{\pmb{x}_n\sim p_t(\pmb{x})\}$  from the target domain to adjust for the shift in the prior over meta-
labels, as we discuss in Section 3.2. Importantly, although our approach relies on the assumption that  $p(\boldsymbol{x} \mid y, z)$  is preserved across distribution shifts, it is based on learning a discriminative base model  $p_s(y, z, | \boldsymbol{x})$ , which we adjust to the target distribution  $p_t(y \mid \boldsymbol{x})$ , as we explain in Section 3.1. Thus we do not need to fit a generative model to the data. We do need access to labeled examples of the confounding factor  $z$  at training time, but such data is often collected anyway (albeit in limited quantities) especially for protected attributes. Additionally, because it operates at test-time, our method does not require retraining to adapt the base model to multiple target domains. We therefore call our approach Test-Time Label Shift Adaptation (TTLSA).
We evaluate TTLSA on various standard image and text classification benchmarks, as well as the CheXpert chest X-ray dataset. We show that in shifted target domains TTLSA outperforms ERM (empirical risk minimization, which often uses spurious features inappropriately in target domains), as well as methods that train invariant classifiers (which ignore spurious features). Note, however,
Figure 1: Modeling assumptions.  $U$  is a hidden confounder that induces a spurious correlation between the label  $Y$  and other causal factors  $Z$ , which together generate the features  $X$ . The distribution of  $U$  can change between target and source domains, but the distribution of  $X$  given  $Y$ ,  $Z$  is fixed.
2We acknowledge that practitioners may have concerns beyond target domain performance, such as certain notions of equal treatment, in which the reliance on such spurious features is problematic or even forbidden. However, equal treatment criteria are often motivated by a desire to ensure that models perform as well as possible when they are applied in domains with very different population compositions, and our method can achieve this goal [see Makar and D'Amour, 2022, for related discussion].
that it has been shown that no single adaptation method can work under all forms of distribution shift [Veitch et al., 2021, Kaur et al., 2022]. Our assumptions capture certain kinds of shift, but certainly not all. In particular, our method is unlikely to help with the kinds of problems studied in the domain adaptation literature, where there is covariate shift (i.e., a change from  $p_s(\pmb{x})$  to  $p_t(\pmb{x})$ ) which is not captured by a change in the distribution over the causal factors  $(y,z)$ .
Spurious correlations, invariant learning, and worst-group performance Spurious correlations have mostly been studied in the domain generalization literature (see e.g., [Koh et al., 2020, Gulrajani and Lopez-Paz, 2021]), in which a model is expected to generalize (i.e., achieve acceptable performance) in a target domain without access to any data from that target domain. In this problem setup, building predictors based on the principle of invariance or minimax optimality with respect to spurious correlation shifts is a natural approach.
Many of these methods make the same modeling assumptions as we do (shown in Figure 1); this is often referred to as an anti-causal prediction setting, since the labels cause (generate) the features rather than vice versa [Schoelkopf et al., 2012, Veitch et al., 2021, Makar et al., 2022, Puli et al., 2022, Zheng and Makar, 2022]. These methods use the fact that an invariant predictor will satisfy certain conditional independencies, and employ regularizers that encourage the desired conditional independence. A related line of work tries to minimize the worst-case loss across groups (values of  $m = (y,z)$ ) using distributionally robust optimization [see, e.g., Sagawa et al., 2020, Liu et al., 2021, Nam et al., 2022, Lokhande et al., 2022]). Recently, Idrissi et al. [2022] showed empirically that the "SUBG" method, which uses ERM on a group-balanced version of the data achieved by subsampling, can be an effective approach to learning worst-group-robust predictors in this setting.
Unsupervised domain adaptation and label shift In the Unsupervised Domain Adaptation (UDA) literature, we assume that unlabeled data from the target domain is available, either at training or test time. To make optimality guarantees, UDA methods require assumptions about the structure of the distribution shift. Methods are often categorized into whether they require a covariate shift assumption (i.e., that the discriminative distribution  $p(y \mid x)$  is preserved [see, e.g., Shimodaira, 2000]), or a label shift assumption (i.e., that generative distribution  $p(x \mid y)$  is preserved [Lipton et al., 2018, Garg et al., 2020]).
In the setting of Figure 1, neither the standard covariate shift nor label shift assumptions hold; however, the label shift assumption does hold with respect to the meta-label  $m = (y,z)$ . Our primary contribution is to show that label shift adaptation techniques, such as Alexandari et al. [2020], when applied with respect to the meta-label, can also be used to adapt to changes due to spurious correlations.
Interestingly, in the "purely spurious" setting described in Makar et al. [2022], which is a special case of Figure 1, we can obtain a risk-invariant model as a special case, by adapting to a target distribution for which  $p_t(y,z) = p_s(y)p_s(z)$  [Veitch et al., 2021, Makar and D'Amour, 2022]. Based on this observation, we show that logit adjustment [Menon et al., 2021], a set of test- and training-time methods developed for long-tail learning, can be repurposed to do invariant learning when applied to the meta-label  $m = (y,z)$ .
Methods using the expanded label shift assumption The "generative multi-task learning" or GMTL method of Makino et al. [2022] makes the same expanded label shift assumption that we do. However, instead of estimating  $p_t(y,z)$  from unsupervised target data, they instead assume that there is some value  $\alpha$  such that  $p_t(y,z) = p_s(y,z)^{1 - \alpha}$ . They state that choosing  $\alpha$  is an open problem, and therefore they report results for a range of possible  $\alpha$ 's. By contrast, we provide a way to estimate  $p_t(y,z)$ , and we do not restrict it to have the above functional form.
<sup>3</sup>The "purely spurious" term was coined in Veitch et al. [2021], but for brevity we describe using formalism from Makar et al. [2022]. Within the model of Figure 1, an association is "purely spurious" if there is a sufficient statistic of  $X$ ,  $e(X)$ , such that  $Y \perp X \mid e(X)$ ,  $Z$  and  $e(X) \perp Z \mid Y$ . This occurs when the influence of  $Y$  on  $X$  can be localized in some features  $e(X)$  that are not affected by  $Z$ , or intuitively, when the influences of  $Y$  and  $Z$  on  $X$  are disentangleable. The results of Idrissi et al. [2022], where spurious correlations are neutralized by data balancing, suggests that pure spuriousness is a common case, at least among worst-group benchmarks. We discuss pure spuriousness, worst-group performance, and invariance in more detail in the supplement.
In Jiang and Veitch [2022], they propose a method called Anti-Causal Transportable and Invariant Representation or "ACTIR" which also makes the same expanded label shift assumption that we do, and further relaxes the assumption that  $Z$  is observed during source training. However, they require access to examples from multiple source distributions, which they use to learn a domain invariant classifier. Furthermore, then require labeled  $(x,y)$  examples to adapt their classifier at test time.
Test time adaptation There is a growing Test Time Adaptation (TTA) literature that explores strategies for adapting a trained model at test time using unlabeled data from the target domain. These methods are based on various heuristics to adapt discriminative models without specifying strong assumptions on the distribution shift structure. For example, TENT [Wang et al., 2021] uses entropy minimization to update the batch normalization layers of a CNN, and MEMO [Zhang et al., 2021] uses ensembles of predictions for different augmentations of a test sample. Most of the TTA literature has focused on models that work on images (e.g., by leveraging data augmentation), so these techniques cannot be applied to embeddings or other forms of input. By contrast, our method can be applied to any classifier, even non-neural ones, such as random forests. For a more comprehensive review of TTA methods, see [Liang et al., 2023].
3 Test-Time Label Shift Adaptation
Our aim is to construct a Bayes optimal predictor for the target distribution,  $f_{t}(\pmb{x})$ , using a large labeled dataset from the source distribution,  $\mathcal{D}_s^{xyz} = \{(x_n,y_n,z_n) \sim p_s(x,y,z)\}$ , and a small unlabeled dataset from the target distribution,  $\mathcal{D}_t^{\pmb{x}} = \{x_n \sim p_t(\pmb{x})\}$ . The optimal prediction for the class label in the target domain is given by
$$
f _ {t} (\boldsymbol {x}) = \arg \max  _ {y} p _ {t} (y | \boldsymbol {x}) = \arg \max  _ {y} \sum_ {z} p _ {t} (y, z | \boldsymbol {x}). \tag {1}
$$
where the joint posterior over the class label  $y$  and the nuisance factor  $z$  is given by
$$
p _ {t} (y, z \mid \boldsymbol {x}) \propto p _ {t} (\boldsymbol {x} \mid y, z) p _ {t} (y, z) = p _ {s} (\boldsymbol {x} \mid y, z) p _ {t} (y, z) \tag {2}
$$
where the first step follows from Bayes' rule, and the second step follows from our causal stability assumption that  $p_t(\pmb{x} | y, z) = p_s(\pmb{x} | y, z) = p(\pmb{x} | y, z)$ .
The first key insight of our method is that we can use the EM algorithm (or other optimization methods) to estimate  $p_t(y,z)$  by maximizing the likelihood of the unlabeled target dataset, as explained in Section 3.2. The second key insight is that we can estimate the likelihood  $p(\boldsymbol{x}|y,z)$  up to a constant of proportionality by fitting a discriminative model on the source dataset, and then dividing out by the source prior:
$$
p _ {s} (\boldsymbol {x} | y, z) \propto \frac {p _ {s} (y , z | \boldsymbol {x})}{p _ {s} (y , z)} \tag {3}
$$
This is known as the "scaled likelihood trick" [Renals et al., 1994], and avoids us having to fit a generative model. We can estimate  $p_{s}(y,z|\boldsymbol{x})$  using any supervised learning method. However, to get good performance in practice, we need to ensure the probabilities are calibrated, as we explain in Section 3.1. Finally, we can estimate the source prior  $p_{s}(y,z)$  just by counting how often each  $(y,z)$  combination occurs in  $\mathcal{D}_s^{xyz}$  and then normalizing. Combining Equation (2) with Equation (3) we can compute the posterior distribution over augmented labels from the unlabeled target data as follows:
$$
p _ {t} (y, z | \boldsymbol {x}) \propto \frac {p _ {s} (y , z | \boldsymbol {x})}{p _ {s} (y , z)} p _ {t} (y, z). \tag {4}
$$
We can then compute  $f_{t}(\pmb{x})$  using (1). In summary, our TTLSA method consists of the following two steps:
1. Train on source. Train a model  $p_s(y, z|x)$  using supervised learning (and calibration) applied to  $\mathcal{D}_s^{xyz}$ , as explained in Section 3.1. Also estimate  $p_s(y, z)$  from  $\mathcal{D}_s^{xyz}$  by counting.
2. Adapt to target. Estimate  $p_t(y, z)$  using EM applied to  $\mathcal{D}_t^x$ , as explained in Section 3.2. Then compute  $p_t(y, z | x)$  using (2) and  $f_t(x)$  using (1).
We discuss these steps in more detail below.
3.1 Fit model to source distribution
In the first step, we fit a discriminative classifier  $p_{s}(y,z|\boldsymbol{x})$  on the source dataset. This is just like a standard classification problem, except we use an augmented output space, consisting of the class label  $y$  and group label  $z$ ; we denote this joint label by  $m \coloneqq (y,z)$ .
**Calibration** Our adaptation procedure crucially relies on the fact that the output of the classifier,  $p_s(m|x)$ , can be interpreted as calibrated probabilities. Since modern neural networks are often uncalibrated [see, e.g., Guo et al., 2017], we have found it important to perform an explicit calibration step. (See the supplementary for an ablation study where we omit the calibration step.) Specifically, we follow Alexandari et al. [2020] and use their "bias corrected temperature scaling" (BCTS) method, which is a generalization of Platt scaling to the multi-class case. In particular, let  $l(x)$  be the vector of  $M$  logits. We then modify  $p_s(m|x)$  as follows:
$$
p _ {s} (m | \boldsymbol {x}) \propto \exp \left(\frac {l (\boldsymbol {x}) _ {m}}{T} + b _ {m}\right) \tag {5}
$$
where  $T \geq 0$  is a learned temperature parameter, and  $b_{m}$  is a learned bias. This calibration is done after the source classifier is trained using a labeled validation set from the source distribution.
Logit adjustment The calibration method is a post-hoc method. However, if the source distribution has a "shortcut", based on spurious correlations between  $z$  and  $y$ , the discriminative model may exploit this by overfitting to it. The resulting model will not learn robust features, and will therefore be hard to adapt, even if we use calibration. To tackle this, we take inspiration from Proposition 1 of Makar et al. [2022], which showed that a classifier that is trained on the unconfounded or balanced distribution  $p_b$ , such that  $p_b(y,z) = p_s(y)p_s(z)$ , will not learn any "shortcut" between the confounding factor  $z$  and the target label  $y$ . Such a model will therefore have a risk which is invariant to changes in the  $p(z|y)$  distribution, and, in "purely spurious" anti-causal settings (a special case of Figure 1, see footnote 3), they showed that this is indeed the optimal one among all risk-invariant predictors (see discussion in Section 2). If both marginals are uniform, then  $p_b(y,z) \propto 1$ , so the corresponding balanced version of the source distribution becomes  $p_b(m|\pmb{x}) \propto \frac{p_s(m|\pmb{x})}{p_s(m)}$ , which in log space becomes  $\log(p_b(m|\pmb{x})) \equiv \log(p_s(m|\pmb{x})) - \log(p_s(m))$ . We train our classifier to maximize this balanced objective; this is known as "logit adjustment" [Menon et al., 2021]. After fitting, we can recover the classifier for the original source distribution using  $p_s(y,z|\pmb{x}) \propto p_b(y,z|\pmb{x})p_s(y,z)$ . Finally we apply calibration to  $p_s(y,z|\pmb{x})$ , as explained above. (Note that when the marginals are not uniform, one can apply a similar adjustment, subtracting off  $\log(p_s(y,z)/(p_s(y)p_s(z))$ .
3.2 Adapt model to target distribution
Given some unlabeled samples from the target distribution,  $\mathcal{D}_t^x = \{\pmb{x}_n\sim p_t(\pmb{x}):n = 1:N\}$ , our goal is to estimate the new prior,  $p_t(y,z) = p_t(m) = \pi_m$ . This is a standard subroutine in label shift adaptation methods [Lipton et al., 2018, Garg et al., 2020]; our innovation is to do this with respect to the meta-label  $m$ . Here, we use an EM approach [Alexandari et al., 2020], and maximize the log likelihood of the unlabeled data,  $\mathcal{L}(\pmb{x};\pmb{\pi},\pmb{\theta}) = \sum_{\pmb{x}_n}\log p_t(\pmb{x}_n;\pmb{\pi},\pmb{\theta})$ , wrt  $\pi$ , where
$$
p _ {t} (\boldsymbol {x}; \boldsymbol {\pi}, \boldsymbol {\theta}) = \log \left(\sum_ {m} \pi_ {m} p _ {t} (\boldsymbol {x} | m; \boldsymbol {\theta})\right) = \log \left(\sum_ {m} \pi_ {m} \frac {p _ {s} (m | \boldsymbol {x} ; \boldsymbol {\theta})}{p _ {s} (m ; \boldsymbol {\theta})}\right) + \text {c o n s t} \tag {6}
$$
where  $\theta$  are the parameters estimated from the source distribution. This objective is a sum of logs of a linear function of  $\pi$ , and is maximized subject to the affine constraints  $\pi_{m} \geq 0$  and  $\sum_{m=1}^{M} \pi_{m} = 1$ . Thus, under weak conditions on the ratios  $p_{s}(m|\boldsymbol{x};\boldsymbol{\theta}) / p_{s}(m;\boldsymbol{\theta})$  (these implied when  $\boldsymbol{x}$  contains some information that can discriminate between levels of  $m$ ) the problem is concave, with a unique global optimum, implying the parameters are identifiable [Alexandari et al., 2020].
A simple way to compute this optimum is to use EM, which automatically satisfies the constraints. Let  $\pi_m^j$  be the estimate of  $\pi_{m}$  at iteration  $j$ . We initialize with  $\pi_m^0 = p_s(m)$  and run the following
(b)
Figure 2: Test set AUC performance on (a) Colored MNIST and (b) CheXpert as we vary the correlation between  $y$  and  $z$  in the target distribution. The vertical dotted line marks the source distribution. There are three performance regimes: (1) the unadapted ERM model (blue line) degrades smoothly under shift; (2) the invariant models (yellow, green, and red lines) have flat performance; and (3) the adapted models (purple and brown lines) out-perform invariant models, yielding a U-shape.
iterative procedure below:
$$
p _ {t} (m | \boldsymbol {x} _ {n}; \boldsymbol {\pi} ^ {j}, \boldsymbol {\theta}) \propto p _ {t} (\boldsymbol {x} _ {n} | m; \boldsymbol {\theta}) \pi_ {m} ^ {j} \propto \frac {p _ {s} (m | \boldsymbol {x} ; \boldsymbol {\theta})}{p _ {s} (m ; \boldsymbol {\theta})} \pi_ {m} ^ {j} / / E s t e p \tag {7}
$$
$$
\pi_ {m} ^ {j + 1} = \frac {1}{N} \sum_ {n = 1} ^ {N} p _ {t} \left(m \mid \boldsymbol {x} _ {n}; \boldsymbol {\pi} ^ {j}, \boldsymbol {\theta}\right) / / \mathrm {M} \tag {8}
$$
We then set  $p_t(m) = \pi_m^J$ . We can also modify this to compute a MAP estimate, instead of the MLE, by using a Dirichlet prior. See the supplementary for a detailed derivation.
After estimating  $p_t(y, z)$  on  $\mathcal{D}_t^x$ , we can compute  $p_t(y, z | x)$  for the examples in  $\mathcal{D}_t^x$  using Equation (4).
In this section, we provide an experimental comparison of our method with various baseline methods on a variety of datasets. In particular, we compare the following methods:
ERM This corresponds to training a model on the source distribution, and then applying the same model to the target distribution without any adaptation, i.e., we assume  $p_t(y|\mathbf{x}) = p_s(y|\mathbf{x})$ .
gDRO The group DRO method of [Sagawa et al., 2020] is designed to optimize the performance of the worst performing group. (We only use this method for the worst-group benchmarks in Section 4.3.)