Robust Domain Generalization under Divergent Marginal and Conditional Distributions
Jewon Yeom, Kyubyung Chae, Hyunggyu Lim, Yoonna Oh, Dongyoon Yang, Taesup Kim
TL;DR
The paper addresses the challenge of domain generalization when both the marginal label distribution $P(Y)$ and the class-conditional distribution $P(X|Y)$ shift across domains. It derives a theoretical bound that decomposes the unseen-domain risk into prior-shift and feature-shift components and connects the feature-shift term to a practical Domain-Class Distribution Alignment loss, enabling tractable optimization. Building on this bound, it introduces RC-Align, a meta-learning framework using a two-stage MAML-style procedure with CE and DA losses (and Manifold Mixup) to minimize the bound across source domains, achieving state-of-the-art results on standard DG benchmarks and superior robustness in Multi-Domain Long-Tailed (MDLT) settings. The approach demonstrates that explicit, bound-driven alignment of domain-class distributions yields stronger generalization than ERM or domain-invariant methods alone, with empirical and theoretical support. The work promises practical impact for real-world deployments facing evolving class priors and appearance distributions, while acknowledging limitations such as computational overhead and centroid estimation in highly imbalanced regimes.
Abstract
Domain generalization (DG) aims to learn predictive models that can generalize to unseen domains. Most existing DG approaches focus on learning domain-invariant representations under the assumption of conditional distribution shift (i.e., primarily addressing changes in $P(X\mid Y)$ while assuming $P(Y)$ remains stable). However, real-world scenarios with multiple domains often involve compound distribution shifts where both the marginal label distribution $P(Y)$ and the conditional distribution $P(X\mid Y)$ vary simultaneously. To address this, we propose a unified framework for robust domain generalization under divergent marginal and conditional distributions. We derive a novel risk bound for unseen domains by explicitly decomposing the joint distribution into marginal and conditional components and characterizing risk gaps arising from both sources of divergence. To operationalize this bound, we design a meta-learning procedure that minimizes and validates the proposed risk bound across seen domains, ensuring strong generalization to unseen ones. Empirical evaluations demonstrate that our method achieves state-of-the-art performance not only on conventional DG benchmarks but also in challenging multi-domain long-tailed recognition settings where both marginal and conditional shifts are pronounced.