Navigation

Introduction to AI

Machine Learning

Deep Learning

Generative AI

Tools & Frameworks

General

Domain Generalization

Domain generalization (DG) is the problem of training a model on data from multiple source domains such that it performs well on a new, unseen target domain — without any access to target data during training. This distinguishes DG from domain adaptation (which has access to unlabeled or labeled target examples) and standard supervised learning (which assumes train and test are i.i.d.).

The core challenge is learning features that are invariant across the observed source domains, with the hope that such features also transfer to the target domain.

Problem Setup

Let $\mathcal{E}$ be a set of environments (domains), each providing a data distribution $P^e(X, Y)$. Training domains are $\mathcal{E}\text{train} = {e_1, \ldots, e_k}$; the test domain $e\text{test}$ is unseen. The goal is to find a predictor $f$ that minimizes:

$$\min_f \max_{e \in \mathcal{E}\text{all}} \mathcal{R}^e(f) = \mathbb{E}{(x,y) \sim P^e}[\ell(f(x), y)]$$

where $\mathcal{E}_\text{all}$ includes training and unseen test domains. This is a worst-case objective over all possible environments.

Empirical Risk Minimization: Why It Fails

The naive approach — pooling data from all source domains and applying standard ERM — fails because it exploits spurious correlations: statistical associations in the training data that do not hold in general.

Classic Example: Cow vs. Background

A classifier trained to distinguish cows from camels may learn to use the background (green grass → cow, sandy desert → camel). In training data, cows appear on grass and camels in deserts with high correlation. Under ERM, the model achieves high training accuracy by relying on the spurious background feature rather than the animal shape.

In deployment on a new domain (e.g., a camel on grass), the model fails catastrophically. The causal feature — animal shape — is invariant across environments; the background feature is spurious.

Invariant Risk Minimization (IRM)

IRM (Arjovsky et al., 2019) formalizes the causal feature requirement as an invariance constraint: a representation $\Phi: X \to Z$ is invariant if there exists a single linear classifier $w$ on top of $\Phi$ that is simultaneously optimal across all training environments:

$$\min_{\Phi, w} \sum_{e \in \mathcal{E}_\text{train}} \mathcal{R}^e(w \circ \Phi)$$

subject to $w \in \arg\min_{\bar{w}} \mathcal{R}^e(\bar{w} \circ \Phi) \quad \forall e \in \mathcal{E}_\text{train}$

The constraint says: $w$ must be simultaneously optimal in every environment, which forces $\Phi$ to extract only features that are useful everywhere — ruling out spurious features that are useful only in some environments.

IRMv1: A Practical Relaxation

The exact IRM constraint is a bilevel optimization problem. IRMv1 relaxes it using a gradient penalty:

$$\mathcal{L}\text{IRM}(\Phi) = \sum{e \in \mathcal{E}\text{train}} \mathcal{R}^e(\Phi) + \lambda \sum{e \in \mathcal{E}\text{train}} |\nabla{w|w=1} \mathcal{R}^e(w \cdot \Phi)|^2$$

The gradient norm $|\nabla_{w|w=1} \mathcal{R}^e(w \cdot \Phi)|^2$ measures how far $w = 1$ is from being locally optimal for environment $e$. Minimizing this across environments penalizes representations where any single environment would “prefer” a different classifier, nudging $\Phi$ toward invariance.

Limitations of IRM

Extensive empirical evaluation (DomainBed) showed that IRMv1 frequently fails to outperform standard ERM on real-world benchmarks. Key failure modes:

  • Finite environment problem: With few source domains ($k < d_\text{spurious}$), there are not enough constraints to identify invariant features.
  • Optimization instability: The gradient penalty makes training noisy and sensitive to $\lambda$.
  • Limited nonlinearity: IRM theory guarantees apply for linear classifiers; deep networks violate the assumptions.

CORAL: Domain Alignment

CORAL (Sun & Saenko, 2016; DeepCORAL 2016) aligns the second-order statistics (covariance matrices) of feature distributions across domains:

$$\mathcal{L}_\text{CORAL} = \frac{1}{4d^2} |C_S - C_T|_F^2$$

where $C_S, C_T \in \mathbb{R}^{d \times d}$ are the feature covariance matrices of source and target domains, and $|\cdot|_F$ is the Frobenius norm.

GroupDRO (Sagawa et al., 2020) minimizes the worst-group empirical risk, upweighting underperforming groups (environment + label combinations) during training. It is effective when spurious correlations cause differential performance across demographic subgroups.

Mixup and Domain-Invariant Augmentation

Mixup interpolates between samples from different domains:

$$\tilde{x} = \lambda x^{e_i} + (1-\lambda) x^{e_j}, \quad \tilde{y} = \lambda y^{e_i} + (1-\lambda) y^{e_j}$$

Training on interpolated points regularizes the model to vary smoothly between domain distributions, reducing reliance on domain-specific texture or color statistics. Domain-invariant augmentation strategies (e.g., style transfer, random domain randomization) pursue the same goal by synthetically creating new domain views during training.

DomainBed: A Rigorous Benchmark

DomainBed (Gulrajani & Lopez-Paz, 2021) is a benchmark and evaluation framework designed to rigorously compare DG algorithms. Key contributions:

  • Standardized train/val/test splits that prevent data leakage across domains.
  • Fixed hyperparameter selection using a held-out source domain (not the test domain).
  • Datasets: PACS (Photo/Art/Cartoon/Sketch), VLCS, OfficeHome, Terra Incognita, DomainNet.

DomainBed Findings

After careful evaluation, DomainBed’s landmark finding was that ERM with strong data augmentation and proper tuning outperforms most specialized DG algorithms on most benchmarks. This challenged the widespread assumption that invariance-based methods were substantially better than ERM.

Key factors that matter more than the DG algorithm:

  • Data augmentation (RandAugment, CutMix, color jitter).
  • Backbone architecture (ResNet-50 → ViT-Large is a bigger gain than any algorithmic change).
  • Hyperparameter selection protocol (test-domain validation vs. source-domain validation changes rankings substantially).

Causal View of Domain Generalization

The theoretical justification for invariant feature learning comes from a structural causal model (SCM) perspective. Consider the causal graph:

$$C \to X, \quad S \to X, \quad C \to Y$$

where $C$ is the causal feature (e.g., animal shape), $S$ is the spurious feature (e.g., background), and $Y$ is the label. Across environments, the mechanism $P(Y \mid C)$ is invariant, but $P(S)$ and $P(S \mid C)$ may shift.

A predictor that uses only $C$ will generalize; one that uses $S$ will fail when $P(S)$ shifts at test time. IRM, GroupDRO, and related methods attempt to operationalize this causal intuition — learning representations that capture $C$ rather than $S$ — but doing so reliably without causal graph knowledge remains an open problem.

Practical Recommendations

Based on DomainBed and subsequent work, a practical DG workflow:

  1. Start with ERM + strong augmentation — color jitter, random grayscale, Gaussian noise, RandAugment.
  2. Use a large pretrained backbone — ViT-L/14 CLIP features are extremely robust to domain shift because CLIP pretraining already saw diverse visual styles.
  3. Apply GroupDRO if domain/group labels are available and worst-case performance matters (fairness-critical applications).
  4. Tune on a held-out source domain, not the target domain, to avoid overfitting the evaluation.
  5. Consider domain-invariant augmentation (style transfer, domain randomization) for specific domain shift types (texture → shape).

Open Challenges

  • Identifying environments: Most DG methods assume environment labels are available. In practice, domains may be latent or ambiguous.
  • Continuous domain shift: Real-world shift is often gradual (temporal covariate shift) rather than discrete domain jumps.
  • Beyond covariate shift: Most DG theory assumes $P(Y \mid C)$ is invariant. When the causal mechanism itself shifts, standard DG fails.
  • Scaling to large language models: DG for NLP involves linguistic style, dialect, and topic shift — structurally different from visual domain shift but equally practically important.

Summary

Domain generalization trains models that transfer to unseen test domains by learning features invariant across source environments. IRM provides a principled causal framework for invariance but struggles in practice due to the finite environment problem. CORAL and GroupDRO offer more robust baselines through covariance alignment and worst-group optimization. DomainBed’s rigorous benchmarking revealed that strong data augmentation and large pretrained backbones often outperform specialized DG algorithms — shifting the field’s focus toward understanding when and why invariance methods provide genuine gains over well-tuned ERM.