Beyond Invariance: Test-Time Label-Shift Adaptation for Addressing "Spurious" Correlations
Authors: Qingyao Sun
Kevin Patrick Murphy
Sayna Ebrahimi
Alexander D'Amour
Venue: NeurIPS 2023 poster (2023)
Keywords: Distribution shift
Spurious correlation
Group robustness
TL;DR: A generic test-time adaptation method for improving robustness of classifiers in the presence of changes in latent group shifts

Abstract: 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.

Beyond Invariance: Test-Time Label-Shift A Adaptation
for Addressing "Spurious" Correlations
Qingyao Sun' Corelll University gs2340cornel1 edu
Kevin Murphy Google DeepMind spmurphy@google..cm
Sayna Ebrahimi Google Cloud AI Research saynae@google. com
Alexander D'Amour
Google DeepMind al gleeddmour@googlee.com com
Abstract
Changes the data distribution at test time can have deleterious effects on the performance of predictive models p(y]z)) We consider situations where there are additional meta-data labels (such as group labels), denoted by that can account for such changes in the distribution. In particular, we assume that the prior distribution ply, 2), which models the dependence between the class label and the "nuisance" factors 2, may change across domains, either due change in the correlation between these terms, change in one of their marginals. However, we assume that the generative model for features p(z [z/y,, 2) 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 Based on this observation, we propose test-time label shift correction that adapts changes in the joint distribution p(y. =) using EM applied to unlabeled are able
samples from the target domain distribution, Ps (z). Importantly, avoid fitting a generative model p(rly, z), and merely need to reweight the outputs of a discriminative model (3 2/2: trained the source distribution. We evaluate our method, which we call "Test- Time Label-Shift Adaptation" (TTLSA), on several standard image and text datasets, well as the CheXpert chest X-ray dataset, and show that improves performance over methods that target invariance changes in the distribution, methods. Code for reproducing experiments is available at boctps://githhe com/ com/nalzok/teee test-ttstteee leeee-laelllssse
well baseline empirical risk minimization
Introduction
Machine learning systems known 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 result, models that rely on such spurious correlations often have worse performance when they are deployed on targer domain that is distinct from the source domain on which they were trained [Geirhos et al. 2020 Izmailoy 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.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 al. 2022 Puli 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.c., 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 's probability of having the disease given positive test may be low if the disease is rare given that patient's 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-relever 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 f
spurious features provide
Motivated by the above, in this paper we propose a test-time approach for optimally adapting distribution shifts which arise due to changes in the underlying joint prior between the class labels y and the nuisance labels 2. We can view these changes as due to a hidden common cause u, such as the location of specific hospital. Thus we assume Pa (u) of pr (v)). where pa is the source distribution, and Pr is the target distribution. Consequently, pr(y,2) Σ. P(y, z|u)p,(ul (u) will change across domains i. However, we assume that the generative model of the features is invariant across domains, so p: 2) = p( z). See Figure | for an illustration of our modeling assumptions.
U
Y
Z
X
Figure I: Modeling assumptions. U is 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 dis- tribution of X given Y, Z is fixed.
The key observation behind our method is that our assump- tions are equivalent the standard "label shift assump- tion", except it is defined with respect to an expanded label m - (y, 2), which we call the meta-label. We call this the "expanded label shift assumption". This lets use existing la- bel shift techniques. such as Alexandari al. [2020). Lipton eet al. [2018), Garg et al. [2020], to adapt our model using small sample of unlabeled data {x,, ~ Pr (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( y, 2) is preserved across distribution shifts, it is based on learning discriminative base model (y, 2, æ). which we adjust to the target distribution Pr x²), as we explain in Section 3.1 Thus we do not need to fit generative model to the data. We do need access to labeled examples of the confounding factor = 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,
"We acknowledge that practitioners may have concerns beyond target domain performance, such certain notions of equal treatment, in which the reliance on such spurious features 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).
2

---

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 ps (a: to Pe (z)) which is not
captured by a change in the distribution over the causal factors (y+ 2).
2 Related work
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 Gul- rajani and Lopez-Paz, 202T]) ; which model is expected to generalize (i.e., achieve acceptable performance) in 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 natural approach.
Many of these methods make the same modeling assumptions as we do (shown in Figure D:: 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 satisty 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,2) 2) 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-robet 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 x) is preserved [see, e.g., (Shimodaira 2000), label shift assumption (i.e., that generative distribution p(æ y) is preserved [Lipton ct al. 2018 Garg et al. 2.20)).
In the setting of Figure D 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.2). 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 changes due to spurious correlations.
Interestingly, in the "purely spurious" setting described in Makar et al. [2022] which is a special case of Figure [] we can obtain a risk-invariant model as a special case, by adapting to a target distribution for which pe (r,,, =) = ps (=(()) |Veitch et al. 2021 Makar and D' Amour 2022] 3 Based on this observation, we show that logit adjustment [Meon el al.) 2021]. I. 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 = (g.2). 2).
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 Pr (r 2) from unsupervised target data, they instead assume that there is some value a such that pt (y, 2) = p=(y, (y,,)))aa They state that choosing a is an open problem, and therefore they report results for range of possible a's. By contrast, we provide way to estimate Pe(y, =), and we do not restrict it to have the above functional form.
'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 an association is "purely spurious" if there is sufficient statistic of X, e(X), such that y I. X e(X). Z and e(X) I. Z | 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 Idrissi et al. [2022), where spurious correlations are neutralized by data balancing, suggests that pure spuriousness is a common case, at least among worst-grous group benchmarks. We discuss pure spuriousness, gorst-group performance, and invariance in more detail in the supplement.
3

---

In Jiang and Veitch [022], they propose a method called Representation or "ACTII" 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 (z, v) examples to adapt their classifier at test time.
Transportable and Invariant
Test time adaptation There 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 heuristies to adapt discriminative models without specifying strong assumptions on the distribution shift structure. For example, TENT (Wang et al. [222]] uses entropy minimization to update the hatch normalization layers of CNN, and MEMO [Zhang et al., 2021] uses ensembles of predictions for different augmentations of test sample. Most of the TTA literature has focused on models that work on images (c.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 Time Label Shift Adaptation
Our aim is to construct a Bayes optimal predictor for the target distribution, fe(x)), using a large labeled dataset from the source distribution, Day= { yay zn. ~ pa( y, =)}. and small unlabeled dataset from the target distribution, Dr = far ~ (x)) The optimal prediction for the class label in the target domain is given by
(x) = arg max (x = arg max Σ Pr (u, =/z))
where the joint posterior over the class label y and the nuisance factor 2 given by
(y ( x (r 2)P: (y, 2) = ps (x (
(2)
where the first step follows from Bayes' rule, and the second step follows from our causal stability assumption that pa (æly, =) (xll,, z) p(ær y,, 2).
The first key insight of our method is that we can use the EM algorithm (or other optimization methods) to estimate (r(y,2) 2) 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(xlly, 2) up a constant of proportionality by fitting a discriminative model on the source dataset, and then dividing out by the source prior:
(v 2 (x)
Ps (xly, 2) x
Ps (s(y, 2)
This is known as the "scaled likelihood trick" [Rnnlls et al. 1994], and avoids us having to fit a generative model. We can estimate Ps(y, 2|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 estimate the source prior pv(y 2) just by counting how often each (y. combination occurs in Day= and then normalizing. Combining Equation (2) with Equation (3) we can compute the posterior distribution over augmented labels from the unlabeled target data follows:
P=(U. 2|xx) (v(y.2). 2).
Pe (c 2/2)
(4)
(y 2)
We can then compute fz( using In summary, our TTLSA method consists of the following two steps:
1. Train on source. Train a model (x 2 z|x)) using supervised learning (and calibration)
applied to Deye, as explained in Section 3.1 Also estimate (v, 2) from Dz by counting.
2. Adapt to target. Estimate Pt (y),2) using EM applied to DJ, as explained in Section 3.2 Then compute Pe(y, 2|æ) using (2) and f( using (I.
We discuss these steps in more detail below.

---

3.1 Fit model to source distribution
In the first step, we fit discriminative classifier (> =la:) on the source dataset. This is just like 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 == (y, 2).
Calibration Our adaptation procedure erucially relies the fact that the output of the classifier, (.)m)aa az), can be interpreted as calibrated probabilities. Since modern neural networks are often uncalibrated [see, e.g.. Guo et al., 22077] we have found it important to perform explicit calibration step. (See the supplementary for an ablation study where we omit the calibration step.) Specifically, we follow Alexandarr 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 ((z) be the vector of M logits. We then modily ps (m]x) as follows:
((z)m T
Ps (m))a:) x exp
bm
(5)
where T 0 is learned temperature parameter, and b, is 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 post-hoc method. However, if the source distribution has "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 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 pb. such that (s(y, =) = ps (z),.(z) will not learn any "shorteut" between the confounding factor and the target label y. Such model will therefore have a risk which is invariant to changes in the p= (= ly) 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 (b 2) x 1, so the
corresponding balanced version of the source distribution becomes Ps (m|mz) x P=(m[a) which in log p.(m)
space becomes log(ps (m (m]x)) = log(ps(m) x)) log(p.(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. (y, x x (b 2 2|2x)
p> ps (y. 2). Finally we apply calibration to (y. 2(x),, as explained above. (Note that when the marginals are not uniform, one can apply a similar adjustment, subtracting off log(p. (p)
3.2
Adapt model to target distribution
Given some unlabeled samples from the target distribution, D {zx Pr(x) 1: N}.
goal is to estimate the new prior, Pe (y. 2) = Pr (m) = Tm- This is a standard subroutine in label shift adaptation methods [Lipton et al. 2018 Garg et al. 2020 ; 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 log Pa (Tn; म. 0) wrt π,, where
likelihood of the unlabeled data, L(x; π 0) Ea.
Ps (m| 0) p 0)
Pr (x;; π,, 0) = loog Σ Tmpi (mpp(c//:
log
Σ Tm
const
(6)
where 0 are the parameters estimated from the source distribution. This objective is sum of logs of a linear function of π, and maximized subject to the affine constraints म. 2 and ΣM= nm = 1. Thus, under weak conditions some information that can discriminate between levels of m) the problem is concave, with a unique global optimum, implying the parameters are identifiable AAexandari et al. 22020].
the ratios P. (m|æ; pp (m: 0) (these implied when contains
simple way to compute this optimum is to use FM,, which automatically satisfies the constraints. Let ml, be the estimate of Km at iteration j. We initialize with Th p=(m) and run the following
5

---

2232
(a)
(b)
Figure Test set AUC performance on (a) Colored MNIST and (b) CheXpert as we vary correlation between and 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:
Ps(m (:mm)::0 T. // E step Pm 0)
Ps (m Rni x (e((aal 0) 0)Tl
(7)
x j++
(e T 0) // M step
(8)
N
We then set (m) The We can also modify this compute a MAP estimate, instead of the MLE, by using a Dirichlet prior. See the supplementary for a detailed derivation.
After estimating 2)) Dj, compute Pr(y, 2/x) for the examples in Dj using Equa- tion A
Experiments
In this section, provide 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 model on the source distribution, and then applying the same model to the target distribution without any adaptation, i.c., we assume Pe (g) = p. (v)=))
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)
SUBG The "SUBG" method of [Idrissi et 2022 subsamples the data so there an equal number of examples in each group m (z. 2). then trains classifier by standard ERM. This is simpler alternative to gDRO yet often achieves comparable performance.
LA This is logit adjustment method of Section 3.1 which approximates 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 TTLSA, but we replace the EM procedure with the ground truth target meta-label prior (v, 2). This gives an upper bound on performance. (We only use this for the CMNIST experiments in Section 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 AArovssy et al. 2019 Guirajani
and Lopez-Pau Paz 2021

---

0
1
0.5
0
0
0.5
y
0
0.5
1 0.5
Table 1: The two "anchor" distributions, reflecting total positive and negative correlation between the class label y and the confounding factor 2. Left: Po. Right: P1- From these distributions, we can create a family of target distributions pA Xpo + (1 A)pI.
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 = U corresponds to digits 0-4, and y = corresponds to 5-9. We then sample random color = E {0. 13. corresponding to red or green, for each image, with probability that depends on the target distribution. See the supplementary material for a visualizaiton of this dataset. Since the color is easier to recognize than the shape, color can act as "shortcut" to predicting the class label, even though this is not robust (domain invariant) feature.
We create a set of 21 target distributions pA = (1 A\Po + Ap1, where À E {0, 0.,0........ 0.995, 1.0], and Pu (y,, 2) and P1 (y, 2) are two anchor distributions in which y and 2 are correlated and anti- correlated, respectively (see [.. By changing X, we can control the dependence between y and 2. We train the classifier on a source domain exhibiting a strong spurious correlation (we choose A 0.05), and then apply the model to each target domain, pa, for À E {0, 0.05..... ... 0.95, 1.0]. We measure performance using Area Under the ROC curve (AUC).
Training procedure During training, we fit LeNet CNN on the training set by using m (y. 2) as the label. We use AdamW with a batch size of 64 and a learning rate of 10 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 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) 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 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 tangent the invariant line when the target distribution is unconfounded (A 0.5); in that case, there is no prior information from (y, 2) 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 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 partieularly relevant application area for our method because sensitive attributes, such as patient sex or self-reporte 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 visualizaiton of this dataset. Each image is associated with 14 discase labels derived
7

---

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. [020], 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 disease), and class as "positive" (representing the presence of disease): images labeled "uncertain" " for the attribute of interest are discarded. Following [Glocker et al. ] we focus on predicting the label y = "Pleural Effuusio" and use sex as the confounding variable 2. See Figure 4 for some samples from the dataset. For the input features æ, we either work with the raw gray-scale images, rescaled to size 224 x 224, or we work with 1376-dim feature embeddings derived from the pretrained CXR model [Sellergren et al., 2202). 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- rre-training dataset for CXR is distinct from the CheXpert dataset use in our experiments.
To compare performance under distribution shift, we created set of 21 distributions, PA (1 A)po + Ap1, where X E {0, 0.05, 0.95, 1.0} as before. We train using 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 linear logistic regression model, following Sellergren et al. (2022). due its simplicity and its good performance. To train this, we use Adam W with batch size of 64 and learning rate of 10 33 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. 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 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 similar to the embedding case, but the absolute performance for all methods is better when using embeddings, as was previously shown in Sellergren al. [2022].
Discussion Glocker al. [2022] point out that embeddings derived from X-ray classification
models may contain information about sensitive attributes 2, such as sex and age. We confirmed this result, and were able classify sex with accuracy of over 95% just using logistic regression on the CXR embeddings, as shown in Table 3 Glocker et al. [022] argue that this may be harmful, since it can cause bias the predictions of the primary label y of interest (disease status). As counterpoint, our results also show that, if we can predict both z and y from the embeddings, this information can be used make optimal adjustments for target populations featuring very different demographic makeups, 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 al. 2020], and have since been widely used in the literature group robustness [see, e.g. Idrissi et al.
2022 on data balancing]. The four datasets are as follows.
CelebA Introduced in [Lu al. 2015 this an image dataset of celebrity faces. The class label y is hair color (blond / not blond) and the group attribute label is sex (male / female).
Waterbirds Introduced in Wah et al. 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 a is background type (land or water).
2011 Sagawa et al. 2020 this is an image dataset of birds
8

---

MultNLI Introduced in [Williams et al. 2018). this is a dataset of sentence pairs, ($1, where the goal predict s entails $2- The class label y corresponds entailment, contradiction or neutral, and the group / attribute label 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 2 represents whether the content is related to minority group (such as LGBT) or not.
Training procedure We use the code and syper-parameters specified in Idrissi et al. [022 for the ERM. gDRO, and SUBG baselines. We also use these same hyper-paramedt values as ERM when fitting our own model TTLSA, which is trained to predict the augmented labels m (y. s) 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 hour and 10 hours on an A100. In total, the experiments took 415 GPU hours.
For each method, for worst-group accuracy evaluations, 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, 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) The Idrissi et al. [2022] paper does not report average group performance, but we computed these results by modifying their code.
Data
m.
my
gDRO
LA/TTLSA
CelebA
0.44
0.49
80.83 (1.46) (0:03)
87.36 (0.47)7 94.68 (0.07)
(1.26))94)9 (0.19)
(0.78
(0.09)
Waterbirds
0.73
0.39
85.78 /93.1 (6.16)
87.98 (0.86) / 93.00 (0.62)
88.87 (0.14))93345 (0.11)
88.38 (0.36) 93.23 (0.34)
MulaNLI
0.49
0.49
68.60 (0.40) 82.70 (0.02)
76.79 (1.24) 88.16 (0.07)
67.89 /2.11 (0.25)
76.33 (1.45) 82.60 (0.04)
CC
0.55
0.65
(8.16((.03) 88.00 (0.03)
79.66 (0.17) 84.46 (0.43)
76.56 (0.25) 79.5/799666 (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, = max(n,) and m. = max( as the maximum probability of (y, =) 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 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 water- birds on land and mird- 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 representation that is optimized for single majority class. In terms of worst-group accuracy, the unadapted LA baseline (which is special case of TTLSA where we do not use EM adaptation) achieves competitive performance 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(æ y, 2) be preserved across domains. This assumption becomes more plausible 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 ct al. [2022 despite using their code. The reason is that they binarize the label during training, use the line grained dataset with 9 possible values during evaluation. Our method requires that 2 have the same set of possible values in train and test, so we use the coarse gaasnne dataset for both training and evaluation, and thus performance metric 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 relax this assumption, potentially by using semi-supervised methods sc.f.,Sohoor et al. 2021 Lokhande et al. 2022 Nam et al. [022) that combine small fully labeled datasets with large partially labeled datasets We also plan 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.
Acknowledgements
We would like to thank Andrew Sellergren at Google Health for his help with the CheXpert CXR embeddings, and Jonathan Caton Google TPU Research Cloud for providing computational resources.
References
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 httpttt::ppppceeeeiiiiieeee alexandari20e html
mlr pwwss/vviss
M. Arjovsky, I.. 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 Opportu- nities. fairmlbook.org 2019. http://wwx ffirmllook. org
D. Borkan, Dixon, J. Sorensen, N. Thain, and Vasserman. Nuanced metries 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 httptttts://ii.rrg/////
1145/3308560 3317593
M. J. Bours. Bayes' rule in diagnosis. Journal of Clinical Epidemiology. 131:158- 160, 2021.
Garg, Y. Wu, S. Balakrishnan, and Z. Lipton. A unified view of label shift estima- tion. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Ad- vances in Neural Information Processing Systems, volume 33, pages 3290-3300. Curran Associates, Inc., 2020. URL https: ttttops:/pprceedinngge www.nenriip... wtwp:/inn..uurusscceceecoe neeuripp.cccppape/222////11. Paper pdf
R. Geirhos, J. Jacobsen, C. Michaelis, R. Zemel, W. Brendel, M. Bethge, and F. A. Wichmann. Shortcut learning in deep neural networks. Nature Muchine Intelligence, 2(1 6115)6666677, Nov. 2020. doi: 10. (0.1008//42256-022 022-02255--00-025--22
http:: org/ 10. 1038/ s42256 -20--0200000022
J. W. Gichoya, I. Banerjee, A. R. Bhimireddy, 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: modelling study. The Lancet Digital Health, 4(6):4006e441 14, 2022.
B. Glocker, C. Jones, M. Bernhardt, and S. Winzeck, Risk of bias in chest x-ray foundation models. Sept. 2022. URL
wtwpg///w..//...oa 02965
I. Gulrajani and D. Lopez-Paz. In search of lost domain d In International Conference on Learning Representationns 2021. htttpp:://peenrrrii....tt.// htw. 1ddXeXDoWtI
'We have conducted a preliminary experiment in which we fit (2/2z) on a small fully labeled subset, D then use this to impute missing values on the larger Diy, We then train a new model p. (3) a|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 labels are missing (see Appendix C.5 in the supplementary material). However, we leave more detailed evaluation of this method to future work.
10

---

C. Guo, G. Pleiss, Y. Sun, and K. Q. Weinberger. On calibration of modern neural networks. 2017.
B. Y. Idrissi, M. Arjovsky, M. Pezeshki, and D. Lopez-Paz. Simple data balancing achieves competi- tive nonst-group-accuracy. Conference on Causal Learning and Reasoning, pages 336-351. PMLR, 2022.
J. Irvin, P. Rajpurkar, M. Ko, Y. Yu, S. Ciurea-Ilcus, C. Chute, H. Marklund, B. Haghgoo, R. Ball, K. Shpanskaya. J. Seekins, D. A. Mong, S. Halabi, J. K. Sandberg, R. Jones, D. B. Larson, C. P. Langlotz, B. N. Patel, M. P. Lungren, and A. Y. Ng. Chexpert: A large chest radiograph dataset with uncertainty labels and expert comparison. Proceedings of the AAAI Conference Artificial Intelligence. 33(01):590 597, July 2019. doi: wtt./////aavv33311..... https://ojs aaai org/ index. php/ AAAI/ httpa//arrcceeeeeeee//eo /vieee/vve///344
URL
P. Izmailov, P. Kirichenko, N. Gruver, and A. G. Wilson. On feature learning in the presence of spurious correlations. In NIPS, Oct. 2022. http:///aar org/
11369
S. Jabbour, D. Fouhey, E. Kazerooni, M. W. Sjoding, and J. Wiens. Deep learning applied chest x-rays: Exploiting and preventing shortcuts. In F. Roshi-Velet J. Fackler, K. Jung, D. Kale, R. Ranganath, B. Wallace, and J. Wiens, editors, Proceedings of the 5th Machine Learning for Healthcare Conference, volume 126 of Proceedings of Machine Learning Research, pages 750-782, PMLR, Aug. 2020. BtKLLtttp::/ppoceedddiie http:///wws.......ess/ss///eeeee wwwps/vvee6....p jabourr2aabb. hhmll
Y. Jiang and V. Veitch. Invariant and transpoortable repressntattio for Anti-Causal domain shifts. July 2022.
J. N. Kaur, E. Kiciman, and A. Sharma. Modeling the data-generating process is necessary for out-of-distribution generalization. In Workshop on Spurious Correlations, Invariance, and Stability, ICML 2022, 2022.
P. W. Koh, S. Sagawa, H. Marklund, S. M. Xie, M. Zhang, A. Balsubramani, W. Hu, M. Yasunaga, R. L. Phillips, S. Beery, J. Leskovec, A. Kundaje, E. Pierson, S. Levine, C. Finn, and P. Liang. WILDS: A benchmark of in-the-wild distribution shifts. Dec. 2020. URL http:://axiv org/ Abs/2012.07422
J. Liang, R. He, and T. Tan. A comprehensive survey on test-time adaptation under distribution shifts. arXiv preprini 2rXiv:2303./536/ 2023.
Z. Lipton, Y.-X. Wang, and Smola. Detecting and correcting for label shift with black box predictors. In J. Dy and A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning. volume 80 of Proceedings of Machine Learning Research, pages 3122-3130. PMLR, July 2018. URL
mlr. http:///ww...a..pppo.//oo///////
ppress/88//iitton/
html
E. Z. Liu, B. Haghgoo, A. S. Chen, A. Raghunathan, P. W. Koh, S. Sagawa, P. Liang, and C. Finn. Just train twice: Improving group robustness without training group information. ICML, July 2021. htt/htpp:
www.:org///aoe httg/abe/21011 09044
09044
Z. Liu, P. Luo, X. Wang, and X. Tang. Deep learning face attributes in the wild. In ICCV, 2015.
V. S. Lokhande, K. Sohn, J. Yoon, M. Udell, C.-Y. Lee, and T. Pfister. Towards group robustness in the presence of partial group labels. In ICML Workshop on Spurious Correlations, Invariance and Stability, Jan. 2022. http:
org/ abs/abas//220 2201. 002222210660 03668
M. Makar and A. D' Amour. Faimess and robustness in anti-causal prediction. Sept. 2022. URL httt:://axivt
http:/ab/2200..0000
M. Makar, B. Packer, D. Moldovan, D. Blalock, Y. Halpern, and A. A'mour. Causally motivated shortcut removal using auxiliary labels. In AISTATS, volume 151, pages 739-766, 2022. URL
https://proceeddnng
mlr http:ppss///..///aaaa///////aa/
pdf
T. Makino, K. J. Geras, and K. Cho. Generative multitask learning mitigates target-causing confound- ing. Feb. 2022. http:://t//axxi wwtp:iii.orr/bb/////0000 04136
A. K. Menon, S. Jayasumana, A. S. Rawat, Jain, A. Veit, and Kumar. Long-tail learning via logit adjustment. In ICLR, 2021. URL http: /tap:::///xiiiv/ arxiv org/ abs 2007 20007007777
11

---

K. P. Murphy. Probabilistic Machine Learning: An introduction. MIT Press, 2022.
J. Nam, J. Kim, J. Lee, and J. Shin. Spread spurious attribute: Improving worst-group accuracy with spurious attribute estimation. In ICLR, Apr. 2022. URL httt://///aaxiiiv http:///aaxxii.rgaaba///////////
02070
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Pretten- hofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12:2825-2830, 2011.
M. Pham, M. Cho, A. Joshi, and C. Hegde. Revisiting self-distillation, 2022.
A. M. Puli, L. H. Zhang, E. K. Oermann, and R. Ranganath. Out-of-distribution generalization in the presence of Nuisance-Inducee spurious correlations. In ICLR, May 2022. URL https //operreview aepanreviev.nettforumnuuuee
S. Renals, N. Morgan, H. Bourlard, M. Cohen, and H. Franco. Connectionist probability estimators in HMM speech recognition. IEEE Trans. Audio Speech Lang. Processing, 2(1):161-174, Jan. 1994. BtLp:tt:///de
doi otg.://oooooooogoooooooo 10, 1109/89 89. 260359
M. Saerens, P. Latinne, and C. Decaestecker. Adjusting the Outputs of a Classifier to New a Priori Prob- abilities: A Simple Procedure. Neural Computation, 14(1):21 01 2002. ISSN 0899-7667. doi: 10.1162/89997662252844444
URL https://doi. org/10 1162/ (Tt...11///0697777777888880
S. Sagawa, P. W. Koh, T. B. Hashimoto, and Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for Worst-Case generalization. In ICLR, 2020.
URL tattp://arxiig org/ wwwp/////////////e///////eee/ee/ abs/ 1911/ 08731
B. Schoelkopf, D. Janzing. Peters, E. Sgouritsa, Zhang, and J. Mooij. On causal and anticausal learning. In J. Langford and J. Pineau, editors, Proceedings of the 29th International Conference on Machine Learning (ICML-12), ICML 12, pages 1255-1262, New York, NY, USA, July 2012. Omnipress. ISBN 1778---45003-285--11 URL https::/ttcceo :/ipll:///cmmm..m cc cc//01//e//// papers/ 625. pdd
A. B. Sellergren, Chen, Z. Nabulsi, Y. Li, A. Maschinot, A. Sarna, J. Huang, C. Lau, S. R. Kalidindi, M. Etemadi, Garcia- Vicente, D. Melnick, Y. Liu, K. Eswaran, D. Tse, N. Beladia, D. Krishnan, and S. Shetty. Simplified transfer learning for chest radiography models using less data. Radiology. 305(2):454-465, Nov. 2022. URthttpp://ds doi. org/10. 10. 1148/radiol. 212482
H. Shimodaira. Improving predictive inference under covariate shift by weighting the log-likelihood function. Journal of statistical planning and inference, 90(2):227-244 2000.
N. S. Sohoni, M. Sanjabi, N. Ballas, A. Grover, S. Nie, H. Firooz, and C. Ré. BARACK: Partially supervised group robustness with guarantees. 2021.
V. Veitch, A. D' Amour, S. Yadlowsky, and J. Eisenstein. Counterfactual invariance to spurious correlations in text classification. In NIPS, Nov. 2021. WBLlittppe//pppnneiiee...ee
id=BdkxQpOiBi8
C. Wah, S. Branson, P. Welinder, P. Perona, and S. Belongie. The Caltech-UCSD birds 2irds-200-201 dataset. Technical report, 2011.
D. Wang, E. Shelhamer, S. Liu, B. Olshausen, and T. Darrell. Tent: Fully test-time adaptation by entropy minimization. In International Conference Learning Representations, 2021. URL htttp:///penreviee..ett// wiw.ormmmmmddxx3332rrr2eee
O. Wiles, S. Gowal, F. Stimberg, S.-A. Rebuffi, I. Ktena, K. D. Dvijotham, and A. T. Cemgil. A fine- grained analysis on distribution shift. In International Conference on Learning Representations, 2022. URL htpp:ttte///ooorrrriio
D14LetuLdyK
A. Williams, N. Nangia, and S. Bowman. A Broad- Proad-Coverage challenge corpus for sentence under- standing through inference. In NAACL, pages 11 1112-1122, New Orleans, Louisiana, June 2018. Association for Computational Linguistics. URL
M. Yi, R. Wang. J. Sun, 7. Li, and 7.-M. Ma. Breaking correlation shift via conditional invariant regularizer. In The Eleventh International Conference Learning Representations.
12

---

M. Zhang, S. Levine, and C. Finn. Memo: Test time robustness via adaptation and augmentation. arXiv preprint arXiv:2100009566, 2021.
Zheng and M. Makar. Causally motivated multi-shortcut identification and removal. In NIPS, Oct. 2022. URL BtRthttss:// btstpp://ppenreviee..eee/t//
13

---

A Derivation of the EM algorithm
In this section describe how to estimate the label prior the target distribution, pe(y, Pe(m) = 7 using the unlabeled data D, There are several approaches to this, including a moment matching method called black box shift learning (Lipton et al. 2018 and 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 al. 2020 for unified
analysis of these two approaches.)
Since our augmented label space is expanded to include both class labels and meta-data 2, 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 Dirichlet prior of the form
M II
Dir(T
Cm m
B(a)
m=1
where B(x)) is the normalization constant. Note that the MLE solution can be recovered by setting Oz 1, which represents a uniform prior.
The goal to maximize the (unnormalized) log posterior of æ given the unlabeled target data X:
L( L(X; πτ)
Pe X)
(10)
lo
Pa X TT) + log Dir( a))
(11)
N
( + log Dir(Traa)
(12)
Σ log
Σ
Tm (2 /m)
+ log Dir(T(aa)
(13)
The first term can be rewritten as
M
M
(mmmxxn))) (m
Σ n
log
TmPo(Enlin
= Σ
log
(14)
Ps (m)
m==1
n
m==1
Pa (m( x))
log
Tm + const
(15)
(m)
m
This objective is of logs of linear function of n, is the log prior. This needs to be maximized subject to the affine constraints n. 2 and 5 M
Tm L, so the problem concave, with
m==
unique global optimum (Alexandari et al. 2020
One way to compute this optimum is EMM Let xe be the estimate of at iteration j; initialize with The (m). First note that
M II m==1
(r(an, m.) = (x (m.
[p (x,, m)) (mm)
(16)
Hence the complete data log posterior is given by
M
=
Σ
I(m, = (x + log Dir(T(a)
(17)
14

---

so the expected complete data log posterior is
Q Q (TT = M;; T (1)
(18)
N
pmm, m X n log log(TTpp((xmmmm) log (ir(zr
(19)
m=1
No log T Po (xm(mm))
(am 1) log Tm log B(())
(20)
M
M
M
NJ log Tm
NJ in
log (xnn/m))
+
Σ m=1
(am - 1) log n + const (21)
const
where we drop constants wrt ग, and where we defined the expected counts to be
NJ p(m, m mn, T²)
(22)
Hence in the step we just need to compute the posterior responsibilities for each label:
nt (m)pss (2rp)ss(xn//)))
p(m, = mm T²) = Em= n (m)))) (m))mm)
(23)
Σm лe (mn)),((rr)) (ml (p=)/pmmmm))
We plug this into Equation (22) and then maximize Equation (2T). 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):
(24)
Em'=1
M
No
where Ño are the prior pseudo counts plus the expected empirical counts:
No No.
am
(25)
At convergence, we have
Pe (r, 2) Ay.- assume that the class label prior is constant, and only the distribution of auxiliary labels has changed, then we can write
(26)
pr(u, 2) = (=()))((2)))
(27)
where
())
p=(=ly) Σ. Pe(y, =²)
(28)
However, do not make this fixed label assumption in experiments.
15

---

8
9
Figure 3: Samples from ColoredMNIST. (a): y 1, 0. (b)y=
AOKER
Figure 4: Samples from CheXpert. Left: Female patient without effusion. Right: Male patient with effusion.
Datasets
In this section we discuss the datasets in more detail.
B.1 Colored MNIST
We show some sample images Figure
B.2 CheXpert
We show some sample images in Figure We list all the target attributes in Table To test the difficult of each task, we train logistic regression model for each attribute on the embeddins. (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 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 often visible in female patients, and this is often easier to detect visually than detecting the disease itself (see Figure 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 ADEMA
0.848
0.696
CONSOLIDATION
0.683
0.911
PNEUMONIA
0.742
0.973
ATELECTASIS
0.694
0.815
16

---

Attribute
AUC
Prob.
PNEUMOTHORAX
0.883
0.875
EFFUSION
0.861
0.508
PLEURAL_ OTHER
0.752
0.987
FRACTURE
0.784
0.962
SUPPORT DEVICES GENDER AGE_AT_CHE CXR
0,900
0.420
0,973
0.586
0.914
0.492
PRIMARY RACE ETHNICITY
0.731
0.459
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.
17

---

1.00
0.98
(withont calibratio)
0.94
0.90
0.0
0.2
0.4
0.6
0.3
1.0
Shift parameter
+
ERM
TTLSA (omeele)
TTLSA (batch size 35000)
Figure 5: Performance on Colored MNIST using an uncalibrated tree classifier. TTSLA still improves 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 Gradient Boosting Classification Tree as our base classifier, instead of a DNN. In particular, we use the HistGradient inttradient.aeee inw.doostttottt...ommmmoooooooo. cepoosstinCCl.m ingClasse assifisee ffe 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 0 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 all the methods worse, and the variance larger. However, the rank ordering of the methods is the same. It is notable that large gap opens up between the Oracle curve and the TTLSA implementations. This suggests that calibration primarily improves estimation of Pr(y, 2) estimation via EM, because the Oracle curve in this subfigure corresponds to using the correct weights with the uncalibrated Pz(y (z) 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 ResNet-50 that was pretrained Imagenet, which we then fine tune on CheXpert images by replicating the gray-scale image along all 3 RGB channels. The qualitative conclusions the same in the embedding case.
C.4 More results the benchmark datasets
In Table and Table we report the per-group accuracy the benchmark datasets.
Foor a binary classification problem, calibration will not change the AUC, but since we derive the posterior over class labels marginalizing 4-way joint, p(y|æ) Σ1 p(y, z|æe), calibration can help.
18

---

0.95
AUC (without calibration)
0.90
0.65
0.80
0.75
0.70
0.2
0.4
1.0
Shift parameter
ERM
Logit Adjustmeut SUBG
TTLSA (ornele)
TTLSA (batch size 512)
TTLSA (batch size 64)
Figure 6: Performance across target domains CheXpert embeddings, following the setup of Figure 2a (a) Results using calibration. Performance mirrors those in Figure 2a (b) Results without calibration. We that calibration both improves performance and decreases variability between runs.
0.96
0.90
0.80
0.75
0.0
0.2
0.A
0.6
0.8
Stiit parameter
Logit SUBG
(onacle) (botch
ITLSA (batch 64)
Figure 7: Performance on CheXpert using raw image (pixel) input instead embeddings. These results with calibration.
19

---

Group label (Y, 2)
Data
Method
(0. 0)
(0. 1)
(1.
(1.1)
(2. 0)
(2. 1)
CelebA
ERM
(0.67)
92.65 (0.79)
96.80 (0.188
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.26)
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 (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)
MulsiLI
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)
$2.48 (0.59)
81.21 (1.30)
79.39 (1.34)
76.87 (1.28)
SURG
69.63 (0.17)
82.85 (0.19)
(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.
Group lahel (Y, 2)
Data
Method
(0. 0)
(0. 1)
(1. 0)
(1. 1)
(2, 0)
(2. 1)
CelebA
ERM
96.34
(0.21)
99.38 (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
Waterbirds
ERM
99. (99422(0 (0.11)
90.27 (1.24)
80.61 (2.51)
94.16
DRO
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)
MultiLII
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)
gRO
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
79.40
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)
(1.33)
SUBG
81.43 (1.09)
80.67 (1.28)
80.80 (1.34)
76.05 (0.44)
84,45 (0.16)
79.27 (1.17)
83.00 (0.95)
84,20 (0.99)
TILSA
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 predict z on samples with group labels 2, calculate p(z[z) for training samples with missing 2, and then fit new =[x) model on the augmented data. In particular, we represent each (y, =) target as a one-hot vector when is known, and use soft (predicted) encoding when is unknown. We train with cross entropy loss. The use of soft labels may have the benefits self-distillation |Phm et al. [2022] The validation set always fully labeled for the purpose of hyperparameter tuning.
The results (on the benchmark datasets) are shown in Table The accuracy barely drops as missingness increases, which means our method is robust to the deficiency in group labels 2.
20

---

Missingness
Data
0.5
0.75
0.875
0.9375
CelebA
84.72/ 95.55
78.33 95.68
77.78 95.37
79.44 99..44
77.22/ 95.20
Waterbirds
88.38 / 95.23
87.63
93.98
88.79
94.41
88.65 194.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, 2) group on the benchmark datasets with partial training
z labels, where model selection is based on average accuracy. The Missingness columns stand for the proportion of training with missing labels, e.g. 0.75 means only 25% of the training samples have z labels.
21

---

D Potential negative societal impacts
The proposed method in this work yields model that can adapt to 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,2) 2) 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 this type require caution, especially given the limitations discussed in Section 5 Further, as we discuss a footnote the main text, our method does not address concerns about cases where making decisions on the basis of is discouraged or forbidden for priori reasons. Given these limitations, there is 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, hope researchers and practitioners will instead acknowledge that, while beneficial use cases information exist, (1) there is need to validate empirically that particular use of information actually socially beneficial, and (2) there are valid reasons why one might want to avoid using information altogether. Further, there is potential risk that if the measurement quality of the labels y, shift across distributions, such that they measure distinct concepts, or exhibit substantially different noise properties (i.e., become biased, exhibit more outliers), our framework might absorb them during adaptation and eventually the outcomes of the system might be biased as well.
22

---

E Invariance Equivalences and Conditions
In this section, we review connections that have been established between risk invariance, ERM balanced data, "separation" between a predictor f(X) and the spurious factor Z, and worst-(y, 2)- group performance. These results are useful for understanding why the application of logit adjustment at training time often yields 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-invariam with respect to a loss function & and a family of
test distributions Q iff it has the same risk Ex]e(f(X), Y)] for each Q E Q. The results
discuss apply to test distribution families that preserve both the generative distribution and the label distribution of the source distribution; that is, Q is the set of distributions such that Q(Y) P(Y)
and Q(X Y, Z) = P(X Y. Z) for each Q E Q. This formulation allows Q(Z Y) can change. This is is the family is considered in Makar et al. [2022] and Makar and D A'mour [2022], , and is
called a "causally compatible" family in Veitch et al. [021]. or correlation shift in Yi et al.
Pure spuriousness The data generating process in Figure is purely spurious if there exists some sufficient statistic e(X) such that (1) 1 X e(X) and (2) e(X) 1 Z 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 e(X).
Veitch et al. 202 coined the term "purely spurious" in a context of 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" the work in Veitch et al. [2021] was concurrent; Makar and D A'mour [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 descructive interfecence between and 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 the conditional independence f(X)z 1 IY. 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 alone.
Data balancing [diissi et al. 2022 study predictors trained on data subsampled so that the (Y distribution uniform: they call this data-balancing. Makar et al. [2022] and Makar and D'Amour [022] study a similar predictors optimized on a similar "ideal" " distribution, where Z) ((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 al. 2020 conditional risk is R2,, (x)) Z 2, Y = y). Note that for all families of test sets that consider, the group-come conuicoondiiional risks are equal for all Q. Worst group risk minimization attempts to minimize the group conditional risk of the worst subgroup. SSarenn et al.) [2002) propose a distributionally robust optimization algorithm for performing this minimization.
define groups in terms of (z. y) values. The group
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.
23

---

et al. establish that for label distribution preserving target families, a predictor f that satisfies separation f(X) 1 will have invariant risk across the family 2 defined above. Notably, this result does not require pure spuriousness.
Under pure spuriousness, the separation condition achieves certain optimality. Veitch al. 2021
Theorem 4.3 establishes that in the purely spurious case, the minimax optimal across the family Q satisfies separation f(X) I Z Y. Similarly, under pure spuriousness, Makar and D Amour [2022] Proposition 2, establishes that the optimal risk-invariant predictor satisties separation.
Interestingly, this result establishes connection between optimality under balanced data, separation, and optimal risk invariance. Specifically, |Mkar et al.] [2022]. Proposition establishes that the optimal model for the "ideal" uncorrelated distribution for which Q(Y, Z) PPY) achieves risk invariance across the family Q. Thus, minimizing risk under a separation constraint targets a similar predictor to the predictor that would target simply optimizing balanced data. |Makar and D' A'mour [022) shows that the near-equivalence holds up empirically, such that learning algorithms targeted 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 minimize risk balanced data
also yield favorable worst-group performance, showing that subsamping can be particularly effective. Sagawa et al. [2020] explore similar ideas, focusing on reweighting strategies, which both they and et al. [2022] find to work relatively poorly with neural models the data regimes they study. (Sagawa al. (202)] 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 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) 1 Y.
2. For all Q E Q. the group-specifes risks are equal within labels, i.e., (X))( YY= y, 2 = =| = (X))) Y = y, = z'| for all y.
The latter fact does not imply that (X) also optimizes worst-group risk, but it does imply that the worst group cannot be the worst due to spurious correlation between Y and Z. This is because, for fixed label value y, the risks of (y, 2) subgroups are the same.
24


Reviews:
[{"review_id": "jZvYMfkwcQ", "reviewer": "Reviewer_KyoC", "summary": "- This paper tackles the test-time adaptation problem to mitigate spurious correlation. By assuming the pervasiveness of $p(x|y, z)$, the authors derived an alternative method for predicting the target distribution and proposed Test-Time Label-Shift Adaptation (TTLSA) to optimize it. The TTLSA involves two steps: training on the source with calibration and logit adjustment and adaptation step using the EM algorithm. The method is evaluated on Colored MNIST, CheXpert, and four popular worst-group benchmark datasets, demonstrating superior adaptation performance (Figure 2) and compatible benchmark performances (Table 2).\n\nMinor corrections: Some references to figures are presented as double question marks (??).", "questions": "In [2], it is claimed that Group DRO performs well because it improves the last linear layer rather than the underlying feature representations. How does TTLSA address this aspect? ", "limitations": "The proposed method requires group information in the training dataset, which can be costly.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper has a coherent flow, making it easy to follow, and the details and evidence are well-presented.\n- The motivation for test-time adaptation is solid, and the experimental results in Figure 2 provide strong support.\n- The ERM performances in Table 2 are reliable, as they closely match the accuracies achieved in [1].\n\n[1] Idrissi, Badr Youbi, et al. \"Simple data balancing achieves competitive worst-group-accuracy.\" *Conference on Causal Learning and Reasoning*. PMLR, 2022.", "weaknesses": "- One of my concerns is the assumption of having an annotated training dataset. The field is evolving towards building algorithms without relying on group information.\n- Another concern is in comparison with Group DRO. Both methods assume that group information is available in the source dataset. TTLSA shows better test accuracy in this experiment, but not for worst-group accuracy.\n- The Group DRO performance is lower than what is presented in [2].\n\n[2] Izmailov, Pavel, et al. \"On feature learning in the presence of spurious correlations.\" *Advances in Neural Information Processing Systems* 35 (2022): 38516-38532.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. Response to your points is below.\n\n“**One of my concerns is the assumption of having an annotated training dataset.**\"\n\nWe agree this is a weakness. However, as we show in Appendix C.5 (in supplemental), it is possible to apply our method with a very small amount of labeled data. In particular, we get good results on the worst group benchmarks even when up to 90% of the z labels are missing. The trick is to use label imputation to infer the missing z labels before applying our algorithm. We will clarify this in the revised paper.\n\n“**TTLSA shows better test accuracy in this experiment [than gDRO], but not for worst-group accuracy.**”\n\nIn table 2, our worst group performance is actually on par with gDRO for 3 out of 4 datasets, with the exception being CelebA. For example, on Waterbirds, gDRO gets 87.98 and we get 88.38. However, the main focus of our paper is on average performance.\n\n“**The Group DRO performance is lower than what is presented in Izmailov**”.\n\nWe ran the gDRO code by Idrissi (with their published hyper-parameters) and reported our findings, which agrees with the original version by Sagawa. Recently, Izmailov provides another implementation of gDRO, but we did not use their code, nor do we know why they reported different numbers for their baseline.\n\n“**In [2], it is claimed that Group DRO performs well because it improves the last linear layer rather than the underlying feature representations. How does TTLSA address this aspect?**”\n\nThe purpose of TTLSA is to make better predictions in the presence of target label shift, not learning better representations. In fact, both logit adjustment and TTLSA operate ad hocly on the output logits, and are thus oblivious to the model internals (which is why it applies to both neural networks and gradient boosted trees). That being said, asking the model to predict (y, z) jointly instead of only y may force it to learn a better presentation for features relevant to z."}, {"review_id": "qI0At95xwR", "reviewer": "Reviewer_wAGR", "summary": "**A brief summary:**\nThe paper proposed \"Test-Time Label-Shift Adaptation\" (TTLSA), that utilizes rather than eliminates spurious correlations. TTLSA adapts to changes in marginal distribution using the Expectation Maximization algorithm. Experiments on several datasets show that TTLSA outperformed traditional invariance methods and baseline empirical risk minimization. \n\n**Main Contributions:**\nThe authors proposed a novel test-time adaptation method (TTLSA that capitalizes on spurious correlations instead of trying to eliminate them. Moreover, they expand on the label shift assumption, incorporating nuisance factors into the labels, train a discriminative classifier to predict the source distribution, and adapt to changes in the marginal distribution using the Expectation Maximization (EM) algorithm.", "questions": "* (1) The introduced model assumption serves as a foundation for the proposed method. Is it possible to empirically verify or illustrate more about the reasoning behind adopting such an assumption? \n\n* (2) It would be better to distinguish the notation of $x_n$, which appears in both the source and target distributions (lines 145-146)\n\n* (3) In Figure 2, can the authors explain more about the $U-shape$ performance of TTLSA? Specifically, why the performance of TTLSA under shift parameter (0.0) might be worse than that of shift parameter (1.0)?\n\n* (4) It seems that certain references to figures are not appropriately included, i.e., line 311.\n\n* (5) It would be better if the authors could discuss the differences with two closely related works, i.e., [1, 2].\n\nReferences:\n\n[1] Coping with label shift via distributionally robust optimization. [ICLR'21] --> propose an objective to cope with label shift, and provide an adversarial algorithm to effectively optimize it.\n\n[2] Distributionally Robust Post-hoc Classifiers under Prior Shifts. [ICLR'23] --> a method for scaling the model predictions at test-time for improved distribution robustness to prior shifts.", "limitations": "Same as mentioned weaknesses.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Strengths are summarized below:\n\n* The paper introduced an assumption that there exists \"a hidden confounder that induces a spurious correlation between the label $Y$ and other causal factors $Z$, which together generate the features $X$,\" such that the label space in the DRO setting could be extended to $m=(y, z)$, hence treat like classical learning paradigm.\n\n* Experiments on several datasets show that TTLSA achieves promising results compared with traditional invariance methods.", "weaknesses": "Main weaknesses are two folds:\n\n* (1) Although the proposed method performs the label shift adaptation at test time, it still relies on a trained discriminative classifier, which is used to predict $p(y, z|x)$ on the source distribution.\n\n* (2) The training on the discriminative classifier requires a large labeled dataset from the source distribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. Response to your points is below.\n\n\n“**The training on the discriminative classifier requires a large labeled dataset from the source distribution….  still relies on a trained discriminative classifier, which is used to predict p(y, z|x)**”\n\n We agree this is a weakness. However, as we show in Appendix C.5 (in supplemental), it is possible to apply our method with a very small amount of labeled data. In particular, we get good results on the worst group benchmarks even when up to 90% of the z labels are missing. The trick is to use label imputation to infer the missing z labels before applying our algorithm. We will clarify this in the revised paper.\n\n“**The introduced model assumption serves as a foundation for the proposed method. Is it possible to empirically verify or illustrate more about the reasoning behind adopting such an assumption?**”\n\nOur method relies on an augmented version of the label shift assumption. The fact that we get better performance on a variety of datasets is empirical demonstration of the reasonableness of our assumption. However, we do not claim that it captures all (or even most) forms of distribution shift - it depends on the richness of the latent factors z.  \n\n“**In Figure 2, can the authors explain more about the U-shape performance of TTLSA? Specifically, why the performance of TTLSA under shift parameter (0.0) might be worse than that of shift parameter (1.0)?**”\n\nPerformance for 0.0 is in fact the same as for 1.0, modulo experimental noise. To see why, note that if there is 100% correlation  or anti-correlation between z and y, then it becomes easy to predict the class label y by simply learning to predict z and then setting y=z or y=not(z). Conversely, if there is no correlation (0.5 on x-axis), there is no shortcut to exploit, and the model reduces to the performance of ERM predicting p(y|x). This is why there is a symmetric U-shaped curve.  \n\n“**It would be better if the authors could discuss the differences with two closely related works, i.e., [1, 2].**”\n\nWe will add discussion of these papers.\n\n"}, {"review_id": "fD0vyKFFck", "reviewer": "Reviewer_1Li8", "summary": "This paper introduces a new method named Test-Time Label Shift Adaptation(TTLSA) to manage distribution shifts in machine learning, specifically focusing on spurious correlations. \nRather than eliminating these spurious relations as mentioned in other papers, TTLSA utilizes the spurious features to adjust the model and argues that this could substantially outperform invariant predictors. The method involves augmenting the label space with a \"nuisance factor\", z, which encapsulates the spurious relations in the data.\n\nTTLSA method involves two main steps:\n1. Train on Source: The model is initially trained on the source data with augmented labels that include both the target labels (y) and the nuisance variables (z). This step helps to capture the spurious correlations that exist in the source data.\n2. Adapt to Target: During testing, TTLSA uses an EM algorithm to adapt to the target data. This step helps the model to adapt to the shift in label distribution between the training and testing data.\n\nThe authors tested TTLSA on several datasets, including the CheXpert and colored MNIST datasets. In both cases, TTLSA outperformed traditional methods in terms of AUC under shifted distributions. However, the authors acknowledge the approach's limitations: it requires the generative distribution to be preserved across domains and needs access to labeled examples of Z during training.", "questions": "1. The paper presents great results using the proposed TTLSA method, but the assumption that the generative distribution p(x|y, z) is preserved across domains seems to rest heavily on the presence of one auxiliary feature. Can the authors provide additional insights into scenarios where this assumption is likely to hold? Are there particular types of data or specific domains where this assumption is more valid?\n\n2. Related to the first question, in real-world scenarios, there can be multiple factors affecting the distribution. How would the TTLSA method behave in such complex situations, given its dependency on a single auxiliary feature for its assumptions? Can the authors share their thoughts on extending their methodology to accommodate multiple auxiliary features?\n\n3. The paper mentioned future work that plans to relax the assumption of requiring labeled examples of Z during training. Could the authors provide some preliminary insights into how they plan to achieve this?\n", "limitations": "The authors have very adequately addressed the limitations.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "**Originality**: The paper introduces a novel method, Test Time Label Shift Adaptation (TTLSA), to handle the problem of distribution shifts in machine learning models. It leverages the spurious relations in the data rather than attempting to eliminate them. This unique approach sets it apart from previous methods.\n\n**Quality**: The paper is of good quality. The authors' method is well-defined, and the experiments are carefully designed, involving various benchmark datasets to demonstrate the effectiveness of the proposed method.\n\n**Clarity**: The paper is well-written and clear, providing a detailed explanation of the methodology, experimental design, and results. \n\n**Significance**: The paper's results on various datasets indicate that TTLSA has the potential to improve machine learning model robustness in the presence of distribution shifts and spurious relations. However, the significance of the method is limited by its assumptions, notably the need for the generative distribution to remain constant across domains. This condition is only based on one additional feature, limiting the method's scope. The results, although promising, should be taken with a degree of caution until further research is conducted to validate these assumptions in more complex scenarios. Overall, the paper holds promise but requires additional exploration and validation.\n", "weaknesses": "1. Assumptions: The TTLSA method assumes that the generative distribution p(x|y, z) is preserved across domains. Here, the authors only address z as a one-dimensional feature. This assumption can be restrictive and may not hold true in many real-world scenarios.\n\n2. Data Requirement: The proposed method requires access to labeled examples of the nuisance factor Z during training. This could be a limitation in scenarios where such labeled data might not be readily available or where it's challenging to identify and label the nuisance factors.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review. Response to your points is below.\n\n“**The TTLSA method assumes that the generative distribution p(x|y, z) is preserved across domains. Here, the authors only address z as a one-dimensional feature. This assumption can be restrictive and may not hold true in many real-world scenarios.**”\n\nWe 100% agree. We leave it to future work to consider multi-variate z. For low-dimensional z, one naive approach is to simply take a cartesian product of its components.\n\n“**The proposed method requires access to labeled examples of the nuisance factor Z during training**”\n\n We agree this is a weakness. However, as we show in Appendix C.5 (in supplemental), it is possible to apply our method with a very small amount of labeled data. In particular, we get good results on the worst group benchmarks even when up to 90% of the z labels are missing. The trick is to use label imputation to infer the missing z labels before applying our algorithm. We will clarify this in the camera ready.\n\n“**The assumption that the generative distribution p(x|y, z) is preserved across domains seems to rest heavily on the presence of one auxiliary feature. Can the authors provide additional insights into scenarios where this assumption is likely to hold?**“\n\nAs mentioned above, we will likely need to make z be a vector of factors (not just a single categorical variable) to make the method more general. We hope that datasets with richer meta-data may provide a source for such structured z’s. Alternatively we may be able to use techniques from the disentangled representation learning literature to avoid the need for extra annotations. However we leave this to future work.\n\n“**Can the authors share their thoughts on extending their methodology to accommodate multiple auxiliary features?**”\n\nWe have indeed thought about this. One idea would be to make p(y, z(1:K) | x) be some kind of structured model, such as a (sparse) conditional random field. The key is to capture the dependence between each z(k) and y, conditioned on x. However we leave this to future work.\n\n“**The paper mentioned future work that plans to relax the assumption of requiring labeled examples of Z during training. Could the authors provide some preliminary insights into how they plan to achieve this?**”\n\nSee comment above about our experimental results in Appendix C.5 in the supplemental.\n\n\n"}, {"review_id": "JSLqHVN6aI", "reviewer": "Reviewer_C4bL", "summary": "This paper proposes to perform test-time adaptation to exploit the existence of spurious correlations in a setting where $p(x | y, z)$ remains constant, but when the prior $p(y, z)$ changes based on the target task. \n\nThe authors evaluate this new method on various tasks, including standard image/text benchmarks and a medical imaging dataset on chest X-rays. Their empirical results show that their proposed method achieves better performance than standard ERM and the (approximation of a) invariance-based method.\n", "questions": "In practice, is this stronger label-shift assumption violated? Some additional discussion about this would be appreciated.\n", "limitations": "Limitations are sufficiently discussed.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The empirical results show that the method achieves better performance than standard ERM and the approximation of an invariant method that assigns a uniform prior over the spurious features on ColoredMNIST and CheXpert.  ", "weaknesses": "Their method relies on a more complex version of the label shift assumption, which must now hold on a much larger set of configurations (|Z| \\timex |Y|), which is a weakness (that the authors do indeed point out). However, it’s unclear how reasonable this assumption is to make.\n\nThere are no comparisons to other test-time adaptation methods, so the additional amount of unlabeled data may be inflating the benefits in the reported results.\n\nSome small issues with the writing: missing references in the Results and Discussion paragraphs of subsection 4.2, text in figures being difficult to read.\n\nAs a side note, I feel that the definition of spurious features used in this paper is a bit counterintuitive. These are useful features for a downstream task, which aren’t what I think of as the standard notion of spurious features.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. Response to your points is below.\n\n\n“**There are no comparisons to other test-time adaptation methods.**”\n\nAll the test-time adaptation methods we are aware of are focused on covariate shift, or label shift, but no methods (AFAIK) can handle shift in the correlation between latent factors and the class prior.  Thus these methods are unsuitable for our problem setting.\n\n“**I feel that the definition of spurious features used in this paper is a bit counterintuitive. These are useful features for a downstream task, which aren’t what I think of as the standard notion of spurious features.**”\n\nWhile it is true that “spurious” suggests the features have no predictive value for the class, in most setups (eg colored MNIST), these features (e.g. background) are in fact predictive of the class, just not in a stable way. Our observation is that we can model this change in correlation, rather than be invariant to it, and thus get better performance.  \n\n\n“**In practice, is this stronger label-shift assumption violated?**“\n\nThe fact that we get better performance on a variety of datasets is empirical demonstration of the reasonableness of our assumption. However, we do not claim that it captures all (or even most) forms of distribution shift.\n\n“**The method relies on a more complex version of the label shift assumption, which must now hold on a much larger set of configurations ($|Z| \\times |Y|)$, which is a weakness (that the authors do indeed point out). However, it’s unclear how reasonable this assumption is to make.**”\n\nThis assumption is actually usually weaker than the standard label shift assumption. Specifically, if we can hold more factors of variation fixed, then we can eliminate potential confounding variables that would make the generative distribution of X (p(X | Y) or p(X | Y, Z)) differ between source and target distributions. The tradeoff is that the discriminative model that we need to learn to use our method p(Y, Z | X) may be harder to learn than p(Y | X). However, this is something we can evaluate empirically.\n"}]

Author rebuttal:
We respond to each reviewer separately below. However, a common concern was our dependence on a training set which has y and z labels. We agree this is a weakness. However, as we show in Appendix C.5 (in the supplemental material), it is possible to apply our method with a very small amount of labeled data. In particular, we get good results on the worst group benchmarks even when up to 90% of the z labels are missing. The trick is to use label imputation to infer the missing z labels before applying our algorithm. We will clarify this in the revised paper. (Note that the theoretical work in Yong'22 shows that some amount of environment labeling is required, absent further assumptions.)

L. I. N. Yong, S. Zhu, L. Tan, and P. Cui, “ZIN: When and How to Learn Invariance Without Environment Partition?,” in Advances in Neural Information Processing Systems, May 2022 [Online]. Available: https://openreview.net/forum?id=pUPFRSxfACD.


Meta-review: The rebuttal has (at least partially) addressed the reviewers' concerns and all reviewers have reached an agreement to accept the paper. Therefore, I would recommend acceptance of the paper. Please revise the paper accordingly in the camera-ready version.
Decision: Accept (poster)