Source-linked AI summary
A Fine-Grained Analysis on Distribution Shift
Olivia Wiles, Sven Gowal, Florian Stimberg, Sylvestre Alvise-Rebuffi, Ira Ktena, Krishnamurthy Dvijotham, Taylan Cemgil
TL;DR
Robustness to distribution shifts is important for reliable machine-learning deployment, yet the mechanisms behind shifts and comparative robustness across shifts remain insufficiently understood. The paper introduces an extensible fine-grained framework and evaluates 19 methods across synthetic and real-world settings. It finds that some methods outperform ERM in particular settings, but performance is inconsistent across datasets and shifts.
Problem
The paper addresses limited understanding of the mechanisms behind distribution shifts and limited systematic evidence about how methods perform across different shifts.
Method
The paper defines an attribute-based robustness framework and evaluates 19 methods across three distribution shifts, two additional conditions, and synthetic and real-world datasets.
Results
Some methods outperform an ERM baseline in particular settings, but results are not consistent across datasets and distribution shifts.
Takeaways & Limitations
Method selection requires evaluation across multiple conditions because no approach is reliably best from the dataset alone.
Takeaways & Limitations
The best method cannot be selected a-priori from only the dataset, and the precise distribution shift in an application may be unclear.
Abstract
from arXiv · showhide
Robustness to distribution shifts is critical for deploying machine learning models in the real world. Despite this necessity, there has been little work in defining the underlying mechanisms that cause these shifts and evaluating the robustness of algorithms across multiple, different distribution shifts. To this end, we introduce a framework that enables fine-grained analysis of various distribution shifts. We provide a holistic analysis of current state-of-the-art methods by evaluating 19 distinct methods grouped into five categories across both synthetic and real-world datasets. Overall, we train more than 85K models. Our experimental framework can be easily extended to include new methods, shifts, and datasets. We find, unlike previous work~\citep{Gulrajani20}, that progress has been made over a standard ERM baseline; in particular, pretraining and augmentations (learned or heuristic) offer large gains in many cases. However, the best methods are not consistent over different datasets and shifts.
1 INTRODUCTION
The paper studies how to make machine-learning models robust to distribution shifts by defining shifts systematically and evaluating methods across realistic conditions. It finds that augmentation, pretraining, and some domain-generalization methods can help, but the best approach varies across datasets and shifts.
- Distribution-shift robustness is essential for deploying models in critical applications because failures can undermine system trust.Examples include changing hospitals, lighting conditions, object poses, and demographic subsets.
- Existing work leaves limited understanding of when and why models generalize across realistic distribution shifts.The paper frames this gap as a need to define important shifts and evaluate robustness systematically.
- The framework models data through attributes and studies generalization to unseen attribute values or changed attribute distributions.A shape-prediction model, for example, should generalize from observed colors to unseen colors or altered color frequencies.
- The evaluation covers spurious correlation, low-data drift, unseen data shift, label noise, and dataset size across synthetic and real-world settings.These shifts are chosen as real-world-inspired conditions that can harm generalization and serve as building blocks for more complex shifts.
- The study compares 19 methods spanning five approaches, including architecture choice, augmentation, domain generalization, adaptive algorithms, and representation learning.The comparison trains more than 85K models and is designed to be extendable to additional methods, shifts, and datasets.
- Simple techniques such as augmentation and pretraining are often effective, while no single method performs best across all datasets and shifts.Domain-generalization algorithms also work for certain datasets and distribution shifts, but results remain inconsistent.
2 FRAMEWORK TO EVALUATE GENERALIZATION
The framework models distribution shifts as changes in attribute distributions while sharing the conditional data-generating process, then evaluates robustness across three shifts and two additional conditions.
- Latent factorisation: The framework represents data using inputs and multiple finite-valued attributes, including one prediction label and nuisance attributes.The classifier minimizes empirical risk because only a finite training sample is available.
- Latent factorisation: Training and test distributions may differ because data collection concentrates on different subsets of hospitals or other attribute combinations.This mismatch can reflect selection bias or confounding factors in data collection.
- Latent factorisation: The latent-factor model assumes distribution shifts change attribute marginals while the conditional generative model remains shared across distributions.The true distribution is factorized into attribute marginals and a conditional generative model.
- Robustness framework: Robustness methods include weighted resampling, data augmentation, and representation learning to address changes between training and test distributions.Weighted resampling uses importance weights, while augmentation and representation learning target the underlying generative structure.
- Distribution shifts: The evaluation considers spurious correlation, low-data drift, unseen data shift, label noise, and dataset size across real-world-inspired settings.The three named shifts are treated as fundamental building blocks for more complex distribution shifts.
- Distribution shifts: Spurious correlation links attributes during training but not testing, whereas low-data drift changes uneven attribute frequencies between training and testing.The framework’s figures visualize these shifts and compare methods relative to a ResNet baseline.
3 MODELS EVALUATED
The study evaluates a broad set of robustness approaches spanning architecture, augmentation, domain generalization, adaptive methods, and representation learning.
- Evaluation scope: The benchmark evaluates 19 algorithms across architecture choice, data augmentation, domain adaptation, adaptive approaches, and representation learning.The methods are intended to cover broad strategies for improving robustness to distribution shifts.
- Architecture choice: Architecture experiments compare ResNet18, ResNet50, ResNet101, ViT, and MLP models with weighted resampling.Weighted resampling oversamples regions with lower training probability.
- Data augmentation: Heuristic augmentation methods include ImageNet augmentation, AugMix without JSD, RandAugment, and AutoAugment.These methods approximate the underlying generative model through fixed augmentation procedures.
- Data augmentation: Learned augmentation uses CYCLEGAN to transform images toward new nuisance attributes while keeping other attributes fixed.The evaluation isolates learned augmentation rather than the SGDRO objective.
- Domain generalization: Domain generalization methods seek representations independent of nuisance attributes and include IRM, DeepCORAL, domain MixUp, DANN, and SagNet.Their performance depends on the invariance of the learned representation.
- Adaptive and representation methods: Adaptive approaches JTT and BN-Adapt dynamically modify reweighting, while β-VAE and ImageNet pretraining target robust representations.The adaptive methods do not provide performance guarantees.
4 EXPERIMENTS
The experiments compare 19 methods across six vision datasets, three distribution shifts, label noise, and dataset size. Results show gains over ERM, but effectiveness varies substantially with datasets, shifts, and attributes.
- Experimental setup: 19 methods are evaluated across six vision datasets, three distribution shifts, varying label noise, and dataset size.The methods span architecture choice, augmentation, domain generalization, adaptive approaches, and representation learning.
- Results: Pretraining often provides a strong performance boost across different shifts and datasets, although it is not consistently helpful on CAMELYON17 and IWILDCAM.The authors attribute this to useful invariances in the pretrained representation.
- Results: ImageNet augmentation generally improves performance, whereas RandAugment, AugMix, and AutoAugment show more variable results.Augmentations work best when they approximate the true underlying generative model; color jitter, for example, harms CAMELYON17 but is essential on SHAPES3D and MPI3D.
- Results: Learned data augmentation is highly effective for spurious correlation and also helps under low-data and unseen data shifts.Its gains under the latter two shifts are smaller than those from pretraining.
- Results: Domain generalization methods offer limited improvement, with DANN helping most notably under low-data drift and unseen data shift.Their performance depends on the dataset and is rarely much better than heuristic augmentation.
- Results: The best algorithm depends on precise conditions: pretraining and learned augmentation remain robust as dataset size decreases, while heuristic augmentation performs poorly.Changing which attributes define a shift can also change the outcome; on DSPRITES, reversing label and auxiliary attributes makes all methods generalize perfectly in the unseen data shift setting.
5 DISCUSSION
No single method performs best across all shifts, datasets, or attributes, making a-priori selection difficult. The authors therefore emphasize shift-aware, adaptable methods and comprehensive evaluation.
- No one method performs best over all shifts, and performance depends on the precise attribute being considered.
- There is no way to decide a-priori on the best method given only the dataset.
- Adaptable algorithms that use auxiliary information are preferable when the distribution shift is known.
- Performance varies with the number of examples, amount of noise, and dataset size.
- Comprehensive evaluations provide a more realistic view of models’ relative performance in practice.
6 RELATED WORK
Prior work has questioned whether domain-generalization methods improve over ERM, while related studies examine augmentation, spurious correlation, low-data drift, and disentangled representations. These lines of work motivate broader benchmarking of robustness across distribution shifts.
- Gulrajani and Lopez-Paz found that no evaluated method significantly outperformed a strong ERM baseline across varied datasets.
- Better augmentation, larger models, and pretraining have been reported to provide a sizeable boost in out-of-distribution performance.
- Fairness and bias studies document the impact of low-data conditions and spurious correlations in applications including face recognition, medical imaging, and conservation.
- Disentangled representation learning aims to separate factors of variation so models can generalize to unseen data.
7 CONCLUSIONS
The paper presents a comprehensive framework for analyzing distribution shifts and evaluates many methods across shifts, datasets, and conditions. Results show that some methods beat ERM in some settings, but performance is inconsistent.
- The framework analyzes 19 methods across spurious correlation, low-data drift, unseen data shift, label noise, and dataset size.
- Some methods perform better than an ERM baseline in some settings, while results remain inconsistent across datasets and methods.
- The appendix includes complete method-by-shift results, RandAugment analyses, alternative label attributes, and ID-validation results.
- Complete results for every dataset and shift are reported in figures 10–12 using each method’s mean and standard deviation.
B.2 A DETAILED ANALYSIS ON THE IMPACT OF AUGMENTATION
The augmentation analysis shows that individual heuristic augmentations affect robustness differently across datasets. No augmentation consistently improves performance, and some augmentations help one dataset while harming another.
- Different heuristic augmentation methods vary in performance, with some outperforming ERM and others not.
- No augmentation always provides a strong performance boost across datasets under unseen data shift.
- Invert improves performance on DSPRITES and SHAPES3D but harms performance on MPI3D, SMALLNORB, and IWILDCAM.
- Color improves performance on most datasets but harms performance on CAMELYON17.
B.3 RESULTS WITH DIFFERENT ATTRIBUTES
Changing the predicted attribute on DSPRITES yields near-perfect unseen-shape generalization, while broader shift evaluations compare accuracy across datasets and sample counts. Validation-set choice preserves relative rankings but does not consistently maximize performance.
- Unseen data shift: All methods generalize to unseen shapes with approximately perfect color-prediction scores on DSPRITES.The experiment reverses the usual label–attribute assignment by predicting color while withholding shapes during testing.
- Validation selection: Neither OOD nor ID validation consistently produces the highest performance on CAMELYON17 and IWILDCAM.OOD validation performs somewhat better on CAMELYON17, whereas ID validation performs best on IWILDCAM.
- Shift comparisons: Figures 10–12 compare top-1 accuracy across datasets for unseen data, spurious-correlation, and low-data-drift shifts.Figures 11 and 12 vary N, the number of samples from the independent distribution; results are resorted for each N.
- Benchmark scope: The benchmark comparison emphasizes controllable shifts, real-world-motivated datasets, and broader method coverage than similar works.Table 1 contrasts this framework with three related benchmarks.
C LITERATURE REVIEW
The literature review situates the framework among shift benchmarks and five robustness approaches, highlighting both the promise and limits of invariance, augmentation, pretraining, and model scaling.
- Additional approaches: The review also covers adaptive approaches, label-noise evaluation, disentanglement, and dataset-specific implementation details.These topics complete the paper’s comparison of robustness methods and benchmark construction.
- Datasets: Existing benchmarks often use synthetic or repurposed datasets because real-world distribution-shift data are difficult and expensive to obtain.The framework is presented as complementary by providing fine-grained control over shift type and magnitude.
- Architecture choice: Model capacity has no strict relationship with robustness: deeper ResNets sometimes win, while ViT is best only on IWILDCAM.The authors caution that additional data or pretraining could alter these relative comparisons.
- Data augmentation: Augmentation can promote invariance to the chosen transformation, but irrelevant transformations waste capacity and learned augmentation depends on generator quality.This limitation applies to both heuristic and learned augmentation approaches.
- Domain generalization: Domain-generalization methods seek domain-invariant features, but enforcing invariance can be overly strict; DANN performs consistently best in this study.Performance still varies across datasets and distribution shifts.
- Representation learning: Pretraining helps in many cases but does not universally improve performance, including spurious correlation on CAMELYON17.On CAMELYON17 and IWILDCAM, pretraining did not improve performance across all shifts.
D.3 SAMPLES FROM THE DIFFERENT DISTRIBUTIONS
The appendix provides sample distributions for each dataset and records the precise shift configurations used across labels, nuisance attributes, and experimental conditions.
- Dataset coverage: The sample-distribution figures include SMALLNORB, CAMELYON17, and IWILDCAM in addition to the synthetic datasets.These figures show the distributions used across both synthetic and real-world data.
- Shift configurations: Table 2 specifies each dataset’s label and nuisance attribute together with the spurious-correlation, low-data-drift, and unseen-data-shift setups.It also lists hyperparameters for label-noise and fixed-dataset-size conditions.
- IWILDCAM: IWILDCAM contains 72 locations.This provides the dataset-specific location count reported in the appendix.
E METHOD
The method section organizes robustness techniques by how they modify architectures, data, representations, weighting, or adaptation, and relates each approach to the robustness framework.
- Approach categories: The study compares standard architectures, heuristic and learned augmentation, domain generalization, adaptive approaches, and representation learning.Each approach is described in terms of the robustness guarantee it is intended to provide.
- Backbone architecture: The standard setup trains a model to predict the true label without using nuisance-attribute information and may oversample lower-probability distribution regions.This weighting procedure is called preweight.
- Data augmentation: Heuristic augmentation creates augmented training data with manually selected transformations, whereas learned augmentation transforms attribute values while keeping other attributes fixed.The learned procedure approximates a conditional generative model from data.
- Domain generalization: Domain generalization targets representations independent of the nuisance attribute so the task classifier is independent of that attribute.Its success depends on whether the method can learn the required invariance.
- Adaptive approaches: Adaptive methods modify the reweighting distribution through multistage training, including weighting difficult samples more heavily.JTT uses an initial classifier to estimate sample difficulty before the second training stage.
- Representation learning: Representation learning initializes models with properties intended to improve robustness, including pretraining on large auxiliary datasets such as ImageNet.The approach is evaluated alongside the other robustness strategies.
- Evaluation scope: The reported results are representative rather than claims of each method’s best possible performance because exhaustive hyperparameter sweeps were infeasible.Architectures, implementations, training details, and sweeps are documented separately.
F.1 BASE ARCHITECTURES
The experiments compare multiple model architectures, augmentation strategies, and distribution-shift methods under standardized training and hyperparameter-selection procedures. The evaluated methods include invariant, feature-alignment, adversarial, reweighting, adaptation, and representation-learning approaches.
- Model architectures: Three model families—ResNets, a four-layer MLP, and ViT—provide architectures with different capacities for the experiments.The ResNet configurations are ResNet18, ResNet50, and ResNet101; the MLP uses 256 hidden units, while ViT settings vary by image size.
- Augmentation methods: ImageNet augmentation, AugMix, RandAugment, and AutoAugment implement heuristic or learned image transformations with dataset- and image-size-specific settings.AugMix uses k = 3 sequences, RandAugment uses N = 3 and M = 5, and AutoAugment selects CIFAR-10 or ImageNet policies by image size.
- Learned augmentation: CycleGAN generates augmented samples by transforming images between attribute values while keeping labels fixed, using reconstruction, classifier, and adversarial losses.The two-stage procedure first learns attribute transformations and then applies randomly selected target attributes to produce augmented images.
- Distribution-shift methods: IRM, DeepCORAL, Domain MixUp, and DANN target distribution shift through shared optimal classifiers, aligned feature statistics, interpolated domains, or adversarially invariant representations.IRM uses a risk term and invariance constraint; DeepCORAL aligns means and covariances; Domain MixUp interpolates images and labels; DANN trains an attribute-prediction adversary.
- Representation and adaptation methods: SagNet, JTT, BN-Adapt, and β-VAE address style invariance, challenging-sample reweighting, attribute-conditioned adaptation, or constrained latent representations.β-VAE varies normalized β and latent size L, then fixes its encoder while training a downstream MLP classifier.
- Training and evaluation: Training uses Adam, early stopping, and validation accuracy, while each seed selects the best model from a hyperparameter sweep for test-time evaluation.ResNets and MLPs use batch size 128, ViTs use batch size 1024, and the sweep allows up to eight models per seed because of capacity limits.