Generative Classifiers Avoid Shortcut Solutions

TL;DR

Discriminative classifiers often latch onto spurious correlations, causing fragility under distribution shift. The authors advocate generative classifiers that model $p_\theta(x|y)$ using diffusion-based (images) or autoregressive (text) models and apply Bayes' rule for prediction, achieving state-of-the-art robustness across five benchmarks without specialized data augmentations. They provide empirical evidence and mechanistic insights, including gradient analyses and Gaussian toy experiments, showing that learning the full input distribution yields more consistent signals and better generalization, especially under shift. The work also characterizes when generative classifiers outperform discriminative ones via generalization phase diagrams and discusses practical considerations like model size and unconditional objectives, indicating meaningful real-world impact for robust classification in domains with spurious correlations.

Abstract

Discriminative approaches to classification often learn shortcuts that hold in-distribution but fail even under minor distribution shift. This failure mode stems from an overreliance on features that are spuriously correlated with the label. We show that generative classifiers, which use class-conditional generative models, can avoid this issue by modeling all features, both core and spurious, instead of mainly spurious ones. These generative classifiers are simple to train, avoiding the need for specialized augmentations, strong regularization, extra hyperparameters, or knowledge of the specific spurious correlations to avoid. We find that diffusion-based and autoregressive generative classifiers achieve state-of-the-art performance on five standard image and text distribution shift benchmarks and reduce the impact of spurious correlations in realistic applications, such as medical or satellite datasets. Finally, we carefully analyze a Gaussian toy setting to understand the inductive biases of generative classifiers, as well as the data properties that determine when generative classifiers outperform discriminative ones.

Figures (18)

• Figure 1: Generative classifiers. We repurpose today's best generative modeling algorithms for classification. Generative classifiers predict $\mathop{\mathrm{arg\,max}}\limits_{y} p_\theta(x \mid y) p(y)$. We use diffusion-based generative classifiers on image tasks and autoregressive generative classifiers on text tasks, and find that they scale better out-of-distribution than discriminative approaches.
• Figure 2: In-distribution vs out-of-distribution accuracy for each dataset. Each point corresponds to a separate training run, other than the diffusion-based generative classifier results, which are checkpoints of a run with default training hyperparameters. We observe better OOD scaling trends (i.e., effective robustness) for generative classifiers on CelebA, CivilComments, and potentially Camelyon17, although results are noisy for this dataset (the red line in Camelyon17 denotes a linear fit for the relationship between ID and OOD accuracy for discriminative models). On the remaining two datasets, they follow the same trend and do better both ID and OOD.
• Figure 3: Gradient norms. Gradient for disc. model rapidly decays to 0, so its learning signal is reduced, while it does not decay for generative classifier.
• Figure 4: Scaling disc. model size does not improve accuracy on Waterbirds. This shows that model size is not a confounder in our experiments.
• Figure 5: Visualization of features (noise dims not shown).