SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
p \left(m _ {n} = m \mid \boldsymbol {x} _ {n}, \boldsymbol {\pi} ^ {j}\right) = \frac {\boldsymbol {\pi} ^ {j} (m) p _ {s} \left(\boldsymbol {x} _ {n} \mid m\right)}{\sum_ {m ^ {\prime} = 1} ^ {M} \boldsymbol {\pi} ^ {j} \left(m ^ {\prime}\right) p _ {s} \left(\boldsymbol {x} _ {n} \mid m ^ {\prime}\right)} = \frac {\boldsymbol {\pi} ^ {j} (m) p _ {s} \left(m \mid \boldsymbol {x} _ {n}\right) / p _ {s} (m)}{\sum_ {m ^ {\prime} = 1} ^ {M} \boldsymbol {\pi} ^ {j} \left(m ^ {\prime}\right) p _ {s} \left(m ^ {\prime} \mid \boldsymbol {x} _ {n}\right) / p _ {s} \left(m ^ {\prime}\right)} \tag {23}
$$
We plug this into Equation (22) and then maximize Equation (21), using a Lagrange multiplier to enforce the sum to one constraint. We then get the following (see e.g., Sec 4.2.4 of Murphy [2022] for the derivation):
$$
\hat {\boldsymbol {\pi}} _ {m} ^ {j + 1} = \frac {\tilde {N} _ {m} ^ {j}}{\sum_ {m ^ {\prime} = 1} ^ {M} \tilde {N} _ {m ^ {\prime}} ^ {j}} \tag {24}
$$
where  $\tilde{N}_m^j$  are the prior pseudo counts plus the expected empirical counts:
$$
\tilde {N} _ {m} ^ {j} = N _ {m} ^ {j} + \alpha_ {m} - 1 \tag {25}
$$
At convergence, we have
$$
p _ {t} (y, z) = \hat {\pi} _ {y, z} ^ {J} \tag {26}
$$
If we assume that the class label prior is constant, and only the distribution of auxiliary labels has changed, then we can write
$$
p _ {t} (y, z) = p _ {s} (y) p _ {t} (z | y) \tag {27}
$$
$$
p _ {t} (z | y) = \frac {p _ {t} (y , z)}{\sum_ {z ^ {\prime}} p _ {t} (y , z ^ {\prime})} \tag {28}
$$
However, we do not make this fixed label assumption in our experiments.
Figure 3: Samples from ColoredMNIST. (a):  $y = 1$ ,  $z = 0$ . (b)  $y = 1$ ,  $z = 1$ .
Figure 4: Samples from CheXpert. Left: Female patient without effusion. Right: Male patient with effusion.
In this section we discuss the datasets in more detail.
B.1 Colored MNIST
We show some sample images in Figure 3.
B.2 CheXpert
We show some sample images in Figure 4. We list all the target attributes in Table 3. To test the difficult of each task, we train a logistic regression model for each attribute on the embeddings. (We get similar results using an MLP.) The resulting AUC scores are shown in Table 3. This shows we can reliably predict all the attributes from the embeddings. The table also shows the marginal distribution of each attribute. Many labels are highly skewed, which means accuracy would be a poor measure of the predictive performance.
Interestingly, we see that we can predict sex with an AUC of 0.973, which is higher than the AUC for effusion (0.861). To understand why, note that we only use frontal scans; consequently breasts are often visible in female patients, and this is often easier to detect visually than detecting the disease itself (see Figure 4), providing a possible "shortcut" for models to exploit.
Attribute	AUC	Prob.
NO_FINDING	0.873	0.909
ENLARGED_CARDIOMEDIASTINUM	0.652	0.942
CARDIOMEGALY	0.843	0.867
AIRSPACE_OPACITY	0.711	0.480
LUNG_LESION	0.761	0.963
PULMONARY_EDEMA	0.848	0.696
CONSOLIDATION	0.683	0.911
PNEUMONIA	0.742	0.973
ATELECTASIS	0.694	0.815
PNEUMOTHORAX	0.883	0.875
EFFUSION	0.861	0.508
PLEURAL_OTHER	0.752	0.987
FRACTURE	0.784	0.962
SUPPORT_DEVICESC	0.900	0.420
GENDER	0.973	0.586
AGE_AT_CXR	0.914	0.492
PRIMARY_RACE	0.731	0.459
ETHNICITY	0.681	0.728
Table 3: Metrics for all the attributes in the CheXpert dataset. (a) AUC using Logistic Regression on CXR embeddings. (b) Baseline prior probability for each attribute, illustrating the severe class imbalance for many attributes.
Figure 5: Performance on Colored MNIST using an uncalibrated tree classifier. TTSLA still improves the performance of the base model.
C Extra results
In this section, we include some extra experimental results.
C.1 Colored MNIST using gradient boosted tree classifier
In Figure 5, we show the results of various methods on the Colored MNIST dataset, where we use a Gradient Boosting Classification Tree as our base classifier, instead of a DNN. In particular, we use the HistGradientBoostingClassifier from scikit-learn [Pedregosa et al., 2011] with default parameters. The results are qualitatively similar to the DNN case.
C.2 The benefits of calibration
In Figure 6 we show the results on CheXpert if we remove the calibration step for our base classifier. Compared to Figure 2b, we see that the overall AUC of all the methods is worse, and the variance is larger. However, the rank ordering of the methods is the same. It is notable that a large gap opens up between the Oracle curve and the TTLSA implementations. This suggests that calibration primarily improves estimation of  $p_t(y,z)$  estimation via EM, because the Oracle curve in this subfigure corresponds to using the correct weights with the uncalibrated  $p_s(y,z|\boldsymbol{x})$  model.
C.3 CheXpert using CNN on raw pixels
In Figure 7 we show the result of various methods when applied to CheXpert images, as opposed to using embeddings. We use a ResNet-50 that was pretrained onImagenet, which we then fine tune on CheXpert images by replicating the gray-scale image along all 3 RGB channels. The qualitative conclusions are the same as in the embedding case.
C.4 More results on the benchmark datasets
In Table 4 and Table 5 we report the per-group accuracy on the benchmark datasets.
For a binary classification problem, calibration will not change the AUC, but since we derive the posterior over class labels by marginalizing a 4-way joint,  $p(y|\boldsymbol{x}) = \sum_{z=0}^{1} p(y,z|\boldsymbol{x})$ , calibration can help.
Figure 6: Performance across target domains on CheXpert embeddings, following the setup of Figure 2a. (a) Results using calibration. Performance mirrors those in Figure 2a. (b) Results without calibration. We see that calibration both improves performance and decreases variability between runs.
Figure 7: Performance on CheXpert using raw image (pixel) input instead of embeddings. These results are with calibration.
Data	Method	Group label (Y,Z)
(0,0)	(0,1)	(1,0)	(1,1)	(2,0)	(2,1)
CelebA	ERM	86.71 (0.67)	92.65 (0.79)	96.80 (0.18)	80.83 (1.46)
gDRO	92.71 (0.05)	92.33 (0.09)	92.63 (0.35)	87.36 (0.47)
SUBG	91.76 (0.22)	91.91 (0.55)	90.96 (0.33)	87.22 (1.38)
LA	91.43 (0.11)	94.77 (0.08)	95.88 (0.12)	84.72 (0.58)
TTLSA	97.59 (0.18)	98.78 (0.03)	80.70 (1.22)	51.25 (0.27)
Waterbirds	ERM	98.95 (0.13)	86.57 (0.48)	86.02 (0.13)	95.76 (0.16)
gDRO	93.46 (0.22)	88.00 (0.88)	90.15 (0.10)	92.06 (0.34)
SUBG	90.76 (0.69)	88.96 (0.19)	91.28 (0.35)	91.36 (0.24)
LA	94.43 (1.63)	88.38 (0.36)	91.32 (0.43)	93.15 (0.54)
TTLSA	94.59 (0.44)	93.68 (0.73)	95.72 (0.29)	97.12 (0.19)
MultiNLI	ERM	80.75 (0.79)	94.94 (0.11)	83.18 (0.47)	78.05 (1.31)	81.98 (0.48)	68.60 (0.40)
gDRO	80.36 (0.63)	85.27 (0.25)	82.48 (0.59)	81.21 (1.30)	79.39 (1.34)	76.87 (1.28)
SUBG	69.63 (0.17)	82.85 (0.19)	74.39 (0.21)	79.68 (0.17)	69.84 (0.48)	68.40 (1.33)
LA	81.63 (1.15)	87.79 (1.99)	84.36 (0.89)	80.95 (2.03)	78.77 (1.09)	76.33 (1.45)
TTLSA	80.24 (0.87)	94.74 (0.58)	81.73 (2.45)	73.90 (1.72)	82.40 (1.49)	63.76 (2.15)
CivilComments	ERM	92.23 (0.42)	90.38 (0.46)	68.57 (1.07)	68.32 (0.97)
gDRO	83.94 (0.70)	79.92 (0.33)	80.97 (0.63)	81.09 (0.42)
SUBG	79.79 (0.56)	79.14 (0.34)	82.52 (0.46)	76.56 (0.25)
LA	84.45 (0.16)	79.27 (1.17)	83.00 (0.95)	84.20 (0.99)
TTLSA	85.53 (1.35)	74.94 (1.96)	84.34 (2.36)	84.61 (2.21)
Table 4: Per-group accuracy on the benchmark datasets, where model selection is based on the worst  $(Y,Z)$  group accuracy on a validation set. Numbers in parentheses signify the standard error calculated based on 4 replication runs.
Data	Method	Group label (Y,Z)
(0,0)	(0,1)	(1,0)	(1,1)	(2,0)	(2,1)
CelebA	ERM	96.54 (0.21)	99.58 (0.05)	86.47 (0.89)	40.28 (2.24)
gDRO	95.48 (0.14)	96.76 (0.14)	87.10 (0.34)	68.75 (1.14)
SUBG	95.45 (0.39)	95.93 (0.44)	79.81 (1.30)	67.64 (4.66)
LA	95.25 (0.55)	98.96 (0.49)	88.77 (1.81)	43.47 (12.05)
TTLSA	97.68 (0.13)	98.99 (0.07)	80.21 (1.04)	47.36 (1.81)
Waterbirds	ERM	99.42 (0.11)	90.27 (1.24)	80.61 (2.51)	94.16 (0.84)
gDRO	97.04 (1.44)	92.84 (1.07)	83.84 (2.47)	89.10 (0.82)
SUBG	96.98 (0.29)	95.88 (0.42)	82.87 (1.51)	83.33 (1.86)
LA	98.23 (0.15)	92.42 (0.36)	85.98 (1.06)	92.91 (0.19)
TTLSA	99.06 (0.11)	93.61 (1.04)	87.66 (0.40)	95.02 (0.92)
MultiNLI	ERM	82.43 (0.06)	95.47 (0.08)	83.62 (0.03)	77.14 (0.16)	80.45 (0.09)	67.36 (0.54)
gDRO	80.37 (0.82)	86.32 (0.64)	81.06 (0.72)	78.22 (0.60)	81.22 (0.22)	78.83 (0.29)
SUBG	68.30 (2.00)	83.95 (2.28)	75.72 (1.68)	79.40 (1.37)	69.91 (1.55)	66.44 (2.23)
LA	82.74 (0.06)	92.92 (0.34)	83.97 (0.44)	79.88 (0.55)	79.77 (0.34)	71.49 (0.95)
TTLSA	81.86 (0.19)	96.52 (0.11)	83.89 (0.37)	76.07 (0.71)	80.51 (0.17)	56.60 (1.36)
CivilComments	ERM	96.00 (0.38)	95.63 (0.53)	55.27 (1.88)	52.21 (2.43)
gDRO	89.59 (0.68)	86.60 (0.86)	71.56 (1.64)	71.94 (1.33)
SUBG	81.43 (1.09)	80.67 (1.28)	80.80 (1.34)	76.05 (0.44)
LA	84.45 (0.16)	79.27 (1.17)	83.00 (0.95)	84.20 (0.99)
TTLSA	91.53 (0.58)	87.18 (1.70)	70.14 (2.69)	71.32 (2.82)
Table 5: Per-group accuracy on the benchmark datasets, where model selection is based on the average  $(Y,Z)$  group accuracy on a validation set. Numbers in parentheses signify the standard error calculated based on 4 replication runs.
C.5 Training with partial group labels
In this section, we evaluate an extension of our method where not all training samples have group labels  $z$ . In particular, we first train an ERM model to predict  $z$  on samples with group labels  $z$ , calculate  $p(z|x)$  for training samples with missing  $z$ , and then fit a new  $p(y,z|x)$  model on the augmented data. In particular, we represent each  $(y,z)$  target as a one-hot vector when  $z$  is known, and use a soft (predicted) encoding when  $z$  is unknown. We train with cross entropy loss. The use of soft labels may have the benefits of self-distillation Pham et al. [2022]. The validation set is always fully labeled for the purpose of hyperparameter tuning.
The results (on the 4 benchmark datasets) are shown in Table 6. The accuracy barely drops as missingness increases, which means our method is robust to the deficiency in group labels  $z$ .
Data	0	Missingness
0.5	0.75	0.875	0.9375
CelebA	84.72 / 95.55	78.33 / 95.68	77.78 / 95.37	79.44 / 94.44	77.22 / 95.20
Waterbirds	88.38 / 95.23	87.63 / 93.98	88.79 / 94.41	88.65 / 94.67	91.28 / 95.05
MultiNLI	76.33 / 82.60	74.87 / 79.55	74.72 / 82.49	76.05 / 82.61	78.75 / 81.72
CivilComments	79.27 / 85.03	76.26 / 85.87	73.87 / 83.55	73.41 / 84.57	66.64 / 80.36
Table 6: Accuracy of the worst / average  $(y, z)$  group on the benchmark datasets with partial training  $z$  labels, where model selection is based on average  $z$  accuracy. The Missingness columns stand for the proportion of training set with missing labels, e.g. 0.75 means only  $25\%$  of the training samples have  $z$  labels.
D Potential negative societal impacts
The proposed method in this work yields a model that can adapt to a new distribution and improves the performance at test time by exploiting spurious correlations to create a label shift correction technique that adapts to changes in the marginal distribution  $p(y, z)$  using unlabeled samples from the target domain. In this way, there are potential societal benefits to our method, especially when  $z$  corresponds to a socially salient attribute, such as a protected class. However, use cases of this type require caution, especially given the limitations discussed in Section 5. Further, as we discuss in a footnote in the main text, our method does not address concerns about cases where making decisions on the basis of  $z$  is discouraged or forbidden for a priori reasons. Given these limitations, there is a potential that the existence of adaptation methods of this type could be used to downplay the potential dangers of misusing sensitive information in machine learning systems. Here, we hope researchers and practitioners will instead acknowledge that, while beneficial use cases of  $z$  information exist, (1) there is a need to validate empirically that a particular use of  $z$  information is actually socially beneficial, and (2) there are valid reasons why one might want to avoid using  $z$  information altogether. Further, there is a potential risk that if the measurement quality of the labels  $y, z$  shift across distributions, such that they measure distinct concepts, or exhibit substantially different noise properties (i.e., become biased, or exhibit more outliers), our framework might absorb them during adaptation and eventually the outcomes of the system might be biased as well.
E Invariance Equivalences and Conditions
In this section, we review connections that have been established between risk invariance, ERM on balanced data, "separation" between a predictor  $f(X)$  and the spurious factor  $Z$ , and worst- $(y, z)$ -group performance. These results are useful for understanding why the application of logit adjustment at training time often yields a predictor that exhibits approximately invariant risk across the test sets that we study in our experiments.
E.1 Key Concepts
Risk invariance A predictor is risk-invariant with respect to a loss function  $\ell$  and a family of test distributions  $\mathcal{Q}$  iff it has the same risk  $E_{\mathcal{Q}}[\ell (f(X),Y)]$  for each  $Q\in \mathcal{Q}$ . The results we discuss apply to test distribution families that preserve both the generative distribution and the label distribution of the source distribution; that is,  $\mathcal{Q}$  is the set of distributions such that  $Q(Y) = P(Y)$  and  $Q(X\mid Y,Z) = P(X\mid Y,Z)$  for each  $Q\in \mathcal{Q}$ . This formulation allows  $Q(Z\mid Y)$  can change. This is the family considered in Makar et al. [2022] and Makar and D'Amour [2022], and is called a "causally compatible" family in Veitch et al. [2021], or a correlation shift in Yi et al..
Pure spuriousness The data generating process in Figure 1 is purely spurious if there exists some sufficient statistic  $e(X)$  such that (1)  $Y \perp X \mid e(X)$  and (2)  $e(X) \perp Z \mid Y$ . In words, if we know  $e(X)$ , there is no further dependence between  $Y$  and  $X$ , and further,  $e(X)$  does not depend on the spurious factor  $Z$  except through  $Z$ 's marginal dependence with  $Y$ . This is consistent with a causal model where the influence of  $Y$  on  $X$  is totally mediated by  $e(X)$ , and  $Z$  has no causal effect on  $e(X)$ .
Veitch et al. [2021] coined the term "purely spurious" in a context of a full counterfactual model of data generation, to refer to data generating processes where the portions of  $X$  that are causally related to  $Y$  and  $Z$  can be separated in a specific sense. Makar et al. [2022] consider the special case of pure spuriousness in the context of the anti-causal model in Figure 1. (They do not use the term "purely spurious" as the work in Veitch et al. [2021] was concurrent; Makar and D'Amour [2022] makes the connection explicit.) Here, we use formalism from Makar et al. [2022] to present the idea to minimize conceptual overhead.
Note that when the data  $X$  is rich, such as images are long passes of text, pure spuriousness is more plausible (or a better approximation to reality) because there is less possibility of destructive interference between  $Y$  and  $Z$  in the generation of  $X$ . Specifically, the simplest examples where pure spuriousness fails are ones where  $X$  is very low-content: e.g.,  $Y$  and  $Z$  are binary, and  $X := Y$  OR  $Z$ .
Separation Separation is a concept popularized in the literature on ML fairness [Barocas et al., 2019, Chapter 3], which stipulates that the predictor  $f(X)$  should satisfy the conditional independence  $f(X)Z \perp |Y$ . When  $Z$  is a sensitive attribute, this condition stipulates that the predictor  $f(X)$  should contain no more information about  $Z$  than one could glean from knowing  $Y$  alone.
Data balancing Idrissi et al. [2022] study predictors trained on data subsampled so that the  $(Y,Z)$  distribution is uniform; they call this data-balancing. Makar et al. [2022] and Makar and D'Amour [2022] study a similar predictors optimized on a similar "ideal" distribution, where  $Q(Y,Z) = P(Y)P(Z)$  for some source distribution  $P$ . This distribution does not "balance" the marginals of  $Y$  and  $Z$ , but it eliminates the marginal correlation between  $Y$  and  $Z$ .
Worst group performance Sagawa et al. [2020] define groups in terms of  $(z,y)$  values. The group conditional risk is  $R_{z,y} = E_Q[\ell (f(X),Y)\mid Z = z,Y = y]$ . Note that for all families of test sets that we consider, the group-conditional risks are equal for all Q. Worst group risk minimization attempts to minimize the group conditional risk of the worst subgroup. Saerens et al. [2002] propose a distributionally robust optimization algorithm for performing this minimization.
E.2 Connections
In the purely spurious setting, there are several connections and near-equivalences between risk invariance, separation, optimality on balanced data, and worst group risk minimization.
Yi et al. establish that for label distribution preserving target families, a predictor  $f(X)$  that satisfies separation  $f(X) \perp Z \mid Y$  will have invariant risk across the family  $\mathcal{Q}$  defined above. Notably, this result does not require pure spuriousness.
Under pure spuriousness, the separation condition achieves a certain optimality. Veitch et al. [2021], Theorem 4.3 establishes that in the purely spurious case, the minimax optimal across the family  $\mathcal{Q}$  satisfies separation  $f(X)\perp Z\mid Y$  . Similarly, under pure spuriousness, Makar and D'Amour [2022], Proposition 2, establishes that the optimal risk-invariant predictor satisfies separation.
Interestingly, this result establishes a connection between optimality under balanced data, separation, and optimal risk invariance. Specifically, Makar et al. [2022], Proposition 1 establishes that the optimal model for the "ideal" uncorrelated distribution for which  $Q(Y,Z) = P(Y)P(Z)$  achieves risk invariance across the family  $\mathcal{Q}$ . Thus, minimizing risk under a separation constraint targets a similar predictor to the predictor that one would target simply optimizing on balanced data. Makar and D'Amour [2022] shows that the near-equivalence holds up empirically, such that learning algorithms targeted at efficiently learning the optimal predictor on balanced data can satisfy both risk invariance and separation criteria.
Idrissi et al. [2022] establish that, empirically, models trained to minimize risk on balanced data also yield favorable worst-group performance, showing that subsampling can be particularly effective. Sagawa et al. [2020] explore similar ideas, focusing on reweighting strategies, which both they and Idrissi et al. [2022] find to work relatively poorly with neural models in the data regimes they study. Sagawa et al. [2020] further establish that under certian convexity conditions, there does exist a reweighting of the data that optimizes worst-group performance, but provide a counterexample showing that this is not always the case with non-convex losses.
Based on the above results, in the purely spurious case, one can establish the following, for  $\mathcal{Q}$  with a uniform distribution on  $Y$ :
1. There exists a predictor  $f^{*}(X)$  that is optimal on the ideal balanced data, is the optimal risk-invariant predictor, and satisfies separation  $f(X) \perp Z \mid Y$ .
2. For all  $Q \in \mathcal{Q}$ , the group-specific risks are equal within labels, i.e.,  $E_{Q}[\ell(f^{*}(X), Y) \mid Y = y, Z = z] = E_{Q}[\ell(f^{*}(X), Y) \mid Y = y, Z = z']$  for all  $y$ .
The latter fact does not imply that  $f^{*}(X)$  also optimizes worst-group risk, but it does imply that the worst group cannot be the worst due to a spurious correlation between  $Y$  and  $Z$ . This is because, for a fixed label value  $y$ , the risks of  $(y,z)$  subgroups are the same.