Learning Subset-Shared Invariances for Domain Generalization with Mixture-of-Experts
For domain generalization researchers, this work addresses the limitation of global invariance by introducing subset-shared invariance, offering a more flexible approach that improves robustness under increasing domain heterogeneity.
This paper shows that enforcing invariance across all source domains in domain generalization discards transferable predictive factors, and proposes subset-shared invariance via a mixture-of-experts architecture to model partially shared structure, achieving improved out-of-domain generalization on DomainBed benchmarks.
Domain generalization (DG) aims to learn a model from one or more source domains that generalizes to an unseen target domain without accessing target data during training. A common approach enforces invariance of representations across all source domains, assuming predictive structure is globally shared. However, we demonstrate that enforcing invariance across more domains gradually restricts the feasible representation space, discarding transferable predictive factors that are not universally shared. To address this limitation, we propose subset-shared invariance, where predictive structure is assumed stable only within domain subsets. We implement this principle with a mixture-of-experts architecture, where each expert aligns the specific domains it serves and a routing mechanism composes subset-invariant components for prediction. This creates a routing-conditioned invariance, jointly learned with the representation. To facilitate effective decomposition, we develop training objectives that encourage selective alignment, confident and balanced routing, and diverse expert specialization. Experiments on DomainBed benchmarks demonstrate improved out-of-domain generalization and greater robustness under increasing domain heterogeneity. Our results suggest that DG should move beyond enforcing a single global invariance and instead model invariance through partially shared structure across domain subsets.