SlowGuess/NeurIPS_2023_OCR huggingface dataset:
SUBG The "SUBG" method of [Idrissi et al., 2022] subsamples the data so there is an equal number of examples in each group  $m = (y,z)$ , then trains a classifier by standard ERM. This is a simpler alternative to gDRO yet often achieves comparable performance.
LA This is our logit adjustment method of Section 3.1, which approximates a domain invariant classifier by targeting a uniform prior on the source domain, thus avoid overfitting to spurious correlations.
TTLSA This is our EM method of Section 3.2 that adapts the LA-based classifier using unlabeled data.
Oracle This is similar to TTLSA, but we replace the EM procedure with the ground truth target meta-label prior  $p_t(y,z)$ . This gives an upper bound on performance. (We only use this for the CMNIST experiments in Section 4.1 and the CheXpert experiments in Section 4.2, where we artificially control the degree of distribution shift.)
4.1 Colored MNIST
In this section, we apply our method to the Colored MNIST dataset [Arjovsky et al., 2019, Gulrajani and Lopez-Paz, 2021].
Table 1: The two "anchor" distributions, reflecting total positive and negative correlation between the class label  $y$  and the confounding factor  $z$ . Left:  $p_0$ . Right:  $p_1$ . From these distributions, we can create a family of target distributions  $p_\lambda = \lambda p_0 + (1 - \lambda)p_1$ .
Dataset We construct the dataset in a manner similar to [Arjovsky et al., 2019, Gulrajani and Lopez-Paz, 2021]. Specifically, we create a binary classification problem, where label  $y = 0$  corresponds to digits 0-4, and  $y = 1$  corresponds to 5-9. We then sample a random color  $z \in \{0,1\}$ , corresponding to red or green, for each image, with a probability that depends on the target distribution. See the supplementary material for a visualization of this dataset. Since the color is easier to recognize than the shape, color can act as a "shortcut" to predicting the class label, even though this is not a robust (domain invariant) feature.
We create a set of 21 target distributions  $p_{\lambda} = (1 - \lambda)p_{0} + \lambda p_{1}$ , where  $\lambda \in \{0, 0.05, \dots, 0.95, 1.0\}$  and  $p_{0}(y,z)$  and  $p_{1}(y,z)$  are two anchor distributions in which  $y$  and  $z$  are correlated and anticorrelated, respectively (see Table 1). By changing  $\lambda$ , we can control the dependence between  $y$  and  $z$ . We train the classifier on a source domain exhibiting a strong spurious correlation (we choose  $\lambda = 0.05$ ), and then apply the model to each target domain,  $p_{\lambda}$ , for  $\lambda \in \{0, 0.05, \dots, 0.95, 1.0\}$ . We measure performance using Area Under the ROC curve (AUC).
Training procedure During training, we fit a LeNet CNN on the training set by using  $m = (y,z)$  as the label. We use AdamW with a batch size of 64 and a learning rate of  $10^{-3}$  for a maximum of 5000 epochs, and run calibration for a maximum of 1000 epochs with the same learning rate and batch size on a  $10\%$  holdout set. However, training typically terminates much earlier since we also calculate the exponential moving average of the validation loss with a decay rate of 0.1, and stop early when the smoothed validation loss does not decrease for 5 consecutive epochs. During evaluation, we infer the target label prior for each target distribution using EM applied to an unlabeled test set of size 64 or 512, and we then predict the class labels on this test set. Finally we repeat this across all the unlabeled minibatches in the target distribution, and report the overall AUC.
Results Figure 2a shows our results, with error bars showing standard error of the mean across 4 trials. There are three main groups of curves. (1) The ERM model (black line) shows performance that gets steadily worse as the target distribution shifts away from the source. (2) The invariant models (SUBG in dashed orange and logit adjustment in solid red) are approximately constant across domains, as expected. (In this experiment, we see that logit adjustment outperforms SUBG.) 3) Our adaptive TTLSA models (blue curves), corresponding to TTLSA with 64 or 512 samples and the oracle method, all show a U-shaped curve, which is an upper bound on all the other curves. In particular, we see that the U-shaped curve is tangent to the invariant line when the target distribution is unconfounded  $(\lambda = 0.5)$ ; in that case, there is no prior information from  $p_t(y,z)$  to exploit. However, in all other target domains, the adapted prior information improves performance. As we increase the size of the unlabeled target dataset, we see that the performance of TTLSA approaches that of the oracle. To illustrate the fact that our method can also be applied to non-neural net classifiers we also applied TTLSA on top a gradient boosted tree classifier. The results (shown in the supplementary) are qualitatively similar to those in Figure 2a.
4.2 CheXpert
Next, we apply our method to the problem of disease classification using chest X-rays based on the CheXpert [Irvin et al., 2019] dataset. Chest X-rays are a particularly relevant application area for our method because sensitive attributes, such as patient sex or self-reported race, can be readily predicted from chest X-rays [Gichoya et al., 2022]. Recent work, e.g., Jabbour et al. [2020] and Makar et al. [2022], has confirmed the potential for classifiers to exploit these features as spurious features.
Dataset CheXpert has 224,316 chest radiographs of 65,240 patients. See the supplementary material for a visualization of this dataset. Each image is associated with 14 disease labels derived
from radiology reports, and 3 potentially confounding attributes (age, sex, and race), as listed in the supplementary. We binarized the attributes as in Jabbour et al. [2020], taking age to be 0 if below the median and 1 if above, and sex to be 0 if female and 1 if male. As for the class labels, we define class 0 as “negative” (corresponding to no evidence of a disease), and class 1 as “positive” (representing the presence of a disease); images labeled “uncertain” for the attribute of interest are discarded. Following [Glocker et al., 2022], we focus on predicting the label  $y =$  “Pleural Effusion” and use sex as the confounding variable  $z$ . See Figure 4 for some samples from the dataset. For the input features  $x$ , we either work with the raw gray-scale images, rescaled to size  $224 \times 224$ , or we work with 1376-dim feature embeddings derived from the pretrained CXR model [Sellergren et al., 2022]. This embedding model was pre-trained on a large set of weakly labeled X-rays from the USA and India. Note, however, that the pre-training dataset for CXR is distinct from the CheXpert dataset we use in our experiments.
To compare performance under distribution shift, we created a set of 21 distributions,  $p_{\lambda} = (1 - \lambda)p_{0} + \lambda p_{1}$ , where  $\lambda \in \{0, 0.05, \dots, 0.95, 1.0\}$  as before. We train using  $\lambda = 0.05$ , representing a strong spurious correlation at training time, and test with all 21 values. The test images are distinct from the training, and each test distribution has 512 samples in total. Each patient may have multiple images associated with them, but there is no patient overlap in the training and test distributions.
**Training procedure** When working with embeddings, we use a linear logistic regression model, following Sellargren et al. [2022], due to its simplicity and its good performance. To train this, we use AdamW with a batch size of 64 and a learning rate of  $10^{-3}$  for a maximum of 5000 epochs, and run calibration for a maximum of 1000 epochs with the same learning rate and batch size on a  $10\%$  holdout set. However, training typically terminates much earlier since we also calculate the exponential moving average of the validation loss with a decay rate of 0.1, and stop early when the smoothed validation loss does not decrease for 5 consecutive epochs. When working with pixels, we use a CNN. See the supplementary for details.
Results Our results for the embedding version of the data are shown in Figure 2b. The trends are essentially the same as in the colored MNIST case. In particular, we see that our TTLSA method outperforms the invariant baselines, and both adaptive and invariant methods outperform ERM. These results for the pixel version of the data are shown in Appendix C.3 in the supplementary. We find that the relative performance of the methods is similar to the embedding case, but the absolute performance for all methods is better when using embeddings, as was previously shown in Sellergren et al. [2022].
Discussion Glocker et al. [2022] point out that embeddings derived from X-ray classification models may contain information about sensitive attributes  $z$ , such as sex and age. We confirmed this result, and were able to classify sex with an accuracy of over  $95\%$  just using logistic regression on the CXR embeddings, as shown in Table 3. Glocker et al. [2022] argue that this may be harmful, since it can cause bias in the predictions of the primary label  $y$  of interest (disease status). As a counterpoint, our results also show that, if we can predict both  $z$  and  $y$  from the embeddings, this information can be used to make optimal adjustments for target populations featuring very different demographic makeup, in ways that can be beneficial for all groups. However, our results also confirm such information does need to be handled with care.
4.3 Worst-group vision and text benchmark datasets
In this section, we apply our method to four benchmark datasets that were first introduced in the group DRO paper [Sagawa et al., 2020], and have since been widely used in the literature on group robustness [see, e.g. Idrissi et al., 2022, on data balancing]. The four datasets are as follows.
CelebA Introduced in [Liu et al., 2015], this is an image dataset of celebrity faces. The class label  $y$  is hair color (blond / not-blond) and the group / attribute label  $z$  is sex (male / female).
Waterbirds Introduced in [Wah et al., 2011, Sagawa et al., 2020], this is an image dataset of birds synthetically pasted onto two different kinds of backgrounds. The class label  $y$  is bird type (land bird or water bird), and the group / attribute label  $z$  is background type (land or water).
MultNLI Introduced in [Williams et al., 2018], this is a dataset of sentence pairs,  $(s_1,s_2)$ , where the goal is to predict if  $s_1$  entails  $s_2$ . The class label  $y$  corresponds to entailment, contradiction or neutral, and the group / attribute label  $z$  indicates presence / absence of negation words.
CivilComments (CC) Introduced in [Borkan et al., 2019], This is a dataset of sentences from online forums. The class label  $y$  represents if the comment is toxic or not, and the group/attribute label  $z$  represents whether the content is related to a minority group (such as LGBT) or not.
Training procedure We use the code and hyper-parameters specified in Idrissi et al. [2022] for the ERM, gDRO, and SUBG baselines. We also use these same hyper-parameter values as ERM when fitting our own model TTLSA, which is trained to predict the augmented labels  $m = (y,z)$  using logit adjustment. We run the experiments on sixteen Nvidia A100 40GB GPUs. Depending on the dataset and method used, each experiment takes somewhere between 1 hour and 10 hours on an A100. In total, the experiments took 415 GPU hours.
For each method, for worst-group accuracy evaluations, we tune the model using worst-group accuracy in a validation set, whereas for average accuracy evaluations, we tune the model using average validation set accuracy. For worst-group evaluations, we use the LA baseline, without adjusting to the test set distribution. This effectively targets a balanced group distribution, similar to SUBG. For average target accuracy evaluations, we use the full TTLSA method with test-time EM adjustment.
Results We summarize our results in Table 2. We report the worst group and average group accuracy, averaged across 4 replication runs. Our worst group numbers for the baseline methods are within error bars for those reported in Idrissi et al. [2022].<sup>4</sup> The Idrissi et al. [2022] paper does not report average group performance, but we computed these results by modifying their code.
Data	ms	mt	ERM	gDRO	SUBG	LA/TTLSA
CelebA	0.44	0.49	80.83 (1.46) / 95.93 (0.03)	87.36 (0.47) / 94.68 (0.07)	87.10 (1.26) / 93.44 (0.19)	84.72 (0.58) / 95.55 (0.09)
Waterbirds	0.73	0.39	85.78 (0.24) / 93.19 (0.16)	87.98 (0.86) / 93.06 (0.62)	88.87 (0.14) / 93.48 (0.11)	88.38 (0.36) / 95.23 (0.34)
MultiNLI	0.49	0.49	68.60 (0.40) / 82.70 (0.02)	76.79 (1.24) / 81.16 (0.07)	67.89 (0.91) / 72.15 (0.25)	76.33 (1.45) / 82.60 (0.04)
CC	0.55	0.65	68.16 (1.03) / 88.00 (0.03)	79.66 (0.17) / 84.46 (0.43)	76.56 (0.25) / 79.56 (0.77)	79.27 (1.17) / 85.03 (0.71)
Table 2: Accuracy of the worst / average  $(y,z)$  group on the benchmark datasets. We define  $m_{s} = \max (\pi_{s})$  and  $m_{t} = \max (\pi_{t})$  as the maximum probability of a  $(y,z)$  group in the source and target distributions. The difference between these values reflects the degree of distribution shift.
Overall, the results show that TTLSA offers a unified approach to achieve a variety of robustness objectives. In terms of average group accuracy, we find that the performance of TTLSA relative to ERM depends on the nature of the shift. For CelebA and MultiNLI, there is no significant distribution shift, so TTLSA is similar to ERM, as expected. For Waterbirds, the majority class in the source becomes much less common in the target (reflecting the fact that the rare combinations of waterbirds on land and land-birds on water become more frequent). TTLSA is able to adapt to this, and outperforms ERM. For CivilComments, the majority class becomes even more common in the target distribution. Although TTLSA can adapt to this change, it cannot match the fact that ERM has been "rewarded" for learning a representation that is optimized for a single majority class. In terms of worst-group accuracy, the unadapted LA baseline (which is a special case of TTLSA where we do not use EM adaptation) achieves competitive performance to the dedicated gDRO and SUBG methods.
5 Conclusions and future work
We have shown that adapting to changes in the nuisance factors  $Z$  can give better results than using classifiers that are designed to be invariant to such changes. However, a central weakness of our approach is that it requires that the generative distribution  $p(x|y,z)$  be preserved across domains. This assumption becomes more plausible if the more factors of variation are included in  $Z$ . However, as  $Z$  becomes high dimensional, each step of our TTLSA method, as well as access to appropriate
Our worst group accuracy results for CivilComments are very different from those reported in Idrissi et al. [2022] despite using their code. The reason is that they binarize the  $z$  label during training, but use the fine-grained dataset with 9 possible  $z$  values during evaluation. Our method requires that  $z$  have the same set of possible values in train and test, so we use the coarse-grained dataset for both training and evaluation, and thus our performance metric is incomparable to theirs.
data, becomes more challenging. In this vein, another weakness is that we require access to labeled examples of  $Z$  during training. In the future, we would like to relax this assumption, potentially by using semi-supervised methods [c.f., Sohoni et al., 2021, Lokhande et al., 2022, Nam et al., 2022] that combine small fully labeled datasets with large partially labeled datasets. We also plan to explore the use of fully unsupervised estimates of the confounding factors  $Z$ , based on generative models, or by leveraging multiple source domains, similar to [Jiang and Veitch, 2022].
We discuss potential societal impacts in the appendix.
We would like to thank Andrew Sellergren at Google Health for his help with the CheXpert CXR embeddings, and Jonathan Caton at Google TPU Research Cloud for providing computational resources.
A. Alexandari, A. Kundaje, and A. Shrikumar. Maximum likelihood with bias-corrected calibration is hard-to-beat at label shift adaptation. In H. D. III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 222-232. PMLR, July 2020. URL https://proceedings.mlr.press/v119/alexandari20a.html.
M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893, 2019.
S. Barocas, M. Hardt, and A. Narayanan. *Fairness and Machine Learning: Limitations and Opportunities*. fairmlbook.org, 2019. http://www.fairmlbook.org.
D. Borkan, L. Dixon, J. Sorensen, N. Thain, and L. Vasserman. Nuanced metrics for measuring unintended bias with real data for text classification. In WWW, WWW '19, pages 491-500, New York, NY, USA, May 2019. Association for Computing Machinery. URL https://doi.org/10.1145/3308560.3317593.
M. J. Bours. Bayes' rule in diagnosis. Journal of Clinical Epidemiology, 131:158-160, 2021.
S. Garg, Y. Wu, S. Balakrishnan, and Z. Lipton. A unified view of label shift estimation. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 3290-3300. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/2020/file/219e052492f4008818b8adb6366c7ed6-Paper.pdf.
R. Geirhos, J.-H. Jacobsen, C. Michaelis, R. Zemel, W. Brendel, M. Bethge, and F. A. Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2(11):665-673, Nov. 2020. doi: 10.1038/s42256-020-00257-z. URL https://doi.org/10.1038/s42256-020-00257-z.
J. W. Gichoya, I. Banerjee, A. R. Bhimireddy, J. L. Burns, L. A. Celi, L.-C. Chen, R. Correa, N. Dullerud, M. Ghassemi, S.-C. Huang, et al. Ai recognition of patient race in medical imaging: a modelling study. The Lancet Digital Health, 4(6):e406-e414, 2022.
B. Glocker, C. Jones, M. Bernhardt, and S. Winzeck. Risk of bias in chest x-ray foundation models. Sept. 2022. URL http://arxiv.org/abs/2209.02965.
I. Gulrajani and D. Lopez-Paz. In search of lost domain generalization. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=1QdXeXDoWtI.
We have conducted a preliminary experiment in which we fit  $p_s(z|\pmb{x})$  on a small fully labeled subset,  $\mathcal{D}_s^{xyz}$ , and then use this to impute the missing  $z$  values on the larger  $\mathcal{D}_s^{xy}$ . We then train a new model  $p_s(y,z|\pmb{x})$  on the soft predicted  $z$  labels and the hard observed  $y$  labels. We get good results on the worst group benchmarks even when up to 90% of the  $z$  labels are missing (see Appendix C.5 in the supplementary material). However, we leave more detailed evaluation of this method to future work.
M. Zhang, S. Levine, and C. Finn. Memo: Test time robustness via adaptation and augmentation. arXiv preprint arXiv:2110.09506, 2021.
J. Zheng and M. Makar. Causally motivated multi-shortcut identification and removal. In NIPS, Oct. 2022. URL https://openreview.net/forum?id=-ZQ0x6yaVa-.
A Derivation of the EM algorithm
In this section we describe how to estimate the label prior on the target distribution,  $p_t(y,z) = p_t(m) = \pi_m$ , using the unlabeled data  $\mathcal{D}_t^x$ . There are several approaches to this, including a moment matching method called black box shift learning [Lipton et al., 2018] and an MLE approach based on the EM algorithm [Saerens et al., 2002]. In [Alexandari et al., 2020], they show that the MLE approach is much better, provided the classifier is calibrated. (See also [Garg et al., 2020] for a unified analysis of these two approaches.)
Since our augmented label space is expanded to include both class labels  $y$  and meta-data  $z$ , the number of labels  $M$  can be large, which can result in problems when computing the MLE. We therefore expand the previous approach to compute the MAP estimate, using a Dirichlet prior of the form
$$
\operatorname {D i r} (\boldsymbol {\pi} | \boldsymbol {\alpha}) = \frac {1}{B (\boldsymbol {\alpha})} \prod_ {m = 1} ^ {M} \boldsymbol {\pi} _ {m} ^ {\alpha_ {m} - 1} \tag {9}
$$
where  $B(\alpha)$  is the normalization constant. Note that the MLE solution can be recovered by setting  $\alpha = 1$ , which represents a uniform prior.
The goal is to maximize the (unnormized) log posterior of  $\pi$  given the unlabeled target data  $\mathbf{X}$ :
$$
\begin{array}{l} \mathcal {L} (\mathbf {X}; \boldsymbol {\pi}) = \log p _ {t} (\boldsymbol {\pi}, \mathbf {X}) (10) \\ = \log p _ {t} (\mathbf {X} | \boldsymbol {\pi}) + \log \operatorname {D i r} (\boldsymbol {\pi} | \boldsymbol {\alpha}) (11) \\ = \sum_ {n = 1} ^ {N} \log p _ {t} \left(\boldsymbol {x} _ {n} \mid \boldsymbol {\pi}\right) + \log \operatorname {D i r} (\boldsymbol {\pi} | \boldsymbol {\alpha}) (12) \\ = \sum_ {n = 1} ^ {N} \log \left[ \sum_ {m = 1} ^ {M} \pi_ {m} p _ {t} \left(\boldsymbol {x} _ {n} \mid m\right) \right] + \log \operatorname {D i r} (\boldsymbol {\pi} | \boldsymbol {\alpha}) (13) \\ \end{array}
$$
The first term can be rewritten as
$$
\begin{array}{l} \sum_ {n} \log \left[ \sum_ {m = 1} ^ {M} \pi_ {m} p _ {s} \left(\boldsymbol {x} _ {n} \mid m\right) \right] = \sum_ {n} \log \left[ \sum_ {m = 1} ^ {M} \pi_ {m} \frac {p _ {s} (m \mid \boldsymbol {x} _ {n}) p _ {s} \left(\boldsymbol {x} _ {n}\right)}{p _ {s} (m)} \right] (14) \\ = \sum_ {n} \log \sum_ {m} \frac {p _ {s} (m | \boldsymbol {x})}{p _ {s} (m)} \boldsymbol {\pi} _ {m} + \text {c o n s t} (15) \\ \end{array}
$$
This objective is a sum of logs of a linear function of  $\pi$ , as is the log prior. This needs to be maximized subject to the affine constraints  $\pi_{m} \geq 0$  and  $\sum_{m=1}^{M} \pi_{m} = 1$ , so the problem is concave, with a unique global optimum [Alexandari et al., 2020].
One way to compute this optimum is to use EM. Let  $\pi^j$  be the estimate of  $\pi$  at iteration  $j$ ; we initialize with  $\pi_m^0 = p_s(m)$ . First note that
$$
p _ {t} \left(\boldsymbol {x} _ {n}, m _ {n}\right) = p _ {s} \left(\boldsymbol {x} _ {n} \mid m _ {n}\right) p _ {t} \left(m _ {n}\right) = \prod_ {m = 1} ^ {M} \left[ p _ {s} \left(\boldsymbol {x} _ {n} \mid m\right) \boldsymbol {\pi} (m) \right] ^ {\mathbb {I} \left(m _ {n} = m\right)} \tag {16}
$$
Hence the complete data log posterior is given by
$$
\mathcal {L} (\mathbf {X}, \mathbf {M}; \boldsymbol {\pi}) = \sum_ {n = 1} ^ {N} \sum_ {m = 1} ^ {M} \mathbb {I} (m _ {n} = m) \log \left[ \boldsymbol {\pi} _ {m} p _ {s} \left(\boldsymbol {x} _ {n} \mid m\right) \right] + \log \operatorname {D i r} (\boldsymbol {\pi} | \boldsymbol {\alpha}) \tag {17}
$$
so the expected complete data log posterior is
$$
\begin{array}{l} Q (\boldsymbol {\pi}, \boldsymbol {\pi} ^ {(j)}) = E _ {\mathbf {M}} [ \mathcal {L} (\mathbf {X}, \mathbf {M}; \boldsymbol {\pi}) | \mathbf {X}, \boldsymbol {\pi} ^ {(j)} ] (18) \\ = \sum_ {n = 1} ^ {N} \sum_ {m = 1} ^ {M} p \left(m _ {n} = m \mid \mathbf {X}, \boldsymbol {\pi} ^ {j}\right) \log \left(\boldsymbol {\pi} _ {m} p _ {s} \left(\boldsymbol {x} _ {n} \mid m\right)\right) + \log \operatorname {D i r} (\boldsymbol {\pi} | \boldsymbol {\alpha}) (19) \\ = \sum_ {m = 1} ^ {M} N _ {m} ^ {j} \log \left(\pi_ {m} p _ {s} \left(\boldsymbol {x} _ {n} \mid m\right)\right) + \sum_ {m = 1} ^ {M} \left(\alpha_ {m} - 1\right) \log \pi_ {m} - \log B (\boldsymbol {\alpha}) (20) \\ = \sum_ {m = 1} ^ {M} N _ {m} ^ {j} \log \pi_ {m} + \underbrace {\sum_ {m = 1} ^ {M} N _ {m} ^ {j} \log p _ {s} \left(\boldsymbol {x} _ {n} \mid m\right)} _ {\text {c o n s t}} + \sum_ {m = 1} ^ {M} \left(\alpha_ {m} - 1\right) \log \pi_ {m} + \text {c o n s t} (21) \\ \end{array}
$$
where we drop constants wrt  $\pi$ , and where we defined the expected counts to be
$$
N _ {m} ^ {j} = \sum_ {n = 1} ^ {N} p \left(m _ {n} = m \mid \boldsymbol {x} _ {n}, \boldsymbol {\pi} ^ {j}\right) \tag {22}
$$
Hence in the E step we just need to compute the posterior responsibilities for each label: