A Fine-Grained Analysis on Distribution Shift

Introduction

• The work addresses limited understanding of when and why models generalize, despite examples of cross-dataset generalization and concerns about existing domain-generalization methods.Robustness matters because distribution shifts can cause unexpected failures in critical applications and undermine trust.Source: §1 · p. 1
• The framework treats data as combinations of attributes and expects models to learn invariance that supports generalization to unseen attribute values or changed attribute distributions.For example, a model trained to predict shape from red and blue objects should generalize to green objects or a different color balance.Source: §1 · pp. 1–2
• The framework evaluates spurious correlation, low-data drift, and unseen data shifts, alongside varying label noise and training-set size.These shifts are presented as real-world-inspired building blocks for more complex shifts.Source: §1 · p. 2
• The evaluation compares 19 methods across five approaches and more than 85K trained models, using synthetic and real datasets.The framework can be extended to evaluate additional shifts, datasets, or methods.Source: §1 · p. 2
• Augmentation and pretraining are often effective, but the best approach cannot be selected a priori and results vary across datasets and attributes.The authors also report that domain generalization algorithms work for certain datasets and shifts, while emphasizing that improving real-world robustness remains an open challenge.Source: §1 · p. 2

2.1 LATENT FACTORISATION

• The framework treats train and test data as sharing a conditional generative model while their attribute marginals may differ, and seeks low risk on the test distribution using training data.Robustness strategies include attribute-weighted resampling, synthetic data augmentation, and unsupervised representation learning with a task-specific head.Source: §2.1 · pp. 3–4

2.2 DISTRIBUTION SHIFTS

• The framework defines spurious correlation, low-data drift, and unseen data shift as changes in attribute relationships or coverage, which respectively affect p(yl|ya) and p(ya).Unseen data shift is a special case of low-data drift, but requires multiple attribute values to be observed in training; together, the shifts can compose distribution changes for label and nuisance attributes.Source: §2.2 · pp. 4–5
• CYCLEGAN performs consistently best for spurious correlation, while ImageNet pretraining performs consistently best for low-data drift; ImageNet augmentation also boosts performance in both settings.For low-data drift, most domain-generalization methods also provide some boost; the figures average percentage changes over seeds and datasets.Source: Figure 3 · p. 5, Figure 4 · p. 5

2.3 CONDITIONS

• The framework separately studies label noise by corrupting observed labels and dataset size by limiting the total number of training samples.These conditions model labeler disagreement or error and settings where collecting more data is expensive or unrealistic.Source: §2.3 · p. 5

Architecture choice

• The architecture comparison covers ResNet18, ResNet50, ResNet101, ViT, and an MLP, with performance depending on representation robustness to distribution shift.The setup also uses weighted resampling preweight to oversample regions less likely to be sampled under ptrain.Source: §3 · p. 6

Heuristic data augmentation

• Heuristic augmentation methods are evaluated as approximations of p(x|y1:K): ImageNet augmentation, AugMix without JSD, RandAugment, and AutoAugment.Performance depends on how well the heuristic augmentations approximate the true generative model.Source: §3 · p. 6

Learned data augmentation

• Learned augmentation uses nuisance-conditioned transformations to change an image’s attribute while keeping others fixed; the evaluation follows Goel et al. (2020) with CYCLEGAN but omits their SGDRO objective.This isolates learned data augmentation, whose performance depends on how well it approximates the true generative model.Source: §3 · p. 6

Domain generalization

• Domain generalization evaluates IRM, DeepCORAL, domain MixUp, DANN, and SagNet to learn representations z independent of the attribute.The target is p(ya, z) = p(ya)p(z), and performance depends on the invariance of z.Source: §3 · p. 6

Adaptive approaches

• Adaptive approaches evaluate JTT and BN-Adapt, which dynamically modify p reweight and do not provide performance guarantees.Source: §3 · p. 6

Representation learning

• Representation learning evaluates β-VAE and ImageNet pretraining, with performance depending on representation quality for the specific task.These methods aim to learn a robust representation z that describes the true prior.Source: §3 · p. 6

4 EXPERIMENTS

• The experiments evaluate 19 methods on six vision-classification datasets across three distribution shifts, varying label noise, and dataset size.The paper presents aggregate results in figures 3–7 and complete results in appendix figures 10–12.Source: §4 · p. 6

4.1 TAKEAWAYS

• No method consistently wins: CycleGAN leads under spurious correlation, pretraining leads under low-data drift, and dataset-specific winners include augmentation/DANN or ViT/MLP.Pretraining is also among the strongest methods for unseen data shift, but the leading method depends on the dataset and shift.Source: §4.1 · p. 7, Figure 5 · p. 7
• Pretraining often boosts performance, plausibly because its learned representations capture useful invariances, though it does not help on every dataset.The paper notes exceptions on CAMELYON17 and IWILDCAM, while highlighting useful properties such as scale, shift, and color invariance.Source: §4.1 · p. 7
• Heuristic augmentation helps when it reflects data attributes: ImageNet augmentation generally improves results, while RandAugment, AugMix, and AutoAugment vary more.Color jitter harms CAMELYON17 but is essential on SHAPES3D and MPI3D; mismatched augmentations may waste model capacity.Source: §4.1 · p. 7
• Learned augmentation works especially well under spurious correlation and can also help with low-data and unseen shifts, though less than pretraining there.Perfectly learned augmentations can produce samples from the true underlying generative model and cover missing parts of the distribution.Source: §4.1 · p. 7
• Domain-generalization methods offer limited gains, with DANN most helpful in some low-data and unseen-shift settings and performance varying by dataset.Their performance is rarely much better than heuristic augmentation.Source: §4.1 · p. 7
• With varying label noise, relative method performance is reasonably consistent; as datasets shrink, heuristic augmentation performs poorly, while pretraining and learned augmentation remain robust.The chosen attributes also matter: on DSPRITES, using color as the label and shape as the other attribute yields perfect unseen-shift generalization, unlike using shape as the label.Source: §4.1 · p. 8, Figure 6 · p. 8, Figure 7 · p. 8

4.2 PRACTICAL TIPS

• Practically, start with heuristic augmentations that approximate the data, learn augmentations if those fail, use pretraining, and expect limited gains from complex methods.The paper identifies choosing heuristic augmentations efficiently and learning the underlying generative model directly from data as open research questions.Source: §4.2 · p. 8

• No method performs best across all shifts; performance depends on the attribute and dataset considered.Source: §5 · p. 8
• Practitioners cannot select the best method from the dataset alone, and identifying the precise shift remains an important area for future investigation.Source: §5 · p. 8
• When the shift is known, adaptable algorithms that use auxiliary information should be evaluated in the context where they will be used.Source: §5 · pp. 8–9
• Comprehensive evaluations should vary examples, noise, and dataset size to give a more realistic view of methods’ relative performance.Source: §5 · p. 9

Benchmarking robustness to out of distribution (OOD) generalization

• Gulrajani & Lopez-Paz (2021) found no method significantly outperformed strong ERM, while Hendrycks et al. (2021) reported gains from larger models, better augmentation, and pretraining.On the Koh et al. (2020) benchmark, the largest boosts came from larger models and better augmentation; this work compares a broader range of approaches across synthetic and real-world datasets and shifts.Source: §6 · p. 9

Benchmarking spurious correlation and low-data drift

• Prior work links low-data and spurious-correlation failures to fairness and bias, and Arjovsky et al. (2019) hypothesized spurious correlation underlies poor generalization to unseen data.The authors report that systematic multi-dataset study with fine-grained control over shift was lacking, and introduce a framework for creating shifts controllably.Source: §6 · p. 9

Benchmarking disentangled representations

• Locatello et al. (2019) found disentanglement unreliable without supervision or implicit assumptions, while Locatello et al. (2020) found weak supervision sufficient.Studies by Dittadi et al. (2021), Schott et al. (2021), and Montero et al. (2020) found representations do not reliably interpolate, extrapolate, or compose complex properties across multiple datasets.Source: §6 · p. 9

7 CONCLUSIONS

• The framework evaluates 19 methods across spurious correlation, low-data drift, and unseen data shift, as well as label noise and dataset size.Source: §7 · p. 9
• Results vary across datasets and methods, although some methods outperform an ERM baseline in some settings.Source: §7 · p. 9
• The authors offer practical tips, promising directions, and open research questions, and propose the framework as a tool for evaluating methods under different conditions and shifts.Source: §7 · p. 9

A OVERVIEW

• The appendices also cover the literature, datasets and shift setup, evaluated methods, and implementation details.Source: A OVERVIEW · p. 16

B RESULTS

• Appendix B provides complete results, analyzes RandAugment augmentation types, and tests alternative labels and ID rather than OOD validation on IWILDCAM and CAMELYON17.Source: B RESULTS · p. 16

B.1 COMPLETE RESULTS

• Figures 10–12 show complete per-dataset results for three shifts, plotting each method’s mean and standard deviation and sorting methods by performance.The shifts are spurious correlation, low-data drift, and unseen data shift.Source: §B.1 · p. 16

B.2 A DETAILED ANALYSIS ON THE IMPACT OF AUGMENTATION

• The analysis randomly applies each RandAugment technique or the identity function to measure each technique’s contribution to model performance.It examines deviations from the mean across datasets under the unseen data shift setting.Source: §B.2 · p. 16
• No single RandAugment augmentation consistently improves performance across datasets.Invert improves performance on DSPRITES and SHAPES3D but harms it on MPI3D, SMALLNORB, and IWILDCAM; color improves performance on most datasets but harms it on CAMELYON17.Source: Figure 8 · p. 16, §B.2 · p. 16

B.3 RESULTS WITH DIFFERENT ATTRIBUTES

• The DSPRITES experiment tests whether unseen-shape results depend on the attribute being predicted, using color labels instead of shape labels.The setup leaves out some shapes and tests color prediction on those unseen shapes.Source: §B.3 · pp. 16–17
• All methods generalize to approximately perfect OOD accuracy when predicting color from unseen shapes on DSPRITES.This tests unseen-shape shift with color as the label rather than shape.Source: §B.3 · p. 17, Figure 9 · p. 17

B.4 RESULTS USING DIFFERENT, OOD VALIDATION SETS

• CAMELYON17 performs a bit better with OOD validation, whereas IWILDCAM performs best with ID validation.Thus, neither validation-set choice consistently produces the higher maximum performance.Source: §B.4 · p. 17
• The relative performance of models is comparable with ID and OOD validation sets on CAMELYON17 and IWILDCAM.Figure 10 compares top-1 accuracy across datasets.Source: §B.4 · p. 17, Figure 10 · p. 18
• The validation-set comparison evaluates IWILDCAM and CAMELYON17 to determine whether ID or OOD validation is preferable for best performance.The passage reports that neither choice performs best overall.Source: §B.4 · p. 17
• Figures 11 and 12 plot top-1 accuracy across values of N for spurious-correlation and low-data-drift shifts, respectively.N is the number of samples from the independent distribution; results are shown across different datasets.Source: Figure 11 · p. 19, Figure 12 · p. 20

C LITERATURE REVIEW

• Benchmarks include synthetic shifts, cross-dataset and temporal shifts, and real-world tasks; WILDS collects real-world shifts with defined OOD tasks (Koh et al., 2020).The framework complements existing datasets with fine-grained control over shift type and amount, extensible approaches and datasets, and a robustness framework inspired by disentangling literature; it can also extend to complex label noise.Source: C LITERATURE REVIEW · p. 21
• Model capacity had no strict performance rule, augmentation helped when its invariance was useful, and DANN performed consistently best while varying across datasets and shifts.The highest-capacity ViT often performed worse, except on IWILDCAM; learned augmentation was also limited by generative-model quality.Source: C LITERATURE REVIEW · pp. 21–22
• Pretraining helped in many cases but not all, JTT and BN-Adapt did not consistently beat baselines, and disentanglement methods need further robustness research.Pretraining did not improve performance across all shifts on CAMELYON17 and IWILDCAM, and was unhelpful under spurious correlation on CAMELYON17.Source: C LITERATURE REVIEW · pp. 22–23

D DATASET

• Appendix D provides dataset details, shift and condition setups, and sample examples for the datasets and shifts.Source: D DATASET · p. 23

D.1 DETAILS

• DSPRITES, MPI3D, and SHAPES3D use shape labels and color attributes; SMALLNORB uses object-category labels with azimuth or lighting attributes across its settings.CAMELYON17 uses tumor presence and hospital, while IWILDCAM uses animal presence or identity and camera-trap location; their shift setups include resplits and held-out groups as applicable.Source: §D.1 · pp. 23–24

D.2 EVALUATED SHIFTS AND CONDITIONS

• The framework controls correlation, low-data sample counts and sampling probabilities, label noise, and total dataset size; spurious-correlation training combines correlated samples with N >= 1 uncorrelated samples.Its test set samples from the uncorrelated distribution and is disjoint from training samples.Source: §D.2 · p. 24
• Low-data drift limits samples to N for selected attribute values while retaining all samples for others; unseen data shift is the special case N = 0.Source: §D.2 · p. 24
• The noisy-label condition varies uniform label-noise probability p at fixed N, while the dataset-size condition varies total training samples n in the low-data setting.Source: §D.2 · p. 24

D.3 SAMPLES FROM THE DIFFERENT DISTRIBUTIONS

• Sample distributions are shown for DSPRITES, MPI3D, SHAPES3D, and SMALLNORB.Figures 13–16 provide the samples for these four datasets.Source: §D.3 · p. 24, Figure 13 · p. 25, Figure 14 · p. 26
• Sample distributions are also shown for CAMELYON17 and IWILDCAM.Figures 17–18 provide the samples for these two datasets.Source: §D.3 · p. 24, Figure 17 · p. 29, Figure 18 · p. 30
• The dataset shift descriptions specify labels, nuisance attributes, spurious-correlation and low-data-drift/unseen-data-shift settings, plus low-data hyperparameters for label-noise and fixed-size conditions.For IWILDCAM, animals are correlated with their most common locations; it has 72 locations, and IWILDCAM and CAMELYON17 use the OOD set from Koh et al. (2020) as Aᵃᶜ.Source: Table 2 · p. 31

E METHOD

• Heuristic augmentation approximates p(x|y1:K) with manually chosen functions, but these may not match the true generative model and use no additional attribute information.The models are trained to predict labels, as in the standard setup.Source: E METHOD · p. 32
• Learned augmentation transforms images between attribute values and generates samples under a uniform attribute distribution, but depends strongly on generative-model quality.Following Goel et al. (2020), the approach uses STARGAN to handle more than two attribute values and excludes the additional SGDRO objective to isolate augmentation.Source: E METHOD · p. 32
• Domain-generalization methods seek representations z independent of the attribute ya, so the task classifier p(yl|z) is attribute-independent if that invariance is achieved.Their effectiveness depends on the method's ability to learn invariance.Source: E METHOD · p. 32
• JTT upweights difficult samples in a second training stage, while BN-Adapt trains attribute-specific models by modifying batch-normalization parameters; neither offers strong guarantees for the final model.JTT uses an initial classifier to estimate difficulty and weights difficult samples by λ; BN-Adapt trains |Aa| models with weights focused on one attribute value each.Source: E METHOD · p. 32
• Representation-learning approaches include ImageNet pretraining and β-VAE disentanglement, but their robustness depends on representation quality for the specific task.The intended representation captures underlying generative factors in the β-VAE approach.Source: E METHOD · p. 32
• ViT had the largest capacity but was hardest to train and performed best only on IWILDCAM, while ResNets performed best on the other datasets.The standard setup trains on ptrain to predict labels without using additional attribute information, with weighted resampling from lower-probability distribution regions.Source: E METHOD · p. 32, Table 3 · p. 33

F IMPLEMENTATION

• Implementation details cover architectures, approach-specific methods, training, and hyperparameter sweeps; the reported results are representative, not claimed to be each method's best possible.The paper notes that much larger sweeps would be needed to establish best-possible results across all methods.Source: F IMPLEMENTATION · p. 33

F.1 BASE ARCHITECTURES

• Experiments use ResNets, a four-layer MLP, and a ViT, with ViT patch and transformer settings adapted to image size.The ResNets are ResNet18, ResNet50, and ResNet101 (He et al., 2016); the MLP has 256 hidden units. ViT settings include patch sizes 4, 12, or 16 for the stated image sizes, with hidden size 256, eight heads, and eight layers for smaller images (Dosovitskiy et al., 2021).Source: §F.1 · p. 33

F.2 HEURISTIC AUGMENTATION

• Heuristic augmentation includes ImageNet augmentation, AugMix, RandAugment, and AutoAugment, with policies and parameters adapted to dataset image size.AugMix uses k = 3; RandAugment uses N = 3 and M = 5. AutoAugment uses the CIFAR-10 policy below 128 pixels and the ImageNet policy otherwise; image-size-dependent cutout and translation parameters are interpolated between ImageNet and CIFAR-10 settings.Source: §F.2 · p. 33

F.3 LEARNED AUGMENTATION

• Learned augmentation uses a two-stage StarGAN procedure to transform an image to a randomly selected attribute while keeping its label fixed.The first stage combines reconstruction and classifier losses with an adversarial loss, sweeping λa while setting λc = λr = 1; the second uses the pretrained model and α = 0, without real samples. The passage describes the approach as CycleGAN (Goel et al., 2020; Choi et al., 2018; Zhu et al., 2017) and references StarGAN (Choi et al., 2018).Source: §F.3 · pp. 33–34

F.4 DOMAIN GENERALIZATION

• Domain-generalization methods target cross-domain risk, feature alignment, smooth predictions, or attribute-invariant representations.IRM uses a shared optimal classifier across domains; DeepCORAL aligns feature means and covariances; Domain MixUp interpolates across domains and labels; DANN uses an adversary to predict attributes; SagNet uses content and style predictors to encourage style-invariant features. Domain MixUp uses λ = 0.2, and SagNet uses λa = 1.Source: §F.4 · p. 34

F.5 ADAPTIVE APPROACHES

• Adaptive approaches include JTT, which reweights challenging samples found by a pretrained classifier, and BN-Adapt, which adjusts parameters using the attribute label.JTT uses λ = 20 and trains its first model for half the training time; BN-Adapt uses r = 10 to control the relative strength of the original and adapted models.Source: §F.5 · p. 34

F.6 REPRESENTATION LEARNING

• Representation learning evaluates a fixed β-VAE encoder with a downstream MLP and also fine-tunes a full ImageNet-pretrained network.The β-VAE uses resized 64x64 images and sweeps latent size L and normalized β; the downstream classifier is an MLP. ImageNet-pretrained models receive resized 224x224 images.Source: §F.6 · pp. 34–35

F.7 TRAINING.

• ResNets and the MLP use Adam with batch size 128, while the ViT uses batch size 1024; all models use validation-accuracy early stopping.Training runs for up to 100K steps on synthetic datasets and 200K steps on CAMELYON17 and IWILDCAM; the ViT did not converge at smaller batch sizes.Source: §F.7 · p. 35

F.8 SWEEPS.

• For each method and seed, the sweep evaluates at most eight models and selects the best by ID or OOD validation performance for test evaluation.The sweep configurations are reported in Table 4.Source: §F.8 · p. 35, Table 4 · p. 36