PaperlayerView exact revision on arXiv

Source-linked AI summary

Towards Principled Disentanglement for Domain Generalization

Hanlin Zhang, Yi-Fan Zhang, Weiyang Liu, Adrian Weller, Bernhard Schölkopf, Eric P. Xing

arXiv:2111.13839v4cs.LGcs.CV

TL;DR

OOD generalization is hindered by spurious correlations and limited access to target-domain information. DDG formulates disentanglement-constrained domain generalization and jointly learns semantic and variation representations, consistently outperforming baselines while producing interpretable structures.

  • Problem

    OOD generalization remains challenging because neural networks can learn spurious correlations, while target-domain data and practical domain labels may be unavailable.

  • Method

    DDG formulates disentanglement between semantic and variation factors as constraints and jointly learns semantic and variation representations through a primal-dual framework.

  • Results

    DDG consistently outperforms baselines across domain-generalization benchmarks, with larger gains in worst-case and variation-rich settings, while uncovering interpretable structures.

  • Takeaways & Limitations

    Disentangling semantic and variation representations offers a principled route to domain generalization with interpretable and controllable generative results.

  • Takeaways & Limitations

    DDG addresses only disentanglement between semantic and variation factors, leaving disentanglement among different variation factors with limited supervision open.

Abstract

from arXiv · show

A fundamental challenge for machine learning models is generalizing to out-of-distribution (OOD) data, in part due to spurious correlations. To tackle this challenge, we first formalize the OOD generalization problem as constrained optimization, called Disentanglement-constrained Domain Generalization (DDG). We relax this non-trivial constrained optimization problem to a tractable form with finite-dimensional parameterization and empirical approximation. Then a theoretical analysis of the extent to which the above transformations deviates from the original problem is provided. Based on the transformation, we propose a primal-dual algorithm for joint representation disentanglement and domain generalization. In contrast to traditional approaches based on domain adversarial training and domain labels, DDG jointly learns semantic and variation encoders for disentanglement, enabling flexible manipulation and augmentation on training data. DDG aims to learn intrinsic representations of semantic concepts that are invariant to nuisance factors and generalizable across domains. Comprehensive experiments on popular benchmarks show that DDG can achieve competitive OOD performance and uncover interpretable salient structures within data.

1. Introduction

The paper frames OOD generalization as learning representations that preserve class semantics while remaining invariant to cross-domain variation. It proposes DDG, a disentanglement-constrained framework with theoretical guarantees, controllable augmentation, and competitive benchmark performance.

  • Spurious correlations cause modern neural networks to learn shortcuts that undermine out-of-distribution generalization, motivating representations that capture intrinsic class semantics and cross-domain invariance.
  • DDG formalizes disentanglement between class semantics and intra- and inter-domain variations as constraints within the domain-generalization problem.
  • DDG learns disentangled representations that support domain-agnostic, controllable data augmentation and uncover salient structures, achieving competitive performance on RotatedMNIST, VLCS, PACS, and WILDS.The augmentation network requires no domain-specific knowledge or gradient estimation.
  • DDG provides rigorous theoretical guarantees on the empirical duality gap through a principled constrained-learning formulation.
  • The method promotes semantic invariance without adversarial training or domain labels and introduces no additional computational overhead for modeling variations.

2. Related Work

Prior work addresses domain generalization through invariant representations, disentangled factors, data augmentation, and constrained optimization. The paper positions DDG as using disentangled effects to learn invariant representations for realistic OOD generalization tasks.

  • Domain Generalization: Domain generalization seeks representations invariant across domains, but IRM involves challenging bi-level optimization and can fail when test data differ sufficiently from training distributions.IRM and its variants were proposed for this challenge, alongside other domain-generalization approaches.
  • Disentangled Representation Learning: Disentangled representation learning models distinct explanatory factors, and prior work links disentanglement to improved OOD performance and stronger domain-adaptation generalization.Existing models are described as insufficient for supporting compositionality, while related methods may be limited to single-dimensional latent codes and toy datasets.
  • Disentangled Representation Learning: The paper uses disentangled effects to learn invariant representations for realistic OOD generalization rather than restricting them to generation or interpretability.The cited prior methods are characterized as difficult to scale beyond toy datasets and adapt to complicated domain-generalization tasks.
  • Data Augmentation: Data augmentation increases training-distribution diversity and can improve OOD generalization and robustness to spurious correlations.Prior work also learns causal image manipulations through active learning to enrich observational datasets and improve generalization.

3. Disentanglement-constrained Optimization for Domain Generalization

DDG recasts domain generalization as disentanglement-constrained optimization, enforcing semantic invariance while allowing variation factors to change. Finite-dimensional parameterization and empirical approximation make the problem tractable, with theoretical gaps governed by approximation, data, and hypothesis-class complexity.

  • Disentanglement assumptions: Semantic invariance requires reconstructing x from its semantic representation and another sample’s variation representation, while predictions depend only on semantic factors.The domain-shift assumption treats variation in P(X) as the source of distribution shift.
  • Constrained domain generalization: DDG replaces domain-label-dependent worst-domain optimization with inequality constraints that enforce invariance to variation factors.The constrained formulation addresses the difficulty and cost of obtaining domain labels.
  • Parameterization and empirical approximation: The infinite-dimensional constrained problem is made tractable through finite-dimensional parameterization and an empirical dual problem based on n training samples.The parameterization framework includes neural networks and reproducing kernel Hilbert spaces as examples.
  • Theoretical guarantees: The parameterization gap is dominated by semantic-function approximation and reconstruction-based transformations on perturbed inputs.This highlights the role of disentanglement in controlling the deviation between statistical and finite-dimensional problems.
  • Theoretical guarantees: The empirical dual approximation depends on sample size, learning difficulty, and hypothesis-class richness, while disentanglement also supports hard-point augmentation for improved invariance.Larger-capacity neural networks or more training data can improve the algorithm according to the theorem discussion.

4. Experiments

Experiments across four domain-generalization benchmarks show that DDG learns disentangled representations supporting controllable generation and augmentation. DDG also improves worst-case performance and converges more reliably than several baseline approaches.

  • Experimental setup: DDG is evaluated on Rotated MNIST, PACS, VLCS, and WILDS against established domain-generalization baselines.The baselines include ERM, IRM, GDRO, Mixup, MLDG, CORAL, MMD, DANN, CDANN, and AugMix.
  • Disentanglement results: Qualitative results show that DDG disentangles semantic content from variation factors such as background and color.Swapping and interpolating semantic or variation codes enables controllable generation and manipulation without changing object semantics.
  • Domain generalization: DDG consistently outperforms the baselines in most single domains and on average, with larger gains in worst-case PACS domains.The results emphasize that average performance may not adequately reflect OOD generalization when worst-case performance is poor.
  • Data augmentation: DDG’s learned augmentation network improves performance over ERM and AugMix, especially on the worst-case PACS target domain.The experiments connect this gain to the constrained optimization underlying DDG’s augmentation procedure.
  • Convergence analysis: DDG converges more reliably than DANN, CDANN, and IRM on WILDS, while adversarial training methods are unstable or difficult to converge.The primal-dual algorithm produces training dynamics resembling ERM more closely than the less stable comparison methods.
  • Qualitative augmentation comparison: Compared with AugMix, DDG generates more diverse styles, supporting its effectiveness as an automatic data-augmentation method.The qualitative comparison indicates that heuristic-based augmentation has greater difficulty producing stylistically diverse training samples.

5. Concluding Remarks · Appendix

DDG presents a theoretically analyzed and practically implementable framework for domain generalization that separates semantic and variation representations while enforcing invariance. The framework is reported to improve OOD performance and empirical convergence, while enabling interpretable and controllable generation.

  • 5. Concluding Remarks: DDG is a disentangled learning framework for domain generalization with theoretical analyses and a practical algorithmic implementation.The framework is the paper’s central methodological contribution.
  • 5. Concluding Remarks: DDG separates semantic and variation representations into different subspaces.This separation structures the learned representations around semantic and variation factors.
  • 5. Concluding Remarks: The framework enforces invariance constraints during representation learning.Invariance constraints are combined with the subspace separation to support domain generalization.
  • 5. Concluding Remarks: DDG yields superior OOD performance relative to the framework’s stated domain-generalization objective.The supplied passage reports this outcome without specifying a numerical comparison or baseline.
  • 5. Concluding Remarks: The method achieves improved empirical convergence.The passage identifies convergence improvement as an empirical benefit of DDG.
  • 5. Concluding Remarks: DDG produces interpretable and controllable generative results.The learned disentanglement supports both interpretation and manipulation of generated outcomes.

A. Proofs · A.1. Proof for parameterization gap

The proof establishes the parameterization gap by relating the constrained primal and dual objectives through strong duality, then bounding deviations using function-class inclusion, Lipschitz continuity, and parameterization errors. It constructs the corresponding primal problem and identifies the perturbation term governing the resulting bound.

  • A.1. Proof for parameterization gap: Strong duality under feasibility yields a saddle-point condition equating the primal and dual optimal objectives.The condition is stated for all admissible ˜λ, ˜fs, and ˜fv.
  • A.1. Proof for parameterization gap: Because Hθ ⊆ F, the proof derives a lower bound from the definition of D∗.The lower-bound step uses the inclusion relation between the parameterized and unrestricted function classes.
  • A.1. Proof for parameterization gap: The upper-bound argument adds and subtracts R ≡ R(fs, fv, λ), separating objective and constraint discrepancies.This decomposition is introduced in Eq. (16).
  • A.1. Proof for parameterization gap: The combined discrepancy is bounded by (1 + |λ|) times the maximum of loss and constraint differences.The proof rewrites both terms as expectations over the data and auxiliary samples.
  • A.1. Proof for parameterization gap: Hölder, Jensen, and Lipschitz inequalities reduce the discrepancy to representation and decoder deviations weighted by Lℓ and Ld.The final inequality uses the ϵs and ϵg parameterization errors for fs and D.
  • A.1. Proof for parameterization gap: To apply strong duality, the proof constructs a primal constraint-learning problem whose perturbation function is m = max {Lℓϵs, Ldϵg}.The construction is specified for λ > 0 and identifies the parameterization-induced perturbation.
  • A.1. Proof for parameterization gap: After expanding and rearranging the dual objective, feasibility identifies it as the dual of the constructed primal problem.The rearranged expression includes the additive perturbation terms involving m.
  • A.1. Proof for parameterization gap: Regularity and the saddle-point condition identify primal solutions, dual variables, and the final optimality relation, with the last step following from Eq. (14).The argument applies the condition for all admissible function pairs before selecting the optimal solution.

A.2. Proof for empirical gap

The proof bounds the empirical gap by applying KKT and complementary slackness conditions, reducing the objective to its unconstrained terms and then invoking VC-dimension bounds. The resulting inequality holds with probability 1 −δ and completes the proof.

  • KKT and complementary slackness conditions establish the key relation used to upper-bound D∗.
  • The constraint-related objective terms vanish, reducing the empirical gap to the remaining objective term.
  • Classical VC-dimension bounds control the remaining terms, yielding the final bound with probability 1 −δ.

A.3. Proof for empirical duality gap

The empirical duality-gap proof follows by combining the parameterization gap and empirical gap using the triangle inequality.

  • A.3. Proof for empirical duality gap: The proof combines the parameterization and empirical gaps via the triangle inequality.This completes the proof of the empirical duality gap.

B. Domain Generalization by Learning on Fictitious Distributions · C. Experimental Settings · C.1. Other Training Details

The paper extends domain adaptation to domain generalization through a fictitious source distribution, decomposing unseen-domain risk into empirical error, divergence, and source-data diversity. It argues that disentangled semantic and variation representations tighten this bound, while experiments use Adam optimization and publicly available implementation details.

  • B. Domain Generalization by Learning on Fictitious Distributions: B. Domain Generalization by Learning on Fictitious Distributions: A fictitious distribution extends classical domain adaptation to domain generalization and decomposes target risk into empirical error, domain divergence, and source-domain diversity.The analysis connects these terms to a tighter risk upper bound from learning disentangled representations.
  • B. Domain Generalization by Learning on Fictitious Distributions: B. Domain Generalization by Learning on Fictitious Distributions: The diversity term is minimized when source-domain mixtures can approximate unseen targets, requiring diverse data or strong augmentation.The passage identifies source-domain coverage as the key condition for reducing this part of the bound.
  • B. Domain Generalization by Learning on Fictitious Distributions: B. Domain Generalization by Learning on Fictitious Distributions: The upper bound favors feature spaces where source-domain labels are hard to distinguish, corresponding to zero pairwise domain divergence.This criterion is stated as dH[DS,l, DS,k] = 0 for all source-domain pairs.
  • C. Experimental Settings: C. Experimental Settings: The paper reports that compositional generalization and extrapolation can improve when training-domain data are sufficiently rich.It notes that predetermined augmentations such as rotation, cropping, Gaussian blur, and color jitter require prior knowledge and domain-specific expertise.
  • B. Domain Generalization by Learning on Fictitious Distributions: B. Domain Generalization by Learning on Fictitious Distributions: The derived bound combines empirical loss, source-domain divergence minimization, and diversity or coverage of the source-domain mixture.The result motivates separate semantic and variation encoders together with a disentanglement term.
  • C.1. Other Training Details: C.1. Other Training Details: All models are optimized using Adam, with detailed hyperparameter settings provided in the publicly available DDG code.The implementation is available at https://github.com/hlzhang109/DDG.

C.2. Dataset Statistics and Visualization

The section visualizes PACS, VLCS, and WILDS to compare their diverse image styles and domain structures. These differences make knowledge transfer across datasets challenging.

  • PACS: PACS comprises four domains: art, cartoons, photos, and sketches.The figure labels these domains A, C, P, and S.
  • VLCS: VLCS comprises four domains: Caltech101, LabelMe, SUN09, and VOC2007.The figure labels these domains C, L, S, and V.
  • WILDS: WILDS includes data from five different medical centers, while the datasets exhibit diverse image styles that challenge knowledge transfer.The samples provide an intuitive comparison across the datasets and their domains.

D. Additional Experimental Results

Additional experiments show that DDG supports diverse, semantically controlled interpolation and disentangles intra- and inter-domain variations while preserving semantic information. On PACS, DDG consistently outperforms MBDG and its variant except on domain S, attributed to stronger variation modeling without domain labels.

  • Qualitative Results: Compared with AugMix, DDG produces more diverse training samples through interpolation than the heuristic-based method.The comparison is presented qualitatively in the referenced figures.
  • Qualitative Results: DDG generates samples with intermediate variation states by combining one image’s semantic code with mixtures of two variation codes.This demonstrates controllable interpolation through the learned semantic and variation representations.
  • Qualitative Results: Qualitative swaps across RotatedMNIST, PACS, and WILDS demonstrate DDG’s disentanglement of semantic and variation factors.On RotatedMNIST, DDG captures intra-domain thickness and inter-domain rotated-angle variations while maintaining semantic information such as feature color.
  • Numerical Comparison: On PACS, DDG consistently outperforms MBDG and its variant except on domain S.The comparison uses the authors’ source code under a test-domain validation protocol.
  • Numerical Comparison: DDG’s performance gain is attributed to capturing within-data variations through random sampling without domain labels and improving robustness to inter- and intra-domain nuisance factors.The method parameterizes and constrains semantic and variation encoders through disentanglement, whereas MBDG relies on a pretrained generator for inter-domain variation simulation.
Loading 2111.13839v4…