Learning Subset-Shared Invariances for Domain Generalization with Mixture-of-Experts

Domain generalization (DG) typically aims to learn a single representation that is invariant across all source domains. The authors investigate whether this global invariance assumption is inherently flawed, arguing that enforcing it across increasingly heterogeneous domains restricts the feasible representation space and discards valuable predictive information that is only shared among subsets of domains.

To address the limitations of global invariance, the authors introduce the concept of subset-shared invariance. They formalize this by showing that as the number of domains increases, the mutual information between the representation and the label decreases under global invariance constraints. To implement this, they propose a Mixture-of-Experts (MoE) architecture where a routing mechanism dynamically partitions data into subsets. Each expert is trained to capture predictive features that are invariant only within its assigned subset of domains. The framework includes specific training objectives to ensure effective decomposition: routing-conditioned optimal transport for selective alignment, a diversity regularizer to prevent expert redundancy, and routing constraints to promote both confident and balanced expert utilization.

The authors demonstrate that global invariance leads to a monotonic shrinkage of the invariant feature subspace, which can degrade performance when domains are highly heterogeneous. Their proposed method, which they call MESSI (Mixture-of-Experts for Subset-Shared Invariance), consistently outperforms standard global-invariance baselines on DomainBed benchmarks. The results suggest that the model successfully discovers meaningful domain groupings through its routing mechanism, confirming that invariant structure in real-world data is often distributed across subsets rather than being globally shared.