Title: Beyond Invariance: Test-Time Label-Shift Adaptation for Addressing "Spurious" Correlations
Venue: NeurIPS 2023
Accepted: yes
Keywords: Distribution shift

Spurious correlation

Group robustness
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 Adaptation for Addressing 'Spurious' Correlations

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 .

1 Introduction

Machine learning systems are known to be sensitive to so-called 'spurious correlations' [Geirhos et al., 2020, Wiles et al., 2022, Arjovsky et al., 2019] between irrelevant features of the inputs and the predicted output label. These features and their associated correlations are called 'spurious' because they are expected to change across real-world distribution shifts. As a result, models that rely on such spurious correlations often have worse performance when they are deployed on a target domain that is distinct from the source domain on which they were trained [Geirhos et al., 2020, Izmailov et al., 2022]. For example, Jabbour et al. [2020] show that neural networks trained to recognize conditions like pneumonia from chest X-rays can learn to rely on features that are predictive of patient demographics rather than the medical condition itself. When the correlation between demographic

factors and the condition change (e.g., when the model is deployed on a different patient population), the performance of such models suffers.

To address this issue, recent work has focused on learning predictors that are invariant to changes in spurious correlations across source and target domains, either using data from multiple environments [e.g., Arjovsky et al., 2019], or data where the labels for the nuisance factors are available at training time [Veitch et al., 2021, Makar et al., 2022, Puli et al., 2022, Makino et al., 2022].

However, 'spurious' correlations can sometimes provide valuable prior information for examples where the input is ambiguous. Consider the example of calculating the probability that a patient has a particular disease given a positive test. It is well known that the underlying prevalence of the disease (i.e., prior probability it is present) in the patient population is highly informative for making this diagnosis; even if the test is highly sensitive and specific, the patient's probability of having the disease given a positive test may be low if the disease is rare given that patient's background [see, e.g., Bours, 2021]. A similar logic applies in many prediction problems: if we can predict, say, patient demographics from the input, and the prevalence of the target label differs between demographic groups in the current patient population, then the demographic-relevant spurious features provide prior information that supplements the information in the target-relevant features of the input. Thus, if a model could be adjusted to use spurious features optimally in downstream target domains, it could substantially out-perform invariant predictors, as we show in this paper. 2

̸

Motivated by the above, in this paper we propose a test-time approach for optimally adapting to distribution shifts which arise due to changes in the underlying joint prior between the class labels y and the nuisance labels z . We can view these changes as due to a hidden common cause u , such as the location of a specific hospital. Thus we assume p s ( u ) = p t ( u ) , where p s is the source distribution, and p t is the target distribution. Consequently, p i ( y, z ) = ∑ u p ( y, z | u ) p i ( u ) will change across domains i . However, we assume that the generative model of the features is invariant across domains, so p i ( x | y, z ) = p ( x | y, z ) . See Figure 1 for an illustration of our modeling assumptions.

The key observation behind our method is that our assumptions are equivalent to the standard 'label shift assumption', except it is defined with respect to an expanded label m = ( y, z ) , which we call the meta-label. We call this the 'expanded label shift assumption'. This lets use existing label shift techniques, such as Alexandari et al. [2020], Lipton et al. [2018], Garg et al. [2020], to adapt our model using a small sample of unlabeled data { x n ∼ p t ( x ) } from the target domain to adjust for the shift in the prior over meta-

U

Y

Z

X

labels, as we discuss in Section 3.2. Importantly, although our approach relies on the assumption that p ( x | y, z ) is preserved across distribution shifts, it is based on learning a discriminative base model p s ( y, z, | x ) , which we adjust to the target distribution p t ( y | x ) , as we explain in Section 3.1. Thus we do not need to fit a generative model to the data. We do need access to labeled examples of the confounding factor z at training time, but such data is often collected anyway (albeit in limited quantities) especially for protected attributes. Additionally, because it operates at test-time, our method does not require retraining to adapt the base model to multiple target domains. We therefore call our approach Test-Time Label Shift Adaptation (TTLSA).

We evaluate TTLSA on various standard image and text classification benchmarks, as well as the CheXpert chest X-ray dataset. We show that in shifted target domains TTLSA outperforms ERM (empirical risk minimization, which often uses spurious features inappropriately in target domains), as well as methods that train invariant classifiers (which ignore spurious features). Note, however,

that it has been shown that no single adaptation method can work under all forms of distribution shift [Veitch et al., 2021, Kaur et al., 2022]. Our assumptions capture certain kinds of shift, but certainly not all. In particular, our method is unlikely to help with the kinds of problems studied in the domain adaptation literature, where there is covariate shift (i.e., a change from p s ( x ) to p t ( x ) ) which is not captured by a change in the distribution over the causal factors ( y, z ) .

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, Gulrajani and Lopez-Paz, 2021]), in which a model is expected to generalize (i.e., achieve acceptable performance) in a target domain without access to any data from that target domain. In this problem setup, building predictors based on the principle of invariance or minimax optimality with respect to spurious correlation shifts is a natural approach.

Many of these methods make the same modeling assumptions as we do (shown in Figure 1); this is often referred to as an anti-causal prediction setting, since the labels cause (generate) the features rather than vice versa [Schoelkopf et al., 2012, Veitch et al., 2021, Makar et al., 2022, Puli et al., 2022, Zheng and Makar, 2022]. These methods use the fact that an invariant predictor will satisfy certain conditional independencies, and employ regularizers that encourage the desired conditional independence. A related line of work tries to minimize the worst-case loss across groups (values of m = ( y, z ) ) using distributionally robust optimization [see, e.g., Sagawa et al., 2020, Liu et al., 2021, Nam et al., 2022, Lokhande et al., 2022]). Recently, Idrissi et al. [2022] showed empirically that the 'SUBG' method, which uses ERM on a group-balanced version of the data achieved by subsampling, can be an effective approach to learning worst-group-robust predictors in this setting.

Unsupervised domain adaptation and label shift In the Unsupervised Domain Adaptation (UDA) literature, we assume that unlabeled data from the target domain is available, either at training or test time. To make optimality guarantees, UDA methods require assumptions about the structure of the distribution shift. Methods are often categorized into whether they require a covariate shift assumption (i.e., that the discriminative distribution p ( y | x ) is preserved [see, e.g., Shimodaira, 2000]), or a label shift assumption (i.e., that generative distribution p ( x | y ) is preserved [Lipton et al., 2018, Garg et al., 2020]).

In the setting of Figure 1, neither the standard covariate shift nor label shift assumptions hold; however, the label shift assumption does hold with respect to the meta-label m = ( y, z ) . Our primary contribution is to show that label shift adaptation techniques, such as Alexandari et al. [2020], when applied with respect to the meta-label, can also be used to adapt to changes due to spurious correlations.

Interestingly, in the 'purely spurious' setting described in Makar et al. [2022], which is a special case of Figure 1, we can obtain a risk-invariant model as a special case, by adapting to a target distribution for which p t ( y, z ) = p s ( y ) p s ( z ) [Veitch et al., 2021, Makar and D'Amour, 2022]. 3 Based on this observation, we show that logit adjustment [Menon et al., 2021], a set of test- and training-time methods developed for long-tail learning, can be repurposed to do invariant learning when applied to the meta-label m = ( y, z ) .

Methods using the expanded label shift assumption The 'generative multi-task learning' or GMTL method of Makino et al. [2022] makes the same expanded label shift assumption that we do. However, instead of estimating p t ( y, z ) from unsupervised target data, they instead assume that there is some value α such that p t ( y, z ) = p s ( y, z ) 1 -α . They state that choosing α is an open problem, and therefore they report results for a range of possible α 's. By contrast, we provide a way to estimate p t ( y, z ) , and we do not restrict it to have the above functional form.

In Jiang and Veitch [2022], they propose a method called Anti-Causal Transportable and Invariant Representation or 'ACTIR' which also makes the same expanded label shift assumption that we do, and further relaxes the assumption that Z is observed during source training. However, they require access to examples from multiple source distributions, which they use to learn a domain invariant classifier. Furthermore, then require labeled ( x, y ) examples to adapt their classifier at test time.

Test time adaptation There is a growing Test Time Adaptation (TTA) literature that explores strategies for adapting a trained model at test time using unlabeled data from the target domain. These methods are based on various heuristics to adapt discriminative models without specifying strong assumptions on the distribution shift structure. For example, TENT [Wang et al., 2021] uses entropy minimization to update the batch normalization layers of a CNN, and MEMO [Zhang et al., 2021] uses ensembles of predictions for different augmentations of a test sample. Most of the TTA literature has focused on models that work on images (e.g., by leveraging data augmentation), so these techniques cannot be applied to embeddings or other forms of input. By contrast, our method can be applied to any classifier, even non-neural ones, such as random forests. For a more comprehensive review of TTA methods, see [Liang et al., 2023].

3 Test-Time Label Shift Adaptation

Our aim is to construct a Bayes optimal predictor for the target distribution, f t ( x ) , using a large labeled dataset from the source distribution, D xyz s = { ( x n , y n , z n ) ∼ p s ( x , y, z ) } , and a small unlabeled dataset from the target distribution, D x t = { x n ∼ p t ( x ) } . The optimal prediction for the class label in the target domain is given by

where the joint posterior over the class label y and the nuisance factor z is given by

where the first step follows from Bayes' rule, and the second step follows from our causal stability assumption that p t ( x | y, z ) = p s ( x | y, z ) = p ( x | y, z ) .

The first key insight of our method is that we can use the EM algorithm (or other optimization methods) to estimate p t ( y, z ) by maximizing the likelihood of the unlabeled target dataset, as explained in Section 3.2. The second key insight is that we can estimate the likelihood p ( x | y, z ) up to a constant of proportionality by fitting a discriminative model on the source dataset, and then dividing out by the source prior:

This is known as the 'scaled likelihood trick' [Renals et al., 1994], and avoids us having to fit a generative model. We can estimate p s ( y, z | x ) using any supervised learning method. However, to get good performance in practice, we need to ensure the probabilities are calibrated, as we explain in Section 3.1. Finally, we can estimate the source prior p s ( y, z ) just by counting how often each ( y, z ) combination occurs in D xyz s and then normalizing. Combining Equation (2) with Equation (3) we can compute the posterior distribution over augmented labels from the unlabeled target data as follows:

We can then compute f t ( x ) using (1). In summary, our TTLSA method consists of the following two steps:

Train on source. Train a model p s ( y, z | x ) using supervised learning (and calibration) applied to D xyz s , as explained in Section 3.1. Also estimate p s ( y, z ) from D xyz s by counting.

Adapt to target. Estimate p t ( y, z ) using EM applied to D x t , as explained in Section 3.2. Then compute p t ( y, z | x ) using (2) and f t ( x ) using (1).

We discuss these steps in more detail below.

3.1 Fit model to source distribution

In the first step, we fit a discriminative classifier p s ( y, z | x ) on the source dataset. This is just like a standard classification problem, except we use an augmented output space, consisting of the class label y and group label z ; we denote this joint label by m := ( y, z ) .

Calibration Our adaptation procedure crucially relies on the fact that the output of the classifier, p s ( m | x ) , can be interpreted as calibrated probabilities. Since modern neural networks are often uncalibrated [see, e.g., Guo et al., 2017], we have found it important to perform an explicit calibration step. (See the supplementary for an ablation study where we omit the calibration step.) Specifically, we follow Alexandari et al. [2020] and use their 'bias corrected temperature scaling' (BCTS) method, which is a generalization of Platt scaling to the multi-class case. In particular, let l ( x ) be the vector of M logits. We then modify p s ( m | x ) as follows:

where T ≥ 0 is a learned temperature parameter, and b m is a learned bias. This calibration is done after the source classifier is trained using a labeled validation set from the source distribution.

Logit adjustment The calibration method is a post-hoc method. However, if the source distribution has a 'shortcut', based on spurious correlations between z and y , the discriminative model may exploit this by overfitting to it. The resulting model will not learn robust features, and will therefore be hard to adapt, even if we use calibration. To tackle this, we take inspiration from Proposition 1 of Makar et al. [2022], which showed that a classifier that is trained on the unconfounded or balanced distribution p b , such that p b ( y, z ) = p s ( y ) p s ( z ) , will not learn any 'shortcut' between the confounding factor z and the target label y . Such a model will therefore have a risk which is invariant to changes in the p ( z | y ) distribution, and, in 'purely spurious' anti-causal settings (a special case of Figure 1, see footnote 3), they showed that this is indeed the optimal one among all risk-invariant predictors (see discussion in Section 2). If both marginals are uniform, then p b ( y, z ) ∝ 1 , so the corresponding balanced version of the source distribution becomes p b ( m | x ) ∝ p s ( m | x ) p s ( m ) , which in log space becomes log( p b ( m | x )) ≡ log( p s ( m | x )) -log( p s ( m )) . We train our classifier to maximize this balanced objective; this is known as 'logit adjustment' [Menon et al., 2021]. After fitting, we can recover the classifier for the original source distribution using p s ( y, z | x ) ∝ p b ( y, z | x ) p s ( y, z ) . Finally we apply calibration to p s ( y, z | x ) , as explained above. (Note that when the marginals are not uniform, one can apply a similar adjustment, subtracting off log( p s ( y, z ) / ( p s ( y ) p s ( z )) .)

3.2 Adapt model to target distribution

Given some unlabeled samples from the target distribution, D x t = { x n ∼ p t ( x ) : n = 1 : N } , our goal is to estimate the new prior, p t ( y, z ) = p t ( m ) = π m . This is a standard subroutine in label shift adaptation methods [Lipton et al., 2018, Garg et al., 2020]; our innovation is to do this with respect to the meta-label m . Here, we use an EM approach [Alexandari et al., 2020], and maximize the log likelihood of the unlabeled data, L ( x ; π , θ ) = ∑ x n log p t ( x n ; π , θ ) , wrt π , where

where θ are the parameters estimated from the source distribution. This objective is a sum of logs of a linear function of π , and is maximized subject to the affine constraints π m ≥ 0 and ∑ M m =1 π m = 1 . Thus, under weak conditions on the ratios p s ( m | x ; θ ) /p s ( m ; θ ) (these implied when x contains some information that can discriminate between levels of m ) the problem is concave, with a unique global optimum, implying the parameters are identifiable [Alexandari et al., 2020].

A simple way to compute this optimum is to use EM, which automatically satisfies the constraints. Let π j m be the estimate of π m at iteration j . We initialize with π 0 m = p s ( m ) and run the following

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Shift parameter

0

.

995

0

.

996

0

.

997

0

.

998

0

.

999

1

.

000

AUC

AUC on coloredmnist

ERM

gDRO

SUBG

LA

TTLSA

TTLSA (Oracle)

(a)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Shift parameter

0

.

4

0

.

5

0

.

6

0

.

7

0

.

8

0

.

9

1

.

0

AUC

AUC on chexpert-embedding

ERM

gDRO

SUBG

LA

TTLSA

TTLSA (Oracle)

(b)

iterative procedure below:

We then set p t ( m ) = π J m . We can also modify this to compute a MAP estimate, instead of the MLE, by using a Dirichlet prior. See the supplementary for a detailed derivation.

After estimating p t ( y, z ) on D x t , we can compute p t ( y, z | x ) for the examples in D x t using Equation (4).

4 Experiments

In this section, we provide an experimental comparison of our method with various baseline methods on a variety of datasets. In particular, we compare the following methods:

ERM This corresponds to training a model on the source distribution, and then applying the same model to the target distribution without any adaptation, i.e., we assume p t ( y | x ) = p s ( y | x ) .

gDRO The group DRO method of [Sagawa et al., 2020] is designed to optimize the performance of the worst performing group. (We only use this method for the worst-group benchmarks in Section 4.3.)

SUBG The 'SUBG' method of [Idrissi et al., 2022] subsamples the data so there is an equal number of examples in each group m = ( y, z ) , then trains a classifier by standard ERM. This is a simpler alternative to gDRO yet often achieves comparable performance.

LA This is our logit adjustment method of Section 3.1, which approximates a domain invariant classifier by targeting a uniform prior on the source domain, thus avoid overfitting to spurious correlations.

TTLSA This is our EM method of Section 3.2 that adapts the LA-based classifier using unlabeled data.

Oracle This is similar to TTLSA, but we replace the EM procedure with the ground truth target meta-label prior p t ( y, z ) . This gives an upper bound on performance. (We only use this for the CMNIST experiments in Section 4.1 and the CheXpert experiments in Section 4.2, where we artificially control the degree of distribution shift.)

4.1 Colored MNIST

In this section, we apply our method to the Colored MNIST dataset [Arjovsky et al., 2019, Gulrajani and Lopez-Paz, 2021].

z

0

1

y

0

0.5

0

1

0

0.5

z

0

1

y

0

0

0.5

1

0.5

0

Dataset We construct the dataset in a manner similar to [Arjovsky et al., 2019, Gulrajani and LopezPaz, 2021]. Specifically, we create a binary classification problem, where label y = 0 corresponds to digits 0-4, and y = 1 corresponds to 5-9. We then sample a random color z ∈ { 0 , 1 } , corresponding to red or green, for each image, with a probability that depends on the target distribution. See the supplementary material for a visualizaiton of this dataset. Since the color is easier to recognize than the shape, color can act as a 'shortcut' to predicting the class label, even though this is not a robust (domain invariant) feature.

We create a set of 21 target distributions p λ = (1 -λ ) p 0 + λp 1 , where λ ∈ { 0 , 0 . 05 , . . . , 0 . 95 , 1 . 0 } , and p 0 ( y, z ) and p 1 ( y, z ) are two anchor distributions in which y and z are correlated and anticorrelated, respectively (see Table 1). By changing λ , we can control the dependence between y and z . We train the classifier on a source domain exhibiting a strong spurious correlation (we choose λ = 0 . 05 ), and then apply the model to each target domain, p λ , for λ ∈ { 0 , 0 . 05 , . . . , 0 . 95 , 1 . 0 } . We measure performance using Area Under the ROC curve (AUC).

Training procedure During training, we fit a LeNet CNN on the training set by using m = ( y, z ) as the label. We use AdamW with a batch size of 64 and a learning rate of 10 -3 for a maximum of 5000 epochs, and run calibration for a maximum of 1000 epochs with the same learning rate and batch size on a 10% holdout set. However, training typically terminates much earlier since we also calculate the exponential moving average of the validation loss with a decay rate of 0.1, and stop early when the smoothed validation loss does not decrease for 5 consecutive epochs. During evaluation, we infer the target label prior for each target distribution using EM applied to an unlabeled test set of size 64 or 512, and we then predict the class labels on this test set. Finally we repeat this across all the unlabeled minibatches in the target distribution, and report the overall AUC.

Results Figure 2a shows our results, with error bars showing standard error of the mean across 4 trials. There are three main groups of curves. (1) The ERM model (black line) shows performance that gets steadily worse as the target distribution shifts away from the source. (2) The invariant models (SUBG in dashed orange and logit adjustment in solid red) are approximately constant across domains, as expected. (In this experiment, we see that logit adjustment outperforms SUBG.) 3) Our adaptive TTLSA models (blue curves), corresponding to TTLSA with 64 or 512 samples and the oracle method, all show a U-shaped curve, which is an upper bound on all the other curves. In particular, we see that the U-shaped curve is tangent to the invariant line when the target distribution is unconfounded ( λ = 0 . 5 ); in that case, there is no prior information from p t ( y, z ) to exploit. However, in all other target domains, the adapted prior information improves performance. As we increase the size of the unlabeled target dataset, we see that the performance of TTLSA approaches that of the oracle. To illustrate the fact that our method can also be applied to non-neural net classifiers we also applied TTLSA on top a gradient boosted tree classifier. The results (shown in the supplementary) are qualitatively similar to those in Figure 2a.

4.2 CheXpert

Next, we apply our method to the problem of disease classification using chest X-rays based on the CheXpert [Irvin et al., 2019] dataset. Chest X-rays are a particularly relevant application area for our method because sensitive attributes, such as patient sex or self-reported race, can be readily predicted from chest X-rays [Gichoya et al., 2022]. Recent work, e.g., Jabbour et al. [2020] and Makar et al. [2022], has confirmed the potential for classifiers to exploit these features as spurious features.

Dataset CheXpert has 224,316 chest radiographs of 65,240 patients. See the supplementary material for a visualizaiton of this dataset. Each image is associated with 14 disease labels derived

from radiology reports, and 3 potentially confounding attributes (age, sex, and race), as listed in the supplementary. We binarized the attributes as in Jabbour et al. [2020], taking age to be 0 if below the median and 1 if above, and sex to be 0 if female and 1 if male. As for the class labels, we define class 0 as 'negative' (corresponding to no evidence of a disease), and class 1 as 'positive' (representing the presence of a disease); images labeled 'uncertain' for the attribute of interest are discarded. Following [Glocker et al., 2022], we focus on predicting the label y = 'Pleural Effusion' and use sex as the confounding variable z . See Figure 4 for some samples from the dataset. For the input features x , we either work with the raw gray-scale images, rescaled to size 224 × 224 , or we work with 1376-dim feature embeddings derived from the pretrained CXR model [Sellergren et al., 2022]. This embedding model was pre-trained on a large set of weakly labeled X-rays from the USA and India. Note, however, that the pre-training dataset for CXR is distinct from the CheXpert dataset we use in our experiments.

To compare performance under distribution shift, we created a set of 21 distributions, p λ = (1 -λ ) p 0 + λp 1 , where λ ∈ { 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 a linear logistic regression model, following Sellergren et al. [2022], due to its simplicity and its good performance. To train this, we use AdamW with a batch size of 64 and a learning rate of 10 -3 for a maximum of 5000 epochs, and run calibration for a maximum of 1000 epochs with the same learning rate and batch size on a 10% holdout set. However, training typically terminates much earlier since we also calculate the exponential moving average of the validation loss with a decay rate of 0.1, and stop early when the smoothed validation loss does not decrease for 5 consecutive epochs. When working with pixels, we use a CNN. See the supplementary for details.

Results Our results for the embedding version of the data are shown in Figure 2b. The trends are essentially the same as in the colored MNIST case. In particular, we see that our TTLSA method outperforms the invariant baselines, and both adaptive and invariant methods outperform ERM. These results for the pixel version of the data are shown in Appendix C.3 in the supplementary. We find that the relative performance of the methods is similar to the embedding case, but the absolute performance for all methods is better when using embeddings, as was previously shown in Sellergren et al. [2022].

Discussion Glocker et al. [2022] point out that embeddings derived from X-ray classification models may contain information about sensitive attributes z , such as sex and age. We confirmed this result, and were able to classify sex with an accuracy of over 95% just using logistic regression on the CXR embeddings, as shown in Table 3. Glocker et al. [2022] argue that this may be harmful, since it can cause bias in the predictions of the primary label y of interest (disease status). As a counterpoint, our results also show that, if we can predict both z and y from the embeddings, this information can be used to make optimal adjustments for target populations featuring very different demographic 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 et al., 2020], and have since been widely used in the literature on group robustness [see, e.g. Idrissi et al., 2022, on data balancing]. The four datasets are as follows.

CelebA Introduced in [Liu et al., 2015], this is an image dataset of celebrity faces. The class label y is hair color (blond / not-blond) and the group / attribute label z is sex (male / female).

Waterbirds Introduced in [Wah et al., 2011, Sagawa et al., 2020], this is an image dataset of birds synthetically pasted onto two different kinds of backgrounds. The class label y is bird type (land bird or water bird), and the group / attribute label z is background type (land or water).

MultNLI Introduced in [Williams et al., 2018], this is a dataset of sentence pairs, ( s 1 , s 2 ) , where the goal is to predict if s 1 entails s 2 . The class label y corresponds to entailment, contradiction or neutral, and the group / attribute label z indicates presence / absence of negation words.

CivilComments (CC) Introduced in [Borkan et al., 2019], This is a dataset of sentences from online forums. The class label y represents if the comment is toxic or not, and the group/attribute label z represents whether the content is related to a minority group (such as LGBT) or not.

Training procedure We use the code and hyper-parameters specified in Idrissi et al. [2022] for the ERM, gDRO, and SUBG baselines. We also use these same hyper-parameter values as ERM when fitting our own model TTLSA, which is trained to predict the augmented labels m = ( y, z ) using logit adjustment. We run the experiments on sixteen Nvidia A100 40GB GPUs. Depending on the dataset and method used, each experiment takes somewhere between 1 hour and 10 hours on an A100. In total, the experiments took 415 GPU hours.

For each method, for worst-group accuracy evaluations, we tune the model using worst-group accuracy in a validation set, whereas for average accuracy evaluations, we tune the model using average validation set accuracy. For worst-group evaluations, we use the LA baseline, without adjusting to the test set distribution. This effectively targets a balanced group distribution, similar to SUBG. For average target accuracy evaluations, we use the full TTLSA method with test-time EM adjustment.

Results We summarize our results in Table 2. We report the worst group and average group accuracy, averaged across 4 replication runs. Our worst group numbers for the baseline methods are within error bars for those reported in Idrissi et al. [2022]. 4 The Idrissi et al. [2022] paper does not report average group performance, but we computed these results by modifying their code.

Overall, the results show that TTLSA offers a unified approach to achieve a variety of robustness objectives. In terms of average group accuracy, we find that the performance of TTLSA relative to ERMdepends on the nature of the shift. For CelebA and MultiNLI, there is no significant distribution shift, so TTLSA is similar to ERM, as expected. For Waterbirds, the majority class in the source becomes much less common in the target (reflecting the fact that the rare combinations of waterbirds on land and land-birds on water become more frequent). TTLSA is able to adapt to this, and outperforms ERM. For CivilComments, the majority class becomes even more common in the target distribution. Although TTLSA can adapt to this change, it cannot match the fact that ERM has been 'rewarded' for learning a representation that is optimized for a single majority class. In terms of worst-group accuracy, the unadapted LA baseline (which is a special case of TTLSA where we do not use EM adaptation) achieves competitive performance to the dedicated gDRO and SUBG methods.

5 Conclusions and future work

We have shown that adapting to changes in the nuisance factors Z can give better results than using classifiers that are designed to be invariant to such changes. However, a central weakness of our approach is that it requires that the generative distribution p ( x | y, z ) be preserved across domains. This assumption becomes more plausible 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

data, becomes more challenging. In this vein, another weakness is that we require access to labeled examples of Z during training. In the future, we would like to relax this assumption, potentially by using semi-supervised methods [c.f., Sohoni et al., 2021, Lokhande et al., 2022, Nam et al., 2022] that combine small fully labeled datasets with large partially labeled datasets. 5 We also plan to explore the use of fully unsupervised estimates of the confounding factors Z , based on generative models, or by leveraging multiple source domains, similar to [Jiang and Veitch, 2022].

We discuss potential societal impacts in the appendix.

References

Alexandari, A. Kundaje, and A. Shrikumar. Maximum likelihood with bias-corrected calibration is hard-to-beat at label shift adaptation. In H. D. III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 222-232. PMLR, July 2020. URL https://proceedings.mlr.press/v119/ alexandari20a.html .

Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893 , 2019.

Barocas, M. Hardt, and A. Narayanan. Fairness and Machine Learning: Limitations and Opportunities . fairmlbook.org, 2019. http://www.fairmlbook.org .

Borkan, L. Dixon, J. Sorensen, N. Thain, and L. Vasserman. Nuanced metrics for measuring unintended bias with real data for text classification. In WWW , WWW '19, pages 491-500, New York, NY, USA, May 2019. Association for Computing Machinery. URL https://doi.org/10. 1145/3308560.3317593 .

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 estimation. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 3290-3300. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/2020/file/ 219e052492f4008818b8adb6366c7ed6-Paper.pdf .

Geirhos, J.-H. Jacobsen, C. Michaelis, R. Zemel, W. Brendel, M. Bethge, and F. A. Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence , 2(11):665-673, Nov. 2020. doi: 10.1038/s42256-020-00257-z. URL https://doi.org/10.1038/s42256-020-00257-z .

W. Gichoya, I. Banerjee, A. R. Bhimireddy, J. L. Burns, L. A. Celi, L.-C. Chen, R. Correa, N. Dullerud, M. Ghassemi, S.-C. Huang, et al. Ai recognition of patient race in medical imaging: a modelling study. The Lancet Digital Health , 4(6):e406-e414, 2022.

Glocker, C. Jones, M. Bernhardt, and S. Winzeck. Risk of bias in chest x-ray foundation models. Sept. 2022. URL http://arxiv.org/abs/2209.02965 .

Gulrajani and D. Lopez-Paz. In search of lost domain generalization. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=lQdXeXDoWtI .

Guo, G. Pleiss, Y. Sun, and K. Q. Weinberger. On calibration of modern neural networks. 2017.

Y. Idrissi, M. Arjovsky, M. Pezeshki, and D. Lopez-Paz. Simple data balancing achieves competitive worst-group-accuracy. In Conference on Causal Learning and Reasoning , pages 336-351. PMLR, 2022.

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. 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 on Artificial Intelligence , 33(01):590-597, July 2019. doi: 10.1609/aaai.v33i01.3301590. URL https://ojs.aaai.org/index.php/AAAI/article/view/3834 .

Izmailov, P. Kirichenko, N. Gruver, and A. G. Wilson. On feature learning in the presence of spurious correlations. In NIPS , Oct. 2022. URL http://arxiv.org/abs/2210.11369 .

Jabbour, D. Fouhey, E. Kazerooni, M. W. Sjoding, and J. Wiens. Deep learning applied to chest x-rays: Exploiting and preventing shortcuts. In F. Doshi-Velez, 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. URL https://proceedings.mlr.press/v126/jabbour20a.html .

Jiang and V. Veitch. Invariant and transportable representations for Anti-Causal domain shifts. July 2022. URL http://arxiv.org/abs/2207.01603 .

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.

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://arxiv.org/ abs/2012.07421 .

Liang, R. He, and T. Tan. A comprehensive survey on test-time adaptation under distribution shifts. arXiv preprint arXiv:2303.15361 , 2023.

Lipton, Y.-X. Wang, and A. 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 https://proceedings.mlr.press/v80/lipton18a.html .

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. In ICML , July 2021. URL http://arxiv.org/abs/2107.09044 .

Liu, P. Luo, X. Wang, and X. Tang. Deep learning face attributes in the wild. In ICCV , 2015.

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. URL http://arxiv.org/abs/2201.03668 .

Makar and A. D'Amour. Fairness and robustness in anti-causal prediction. Sept. 2022. URL http://arxiv.org/abs/2209.09423 .

Makar, B. Packer, D. Moldovan, D. Blalock, Y. Halpern, and A. D'Amour. Causally motivated shortcut removal using auxiliary labels. In AISTATS , volume 151, pages 739-766, 2022. URL https://proceedings.mlr.press/v151/makar22a/makar22a.pdf .

Makino, K. J. Geras, and K. Cho. Generative multitask learning mitigates target-causing confounding. Feb. 2022. URL http://arxiv.org/abs/2202.04136 .

K. Menon, S. Jayasumana, A. S. Rawat, H. Jain, A. Veit, and S. Kumar. Long-tail learning via logit adjustment. In ICLR , 2021. URL http://arxiv.org/abs/2007.07314 .

P. Murphy. Probabilistic Machine Learning: An introduction . MIT Press, 2022.

Nam, J. Kim, J. Lee, and J. Shin. Spread spurious attribute: Improving worst-group accuracy with spurious attribute estimation. In ICLR , Apr. 2022. URL http://arxiv.org/abs/2204.02070 .

Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, 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.

Pham, M. Cho, A. Joshi, and C. Hegde. Revisiting self-distillation, 2022.

M. Puli, L. H. Zhang, E. K. Oermann, and R. Ranganath. Out-of-distribution generalization in the presence of Nuisance-Induced spurious correlations. In ICLR , May 2022. URL https: //openreview.net/forum?id=12RoR2o32T .

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. URL http://dx.doi.org/10.1109/89.260359 .

Saerens, P. Latinne, and C. Decaestecker. Adjusting the Outputs of a Classifier to New a Priori Probabilities: A Simple Procedure. Neural Computation , 14(1):21-41, 01 2002. ISSN 0899-7667. doi: 10.1162/089976602753284446. URL https://doi.org/10.1162/089976602753284446 .

Sagawa, P. W. Koh, T. B. Hashimoto, and P. Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for Worst-Case generalization. In ICLR , 2020. URL http://arxiv.org/abs/1911.08731 .

Schoelkopf, D. Janzing, J. Peters, E. Sgouritsa, K. 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 978-1-4503-1285-1. URL https://icml.cc/2012/papers/625.pdf .

B. Sellergren, C. Chen, Z. Nabulsi, Y. Li, A. Maschinot, A. Sarna, J. Huang, C. Lau, S. R. Kalidindi, M. Etemadi, F. 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. URL http://dx.doi.org/10.1148/radiol.212482 .

Shimodaira. Improving predictive inference under covariate shift by weighting the log-likelihood function. Journal of statistical planning and inference , 90(2):227-244, 2000.

S. Sohoni, M. Sanjabi, N. Ballas, A. Grover, S. Nie, H. Firooz, and C. Ré. BARACK: Partially supervised group robustness with guarantees. 2021.

Veitch, A. D'Amour, S. Yadlowsky, and J. Eisenstein. Counterfactual invariance to spurious correlations in text classification. In NIPS , Nov. 2021. URL https://openreview.net/forum? id=BdKxQp0iBi8 .

Wah, S. Branson, P. Welinder, P. Perona, and S. Belongie. The Caltech-UCSD birds-200-2011 dataset. Technical report, 2011.

Wang, E. Shelhamer, S. Liu, B. Olshausen, and T. Darrell. Tent: Fully test-time adaptation by entropy minimization. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=uXl3bZLkr3c .

Wiles, S. Gowal, F. Stimberg, S.-A. Rebuffi, I. Ktena, K. D. Dvijotham, and A. T. Cemgil. A finegrained analysis on distribution shift. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=Dl4LetuLdyK .

Williams, N. Nangia, and S. Bowman. A Broad-Coverage challenge corpus for sentence understanding through inference. In NAACL , pages 1112-1122, New Orleans, Louisiana, June 2018. Association for Computational Linguistics. URL https://aclanthology.org/N18-1101 .

Yi, R. Wang, J. Sun, Z. Li, and Z.-M. Ma. Breaking correlation shift via conditional invariant regularizer. In The Eleventh International Conference on Learning Representations .

Zhang, S. Levine, and C. Finn. Memo: Test time robustness via adaptation and augmentation. arXiv preprint arXiv:2110.09506 , 2021.

Zheng and M. Makar. Causally motivated multi-shortcut identification and removal. In NIPS , Oct. 2022. URL https://openreview.net/forum?id=-ZQOx6yaVa-.

A Derivation of the EM algorithm

In this section we describe how to estimate the label prior on the target distribution, p t ( y, z ) = p t ( m ) = π m , using the unlabeled data D x t . There are several approaches to this, including a moment matching method called black box shift learning [Lipton et al., 2018] and an MLE approach based on the EM algorithm [Saerens et al., 2002]. In [Alexandari et al., 2020], they show that the MLE approach is much better, provided the classifier is calibrated. (See also [Garg et al., 2020] for a unified analysis of these two approaches.)

Since our augmented label space is expanded to include both class labels y and meta-data z , the number of labels M can be large, which can result in problems when computing the MLE. We therefore expand the previous approach to compute the MAP estimate, using a Dirichlet prior of the form

where B ( α ) is the normalization constant. Note that the MLE solution can be recovered by setting α = 1 , which represents a uniform prior.

The goal is to maximize the (unnormalized) log posterior of π given the unlabeled target data X :

The first term can be rewritten as

This objective is a sum of logs of a linear function of π , as is the log prior. This needs to be maximized subject to the affine constraints π m ≥ 0 and ∑ M m =1 π m = 1 , so the problem is concave, with a unique global optimum [Alexandari et al., 2020].

One way to compute this optimum is to use EM. Let π j be the estimate of π at iteration j ; we initialize with π 0 m = p s ( m ) . First note that

Hence the complete data log posterior is given by

so the expected complete data log posterior is

where we drop constants wrt π , and where we defined the expected counts to be

Hence in the E step we just need to compute the posterior responsibilities for each label:

We plug this into Equation (22) and then maximize Equation (21), using a Lagrange multiplier to enforce the sum to one constraint. We then get the following (see e.g., Sec 4.2.4 of Murphy [2022] for the derivation):

where ˜ N j m are the prior pseudo counts plus the expected empirical counts:

If we assume that the class label prior is constant, and only the distribution of auxiliary labels has changed, then we can write

where

However, we do not make this fixed label assumption in our experiments.

At convergence, we have

UPAISHT

B Datasets

In this section we discuss the datasets in more detail.

B.1 Colored MNIST

We show some sample images in Figure 3.

B.2 CheXpert

We show some sample images in Figure 4. We list all the target attributes in Table 3. To test the difficult of each task, we train a logistic regression model for each attribute on the 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 a poor measure of the predictive performance.

Interestingly, we see that we can predict sex with an AUC of 0.973, which is higher than the AUC for effusion (0.861). To understand why, note that we only use frontal scans; consequently breasts are often visible in female patients, and this is often easier to detect visually than detecting the disease itself (see Figure 4), providing a possible 'shortcut' for models to exploit.

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Shift parameter

0

.

90

0

.

92

0

.

94

0

.

96

0

.

98

1

.

00

AUC (without calibration)

TTLSA (oracle)

TTLSA (batch size 35000)

ERM

C Extra results

In this section, we include some extra experimental results.

C.1 Colored MNIST using gradient boosted tree classifier

In Figure 5, we show the results of various methods on the Colored MNIST dataset, where we use a Gradient Boosting Classification Tree as our base classifier, instead of a DNN. In particular, we use the HistGradientBoostingClassifier from scikit-learn [Pedregosa et al., 2011] with default parameters. The results are qualitatively similar to the DNN case.

C.2 The benefits of calibration

In Figure 6 we show the results on CheXpert if we remove the calibration step for our base classifier. Compared to Figure 2b, we see that the overall AUC of all the methods is worse, and the variance is larger. 6 However, the rank ordering of the methods is the same. It is notable that a large gap opens up between the Oracle curve and the TTLSA implementations. This suggests that calibration primarily improves estimation of p t ( y, z ) estimation via EM, because the Oracle curve in this subfigure corresponds to using the correct weights with the uncalibrated p s ( y, z | x ) model.

C.3 CheXpert using CNN on raw pixels

In Figure 7 we show the result of various methods when applied to CheXpert images, as opposed to using embeddings. We use a ResNet-50 that was pretrained on Imagenet, which we then fine tune on CheXpert images by replicating the gray-scale image along all 3 RGB channels. The qualitative conclusions are the same as in the embedding case.

C.4 More results on the benchmark datasets

In Table 4 and Table 5 we report the per-group accuracy on the benchmark datasets.

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Shift parameter

0

.

70

0

.

75

0

.

80

0

.

85

0

.

90

0

.

95

1

.

00

AUC (without calibration)

TTLSA (oracle)

TTLSA (batch size 512)

TTLSA (batch size 64)

Logit Adjustment

SUBG

ERM

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Shift parameter

0

.

70

0

.

75

0

.

80

0

.

85

0

.

90

0

.

95

1

.

00

AUC

TTLSA (oracle)

TTLSA (batch size 512)

TTLSA (batch size 64)

Logit Adjustment

SUBG

ERM

Group label

(

Y, Z

)

Data

Method

(0, 0)

(0, 1)

(1, 0)

(1, 1)

(2, 0)

(2, 1)

CelebA

ERM

86.71 (0.67)

92.65 (0.79)

96.80 (0.18)

80.83 (1.46)

gDRO

92.71 (0.05)

92.33 (0.09)

92.63 (0.35)

87.36 (0.47)

SUBG

91.76 (0.22)

91.91 (0.55)

90.96 (0.33)

87.22 (1.38)

LA

91.43 (0.11)

94.77 (0.08)

95.88 (0.12)

84.72 (0.58)

TTLSA

97.59 (0.18)

98.78 (0.03)

80.70 (1.22)

51.25 (0.27)

Waterbirds

ERM

98.95 (0.13)

86.57 (0.48)

86.02 (0.13)

95.76 (0.16)

gDRO

93.46 (0.22)

88.00 (0.88)

90.15 (0.10)

92.06 (0.34)

SUBG

90.76 (0.69)

88.96 (0.19)

91.28 (0.35)

91.36 (0.24)

LA

94.43 (1.63)

88.38 (0.36)

91.32 (0.43)

93.15 (0.54)

TTLSA

94.59 (0.44)

93.68 (0.73)

95.72 (0.29)

97.12 (0.19)

ERM

80.75 (0.79)

94.94 (0.11)

83.18 (0.47)

78.05 (1.31)

81.98 (0.48)

68.60 (0.40)

gDRO

80.36 (0.63)

85.27 (0.25)

82.48 (0.59)

81.21 (1.30)

79.39 (1.34)

76.87 (1.28)

SUBG

69.63 (0.17)

82.85 (0.19)

74.39 (0.21)

79.68 (0.17)

69.84 (0.48)

68.40 (1.33)

LA

81.63 (1.15)

87.79 (1.99)

84.36 (0.89)

80.95 (2.03)

78.77 (1.09)

76.33 (1.45)

TTLSA

80.24 (0.87)

94.74 (0.58)

81.73 (2.45)

73.90 (1.72)

82.40 (1.49)

63.76 (2.15)

CivilComments

ERM

92.23 (0.42)

90.38 (0.46)

68.57 (1.07)

68.32 (0.97)

gDRO

83.94 (0.70)

79.92 (0.33)

80.97 (0.63)

81.09 (0.42)

SUBG

79.79 (0.56)

79.14 (0.34)

82.52 (0.46)

76.56 (0.25)

LA

84.45 (0.16)

79.27 (1.17)

83.00 (0.95)

84.20 (0.99)

TTLSA

85.53 (1.35)

74.94 (1.96)

84.34 (2.36)

84.61 (2.21)

C.5 Training with partial group labels

In this section, we evaluate an extension of our method where not all training samples have group labels z . In particular, we first train an ERM model to predict z on samples with group labels z , calculate p ( z | x ) for training samples with missing z , and then fit a new p ( y, z | x ) model on the augmented data. In particular, we represent each ( y, z ) target as a one-hot vector when z is known, and use a soft (predicted) encoding when z is unknown. We train with cross entropy loss. The use of soft labels may have the benefits of self-distillation Pham et al. [2022]. The validation set is always fully labeled for the purpose of hyperparameter tuning.

The results (on the 4 benchmark datasets) are shown in Table 6. The accuracy barely drops as missingness increases, which means our method is robust to the deficiency in group labels z .

D Potential negative societal impacts

The proposed method in this work yields a model that can adapt to a new distribution and improves the performance at test time by exploiting spurious correlations to create a label shift correction technique that adapts to changes in the marginal distribution p ( y, z ) using unlabeled samples from the target domain. In this way, there are potential societal benefits to our method, especially when z corresponds to a socially salient attribute, such as a protected class. However, use cases of this type require caution, especially given the limitations discussed in Section 5. Further, as we discuss in a footnote in the main text, our method does not address concerns about cases where making decisions on the basis of z is discouraged or forbidden for a priori reasons. Given these limitations, there is a potential that the existence of adaptation methods of this type could be used to downplay the potential dangers of misusing sensitive information in machine learning systems. Here, we hope researchers and practitioners will instead acknowledge that, while beneficial use cases of z information exist, (1) there is a need to validate empirically that a particular use of z information is actually socially beneficial, and (2) there are valid reasons why one might want to avoid using z information altogether. Further, there is a potential risk that if the measurement quality of the labels y, z shift across distributions, such that they measure distinct concepts, or exhibit substantially different noise properties (i.e., become biased, or exhibit more outliers), our framework might absorb them during adaptation and eventually the outcomes of the system might be biased as well.

E Invariance Equivalences and Conditions

In this section, we review connections that have been established between risk invariance, ERM on balanced data, 'separation' between a predictor f ( X ) and the spurious factor Z , and worst-( y, z ) -group performance. These results are useful for understanding why the application of logit adjustment at training time often yields a predictor that exhibits approximately invariant risk across the test sets that we study in our experiments.

E.1 Key Concepts

Risk invariance A predictor is risk-invariant with respect to a loss function ℓ and a family of test distributions Q iff it has the same risk E Q [ ℓ ( f ( X ) , Y )] for each Q ∈ Q . The results we discuss apply to test distribution families that preserve both the generative distribution and the label distribution of the source distribution; that is, Q is the set of distributions such that Q ( Y ) = P ( Y ) and Q ( X | Y, Z ) = P ( X | Y, Z ) for each Q ∈ 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'Amour [2022], and is called a 'causally compatible' family in Veitch et al. [2021], or a correlation shift in Yi et al..

Pure spuriousness The data generating process in Figure 1 is purely spurious if there exists some sufficient statistic e ( X ) such that (1) Y ⊥ X | e ( X ) and (2) e ( X ) ⊥ 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 on e ( X ) .

Veitch et al. [2021] coined the term 'purely spurious' in a context of a full counterfactual model of data generation, to refer to data generating processes where the portions of X that are causally related to Y and Z can be separated in a specific sense. Makar et al. [2022] consider the special case of pure spuriousness in the context of the anti-causal model in Figure 1. (They do not use the term 'purely spurious' as the work in Veitch et al. [2021] was concurrent; Makar and D'Amour [2022] makes the connection explicit.) Here, we use formalism from Makar et al. [2022] to present the idea to minimize conceptual overhead.

Note that when the data X is rich, such as images are long passes of text, pure spuriousness is more plausible (or a better approximation to reality) because there is less possibility of descructive interfecence between Y and Z in the generation of X . Specifically, the simplest examples where pure spuriousness fails are ones where X is very low-content: e.g., Y and Z are binary, and X := Y OR Z .

Separation Separation is a concept popularized in the literature on ML fairness [Barocas et al., 2019, Chapter 3], which stipulates that the predictor f ( X ) should satisfy the the conditional independence f ( X ) Z ⊥ | Y . When Z is a sensitive attribute, this condition stipulates that the predictor f ( X ) should contain no more information about Z than one could glean from knowing Y alone.

Data balancing Idrissi et al. [2022] study predictors trained on data subsampled so that the ( Y, Z ) distribution is uniform; they call this data-balancing. Makar et al. [2022] and Makar and D'Amour [2022] study a similar predictors optimized on a similar 'ideal' distribution, where Q ( Y, Z ) = P ( Y ) P ( Z ) for some source distribution P . This distribution does not 'balance' the marginals of Y and Z , but it eliminates the marginal correlation between Y and Z .

Worst group performance Sagawa et al. [2020] define groups in terms of ( z, y ) values. The group conditional risk is R z,y = E Q [ ℓ ( f ( X ) , Y ) | Z = z, Y = y ] . Note that for all families of test sets that we consider, the group-conditional risks are equal for all Q. Worst group risk minimization attempts to minimize the group conditional risk of the worst subgroup. Saerens et al. [2002] propose a distributionally robust optimization algorithm for performing this minimization.

E.2 Connections

In the purely spurious setting, there are several connections and near-equivalences between risk invariance, separation, optimality on balanced data, and worst group risk minimization.

Yi et al. establish that for label distribution preserving target families, a predictor f ( X ) that satisfies separation f ( X ) ⊥ Z | Y will have invariant risk across the family Q defined above. Notably, this result does not require pure spuriousness.

Under pure spuriousness, the separation condition achieves a certain optimality. Veitch et al. [2021], Theorem 4.3 establishes that in the purely spurious case, the minimax optimal across the family Q satisfies separation f ( X ) ⊥ Z | Y . Similarly, under pure spuriousness, Makar and D'Amour [2022], Proposition 2, establishes that the optimal risk-invariant predictor satisfies separation.

Interestingly, this result establishes a connection between optimality under balanced data, separation, and optimal risk invariance. Specifically, Makar et al. [2022], Proposition 1 establishes that the optimal model for the 'ideal' uncorrelated distribution for which Q ( Y, Z ) = P ( Y ) P ( Z ) achieves risk invariance across the family Q . Thus, minimizing risk under a separation constraint targets a similar predictor to the predictor that one would target simply optimizing on balanced data. Makar and D'Amour [2022] shows that the near-equivalence holds up empirically, such that learning algorithms targeted at efficiently learning the optimal predictor on balanced data can satisfy both risk invariance and separation criteria.

Idrissi et al. [2022] establish that, empirically, models trained to minimize risk on balanced data also yield favorable worst-group performance, showing that subsamping can be particularly effective. Sagawa et al. [2020] explore similar ideas, focusing on reweighting strategies, which both they and Idrissi et al. [2022] find to work relatively poorly with neural models in the data regimes they study. Sagawa et al. [2020] further establish that under certian convexity conditions, there does exist a reweighting of the data that optimizes worst-group performance, but provide a counterexample showing that this is not always the case with non-convex losses.

Based on the above results, in the purely spurious case, one can establish the following, for Q with a uniform distribution on Y :

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 ) ⊥ Z | Y .

For all Q ∈ Q , the group-specific risks are equal within labels, i.e., E Q [ ℓ ( f ∗ ( X ) , Y ) | Y = y, Z = z ] = E Q [ ℓ ( f ∗ ( X ) , Y ) | Y = y, Z = z ′ ] for all y .

The latter fact does not imply that f ∗ ( X ) also optimizes worst-group risk, but it does imply that the worst group cannot be the worst due to a spurious correlation between Y and Z . This is because, for a fixed label value y , the risks of ( y, z ) subgroups are the same.

Peer reviews:
comments: 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. 

The 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.

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.

There 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.

Some small issues with the writing: missing references in the Results and Discussion paragraphs of subsection 4.2, text in figures being difficult to read.

As 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.

QUESTIONS:
In practice, is this stronger label-shift assumption violated? Some additional discussion about this would be appreciated.

LIMITATIONS:
Limitations are sufficiently discussed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-25

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n**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\nWEAKNESSES:\n1. 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\nQUESTIONS:\n1. 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\nLIMITATIONS:\nThe authors have very adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n**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.\n\nSTRENGTHS:\nStrengths 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.\n\nWEAKNESSES:\nMain 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.\n\nQUESTIONS:\n* (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.\n\nLIMITATIONS:\nSame as mentioned weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- 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 (??).\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nIn [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\nLIMITATIONS:\nThe proposed method requires group information in the training dataset, which can be costly.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}