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

1. Problem Statement

Domain Generalization (DG) aims to train models on one or more source domains that generalize to unseen target domains without accessing target data during training. The prevailing approach in DG relies on learning domain-invariant representations, typically by enforcing the conditional distribution $P(Z|Y)$ to remain consistent across all source domains. Methods such as feature alignment, Invariant Risk Minimization (IRM), and gradient matching operate under the assumption that a single, globally shared predictive structure exists across all environments.

However, the authors argue that this assumption is often too restrictive. In real-world scenarios, domains may only partially overlap in their predictive structures; specific features might be stable in some domains but vary in others. Enforcing global invariance across an increasing number of heterogeneous domains progressively restricts the feasible representation space. Theoretical analysis and empirical observations (Figure 1) demonstrate that as the number of source domains grows, the mutual information between the learned representation and the label ($I(Z; Y)$) monotonically decreases. Consequently, valuable predictive factors that are stable only within specific subsets of domains are discarded, leading to degraded generalization performance when domain heterogeneity is high.

2. Methodology: MESSI

To address the limitations of global invariance, the paper proposes MESSI (Mixture-of-Experts for Subset-Shared Invariances). The core hypothesis is that invariant structure should be modeled as subset-shared rather than globally shared. The method employs a Mixture-of-Experts (MoE) architecture to learn a structured decomposition of domain relationships.

2.1 Architecture

MESSI utilizes a shared backbone encoder (e.g., DeiT) followed by a router and multiple expert networks:

1. Shared Encoder: Produces a feature representation $u = b_\phi(x)$.
2. Router: Computes a soft distribution over experts $\pi(x) = \text{softmax}(g_\theta(u))$, effectively approximating a latent subset variable $S$.
3. Experts: Each expert $m$ produces a specific representation $z^{(m)} = h_m(u)$.
4. Aggregation: The final representation is a weighted sum $z(x) = \sum \pi_m(x) z^{(m)}$, allowing for the adaptive composition of subset-invariant components.

2.2 Training Objectives

The model is trained to optimize a composite loss function (Equation 5):
$$\mathcal{L} = \mathcal{L}_{cls} + \lambda_{ssi}\mathcal{L}_{OT}^{ssi} + \lambda_{sp}\mathcal{L}_{sp} + \lambda_{bal}\mathcal{L}_{bal} + \lambda_{div}\mathcal{L}_{div}$$

• Routing-Conditioned Subset Alignment ($\mathcal{L}_{OT}^{ssi}$): Instead of aligning all domain pairs, invariance is enforced selectively. For each expert $m$, the method aligns class-conditional feature distributions only between domain pairs $(i, j)$ where the expert has high routing responsibility for a specific class $c$. The alignment weight $a_{ijc}^{(m)}$ is derived from the routing probabilities. The discrepancy is measured using Optimal Transport (OT) (or an MMD variant), weighted by these routing masses. This creates a "routing-induced invariance" where the model learns where invariance should hold.
• Expert Specialization ($\mathcal{L}_{div}$): To prevent experts from collapsing into redundant representations, a diversity regularizer minimizes the cosine similarity between the feature directions of different experts.
• Confident and Balanced Routing ($\mathcal{L}_{sp}, \mathcal{L}_{bal}$):
  • $\mathcal{L}_{sp}$ encourages sparse routing at the instance level (low entropy), ensuring each sample activates a specific subset of experts.
  • $\mathcal{L}_{bal}$ ensures balanced expert utilization at the dataset level, preventing a few experts from dominating the traffic.

3. Key Contributions

The paper outlines three primary contributions:

1. Identification of a Fundamental Limitation: The authors demonstrate that enforcing invariance across an increasing number of domains reduces predictive information. They introduce subset-shared invariance, reframing DG as a problem of learning structured, subset-dependent invariances rather than a single global invariant representation.
2. Routing-Based MoE Framework: They propose a method that uses routing to induce a data-driven decomposition of domain relationships. By aligning experts only with the domain subsets they serve, the method enforces invariance selectively, enabling the adaptive composition of invariant components without requiring explicit subset labels.
3. Empirical Validation: The approach demonstrates improved out-of-domain generalization on standard benchmarks, particularly under heterogeneous domain shifts and as the number of source domains increases.

4. Experimental Results

Experiments were conducted on standard DG benchmarks (DomainBed: PACS, OfficeHome, TerraIncognita, DomainNet) and a controlled Rotated-Colored MNIST dataset designed to test domain growth.

• Standard Benchmarks: MESSI (specifically the DeiT-S/16 variant with OT) achieved state-of-the-art average accuracy (66.2%) on DomainBed, outperforming strong baselines including ERM, CORAL, Fishr, SAGM, and other MoE-based methods like GMoE and OMoE.
• Robustness to Domain Growth: On the Rotated-Colored MNIST fixed-budget domain-growth protocol, MESSI showed superior robustness. While global alignment methods (like CORAL) suffered performance drops as the number of source domains increased (due to over-constraint), MESSI maintained high accuracy with a minimal drop (1.86% drop vs. 2.58% for the best CORAL variant).
• Ablation Studies:
  • Removing the subset-conditioned alignment ($\mathcal{L}_{ssi}$) caused the largest performance drop, confirming it as the critical component.
  • Removing diversity or routing regularizers led to expert redundancy and diffuse routing, reducing performance.
  • Comparisons with controls (Global-MoE, Random-Subset-MoE) showed that MESSI learns a distinct alignment structure rather than simply sparsifying a global objective.
• Efficiency: While the OT-based variant introduces training overhead (approx. 12.39x ERM time), the MMD variant reduces this to ~4.24x. Inference costs remain comparable to other MoE methods.

5. Significance and Claims

The paper claims that Domain Generalization should move beyond the paradigm of enforcing a single global invariance. Instead, it argues that predictive structure is inherently subset-dependent.

• Theoretical Insight: The work provides a theoretical justification (via mutual information analysis) for why global invariance fails as domain heterogeneity increases: it forces the model to discard predictive factors that are not universally shared.
• Practical Implication: By modeling invariance through partially shared structures across domain subsets, MESSI achieves greater robustness. The routing mechanism allows the model to discover these subsets directly from data without requiring explicit subset labels or predefined domain groupings.
• Limitations: The authors modestly acknowledge that the method relies on the router's ability to specialize; if source domains are few or class coverage is imbalanced, the routing signals may be noisy. Additionally, the computational cost of pairwise alignment, particularly with Optimal Transport, is higher than standard global alignment methods.

In conclusion, MESSI offers a mechanism to preserve and compose subset-level invariances, suggesting that the future of DG lies in structured, adaptive decomposition rather than uniform global alignment.