Ask Your Distribution Shift if Pre-Training is Right for You (2403.00194v2)

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. Code is available at https://github.com/MadryLab/pretraining-distribution-shift-robustness.

• The paper demonstrates that pre-training notably improves extrapolation in models, enhancing robustness against shifts beyond the training distribution.
• It reveals that pre-training is less effective in mitigating biases from spurious training features, emphasizing the need for bias-specific interventions.
• The study validates combined strategies like pre-training with Deep Feature Reweighting to achieve superior performance on tasks with real-world distribution shifts.

An Analytical Perspective on the Robustness of Pre-Trained Models Under Distribution Shifts

The paper "Ask Your Distribution Shift if Pre-Training is Right for You" by Cohen-Wang et al. presents a nuanced inquiry into the efficacy of pre-training as a method for enhancing the robustness of machine learning models against distribution shifts. It addresses the inconsistent success of pre-training, which has yielded substantial robustness improvements in some scenarios while providing negligible benefits in others. This disparity motivates the core research question: Under what conditions does pre-training affect model robustness against distribution shifts?

The paper identifies two primary failure modes associated with distribution shifts. The first is poor extrapolation, where models struggle to generalize beyond the reference distribution. The second involves biases embedded within the training data that lead to reliance on spurious features. Through theoretical exploration and empirical validation, the authors establish that pre-training primarily counters the former failure mode by aiding in extrapolation. However, it does little to mitigate the latter, specifically the biases present in training datasets.

The theoretical underpinning, elucidated in a logistic regression setting, convincingly illustrates that pre-training can shape a model's decision boundary, particularly beyond the reference distribution's support. This extrapolative power is vividly demonstrated through controlled experiments involving synthetic distribution shifts like altered color tints and geometric transformations applied to datasets such as ImageNet. These experiments confirm that while pre-training enhances robustness in cases necessitating extrapolation, it fails to resolve issues stemming from dataset biases, such as those seen in spurious correlation scenarios.

A significant implication of this understanding is the complementary nature of pre-training combined with bias-specific interventions. The empirical demonstrations using interventions like Deep Feature Reweighting (DFR) exemplify how distinct strategies can collectively bolster model robustness, addressing different facets of failure. For instance, in scenarios like the WILDS-FMoW task, which articulates shifts in satellite imagery datasets, a synthesis of pre-training with DFR delivers notable robustness benefits over employing either strategy independently.

Further explorations reveal that a robust pre-trained model can be fine-tuned effectively on a de-biased dataset of limited size and diversity without relinquishing performance benefits. This insight holds particular promise for applications where de-biasing the entire dataset may be unfeasible due to resource constraints.

In conclusion, this paper contributes a rigorous framework for understanding the limitations and capacities of pre-trained models under distribution shifts. It advocates for a targeted application of pre-training, emphasizing the need for refined approaches in addressing dataset biases. The findings are particularly significant when considering the heterogeneous nature of real-world data operational contexts, offering guidance on optimizing model robustness through pre-training when tasked with unseen or dynamically evolving environments. Additionally, this work sets the stage for future investigations into devising and fine-tuning pre-training paradigms tailored to specific distribution shifts, potentially informing the development of more resilient artificial intelligence models in various domains.