arXiv CSopen access

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

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

arXiv:2606.25665v1 [cs.LG] 24 Jun 2026

Tien-Hung Nguyen∗, Tien-Dat Tran∗, M.-Duong Nguyen, Kok-Seng Wong† VinUniversity, Vietnam {26hung.nt, duong.nm2, dat.tt6, wong.ks}@vinuni.edu.vn

Abstract 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.

1

Introduction

Domain generalization (DG) refers to the challenge of training models that perform well on unseen domains, without having access to data from these target domains during training [46, 23, 34]. A widely used approach in DG is to develop domain-invariant representations by ensuring that the conditional distribution P (Z|Y ) remains consistent across different domains. This idea is central to many methods, including feature alignment [57, 51] and invariant learning frameworks such as invariant risk minimization (IRM) [1] and its extensions [53, 56]. Nonetheless, applying invariance uniformly to all domains may be too restrictive (see Figure 1). In real-world settings, domains may only partially overlap in structure. For instance, specific features could be stable in some domains but vary in others [52, 68]. Therefore, imposing global invariance can eliminate useful predictive information, especially when domain shifts are heterogeneous [21, 62]. This observation indicates that predictive structure is not universally shared, but instead distributed among subsets of domains. Similar patterns of subset-based heterogeneity are also found in multimodal learning, where individual samples may use different modality combinations [71]. To formalize this limitation, we analyze the effect of enforcing invariance across an increasing number of domains and show that it progressively reduces the mutual information between learned representations and labels. This reveals an inherent trade-off between invariance and predictive information. As invariance constraints become stronger, the representation becomes less informative ∗ Co-first Authors. † Corresponding Author: [email protected]

for prediction. This observation suggests that DG should not enforce invariance globally, but should instead selectively preserve predictive structure that is stable only within subsets of domains. We therefore argue that invariance in DG should be structured rather than global. Specifically, we introduce subset-shared invariance, where invariant structure is assumed to hold within subsets of domains rather than across all domains. This reframes DG as a structured decomposition problem. Instead of learning a single invariant representation, the model should capture multiple invariant components, each corresponding to a subset of domains, and combine them adaptively for prediction. To instantiate this principle, we leverage a mixture-of-experts (MoE) architecture [30, 55] as a mechanism for learning such a decomposition. A routing function induces a soft partition over inputs, which in turn defines subsets of domains associated with each expert. We enforce invariance only within the subsets defined by routing, allowing the model to learn in which domains invariance is appropriate. This results in a routing-induced invariance, where the model learns where invariance should hold, rather than enforcing it uniformly across all domains. MoE models have been used in DG to boost both capacity and specialization [35, 9, 14]. However, earlier work relies on experts to act as feature extractors or domain-specific modules. Our method instead uses routing to directly shape how invariance is applied, so that expert specialization helps find invariant rela- Figure 1: Test accuracy changes as the numtionships specific to certain subsets of domains. ber of source domains increases under a fixed To operationalize this idea, we propose a subset- training budget. Adding more domains at conditioned invariance objective that aligns class- first improves performance, but after a cerconditional feature distributions only across domain tain point, accuracy begins to decline. This pairs selected by routing. This alignment is im- pattern shows an important limitation of enplemented using an optimal transport formulation forcing global invariance. When more hetweighted by expert-specific routing mass, which in- erogeneous domains are included, the shared duces soft domain assignments. To encourage mean- invariant signal weakens and valuable predicingful decomposition, we introduce a diversity regu- tive information is lost. In contrast, MESSI larizer across experts, along with routing objectives stays robust, implying that preserving subsetthat promote both confident (instance-level sparse) specific structure is crucial for generalization and balanced (dataset-level uniform) expert utiliza- when domains differ significantly. tion. We evaluate our approach on standard DG benchmarks and demonstrate consistent improvements over strong baselines. Further analysis shows that the learned routing patterns correspond to meaningful domain groupings, supporting the hypothesis that invariant structure is inherently subset-dependent rather than globally shared. Our main contributions are as follows. • We identify a fundamental limitation of global invariance in DG and show that enforcing invariance across increasing numbers of domains can reduce predictive information. Based on this insight, we introduce subset-shared invariance, which reframes DG as learning structured, subset-dependent invariances rather than a single globally invariant representation. • We propose a routing-based MoE framework that learns where invariance should hold. By using routing to induce a data-driven decomposition of domain relationships, our method enforces invariance selectively across domain subsets, enabling adaptive composition of invariant components. • We demonstrate that the proposed approach improves out-of-domain generalization on standard DG benchmarks, particularly under heterogeneous domain shifts, and provides robustness as the number of source domains increases.

2

Related Work

Domain Generalization and Invariant Representation Learning. DG is commonly approached by learning representations that are invariant across domains, typically by enforcing the conditional 2

distribution P (Z|Y ) to be consistent across environments [46, 18, 42]. This principle underlies a wide range of methods, including feature alignment approaches [57, 51, 39] and invariant learning frameworks such as IRM and its variants [1, 53, 56]. Subsequent work has explored different mechanisms to approximate invariance, including data augmentation [54, 73], episodic and metalearning strategies [2, 17], and self-supervised learning [6, 43]. More recent approaches refine invariance through causal reasoning [44, 43, 69, 27, 36], disentangled representations [72, 11, 50], and indirect alignment objectives [63]. Despite their differences, these approaches largely share the assumption that a single representation can be made invariant across all domains. In contrast, we argue that invariant structure is often only partially shared, and should be modeled at the level of domain subsets. Optimization and Structural Constraints. A complementary line of work studies DG from an optimization perspective, emphasizing that generalization depends not only on representation invariance but also on the solutions reached during training. Methods in this direction improve robustness through gradient regularization [53, 56, 47, 16], geometry-aware optimization, and flatness constraints [7, 61, 45]. More recent approaches explicitly shape training trajectories [3, 13, 15] or constrain interference with pretrained knowledge [29, 70, 31, 66]. While these methods improve generalization through optimization dynamics, they do not explicitly address how invariant structure should be distributed across domains. Our work is complementary, focusing on the structure of invariance itself rather than the optimization process. Modular Experts and Parameter-Efficient Adaptation. Recent DG work has moved toward modular architectures instead of a single shared model. Expert-based methods, such as sparse experts [35, 10] and multi-expert distillation [9], use specialized modules to capture transferable patterns. Parameter-efficient approaches such as prompt-based [65, 41, 38], adapter-based [32], low-rank [40], multimodal [67], and regularization-based [8, 64, 12, 48] are focusing on targeted adaptation. While these methods improve robustness through specialization, they rarely model predictive factors shared only by subsets of domains or use routing for invariance. Recent MoE-based work, such as FlexMoE [71], addresses subset-shared structure in multimodal settings with predefined indices, but requires explicit subset labels that are usually unavailable in DG. Difference from prior work. Unlike prior expert-based or parameter-efficient DG methods, which primarily exploit specialization or constrained adaptation, our approach by contrast, models which domains should be aligned. We introduce subset-shared invariance and use routing-conditioned expert alignment, so that experts capture predictive factors stable within only certain domain subsets. This avoids forcing all information into a single invariant representation. Unlike FlexMoE, which require predefined subset structures, our method discovers subset structure directly from data. The routing mechanism adaptively determines which domains to align, enabling data-driven decomposition of invariant structure without explicit labels for subset membership.

3

Rethinking Feature Invariance in Domain Generalization

We revisit the concept of feature invariance in DG. We show that enforcing invariance across an increasing number of domains progressively imposes stronger constraints on the learned representation, leading to a monotonic shrinkage of the invariant feature subspace. Consequently, globally invariant representations may discard predictive information that is shared only across subsets of domains. 3.1

Problem Setup

We consider the multi-source DG setting with K source domains and no access to target-domain data during training. Formally, let X ∈ X denote the input, Y ∈ Y the label, and D ∈ {1, . . . , K} the domain index, where {Dk }K k=1 are the source-domain distributions over (X, Y ). A predictor is parameterized as ŷ = fθ (z) with representation z = gϕ (x), where gϕ : X → Rd is a feature extractor and fθ is a classifier. Training then minimizes the empirical risk across source domains [46, 23, 59]: Lcls =

K   1 X E(x,y)∼Dk ℓ(fθ (gϕ (x)), y) . K k=1

3

(1)

3.2

Monotonic Shrinkage in Globally Learned Invariance

To analyze the multi-DG problem, we first reformulate the general DG optimization problem as follows. Let Z = gϕ (X) denote the learned representation. For a subset of source domains K ⊆ {1, . . . , K}, define the pairwise index set P(K) = {(i, j) ∈ K × K : i < j}. We consider the following information-theoretic objective: X max JK (ϕ) = I(Z; Y ) − λ I(Z; D | Y, D ∈ {i, j}) , λ > 0. (2) ϕ

(i,j)∈P(K)

Theoretically, the first term encourages label-predictive representations, while the second term penalizes domain-discriminative information within each domain pair after conditioning on the label. We denote the optimal value by VK = supϕ JK (ϕ). From the above objective function, we have the following propositions. Proposition 3.1. Let K ⊆ K′ . Then, for every ϕ, we have JK′ (ϕ) ≤ JK (ϕ). Consequently, the following inequality holds: VK′ ≤ VK . Proposition 3.2. For any (i, j) ∈ P(K), I(Z; D | Y, D ∈ {i, j}) = 0 if and only if the conditional distributions satisfy P (Z|Y, D = i) = P (Z|Y, D = j). Proposition 3.1 formalizes that enlarging the domain set introduces additional nonnegative pairwise invariance penalties. As a result, the objective decreases pointwise for every ϕ, and hence the optimal value is non-increasing. Proposition 3.2 shows that each pairwise term vanishes if and only if P (Z|Y, D = i) = P (Z|Y, D = j). Enforcing this over all (i, j) ∈ P(K) yields Z ⊥ D|Y , i.e., global conditional invariance across domains. From the Propositions 3.1 and 3.2, we have the following theorems. Theorem 3.3. Assume that, for a domain set K, we have I(Z; D | Y, D ∈ {i, j}) = 0, ∀(i, j) ∈ P(K). Then, we have P (Z|Y, D = i) = P (Z|Y, D = j)

∀i, j ∈ K,

(3)

and thus, Z ⊥ D|Y over D ∈ K. Theorem 3.4. Suppose that for every domain set K, there exists an optimizer ϕ∗K of equation 2 satisfying exact pairwise conditional invariance: ∗ I(ZK ; D | Y, D ∈ {i, j}) = 0

∀(i, j) ∈ P(K),

(4)

∗ ∗ ∗ where ZK = gϕ∗K (X). Then, for K ⊆ K′ , we have I(ZK ′ ; Y ) ≤ I(ZK ; Y ).

Theorem 3.3 shows that each pairwise term enforces equality of conditional distributions across domains, and that vanishing penalties over all pairs imply global invariance Z ⊥ D|Y . Theorem 3.4 further shows that enlarging the domain set introduces additional invariance constraints, reducing the optimal objective value and, under exact invariance, the achievable label information I(Z; Y ). Key takeaway. The above analysis reveals an inherent trade-off between domain invariance and ∗ predictive information. For a moderate number of domains, the optimal representation ZK can ∗ retain sufficient label-relevant information while satisfying invariance constraints, i.e., I(ZK ;Y ) remains above the level required for accurate prediction. However, as the number of domains increases, the invariance constraints become progressively stronger, leading to a monotonic decrease in ∗ I(ZK ; Y ). When this quantity falls below the level necessary for reliable prediction, the representation becomes overly invariant and loses discriminative information, resulting in degraded generalization performance. Empirically, we observe a non-monotonic pattern under the fixed-budget domaingrowth setting (see Appendix D.3). Specifically, increasing K initially improves target accuracy, but beyond a certain point leads to performance degradation. This behavior is consistent with the predicted trade-off: moderate domain diversity encourages the suppression of spurious correlations and improves invariance, while excessive heterogeneity enforces overly strong invariance constraints that suppress label-relevant information. 3.3

Subset-Conditioned Invariance

The analysis above suggests that enforcing invariance uniformly across all domains may be overly restrictive when predictive structure is only partially shared. We therefore consider a structured formulation in which invariance is conditioned on latent subsets of domains. 4

Definition 3.5 (Subset-Conditioned Invariance). Let S ∈ {1, . . . , M } denote a latent subset variable (m) and let Z (m) = gϕ (X) be the representation associated with subset m. We say that Z (m) satisfies subset-conditioned invariance if I(Z (m) ; D|Y, S = m) = 0. Equivalently, conditioned on the label and subset assignment, the representation is invariant to domain-specific variation, i.e., P (Z (m) |Y, D = i, S = m) = P (Z (m) |Y, D = j, S = m), for all domains i, j satisfying P (D = i|S = m) > 0 and P (D = j|S = m) > 0. Given Definition 3.5, we consider the following objective: max ϕ,π

M X

I(Z (m) ; Y |S = m),

s.t. I(Z (m) ; D|Y, S = m) = 0,

∀m.

m=1

This formulation encourages each component Z (m) to capture predictive features that are invariant within a subset of domains. Lemma 3.6 (Subset-Conditioned Invariance). For a fixed subset m, the condition I(Z (m) ; D|Y, S = m) = 0 holds if and only if, for all domains i, j such that P (D = i|S = m) > 0 and P (D = j|S = m) > 0, P (Z (m) |Y, D = i, S = m) = P (Z (m) |Y, D = j, S = m). Corollary 3.7. Subset-conditioned invariance can be enforced through pairwise alignment across domains within each subset. Theorem 3.8 (Optimality of Subset-Conditioned Invariance). Suppose there exist subsets {Dm }M m=1 such that for each subset m, P (Y |X, D = i) = P (Y |X, D = j), ∀i, j ∈ Dm , and the predictive structures differ across subsets. Then any representation satisfying global invariance I(Z; D|Y ) = 0 may discard predictive information, whereas a subset-conditioned representation {Z (m) } satisfying I(Z (m) ; D|Y, S = m) = 0 for each m can preserve all subset-specific predictive factors. The above results show that enforcing I(Z (m) ; D|Y, S = m) = 0 is equivalent to aligning classconditional feature distributions across domain pairs within each subset. This provides a practical route to approximate the objective via pairwise alignment weighted by subset assignment probabilities. In the next section, we develop a method that instantiates this formulation using a MoE architecture, where experts represent subset-invariant components and routing approximates the latent subset. Proofs of all propositions, lemmas, and theorems are provided in Appendix A.

4

Proposed Method

To avoid the over-constraint induced by global invariance, we propose a routing-based MoE framework that learns multiple subset-invariant components and composes them adaptively for prediction (Figure 2). Motivated by the subset-conditioned invariance formulation in Section 3, each expert is designed to capture a distinct subset-specific predictive mechanism, while the routing function approximates the latent subset assignment variable S through input-dependent expert selection. Given an input x, a shared encoder produces a feature representation u = bϕ (x), the router outputs a distribution over experts π(x) = softmax(gθ (u)), and each expert produces an expert-specific representation z (m) = hm (u). The final representation is formed by aggregating expert outputs PM (m) according to the routing distribution z(x) = . This allows different subsetm=1 πm (x)z invariant components to be activated adaptively for different samples. This design introduces three challenges: (1) determining which domain pairs should be aligned for each expert, (2) preventing experts from collapsing to redundant representations, and (3) ensuring routing remains both confident at the instance level and balanced across the dataset. To address these challenges, we optimize the following objective: L = Lcls + λssi LOT ssi + λsp Lsp + λbal Lbal + λdiv Ldiv .

(5)

Here, Lcls denotes the classification loss, LOT ssi enforces routing-conditioned subset alignment, Lsp encourages confident routing, Lbal promotes balanced expert utilization, and Ldiv encourages diverse expert representations. We describe each component below. Architectural details are provided in Appendix B.1. Full optimization details and hyperparameters are provided in Appendices F.2 and F.3.

5

Source domains

D1

Expert 1 (h1)

Router π(x)

D2

D1 D2 D3 D4

Expert 1

5

6

Routing scores π(x)

Expert M (hM)

Class c across multiple domains (Expert m)

D1

D4

Z(m)j,c

Z(m)i,c

...

D2 D3

Low

Expert 4 1 2 3 4

Expert 4 (h4)

High

Expert 2 Expert 3

D4

Expert 3 (h3)

A. Routing-Conditioned Subset Alignment

Shared DeiT Backbone

D3

Expert 2 (h2)

Router responsibility of expert m on class c in domain k

Geometry-aware alignment on selected domain pairs

Feature space

Classifier

B. Expert Specialization

Target (unseen)

1

Forward data flow Loss contribution

Expert 1

H1

1

Expert 2

H2

2

⋮

⋮

Expert M

HM

C. Balanced and Confident Routing

M

Σ

2

...

M

After

Before

1

ŷ

1 1 1

Minimize off-diagonal correlation

Feature space

Training Objective

Figure 2: Overview of MESSI. An input is encoded by a shared backbone and routed to multiple experts, whose outputs are composed into the final representation for prediction. The training objective combines classification loss with subset-conditioned invariance, confident and balanced routing, and expert specialization terms. 4.1 Routing-Conditioned Subset Alignment Section 3 shows that enforcing conditional invariance uniformly across all domain pairs can become overly restrictive as the number and heterogeneity of source domains increase. In particular, equation 2 penalizes domain-discriminative information across all pairs (i, j), which leads to progressively stronger constraints as the domain set grows. To avoid this effect, we relax the objective by enforcing invariance selectively over a subset of domain pairs. Specifically, we enforce invariance at the level of expert-specific representations. For each expert m, domain k, and class c, we compute the average P (m) routing responsibility ρk,c = |D1k,c | (x,y)∈Dk , y=c πm (x), which measures how strongly expert m is associated with class c in domain k. Using these routing statistics, we define pairwise alignment (m) (m) (m) weights aijc = σ(αρi,c ) · σ(αρj,c ), which become large when expert m assigns high routing mass to class c in both domains i and j. Consequently, the invariance objective is: LOT ssi =

M XX   X (m) (m) (m) aijc Wε Zi,c , Zj,c ,

(6)

m=1 c∈Y i<j

s.t. Wε (P, Q) =

min γ∈Π(P,Q)

E∥z − z̃∥22 + ε KL(γ ∥ P ⊗ Q).

Our primary formulation uses entropic OT as the discrepancy measure (the motivation is demonstrated in Appendix B.3). An MMD-based variant that uses the same routing-conditioned pair-selection mechanism is described in Appendix B.2. 4.2

Expert Specialization

A final failure mode is redundancy across experts: even with well-behaved routing, different experts may still converge to similar outputs. To discourage this, let Hm ∈ RB×r be the batch output matrix of expert m, and let H̃m = Hm /(∥Hm ∥F + ε) be its normalized version. Specifically, we minimize 2

Ldiv =

X m̸=n

1 ⊤ H̃ H̃n . B m F

(7)

⊤ The matrix B1 H̃m H̃n measures the correlation between the feature directions of experts m and n over the minibatch. If two experts produce similar representations, this matrix has large magnitude and equation 7 increases. Minimizing Ldiv therefore drives experts toward complementary predictive directions rather than redundant ones.

Recent works have explored enforcing orthogonality among experts to promote functional diversity [29, 28, 49, 19, 24]. Methods such as MOORE [28] and OMOE [19] project expert outputs into an 6

orthogonal space via Gram-Schmidt, but this may not preserve task-relevant information and can discard salient features. In contrast, the orthogonality loss in [24] admits a trivial minimum where expert outputs collapse to zero, reducing the penalty without ensuring meaningful diversity. Our specialization loss instead promotes diversity by minimizing cosine similarity between expert outputs, encouraging directional decorrelation while avoiding scale-related degeneracies. 4.3

Balanced and Confident Routing

The subset-conditioned invariance loss determines which class-conditional domain pairs each expert should align, but it does not by itself guarantee that routing produces a meaningful decomposition. In particular, two degenerate regimes can arise. If routing is too diffuse, each input weakly activates many experts, so the representation remains effectively shared and expert specialization does not emerge. If routing collapses globally, only a small subset of experts receives most of the traffic, while the remaining experts are rarely trained and therefore cannot capture distinct subset-shared factors. We address these two failure modes by regularizing routing at both the instance and dataset levels: " M # 2 M  X X 1 Lsp = −Ex πm (x) log πm (x) , Lbal = Ex [πm (x)] − . (8) M m=1 m=1 The two objectives regularize routing at complementary scales. The sparsity term Lsp acts locally on each routing distribution π(x), encouraging low-entropy expert assignment for individual samples. The balancing term Lbal instead operates on the aggregated routing statistics Ex [πm (x)], preventing a small subset of experts from dominating across the dataset.

5

Experiments

Our experiments are organized around three claims. C 1 : Subset-conditioned expert alignment improves standard DG performance. C 2 : MESSI is more robust than globally aligned methods as the number of source domains increases. C 3 : The gains arise from routing-induced subset-conditioned alignment rather than additional MoE capacity, global alignment, or random sparse alignment. Table 1: OOD accuracy (%) under the DomainBed training-domain validation criterion, averaged over three seeds. MESSI-Ti and MESSI-S use DeiT-Ti/16 and DeiT-S/16 backbones, respectively. Best results are in bold; second-best results are underlined. Method

Backbone

PACS ↑

OfficeHome ↑

TerraInc ↑

DomainNet ↑

Avg. ↑

ERM [59] CORAL [57] Fishr [53] SAGM [61] LFME [9]

ResNet-50 ResNet-50 ResNet-50 ResNet-50 ResNet-50

83.8±0.8 83.5±0.4 85.5±0.2 86.4±1.2 84.9±0.4

66.6±0.4 66.0±0.5 68.6±0.2 69.4±0.2 68.5±0.2

47.2±0.8 43.8±0.9 47.4±1.6 48.8±1.0 49.5±0.8

41.6±0.2 39.2±0.3 41.7±0.3 43.2±0.4 38.6±0.2

59.8 58.1 60.8 61.9 60.4

ERM [59] DynMoE [25] GMoE [35] OMoE [19]

DeiT-S/16 DeiT-S/16 DeiT-S/16 DeiT-S/16

86.2±0.1 85.2±0.4 87.3±0.1 87.1±0.3

72.2±0.4 73.4±0.3 73.5±0.1 73.2±0.1

42.0±0.8 44.5±0.6 48.2±0.5 46.0±0.5

47.3±0.2 45.9±0.4 47.8±0.5 46.1±0.2

61.9 62.3 64.2 63.1

MESSI-Ti-OT (Ours) MESSI-S-MMD (Ours) MESSI-S-OT (Ours)

DeiT-Ti/16 DeiT-S/16 DeiT-S/16

86.0±0.3 89.3±0.6 90.9±0.7

69.6±0.2 74.0±0.3 76.1±0.2

42.5±1.1 47.5±0.4 49.3±0.2

43.3±0.4 46.2±0.4 48.6±0.3

60.4 64.3 66.2

5.1

Experimental Setup

Benchmarks. We evaluate standard OOD generalization on DomainBed [23], including PACS [37], OfficeHome [60], TerraIncognita [4], and DomainNet [51], and test domain growth on RotatedColored MNIST (see Appendix C.2 for details), a controlled variant inspired by Rotated MNIST and Colored MNIST [22, 1]. We follow the DomainBed leave-one-domain-out protocol described in Appendix C.1, while Rotated-Colored MNIST allows us to vary the number of source domains to 7

validate the effect of monotonic shrinkage. Details of Rotated-Colored MNIST and the fixed-budget domain-growth protocol are provided in Appendices C.2 and C.3. Implementation. MESSI uses a pretrained DeiT [58] encoder followed by an M =6 expert MoE head, where each expert is a two-layer MLP and the router computes soft expert weights from the final CLS feature. We evaluate both DeiT-Ti/16 and DeiT-S/16 backbones under the same input resolution, augmentations, optimization schedule, and model-selection criterion as the baselines. All experiments are run on NVIDIA H100 GPUs. Baselines and parameter counts are provided in Appendix F. 5.2

DomainBed Results

Table 1 evaluates C 1 by reporting OOD accuracy on DomainBed benchmarks. Reported baselines provide broader DG context, while MESSI-S enables a same-backbone comparison to transformerbased methods. Additional backbone comparisons are presented in Appendix D.1. The samebackbone comparison in Table 1 and the model-size analysis in Appendix F.6 rule out the simple explanation that MESSI improves by using a larger MoE model. MESSI-S uses fewer methodspecific parameters and lower inference cost than GMoE/OMoE, while achieving higher average OOD accuracy. However, parameter count alone does not isolate the training mechanism. We therefore further compare against same-architecture controls in Ap- Table 2: Fixed-budget domain-growth results. Peak is pendix E.1, including a classification- the best test accuracy across source-domain budgets, and only MoE and variants that remove Drop = Peak − Acc(Kmax ). Best results are shown subset-conditioned alignment, routing in bold. Full curves and setup details are provided in sparsity, load balancing, or diversity. Appendix D.3. These controls test whether the gain remains after model capacity, backbone, Method Peak ↑ Acc(Kmax ) ↑ Drop ↓ Avg ↑ routing module, and training budget CORAL (γ=1) 71.96 69.38 2.58 65.38 are fixed. Per-target-domain accuracies CORAL (γ=10) 71.64 68.77 2.87 65.20 for PACS, OfficeHome, TerraIncog- CORAL (γ=100) 70.10 66.46 3.64 63.88 nita, and DomainNet are reported CORAL (γ=250) 68.75 63.67 5.08 62.36 in Appendix D.2. As a supporting CORAL (γ=500) 67.69 59.78 7.91 60.62 many-domain evaluation, we report MESSI (Ours) 75.31 73.45 1.86 68.32 WILDS-iWildCam [34, 5] results in Appendix D.4, using the protocol described in Appendix C.4. 5.3

In-depth Analysis

MESSI is more robust under domain growth. Table 2 evaluates C 2 on Rotated-Colored MNIST under a fixed-budget domain-growth protocol. The target domain and total source training budget are fixed, while the number of source domains increases. This setting separates the effect of increasing domain heterogeneity from the effect of adding more training data. MESSI achieves the best peak accuracy, the best accuracy at Kmax , and the smallest drop from peak accuracy. The strongest CORAL variant reaches 71.96% peak accuracy and drops by 2.58 points at Kmax , whereas MESSI reaches 75.31% and drops by only 1.86 points. This supports C 2 : routing-conditioned subset alignment is less affected by the over-constraint induced by global moment alignment when source-domain heterogeneity grows. The fixed-budget protocol and full curves are provided in Appendices C.3, and D.3. Table 3: C 3 diagnostic on PACS. Routing-conditioned subset alignment explains the gain. Table 4 evaluates C 3 using sameControl MESSI architecture controls on PACS. Both global and ran- Metric dom sparse alignment reduce accuracy relative to Resp. corr. r +1.00 −0.20 ERM-MoE, showing that simply adding alignment Slot overlap J 0.95 0.04 to an MoE architecture is insufficient. Domain-only Discrepancy ↓ 0.234 ± 0.020 0.096 ± 0.014 routing also underperforms, indicating that routing mass alone is too coarse without class-domain responsibility. In contrast, the routing-conditioned class-domain selector avoids the degradation observed with global and random alignment, while the full objective achieves the best performance. 8

These results support C 3 : the improvement depends on where alignment is applied, whereas sparse and balanced routing make the subset decomposition more effective. Table 4: Mechanism ablation under the same MoE architecture. The first group isolates the alignment selector. The last two rows separate routing-conditioned class-domain alignment from routing regularization. Best results are shown in bold. Variant

Alignment selector

ERM-MoE Global-MoE Random-Subset-MoE Domain-Only Routing MESSI w/o Lsp , Lbal MESSI (Ours)

none all class-domain slots matched random slots routing domain mass routing class-domain mass routing class-domain mass

Class-aware

Routing-cond.

PACS ↑

– ✓ ✓ – ✓ ✓

– – – ✓ ✓ ✓

89.1±0.3 87.1±0.2 87.0±0.1 88.2±0.2 89.0±0.2 90.9±0.7

MESSI changes the alignment structure. We further test whether MESSI merely reweights a global alignment pattern or learns a different subset structure. A slot is a class-domain-expert tuple (m, i, j, c), indicating that expert m aligns class c between domains i and j. We compare selectors using the Pearson correlation r between flattened slot weights and the Jaccard overlap J between top-selected slot sets. Higher values indicate similar alignment structures. Table 3 shows that matched random sparse selection remains close to global alignment, with r = 1.00 and J = 0.95. MESSI instead selects a substantially different set of slots, with r = −0.20 and J = 0.04. On the same MESSI-selected slots SM , MESSI also reduces class-conditional discrepancy from 0.234 ± 0.020 to 0.096 ± 0.014. Thus, the gain is not explained by MoE capacity or sparsity alone; MESSI learns a different alignment structure and aligns those selected slots more effectively. Ablation test. Objective ablations are reported in Appendix E.1. Additional diagnostics on expert specialization, routing behavior, and routing-aware alignment controls are provided in Appendices E.2, E.3, and E.4. Appendix E.5 studies the pairwise-to-global invariance trade-off by sweeping the invariance weight and measuring pairwise discrepancy, domain predictability, and target accuracy. Training Cost and Inference Efficiency Table 5 shows that MESSI introduces training-time overhead mainly through the subset-conditioned discrepancy. MESSI-S-MMD costs 4.24× ERM and is close to GMoE+SAGM, while MESSI-S-OT costs 12.39× ERM. The gap between the two MESSI-S variants isolates the cost of the entropic OT solver. MMD gives a non-iterative discrepancy, whereas OT runs Sinkhorn iterations over selected class-conditional domain pairs.

6

Table 5: Training and inference cost on PACS, measured on NVIDIA H100 GPUs. Train time is measured in wall-clock seconds per 1K optimization steps after a 100-step warm-up. Relative time is normalized by ERM. Inference is measured with all training-only losses disabled. Method

Backbone

ERM SAGM GMoE GMoE+SAGM MESSI-Ti-OT MESSI-S-MMD MESSI-S-OT

ResNet-50 103.78±0.50 1.00× ResNet-50 211.37±1.26 2.04× DeiT-S/16 209.71±0.49 2.02× DeiT-S/16 429.88±15.86 4.14× DeiT-Ti/16 1200.86±22.38 11.57× DeiT-S/16 439.68±5.35 4.24× DeiT-S/16 1286.26±28.57 12.39×

Train

Rel.

Infer. 32.42±0.01 32.44±0.01 74.22±0.11 74.17±0.08 23.50±0.01 61.40±0.01 61.41±0.01

Conclusion

We proposed MESSI, a domain generalization framework based on subset-shared invariance. Our analysis shows that enforcing invariance across many domains can restrict the feasible predictive representation and discard factors that are stable only across domain subsets. MESSI addresses this issue with a routing-conditioned MoE model, where each expert aligns the class-conditional domain pairs it actively serves and the routed mixture composes the resulting subset-invariant components for prediction. Experiments on DomainBed and controlled source-domain expansion support the claim that DG methods should preserve and compose subset-level invariances rather than always enforce a single invariant space across all domains. We discuss failure modes, including reliance on reliable routing specialization and additional alignment cost, in Appendix G.

9

References [1] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893, 2019. [2] Yogesh Balaji, Swami Sankaranarayanan, and Rama Chellappa. Metareg: Towards domain generalization using meta-regularization. Advances in neural information processing systems, 31, 2018. [3] Aristotelis Ballas and Christos Diou. Gradient-guided annealing for domain generalization. In Proceedings of the Computer Vision and Pattern Recognition Conference, pages 20558–20568, 2025. [4] Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In Proceedings of the European conference on computer vision (ECCV), pages 456–473, 2018. [5] Sara Beery, Arushi Agarwal, Elijah Cole, and Vighnesh Birodkar. The iwildcam 2021 competition dataset. arXiv preprint arXiv:2105.03494, 2021. [6] Fabio M Carlucci, Antonio D’Innocente, Silvia Bucci, Barbara Caputo, and Tatiana Tommasi. Domain generalization by solving jigsaw puzzles. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 2229–2238, 2019. [7] Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, and Sungrae Park. Swad: Domain generalization by seeking flat minima. Advances in Neural Information Processing Systems, 34:22405–22418, 2021. [8] Junbum Cha, Kyungjae Lee, Sungrae Park, and Sanghyuk Chun. Domain generalization by mutual-information regularization with pre-trained models. In European conference on computer vision, pages 440–457, 2022. [9] Liang Chen, Yong Zhang, Yibing Song, Zhiqiang Shen, and Lingqiao Liu. Lfme: A simple framework for learning from multiple experts in domain generalization. Advances in Neural Information Processing Systems, 37:102919–102947, 2024. [10] Xuweiyi Chen, Wentao Zhou, Aruni RoyChowdhury, and Zezhou Cheng. Point-moe: Largescale multi-dataset training with mixture-of-experts for 3d semantic segmentation. In The Fourteenth International Conference on Learning Representations, 2026. [11] De Cheng, Zhipeng Xu, Xinyang Jiang, Nannan Wang, Dongsheng Li, and Xinbo Gao. Disentangled prompt representation for domain generalization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 23595–23604, 2024. [12] Dong Kyu Cho, Inwoo Hwang, and Sanghack Lee. Peer pressure: Model-to-model regularization for single source domain generalization. In Proceedings of the Computer Vision and Pattern Recognition Conference, pages 15360–15370, 2025. [13] Sumin Cho, Dongwon Kim, and Kwangsu Kim. One-step generalization ratio guided optimization for domain generalization. In Forty-second International Conference on Machine Learning, 2025. [14] Yongxing Dai, Xiaotong Li, Jun Liu, Zekun Tong, and Ling-Yu Duan. Generalizable person re-identification with relevance-aware mixture of experts. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 16145–16154, 2021. [15] Jingfeng Zhang Di Zhao, Hongsheng Hu, Philippe Fournier-Viger, Gillian Dobbie, and Yun Sing Koh. Unlearning during training: Domain-specific gradient ascent for domain generalization. In The Fourteenth International Conference on Learning Representations, 2026. [16] Khoi Do, Nam-Khanh Le, Quoc-Viet Pham, Binh-Son Hua, Won-Joo Hwang, and Duong Nguyen. Domain generalization via pareto optimal gradient matching. In 28th European Conference on Artificial Intelligence, ECAI 2025. IOS Press BV, 2025. 10

[17] Qi Dou, Daniel Coelho de Castro, Konstantinos Kamnitsas, and Ben Glocker. Domain generalization via model-agnostic learning of semantic features. Advances in neural information processing systems, 32, 2019. [18] Chen Fang, Ye Xu, and Daniel N Rockmore. Unbiased metric learning: On the utilization of multiple datasets and web images for softening bias. In Proceedings of the IEEE international conference on computer vision, pages 1657–1664, 2013. [19] Jinyuan Feng, Zhiqiang Pu, Tianyi Hu, Dongmin Li, Xiaolin Ai, and Huimu Wang. Omoe: Diversifying mixture of low-rank adaptation by orthogonal finetuning. In European Conference on Artificial Intelligence, 2025. [20] Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario March, and Victor Lempitsky. Domain-adversarial training of neural networks. Journal of machine learning research, 17(59):1–35, 2016. [21] Robert Geirhos, Kantharaju Narayanappa, Benjamin Mitzkus, Tizian Thieringer, Matthias Bethge, Felix A Wichmann, and Wieland Brendel. Partial success in closing the gap between human and machine vision. Advances in Neural Information Processing Systems, 34:23885– 23899, 2021. [22] Muhammad Ghifary, W Bastiaan Kleijn, Mengjie Zhang, and David Balduzzi. Domain generalization for object recognition with multi-task autoencoders. In Proceedings of the IEEE international conference on computer vision, pages 2551–2559, 2015. [23] Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. In International Conference on Learning Representations, 2021. [24] Hongcan Guo, Haolang Lu, Guoshun Nan, Bolun Chu, Jialin Zhuang, Yuan Yang, Wenhao Che, Xinye Cao, Sicong Leng, Qimei Cui, et al. Advancing expert specialization for better moe. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. [25] Yongxin Guo, Zhenglin Cheng, Xiaoying Tang, Zhaopeng Tu, and Tao Lin. Dynamic mixture of experts: An auto-tuning approach for efficient transformer models. In The Thirteenth International Conference on Learning Representations, 2025. [26] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016. [27] Zhuo He, Shuang Li, Wenze Song, Longhui Yuan, Jian Liang, Han Li, and Kun Gai. Learning time-aware causal representation for model generalization in evolving domains. In Forty-second International Conference on Machine Learning, 2025. [28] Ahmed Hendawy, Jan Peters, and Carlo D’Eramo. Multi-task reinforcement learning with mixture of orthogonal experts. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=aZH1dM3GOX. [29] Jiajun Hu, Jian Zhang, Lei Qi, Yinghuan Shi, and Yang Gao. Learn to preserve and diversify: Parameter-efficient group with orthogonal regularization for domain generalization. In European Conference on Computer Vision, pages 198–216, 2024. [30] Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79–87, 1991. [31] Saqib Javed, Hieu Le, and Mathieu Salzmann. Qt-dog: Quantization-aware training for domain generalization. In Forty-second International Conference on Machine Learning, 2025. [32] Yuyang Ji, Zeyi Huang, Haohan Wang, and Yong Jae Lee. Customizing domain adapters for domain generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 934–944, 2025. [33] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. 11

[34] Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: A benchmark of in-the-wild distribution shifts. In International conference on machine learning, pages 5637–5664. PMLR, 2021. [35] Bo Li, Yifei Shen, Jingkang Yang, Yezhen Wang, Jiawei Ren, Tong Che, Jun Zhang, and Ziwei Liu. Sparse mixture-of-experts are domain generalizable learners. In The Eleventh International Conference on Learning Representations, 2023. [36] Chen Li, Huiying Xu, Changxin Gao, Zeyu Wang, Yun Liu, and Xinzhong Zhu. Towards singlesource domain generalized object detection via causal visual prompts. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. [37] Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain generalization. In Proceedings of the IEEE international conference on computer vision, pages 5542–5550, 2017. [38] Deng Li, Aming Wu, Yaowei Wang, and Yahong Han. Prompt-driven dynamic object-centric learning for single domain generalization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 17606–17615, 2024. [39] Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex C Kot. Domain generalization with adversarial feature learning. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 5400–5409, 2018. [40] Tao Li, Zhengbao He, Yujun Li, Yasheng Wang, Lifeng Shang, and Xiaolin Huang. Flatlora: Low-rank adaptation over a flat loss landscape. In International Conference on Machine Learning, pages 34549–34563. PMLR, 2025. [41] Xinyao Li, Yinjie Min, Hongbo Chen, Zhekai Du, Fengling Li, and Jingjing Li. Generalizing vision-language models with dedicated prompt guidance. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 40, pages 23239–23247, 2026. [42] Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao. Deep domain generalization via conditional invariant adversarial networks. In Proceedings of the European conference on computer vision (ECCV), pages 624–639, 2018. [43] Fangrui Lv, Jian Liang, Shuang Li, Bin Zang, Chi Harold Liu, Ziteng Wang, and Di Liu. Causality inspired representation learning for domain generalization. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 8046–8056, 2022. [44] Divyat Mahajan, Shruti Tople, and Amit Sharma. Domain generalization using causal matching. In International conference on machine learning, pages 7313–7324. PMLR, 2021. [45] Lucas Mansilla, Rodrigo Echeveste, Diego H Milone, and Enzo Ferrante. Domain generalization via gradient surgery. In Proceedings of the IEEE/CVF international conference on computer vision, pages 6630–6638, 2021. [46] Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In International conference on machine learning, pages 10–18. PMLR, 2013. [47] Trong Binh Nguyen, Duong Minh Nguyen, Jinsun Park, Viet Quoc Pham, and Won-Joo Hwang. Federated domain generalization with data-free on-server matching gradient. In The Thirteenth International Conference on Learning Representations, 2025. [48] Yao Ni, Shan Zhang, and Piotr Koniusz. Pace: Marrying generalization in parameter-efficient fine-tuning with consistency regularization. Advances in Neural Information Processing Systems, 37:61238–61266, 2024. [49] James Oldfield, Markos Georgopoulos, Grigorios Chrysos, Christos Tzelepis, Yannis Panagakis, Mihalis Nicolaou, Jiankang Deng, and Ioannis Patras. Multilinear mixture of experts: Scalable expert specialization through factorization. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https://openreview.net/forum?id= bIa03mAtxQ. 12

[50] Tan Pan, Kaiyu Guo, Dongli Xu, Zhaorui Tan, Chen Jiang, Deshu Chen, Xin Guo, Brian C Lovell, LIMEI HAN, Yuan Cheng, et al. Minimal semantic sufficiency meets unsupervised domain generalization. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. [51] Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In Proceedings of the IEEE/CVF international conference on computer vision, pages 1406–1415, 2019. [52] Vihari Piratla, Praneeth Netrapalli, and Sunita Sarawagi. Efficient domain generalization via common-specific low-rank decomposition. In International conference on machine learning, pages 7728–7738. PMLR, 2020. [53] Alexandre Rame, Corentin Dancette, and Matthieu Cord. Fishr: Invariant gradient variances for out-of-distribution generalization. In International Conference on Machine Learning, pages 18347–18377. PMLR, 2022. [54] Shiv Shankar, Vihari Piratla, Soumen Chakrabarti, Siddhartha Chaudhuri, Preethi Jyothi, and Sunita Sarawagi. Generalizing across domains via cross-gradient training. In International Conference on Learning Representations, 2018. [55] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations, 2017. [56] Yuge Shi, Jeffrey Seely, Philip HS Torr, N Siddharth, Awni Hannun, Nicolas Usunier, and Gabriel Synnaeve. Gradient matching for domain generalization. In 10th International Conference on Learning Representations, ICLR 2022, pages 1–28, 2022. [57] Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In European conference on computer vision, pages 443–450, 2016. [58] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pages 10347–10357. PMLR, 2021. [59] Vladimir Vapnik. Principles of risk minimization for learning theory. Advances in neural information processing systems, 4, 1991. [60] Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 5018–5027, 2017. [61] Pengfei Wang, Zhaoxiang Zhang, Zhen Lei, and Lei Zhang. Sharpness-aware gradient matching for domain generalization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 3769–3778, 2023. [62] Yimu Wang, Yihan Wu, and Hongyang Zhang. Lost domain generalization is a natural consequence of lack of training domains. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 15689–15697, 2024. [63] Wei Wei, Zixiong Li, Jing Yan, Mingwen Shao, and Lin Li. Indirect alignment and relationships preservation for domain generalization. In Proceedings of the Thirty-Fourth International Joint Conference on Artificial Intelligence, pages 2054–2062, 2025. [64] Zhixiang Wei, Lin Chen, Yi Jin, Xiaoxiao Ma, Tianle Liu, Pengyang Ling, Ben Wang, Huaian Chen, and Jinjin Zheng. Stronger fewer & superior: Harnessing vision foundation models for domain generalized semantic segmentation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 28619–28630, 2024. [65] Changsong Wen, Zelin Peng, Yu Huang, Xiaokang Yang, and Wei Shen. Domain generalization in clip via learning with diverse text prompts. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 9559–9569, 2025. 13

[66] Boqian WU. Dynamic sparse training versus dense training: The unexpected winner in image corruption robustness. In International Conference on Learning Representations (ICLR), 2025. [67] Zhipeng Xu, Zilong Wang, XINYANG JIANG, Dongsheng Li, De Cheng, and Nannan Wang. Reasoning-driven multimodal llm for domain generalization. In The Fourteenth International Conference on Learning Representations, 2026. [68] Huaxiu Yao, Xinyu Yang, Xinyi Pan, Shengchao Liu, Pang Wei Koh, and Chelsea Finn. Improving domain generalization with domain relations. In The Twelfth International Conference on Learning Representations, 2024. [69] Naiyu Yin, Hanjing Wang, Yue Yu, Tian Gao, Amit Dhurandhar, and Qiang Ji. Integrating markov blanket discovery into causal representation learning for domain generalization. In European Conference on Computer Vision, pages 271–288, 2024. [70] Seokju Yun, Seunghye Chae, Dongheon Lee, and Youngmin Ro. Soma: Singular value decomposed minor components adaptation for domain generalizable representation learning. In Proceedings of the Computer Vision and Pattern Recognition Conference, pages 25602–25612, 2025. [71] Sukwon Yun, Inyoung Choi, Jie Peng, Yangfan Wu, Jingxuan Bao, Qiyiwen Zhang, Jiayi Xin, Qi Long, and Tianlong Chen. Flex-moe: Modeling arbitrary modality combination via the flexible mixture-of-experts. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https://openreview.net/forum?id=ihEHCbqZEx. [72] Hanlin Zhang, Yi-Fan Zhang, Weiyang Liu, Adrian Weller, Bernhard Schölkopf, and Eric P Xing. Towards principled disentanglement for domain generalization. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 8024–8034, 2022. [73] Kaiyang Zhou, Yongxin Yang, Yu Qiao, and Tao Xiang. Domain generalization with mixstyle. In International Conference on Learning Representations, 2021.

14

A

Proofs

A.1

Proof of Proposition 3.1

Proof. Since K ⊆ K′ , we have P(K) ⊆ P(K′ ).

(9)

Therefore, X

I(Z; D | Y, D ∈ {i, j})

(i,j)∈P(K′ )

=

X

X

I(Z; D | Y, D ∈ {i, j}) +

I(Z; D | Y, D ∈ {i, j}) .

(10)

(i,j)∈P(K′ )\P(K)

(i,j)∈P(K)

Every conditional mutual information term is nonnegative. Hence the second sum is nonnegative, which gives X X I(Z; D | Y, D ∈ {i, j}) ≥ I(Z; D | Y, D ∈ {i, j}) . (11) (i,j)∈P(K′ )

(i,j)∈P(K)

Multiplying by −λ with λ > 0 and adding I(Z; Y ) yields JK′ (ϕ) ≤ JK (ϕ). Taking the supremum over ϕ on both sides gives VK′ ≤ VK .

A.2

(12) (13)

Proof of Proposition 3.2

Proof. Under the restricted event {D ∈ {i, j}}, the domain variable takes only two values. By the standard characterization of conditional mutual information, I(Z; D | Y, D ∈ {i, j}) = 0 (14) if and only if Z ⊥ D|Y, D ∈ {i, j}. (15) This conditional independence means that, for each label value, the conditional distribution of Z is the same in the two-domain subpopulation: P (Z|Y, D = i, D ∈ {i, j}) = P (Z|Y, D = j, D ∈ {i, j}). (16) Since {D = i} and {D = j} are subsets of {D ∈ {i, j}}, this is equivalent to P (Z|Y, D = i) = P (Z|Y, D = j). (17) The converse follows by reversing the same argument. A.3

Proof of Theorem 3.3

Proof. By Proposition 3.2, for every (i, j) ∈ P(K), P (Z|Y, D = i) = P (Z|Y, D = j). (18) Thus all domains in K share the same conditional law of Z given Y . Denote this common law by QY . Then for every i ∈ K, P (Z|Y, D = i) = QY . (19) Marginalizing over domains gives X P (Z|Y ) = P (Z|Y, D = i)P (D = i|Y ) i∈K

=

X

QY P (D = i|Y ) = QY .

(20)

i∈K

Therefore, P (Z|Y, D = i) = P (Z|Y ) which is exactly Z ⊥ D|Y over D ∈ K. 15

∀i ∈ K,

(21)

A.4

Proof of Theorem 3.4

Proof. Under the exact invariance assumption, every pairwise penalty term vanishes at ϕ∗K . Therefore, ∗ VK = JK (ϕ∗K ) = I(ZK ; Y ). (22) Similarly, ∗ VK′ = I(ZK (23) ′ ; Y ). ′ From Proposition 3.1, if K ⊆ K , then VK′ ≤ VK . Substituting the two identities above yields ∗ ∗ I(ZK (24) ′ ; Y ) ≤ I(ZK ; Y ).

A.5

Proof of Lemma 3.6

Proof. Fix a subset index m with P (S = m) > 0. By the standard characterization of conditional mutual information, I(Z (m) ; D|Y, S = m) = 0 (25) if and only if Z (m) ⊥ D|Y, S = m. This conditional independence holds if and only if, for every label value y and every domain i with P (D = i|Y = y, S = m) > 0, P (Z (m) |Y = y, D = i, S = m) = P (Z (m) |Y = y, S = m). Therefore, for any two domains i, j with positive probability under subset m,

(26)

P (Z (m) |Y = y, D = i, S = m) = P (Z (m) |Y = y, S = m) = P (Z (m) |Y = y, D = j, S = m). (27) This proves the forward direction. Conversely, suppose that for all domains i, j with positive probability under S = m, P (Z (m) |Y, D = i, S = m) = P (Z (m) |Y, D = j, S = m).

(28)

(m)

Then the conditional distribution of Z given Y and S = m is the same for all such domains. Denote this common distribution by Qy . For any domain i with positive probability, P (Z (m) |Y = y, D = i, S = m) = Qy .

(29)

Moreover, P (Z (m) |Y = y, S = m) =

X

P (Z (m) |Y = y, D = k, S = m)P (D = k|Y = y, S = m) = Qy .

k

(30) Thus, P (Z (m) |Y = y, D = i, S = m) = P (Z (m) |Y = y, S = m),

(31)

Z (m) ⊥ D|Y, S = m.

(32)

I(Z (m) ; D|Y, S = m) = 0.

(33)

which implies Hence,

A.6

Proof of Corollary

Proof. By Lemma 3.6, subset-conditioned invariance for subset m, I(Z (m) ; D|Y, S = m) = 0, is equivalent to equality of the class-conditional feature distributions

(34)

P (Z (m) |Y, D = i, S = m) = P (Z (m) |Y, D = j, S = m) for all domain pairs i, j that have positive probability in subset m.

(35)

Therefore, enforcing equality of these distributions for every such pair is sufficient to obtain subsetconditioned invariance. Conversely, subset-conditioned invariance implies that all such pairwise class-conditional distributions are equal. Hence subset-conditioned invariance can be enforced through pairwise class-conditional alignment within each subset. 16

A.7

Proof on Theorem 3.8

Proof. We give a constructive example. Consider four domains partitioned into two subsets, D1 = {1, 2},

D2 = {3, 4}.

(36)

Let the label be binary, Y ∈ {−1, +1}, with P (Y = 1) = P (Y = −1) = 1/2. For domains in D1 , suppose the label-relevant feature is X1 = Y,

(37)

while another feature X2 is independent noise. For domains in D2 , suppose the label-relevant feature is X2 = Y,

(38)

while X1 is independent noise. Thus, the predictive mechanism is shared within each subset but differs across subsets. Now impose global conditional invariance on a single representation Z = g(X): I(Z; D|Y ) = 0.

(39)

This requires the conditional distribution P (Z|Y, D) to be the same across all four domains. In particular, for fixed Y = y, the distribution of Z in domains where X1 is predictive must match the distribution of Z in domains where X1 is noise. Therefore, any component of Z that preserves the subset-specific predictive role of X1 in D1 would make P (Z|Y, D) differ between D1 and D2 , violating global invariance. The same argument applies to X2 : preserving its subset-specific predictive role in D2 would distinguish domains in D2 from domains in D1 . Hence, a globally invariant representation cannot simultaneously preserve both subset-specific predictive factors while satisfying I(Z; D|Y ) = 0. It must suppress at least one subset-specific predictive factor to make the conditional representation distribution identical across all domains. By contrast, define a latent subset variable S=1

for D ∈ D1 ,

S=2

for D ∈ D2 .

Z (1) = X1 ,

Z (2) = X2 .

(40)

Let (41)

Within subset D1 , X1 = Y in both domains, so P (Z (1) |Y, D = 1, S = 1) = P (Z (1) |Y, D = 2, S = 1).

(42)

I(Z (1) ; D|Y, S = 1) = 0.

(43)

Therefore, Similarly, within subset D2 , X2 = Y in both domains, so I(Z (2) ; D|Y, S = 2) = 0.

(44)

Thus, the subset-conditioned representations satisfy subset-conditioned invariance while preserving the predictive factor specific to each subset. This establishes that subset-conditioned invariance can preserve subset-specific predictive structure that global invariance may suppress.

B

Details of the MESSI MoE Architecture

B.1

MoE Architecture

MESSI implements subset-shared invariance with a sparse MoE representation learner. Let u = bϕ (x) be the backbone feature. A router produces π(x) ∈ ∆M , and each expert hm produces z (m) = hm (u). The prediction uses the routed mixture z(x) =

M X

πm (x)z (m) ,

m=1

17

ŷ = f (z(x)).

(45)

The subset-conditioned invariance loss is applied to the expert features {z (m) }M m=1 . For each expert, class, and domain pair, the aggregated routing mass defines the alignment weight. Thus, routing determines which class-conditional domain pairs each expert aligns, while the classifier predicts from their routed composition. We follow the sparse MoE design of GMoE [35]. In selected ViT/DeiT blocks, the feed-forward network is replaced by a sparse MoE layer fMoE (x) =

M X

Gm (x)Em (x),

(46)

m=1

where Em is the m-th FFN expert and G(x) is a sparse top-k routing distribution. We use the same cosine router and sparse inference pattern as GMoE. Unless otherwise stated, all MoE-based variants use the same backbone, expert count, routing configuration, and MoE placement, so the ablations isolate the training objective rather than model capacity. B.2

Scalability of Subset-Conditioned OT

The subset-conditioned OT loss is computed on the current source minibatch. It does not use targetdomain samples, memory banks, queues, cached features, cross-batch accumulation, or movingaverage feature estimates. At each step, we concatenate the source-domain minibatches and compute labels y, source-domain indices d, routing probabilities πm (x), and per-expert features hm (x). For expert m, class c, and source domain i, define (m)

Zi,c = {hm (x) : y = c, d = i},

(m)

ni,c = |Zi,c |.

(47)

Only active minibatch cells are used. A class c contributes only if it appears in at least two sourcedomain samples in the concatenated minibatch. A domain pair (i, j) contributes only if ni,c > 0 and nj,c > 0. Missing class-domain cells are skipped. Singleton class-domain groups are allowed, since empirical OT is well-defined for non-empty empirical distributions. For each active tuple (m, i, j, c), the OT term is weighted by the minibatch routing responsibility     X 1 (m) (m) (m) (m) aijc = σ α π̄i,c σ α π̄j,c , π̄i,c = πm (x). (48) ni,c x:y=c,d=i

The routing averages are computed from the current minibatch and detached before weighting the OT discrepancy. Thus, routing selects which class-domain pairs are aligned, while the router itself is trained through the classification objective and routing regularizers. For each active pair, we compute entropic OT with a differentiable log-space Sinkhorn solver, uniform empirical marginals, and a squared Euclidean cost matrix over expert features. The entropic regularization coefficient and the number of Sinkhorn iterations T are fixed. Let CB be the set of active classes in the minibatch, and let PB (c) = {(i, j) : i < j, ni,c > 0, nj,c > 0} be the active source-domain pairs for class c. The minibatch-level complexity is   X X OM T ni,c nj,c  ,

(49)

(50)

c∈CB (i,j)∈PB (c)

plus the same pairwise cost-matrix construction without the factor T . The cost therefore scales with the number and size of active class-domain pairs in the minibatch, not with the total number of dataset classes. This distinction matters for large-class datasets such as DomainNet, where most classes are absent from any single minibatch. The worst case remains expensive when many samples from the same class appear across multiple source domains, which explains the higher training-time cost of the OT variant in Table 5. The OT loss is disabled at test time. Inference uses only the encoder, router, expert heads, and classifier, and therefore requires no Sinkhorn iterations or pairwise class-domain computations. 18

MMD variant for computational comparison. We also implement an MMD variant that keeps the same backbone, router, experts, active-pair selection rule, and auxiliary losses, replacing only the OT discrepancy. For each active domain pair (i, j) ∈ PB (c), the loss is LMMD = ssi

M X X

  (m) (m) (m) aijc MMD2 Zi,c , Zj,c .

X

(51)

m=1 c∈CB (i,j)∈PB (c)

With a Gaussian RBF kernel k, MMD2 (P, Q) = Ez,z′ ∼P k(z, z ′ ) + Ez̃,z̃′ ∼Q k(z̃, z̃ ′ ) − 2Ez∼P,z̃∼Q k(z, z̃).

(52)

For groups of size ni,c and nj,c , MMD requires O(n2i,c + n2j,c + ni,c nj,c ) kernel evaluations and no iterative solver. Entropic OT requires the pairwise cost matrix and T Sinkhorn iterations, with cost O(T ni,c nj,c ). MMD therefore has a lower training-time constant factor, whereas OT provides geometry-aware matching at higher cost. Both variants have the same inference-time architecture because the discrepancy loss is used only during training. B.3

From Latent Subsets to Routing-Weighted Alignment

The formulation above uses an ideal latent subset variable S, whereas in practice subset membership is not observed. We approximate this latent variable with the learned router by interpreting πm (x) ≈ P (S = m|x).

(53)

Under this interpretation, the subset-conditioned invariance constraint I(Z (m) ; D|Y, S = m) = 0

(54)

is equivalent to matching the class-conditional distributions P (Z (m) |Y = c, D = i, S = m) = P (Z (m) |Y = c, D = j, S = m)

(55)

for domains assigned to subset m. Because S is latent, we cannot directly select the domain pairs belonging to each subset. Instead, we estimate the responsibility of expert m for class c in domain k by the average routing mass X 1 (m) ρk,c = πm (x). (56) |Dk,c | (x,y)∈Dk , y=c

The product-based weight (m)

(m)

(m)

aijc = σ(αρi,c )σ(αρj,c )

(57)

then acts as a soft relaxation of the indicator that domains i and j both belong to subset m for class c. Thus, the empirical OT objective   X (m) (m) (m) aijc Wε Zi,c , Zj,c (58) m,c,i<j

can be viewed as a tractable relaxation of enforcing subset-conditioned invariance. Exact mutualinformation constraints are replaced by empirical class-conditional distribution matching, and hard latent subset assignments are replaced by learned soft routing responsibilities.

19

C

Experimental Protocols and Datasets

This appendix specifies the datasets and evaluation protocols used in the main text. We first describe the DomainBed leave-one-domain-out protocol for standard OOD generalization. We then introduce Rotated-Colored MNIST, a controlled benchmark that disentangles shape-based prediction from rotation and color-label shortcut shifts, and define the fixed-budget domain-growth protocol used to vary source-domain heterogeneity without increasing the training budget. Finally, we describe the WILDS-iWildCam [34, 5] protocol for the additional many-domain camera-trap evaluation. C.1

DomainBed Protocol

We evaluate OOD generalization on DomainBed using the standard leave-one-domain-out protocol. For each dataset, one domain is selected as the target domain and all remaining domains are used for training. We follow the training-domain validation criterion: 20% of each source domain is reserved for validation, and the checkpoint with the highest average source-validation accuracy is selected. Target-domain data are not used for training, validation, hyperparameter tuning, or model selection. All results are averaged over three random seeds. For each target split, all methods use the same source domains, target domain, and evaluation protocol. The dataset-level score is the mean accuracy over all target domains. C.2

Rotated-Colored MNIST Construction

We construct Rotated-Colored MNIST as a controlled diagnostic for domain generalization. Each environment is defined by two nuisance factors: a rotation angle applied to the digit image and a color-label correlation applied after label noise. The prediction task is binary digit classification, so digit shape is the label-relevant factor, while rotation and color correlation define the environment shift. Algorithm 1 summarizes the construction. We pool the MNIST training and test sets into N = Algorithm 1: Rotated-Colored MNIST construc70,000 examples. The ten-way label is con- tion. verted into a binary label by mapping digits {0, 1, 2, 3, 4} to class 1 and digits {5, 6, 7, 8, 9} # E: number of candidate environments to class 0. Let ȳ denote this clean binary label. # Dth: rotation step Following Colored MNIST, we corrupt it with # Dp: color-correlation step label noise ηy ∼ Bernoulli(0.25): y = ȳ ⊕ ηy ,

X, y10 = load_mnist() X, y10 = shuffle(X, y10)

(59)

where ⊕ denotes XOR. This prevents the binary task from being perfectly deterministic.

C = round(1.0 / Dp) envs = []

Environment parameterization. We construct candidate environments indexed by e ∈ {0, . . . , E − 1}. Each environment has a rotation angle θe and a color-flip probability pe . The grid is controlled by a rotation step ∆θ and a color step ∆p . Let C = round(1/∆p ) be the number of color levels per rotation. Under the row-major layout, jek θe = ∆θ , pe = ∆p (e mod C). C (60) Environments in the same row share the same rotation and differ only in color-label correlation, whereas moving across rows adds geometric variation. In the fixed-budget domain-growth experiments, we use ∆θ = 45◦ and ∆p = 0.1. Increasing 20

for e in range(E): th = Dth * floor(e / C) p = Dp * (e % C) Xe, ye10 = split_env(X, y10, e) y0 = float(ye10 < 5) y = xor(y0, Bernoulli(0.25)) Xr = rotate(Xe, th) c = xor(y, Bernoulli(p)) Xc = two_channel(Xr) Xc[arange(len(Xc)), 1 - c, :, :] = 0 envs.append((Xc / 255.0, y)) return envs

the number of source environments therefore increases heterogeneity in both rotation and shortcut strength, without increasing the total training budget. Rotation and colorization. For each environment e, images assigned to that environment are first rotated by θe using bilinear interpolation. The rotated grayscale image is then duplicated into two channels. The active color channel is determined by the noisy label y and the environment-specific flip probability pe : c = y ⊕ ηc , ηc ∼ Bernoulli(pe ). (61) The channel indexed by c is retained and the other channel is set to zero. The resulting input has shape 2 × 28 × 28 and is normalized to [0, 1]. The parameter pe controls the color shortcut. When pe = 0, color is perfectly correlated with the noisy label. When pe = 0.5, color is independent of the label. When pe approaches 1, the correlation is reversed. Purpose of the benchmark. This benchmark is used only as a diagnostic. Because the rotation angle and color-label correlation are known by construction, it allows us to test whether a method preserves shape-based predictive structure as source-domain heterogeneity increases, or instead collapses toward an overly restrictive globally aligned representation. C.3

Fixed-Budget Domain-Growth Protocol

We use Rotated-Colored MNIST to study how DG methods behave as the number of source domains increases. The construction above defines the candidate environment pool; this protocol defines how source domains are sampled under a fixed training budget. Held-out target environment. For all domain-growth experiments, we fix the target environment to θtest = 0◦ , ptest = 0.5. (62) Under equation 60, this corresponds to etest = 5 when ∆p = 0.1. Since ptest = 0.5, color is independent of the binary label. The target domain therefore tests whether a method learns shapebased prediction rather than relying on the color shortcut. Fixed-budget source-domain expansion. We vary the number of source domains as Ks ∈ {3, 5, 7, 9, 11, 13, 15, 17}. (63) For each Ks , source domains are sampled from the candidate pool excluding etest . The total number of training examples is fixed to B, so each selected source domain contributes at most ⌊B/Ks ⌋ examples. This separates the effect of increasing source-domain heterogeneity from simply increasing the amount of training data. Sampling, seeds, and model selection. For each Ks , we sample R independent source-domain subsets. Each method is trained with three random seeds per subset, and we report mean and standard deviation over both subset sampling and training seeds. Model selection uses source-validation data only; the target domain is used only for final evaluation. Table 6: Fixed-budget domain-growth protocol on Rotated-Colored MNIST. Setting

Value

Source domains Total training examples Max examples per source domain Target environment Subset samples Training seeds Model selection Target usage

C.4

Ks ∈ {3, 5, 7, 9, 11, 13, 15, 17} Fixed to B for all Ks ⌊B/Ks ⌋ θtest = 0◦ , ptest = 0.5 R per Ks 3 per subset Source-validation only Final evaluation only

WILDS-iWildCam Protocol

Benchmark. We additionally evaluate MESSI on WILDS-iWildCam [34, 5], a real-world manydomain benchmark for camera-trap species recognition. Each example is an image captured by 21

a static camera trap, the label is one of 182 species, and the domain is the camera-trap location. The training split contains 243 source locations with substantial variation in background, viewpoint, illumination, vegetation, and species frequency. This makes iWildCam a natural stress test for global invariance: predictive structure need not be shared uniformly across all camera locations. Splits and model selection. We follow the official WILDS five-split protocol, summarized in Table 7. Models are trained on the train split. The ID validation and ID test splits contain held-out images from source camera locations, whereas the OOD validation and OOD test splits contain images from camera traps disjoint from the training cameras. We select checkpoints using OOD validation and report OOD test macro F1 as the primary generalization metric. Table 7: Split structure of WILDS-iWildCam. The domain is the camera-trap location. Split Camera locations Role Train ID Val ID Test OOD Val OOD Test

Source camera traps Source camera traps Source camera traps Unseen camera traps Unseen camera traps

Optimization In-distribution validation In-distribution test OOD checkpoint selection Final OOD evaluation

Compared methods. We compare three MoE-based methods under the same DeiT-S/16 backbone, preprocessing pipeline, and six-expert configuration. GMoE [35] is the original sparse MoE baseline. OMoE [19] augments the MoE baseline with Gram–Schmidt orthogonalization and the importanceCV2 load-balancing objective. MESSI uses the same backbone family and is trained with the proposed subset-conditioned invariance objective. Preprocessing. Images are resized to 224 × 224 and normalized with ImageNet statistics. During training, we apply random resized cropping, horizontal flipping, color jittering, and random grayscale augmentation. This setup is lighter than the official WILDS-iWildCam high-resolution pipeline. We therefore use this experiment as a controlled comparison among related MoE variants, not as a direct reproduction of leaderboard-scale WILDS results. Many-domain batching. We treat each camera-trap location as a domain. Since the training split contains 243 source locations, using all domains in every minibatch is not tractable. We use stochastic domain-subset batching following the WILDS n_groups_per_batch convention [34, 5]. At each step, we sample K = 4 source locations uniformly without replacement and draw 8 images from each location, forming a balanced minibatch of 32 images. For MESSI, these K = 4 active domains define K(K − 1)/2 = 6 candidate domain pairs for Lssi . Optimization and metric. All methods are trained with Adam using learning rate 3 × 10−5 and zero weight decay. For MESSI, we set λssi = 0.01, λsp = λbal = 0.02, λdiv = 0.02, and routing-pair temperature α = 4.0. For conditional MMD, we use a multi-scale RBF kernel with bandwidths {1, 2, 4, 8, 16}. We report per-class macro F1, the standard WILDS metric for iWildCam. Macro F1 gives equal weight to each species and is less dominated by frequent classes such as empty frames. Training budget. All compared methods are trained for 150K steps. The longer schedule is used because the subset-conditioned invariance objective introduces an additional multi-domain alignment term that slows early optimization before stabilizing. We therefore interpret this experiment as a convergence-aware comparison under a shared backbone and preprocessing pipeline, rather than an equal-step compute benchmark.

22

D

Extended Results

D.1

Extended Results on DomainBed

Table 8: OOD accuracy (%) under the DomainBed training-domain validation criterion, averaged over three seeds. MESSI-Ti and MESSI-S use DeiT-Ti/16 and DeiT-S/16 backbones, respectively. Best results are in bold; second-best results are underlined. Method

Backbone

PACS ↑

OfficeHome ↑

TerraInc ↑

DomainNet ↑

Avg. ↑

ERM [59] IRM [1] MMD [39] DANN [20] MixStyle [73] CORAL [57] Fishr [53] SAGM [61] LFME [9]

ResNet-50 ResNet-50 ResNet-50 ResNet-50 ResNet-50 ResNet-50 ResNet-50 ResNet-50 ResNet-50

83.8±0.8 81.1±0.5 81.4±0.7 79.4±0.3 82.7±0.4 83.5±0.4 85.5±0.2 86.4±1.2 84.9±0.4

66.6±0.4 58.2±0.2 60.1±0.5 59.6±0.4 59.8±0.7 66.0±0.5 68.6±0.2 69.4±0.2 68.5±0.2

47.2±0.8 38.7±0.8 42.2±0.9 38.1±0.8 41.1±1.0 43.8±0.9 47.4±1.6 48.8±1.0 49.5±0.8

41.6±0.2 30.6±1.0 20.5±0.3 31.6±0.1 34.1±0.1 39.2±0.3 41.7±0.3 43.2±0.4 38.6±0.2

59.8 52.2 51.1 52.2 54.4 58.1 60.8 61.9 60.4

ERM [59] DynMoE [25] GMoE [35] OMoE [19]

DeiT-S/16 DeiT-S/16 DeiT-S/16 DeiT-S/16

86.2±0.1 85.2±0.4 87.3±0.1 87.1±0.3

72.2±0.4 73.4±0.3 73.5±0.1 73.2±0.1

42.0±0.8 44.5±0.6 48.2±0.5 46.0±0.5

47.3±0.2 45.9±0.4 47.8±0.5 46.1±0.2

61.9 62.3 64.2 63.1

MESSI-Ti-OT (Ours) MESSI-S-MMD (Ours) MESSI-S-OT (Ours)

DeiT-Ti/16 DeiT-S/16 DeiT-S/16

86.0±0.3 89.3±0.6 90.9±0.7

69.6±0.2 74.0±0.3 76.1±0.2

42.5±1.1 47.5±0.4 49.3±0.2

43.3±0.4 46.2±0.4 48.6±0.3

60.4 64.3 66.2

D.2

DomainBed Per-Domain Results

Tables 9–12 report per-target-domain accuracy. Best results are shown in bold, and second-best results are underlined. The average column matches Table 1. Table 9: Per-domain accuracy on PACS. Method ERM IRM MMD CORAL Fishr SAGM DynMoE GMoE OMoE MESSI-Ti-OT (Ours) MESSI-S-MMD (Ours) MESSI-S-OT (Ours)

Art ↑

Cartoon ↑

Photo ↑

Sketch ↑

Avg. ↑

86.2±0.8 83.0±0.6 84.0±0.8 85.2±0.7 87.4±0.3 87.5±1.6 89.0±0.3 89.9±0.4 89.6±0.4 88.3±0.9 90.4±0.6 91.6±0.3

76.0±0.7 74.5±0.4 74.8±0.6 76.6±1.0 78.4±0.4 80.7±1.4 81.2±0.5 83.7±0.5 83.4±0.5 78.1±1.1 86.7±0.8 88.7±0.9

95.9±0.4 95.8±0.3 95.7±0.3 95.1±0.2 97.7±0.1 96.4±0.6 98.9±0.2 98.9±0.1 98.8±0.2 98.1±0.5 98.7±0.4 99.1±0.3

77.1±0.9 71.1±0.7 71.1±1.0 77.1±0.8 78.5±0.5 81.0±3.2 71.7±0.7 76.7±0.6 76.6±0.7 79.3±2.5 81.4±0.6 84.1±1.4

83.8±0.8 81.1±0.5 81.4±0.7 83.5±0.4 85.5±0.2 86.4±1.2 85.2±0.4 87.3±0.1 87.1±0.3 86.0±0.3 89.3±0.6 90.9±0.7

23

Table 10: Per-domain accuracy on OfficeHome. Method ERM IRM MMD CORAL Fishr SAGM DynMoE GMoE OMoE MESSI-Ti-OT (Ours) MESSI-S-MMD (Ours) MESSI-S-OT (Ours)

Art ↑

Clipart ↑

Product ↑

Real-World ↑

Avg. ↑

61.8±0.8 52.1±0.3 54.0±0.5 61.8±0.4 64.1±0.2 64.8±0.7 71.6±0.3 71.5±0.1 71.1±0.2 66.5±0.5 74.2±0.2 76.1±0.7

52.5±0.5 44.3±0.2 46.0±0.4 52.7±0.6 54.9±0.3 55.3±0.6 59.5±0.4 58.0±0.3 57.8±0.3 54.2±0.3 59.9±0.3 61.2±0.3

75.7±0.4 65.5±0.3 68.2±0.5 74.2±0.5 76.8±0.3 78.1±0.7 79.6±0.2 81.2±0.5 80.9±0.4 77.1±0.7 80.4±0.4 82.2±0.4

76.4±0.4 70.9±0.2 72.2±0.6 75.3±0.5 78.6±0.2 79.4±0.6 82.7±0.3 83.2±0.1 83.0±0.2 80.6±1.2 81.5±0.2 84.9±0.4

66.6±0.4 58.2±0.2 60.1±0.5 66.0±0.5 68.6±0.2 69.4±0.2 73.4±0.3 73.5±0.1 73.2±0.1 69.6±0.2 74.0±0.3 76.1±0.2

Table 11: Per-domain accuracy on TerraIncognita. Method

L100 ↑

L38 ↑

L43 ↑

L46 ↑

Avg. ↑

ERM IRM MMD CORAL Fishr SAGM DynMoE GMoE OMoE MESSI-Ti-OT (Ours) MESSI-S-MMD (Ours) MESSI-S-OT (Ours)

50.8±1.8 42.0±1.0 47.2±1.2 46.1±1.0 52.1±1.6 52.7±3.2 54.8±0.8 56.6±0.7 55.5±0.8 53.6±1.9 57.1±0.5 60.2±0.7

42.5±0.7 31.2±0.7 35.6±0.8 39.6±3.2 42.5±2.1 43.5±4.0 36.4±0.7 36.2±0.9 34.2±1.4 23.1±1.2 36.9±0.5 38.3±0.5

57.9±0.6 46.5±0.6 50.6±0.7 53.3±0.7 56.7±0.8 58.7±0.7 49.0±0.5 58.0±0.6 55.3±0.6 52.2±0.8 54.8±0.4 56.7±0.3

37.6±1.2 35.1±0.8 35.4±1.0 36.2±0.8 38.3±1.5 40.3±1.4 37.8±0.9 42.0±0.8 39.0±0.7 41.2±0.3 41.2±0.3 42.1±0.2

47.2±0.8 38.7±0.8 42.2±0.9 43.8±0.9 47.4±1.6 48.8±1.0 44.5±0.6 48.2±0.5 46.0±0.5 42.5±1.1 47.5±0.4 49.3±0.2

Table 12: Per-domain accuracy on DomainNet. Method

Real ↑

Clipart ↑

Painting ↑

Sketch ↑

Infograph ↑ Quickdraw ↑

Avg. ↑

ERM IRM MMD CORAL Fishr SAGM DynMoE GMoE OMoE MESSI-Ti-OT (Ours) MESSI-S-MMD (Ours) MESSI-S-OT (Ours)

61.2±0.4 48.0±0.8 33.8±0.5 57.6±0.3 60.0±0.2 61.9±0.2 62.4±0.3 68.8±0.3 66.6±0.1 62.4±0.5 67.4±0.3 69.9±0.2

58.8±0.3 42.5±0.7 29.3±0.4 56.1±0.3 58.4±0.1 61.1±0.3 60.1±0.4 66.1±0.4 63.8±0.2 60.1±0.3 64.9±0.3 66.9±0.2

46.7±0.4 34.2±0.8 23.0±0.4 44.8±0.4 47.1±0.2 49.7±0.4 49.9±0.3 55.2±0.4 53.1±0.2 49.1±0.3 53.4±0.5 55.7±0.4

50.1±0.7 36.8±1.0 24.8±0.5 47.6±0.6 50.1±0.2 52.2±0.4 52.0±0.5 55.1±0.5 53.2±0.4 49.5±0.3 53.7±0.3 55.9±0.2

19.4±0.3 12.1±0.5 6.7±0.2 18.3±0.3 20.6±0.1 20.5±0.2 22.4±0.2 24.7±0.5 23.5±0.3 22.1±0.4 24.1±0.5 25.6±0.5

41.6±0.2 30.6±1.0 20.5±0.3 39.2±0.3 41.7±0.0 43.2±0.4 45.9±0.4 47.8±0.5 46.1±0.2 43.3±0.4 46.2±0.4 48.6±0.3

24

13.4±0.2 10.0±0.6 5.4±0.2 10.8±0.3 14.0±0.1 13.8±0.6 28.6±0.4 16.9±0.4 16.1±0.2 16.4±0.3 13.7±0.3 17.8±0.3

D.3

Fixed-Budget Domain-Growth results

For the fixed-budget domain-growth experiment, we vary Ks ∈ {3, 5, 7, 9, 11, 13, 15, 17} while keeping the total source training budget and the held-out target environment fixed. The target has ptest = 0.5, so color is independent of the binary label. Target accuracy therefore reflects whether a method learns shape-based prediction rather than relying on the color shortcut. We summarize robustness using the peak target accuracy and the final accuracy at Kmax = 17. We define Drop = Peak − Acc(Kmax ) and RelDrop = Drop/Peak. RelDrop prevents flat but low-performing methods from appearing robust simply because they never reach high accuracy. Table 13: Fixed-budget domain-growth summary on Rotated-Colored MNIST. Drop is Peak − Acc(Kmax ) with Kmax = 17, and RelDrop is Drop/Peak. Best results are shown in bold. Method CORAL (γ=1) CORAL (γ=10) CORAL (γ=100) CORAL (γ=250) CORAL (γ=500) MESSI (Ours)

Peak Acc. ↑

Acc(Kmax ) ↑

Drop ↓

RelDrop ↓

71.96 71.64 70.10 68.75 67.69 75.31

69.38 68.77 66.46 63.67 59.78 73.45

2.58 2.87 3.64 5.08 7.91 1.86

3.59% 4.01% 5.19% 7.39% 11.69% 2.47%

Table 14: Fixed-budget source-domain expansion on Rotated-Colored MNIST for Ks ∈ {3, 5, 7, 9}. The target environment and total source budget are fixed. Best results are shown in bold. Method CORAL (γ=1) CORAL (γ=10) CORAL (γ=100) CORAL (γ=250) CORAL (γ=500) MESSI (Ours)

Ks = 3 ↑

Ks = 5 ↑

Ks = 7 ↑

Ks = 9 ↑

50.84 50.79 50.79 50.79 50.79 50.79

50.79 50.79 50.82 50.86 50.79 57.14

68.45 68.80 68.85 65.65 59.95 74.00

71.96 71.64 70.10 68.75 64.89 75.31

Table 15: Continuation of Table 14 for Ks ∈ {11, 13, 15, 17}. Best results are shown in bold. Method CORAL (γ=1) CORAL (γ=10) CORAL (γ=100) CORAL (γ=250) CORAL (γ=500) MESSI (Ours)

Ks = 11 ↑

Ks = 13 ↑

Ks = 15 ↑

Ks = 17 ↑

70.44 70.91 69.10 67.30 67.69 74.80

71.00 71.00 68.59 67.84 66.17 74.10

70.20 68.92 66.35 63.99 62.81 74.99

69.38 68.77 66.46 63.67 59.78 73.45

Results. Table 13 shows that MESSI achieves the highest peak accuracy, the highest final accuracy at Kmax = 17, and the smallest relative drop. The full results in Tables 14 and 15 show the same pattern across individual source-domain budgets. At small Ks , all methods are close to chance because the source environments provide limited evidence for separating shape from color. Once source diversity becomes informative, MESSI improves sharply and remains stable as Ks increases, whereas stronger CORAL penalties become increasingly brittle. This supports the role of routing-conditioned subset alignment under increasing domain heterogeneity. D.4

Additional Results on WILDS-iWildCam

Table 16 reports additional results on WILDS-iWildCam [34, 5] under the protocol in Appendix C.4. This experiment is a controlled many-domain comparison among closely related MoE variants, not a leaderboard-oriented WILDS evaluation. Results. Table 16 shows that MESSI achieves the best macro F1 on all four evaluation splits. On the primary OOD test split, MESSI improves over OMoE by +5.20 F1 and over GMoE by +6.48 25

Table 16: Additional results on WILDS-iWildCam. All methods use the same DeiT-S/16 backbone, 224 × 224 preprocessing pipeline, and training budget. We report macro F1 on OOD and ID splits. OOD test is the primary evaluation on unseen camera locations. Best results are shown in bold. OOD F1 ↑ ID F1 ↑ Method Steps Val Test Val Test OMoE GMoE MESSI (Ours)

150K 150K 150K

29.14 29.06 35.49

27.96 26.68 33.16

48.57 48.84 52.86

47.26 47.28 51.14

F1. The same pattern appears on the ID test split, where MESSI improves over OMoE by +3.88 F1 and over GMoE by +3.86 F1. These gains are obtained under the same backbone, preprocessing pipeline, number of experts, and training budget, indicating that the improvement is not explained by generic sparse MoE capacity alone. Interpretation. The absolute scores are not intended to match heavily tuned WILDS leaderboard results, since we use lower-resolution inputs and a lightweight preprocessing pipeline. The purpose of this experiment is narrower: to test whether MESSI retains a relative advantage over closely related MoE baselines in a real many-domain setting. The consistent improvement on both OOD and ID splits supports the role of routing-conditioned subset alignment when domains are numerous, heterogeneous, and only partially share predictive structure.

E

Ablations Isolate Routing-Conditioned Subset Alignment

E.1

Ablation on Training Objectives

The following ablations test whether MESSI’s gains come from routing-induced subset-conditioned invariance rather than MoE capacity, generic feature alignment, or individual regularizers. Unless otherwise stated, all variants use the same backbone, number of experts, optimizer, augmentation pipeline, and model-selection criterion. We report OOD accuracy together with diagnostics tied to the intended role of each term: routing entropy for assignment confidence, expert-load standard deviation for usage imbalance, off-diagonal cosine similarity for expert redundancy, and class-conditional routing JS divergence for routing consistency. Diagnostics are computed on held-out source-validation splits; target-domain data are used only for final accuracy evaluation. Table 17: Objective leave-one-out ablation on PACS, averaged over the four leave-one-domain-out runs. Diagnostic columns are rescaled as indicated in the headers. Best results are shown in bold. Routing Ent. ↓ Load Std. ↓ Offdiag Cos. ↓ Routing JS ↓ Variant PACS Acc. ↑ (×10−2 ) (×10−1 ) (×10−2 ) (×10−4 ) MESSI (Ours) w/o Lssi w/o Lsp w/o Lbal w/o Ldiv w/o Lsp , Lbal , Ldiv

90.90±0.70 86.89±0.75 87.06±0.90 88.82±0.65 88.79±0.70 88.59±0.80

0.7±0.2 1.2±0.3 53.2±5.5 0.3±0.1 5.8±1.2 39.7±4.5

3.6±0.2 3.5±0.2 3.0±0.3 3.7±0.3 1.2±0.2 1.0±0.1

-0.2±0.3 -0.1±0.3 0.0±0.4 -0.1±0.3 35.6±3.0 38.9±3.5

0.1±0.0 2.0±0.6 1.0±0.4 0.1±0.1 33.0±4.0 30.0±3.5

P Metric definitions. Routing entropy is Ex [− m πm (x) log πm (x)]. Load standard deviation is the standard deviation of average expert usage Ex [πm (x)] across experts. Offdiag Cos. is the mean offdiagonal cosine similarity between expert representations. Routing JS is the average Jensen–Shannon divergence between domain-specific class routing prototypes and their class-wise template. Load standard deviation should be read together with routing entropy: a low load standard deviation can also arise from diffuse routing that uses all experts weakly. Results. Table 17 shows that the full objective achieves the highest PACS accuracy. Removing Lssi gives the largest drop, from 90.90% to 86.89%, indicating that routing-conditioned subset alignment is the main accuracy-critical component. Removing Lsp sharply increases routing entropy, 26

so expert assignments become diffuse. Removing Lbal has a smaller effect but slightly increases load imbalance. Removing Ldiv strongly increases off-diagonal cosine similarity and routing JS, showing that experts become redundant and less class-consistent without explicit specialization. Removing all routing and specialization regularizers produces diffuse routing and high expert redundancy, with lower accuracy than the full model. Overall, Lssi drives the OOD gain, while the auxiliary terms stabilize the sparse, balanced, and non-redundant expert decomposition. E.2

Expert Specialization

Figure 3 and Table 18 analyze expert specialization on OfficeHome. The goal of this diagnostic is to test whether MESSI learns distinct expert representations, rather than gaining accuracy only from increased MoE capacity. a) GMoE - cosine similarity

b) MESSI (Ours) - cosine similarity

Figure 3: Expert specialization on OfficeHome. (a) Pairwise cosine similarity between expert representations in GMoE. (b) Pairwise cosine similarity between expert representations in MESSI. Lower off-diagonal similarity indicates stronger expert specialization. Protocol. We follow the DomainBed leave-one-domain-out protocol on OfficeHome, using Art, Clipart, Product, and Real-World as the held-out domains in turn. All diagnostics are computed on the held-out target split, so the analysis reflects expert behavior on unseen domains. We use OfficeHome because its four visually diverse domains and fine-grained label space make expert redundancy easier to observe than on smaller benchmarks such as PACS. Metrics. Both models use M = 6 experts. For each input x, we extract the output of each expert, denoted by hm (x). For MESSI, these are the outputs of the MoE head. For GMoE, we use the per-expert outputs from the final MoE layer before routing aggregation. We quantify specialization with the pairwise cosine matrix C(x)m,n = cos(hm (x), hn (x)) ,

(64)

averaged over samples and then over held-out domains. Lower off-diagonal values indicate less redundancy across experts. Table 18 further reports three summary metrics. Mean Offdiag Cos. is the average off-diagonal entry of the cosine matrix. Effective rank is ! .X X erank = exp − pi log pi , pi = σi σj , (65) i

j

where {σi } are the singular values of the stacked expert-output matrix. Higher effective rank indicates that expert outputs span a more diverse subspace. Dead expert rate is the fraction of experts whose average routing mass falls below a small threshold. Results. Figure 3 shows a clear qualitative difference. GMoE exhibits uniformly high off-diagonal cosine similarity, indicating that its experts produce highly redundant representations. In contrast, 27

Table 18: Quantitative expert-specialization diagnostics on OfficeHome. Mean off-diagonal cosine measures expert redundancy. Effective rank measures the diversity of the expert-output subspace. Dead expert rate is the fraction of experts whose average routing mass falls below a small threshold. Best results are shown in bold. Method OMoE GMoE MESSI w/o Ldiv MESSI (Ours)

Mean Offdiag Cos. ↓

Effective Rank ↑

Dead Expert Rate ↓

0.35 ± 0.04 0.92 ± 0.03 0.50 ± 0.05 0.01 ± 0.01

3.95 ± 0.12 3.64 ± 0.15 3.75 ± 0.14 4.21 ± 0.11

0.45 ± 0.07 0.54 ± 0.06 0.40 ± 0.06 0.38 ± 0.05

MESSI yields a near-diagonal matrix, showing that different experts encode substantially different features. Table 18 confirms this trend quantitatively. MESSI reduces mean off-diagonal cosine similarity from 0.92 in GMoE to 0.01, increases effective rank from 3.64 to 4.21, and lowers the dead expert rate from 0.54 to 0.38. Removing Ldiv weakens this effect: the off-diagonal cosine rises to 0.50 and the effective rank drops to 3.75. OMoE improves over GMoE, but remains substantially less specialized than the full MESSI model. Takeaway. These results show that MESSI does not improve simply by adding MoE capacity. Instead, its routing and diversity objectives encourage experts to become both distinct and active, yielding a less redundant and more expressive decomposition. E.3

Routing Diagnostics

We further analyze whether the learned router induces meaningful expert specialization rather than merely increasing model capacity. For each input, we record the routing distribution over experts from the last MoE block and aggregate it within each source domain. For a domain k, we visualize three complementary routing statistics. First, P (e|d = k) measures the average expert usage of domain k. Second, Pk (e|y) denotes the class-conditioned expert usage estimated only from samples in domain k, i.e., P (e|y, d = k). Third, P (e|d = k, y) provides the same domain-conditioned class view in a format used to compare class-wise routing patterns across methods. These per-domain diagnostics reveal whether the router uses experts selectively, whether expert assignment is related to semantic classes, and whether the learned routing structure is dominated by domain-specific shortcuts. Figures 4–7 compare GMoE and MESSI on each PACS domain. Compared with GMoE, MESSI exhibits sharper and more class-structured routing patterns, indicating stronger expert selectivity and semantic specialization. At the same time, the class-conditioned routing visualizations allow us to check whether this specialization remains meaningful within each domain rather than collapsing into uniform expert usage or pure domain-level memorization.

28

P(e|d) A

e0

e1

e2

Expert

e3

e4

e5

P(e|d)

1.0 0.8 0.6 0.4 0.2 0.0

A

e0

e1

(a) GMoE: P (e|d = A) dog elephant giraffe guitar horse house person

e1

e2 e3 Expert

e4

e5

1.0 0.8 0.6 0.4 0.2 0.0

dog elephant giraffe guitar horse house person

e1

e2 e3 Expert

e4

e5

e0

e1

e2 e3 Expert

e4

e5

1.0 0.8 0.6 0.4 0.2 0.0

(d) MESSI: PA (e|y)

P(e|d,y)

e0

e3

P(e|y)

(c) GMoE: PA (e|y) A/dog A/elephant A/giraffe A/guitar A/horse A/house A/person

Expert

(b) MESSI: P (e|d = A)

P(e|y)

e0

e2

1.0 0.8 0.6 0.4 0.2 0.0

e4

e5

1.0 0.8 0.6 0.4 0.2 0.0

A/dog A/elephant A/giraffe A/guitar A/horse A/house A/person

(e) GMoE: P (e|d = A, y)

P(e|d,y)

e0

e1

e2 e3 Expert

e4

e5

1.0 0.8 0.6 0.4 0.2 0.0

(f) MESSI: P (e|d = A, y)

Figure 4: Routing diagnostics on the PACS Art domain. Rows correspond to routing aggregations and columns correspond to methods. P (e|d = A) measures domain-level expert usage. PA (e|y) denotes class-conditioned expert usage computed only on domain A. P (e|d = A, y) visualizes class-conditioned routing within the Art domain.

29

P(e|d) C

e0

e1

e2

Expert

e3

e4

e5

P(e|d)

1.0 0.8 0.6 0.4 0.2 0.0

C

e0

e1

(a) GMoE: P (e|d = C) dog elephant giraffe guitar horse house person

e1

e2 e3 Expert

e4

e5

1.0 0.8 0.6 0.4 0.2 0.0

dog elephant giraffe guitar horse house person

e1

e2 e3 Expert

e4

e4

e5

e0

e1

e2 e3 Expert

e4

e5

1.0 0.8 0.6 0.4 0.2 0.0

(d) MESSI: PC (e|y)

P(e|d,y)

e0

e3

P(e|y)

(c) GMoE: PC (e|y) C/dog C/elephant C/giraffe C/guitar C/horse C/house C/person

Expert

(b) MESSI: P (e|d = C)

P(e|y)

e0

e2

1.0 0.8 0.6 0.4 0.2 0.0

e5

1.0 0.8 0.6 0.4 0.2 0.0

C/dog C/elephant C/giraffe C/guitar C/horse C/house C/person

(e) GMoE: P (e|d = C, y)

P(e|d,y)

e0

e1

e2 e3 Expert

e4

e5

1.0 0.8 0.6 0.4 0.2 0.0

(f) MESSI: P (e|d = C, y)

Figure 5: Routing diagnostics on the PACS Cartoon domain. Rows correspond to routing aggregations and columns correspond to methods. P (e|d = C) measures domain-level expert usage. PC (e|y) denotes class-conditioned expert usage computed only on domain C. P (e|d = C, y) visualizes class-conditioned routing within the Cartoon domain.

30

P(e|d) P

e0

dog elephant giraffe guitar horse house person

P/dog P/elephant P/giraffe P/guitar P/horse P/house P/person

e1

e2

Expert

e3

e4

e5

P(e|y)

e0

e1

e2 e3 Expert

e4

e5

P(e|d,y)

e0

e1

e2 e3 Expert

e4

e5

P(e|d)

1.0 0.8 0.6 0.4 0.2 0.0

P

e0

1.0 0.8 0.6 0.4 0.2 0.0

dog elephant giraffe guitar horse house person

1.0 0.8 0.6 0.4 0.2 0.0

P/dog P/elephant P/giraffe P/guitar P/horse P/house P/person

e1

e2

Expert

e3

e4

e5

P(e|y)

e0

e1

e2 e3 Expert

e4

e5

P(e|d,y)

e0

e1

e2 e3 Expert

e4

e5

1.0 0.8 0.6 0.4 0.2 0.0

1.0 0.8 0.6 0.4 0.2 0.0

1.0 0.8 0.6 0.4 0.2 0.0

Figure 6: Routing diagnostics on the PACS Photo domain. In each row, the left panel shows GMoE and the right panel shows MESSI. From top to bottom, the rows show P (e|d = P ), PP (e|y), and P (e|d = P, y). These correspond to domain-level expert usage, class-conditioned expert usage computed on domain P , and class-conditioned routing within the Photo domain, respectively.

31

P(e|d) S

e0

dog elephant giraffe guitar horse house person

S/dog S/elephant S/giraffe S/guitar S/horse S/house S/person

e1

e2

Expert

e3

e4

e5

P(e|y)

e0

e1

e2 e3 Expert

e4

e5

P(e|d,y)

e0

e1

e2 e3 Expert

e4

e5

P(e|d)

1.0 0.8 0.6 0.4 0.2 0.0

S

e0

1.0 0.8 0.6 0.4 0.2 0.0

dog elephant giraffe guitar horse house person

1.0 0.8 0.6 0.4 0.2 0.0

S/dog S/elephant S/giraffe S/guitar S/horse S/house S/person

e1

e2

Expert

e3

e4

e5

P(e|y)

e0

e1

e2 e3 Expert

e4

e5

P(e|d,y)

e0

e1

e2 e3 Expert

e4

e5

1.0 0.8 0.6 0.4 0.2 0.0

1.0 0.8 0.6 0.4 0.2 0.0

1.0 0.8 0.6 0.4 0.2 0.0

Figure 7: Routing diagnostics on the PACS Sketch domain. In each row, the left panel shows GMoE and the right panel shows MESSI. From top to bottom, the rows show P (e|d = S), PS (e|y), and P (e|d = S, y). These correspond to domain-level expert usage, class-conditioned expert usage computed on domain S, and class-conditioned routing within the Sketch domain, respectively.

32

E.4

Routing-Aware Alignment Controls

The subset-conditioned invariance objective makes routing determine where alignment is applied. We test this mechanism by comparing MESSI with two controls. Global-MoE uses the same MoE architecture and auxiliary regularizers, but aligns all valid class-conditional source-domain pairs uniformly. Random-Subset-MoE uses the same sparsity level as MESSI, but selects aligned slots at random instead of using routing responsibility. All diagnostics are computed on held-out source-validation splits only; target-domain samples are not used. For expert m, source domain k, and class c, we compute the routing mass X 1 (m) ρk,c = πm (x). |Dk,c |

(66)

(x,y)∈Dk , y=c

This induces a pair responsibility for each class-conditional domain pair: (m)

(m)

(m)

aijc = σ(αρi,c )σ(αρj,c ),

i < j.

(67)

We call each tuple (m, i, j, c) a slot. A slot specifies which expert, domain pair, and class are considered for alignment. (m,A)

Responsibility similarity. For each method A, we flatten the responsibilities {aijc } over all common valid slots into a vector aA . We compare two methods using the Pearson correlation r(A, B) = corr(aA , aB ).

(68)

We also compare the top-q% selected slots. Let SA be the set of slots selected by method A. The Jaccard overlap is J(A, B) =

|SA ∩ SB | . |SA ∪ SB |

(69)

High r(A, B) and J(A, B) indicate that two methods select similar class-domain-expert slots. Alignment on selected slots. A selected-versus-non-selected gap can be misleading because highresponsibility slots may already be easier to align. We therefore evaluate all methods on the same set of MESSI-selected slots. Let SM be the top-q% slots selected by MESSI. For each method A, we compute   X 1 (m,A) (m,A) D Zi,c , Zj,c , (70) ∆A (SM ) = |SM | (m,i,j,c)∈SM

where D(·, ·) is the empirical energy distance between the two class-conditional expert-feature distributions. Lower values indicate better alignment on the slots selected by MESSI. Table 19: Routing-aware alignment diagnostics. r(A, B) is the Pearson correlation between pairresponsibility vectors, and J(A, B) is the Jaccard overlap between top-q% selected-slot sets. ∆A (SM ) is the feature discrepancy of method A evaluated on MESSI-selected slots; lower is better. Ctrl. ∆ averages Global-MoE and Random-Subset-MoE. Diagnostics use held-out source-validation splits only. Dataset

r(G,R) J(G,R) r(M,G) J(M,G) r(M,R) J(M,R)

PACS OfficeHome TerraIncognita DomainNet

+1.00 +1.00 +1.00 +1.00

1.00 0.93 0.95 0.95

−0.20 −0.20 −0.20 −0.20

−0.20 −0.20 −0.20 −0.20

0.04 0.11 0.13 0.11

0.04 0.11 0.14 0.10

Ctrl. ∆ ↓

MESSI ∆ ↓

0.235 ± 0.014 0.289 ± 0.004 0.250 ± 0.011 0.252 ± 0.004

0.096 ± 0.015 0.217 ± 0.006 0.179 ± 0.026 0.143 ± 0.032

Results. Table 19 shows two trends. First, Global-MoE and Random-Subset-MoE induce nearly identical responsibility patterns, with r(G, R) = 1.00 and high Jaccard overlap across datasets. In contrast, MESSI has weak or negative correlation with both controls and much lower selected-slot overlap. Thus, routing-aware alignment changes which expert-domain-class slots are selected for invariance, rather than simply sparsifying a global alignment objective. 33

Second, MESSI achieves lower discrepancy on the same MESSI-selected slots. The reduction is consistent across PACS, OfficeHome, TerraIncognita, and DomainNet. This indicates that the slots selected by MESSI are not only different from the controls, but are also better aligned after training. These diagnostics support the intended mechanism: routing determines where invariance is enforced, and the resulting selected slots receive stronger class-conditional alignment. E.5

Pairwise-to-Global Invariance Sweep

We study how the invariance weight affects alignment, domain information, and target performance. On a diagnostic PACS split, we sweep λinv ∈ {0, 10−3 , 10−2 , 10−1 , 1, 10} and report three quantities: pairwise class-conditional discrepancy, class-conditional domain-probe accuracy, and target-domain accuracy. Figure 8 and Table 20 show three trends. First, increasing λinv steadily reduces pairwise discrepancy, indicating stronger class-conditional alignment across source-domain pairs. Second, the classconditional domain-probe accuracy also decreases, approaching the source-domain chance level at λinv = 10. Thus, stronger invariance suppresses domain-predictive information in the learned representation. Third, target accuracy is not monotonic: it improves at moderate invariance strength, peaks at λinv = 10−2 , and then degrades sharply. At λinv = 10, target accuracy collapses even though pairwise discrepancy is nearly eliminated. This sweep complements the theoretical result that enforcing pairwise conditional invariance over all domain pairs induces global conditional invariance. Empirically, however, stronger invariance is not always better. Moderate alignment improves generalization, whereas overly strong alignment removes label-relevant structure together with domain-specific variation. This supports the need for controlled, routing-conditioned alignment rather than uniformly enforcing global invariance.

0.04

80

0.03 0.02 0.01 0.00

(c) Target accuracy

accuracy (%)

(b) Domain probe accuracy 90 accuracy (%)

discrepancy

(a) Pairwise discrepancy 0.05

70 60 50 40

0

0.

001

0

.01

0.1

lambda_inv

1

10

0

01

0.0

1

0.0

0.1 lambda_inv

1

10

90 80 70 60 50 40 30 20

0

01

0.0

1

0.0

0.1 lambda_inv

1

10

Figure 8: Pairwise-to-global invariance sweep on a PACS diagnostic split. We vary λinv and report pairwise class-conditional discrepancy, class-conditional domain-probe accuracy, and target-domain accuracy. Increasing λinv reduces pairwise discrepancy and domain predictability, but target accuracy peaks at moderate invariance strength and collapses under overly strong alignment. The dashed line in the domain-probe panel denotes the source-domain chance level.

Table 20: Numerical results for the pairwise-to-global invariance sweep on a PACS diagnostic split. Moderate invariance improves target accuracy, with the best result at λinv = 10−2 . In contrast, λinv = 10 nearly eliminates pairwise discrepancy and strongly suppresses domain predictability, but causes severe target-accuracy collapse. λinv

Pairwise Disc. ↓

Domain Probe Acc. ↓

Target Acc. ↑

0 10−3 10−2 10−1 1 10

0.05475 ± 0.0030 0.02620 ± 0.0022 0.02511 ± 0.0019 0.02188 ± 0.0017 0.02106 ± 0.0015 0.00107 ± 0.0004

93.42 ± 0.85 92.51 ± 0.95 89.51 ± 1.10 80.09 ± 1.60 75.51 ± 1.85 34.62 ± 2.20

89.24 ± 0.90 89.24 ± 0.85 91.69 ± 0.31 89.24 ± 1.05 84.11 ± 1.40 16.14 ± 2.80

34

F

Implementation Details

F.1

Architecture

Backbone. MESSI is evaluated with DeiT-Ti/16 and DeiT-S/16 encoders initialized from ImageNet1K pretrained checkpoints [58]. The DeiT-S/16 setting enables comparison with transformer-based and MoE-based baselines at the same encoder scale, whereas DeiT-Ti/16 tests whether subsetconditioned expert specialization remains effective with a lightweight encoder. In both settings, MESSI keeps the Transformer encoder unchanged and attaches a routing-based expert head to the final CLS representation. Thus, the backbone provides general visual features, while the expert head performs subset-conditioned specialization. Expert head. MESSI uses M = 6 experts. Let z = bϕ (x) denote the final CLS representation produced by the encoder. Each expert maps z to a 384-dimensional representation with a two-layer MLP: hm (z) = Wm,2 σ(Wm,1 z),

(71)

where σ is GELU. The router maps the same CLS representation to an expert distribution: π(x) = softmax(Wr z).

(72)

The final representation is the routed expert mixture h(x) =

M X

πm (x)hm (z),

(73)

m=1

and the classifier is applied to h(x). Router input and domain usage. The router receives only the image representation z. It does not receive the domain label. Domain labels are used only during training to compute source-domain alignment losses. At test time, MESSI requires neither domain labels nor target-domain data. F.2

Training Details

Backbone. All MESSI-S variants use a DeiT-Small/16, while MESSI-Ti uses DeiT-Ti/16 encoder initialized from an ImageNet-1K pretrained checkpoint [58]. The encoder has model dimension 384. MESSI keeps the Transformer encoder unchanged and attaches a 6-expert mixture head to the final CLS representation, with expert output dimension 384. Optimizer and schedule. We train with Adam [33] and use a constant learning rate. The learning rate, weight decay, and number of training iterations are dataset-specific. These values are tuned once on source-validation splits following the DomainBed protocol and then fixed for all final runs: Dataset PACS OfficeHome TerraIncognita DomainNet iWildCam (WILDS)

Learning rate −5

3 × 10 1 × 10−5 5 × 10−5 5 × 10−5 3 × 10−5

Weight decay −6

10 10−6 10−4 0 0

Iterations 5,000 10,000 10,000 15,000 150,000

Mini-batch construction. For PACS, OfficeHome, TerraIncognita, and DomainNet, we follow the leave-one-domain-out protocol. Each minibatch contains 32 images per source domain, giving total batch sizes of 96 for PACS, OfficeHome, and TerraIncognita, and 160 for DomainNet. For iWildCam, we treat camera-trap locations as domains. At each step, we sample Kloc = 4 source locations and draw Kimg = 8 images from each location, forming a minibatch of 32 images. These active locations define the domain pairs used by Lssi . Augmentation. Training images use the following augmentation pipeline: • RandomResizedCrop(224, scale=(0.7,1.0)), • RandomHorizontalFlip, 35

• ColorJitter(0.3, 0.3, 0.3, 0.3), • RandomGrayscale(p=0.1). Images are then normalized with ImageNet statistics, µ = (0.485, 0.456, 0.406) and σ = (0.229, 0.224, 0.225). Validation and test images are resized to 224 × 224 and normalized without augmentation. Model selection. For PACS, OfficeHome, TerraIncognita, and DomainNet, we use the DomainBed training-domain validation criterion. A stratified 20% split of each source domain is reserved for validation, and we select the checkpoint with the highest mean source-validation accuracy. Targetdomain data are not used for training, hyperparameter tuning, or model selection. For iWildCam, we follow the WILDS protocol and select checkpoints using the official OOD validation split; final performance is reported on the OOD test split. F.3

Hyperparameters

All five datasets share the same MESSI objective weights and routing configuration. We denote the sub-invariance loss as Lssi in the main method section; in the codebase the corresponding weight is called lambda_inv, i.e. λssi ≡ lambda_inv.

F.4

Hyperparameter

Value

λssi λsp λbal λdiv Routing softmax temperature α Number of experts M

0.01 0.02 0.02 0.02 4.0 6

Notation and Hyperparameter Summary

Table 21: Summary of notation used by MESSI. Domain indices refer to source domains in one leave-one-domain-out training run. Target-domain samples are not used when computing training losses.

F.5

Symbol

Meaning

Ks M c∈Y i, j ∈ {1, . . . , Ks } bϕ (x) hm (bϕ (x)) πm (x) (m) ρi,c (m) aijc (m) Zi,c CB PB (c) Lssi Lsp Lbal Ldiv

Number of source domains in the current DG split Number of experts Class index Source-domain indices Shared encoder representation Representation produced by expert m Router probability assigned to expert m Mean routing mass for expert m, domain i, class c Routing-induced weight for aligning domains i, j, class c, expert m Minibatch expert features from domain i, class c Classes present in the current source minibatch Source-domain pairs where class c appears in both domains Subset-conditioned invariance loss Routing sparsity loss Expert load-balancing loss Expert diversity loss

Baseline Implementations

We compare MESSI against three groups of baselines. The first group contains standard DomainBed baselines implemented with a ResNet-50 [26] backbone. ERM serves as the empirical risk minimization reference. IRM [1] represents invariant risk minimization. MMD [39] and CORAL [57] represent 36

Table 22: Summary of implementation hyperparameters used by MESSI. Hyperparameter Role M α ϵOT TOT λssi λsp λbal λdiv βbal

Number of experts (m) Sharpness of routing-induced pair weights aijc Entropic regularization coefficient in Sinkhorn OT Number of Sinkhorn iterations Weight of subset-conditioned invariance loss Weight of routing sparsity loss Weight of load-balancing loss Weight of expert diversity loss EMA coefficient for dataset-level routing load, if EMA balancing is used

distribution-alignment methods, where MMD matches kernel mean embeddings and CORAL matches second-order feature statistics. DANN [20] and CDANN [42] represent adversarial domain-alignment methods, with CDANN using class-conditional domain discrimination. MixStyle [73] represents feature-statistics-based data augmentation. Fish [56] represents gradient-matching-based DG, and SWAD [7] represents flat-minima-based model selection and weight averaging. Together, these methods form the conventional DomainBed reference set across ERM, invariant learning, feature alignment, adversarial alignment, augmentation, gradient matching, and flat-minima optimization. Unless otherwise stated, we use the DomainBed protocol and the reported or reproduced DomainBedcompatible implementations under the same train-domain validation model-selection rule. The second group contains recent optimization and architecture baselines. SAGM [61] represents sharpness-aware optimization for DG. GMoE [35] represents sparse mixture-of-experts architectures for domain generalization, and GMoE+SAGM combines the GMoE backbone with the SAGM optimizer. LFME [9] is included as a recent expert-based DG method based on multi-expert learning. These baselines test whether the gains of MESSI can be explained by stronger optimization, larger model capacity, or generic expert specialization. The third group contains controlled MoE variants used for mechanism analysis. These variants keep the MESSI backbone, number of experts, routing module, training budget, and model-selection criterion fixed, while changing only the training objective. This group includes classification-only MoE, globally aligned MoE, random-subset-aligned MoE, and leave-one-out variants that remove individual MESSI losses. These controls isolate the effect of routing-conditioned subset alignment from MoE capacity, global alignment, random sparse alignment, and auxiliary regularization. F.6

Model Size

Table 23: Trainable parameter counts in millions, excluding the final dataset-specific classifier. Enc. denotes shared encoder parameters. MoE/Head denotes additional expert, router, and method-specific head parameters. Values marked by ≈ are computed from the implemented architecture or from reported model-size differences. Method

Backbone

Enc.

MoE/Head

Total

ERM SAGM LFME GMoE GMoE + SAGM MESSI-S MESSI-Ti

ResNet-50 ResNet-50 ResNet-50 DeiT-S/16 DeiT-S/16 DeiT-S/16 DeiT-Ti/16

25.60 25.60 25.60 21.70 21.70 21.70 5.70

0 0 0 ≈ 12.10 ≈ 12.10 ≈ 7.09 ≈ 1.78

25.60 25.60 25.60 ≈ 33.80 ≈ 33.80 ≈ 28.79 ≈ 7.48

We exclude the final classifier because its size depends on the number of classes and is small relative to the encoder and expert modules. SAGM adds no trainable module, so it has the same parameter count as its base architecture; for the same reason, GMoE+SAGM has the same size as GMoE. LFME is listed with its ResNet-50 inference backbone for reference, although its training-time teacher/expert construction differs from a single MoE inference model. The GMoE MoE/Head count is estimated 37

from the reported difference between GMoE-S/16 and ViT-S/16. MESSI head counts are computed from our implementation. Model-size discussion. Table 23 shows that the gains of MESSI-S are not explained by a larger parameter budget. Under the DeiT-S/16 encoder scale, MESSI-S has approximately 28.79M trainable parameters, compared with approximately 33.80M for GMoE-S/16. Thus, MESSI-S uses about 15% fewer trainable parameters than GMoE-S/16 while still outperforming GMoE and OMoE in Table 1. Parameter count alone does not establish the mechanism; therefore, Appendix E.1 provides same-architecture ablations where the MoE capacity is fixed and only the objective terms are changed. MESSI-Ti gives a smaller operating point. With a DeiT-Ti/16 encoder, the model has approximately 7.48M trainable parameters, making it about 3.85× smaller than MESSI-S and 4.52× smaller than GMoE-S/16. This variant is intended for settings where a lower parameter footprint is preferred, while retaining the same routing-conditioned subset-alignment design.

G

Limitations

MESSI assumes that predictive structure is shared only across subsets of domains and that the router can recover useful routing patterns from data. When source domains are few, class coverage is limited or imbalanced, or routing fails to specialize, the resulting class-conditional routing-weighted discrepancies may become noisy and provide only weak alignment signals. Conversely, when domains already admit a strong globally invariant representation, the additional routing decomposition may introduce unnecessary complexity and reduce statistical efficiency relative to simpler global alignment methods. The method also introduces additional computational overhead due to pairwise class-conditional discrepancy estimation across experts and domains, particularly when using OT-based alignment objectives. Finally, the latent subsets induced by routing are not identifiable in general, and the learned routing assignments may not always correspond to semantically meaningful domain structure.

38