PaperlayerView exact revision on arXiv

Source-linked AI summary

Model-Based Domain Generalization

Alexander Robey, George J. Pappas, Hamed Hassani

arXiv:2102.11436v5stat.MLcs.AIcs.LG

TL;DR

Out-of-distribution shifts can cause catastrophic failures, motivating domain generalization from related training domains to unseen test domains. The paper introduces MBDG, recasts the problem using transformation-based invariance and constrained learning, and reports improvements of up to 30 percentage points over state-of-the-art baselines.

  • Problem

    Domain generalization addresses catastrophic out-of-distribution failures by learning from related training domains for evaluation on an unseen test domain.

  • Method

    MBDG learns domain transformations, enforces invariance to them, and uses a duality-based unconstrained relaxation of the resulting constrained problem.

  • Results

    Up to 30 percentage points of improvement over state-of-the-art baselines are reported across benchmarks including ColoredMNIST, Camelyon17-WILDS, FMoW-WILDS, and PACS.

  • Takeaways & Limitations

    MBDG provides a model-based and theoretically grounded approach to domain generalization with reported gains across several benchmarks.

  • Takeaways & Limitations

    The theory assumes covariate-shift structure and requires enforcing invariance across all environments, although only finitely many training domains are observed.

Abstract

from arXiv · show

Despite remarkable success in a variety of applications, it is well-known that deep learning can fail catastrophically when presented with out-of-distribution data. Toward addressing this challenge, we consider the domain generalization problem, wherein predictors are trained using data drawn from a family of related training domains and then evaluated on a distinct and unseen test domain. We show that under a natural model of data generation and a concomitant invariance condition, the domain generalization problem is equivalent to an infinite-dimensional constrained statistical learning problem; this problem forms the basis of our approach, which we call Model-Based Domain Generalization. Due to the inherent challenges in solving constrained optimization problems in deep learning, we exploit nonconvex duality theory to develop unconstrained relaxations of this statistical problem with tight bounds on the duality gap. Based on this theoretical motivation, we propose a novel domain generalization algorithm with convergence guarantees. In our experiments, we report improvements of up to 30 percentage points over state-of-the-art domain generalization baselines on several benchmarks including ColoredMNIST, Camelyon17-WILDS, FMoW-WILDS, and PACS.

1 Introduction

The paper addresses out-of-distribution failures by introducing Model-Based Domain Generalization, which combines learned domain transformations with invariance and a duality-based optimization approach. The resulting algorithm improves over state-of-the-art baselines by up to thirty percentage points across several benchmarks.

  • Model-Based Domain Generalization learns transformations between domains and enforces predictor invariance to those transformations.
  • The framework reformulates domain generalization under covariate shift as a semi-infinite constrained statistical learning problem.
  • The approach derives data-dependent duality-gap bounds for an empirical parameterized dual that provides a tight approximation to the constrained problem.
  • Up to thirty percentage points separate the proposed algorithm from state-of-the-art baselines across several benchmarks.
  • The proposed primal-dual algorithm significantly outperforms state-of-the-art baselines on ColoredMNIST, Camelyon17-WILDS, and PACS.

2 Related work

Prior domain generalization work seeks domain-consistent representations or augments training data, including through style transfer and image-to-image translation. These approaches address out-of-distribution performance through representation alignment or modified data.

  • Domain generalization methods commonly learn internal feature representations that remain consistent across domains.
  • Unsupervised domain adaptation is related but assumes unlabeled test-domain data is available during training.
  • Data augmentation methods modify training data through transformations such as style transfer and image-to-image translation.
  • Other augmentation approaches remove textural features to encourage domain invariance rather than generating new data.

3 Domain generalization

Domain generalization learns from finitely many environmental domains while targeting all possible domains, where unseen-domain data is unavailable. The setting is difficult because arbitrary test-domain generalization is impossible without additional structure, including assumptions about covariate or concept shift.

  • Unlike standard supervised learning, observations are collected under environmental conditions that alter or corrupt the measured data.
  • Domain generalization seeks a predictor that performs well across all environments using samples from only finitely many training domains.
  • The objective minimizes worst-case risk over the full family of domains, including unseen domains without available training data.
  • Generalization difficulty depends on how closely unseen-domain distributions resemble the training-domain distributions.
  • Covariate shift changes instance marginals, whereas concept shift changes instance-conditional label distributions.

4 Model-based domain generalization

MBDG models inter-domain variation through a domain transformation function under covariate shift and imposes invariance to that transformation. This yields an equivalent constrained learning problem, while practical use is limited by unknown transformations, finite domains, and difficult invariance constraints.

  • 4 Model-based domain generalization: MBDG restricts its theoretical analysis to inter-domain variation caused solely by covariate shift through an underlying data-generation model.
  • 4 Model-based domain generalization: Under the assumptions, domain generalization is equivalent to a constrained statistical learning problem defined over the underlying distribution and transformations.
  • 4 Model-based domain generalization: The domain transformation model G maps underlying instances X to domain-specific observations Xe while preserving stable instance-conditional label distributions.
  • 4.6 Challenges in solving Problem 4.6: In practice, analytic expressions for G may be unavailable, and invariance over all environments cannot generally be evaluated from finitely many observed domains.
  • 4 Model-based domain generalization: The assumptions exclude concept shift, where labels and environments jointly influence observations, as in ColoredMNIST.
  • 4 Model-based domain generalization: G-invariance requires predictions to remain unchanged when an instance is transformed by G for every environment.

5 Data-dependent duality gap for MBDG

The paper relaxes the difficult G-invariance constraint, parameterizes the resulting problem, and derives an empirical unconstrained dual with a data-dependent duality gap. The gap is small when the invariance margin and parameterization error are small and enough samples are available.

  • Relaxing the invariance constraint: The G-invariance equality constraint is relaxed to an inequality constraint compatible with constrained PAC learning.The relaxation uses a margin γ and a distance metric over probability distributions.
  • Relaxing the invariance constraint: When γ = 0 under mild conditions on d, the relaxed and original problems are equivalent, while for γ > 0 their gap is bounded by Lγ.The result applies to KL-divergence and, more generally, f-divergences.
  • Parameterizing the problem: Finite-dimensional parameterization makes optimization tractable but can introduce nonconvexity and constraint-satisfaction difficulties in deep networks.An ϵ-parameterization approximates each function in F within expected sup-norm error ϵ.
  • Parameterizing the problem: The parameterized dual is unconstrained and therefore more amenable to optimization over deep neural networks than the original constrained problem.Its solution can recover a close approximation of the primal problem when the parameterization gap is small.
  • Data-dependent guarantee: The empirical parameterized dual has a small duality gap when γ is small, Aϵ closely approximates F, and sufficiently many samples are available.The guarantee holds with probability 1 − δ under mild regularity assumptions and finite VC-dimension.

6 Learning domain transformation models from data

MBDG learns domain transformation models from training-domain data using multimodal image-to-image translation networks. These models generate varied images that preserve input characteristics while often generalizing to new domains.

  • Learning transformations: The underlying domain transformation G is usually unknown and difficult to express in closed form, especially for variation in medical-image coloration, brightness, and contrast.This unknown transformation is a central challenge for applying the model-based framework.
  • Learning transformations: The paper learns approximations of G from training-domain instances using multimodal image-to-image translation networks.The networks transform samples from one dataset so they resemble diverse images from another dataset.
  • Learning transformations: These architectures combine a disentangled representation with a generative model that separates instance content from environmental variation and recombines it with a new environment.Replacing the environmental code e with e′ produces an instance in a new domain.
  • Empirical behavior: MUNIT models trained on Camelyon17-WILDS, FMoW-WILDS, and PACS learn accurate and diverse transformations that often generalize to new domains.Generated samples retain characteristic input features while exhibiting clear variation.

7 A principled algorithm for Model-Based Domain Generalization

The proposed algorithm solves the empirical parameterized dual with a primal-dual procedure using learned domain transformations. Under stated conditions, sufficiently many iterations and a small step size yield a solution close to the Model-Based Domain Generalization problem.

  • Algorithmic formulation: MBDG uses a learned transformation model G to solve the unconstrained dual optimization problem through primal-dual iteration.The method is motivated by the theoretical reformulation and is broadly applicable beyond settings with covariate shift.
  • Algorithmic formulation: The iteration alternates a primal update for predictor parameters with a dual update for the constraint variable.The outer maximization is linear in the dual variable, motivating the alternating scheme.
  • Convergence guarantee: With sufficiently many steps and a small enough step size, the primal-dual iteration converges with high probability to a solution close to Problem 4.6.The result assumes bounded loss and distance functions, finite VC-dimension, and mild regularity conditions.
  • Practical implementation: The practical MBDG algorithm replaces exact primal minimization with stochastic gradient descent and uses one dual variable across training domains.These modifications address the cost of finding a ρ-minimizer and the potentially large number of training domains.
  • Practical implementation: Algorithm 1 generates transformed images, computes classification and invariance regularization terms, then performs primal and dual updates until convergence.The implementation uses a sampled environmental latent code in MUNIT and a KL-divergence distance for the invariance term.

8 Experiments

Across domain generalization benchmarks, MBDG improves performance while learned transformations and primal-dual updates enforce the intended invariance constraints. Ablations show that learned transformations and dual ascent outperform standard augmentation and fixed regularization.

  • Camelyon17-WILDS and FMoW-WILDS: MBDG improves more than 20 percentage points over state-of-the-art baselines on Camelyon17-WILDS and around one percentage point on FMoW-WILDS.
  • Camelyon17-WILDS and FMoW-WILDS: Learned transformations that vary coloration and brightness help MBDG enforce invariance to domain changes on Camelyon17-WILDS.
  • Camelyon17-WILDS and FMoW-WILDS: MBDG is robust to changes under G, whereas ERM and IRM are not nearly as robust on Camelyon17-WILDS training data.
  • ColoredMNIST: On ColoredMNIST, MBDG improves over each baseline by nearly thirty percentage points and reaches 68.5% average accuracy despite violated assumptions.
  • ColoredMNIST: The G-invariance constraint forces identical predictions for differently colored versions of the same digit, explicitly eliminating color as a predictive feature.
  • ColoredMNIST: Dual-variable updates enforce the ColoredMNIST constraint by driving distReg(θ) toward γ = 0.025, unlike fixed regularization.
  • ColoredMNIST: Data-augmentation variants significantly outperform baselines but do not perform nearly as well as MBDG, making primal-dual iteration more effective for enforcing invariance.
  • PACS: MBDG achieves 85.6% average accuracy on PACS, nearly two percentage points above baselines, with nearly seven-point gains on Sketch.

9 Conclusion

The paper introduces Model-Based Domain Generalization and connects domain generalization to a semi-infinite constrained learning problem, yielding a duality-based primal-dual algorithm with improvements of up to 30 percentage points over state-of-the-art baselines.

  • Model-Based Domain Generalization is formulated through a natural data-generation model and an associated invariance notion as a semi-infinite constrained statistical learning problem.
  • A duality-based perspective produces a novel primal-dual algorithm for the constrained domain generalization problem.
  • Up to 30 percentage points of improvement over state-of-the-art baselines are reported for the proposed algorithm.

A Further theoretical results and discussion

The appendix develops theoretical conditions for relaxation tightness, equivalence between relaxed constraints and G-invariance, dual perturbation bounds, and learning-theoretic connections to constrained PAC learning. It also contrasts dual-based enforcement with fixed-weight regularization, whose generalization guarantees and hyperparameter choice are limited.

  • A.1 The case when γ = 0: When γ = 0 and the distance metric separates distributions, the relaxed problem is equivalent to the original G-invariant problem.The equivalence follows because zero expected nonnegative distance implies equality almost surely, and the metric condition converts equality into G-invariance.
  • A.2 The case when γ > 0: For γ > 0, relaxation tightness is not guaranteed generally, but Lipschitz continuity of P⋆(γ) yields the bound |P⋆−P⋆(γ)| ≤ Lγ.The perturbation function cannot generally be assumed Lipschitz; under strong duality, its Lipschitz constant is related to the L1 norm of an optimal dual variable.
  • A.3 Constrained learning theory: The relaxed inequality constraints transform hard G-invariance constraints into a form compatible with constrained learning-theoretic analysis.This reformulation is presented as a key step toward deriving learning guarantees for the model-based problem.
  • A.4 Regularization vs. dual ascent: Fixed-weight regularization can converge to local minima, but generalization guarantees apply only to the aggregated objective rather than each constraint term.The penalty weights are nontrivial hyperparameters and often require significant domain knowledge.

B.1 Proof of Proposition 4.3

The theoretical development rewrites domain generalization under covariate-shift and invariant-label assumptions as a constrained problem, then analyzes its dual, perturbations, parameterized approximations, and empirical gap.

  • The proof uses shared conditional label distributions and transformed covariate distributions to rewrite environmental risk in terms of the learned domain transformation G.The key identity expresses environmental risk as E[ℓ(f(G(X,e)),Y)].
  • Under the stated assumptions, domain generalization is equivalent to an optimization problem over predictors constrained to be G-invariant across all environments.The constraint requires f(x) = f(G(x,e)) almost everywhere for every environment e.
  • The theoretical analysis bounds parameterized and empirical dual approximations using perturbation arguments, function approximation, and empirical-gap control.The empirical-gap result assumes bounded nonnegative loss and distance functions and depends on the VC-dimension of the approximating hypothesis class.
  • Strong duality supports saddle-point relations for the relaxed primal and dual problems under strict-feasibility and regularity assumptions.The assumptions require regular loss and distance functions plus a predictor satisfying the relaxed constraints with a specified margin.

C.1 Data augmentation

The data-augmentation ablations compare ERM with learned transformations, MBDG with augmentation, and regularized MBDG against the full MBDG procedure. The reported comparisons show that augmentation can help, but enforcing invariance through the primal-dual procedure is more effective.

  • Data-augmentation variants: MBDA trains ERM with data augmentation generated by the learned domain transformation model G(x,e), without enforcing G-invariance constraints.This variant follows empirical evidence favoring tuned ERM with data augmentation for domain generalization.
  • Data-augmentation variants: MBDG-DA adds data augmentation through G(x,e) to the training objective while retaining the G-invariance constraints.
  • ColoredMNIST results: Nearly 20 percentage points separate MBDA from the baselines on ColoredMNIST, although MBDA remains below MBDG when the -90% domain is the test domain.

D Additional experiments and experimental details

The appendix documents dataset-specific tuning and ablations comparing learned MUNIT transformations with standard image transformations. It also reports competitive MBDG performance on VLCS and describes the evaluation setup for WILDS datasets.

  • Experimental setup: DomainBed experiments on ColoredMNIST, PACS, and VLCS retained default hyperparameters while recording additional MBDG settings and random-generation criteria.Camelyon17-WILDS used its out-of-distribution validation set for classifier hyperparameter tuning, DenseNet-121, Adam, and batch size 200.
  • Camelyon17-WILDS ablations: The Camelyon17-WILDS ablation enforced invariance between clean training images and versions transformed by fixed augmentation methods.The appendix compares Color Jitter, Brightness and contrast, RandAugment, geometric RandAugment, color-based RandAugment, and MUNIT transformations.
  • FMoW-WILDS: FMoW-WILDS experiments reran all baselines and MBDG over a hyperparameter grid because the prior work did not report its search grid.The comparison used the same architecture and optimizer as the prior work and the provided out-of-distribution validation set.
  • VLCS results: MBDG achieved the best result on the VLCS LabelMe subset by nearly two percentage points while remaining competitive with other state-of-the-art methods.The appendix presents the full VLCS baseline results in Table 6.

E Further discussion of domain transformation models

The paper contrasts known transformation models with practical settings where domain variation is difficult to express analytically. It therefore learns multimodal transformation models, while emphasizing that the true transformation model is unnecessary for the algorithm and experiments.

  • Known transformation models: Fixed-rotation domain generalization provides a tractable example in which each rotation angle identifies a distinct domain.The transformation is represented using a one-dimensional rotation matrix parameterized by the angle e.
  • Unknown transformation models: For practical domains, transformation models are usually unknown and difficult to represent with concise mathematical expressions.The paper specifically notes the challenge of modeling coloration, brightness, and contrast variation in Camelyon17-WILDS.
  • Learned transformation models: The paper learns domain transformation models from unlabeled training-domain data and argues that state-of-the-art domain generalization does not require access to the true model G.MUNIT is used to parameterize learned transformations in the experiments.
  • Learned transformation models: MUNIT is selected because its multimodal mapping generates diverse images from different style codes, supporting invariance over a wider class of images than deterministic mappings.Figures 12–14 show generated multimodal samples for Camelyon17-WILDS, FMoW-WILDS, and PACS.
Loading 2102.11436v5…