Ask Your Distribution Shift if Pre-Training is Right for You
Authors: Benjamin Cohen-Wang
Joshua Vendrow
Aleksander Madry
Venue: Accepted by TMLR ()

Abstract: Pre-training is a widely used approach to develop models that are robust to distribution shifts. However, in practice, its effectiveness varies: fine-tuning a pre-trained model improves robustness significantly in some cases but *not at all* in others (compared to training from scratch). In this work, we seek to characterize the failure modes that pre-training *can* and *cannot* address. In particular, we focus on two possible failure modes of models under distribution shift: poor extrapolation (e.g., they cannot generalize to a different domain) and biases in the training data (e.g., they rely on spurious features). Our study suggests that, as a rule of thumb, pre-training can help mitigate poor extrapolation but not dataset biases. After providing theoretical motivation and empirical evidence for this finding, we explore two of its implications for developing robust models: (1) pre-training and interventions designed to prevent exploiting biases have complementary robustness benefits, and (2) fine-tuning on a (very) small, non-diverse but *de-biased* dataset can result in significantly more robust models than fine-tuning on a large and diverse but biased dataset.

Published in Transactions on Machine Learning Research (01/2025)
Ask Your Distribution Shift if Pre- Pre-ninnning is Right for You
Benjamin Cohen- Cohen-Wang
bencoo@mil.ccmm edu
Massachusetts Institute of Technology
Joshua Vendrow
jvendrow@miil. edu
Massachusetts Institute of Technology
Aleksander Mądry
madry@mit edu
Massachusetts Institute of Technology
Reviewed on OpenReview: https: // openreview. net/ forun? id= edULLIVRoC
Abstract
Pre-training is a widely used approach to develop models that are robust to distribution shifts. However, in practice, its effectiveness varies: fine-tuning a pre-trained model improves robustness significantly in some cases but not at all in others (compared to training from scratch). In this work, we seek to characterize the failure modes that pre-training can and connot address. In particular, we focus on two possible failure modes of models under distribution shift: poor extrapolation (e.g. they cannot generalize to a different domain) and biases in the training data (e.g., they rely on spurious fcatures). Our study suggests that, as rule of thumb, pre-training can help mitigate poor extrapolation but not dataset. biases. After providing theoretical motivation and empirical evidence for this finding, we explore two of its implications for developing robust models: (1) pre-training and interventions designed to prevent exploiting biases have complementary robustness benefits, and (2) fine-tuning on a (very) small, non-diverse but de-biased dataset can result in significantly more robust. models than fine-tuning on a large and diverse but biased dataset."
1 Introduction
A common paradigm for developing machine learning models is pre-training them on a large, diverse dataset (c.g., ImageNet (Deng et al., 2009), JFT-300M (Sun et al., 2017), LAION-5B (Schuhmann et al., 2022)) and then fine-tuning them on lask-specific data. Indeed, compared training from scratch, fine-tuning pre- trained model often significantly improves performance and reduces computational costs (Sharif Razavian et al., 2014; Sun et al., 2017; Kornblith et al., 2019).
Yet another benefit that pre-training may offer is distribution shift robusiness. Specifically, machine learning models tend to suffer from distribution shifts, i.e., changes between the reference distribution used to develop the model and the shifted distribution that the model actually encounters when deployed. For example, a tumor identification model trained on tissue slide images from one hospital might perform poorly when deployed at another hospital (Bandi et al., 2018; Koh et al., 2020). Notably, different models (with different architectures, hyperparameters, etc.) tend be similarly sensitive a given distribution shift. However, models pre-trained on auxiliary data and then fine-tuned the reference distribution can break this trend, exhibiting substantially higher performance the shifted distribution than models trained from scratch with the same performance on the reference distribution (Taori et al., 2020; Miller et al., 2020; 2021; Andreassen et al., 2021; Wortsman et. al., 2021).
These robustness benefits pre-training are promising, but they are not universal. In particular, fine- tuming the same pro-trained model can yield significant robustness gains on some distribution shifts but not
"Cdee available https: http:tttw..o..aaaaaaaaapeettnnnn
shift robustness

---

Published in Transactions on Machine Learning Research (01/2025)
others (Section 3). Would solution to attain robustness to the latter shifts then be to fine-tune larger model pre-trained on more data? Or are there fundamental limitations to the robustness that pre-training can provide? To answer these questions, we would like to develop a more fine-griined understanding of when pre-training can improve robustness. Specifically, we ask:
Can we identify and characterize the foilure modes that pre-training con and cannot address?
Recall that under distribution shift, models can fail in number of ways. One of them is their inability to extrapolate effectively outside of the reference distribution (Gulrajani & Lopez-Paz, 2020; Koh et al., 2020). If, for instance, a model is trained only on photos taken during the day, then it might fail when deployed on photos taken at night.
Models can also underperform even when the shifted distribution does not contain anything "new." In par- ticular, they can fail due to biases in the reference distribution. For example, if a certain feature is spurionsly correlated with the label in the reference distribution, a model might learn exploit this relationship and fail on examples encountered during deployment where it does not hold (Arjovsky et al., 2019; Geirhos et al., 2020).
1.1 Our contributions
To identify the failure modes that pre-training address, we study the robustness benefits of pre-training under two types of distribution shifts: (1) shifts where extrapolation is necessary and (2) shifts where extrapolation is not needed. We start by analyzing simple logistic regression setting and illustrate why pre-training might. improve robustness to the former lype of shift, but not the latter (Section 4). We subsequently build on this intuition by measuring the robustness benefits of pre-training on syuthetic and natural distribution shifts of cach type (Section 5). Our results suggest the following rule of thumb: pre- training can help with extrapolation, but does not address other failures, for example, those stemming from dataset biases.
Implications for developing robust models. Guided by this rule of thumb, we explore two related avenues for harnessing pre-training to develop robust models.
1. Combining pre-training with interventions designed to handle bias (Section 6): There number of robustness interventions specifically designed to mitigate biascs present in training dataset (Byrd & Lipton, 2019; Sagawa el al., 2020a; Liu et al., 2021; Kirichenko et al., 2022; Idrissi et al., 2022). Our findings suggest that pre-training and this kind of intervention address two different sources of failures (the former helping with extrapolating and the latter with avoiding datasct biascs) and thus may be viewed as complementary. We indeed find that combining them can yield models with both sets of benefits.
2. Curaling dalasels for fine-luning (Section 7): One possible intervention that aims to address dataset biases is curating a de-biased dataset. In general, however, the de-biasing process might be pro- hibitively expensive. That said, we find that if we leverage pro-training help with extrapolation, we might only need a small, non-diverse fine-tuning dataset; such dalaset might actually be fea- sible to de-bias. For example, we demonstrate that fine-tuning on carefully de-biased hair color classification dataset with only 64 examples yields greater robustness than fine-tuning the entire CclebA dataset (Liu et al., 2015).
2 Background
Fine-tuning a pre-trained model. Methods for fine-tuning a pre-trained model vary: two common strategies are full fine-tuning, in which one continues training the entire model, and linear probing, in which one only finc-tunes the final layer. Some recent pro-trained models with natural language supervision (c.g.. CLIP (Radford et al., 2021), ALIGN (Jia et al., 2021)) can also be adapted to downstream task in a zero-shot context (i.e., without fine-tuning) by specifying the task through text description. In this work,
2

---

Published in Transactions on Machine Learning Research (01/2025)
1.06%
ImageNet Sketch accuracy
Image/NetVP accuracy
85
80
80
70
60
75
50
70
40
65
30
60
55
20
70
75
80
85
70
75
80
85
ImageNet accuracy
ImageNet accuracy
Baseline
Baseline linear fit
Pre-trained
y=x
Figure 1: The robustness benefits of pre-training vary. On the ImageNet-V3. distribution shift (left), different pre-trained models exhibit very little effective robustness (ER), i.e., little improvement over the linear trend models trained from scratch (see Section 2). Meanwhile, on the ImageNet Sketch distribution shift (right), some of these pre-trained models exhibit substantial effective robustness. We report average effective robustness with 95% confidence interval in the top left of each plot.
locus on the full fine-tuning strategy, which typically outperforms linear probing and zero-shot models the reference distribution. We also will sometimes consider linear probing zero-shot adaptation followed by full finc-tuning; this can in some cases improve over full fine-tuning alone in terms of robustness and performance (Kumar et al., 2022). We discuss other fine-tuning strategies in Appendix D.1.
Mcasuring robustness. For many distribution shifts, different models trained from scratch on the refer- ence distribution exhibit similar degrees of robustness to the shift. Specifically, when varying architectures, hyperparameters and training methods there is often a strong linear relationship between the reference ac- curacy and shifted accuracy? (i.c., the accuracics on the reference and shifted distributions, respectively) (Taori et al., 2020; Miller al., 2020; 2021). This relationship, dubbed accuracy on the line, can be visu- alized by plotting shifted accuracies against reference accuracies and finding linear fit.. When this linear trend is strong (i..,, shifted accuracies are highly correlated with reference accuracies), one predict the shifted accuracy of models trained from scratch from their reference accuracy. Furthermore, to quantify the robustness of a model trained with robustness intervention beyond the "baseline" models trained from scratch, one can mcasure the amount by which its shifted accuracy excoods the lincar fit's prediction, a metric known as effective robusiness (ER) (Taori al., 2020). In this work, we choose study distribution shifts for which accuracy on the line holds (i.e., the linear fit is strong) and quantify robustness by computing effective robustness (see, e.g., Figure 1). See Appendix B.1.2 for additional details.
3
The Robustness Benefits of Pre- Training Vary
Our investigation motivated by the following observation:
Pre-training can significantly improve robustness to some distribution shifts but not others.
To illustrate this, we consider two distribution shifts of ImageNet (Deng et al., 2009): ImageNet- V2 (Recht et al., 2019) and ImageNet Sketch (Wang et al., 2019). For each of these shifts, we measure the effective robustness (see Scction 2) of various pro-trained models. Specifically, we first establish bascline for robust- ness by evaluating 78 models trained from scratch ImageNet. (from PyTorch Image Models (Wightman, 2019)). We observe strong linear relationship between their reference and shifted accuracies (see Figure 1). Next, we evaluate 55 pre-trained models that are fine-tuned on ImageNet (also from Pyroh Image
For linear relationship, accuracies are probit- scaled (transformed by the inverse Gaussian CDF).
3

---

Published in Transactions on Machine Learning Research (01/2025)
Positive example
Negative example
Model decision boundary
Optimal decision boundary
44
4
4
(a) Reference dataastt in subspace Wa
(b) Models trained from different initializations
(c) Shift inside of W models equally robust
(d) Shift outside of W
robustness varies
Figure 2: Illustration of logistic regression setting. (a) Consider reference dataset that lies within a subspace Wio of Rd (b) Models trained from different initializations all learn the same (optimal) decision boundary in Wiet, but may behave differently outside of Wiet. (c) Under shifts within Wiet: models with dif- ferent initializations are equally robust. (d) Under shifts outside of Wret, initialization can affect robustnoss.
Models) and measure the improvements in shifted accuracy over the linear trend. See Appendix B.2 for the exact sctup.
We find that while some of the pre-trained models exhibit substantial effective robustness on ImageNet Sketch, they all exhibit very little effective robustness on ImageNet-Ve These pre-trained models represent a wide variety model architectures, pre-training datasets and pre-training algorithms-the largest model has 1 billion parameters and is pre-trained on a dataset of billion image-lext, pairs. Yet, the highest effective robustness attained by any of these models on ImageNet-V V2 is just 1.80%. This suggests that pre-training alone might not sulfice to address certain types of failures that occur under distribution shift. We would like to better understand this limitation; we identify and characterize these types of failures?
4 Studying Pre- Training in a Logistic Regression Setting
Our central goal is to understand the failure modes that pro-training and cannot address. To this end, we first study the robustness benefits of pre-training in simple logistic regression setting (see Figure 2).
Setup. Suppose that we are given access to reference dataset S of input-label pairs, each consisting of d-dimensional input x e gd and binary label y [}, We are concerned with finding weights w gd that minimize the (standard) logistic loss on Steft
Lee((W)
Σ ESww
log(1 + e e- w"
(1)
We assume that the reference dataset Stet satisfies the following conditions:
Inputs Ster lie within a k-dimensional (with k < d) subspace Wrat of Rd, Intuitively, this condition corresponds features lacking certain variation in the reference dataset.
2. The logistic loss Lref has a minimum value. This condition ensures that minimizing L,ef is well-defined. Note that there may be multiple weights that attain this minimum value.
Starting with initial weights wait (which, in our case, are either random or the result of pre-training), suppose that we use gradient descent minimize [w) We would like to understand how well the resulting model performs under distribution shift. In particular, what role does pre-training play through Wint? To answer this question, establish the following theorem (proof in Appendix A):
4

---

Published in Transactions on Machine Learning Research (01/2025)
In-support shift
Out-of-support shift
Label
Label
Label
Label
None
Less
Setting
Setting
More
All
Figure Examples of in-support and out-of-support shifts. One example of in-support shift (left) is shift in which the indoor/outdoor frequencies of animal appearances change, but the possible combinations of animal and sctting remain the same. An example of an out-of-suppor shift (right) is a shift from day to night: the nighttime setting is entirely novel.
Theorem 4.1. Suppose that we start with initial weights Wine E RRd and run gradient descent to minimize Lre((w). With an appropriately chosen learning rate, gradient descent converges to weights that minimize Trof- Purthermore, can be wrillen as
1 = Wre + u rdf
Winst-
(2)
Here, Whet properly of the reference datasel Sraf and lies within the reference subspace Wref- Meanswhile, projw Wist is the component of Winst that is orthogonal Wref-
Theorem 4.1 implies that there are multiple weights that attain the minimum value of Lyef, and that the initial weights Winit determine which of them we learn. Specifically, we can decompose the learned weights à into two terms: wiy and projw eut Winis- Notice that the first term is just a property of the reference dataset and in the reference subspace Wzet, while the second term depends on Wintt and is orthogonal to Wret- As result, the reference dataset itself fully specifies the model's behavior on inputs Weets while the initialization determines how the model extends outside of Wref. Consequently, changing a modcl's initialization (e.g., with pre-training) can affect performance outside of Wref, but not within Wref.
This observation gives rise to an intuition that will guide investigations in the remainder of this work: pre-training can improve robustness to a distribution shift only when the shifted distribution contains "out- of-support" inpuls, that is, inputs that could not be reasonably sampled from the reference dist In other words, pre-training helps specifically with extrapolation outside of the reference distribution.
5 Exploring the Empirical Robustness Benefits of Pre- Fre-Training
In Scction 4, we found that in a simple logistic regression setting, pre-training helps specifically with ex- trapolation. We now want to assess whether this principle holds more broadly. To do so, we measure the robustness benefits of pre-training under two types shifts: in-support shifts, where models cannot fail due to poor extrapolation (but might fail for other reasons, c.g., datasct biases), and out-of-support shifts, where models can fail due to poor extrapolation (see Figure 3). We begin by describing these two types of shifts in more detail and providing intuitions for why pre-training might improve robustness to out-of-support shifts, but not in-support shifts.
In-support shift. A distribution shift is in-support if any input that could be sampled from the shifted distribution could also be reasonably sampled from the reference distribution. In other words, the shifted distribution does not contain anything "new"; however, an in-support shift can still cause failures if, for example, the reference distribution biased. To illustrate this failure mode, consider a cat vs. dog image classification task in which photos are either taken indoors or outdoors. Suppose that in the reference distribution 90% of cats appear indoors and 90% of dogs appear outdoors (i.e., the setting is spuriously correlated with the animal). A model trained this distribution would likely rely (at least in part)
5

---

Published in Transactions on Machine Learning Research (01/2025)
In-support shifts
Out-of-support shifts
Spurious tint shift
Label shift
Unseen tint shift
Flip shift
Referance
Shifted
Reference
Shifted
Reference
Shifted
Reference
Shifted
ER = 0.34 0.87%
80 ER 0.23 0.82%
ER 5.93 0.80%
80
ER 13.42 = 0.74%
Shitted accuracy
80
Shifted accuracy
70
60
60
60
50
30
60
70
80
50
60
70
80
50
60
70
80
50
60
70
80
Reference accuracy
Reference accuracy
Reference accuracy
Reference accuracy
Baseline
Baseline
y-x
(FT)
(LP-FT)
(ZS-FT)
AugRog (FT)
AugReg (LP-FT)
Figure 4: Robustness of pro-trained models to synthetic in-support and out-of-support shifts. For each of two in-support shifts (left) and two out-of-support shifts (right) constructed by modifying Ima- geNet, the reference and shifted accuracies of models trained from scratch (in blue) are linearly correlated. Pre-trained models exhibit little effective robustness (ER), i.e, little improvement over the linear trend (see Section 2), on the in-support shifts, but have significant effective robustness on the out-of-support shifts (averages with 95% confidence intervals in the top left of each plot). Error bars denote 95% confidence intervals over 4 random trials.
indoor vs. outdoor features (Xiao et al., 2020; Geirhos et al., 2020). Thus, under a shift in which the setting/anini /animal correlation is reversed (which would be in-support but out-of-distribution) the model would likely underperform. If pre-training helps specifically with extrapolation, then it would not address this failure mode and, more generally, could not improve robustness to in-support shifts.
distribution shift out-of-support if there exists an input that could be sampled from the shifted distribution but could not. be reasonably sampled from the reference distribution. For example, consider cat dog image classification task in which photos from the reference distribution are taken during the day and photos from the shifted distribution distribution contains images with previonsly unseen lighting conditions. Here, a model trained from scratch might learn features that are sensitive to lighting and thus fail under the shift. Meanwhile, pre-training could provide priors for extrapolating, e.g., by producing features that are agnostic to lighting conditions as a starting point, leading greater robustness.
Out-of-support shift.
taken at night. In this case, the shifted
5.1 Constructing synthetic in-support and out- out-of-support shifts
We now want to measure the robustness gains that pre-training provides in-support and out-of-support shifts. To this end, we explicity construct two shifts of cach type by modifying ImageNet (Deng et al., 2009): (1) "spurious tint shift" in which we add a tint that is spuriously correlated with the label in the reference dataset, but not in the shifted dataset, (2) a "label shift" which the relative frequencies of classes change between the reference and shifted datasots, (3) an "unscen shift" add random tint in the shifted dataset, and (4) a "flip shift" in which we vertically flip images in the shifted dataset (see the top of Figure for visualizations).

---

Published in Transactions on Machine Learning Research (01/2025)
0.14
Out-of-suppot
Image
In-support
Average effective rabustness
0.12
0.10
Sketch In-support
0.08
0.06
pen
honokerchiet
scabbard
bulletproot
0.04
Sketch Out-of-suppor
0.02
0.00
sarang
coffeemaker
rugby ball
slide rule
torch
ImageNet-V2
ImageNet Sketch
ImageNet-R
(a) Random samples from: ImageNet (Lop): the in-
(b) Average effective robustuess of 55 pre-trained models on each split of each of the three shifts. Error bars denote 93% confidence intervals.
support split of ImageNet Sketch (middle) and out-of- support split of ImageNet Sketch (bottom).
Figure 5: Dividing shifts of ImageNet into in-support and out-of-support splits. We divide each of the ImageNet-V2, ImageNet Sketch and ImageNet-R datasets into an in-support split containing examples that look like ImageNet examples and an out-of-support split containing examples that look unlike ImageNet examples (see Appendix B.4 for a description of the splitting method). We display samples from each split of ImageNet Sketch in Figure 5a and report the average effective robustnesses of pre-trained models Figure 5b. See Appendix C.2.3 for scatlerplots of reference vs. shifted accuracy.
For each shift, as a baseline, we train a VIT-B/32 (Dosovitskiy et al., 2021) model from scratch on the reference dataset. We evaluate this model at different epochs and find a strong linear relationship wet reference and shifted accuracy, i.e., the accuracy on the line phenomenon occurs" (see Figure 4). Next, fino-tune pro-trained VrT-B/32 models and measure their effective robustness above this baseline. We consider two pre-trained models: CLIP (Radford et al., 2021) and AugReg (Steiner et al., 2021), and three (full) fine-tuning strategies: standard full fine-tuning (FT), linear probing followed by full fine-tuning (LP- FT) and zero-shot initialization followed by full fine-tuning (ZS-FTT). We select fine-tuning hyperparameters that maximize accuracy on the reference distribution (in Appendix C.1.1, we find that other reasonable hyperparameter choices yield similar robustness).
We observe that pre-trained models exhibit substantial effective robustness on out-of-support shifts, but have close to zero effective robustness in-support shifts (see Figure 4). In Appendix C.1.2, vary the strength of the biases in the in-support shifts and find that the elfective robustness of pre-trained models remains close to zero. See Appendix B.3 for description of the exact setup.
5.2 Dividing natural shifts into in-support and out-of-suppee out-of-support splits
So far, we have const ructed synthetic in-support and out-of-support shifts and observed that pre-training can significantly improve robustness to the latter but not the former. Now, we demonstrate that this principle seems to extend to natural shifts as well. Note that it hard to find natural shifts that are "purely" in- support. After all, under natural shifts the shifted dataset may contain some inputs that are similar those in the reference dataset and some that not. For example, in shift from photos to sketches, some sketches may look more photorealistic but most would probably be clearly distinguishable from photos. To be able to measure robustness lo each lype of shift, we thus divide several natural shifted datasets each into an "in-support split" containing inputs that look like they could have come from the reference dataset and an "out-of-support split" containing the remaining inputs. We do so by training a classifier to distinguish between the reference and shifted datasets and using this classifier to approximate the probability sampling given shifted example from the reference distribution (see Appendix B.4.1 for details).
Tyypically, one evaluates different models find this relationship. We different epachs due computational constraints.

---

Published in Transactions on Machine Learning Research (01/2025)
Effective robustnesses
Overlap of corrected examples (% of shifted worst-performing group)
55
Shifted worst-group accuracy
50
ER = 12.93%
45
40
ER 10.04%
ER = 1.99%
35
7.3%
3.3%
8.8%
30
25
20
4.5%
40
45
50
55
60
65
Reference accuracy
Baseline
Baseline linear fit
y
DFR
Pre-trained
Pre-trained and DFR
Figure 6: Combining pre-training and Deep Feature Reweighting (DFR) the WILDS-FMoW shift. Pre-training and DFR (an intervention designed to handle dataset biases (Kirichenko et. al., 2022)) each yield some effective robustness (ER) and combining these two interventions yields the most effective robustness (left). The examples corrected by applying pre-training and DFR have little overlap (right), indicating that they largely improve performance on different subpopulations. Meanwhile, the examples corrected by combining pre-training with DFR include most of the examples corrected by the individual interventions (right), suggesting that combining pre-training with DFR improves performance on both of these subpopulations. Error bars denote 95% confidence intervals over 64 random trials.
Specifically, consider three natural shifts of the ImageNet dataset: ImageNet-V2 (Recht et al., 2019), which closely resembles ImageNet, ImageNet. Sketch (Wang el al., 2019), which consists of sketches of Im- ageNet classes, and ImageNet-R. (Hendrycks et al., 2020a), which consists of "renditions" (e.g, paintings, sculptures, cartoons) of subset of ImageNet classos. We choose these shifted datasets because they include many inputs that look like they could have come from ImageNet and many that do not (according to our splitting method)4 Tn Figure 5a, we visualize examples from the in-support and out-of-support splits of ImageNet Sketch.
Consistent with our hypothesis that pre-training helps specifically with extrapolation, on the support splits of ImageNet Sketch and ImageNet-R pre-trained models have substantially higher effective robustness than on the respective in-support splits (see Figure 5b). On both ImageNet- V2 splits, however, pre-trained models have very little effective robustness. This may be because ImageNet- V2 is visually similar to Ima- geNet, so poor extrapolation might not be a significant failure mode (instead, the performance drop may be due to an increased presence of "harder" examples, as Recht et al. (2019) suggest). Thus, pre-training helps only with extrapolation, it would not be able substantially improve robustness on the ImageNet-VV V2 sutport examples. See Appendix B.4.2 for a description of the exact setup.
6 Combining Pre- Pre-Training with Interventions for Handling Bias
Our observations in Section 5 suggest that pro-training indeed can help prevent failures caused by poor extrapolation but not those stemming from biases in the reference dataset. How, then, can we develop models that avoid both failure modes? In this section, we explore one possible strategy: combining pre- training with interventions specifically designed to handle dataset biases.
In particular, we investigate the effectiveness of this strategy on WILDS-FMoW (Christie et al., 2018; Koh et al., 2020), a distribution shift benchmark for classifying satellite images (in Appendix C.3.1, we provide a
We also explored ObjectNet (Barbu et al., 2019) and ImageNet- marks fewer than examples from these shifted datasets as "in-support," and thus we cannot reliably measure in-support. accuracy.
SuageNet-Vid-Robuu Vid-Robust (Shankar et al., 2019) but splitting method

---

Published in Transactions on Machine Learning Research (01/2025)
similar analysis for synthetic distribution shift). In WILDS-FMoW, the reference dataset consists of satellite images taken between 2002 and 2012, while the shifted dataset consists of satellite images taken wetween 2016 and 2017. Additionally, the images depict different regions and models typically underperform on underrepresented regions. Following Koh et al. (2020), we evaluate the worst-group accuracy (the minimum accuracy across groups in our case, regions) on the shifted dataset. Hence, robustness to this shift requires being able to both extrapolate to later years and perform consistently across regions (e.g., by avoiding biases that are harmful to performance on some regions).
Aiming to overcome these two challenges, we leverage two types of interventions. To extrapolate better to later years, we initialize the model via pro-training: specifically, we obtain our model by fine-tuning CLIP ResNet-50 model. To handle potential biases in the reference dataset, we employ Deep Feature Reweighting (DFR) (Kirichenko et al., 2022), an intervention intended to de-bias a model by re-training just the final layer on group-balanced data. We measure the effective robustness of each intervention over a baseline of ResNet-50 models trained from scratch. We find that pre-training and DFR each yield some effective robustness and that combining the two yields greater effective robust.ness than applying either individually (see the left side of Figure 6). See Appendix B.5 for a description of the exact setup.
Understanding robustness benefits. We observe that combining pre-training and DFR can be effective for developing robust models, but is this actually because they address different failure modes, we suggest? To answer this question, we consider the corrected eramples of each intervention, i.., the set of test examples that are often classified incorrectly by baseline model but correctly by model with the intervention (on average over 64 trials). We observe that the corrected examples of pre-training and DFR have little overlap (see the right side of Figure 6), suggesting that their benefits are indeed complementary. Meanwhile, the corrected examples of combining pre-training with DFR include most of the corrected examples of the individual interventions. This suggests that combining pre-training with DFR not only yields high effective robustness but in fact leads to models with both sets of benefits.
7 Curating Datasets for Fine-Tuning
In Section 6, we explored pairing pre-training with interventions specifically designed to address dataset biases. We observed that this strategy can be effective for developing models that both extrapolate effectively and avoid undesirable biases present in the reference distribution.
In this section, we highlight one such intervention: training on a carefully curated (and, in particular, de- biased) dataset instead of the original reference dataset. In general, de-biasing a large and diverse dataset may be prohibitively expensive. However, if we can rely on pre-training for extrapolation (as suggested in Section 5), might only need small, non-diverse fine-tuning dataset, which would be more feasible to do-bias. Thus, curating such a datasct and then fine-tuning large pre-trained model on it might be a relatively inexpensive method for developing robust and performant models.
As a case study, we consider the task of predicting hair color (blond vs. non-blond) in the CelebA dataset (Liu et al., 2015). In this dataset, hair color is spuriously correlated with other attributes (espocially gender). For example, 24% of females are blond, while only 2% of males are blond. Following works studying group robustness (Sagawa et al., 2020a; Liu et al., 2021; Kirichenko et al., 2022), we measure worst-group accuracy to assess robustness rather than measuring accuracy on explicit shifted dataset. In this case, the four groups are blond females, non-blond females, blond males and non-blond males. A model exploiting the spurious correlation between gender and hair color would likely perform poorly on the underrepresented group of blond males.
Curating a de-biased datasct. To curate de-biased dataset for hair color classification with n examples, we construct a "counterfactual example" for each of n1/2 CelebA examples by changing the person's hair to a color corresponding to the opposite class (i.e., blond to non-blond and vice versa). We ensure that attributes besides hair color remain unchanged and include both the original and edited images in our datasct. Hence, attributes that are spuriously correlated with hair color in the CelebA dataset (e.g., gender, age) are equally represented in the blond and non-blond populations of our curated dataset. To illustrate that this dataset
9

---

Published in Transactions on Machine Learning Research (01/2025)
CelebA clttaset (n 162K)
Non-blond femaler 44%
Bland female: 14%
Accuracy vs. worst-group accuracy
90
68.35%
Worst-group accuracy
-0 Avg. ER 60.85%
70
Non-blond male: 40%
Bland male:
60
Curated dataset (r> = 64, 8192)
50-
ER 5.74%
Non Blon-blond ferale:: 25%
Bland femai e: 25%
40
Accuracy
Baseline
Pre-trained
Countarfact.ua) Blond Samals: 25%
Counterfachai) Non-blond female: 25%
Baseline linear fit
Curated In 64, .... .1.2)
-
Pre-trained Curated (n 64)
Non-bland
Blond
Counterfactual Jusing image aditing)
(a) CelebA vs. our curated hair color classification dataset. In the CelebA dataset (top), attribules such as gender are spurionsly correlated with the class (blond non-blond). In much smaller curated dataset (bottom), every real image is paired with a synthesized "counterfactual example" of the other class. As result, the primary difference between the blond and non-ble non-blond populations is hair color; other attributes such gender, age and hair style are not predictive. We include only females in our dataset to illustrate that diversity might not be necessary for robustness when fine-tuning.
(b) Fine-tuning on our curated dataset. Fine-tuning pre-trained model the CelebA dalaset (orange) yields little effective robustness over a baseline of models trained from scratch (blue). However, fine-tuning the same pre- trained model on just 64 examples from our curated dataset (red) yields model with both high effective ro- bustness and high accuracy. Training from scratch on our curated dataset (green) also yields high effective ro- bustness but results in substantially lower accuracy than pre-trained models, even with many more examples. Er- ror bars denote 95% confidence intervals over 64 random trials.
Figure Fine-tuning pre-trained model on a small, non-diverse but de-biased dataset (see Figure 7a) yields robust and performant model for hair color classification in CelebA (sce Figure 7b).
does not need to be diverse to yield high robustness and performance when fine-tuning, we restrict the dataset to include only females. See Figure 7a for a visualization of the dataset and Appendix B.6 for the image editing process. In Appendix C.4.2, we consider the simpler curation strategy of balancing the number of samples from each group (Idrissi et al., 2022) and find that counterfactual image editing is more effective.
Fine-tuning on a de-biased dataset. As expected, models trained from scratch on the CelebA dataset exhibit high accuracy but very low worst-group accuracy, likely because they rely on gender to predict hair color (see Figure 7b). Furthermore, a pre-trained CLIP ViT-B/32 model fine-tuned on the CelebA dataset exhibits very little effective robustness above these models trained from scratch, consistent with our hypothesis that pre-training does not mitigate dataset biases. However, we observe that fine-tuning the same pre-trained model on just 64 examples from our curated dataset yiclds model with both high accuracy and effective robustness. Finally, we also Irain models from scratch on our curated dataset and find that they exhibit substantial effective robustness, but require many more examples to attain a comparable accuracy. This suggests that the extrapolation benefits of pro-training are key make effective use of our small, non- diverse curated dataset. In particular, as we illustrate in Appendix C.4.1, pre-trained models extrapolate from the female-only curated dataset to males better than models trained from scratch.
10

---

Published in Transactions on Machine Learning Research (01/2025)
8 Related Work
Characterizing distribution shifts. There exists plcthora of definitions for characterizing distribution shifts, many of which are aligned with the in-support and out-of-support chracterizations that discuss in this work. For example, domain generalization involves shifts in which the reference and shifted distributions are from different domains (Koh et al., 2020; Guirajani & Lopez-Paz, 2020). In a subpopulation shift, subpopulations appear with different frequencies in the reference and shifted distributions (Santurkar et al., 2021; Koh et al.; 2020; Yang et al., 2023). In shifts with spurious corvelations, certain features are predictive in the reference distribution but not in the shifted distribution (Arjovsky et al., 2019; Sagawa et al., 2020b). Two more formal characterizations are covariale shift (Shimodaira, 2000), under which p(ylæ) is fixed, and label shift (Lipton et al., 2018), under which the label distribution may change but p(øly) is fixed. We relate these definitons to in-support and out-of-support shifts in Appendix D.4.
Robustness benefits of pre-training. Several works have suggested that pre-training can be an effective strategy for improving robustness to distribution shifts (Hendrycks et al., 2019; 2020a;b; Tu et al., 2020; Taori et al., 2020; Miller et al., 2021; Wiles et al., 2021; Andreassen et. al., 2021; Bommasani el. al., 2021; Liu et al., 2022b; Ramanujan et al., 2023). In particular, Wiles et al. (2021) define different types of distribution shifts and find that pro-training frequently improves performance under these shifts, while most other interventions primarily help specific settings. In the natural language processing selting, Tu el al. (2020) argue that when pre-training helps with spurious correlations, it is because pre-trained models can generalize better we discuss in Appendix D.5, this is
from the small number of counterexamples to these correlations; consistent with our intuition that pre-training helps specifically with extrapolation. Lastly, Bommasani et al. (2021) discuss failure modes that pre-training is unlikely to address including spurious correlations (both in pre-training and fine-tuning datasets) and extrapolation across time.
9 Conclusion
In this work, we study the failure modes that pre-training alone can and cannot address. Our findings suggest that pre-training can help mitigate failures caused by poor extrapolation (e.g., inability to generalize to a new domain) but might not address other failures, such those stemming from dataset biases. In light of this observation, dataset biases present a fundamental limitation that cannot be overcome by simply leveraging additional pre-training data or larger models. We thus encourage practitioners not to treat pre- training as a panacca for robustness. Instead, they should consider the specific failures modes they might encounter to determine if pre-training can help.
References
Samira Abnar, Mostafa Dehghani, Behnam Neyshabur, and Hanie Sedghi. Exploring the limits of large scale pre-training. arXiv preprint arXiv:2110. 0.00200 2021.
Anders Andreassen, Yasaman Bahri, Behnam Neyshabur, and Rebecca. Roelofs. The evolution of out-of- distribution robustness throughout fine-tuning. arxiv preprint arXiv:2106. 15831, 2021.
Martin Arjovsky, Lóon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprinl orXiv:1907.02560 2019.
Peter Bandi, Oscar Geessink, Quirine Manson, Marcory Van Dijk, Maschenka Balkenhol, Meyke Hermsen, Babak Ehteshami Bejnordi, Byungiae Lee, Kyunghyun Paeng, Aoxiao Zhong, et al. From detection of individual metastases classification of lymph node status at the patient level: the camelyon17 challenge. IEEE Transoctions on Medical Imaging, 2018
Andrei Barbu, David Mayo, Julian Alverio, William Luo, Christopher Wang, Dan Gutfreund, Josh Tenen- baum, and Boris Katz. Objectnet: A large-scale bias-controlled dataset for pushing the limits of object recognition models. In Newral Information Processing Systems (NeurIPS), 2019.
11

---

Published in Transactions on Machine Learning Research (01/2025)
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXic:2108.07258 2021.
Tim Brooks, Aleksander Holynski, and Alexei A Efros. Instructpix2pix: Learning to follow image editing instructions. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 18392 18402, 2023.
Sébastien Bubeck. Theory of convex optimization for machine learning. arXiv preprint arXiv:1405.4988. 15, 2014.
Jonathon Byrd and Zachary Lipton. What is the effect of importance weighting in deep learning? In International conference on machine learning, pp. 872 881. PMLR, 2019.
Gordon Christie, Neil Fendley, James Wilson, and Ryan Mukherjee. Functional map of the world. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 6 2018.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In Computer Vision and Pattern Recognation (CVPR), 2009.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. In International Conference on Learning Rep- resentations (ICLR), 2021.
Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix Wichmann. Shortcut learning in deep neural networks. In Nature Machine In- telligence, 2020.
Karan Goel, Albert Gu, Yixuan Li, and Christopher Ré. Model patching: Closing the subgroup performance gap with data augmentation. arXiv preprint ar Xrxi:20008. 06775, 2020.
Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. arXiv:2007.01/33 2020.
arXie preprint
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In International conference on machine learning, pp. 1321 1330. PMLR, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015.
Kaiming He, Ross Girshick, and Piotr Dollár. Rethinking imagenet pre-training, In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4918-4927, 2019.
Dan Hendrycks, Kimin Lee, and Mantas Mazcika. Using pre-training can improve model robustness and uncertainly. In International Conference on Machine Learning, pp. 2712 2721. PMLR, 2019.
Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, Dawn Song, Jacob Steinhardt, and Justin Gilmer. The many faces of robustness: A critical analysis of out-of-distribution generalization, 2020a.
Dan Hendrycks, Xiaoyuan Liu, Eric Wallace, Adam Dziedzic, Rishabh Krishnan, and Dawn Song. Pretrained transformers improve out-of-distribution robustness. arXiv preprint arXi::204.06100 2020b.
John Hewitt, Xiang Lisa Li, Sang Michael Xie, Benjamin Newman, and Percy Liang. Ensembles and cocktails: Robust finetuning for natural language generation. 2021.
Badr Youbi Idrissi, Martin Arjovsky, Mohammad Pezeshki, and David Lopez-Paz. Simple data balancing achieves competitive vorst-group-accuracyy In Conference on Causal Learning and Reasoning, pp. 336 351- PMI.R, 2022.
12

---

Published in Transactions on Machine Learning Research (01/2025)
Gabriel Ilharco, Mitchell Wortsman, Ross Wightman, Cade Gordon, Nicholas Carlini, Rohan Taori, Achal Dave, Vaishaal Shankar, Hongseok Namkoong, John Miller, Hannaneh Hajishirzi, Ali Farhadi, and Ludwig Schmidt. Openclip, 2021.
Chao Jia, Yinfei Yang, Ye Xia, Yi-Ting Chen, Zarana. Parekh, Hieu Pham, Quoe Le, Yun-Hsuan Sung, Zhen Li, and Tom Duerig, Scaling up visual and vision-language representation learning with noisy text supervision. In International Conference on Machine Learning, pp. 4904-4916. PMLR, 2021.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of GANs for improved quality, stability, and variation. In International Conference Learning Representations, 2018.
Alexander Ke, William Ellsworth, Oishi Banerjee, Andrew Y Ng, and Pranav Rajpurkar. Chextransfer: performance and parameter efficiency of imagenet models for chest x-ray interpretation. In Proceedings of the Conference on Health, Inference, and Learning, pp. 116-124, 2021.
Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training sufficient for robustness spurious correlations. arXiv preprint arXiv:2204. 029337, 2022.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Sara Becry, et al. Wilds: A benchmark of in-the-wild distribution shifts. ar preprinl arXiv:2012.070 07421, 2020.
Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby. Big transfer (bit): General visual representation learning. orXiv preprint orXiv: 11370, 2019.
Simon Kornblith, Jonathon Shlens, and Quoe V Le. Do better imagenet models transfer better? In computer vision and pattern recognition (CVPR), 2019.
Ananya Kumar, Aditi Raghunathan, Robbie Jones, Tengyu Ma, and Percy Liang. Fine-tuning can distort pretrained features and underperform out-of-distrrbution. arXiv preprint arXiv:2202.100:: 2022.
Guillaume Loclere, Andrew Ilyas, Logan Engstrom, Sung Min Park, Hadi Salman, and Aleksander Madry. flev. https:/ thw.. com com//ibffcr/ffer// 2022.
Cheng-Han Lee, Ziwci Liu, Lingyun Wu, and Ping Luo. Maskgan: Towards diverse and interactive facial image manipulation. In Proceedings of the IEEE/CVP Conference on Computer Vision and Paltern Recognition, pp. 5549 5558, 2020.
Zachary Lipton, Yu-Xiang Wang, and Alexander Smola. Detecting and correcting for label shift with black box predictors. In International conference on machine learning, pp. 3122 3130. PMLR, 2018.
Evan 7 Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In International Conference on Machine Learning, 2021.
Zhuang Liu, Hanzi Mao, Chao- Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 11976-11986, 2022a.
Ziquan Liu, Yi Xu, Yuanhong Xu, Qi Qian, Hao Li, Rong Jin, Xiangyang Ji, and Antoni B Chan. An em- pirical study distribution shift robustness from the perspective of pre-training and data augmentation. arXiv preprint ar Xix::2205.. 12753, 2022b.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In International Conference on Computer Vision (ICCV), 2015.
Dhruv Mahajan, Ross Girshick, Vignesh Ramanathan, Kaiming He, Manohar Paluri, Yixuan Li, Ashwin Bharambe, and Laurens van der Maaten. Exploring the limits of weakly supervised pretraining. In Proceedings of the European Conference on Computer Vision (FCCV), 2018.
13

---

Published in Transactions on Machine Learning Research (01/2025)
John Miller, Karl Krauth, Benjamin Recht, and Ludwig Schmidt. The effect of natural distribution shift question answering models. In International Conference on Machine Learning, pp. 6905-6906000 0016- PMT.R, 2020.
John P Miller, Rohan Taori, Aditi Raghunathan, Shiori Sagawa, Pang Wei Koh, Vaishaal Shankar, Percy Liang, Yair Carmon, and Ludwig Schunddt. Accuracy on the line: on the strong correlation between out- of-iistrinttion and in-distribution generalization. In International Conference on Machine Learning, pp. 7721-7735. PMLR, 2021.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkinn Jack Clark, et al. Learning transferable visual models from natural language supervision. In arXiv preprint ar Xiv:21003 000000 2021.
Vivek Ramanujan, Thao Nguyen, Sewoong Oh, Ludwig Schmidt, and Ali Farhadi. On the connection wetween pre-training data diversity and fine-tuning robustness. orXiv preprint ar 122: 2023.
Benjamin Recht, Rebecca Roclofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? In International Conference on Machine Learning (ICML), 2019.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recogrition, pp. 10684 10695, 2022.
Shiori Sagawa, Pang Wei Koh, Tatsunoi B. Hashimoto, and Percy Liang, Distributionally robust neural net- works for group shifts: On the importance of regularization for worst-case generalization. In International Conference Learning Representations, 2020a.
Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overpa- ramoterization exacerbates spurious correlations. In International Conference on Machine Learning, pp. 8346 8356. PMLR, 2020b.
Hadi Salman, Saachi Jain, Andrew Ilyas, Logan Engstrom, Eric Wong, and Aleksander Madry. When does bias transfer in transfer learning? In arXie preprint arXiv:2207.088/4 2022.
Shibani Santurkar, Dimitris Tsipras, and Aleksander Madry. Breeds: Benchmarks for subpopulation shift. In International Conference on Learning Representations (ICLR), 2021.
Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mchdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, el al. Laion-5b: An open large-scale dataset for training next generation image-text models. In arXiv preprint orXiv:2210.08400 2022.
Vaishaal Shankar, Achal Dave, Rebecca Roelofs, Deva Ramanan, Benjamin Recht, and Ludwig Schmidt. Do image classifiers generalize across time? arXiv preprint srXiv:1906.02111 2019.
Ali Sharif Razavian, Hossein Azizpour, Josephine Sullivan, and Stefan Carlsson. Cun features off-the-shelf: an astounding baseline for recognition. In conference on computer vision and pallern recognition (CVPR) workshops, 2014.
Hidetoshi Shimodaira. Improving predictive inference under covariate shift by weighting the log- likellood function. Journal of statistical planning and inference, 2000.
Andreas Steiner, Alexander Kolesníkov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. orXiv preprint arXiv:2106. 10270, 2021.
Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhinav Gupta. Revisiting unrcasonable effectiveness of data in deep learning era. In Proceedings of the IEEE international conference on computer vision, 2017.
14

---

Published in Transactions on Machine Learning Research (01/2025)
Rohan Taori, Achal Dave, Vaishaal Shankar, Nicholas Carlini, Benjamin Recht, and Ludwig Schmidt. Mea- suring robustness to natural distribution shifts in image classification. In Neural Information Processing Systems (NewrIPS), 2020.
Lifu Tu, Garima Lalwani, Spandana Gella, and He He. An empirical study on robustness to spurious corre- lations using pre-trained language models. Transactions of the Association for Computational Linguistics, 8:621 633, 2020.
Haohan Wang, Songwei Ge, Eric P Xing, and Zachary Lipton. Learning robust global representations by penalizing local predictive power. Neural Information. Processing Systems (NewrIPS), 2019. 2019
Ross Wightman. Pytorch image models. http::/hhhw....ooo.aaaaaaa//aaaa
2019.
Olivia Wiles, Sven Gowal, Florian Stimberg, Sylvestre Alvisc-Rebufi, Ira Ktena, Krishnamurthy Dvijotham, and Taylan Cemgil. A fine-grained analysis on distribution shift. arXiv preprial arXiv:21 10. 11328, 2021.
Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim, Mike Li, Simon Kornblith, Rebecca Roelofs, Raphael Gontijo-Lopes, Hannaneh Hajishirzi, Ali Farhadi, Hongseok Namkoong, and Ludwig Schmidt, Robust finc-tuning of zero-shot models. In arXiv preprint Xiv:21100 01903, 2021.
Kai Xiao, Logan Engstrom, Andrew llyas, and Aleksander Madry. Noise or signal: The role of image backgrounds in object recognition. arXiv preprint XrXiv:2006.099 2020.
Michael Xie, Neal Jean, Marshall Burke, David Lobell, and Stefano Ermon. Transfer learning from deep features for remote sensing and poverty mapping. In Thirlieth AAAI Conference on Arlificial Intelligenee 2016.
Yuzhe Yang, Haoran Zhang, Dina Katabi, and Marzyeh Ghassemi. Change is hard: A closer look subpopulation shift. arXie preprint arXin:2902. 12254, 2023.
Nanyang Ye, Kaican Li, Haoyue Bai, Rumpeng Yu, Lanqing Hong, Fengwei Zhou, Zhenguo Li, and Jun Zhu. Ood-bench: Quantifying and understanding two dimensions of out-of-distribution generalization. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 7947-7958, 2022.
15

---

Published in Transactions on Machine Learning Research (01/2025)
Appendices
A Theoretical Results
17
A.1 Proof of Theorem 4.1
B Experiment Details
21
B.1 General
21
B.2 The robustness benefits of pre-training vary
22
B.3 Constructing synthetic in-support and out-of-support shifts
22
B.4
Dividing natural shifts into in-support and out-of-support splits
................................
23
B.5 Combining pre-training with interventions for handling bias
25
B.6 Curating datasets for fine-tuning
25
C Additional Results
28
C.1 Constructing synthetic in-support and out-of-support shifts
28
C.2 Dividing natural shifts into in-support and out-of-support splits
31
C.3 Combining pre-training with interventions for handling bias
C.4 Curating datasets for fine-tuning
35
D Additional Discussion
38
D.1 Alternative fine-tuning strategies
................................
38
D.2 Can pre-training hurt extrapolation?
38
D.3 When does pre-training help with extrapolation?
38
D.4 Relating in-support and out-of-support shifts to existing characterizations
38
D.5 Understanding the robustness of pre-trained language models to spurious correlations
39
D.6 Additional related work
39
16

---

Published in Transactions on Machine Learning Research (01/2025)
A Theoretical Results
A.1 Proof of Theorem 4.1
Setup. Suppose that given access to reference datasel Stet of input-label pairs (x, y), with Rd and y {-1,1}. We decide to learn linear classifier for this task by finding weight that minimizes the (standard) logistic loss on Stef:
[ue(w)
Σ
log(1 + ee- =
(1)
(a.
We assume that the reference dataset Stef satisfies the following conditions:
Inputs Stet lie within a k-dimensional (with < d) subspace Wrer of Rd. Intuitively, this condition represents lack of variation in certain features in the reference dataset.
The logistic loss Lref has a minimum value. This condition ensures that minimizing Lief is well-defined. Note that there may be multiple weights that attain this minimum valuc.
Theorem 4.1. Suppose that we start with initial weights Wint E Rd and run gradient descent to minimize Lmg(w). With approprialely chosen learning rale, gradient descent converges to weights w that mánimize Lef. Purthermore, w can be written as
Wee projw Winst-
(2)
Here, Wief is property of the reference dataset Syef and lies within the reference subspace Weef- Meanwhile, projw-Wio Wint is the component of Wintt that is orthogonal to Wref-
To prove Theorem 4.1, will first show that rumning gradient descent starting from an initialization within Wrer always converges to the same weights wee- We will then show that running gradient descent starting from an arbitrary initialization has the same convergence behavior except for an "offset" term projw Wint representing the component of the initialization that is orthogonal to Wref-
for
A.1.1 Convexity and smoothness of the loss
We begin by providing the gradient and hessian Lyef and using these to establish convexity (Lemma A.1) and smoothness (Lemma A.2) propertics of Lyef- The gradient of Lyef is
V Leee(w)=
Σ ..
(3)
1 e z.y
(a.))E
The Hessian of Leef is
2 (we)(w)
Σ
T
X xDw)XX
(4)
2 e-w a. ew a.y
where X E xl is the matrix of inputs Stet and D(w) [SSa the diagonal matrix with D(w) 2te uT= Note in particular that the non-zero elements of D(w) are in (0, 1/4).
Lemma A.1. The loss Lonf is (1) conver on Rd, (2) strictly convex on W Weers and (3) strongly conver on any closed convex subset of Wref-
Proof. According to Taylor's Theorem, for any u, v E Rdd there exists a E [0, such that
Lie((v) = Lie((u) + Lree(())
T (v u) 2 (v - u) T T 2 Lyoe(v + a (v- u))( - u).
(5)
17

---

Published in Transactions on Machine Learning Research (01/2025)
1. Convexity on R To show that Liet is convex on R", need to show that
Tee((p) 2 [ue Viret(u) (o u)
for any u, E TR.. Using (5), suffices to show that a T T Ieet(wa Recall from (4) that Loo((w) x Thus, have
for any a e Rd and w RR..
2 [nor(w)a = x D(w)Xa
(w)) 1/2 Xal =IDD())122al3
2 0
Strict convexity on Wee- Next, show that Lief is strictly convex on Wief we need show that
Lie((t) > (ue((u) + V Lr (o (0- 2
for any u, Wzef- Using (5), it suffices to show that a L2 > 0 for any non-zero E Wzet and E Wef- We know that a V (w
1/2 Xall3. Since D(w) is diagonal with positive entries along the diagonal, 1/2 Xa 1 > 0 if and only if Xa # O. Recall that Wit is the subspace spanning the rows of X. Hence, since a is non-zero and is in Wref, we know that Xa + 0.
Strong convexity on any closed convex subset of W ret- Finally, show that Lief strongly convex any closed convex subset T of W need to show that there exists > 0 such that
[uo((u) 2 Inor(u) + VLyer(u) (v u) + m 11 ul12
2
for any u, v E T. Using (5), suffices to show that there exists an m > 0 such that T V22 Lrer(w)a m ||all3 for any Wre and E T. Making use of the fact that T is closed, let Amin be the 2 minimum diagonal entry of D(w) for T, that is,
Amin min min D(w).-
WET ie SS.dl
Next, let Cmin be the minimum value of xall3 over unit vectors a in Wref, that is,
Cmin
min
|aall
oeWaot.llala-
We previously established that Xa + 0 for any non-zero Wrefs which means that Cmin > 0. F Finally, we conclude that for m = 2 Amin CCina a L = [D(w) 1/2 Xall3 2 Amin Cmin ||a||| 1 m 2
||a]3
Lemma A.2. The gradient of the loss function Lveg is K K-Lipschitz with K = ||
Proof. To show that V Viret is K-Lipschitz, we need to show that (2I KT. Recall from (4) that LL = X X. Thus, we have
a T 2 Lor(wwa = a TT xTT D(w)Xa
[[D(w) X
< (LD(w)1/11) ||X 11 ||a|3
(| -
In the final step, we use the fact that (ww) is diagonal with non-zero elements in (0, 1/4) conclude that || (2w))/21112 < 1/4.
18

---

Published in Transactions on Machine Learning Research (01/2025)
A.1.2 Convergence of gradient descent within the reference subspace
Next, we establish that there exists a unique minimizer of Lief within the reference subspace ore (Lemma A.3) and that gradient descent converges to these weights (Lemma A.4).
Lemma A.3. There crisls unique Wrof E Weer such that waas E arg min L(w).
Proof. We will first show that there exists W W ref such that wiet arg mig, [u)r(w). Let w E arg ming (w) be arbitrary minimimum point of Liet- By definition, for every (2.y) Stet, E Wiet- Hence, for every such I, w w = projw... w TT.. x. This means that [res("") = Tree (prowwiw.... w"), which implies that oo == projw... "² E arg min Lve((w), as desired. Next, because Leet is strictly convex on Weef (Lemma A.1), Weet is the only minimum point of Lief in Wzel-
Lemma A.4. If we start with Wint E Wref and run gradient descent with 17 = 4/| to minimize Lref(w), the weights will converge to waof
Proof. Suppose that we start with initial weights Winit E Wfer and run gradient descent to minimize Lyef with learning rate n. In particular, let w(() Wiit and w(t+1)- wit) + n (Leer(w(t)), Because Lyef is
convex (Lemma A. V Toe is K-Lipschitz with K = ||x11 (Lemma A.2), and 1) 4/ 4/||X| S 11/K, we know from Theorem 3.2 of Bubeck (2014) that
K
||mmit
w
[ee((ul [ree((Wfet)
(6)
Hence, the loss attained by w!() converges to the optimal loss attained by wiet To show that wol converges to Weer will show that Lief is strongly convex set containing every w(th for t > 0. In particular, consider the set WGn = fw E Wro [|ll-w"ed la S || - wfelll2 containing weights in Wi at least as close to Wiet as Winit. Clearly, WGD contains (0) = Wmitt. We know from Theorem 3.2 of Bubeck (2014) that with each iteration of gradient descent we get closer to a minimum point, that is, [| - w'tell S || wiell- Additionally, because Winit and Lref are in Wref, every w(t) is in Wref- Hence, every w(()) is in WGD. Because WGD is closed and convex, from Lemma A.1 we know that Lief strongly convex on WGD. This means that there exists m such that
(tt)( 2 (wfe) + V (ree(weed)
(witt) (ww
-wfet) + m. ||x(()) wfer 2
Plugging in V Lree(wiet) (ot) = 0 and rearranging, we get
2 (L)e((u(()) - Leer(w(fot)
[w0(0) 1000-wwrlll
w S m
Finally, combining with (6) yields
2.K ||miiit wiet
(l((1) wirll2 S
(7)
(t 1)
which completes our proof.
A.1.3 Proof of Theorem 4.1
We are now ready to prove Thcorem 4.1. Suppose that we start with initial weights Wint and run gradient descent. to minimize Leof with learning rate 7) 4/||X l In particular, let. w(()) = Winit and w(++1) w( + (lee(W(()) for 2 0. We will show that running gradient descent starting with an arbitrary Winit has the same behavior as running gradient descent with Wini projected onto WTof- To be more precise, suppose that we instead start with initial weights projw.., Wini when running gradient descent. In particular, with projw-u u denoting the projection of u onto a subspace W, let Wproj (0) = projw, Wint and Uproj (t+1) = Wproj (+) + + (iree( (1) for t 2 0. Then the trajectory of w(t) is the same as that of Wproj (t) but with an additional component projw Winit = (Wint projw... Wins). That is, cor
projw ees Winit Upros (0) (t
19

---

Published in Transactions on Machine Learning Research (01/2025)
To show that this is the case, we will proceed by induction. As base case,
w
Winit
(Winit projw... Winit) + projw,.. Winit projw Winit + Uproj (0).
For the inductive step, assume that the statement holds for t = k. Then,
w (k-1) = (k) +1 (ere((u(k)
projw: Winit + Wproj (k) + rj Vireet(roojw Winit + Uproj (8))
projw Winit + Uproj (A) + 1) Veret(uprof (8))
projwa Winit + Wprol (k+1)
where in the third step we the fact that Viret(u + v) Virer(u) if W This completes the induction. Because Wpro) (0) projw... Wint Wier, from Lemma A.4 (in particular, from equation 7), we know that
praj w S 2. K. |lprojw Wimit Winit - w
(1-
where K and m are positive constants. Finally, we conclude that
(0) w = [I(projwy Wintt + Wproj (3)) (projw Wint + w
[| (f) w
2 K proo ,or Winit wfell
(t
Hence, w(t) converges to á completing proof.
20

---

Published in Transactions on Machine Learning Research (01/2025)
B
Experiment Details
B.1
General
B.1.1
Model training
All models are trained using the FFCV data-loading library (Leclerc et al., 2022) on a cluster of A100 GPUs.
B.1.2 Measuring effective robustness
Effective robustness. In this work, we quantify the robustness of pre-trained models using effective robusiness (ER), a measure of the robustness a model above the "baseline" of models trained from scratch (Taori et al., 2020). Computing this metric first involves establishing relationship between the accuracies of baseline models (in case, models trained from scratch (M) and Accain (M) denote the accuracies of model M test datasets drawn from the reference and shifted distributions, respectively. Given a set Mhaseline of baseline models, compute linear fit relating & - (Accerce(()))) and & (Accalit (Af),, where Φ is the probit function (i.e., the inverse cumulative distribution function of the standard normal distribution). We compute a linear fit relating probit-scaled accuracies (instead of the accuracies themselves) because this has been empirically observed to improve the strength of the linear relationship (Miller et al., 2021; Taori et al., 2020). Formally, we compute parameters and such that
reference dataset). In particular, let
à, b arg min a
Σ
ll(a (-((Acce((MM) b) (Acccau((()))
Moneeiiiiee
Let Accabif (M) be the resulting function estimating shifted accuracy given reference accuracy, that is Accabift (M) = Φ(a Baae:d-=AAAcc(fAaa "(Accre((M)) + 8).
Then the effective robustness of model M
ER(M) = Acc (M) AcCahitt (M)
Intuitively, effective robustness is the extent to which a model's accuracy on the shifted distribution exceeds the accuracy of a baseline model with the same accuracy on the reference distribution (see Figure 8).
90
y=x
Baseline
Baseline lincar fit
Pre-trained
Shifted accuracy
80
Effective bobustness
(ER)
70
Bye
60
50
60
65
70
80 Reference accuracy
75
90
Figure 8: Visualization of effective robustness. To compute effective robustness (ER), we first estab- lishes a linear relationship between the (probit-scaled) accuracies of baseline models (blue) on the reference and shifted datascts. The effective robustnoss (groen) a pre-trained model (orange) is the amount by which its actual accuracy on the shifted dataset exceeds the prediction of the linear trend.
21

---

Published in Transactions on Machine Learning Research (01/2025)
Establishing baseline for effective robustness. To establish a baseline with respect to which measure effective robustness, we need a set of baseline models trained from scratch on the reference dataset. The set of baseline models we consider varies by experiment. In each of the experiments in which we measure effective robustness, we confirm that a strong linear relationship exists between the probit-scaled accuracies of our baseline models on the reference and shifted datasets (see, e.g., Figure 4).
B.2 The robustness benefits of pre-training vary
In Section 3, we illustrate that pre-trained models exhibit substantial effective robustness on the ImageNet Sketch distribution shift but very little effective robustness on the ImageNet-V2 distribution shift. We consider 78 models trained from scratch on ImageNet and 55 pre-trained models fino-tuned on ImageNet, all taken from PyTorch Image Models (Wightman, 2019). The pre-trained models represent variety of model architectures (c.g., ResNet (He et al., 2015), ConvNeXt (Liu et al., 2022a), ViT (Dosovitskiy et al., 2021)); pre-training datasets (e.g., IG-1B (Mahajan et al., 2018), LAION-2B (Schuhmann et al., 2022), OpenAl's WIT (Radford et al., 2021)), and pre-training algorithms (e.g., supervised learning, CLIP (Radford et al., 2021)). The complete list of models used is available with our code at bttps://gitee thub. com com/MadryLab/ pretraining- weaiingddistriii distribution wiwioon-ssiftrree thift-robustness.
B.3 Constructing synthetic in-support and out-of-suppon shifts
In Section 5.1, we measure the effective robustness of various pre-trained and fine-tuned models on two in-support and two out-of-support shifts syuthetically constructed by modifying ImageNet.
Specifications of hentc shifts. Here, we provide detailed descriptions of the four synthetic distribu- tion shifts (see Figure for visualizations).
1. Spurions tint shift. (in-support): We tint images (i.e., replace each pixel with mix of the original value, with weight 0.75 and a specific color, with weight 0.25) such that the tint is correlated with the label in the reference distribution but not in the shifted distribution (i.e., lint is a spurious feature). Specifically, in the reference distribution we apply tint with a class-specific color to Pspurious = 0.5 of examples and a tint with a random color the remaining 1-Paa PPppurious 0.5 of examples. Mcanwhile, in the shifted distribution we apply a tint with random color universally.
Label shift (in-support): Label shift is a commonly studied type of distribution shift in which the relative frequencies of classes change, but p(xly) is fixed. To construct a label shift, we sub-sample ImageNet such that in the reference distribution, a randomly selected 500 classes are less likely appear than the remaining 500 classes. In particular, the selected classes appear with probability Pminority = 0.2, while the remaining classes appear with probability 1- P-miuorrity = 0.8. In the shifted distribution, these relative frequencies are reversed.
3. Unseen tint shift (out-of-support): We randomly tint images in the shifted distribution (with the same protocol in the spurious tint shift).
4. Flip shift. (out-of-support): We vertically flip images in the shifted distribution.
Shared model specifications. When training, we use the FFCV implementation of RandomResizzddCro DeGodray,, resizing image crops to resolution of 224 x 224. For data augmentation, we the FFCV implementations of HandoomHorisomalPPing When evaluating, we use the FFCV implementation of CenterChopPGBBmmggeeedd with a ratio of 224/256, resizing image crops to a resolution of 224 x 224.
Specifications of baseline models. As a baseline, we Irain a ViT-B/32 model (the implementation of Ilharco et al. (2021)) from scratch on ImageNet. We run AdamW for 100 epochs, using a cosine learning rate schedule with a peak learning rate of 0.003 and 10 warmup epochs, a batch size of 512, a weight decay of 0.1, label smoothing of 0.1 and gradient clipping at global norm 1. To establish a baseline for effective robustness, evaluate this model at epochs 50 through 85 (we stop at 85 because the model's accuracy
22

---

Published in Transactions on Machine Learning Research (01/2025)
later epochs becomes highly correlated). Miller al. (2021) observe that evaluating model trained from scratch at different epochs in this way often exhibit strong linear relationship between their accuracies on the reference and shifted distributions (and the same relationship holds for models with different architectures, hyperparameters, etc.).
Specifications of pre-trained models and fine-tuning strategies. We consider two different pre- trained models: a CLIP (Radford et al., 2021) ViT-B/32 (the implementation of Ilharco et al. (2021)) and AugReg (Steiner et al., 2021) (the implementation of Wightman (2019)). For the AugReg model, we consider full fine-tuning (FT) and linear probing followed by full fine-tuning (LP-FT) (Kumar al., 2022). We perform linear probing by running AdamW for 4 epochs, using a cosine learning rate schedule, peak learning rate of 0.001, a batch size of 512, and without weight decay or gradient clipping. For the CLIP model, we consider zero-shot initialization followed by full fine-tuning (ZS-FT) in addition to these two strategies. We perform zero-shot initialization following Wortsman et al. (2021).
We fully fine-tune models by running AdamW for 8 epochs, using cosine learning rale schedule with 1 warmup epoch. We select the best peak learning rate (in terms of reference accuracy) among 3 x 10-4, 1 x 10-4,3 10-5,1 x 10- 3xxx 10-6, x 10-6, We use batch size of 512, weight decay of 0.1, and gradient clipping at global norm 1.
B.4 Dividing natural shifts into in-support and out-of splits
B.4.1 Splitting Shifted Dataset
To split a shifted datasel into an "in-support split" " and an "out-of-support split", we would ideally measure the reference distribution probability density Pref of inputs in the shifted dataset and assign inputs with small Pref to the out-of-support split. Unfortunately, it is difficult to estimate Pret directly when dcaling with high-dimensional inputs (in this case, images). Instead, we estimate the probability density ralio Pref Prot/Pahift, that is, how much more likely an imput is under the reference distribution than under the shifted distribution. We then assign examples in the shifted dataset with Prei/P < 0.2 to the out-of-support split and examples with Pret/ Pshift 2 0.2 to the in-support split. We visualize examples in Figure 13.
To estimate Pret/ Pret/Pshife, Pshitts we use a classifier trained to distinguish between examples from the reference and shifted datasets. Specifically, let p be a probability mass/ /ensity function over examples that can either be drawn from Dzet Dahint (i.e., p represents the distribution of a dataset created by joining a reference dataset and a shifted dataset). Next, let Uret be the event that an example is drawn from Dret and Yshift. be the event that an example is drawn from Dehift We can express the ratio Pref/ Pshitt as follows:
Estimating
(r)
p(x()ret)
Pahin (x)
p(a ((aly))
p(yret|x)) p(z).
p(yahitt)
plyret) pl (x) -p(x)
p(yredlz) p(yama)
p(yahin( (x) p(yret)
The terms p(yret) and p(yanin) are easy to estimate since they are simply the proportions of reference and shifted examples p. Hence, to estimate Pret/Pshiat we just need to estimate p(yeetar) and p(ysta//)).
To do so, train classiffer to distinguish between reference and shifted examples from p. We construct such a dataset by combining 100K samples from ImageNet with each of the shifted datasets (for ImageNet-R, which contains a subsct of the classes of ImageNet, we restrict the 100K samples to these classes). Next, we fine-tune a CLIP ViT-L/14 pre-trained on LAION-2B from OpenCLIP (Ilhacco et al., 2021) to distinguish between reference and shifted examples. We first fine-tune just the final layer with a learning rate of 0.1 and then fine-tune the entire model with the best learning rate selected from 2 10-4, 1 10- 10-4, 5 x 10-5, 2 x 10-5, 3, 10-5, 5 10-6, x 10-6 and 1 x 10-6 After training the classifier, we calibrate it through temperature scaling (Guo et al., 2017). We then estimate p() and p(ysinl)z))
dataset drawn
23

---

Published in Transactions on Machine Learning Research (01/2025)
by applying sigmoid to its output, from which can estimate Pret/ Pref/Phhift. To estimate this ratio for the entire shifted dataset, we split the dataset into 10 folds and train classifier to estimate Prot/ Pshift on each fold using the remaining 9 folds.
Calibrating the classifiers used for splitting As discussed in Section B.4, our method for dividing a shifted dataset into an in-support split and an out-of-support split requires a calibrated classifier distinguish between examples from the reference and shifted datasets. Recall that to distinguish between examples from the reference and shifted datasets, finc-tune CLIP Radford et al. (2021) ViT-L/14 Dosovitskiy et al. LAION-2B from OpenCLIP (Ilharco et al., 2021). Such over-parameeeerized models
(2021) pre-trained can be overconfident in their predictions (and thus uncalibrated), so we calibrate the classifier by rescaling its (logit) output, a method known as temperature scaling (Guo et al., 2017).
In particular, let I a (potentially uncalibrated) classifier trained to distinguish between examples from the reference and shifted datasets (where the output of f is a logit). We find the scaling parameter a that minimizes the standard logistic loss of on calibration set Seal:
Σ (Sea)
log(1 + e -
a arg min a'
We then define rescaled classifier fcm(()) a f(x) (which is used to estimate the ratio Pret/ Pret/Pahin). Patin). We produce calibration curves of the rescaled classifiers for each of the shifted datasets we split (see Figure 9) and observe that they are indeed well-calibrated.
ImageNet-V2
ImageNet Sketch
ImageNet-R
1.0
1.0
1.0
0.8
0.8
0.8
Positive rate
0.6
0.6
0.6
0.4
0.4
0.4
0.2
0.2
0.2
0.0
0.0
0.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Average predicted probability
Average predicted probability
Average predicted probability
y= (perfectly calibrated)
+ Our model
Figure 9: Calibration curves of classifiers used for splitting. We display calibration curves for the classifiers used to divide ImageNet-V22 V2, ImageNet-Skette and ImageNet-R, into in-support and out-of-support splits. Specifically, we sort the outputs of each classifier on combined dataset of reference and shiftod examples into 100 bins (where bin edges are quantiles). For each bin, compute the actual positive rate (i.e., the proportion of examples from the shifted dataset) and the average predicted probability of example being from the shifted datasct. When we plot the actual positive rates against average predicted probabilities, they are close to equal (close to y = x), suggesting that the classifiers are well-calibrated. Error bars denote 95% Clopper-Pearsee confidence intervals.
B.4.2 Specifications of ImageNet models
To measure the robustness benefits of pro-training on in-support and out-of-support splits of ImageNet distribution shifts, we use the same suite of ImageNet models from PyTorch Image Models (Wightman, 2019) as detailed in Appendix B.2.
24

---

Published in Transactions on Machine Learning Research (01/2025)
B.5 Combining pre-training with interventions for handling bias
Shared model specifications. When training on the WILDS-FMoW dataset, we use the FFCV imple- mentation of HandomdHorizonta/llp.
Specifications of models trained from scratch. We train models from scratch by running SGD for 64 epochs, using a triangular learning rate schedule with a peak learning rate of 0.2 and warmup epochs, a batch size of 128, a weight docay of 10-4 and a momentum of 0.9.
Baseline specifications. To establish a baseline, we train 100 ResNet-50 models from scratch on ran- dom subsets ranging from 25% of the reference dataset to the entire datasct. We increase the number of epochs and warmup epochs inversely with the size of the subset. Miller et al. (2021) observe that models trained from scratch in this way often exhibit a strong linear relationship between their accuracies on the reference and shifted distributions (and the same relationship holds for models with different architectures, hyperparameters, etc.).
Specifications of pre-trained models. The pre-trained model in this experiment is a CLIP ResNet-50 model (the implementation of Ilharco et al. (2021)), adapted using linear probing followed by full fine-tuning, Note that the CLIP ResNet-50 architecture (Radford et al., 2021) deviates from the standard ResNet-50 architecture of He et al. (2015). We perform linear probing by rumning AdamW for 8 epochs, using cosine learning rate schedule, peak learning rate of 0.001, a batch size of 512, and without weight decay or gradient clipping. We fine-tune models by running AdamW for 16 epochs, using a cosine learning rate schedule with a peak learning rate of x 10-4 and warmup epochs, a batch size of 512, a weight decay of 0.1, and gradient clipping at global norm 1.
Our implementation of Deep Feature Reweighting. The Deep Feature Reweighting (DFR) interven- tion proposed by Kirichenko et al. (2022) aims to improve the robustness of model on difficult subpop- ulations by using validation dataset with group labels. The algorithm consists of two steps: (1) train standard model on the original training dataset, and (2) re-train only the final layer of the model (i.e., "ro-weight" the features of the model) on the validation dataset to be more favorable to minority groups. To re-train the final layer, Kirichenko et. al. (2022) repeatedly sample group-balanced subsets of the validation dataset, re-train the final layer on each subset, and then average the resulting re- trainedd final layers. Our im- plementation differs slightly in that we assign sample weights to the validation dataset such that cach group has equal total weight and re-train the final layer on the weighted validation dataset. When applying Deep Feature Reweighting to WILDS-FMMW, we use the out-of-distributiio validation set following Kirichenko et al. (2022).
B.6 Curating datasets for fine-tuning
Image editing to synthesize "counterfactual examples" In order to curate "de-biased" dataset for hair color classification, edit images from CelebA-HQ (Karras et al., 2018), subset of the CelebA dataset with segmentation masks for each attribute provided by CelebAMask-HO (Lee et al., 2020). To change the hair color in a given image, we use InstructPix2Pix (Brooks et al., 2023), a recent image editing model fine-tuned from Stable Diffusion (Rombach et al., 2022). This model accepts an imput image to be edited along with a prompt describing the desired change (e.g., "change the hair color to blond"). We find that InstructPix2Pix able to successfully odit the hair color; however, this model often makes undesired changes to attributes such as skin tone and eye color (see, e.g., the left side of Figure 10). To ensure that we only edit hair color, we use the attribute masks to isolate the pixels in a given image corresponding to the hair region, and ignore any changes made outside of this area. When using a binary mask, this procedure could cause unnatural "edges" along the border of the mask. Thus, we apply a Gaussian blur to the hair mask to smooth the transition when "merging" the original and edited images.
To edit an image from non-blond to blond, we use the prompt "change the hair color to blond." When oditing from blond to non-blond, however, we find that the prompt "change the hair color to non-blond" gives inconsistent results, likely because the instruction is vague. We observe that most non-blond people
25

---

Published in Transactions on Machine Learning Research (01/2025)
in the CelebA datasct have brown black hair, simple heuristic randomly edit each image with either the prompt "change the hair color to brown" or the prompt. "change the hair color to black." See Figure 10 for a visualization of the image editing process.
InstructPix2Pix: "Change the hair color to blond"
Counterfactual Example
Accept changes within segmentation
Unwanted Edits
Reject changes outside of segmentation
Figure 10: Synthesizing counterfactual examples. We edit hair color in CelebA- HQ images using InstructPix2Pix (Brooks et al., 2023). However, this model can also make unwanted changes to attribute other than hair color, c.g., changing eye color (left). To avoid such issucs, in the final image we incorporate only changes within the hair region of the image.
Shared model specifications. Accuracy and worst-group accuracy on the CelebA dataset are sensitive hyperparameter choices. As result, we conduct a grid search select hyperparameters for each type of model. We use classebalanced accuracy as the metric for hyperparameter selection, which empirically better correlates with worst-group accuracy than standard accuracy.
When selecting hyperparameters for curated dataset of a given size, we randomly sample 32 datasets of that size from a pool of 16, 000 images (i.e., 8, 000 CelebA images and their corresponding counterfactual synthesized images) and average the class-balanced accuracies of models trained on each dataset. When evaluating the accuracy and worst-group accuracy of models trained on curated dataset of given size, we similarly randomly sample 64 datasets of that size and report average metrics.
For all models,
the FFCV implementation of Random.HrizznttolFig for data augmentation.
Specifications models trained from scratch. We train ResNet-18 models from scratch by running SGD for 32 epochs, using triangular learning rate schedule with 4 warmup epochs. We use batch size of 128, weight decay of 5 10-1 and momentum of 0.9. We select. the best combination of batch size and learning rate from batch sizes of 64, 128, 256, 512 and learning rates of 0.5, 0.2, 0.1, 0.05, 0.02, 0.01.
When training models from scratch on our curated dataset, we run SGD for 512 epochs and use triangular learning rate schedule with 64 warmup epochs. We use a batch size equal to the total number of examples when it is less than 512 and a batch size of 512 otherwise. We use a weight decay of 5x 10-4 and a momentum of 0.9. We select the best learning rale from 0.5, 0.2, 0.1,0.00,.... 0.05, 0.02, 0.01.
Bascline specifications. To establish bascline, we train 100 ResNet-50 models from scratch on ran- dom subsets ranging from 5% of the reference dataset to the entire dataset. We increase the number of epochs and warmup epochs inversely with the size of the subset. Miller et al. (2021) observe that models trained from scratch this way often exhibit. a strong linear relationship between their accuracics on the reference and shifted distributions (and the same relationship holds for models with different architectures, hyperparameters, etc.)
26

---

Published in Transactions on Machine Learning Research (01/2025)
Specifications of pre-trained models. The pre-trained model in this experiment is CLIP ViT-B/32 model initialized as vero-shot classifier with "blond" and "non-blond" as the class names. We fine-tune models by running AdamW for 16 epochs, using a cosine learning rate schedule with 2 warmup epochs, and a weight docay of 0.1. We select the best combination batch size and learning rate from batch sizes of 64, 128, 256, 512 and learning rales of 10-5, 10-5, 3 x 10-6, 10-6.
When training on our curated dataset, we use a batch size of 64 (the size of the dataset) and select the best learning rate from 3 x 10-5, 1 10-5, 3,3 x 10 10-6 1 x 10-6
27

---

Published in Transactions on Machine Learning Research (01/2025)
C
Additional Results
C.1
Constructing synthetic in-support and out-of-support shifts
C.1.1 How does the choice of fine-tuning hyperparameters affect robustness?
In Section 5.1, we select hyperparameters (in particular, learning rate) for fine-tuning that maximize accu- racy the reference distribution. This reasonably simulates hyperparameter selection in practice because typically only samples from the reference distribution available.
In this scction, we investigate how the choice of hyperparameters affects the robustness of pre-trained models. In particular, we would like to undersland if pre-training yields little effective robustness to in-support shifts and substantial effective robustness to out-of-support shifts across a wider range of hyperparameter choices. We study the spurious tint shift (an in-support shift) and the flip shift (an out-of-support shift)) from Section 5.1 and vary the learning rale, weight decay, number of epochs, and batch size of CLIP ViT-B/32 initialized with zero-shot weights (Figure 11). With zero-shot initialization, the starting point of fine-tuning robust task. Hence, even under an in-support shift, hyperparameter choices that
model that performs well do not change the model substantially (e.g., low learning rate, small number of epochs) result in substantial iffective robustness. However, these hyperparameter choices generally result in lower absolute reference and shifted accuracies, and are thus unreasonable. The hyperparameter choices that are relevant in practice are those with high reference accuracy, and these are the hyperparameters that we use in our experiments.
28

---

Published in Transactions on Machine Learning Research (01/2025)
Varying learning rate
Varying batch size
Baseline Baseline linear fit
80
Baseline
Shiftad accuracy
CLIP zero-shot init.
Shilted accuracy
Baseline linear fit
y=x
CLIP zero-shot init.
0.0003
y ==
60
0,0001
60
(used)
128
50
1e-05
40
3e-06
512 (used)
1c:06
60
70
80
70
Reference accuracy
Reference accuracy
Varying number of epochs
Varying weight decay
Baseline
Baseline
Baseline linear
Baseline linear
80
Shittd accuracy
Shitte accuracy
70
y=
70
y
(used)
60
60
0.01
0,001
8 (used)
0,0001
40
16
40
1e-05
32
0.0
70
Reference accuracy
Reference accuracy
(a) In-support shift. The in-support shift consider is the "spurious tint shift" in which we introduce tint that spurionsly correlated with the label. this in-support shift, learning rate and number epochs influence effective robustness, but the best hyperparameter choices result in model with little effective robustness.
Varying learning rate
Varying batch size
80
Baseline
Baseline linear fit
Baseline
Shittd accuracy
70
CLIP zero-shot init.
Shifte accuracy
70
Baseline linear zero-shot init.
60
y=
0, 0,0003
y
50
50
40
3o-05
40
128
1e-05 (used)
256
30
3e-06
512 (used)
1e-06
Reference accuracy
Reference
Varying number of epochs
Varying weight decay
80
Baseline linear
linear
Shifte accuracy
70
CLIP zero-shot init.
Shitte accuracy
zero-shot init.
60
y
60
0.1 (used)
50
50
0.01
+
4
+
0.001
8 (used)
0.0001
30
16
30
1c-05
32
0.0
50
60
80
50
70
80
Reference accuracy
Reference accuracy
(b) Out-of-support shift. The out-of-support shift we consider is the "flip shift" in which we pad images in the shifted distribution. On this out-of-suppor shift, batch size most significantly affects robustness, while learning rate and number of epochs affect overall performance.
Figure 11: The effects of hyperparameter choices on robustness. We vary hyperparameters when fine-tuning a CLIP ViT-B/32 initialized with zero-shot. weights on synthetic ImageNet shifts from Section 5.1 (different shades green). Varying certain hyperparameters (e.g., learning rate, number of epochs) can affect the effective robustness of pro-trained models even on an in-support shift. In our experiments, wo choose hyperparameters which yield high reference accuracy (purple).
29

---

Published in Transactions on Machine Learning Research (01/2025)
C.1.2 How does the strength of the bias affect robustness to in-support shifts?
In Section 5.1, we consider two in-support shifts under which models might fail due to dataset biases. In particular, in the spurious tint. shift, introduce tint that is spuriously correlated with the label in the reference dataset, but not in the shifted dataset. The probability that an example in the reference dataset has a class-specific lint (as opposed to random tint) is determined by parameter Pspurious (set to 0.5 for the experiments in Figure 4). In the label shift, the relative frequencies of classes change between the reference and shifted datasets. The classes are divided into "majority" and "minority" classes, with "minority" classes appearing with probability Pminority in the reference dataset (set to 0.2 for the experiments in Figure 4). the shifted distribution, the relative frequencies of the classes are reversed.
In this scction, investigate how the strength of the bias, i..,, Pspurious and Puinority affects the robustness of pre-trained models to these in-support shilts. We observe the average effective robustness of pre-trained models largely remains close to zero as we vary these parameters (see Figure 12).
Pspurious 0.5
Pa urious
Pspurious
Pepurious
Pspurious 0.9
ER 0.34 0.87%
1.11%
0.84
1.51%
ER 1.66 1.1%%
ER 7.23 2.62%
80
80
80
70-
80
Shifted
60
60
60
50-
40
20
70
60 70
80
60
70
8D
90
70
80
90
70
90 Reference
Reference
Reference Accuracy
Reference Accuracy
Reference Accuracy
Baseline
y
(FT)
(LP-FT)
(2S:FT)
AugReg (FT)
AugReg (LP.FT)
(a) Spurious tint shift. Across several different probabilities of the spurious class-specifie tint (Pepurions). the average effective robustness of pre-trained models (top left of each plot) on the in-support "spurious tint shift" is close to zero. The one exception is the shift with Pspurious 0.9, where the effective robustness is higher. This may be because this shift is "close" to an out-of-support shift, since the probability of ohserving an example with a random tint (as opposed class-specific tint) is low, Hence, pre-training might help by extrapolating better from the small number of randomly tinted examples.
Pminonity 0.1
Ominoniy 0.15
Pminanit
Pminonty
Pminorky 0.3
ER 33.63 2.3.63222778 2.73%
ER 2.66%
80
ER =
0.2 0.82%
80
ER 0.300
0.89%
80
ER -0.31
Shited Acouracy
70
70
70
70
60
60
60
60
60
50
40
40
30
40
50
70
80
60
70
80
50
60
70
80
50
50
60
70
80
Reference Accuracy
Reference Accuracy
Reference Accuracy
Reference Accuracy
Reference Accuracy
Baseline
y=x
CLP (ZS-FT)
AugReg (FT)
AugReg (LP-FT)
(b) Label shift. Across several different probabilities of the minority classes (Pmimority). the average effective robust- ness pre-trained models (top left of each plot) on the in-support "label shift" is close to zero. We note that the shifts with Paiicrity 0.1 and Painority 0.15, the effective robustness is slightly negative. However, the linear correlation among baseline models weak under these shifts, so these effective robustnesses are less meaningful.
The effects of the strength of the bias on robustness to in-support shifts. We vary the strength of the bias of the two synthetic ImageNet in-support shifts from Scction 5.1. Broadly, the effective robustness of pre-trained models (top left of each plot) is close zero across bias strengths.
Figure 12:
30

---

Published in Transactions on Machine Learning Research (01/2025)
C.2 Dividing natural shifts into in-support and out-of-support splits
C.2.1 Sizes of in-support and out-of-support splits
In Table 1, we report the sizes of the in-support and out-of-support splits we compute for ImageNet- V2, ImageNet Sketch and ImageNet-R. The out-of-support splits are much larger than the in-support splits, perhaps because the large majority of the examples from these shifted datasets look unlike examples from ImageNet.
Table 1: Sizes of in-support and out-of-support splits.
Dataset
In-support split size Out-of-support split size
ImageNet-V2
1920
8080
ImageNet Sketch
162
50727
ImageNet-R
588
29412
C.2.2 Examples from in-support and out-of-support splits
In Figure 13, we provide samples from the in-support and out-of-support splits we compute for ImageNet- V2, Imagehet-Sketch and ImageNet-R.
31

---

Published in Transactions on Machine Learning Research (01/2025)
ImageNer
podestal
suiphur butterfly
broom
Volre
Bernese
gyromitra
Kingsnake
popsicle
V2 In-support
great white
leystick
cabbaae
English Setter
wallet
construction crame
apron
Dobermann
geyser
wild boar
Out- Put-otsuuport
V2
cello
Fox squitrel
pole
freight car
Lape player
frying
low truck
plate rack
eraser
Dealin glon Torrier
Skerch In-support
feather boa
whale
Steamm locomotive
Tallet
couch
Siamese cat
airoraft carrier
ciiif dielling
through arch bridge
rotary dial telephone
Skeetch Out-of-support
Tootball helmet
Arctic
Bile roof
machine
turtle
In-support
aint yo
MannalClean
your
pig
Highland
Collie
strawberry
grasshopper
scorpion
Out-of-uupor
scorpion
revolee
frog
banana
pufferfish
Figure 13: Random samples from ImageNet and from the in-support and out-of-support splits of ImageNet- V2, ImageNet Sketch and ImageNet-R. In ImageNet-V2, V2, it is difficult to distinguish between examples from the in-support and out-of-support splits. In ImageNet Sketch and ImageNet-R, examples from the in-support splits look more realistic (i.e., more like ImageNet examples) than examples from the oft-of-support splits.
C.2.3 Scatter plots of reference vs. shifted accuracy
In Figure 14, we provide scatter plots of accuracy on ImageNet vs. accuracy on the in-support and out-of- support splits of ImageNet-Ve V2, ImageNet Sketch and ImageNet-R.
32

---

Published in Transactions on Machine Learning Research (01/2025)
ImageNet-VA
ImageNet Sketch
ImageNet-
ER
0.17% %
ER
1.56%
ER 1.78%
85
85
in-support accuracy
80
90
80
75
70
80
70
65
70
60
60
55
50
50
75
90
12.30%
13.04%
Out-of-suppon accuracy
85
80
90
80
70
80
60
70
50
60
40
50
65
40
30
60
30
55
20
20
75
80
87.5
90
92.5
95
97.5
ImageNet Accuracy
ImageNet Accuracy
ImageNet (class-subsamplec) Accuracy
Baseline
Baseline linear fit
Pre-trained
y=
Figure 14: Reference vs. shifted accuracy for in-support and out-of-support splits of ImageNet shifts. On each of the three ImageNet shifts we consider, the average effective robustness (ER) pre-trained models (orange) above the bascline of models trained from scratch (blue) on the in-support split (top) is small. Meanwhile, their effective robustness can be very large on the out-of-support split (bottom).
C.2.4 Controlling for difficulty when measuring effective robustness
The significance of a given effective robustness depends on the "difficulty" of' distribution shift. For example, if shift causes an accuracy drop of 5%, an effective robustness of 4% might be considered large, but if a shift that causes drop of 25%, effective robustness of 4% would probably be considered small. When we divide shifted dataset into in-support and out-of-support split, the out-of-support split is typically more difficult than the in-support split. Tf we compare the effective robustness of pre-trained models examples of similar difficulty in the in-support and out-of-support splits, do our findings from Section 5.2 still hold? In particular, do pre-trained models still exhibit substantially higher robustness on out-of-support examples than on in-support examples?
To answer this question, we re-weight examples in out-of-support splits such that the difficulty distribution of the out-of-support split matches that of the in-support split. Specifically, quantify the difficulty of given example in terms of the fraction of baseline models (of 77 total baseline models) that classify it incorrectly. Given an example of difficulty d, we re-weight it by factor of Pin-support Poupootf---ppprr (dpputtttpp//ttttttssspppprr (d)
where Pin-support. is the difficulty probability density function the in-support split and Pont-nf-support is the difficulty probability density function of the out-of-support split. We then compute a "re-weighted" accuracy, which in turn yields a re-weighted effective robustness, on the out-of-support split. Intuitively, this re-weighted effective robustness represents the effective robustness of pre-trained models on out-of-support examples of similar difficulty to in-support examples.
33

---

Published in Transactions on Machine Learning Research (01/2025)
We report the re-weighted effective robustnesses in Figure 15. We observe that the re-weighted effective robustnesses of pre-trained models on out-of-support splits are indeed lower than the original effective ro- bustnesses. However, they are still substantially higher than the effective robustnesses on in-support splits.
0.14
Average effective robustress
0.12
0.10
0.08
0.06
I
0.04
0.02
I
0.00
ImageNet-V2
ImageNet Sketch
ImageNet-R
Out-of-support
Out-of-support (re-weighted)
In-support
Figure 15: Re-weighed effective robustness of pro-trained models on in-support and out-of- support splits of ImageNet shifts. When we re-weight examples in out-of-support splits to match the difficulty distributions of their corresponding in-support splits, the average effective robustnesses of pre- trained models (green) decrease relative to the original effective robustnesses (blue). However, they are still very high on ImageNet. Sketch and ImageNel-R. Meanwhile, the average effective robustnesses of pre-trained models on in-support splits (orange) are consistently low.
C.3
Combining pre-training with interventions for handling bias
C.3.1 Studying a synthetic shift
In this section, we provide an additional experiment in synthetic setting to further illustrate that pre- training and interventions designed to handle dataset. biases can be complementary. In Section we dis- cussed how robustness to the WILDS-FMoW distribution shift requires both extrapolating to later years and performing consistently across regions. We construct synthetic distribution shift using that similarly requires both extrapolating well and avoiding reliance spurious features. Specifically, we combine the tint and pad shifts from Section 5.1. We modify CIFAR-10 such that in the reference distribution, we add a tint that is spuriously correlated with the label: 80% of reference examples have a class-specific tint while the remaining 20% are randomly tinted. Meanwhile, in the shifted distribution, examples are always randomly tinted and are also padded (we add black pixels each side of the original 32 32 CIFAR-10 images).
To extrapolate to padded examples, initialize a CLIP ResNet-50 and perform linear probing followed by full fine-tuning on the reference distribution. To handle the spurious correlation between tint and label, we consider the intervention of training on randomly tinted examples, which we refer to as balancing. This is an "oracle" sorts for handling dataset biases; simply modifies the training distribution such that spurious features are not useful.
As with WILDS-FMoN we find that pre-training and balancing each yield some effective robustness (see the left side of Figure 16). In this case, combining the two does not yield the greatest effective robustness, but does have the highest shifted accuracy. We apply the same methodology as in Section 6 to understand the robustness benefits of pre-training and balancing. Here, we observe a greater overlap between the corrected examples of pre-training and balancing than we did for pre-training and DFR in the case of WILDS-FMoS FoW (see the right side of Figure 16). This may be due to the fact that every example requires both extrapolation and avoiding reliance on the spurious bias. In other words, the failure modes that pre- training and balancing are intended to address cooccur. However, we note that there are still many examples that corrected by one of pre-training and balancing, but not the other, suggest complementary benefits.
34

---

Published in Transactions on Machine Learning Research (01/2025)
Similarly to our observations with WILDS-FMoW! combining pre-training with balancing corrects most of the examples corrected by the individual interventions. These results corroborate our finding that pre- training and interventions designed to handle dataset biases can be complementary.
Effective robustnesses
Overlap of corrected examples (% of shifted dataset)
95
ER = 11.50%
Shifted accuracy
0.2%
0.7%
90
ER 15.23%
1.82%
2.3%
6.8%
2.6%
85
80
1.6%
94
95
96
97
Reference accuracy
Baseline
Baseline linear fit
y =
Balanced
Pre-trained
Pre-trained and Balanced
Figure 16: Combining pre-training and balancing on a synthetic CIFAR-10 distribution shift. Pre-training and balancing (an "oracle" intervention for handling dataset biases) each yield some effective robustness (ER) and combining these two interventions yields a high effective robustnoss and the highest shifted accuracy (left). substantial number of examples are corrected by one of pre-training and balancing, but not. the other (right), indicating that there are different subpopulations where they improve performance. Meanwhile, the examples corrected by combining pre-training with balancing include most of the examples corrected by the individual interventions (right), suggesting that combining pre-training with balancing improves performance both of these subpopulations. Error bars denote 95% confidence intervals over 64 random trials.
C.4 Curating datasets for fine-tuning
C.4.1 Understanding the robustness benefits of pre-training when fine-tuning on curated dataset
In Section 7, we find that fine-tuning curated dataset with only 64 examples can yield performant and robust model for hair color classification. We observe that pre-training is necessary for effective use of the small curated dataset; in particular, training a model from scratch on a curated dataset yields robustness gains, but these gains are smallor and many more examples are required to attain comparable accuracy.
In this section, we shed additional light on how pre-training helps in this setting. Based on our intuition from Sections and that pre-training helps specifically with extrapolation, we hypothesize that pre-training provides two benefits when training small curated dataset. First, pre-trained model may be able to extrapolate better from small number of examples. This would result in both higher accuracy on the original CelebA distribution and higher worst-group accuracy, which we observe in Figure 7b. Second, recall that our curated dataset consists entirely of females, but hair color classification models are expected to perform well on males too. To compare different model's ability to extrapolate along this axis, we plot the balanced accuracy on males against the balanced accuracy on females. In Figure 17, we observe that the pre-trained model indeed generalizes better males than models trained from scratch.
35

---

Published in Transactions on Machine Learning Research (01/2025)
Female balanced vs. male balanced accuracy
Male balanced accuracy
85
80
75
78
80
82
84
86
88
90
Female balanced accuracy
y=x
Curated (n 64, 8192)
Curated (n 64, .... 8192) linear Pre-trained and Curated (n 64)
Figure 17: Comparing extrapolation from females to males of pre-trained models and models trained from scratch. We plot the balanced accuracy males against the balanced accuracy of females of pre-trained model fine-tuned on the curated dataset from Section (red) and models trained from scratch on this dataset (green). Models trained from scratch establish linear relationship between male and female balanced accuracy; however, the pre-trained model outperforms this trend, suggesting that it more effectively extrapolates to males from the femalc-only curated dataset.
C.4.2 Exploring balancing instead of counterfactual editing
In Section 7, choose to curate a dataset by augmenting images from CelebA with "counterfactual ex- amples" in which we edit the hair color to the opposite class. We do so in order to de-bias this dataset as much as possible. In this section, we explore a simpler approach to curating a dataset: balancing classes. Similarly to our curated dataset, we constrain this balanced dataset. to include only females. As with our curated dataset, we observe that fine-tuning pre-trained model on a class-balanced female-only dataset yields robust and performant model for hair color classification (see Figure 18a). We also observe again that pre-training improves over training from scratch by helping with extrapolation from the female-only reference dataset to males (see Figure 18b).
36

---

Published in Transactions on Machine Learning Research (01/2025)
Accuracy vs. worst-group accuracy
Female balanced vs, male balanced accuracy
90
ER 58.04%
90
Worst-group accuracy
Avg. ER 55.58%
70
85
80
Male
70
30
86
77.5 80
82.5
85
87.5
90
92.5
95
Female balanced accuracy
Accuracy
y=
Baseline
Pre-trained Balanced (n 64, 16384)
Balanced (n 64, 16384)
Baseline linear fit
Balanced ... linear Pre Pre-trained and Balanced (n 64)
y-x
Pre-trained and Balanced (n=
(a) Fine-tuning on a balanced female-only dataset. Fine-tuning a pre-trained model on CelebA dataset (orange) yields little effective robustness over baseline of models trained from scratch (blue). However, fine- tuning the same pre-trained model on just 61 examples from balanced female-only dataset (red) yields a model with both high effective robustness and high accuracy. Training from scratch on a balanced female-only dataset (green) also yields high effective robustness, but results in substantially lower accuracy than pre-trained models, with many more examples. Error bars denote 95% confidence intervals over 64 random trials.
(b) Comparing extrapolation from females to males of pre-trained models and models trained from scratch. We plot the balanced accuracy on males against balanced accuracy females of a pre-trained model fine-tuned on balanced female-only dataset (red) and models trained from scratch on this dataset (green). Models trained from scratch establish linear relationship between male and female balanced accuracy; however, the pre-trained model outperforms this trend, suggesting that it more effectively extrapolates to males from the female- only reference dataset.
Figure 18: Finc-tuning a pre-trained model on a small, non-diverse but de-biased datasct (in this case, a balanced femalc-only dataset) yields a robust and performant model for hair color classification in CelebA (see Figure 7b).
37

---

Published in Transactions on Machine Learning Research (01/2025)
D Additional Discussion
D.1 Alternative fine-tuning strategies
In this work, we focus on the common setting in which a pre-trained model is fully fine-tuned. It is important to note that pre-trained models used in a zero-shot context (i.e., without fine-tuning) and partially fine-tuned models (c.g., only the final classification layer is updated) are froquently more robust than fully finc-tuned models (Radford et al., 2021; Miller et al., 2021; Kunar et al., 2022). Such models may have higher effective robustness than fully fine-tuned models or in some cases may even outperform fully fine-tuned models on the shifted distribution. Howover, such models are typically less performant on the reference distribution than fully fine-tuned models.
Several works observe this tradeoff between performance on the reference distribution and robustness and devise methods for mitigating it, i.e, methods for robust fine-tuning (Wortsman et al., 2021; Hewitt et al., 2021; Kumar et al., 2022). For example, Kumar et al. (2022) argue that full fine-tuning "distorts" pre- trained features and propose linear probing before full fine-tuning (LP--FT) to prevent distortion. They also suggest that fine-tuming a model initialized as a zero-shot classifier may have a similar effect. In addition to full fine-tuning, in Section 5.1 we thus consider LP-FT and zero-shot initialization for fine-tuning. On in-support shifts, we observe that LP-FT and zero-shot initialization do not provide effective robustness benefits comparod to full fino-tuning (soc Figure 4), suggesting that these strategics do not help mitigate dataset biases.
Another strategy for robust fine-tuning is to ensemble zero-shot model and fully fine-tuned model. Both weight-space ensembles (Wortsman et al., 2021) and output-space ensembles (Hewitt et al., 2021) have been shown to improve robustness, somelimes even without sacrificing performance on the reference distribution. In fact, this strategy can yield robustness benefits even when dataset biases are a primary failure mode because the zero-shot model is independent of the biased reference dataset. Our work seeks to complement such empirically effective strategics by providing an understanding of when they are necessary. In particular, our findings suggest that ensembling valuable precisely when dataset biases cause failures.
D.2 Can pre-training hurt extrapolation?
In this work, we discuss distribution shifts in which pre-training is beneficial to a model's ability to ex- trapolation outside of the reference distribution. A natural question to consider is whether pre-training can instead hurt it, yielding worse extrapolation than a model trained from scratch. A recent work by Salman et al. (2022) suggests that this is indeed possible. Specifically, they show that biases of pre-trained models can persist during fine-tuning. For example, a model pre-trained on ImageNet and fine-tuned on CIFAR-10 is highly sensitive the presence of tennis balls (which are an ImageNet class but not a CTFAR-10 class). Mcanwhile, a model trained from scratch on CIFAR-10 is not particularly sensitive to tennis balls. Thus, under a hypothetical "tennis ball shift" in which tennis balls appear in images in the shifted distribution, pre-trained model would be less robust than model trained from scratch. In this instance, pre-training provides harmful prior for how to extrapolate.
D.3 When does pre-training help with extrapolation?
In this work, we provide evidence that pre-training can help with extrapolation, but not with other failure modes. A natural question to consider is whether a particular pre-trained model and fine-tuning strategy in fact does help with a given extrapolation task. To this end, Ramanujan et al. (2023) explore how the composition of the pre-training dalaset alfects robustness on the WILDS-WWidCam distribution shift (Koh et al., 2020). We consider further exploration of this question to be a valuable direction for future work.
D.4 Relating in-support and out-of- shifts to existing characterizations
The characterizations relevant in this work, in- in-support shift and out- of support shift, overlap with many existing definitions. Ye et al. (2022) introduce notions of corvelation shift and diversily shift (closely aligned
38

---

Published in Transactions on Machine Learning Research (01/2025)
with in-support and out-of-support shifts, respectively) and provide method for measuring the "amount" of each type of shift in given distribution shift (similar to our method for dividing a distribution shift into in-support and out-of-support splits). Subpopulation shift (and its sub-types), shifts involving spurious correlations, covariate shift, and label shift are typically in-support. Howover, there are exceptions; for example, some works consider subpopulation shifts which subpopulation does not appear in the reference distribution (Santurkar al., 2021; Yang et al., 2023), which are out-of-suppor.. Domain generalization problems are nearly always out-of-support and extrapolating effectively outside of the reference distribution is often a key challenge of these tasks.
D.5 Understanding the robustness of pre-tranned language models to spurious correlations
Tu et al. (2020) study the robustness of pre-trained language models to distribution shifts with spurious correlations. Their central finding is that pre-training can improve performance on shifted datasets in which spurious correlations do not hold. They illustrate that this is because pre-trained models generalize better from the small number of counterexamples to these correlations in the reference dataset. This is a similar phenomenon to our observation from Figure 12a: pre-training can provide some effective robustness on in-support shifts that are "close" to an out-of-support shift. In cases such as those discussed by Tu et al. (2020), we hypot that pre-training can help to a limited exlent by extrapolating better, but cannot mitigating the underlying failure mode of dataset biases.
D.6 Additional related work
Pre-training. Pre-training a model (or taking an existing pre-trained model) and then fine-tuning it on a task-specific dataset common practice when developing machine learning modcls, often significantly improving performance over training a model from seratch (Sharif Razavian et al., 2014; Sun al., 2017: Kornblith et al., 2019; Kolesnikov et al., 2019). Pre-training is unrelated to the pre-training task, suggesting that pre-training yields useful general-purpose features; for example, object classification models trained on ImageNet (Deng et al., 2009) are good initializations for remote sensing (Xie et al., 2016) and medical imaging (Ke et al., 2021) tasks. Although greatly effective, pre-training is not without limitations, In some settings, pre-training does not improve performance over a randomly initialized model trained for long enough (He et al., 2019). Downstream performance can saturate as performance on the pre-training task improves (Abnar al., 2021). Finally, biases of pre-trained models can persist after fine-tuning (Salman et al., 2022).
be effective even when the downstream task
Distribution shift robustness. Machine learning models are often deployed in different environments from those in which they are trained. Such distribution shifts can cause models to significantly underperform (Koh et al., 2020; Gulrajani & Lopez-Paz, 2020; Hendryeks et al., 2020a). Numerous interventions have been proposed to improve the robustness models, often targeting particular types of shifts. These include algorithmic interventions (Arjovsky al., 2019; Byrd & Lipton, 2019; Sagawa et al., 2020a; Liu et al., 2021; Kirichenko ct al., 2022; Idrissi al., 2022) (often requiring group information), data augmentations (Hendrycks al., 2020a; Goel et al., 2020) and pre-training (discussed below). However, interventions proposed thus far have failed to provide consistent benefits across distribution shift benchmarks (Koh et al., 2020; Guirajani & Lopez-Paz, 2020; Hendrycks et al., 2020a; Wiles et al., 2021; Ye et al., 2022), rendering distribution ion shift robustness a persistent challenge.
39


Reviews:
[{"review_id": "tBFnnjk4VE", "reviewer": "Reviewer_eCVF", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**:\n\n- The paper is clearly written and tackles a question of broad interest to the machine learning community.\n- The authors present a straightforward yet illustrative toy example using logistic regression to explain how pre-training can—and cannot—address distribution shifts. Notably, they show that when the training data lies on a low-dimensional subspace, shifts that move data points out of this subspace can either hurt or leave the model unaffected, depending on the initialization.\n- The core message of the paper is compelling and could have useful implications for the community. However, the sufficiency of the experimental evidence to support these conclusions is debatable.\n\n**Weaknesses**:\n\n- The paper lacks theoretical guarantees, which, while challenging to provide in this context, weakens the impact of the claims. It relies almost entirely on experiments to substantiate its arguments.\n- The experimental results may not fully support the authors' broad claims. For example, the paper seems to equate \"in-support\" shifts with learning spurious features. Is this always the case, or could there be other types of in-support shifts? Without theoretical guarantees, it is unclear whether the authors' conclusions generalize beyond their specific experiments.\n- The dataset refinement process described in Section 5.2—such as the removal of sketches from ImageNet—aims to isolate \"out-of-support\" shifts from \"in-support\" shifts. However, this manipulation may undermine the naturalness and generalizability of the results, as the datasets have been altered to fit the proposed scenario more cleanly.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper addresses an important question in machine learning: Why does pre-training a model on a large dataset followed by fine-tuning for a downstream task sometimes enhance robustness to distribution shifts, while in other cases, it does not? Specifically, the authors explore whether there is a fundamental barrier that prevents models from gaining robustness to certain types of distribution shifts. Their answer is \"Yes.\" They argue that pre-training improves robustness when shifts require extrapolation but is less effective when the issue is related to other factors, such as dataset biases. In other words, pre-training helps when the shifted distribution includes “out-of-support” inputs—inputs that couldn't reasonably be sampled from the reference distribution—thus aiding extrapolation beyond the original training data.\n\nThe authors substantiate this claim primarily through a series of diverse experimental validations, with some theoretical exploration in the context of logistic regression. However, the paper lacks rigorous proofs or substantial theoretical results. Experimentally, the authors create two different scenarios: \n\n1. **In-support shifts**: where the training and test datasets share the same support but differ, for instance, in density. In these cases, pre-training often fails to confer additional robustness. \n2. **Out-of-support shifts**: where the test data contains instances that the training distribution cannot generate, requiring extrapolation. Here, pre-training has been shown to enhance robustness.\n\nThe paper also extends its findings by proposing techniques to improve robustness to distribution shifts. The authors suggest combining pre-training with existing intervention methods aimed at mitigating dataset biases, arguing that these approaches complement one another. They demonstrate that fine-tuning on a small debiased dataset can sometimes result in significantly greater robustness than training from scratch on a much larger, biased dataset.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please see Weaknesses.", "broader_impact_concerns": "I do not have any concerns regarding this issue.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "DLxjXV80ux", "reviewer": "Reviewer_9dJB", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n1. The paper studies an interesting question of whether pre-training can help improve robustness under two different types of distribution shifts.\n\n2. Both experiments and theoretical analysis are provided to support the findings in this paper.\n\n\n### Weaknesses\n\n1. The paper’s theoretical analysis is on simplified logistic regression models, there lacks a discussion of connection to the more complex deep models.\n\n2. The definitions of in-support and out-of-support distribution shift is vague to me. Can there be formal mathematical definitions of these shifts? Can the authors then relate the shifts with Theorem 4.1? In other words, how does Theorem 4.1 lead to the conclusion on how pre-training can and cannot mitigate certain types of shifts?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper studies the robustness of pre-trained models under distribution shifts. Specifically, the authors focus on two possible failure modes under distribution shift: \n1. poor extrapolation (e.g., they cannot generalize to a different domain);\n2. biases in the training data (e.g., they rely on spurious features).\n\nThrough empirical evaluations and theoretical analysis of a logistic regression model, the paper concludes that pretraining can mitigate poor extrapolation but not bias in the training data.\n\nTo develop robust models, the authors further suggest that \n1. bias mitigation techniques can further improve the robustness of pre-trained models;\n2. fine-tuning on a small, non-diverse but de-biased dataset leads to more robust models than fine-tuning on a large diverse but biased dataset.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. The authors should discuss how the simplified model in the theoretical analysis relates to the large model used in empirical experiments. Specifically, the considered logistic regression problem is a simple convex problem, with a deterministic training procedure. This model does not carry some key properties of the deep models which require stochastic pre-training and do not necessarily converge to global optimal.\n\n2. The experiments in the paper focus on specific tasks (e.g., image classification), whether studies on these tasks and datasets generalize to different tasks, domains, and model architectures remains a question to be discussed.", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "5vsykef8yn", "reviewer": "Reviewer_EzeB", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pro:\n1. This paper is well-motivated. It starts from a plain fact and analyzes the phenomenon through conducted experiment. Then they propose the main problem and give their solution.\n2. The conducted experiment is detailed and rigorous. The authors try to provide both theoretical and empirical justification for all the claims.\n\nCon:\n1. The logistic regression setting for studying pre-training is confusing. I understand what the authors aim to illustrate, but could the setting be general for all distribution shift scenarios?\n2. Assume two data distributions A and B. Denote the sample as X, and label as Y. For, $P_A(X, Y)=P_A(X)P_A(Y|X)$ and $P_B(X, Y)=P_B(X)P_B(Y|X)$, covariant shift assumes that $P_A(Y|X)$ and $P_B(Y|X)$ are identical. The proposed definition of in-support and out-of-support distribution shift seems to be the same as covariant shift.\n3. I am confused about Theorem 4.1. What is the purpose of coming up with this theorem?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper starts with the phenomenon \"pre-training can significantly improve robustness to some distribution shifts but not others\" and focuses on characterizing the distribution shift. Through a series of experiment, the authors try to explain why the pre-trained model improves robustness under some types of shift while compromise under others. Finally, they propose two manner to develop robust models, i.e., combining pre-trained models with designed bias and fine-tuning on curated datasets.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The motivation of this paper is interesting and this direction is potential. The proposed manners for developing a robust model based on pre-training are reasonable. However, the organization of this paper and some claims are quite confusing. I advise the author to clarify their statement and the story they want to express.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper examines how pre-training affects robustness to distribution shift, finding situations where it helps and others where it does not. It investigates these phenomena through experimentation and analysis, and explores the implications for robust model development. After some discussion with the authors, the reviewers found that any concerns had been addressed and the paper is suitable for publication at TMLR.
Decision: Accept (poster)