Diverse and Sparse Mixture-of-Experts for Causal Subgraph–Based Out-of-Distribution Graph Learning
Authors: Jerry Sun
Mohamed Abubakr Hassan
Yaoyu Zhang
Wanying Zhang
Chi-Guhn Lee
Venue: ICLR 2026 Poster (2026)
Keywords: graph neural network
out-of-distribution learning
mixture-of-experts
Primary area: learning on graphs and other geometries & topologies
TL;DR: We introduce a Mixture-of-Experts framework with diverse experts and sparse gating for graph OOD generalization, with a derived risk bound linking these properties to reduced OOD error.

Abstract: Current state-of-the-art methods for out-of-distribution (OOD) generalization lack the ability to effectively address datasets with heterogeneous causal subgraphs at the instance level. Existing approaches that attempt to handle such heterogeneity either rely on data augmentation, which risks altering label semantics, or impose causal assumptions whose validity in real-world datasets is uncertain. We introduce **DiSCO**, a novel *Mixture-of-Experts (MoE)* framework for **Di**versity- and **S**parsity-driven **C**ausal **O**OD graph learning, designed to model heterogeneous causal subgraphs without relying on restrictive assumptions. Our key idea is to address instance-level heterogeneity by enforcing semantic *diversity* among experts, each generating a distinct causal subgraph, while a learned gate assigns *sparse* weights that adaptively focus on the most relevant experts for each input. Our theoretical analysis shows that these two properties jointly reduce OOD error. In practice, our experts are scalable and do not require environment labels. Empirically, our framework achieves strong performance on the GOOD benchmark across both synthetic and real-world structural shifts.

Published as a conference paper at ICLR 2026
DIVERSE AND SPARSE MIXTURE-OF-EXPERIS FOR CAUSAL SUBGRAPH-BASEI OUT-OF-DISTRIBUTION
GRAPH LEARNING
Jerry Sun & Mohamed A. Hassan & Yaoyu Zhang & Wanying Zhang & Chi-Guhn Lee
Department of Mechanical and Industrial Engineering University of Toronto Toronto, ON M5S IAI, Canada
ABSTRACT
Current state-of-the-arr methods for diitr---ditriuttionn (OOD) generalization lack the ability to effectively address datasets with heterogeneous causal subgraphs at the instance level. Existing approaches that attempt to handle such heterogeneity either rely data augmentation, which risks altering label semantics, impose causal assumptions whose validity in real-world datasets is uncertain. We intro- duce DiSCO, a novel Mixture-of-Experts (MoE) framework for Diversity- and Sparsity-driven Causal OOD graph learning, designed to model heterogeneous causal subgraphs without relying on restrictive assumptions. Our key idea is address instance-level heterogeneity by enforcing semantic diversity among ex- perts, each generating a distinct causal subgraph, while learned gate assigns sparse weights that adaptively focus on the most relevant experts for each input. Our theoretical analysis shows that these two properties jointly reduce OOD error. In practice, our experts are scalable and do not require environment labels. Em- pirically, our framework achieves strong performance on the GOOD benchmark across both synthetic and real-world structural shifts.
INTRODUCTION
Out-of-distribution (OOD) generalization concerns learning models that remain reliable when the test distribution differs from training (Ye et al.) (021). Despite recent progress, it remains funda- mental challenge, especially for graphs (Fan et al. 2023). whose complexity admits shifts not only in node edge attributes but also in structural properties such size, sparsity, motif frequency (Wu et al., 2018). Such structural shifts can mislead models, as spurious topological correlations often dominate learning (Gui et al. 2022).
A dominant paradigm in graph OOD learning is causal subgraph identification: each graph is as- sumed to contain a subgraph G. responsible for the label, while the remaining structure G2 reflects spurious variation (Gui et al.] [003]] Sui et al.] 2025) An et al.) 2024 Yao et al., 2025) Chen et al.) 2022]. In principle, recovering G. should yield predictions robust distribution shifts. In practice, however, existing methods rely on restrictive causal assumptions (c.g., G3 is independent of the la- bel) that may fail in practice (Gui al., 2022, 2023). For example, sentiment analysis, stylistic markers such as word length can frequently track sentiment, meaning G, is not independent of the label.
A further challenge is instance heterogeneity: even within single environment and label class, dif- ferent samples may rely on fundamentally different causal subgraphs. For example, in molecular property prediction, multiple chemotypes can produce the same biological activity. meaning two "active" molecules may depend on entirely different causal subgraphs (Wu et al., (018). Methods that assume single invariant G, across environments or labels cannot capture such variability and therefore struggle under instance-level causal diversity (Sui et al.) 2025). Some approaches attempt to approximate this heterogeneity via data augmentation (Wu al. 2024: Sui et al. 2023). but perturbing graph structure cannot guarantee label correctness and may change the true causal sub- graph-especiall in motif-centric datasets such GOOD-Motif, where labels correspond directly

---

Published as a conference paper at ICLR 2026
to specific subgraphs (Gui et al. 2022 Wu et al. 2018). These limitations motivate methods that directly model causal diversity at the instance level, rather than approximating it through augmenta- tion.
We address this gap by introducing DiSCO, a offture-of-Experee (MoE) framework for Diversity- and Sparsity-driven Causal OOD graph learning. Unlike prior MoE approaches that rely prede- fined shift types or augmentation strategies, DiSCO employs experts as causal subgraph extractors, enabling them specialize in distinct causal mechanisms. A learned gating network adaptively selects the most relevant experts for each instance. Our theoretical analysis shows that (i) semantic diversity is necessary for meaningful specialization, and (ii) instance-leve sparsity naturally fol- lows from diversity via induced loss gaps. Together, these results provide principled justification for MoE in OOD graph learning. Importantly, our framework does not rely on environment labels or causal assumptions, making it assumption-light, Building on this foundation, our contributions are threefold:
1. Theory: We provide a principled justification for MoE in graph OOD learning by deriving formal risk bound, showing that semantic diversity among experts and instance-level sparsity during gating jointly reduce OOD error.
2. Implementation: We design a causal subgraph-based MoE framework that instantiates these principles through decorrelation regularizer for diversity and learnt gating mech- anism for sparsity, without requiring environment labels or strong causal assumptions.
3. Empirics: On the GOOD benchmark (Gui et al. 2022), DiSCO achieves strong perfor- mance across synthetic and real-world structural shifts. Ablations verify the necessity of diversity and sparsity.
2 BACKGROUND
Graph OOD Learning. In graph learning, each input is a graph x with an associated label y. In an OOD setting, models trained on data drawn from certain environments but expected to perform well when evaluated on new environments that were not seen during training. We focus on covariate shift, where the relationship between and y is stable but the distribution of graph inputs changes. This reflects realistic cases where the causal mechanism is preserved, yet nuisance factors such as graph size, sparsity, or motif frequency vary.
Causal subgraph paradigm. A common basis for graph OOD learning is the causal subgraph paradigm: each graph input contains a causal subgraph G. that determines the label y. while the complement G2 captures spurious variation. Successfully identifying G. would yield iiift-invariant
predictions, but since G. unobserved, existing methods rely on assumptions about the underlying
structural causal model (SCM) (Gui et al. [203 Suu et al. [025 Chen al. [022 An et al. 2024 Yao et al. 2025 Wu et al. 2022: Miao et al. 2022). Examples include assuming G. I y (Gui et al., 2023) that Ge invariant across environments or classes (Sui et al.) 2025p. Such assumptions
are often unrealistic: Such assumptions are often unrealistic: they overlook instance heterogeneity,
and they may fail when G, correlates with y (e.g., molecular scaffolds shared by active compounds
(Zhang et al. 2024)).
Instance heterogeneity. Instance heterogeneity refers the setting where the causal subgraph varies not only across environments but also across individual samples within the same dataset and label class. Formally, if each graph G contains a (possibly latent) causal subgraph G. sufficient to predict then the collection of such G.'s 's need not be identical across the dataset; this variability constitutes causal diversity. Such diversity arises naturally in real-world domains and is also present in the GOOD benchmark (Gui et al.) 2022). In GOOD-HIV, multiple chemotypes can yield the same biological activity, and GOOD-SST2. sentiment can be expressed through structurally dis- tinct syntactic patterns. As result, two instances from the same environment with the same label may rely on fundamentally different causal subgraphs, and methods that assume a single invariant Ge across environments or classes cannot account for such causal diversity. While some works at- tempt to approximate heterogeneity via data augmentation (Lu et al., 2024 Yao et al., 2024 Wu et al. 2024: Chen et al.) 2023 Sui et al.) 2023 Miao et al. 2022). perturbations do not guarantee label correctness and may alter the true causal subgraph- especially in centric datasets such
as GOOD-Motif (Gui et al. 2022). These limitations motivate the need for frameworks that directly
2

---

Published as a conference paper at ICLR 2026
model causal diversity at the instance level. We propose alternative perspective: model hetero- geneity directly through a MoE framework by allowing multiple experts to extract diverse causal subgraphs and using instance-specific gating to focus the most relevant ones.
Causal Assumptions. Many OOD graph methods rely on assumptions about the underlying SCM, particularly regarding how the causal subgraph Gz, spurious subgraph Gs, and label y interact. A
common assumption is that G. is conditionally independent of given Ge (Gui et al.) 2023; Sui
al. 2025 Chen al. 2022 An et al. 2024 Yao et al. 2025 Wu al. 2022; Miao al.) (0022)
However, this assumption does not hold under alternative SCMs such as FIIF (fully informative invariant features), where G, is entirely determined by G.. or PIIF (partially informative invariant features), where G3 is partly driven by Ge but still correlated with y. In practice, such independence conditions are rarely guaranteed. In molecular property prediction, for instance, structural scaffolds often correlate with biological activity, violating G. 1. In social or textual graphs, stylistic markers frequently track sentiment, creating direct G3 >> y dependencies. Even controlled benchmarks such the GOOD benchmark (Gui et al., 2022), these SCM-level assumptions are not guaranteed. Consequently, relying on a single causal assumption can lead to brittle behavior. In contrast, our method imposes SCM-level assumptions:
avoid invariance- instance-specifie expert selection, yielding robustness to diverse causal mechanisms.
losses and instead allow
MoE for GNNs. MoE architectures (Jacobs et al.) 1991| Shazeer al. 2017) consist of multiple experts combined through gating function, and in graph domains they have primarily been applied to improve scalability and efficiency (Wang al., 2023a Chen al. 2025 Huu et al. 2022). To the best of our knowledge, GraphMETRO is the only OOD-oriented MoE approach to date, assigning each expert to a predefined shift type and training them as augmentation strategies invariant those shifts (Wu et al. 2024). This approach naturally inherits the label validity risks mentioned above. In contrast, we employ MoE for causal subgraph identification, encouraging experts capture diverse causal hypotheses and directly addressing instance heterogeneity. Additional related work discussed in Appendix
METHOD
In this section, we derive an explicit risk bound that decomposes OOD eгror into coverage and selection terms, and we show how diversity and sparsity jointly serve to reduce it. All proofs are deferred to Appendix A
3.1 PRELIMINARIES AND NOTATION
We consider supervised learning on graphs. where each instance x (V,E, X) with node set V. edge set E, and node features X, together with an associated label y E y. Let X' denote the input space of graphs and D a distribution over (x, y) E X x V. We write Dx for the marginal on X. A predictor h maps to logits zn (ar) E RC, where C [ is the number of classes and c E (I....,C) C)} indexes a class. Predictions are evaluated with the cross-entropy loss (CE y) (Goodfellow et al. 2016), and the risk of h under distribution D is
Rh(h) E(z.o)=Dl(cr LCc (2) (x).
We define OOD generalization with respect to a family of environments M, where each m E M corresponds to a distribution D.,. over X' x Y?. Training occurs on a subset Mininn while evaluation is on unseen test environments Miest with Must \ MMgain * 0. We denote a particular unseen test environment by m', with associated distribution Dmn The objective is control the worst-case risk
ROOD (h) sup RD... (h).
mEMiue We focus on the case of covariate shift, where the marginal distribution over X' varies across envi- ronments while the conditional distribution P(y x) remains stable.
denote K experts, and let (xx) E A be gating distribution (sometimes we omit the dependence on x and write simply т). Each expert h, outputs logits 2x( e RC. The MoE aggregates predictions via
cesture-of-Experti Let [h. =
K
Zmiix (I) = Σ
(x) (x),,
(xx = arg max
Emix [c)
c5 ().....
3

---

Published as a conference paper at ICLR 2026
For expert hi, the crosspentropy loss is (xx, = (c and the MoE risk under distribution D is
(Mo E (< (fC (x); v)]-
We also define the auxiliary mixture-of-lossee quantity
K
E(x, y) T (=)))
which by Jensen's inequality (Jensen (906) upper-bounds the caxture-of-predictton loss.
3.2 DIVERSITY BETWEEN EXPERTS (COVERAGE)
For MoE model be effective, experts must represent distinct hypotheses; otherwise the gate has no meaningful choice and the model collapses. We therefore encourage semantic diversity: experts should emphasize different subgraphs in the input, providing coverage over latent causal mechanisms. As we will show, such diversity not only prevents model collapse but also induces non-trivial specialization, which in turn forces sparse gating-
Definition 3.1 (Semantic diversity). Let expert produce mask probability vector v; on input graph x (nodes/edges indexed by Ig). Standardize
[0,
(x)
(1)
Σ
(w),
(o. (0.0)/2=
=1= Σ (x)
-4
[Tal
0,
uEL.
DEL
We say the experts are with small 0 to avoid division by zero. Let Pij = tha (v/1)) pl0)). semantically diverse if
Σ i#j
(2) Een Dx [1
< Teorr for some Teore >> O.
K(K 1)
By enforcing low correlation between experts" masks, semantic diversity ensures experts attend to different parts the input graph. Under standard GNN encoders whose logits are Lipschitz in masked embeddings (Scarselli al. 2008 Joshi et al. 2023), low mask correlation encourages en- coding of distinct structural signals, discouraging collapse onto the same subgraph. This guarantees coverage of multiple causal hypotheses, giving the MoE leverage identify causal subgraphs under heterogeneity.
3.3 INSTANCE-LEVEL SPARSITY (SELECTION)
one expert predicts the causal mechanism for an input while others do not, the mixture is reli- able only the gate concentrates sufficient mass that expert: the selection effect. This ensures that once distinct causal mechanisms exist among the experts (via semantic diversity), the gate can isolate the correct one. To formalize this, we introduce metric, the loss gap. which measures the discrepancy between the best and next-best expert on an instance.
Definition 3.2 (Loss gap). For (x,, y). let it (x. y) = arg minte[K] li y) denote the best expert, where [K] = [I,...,K] The loss gap is
A(I. min (*,0)
(4(2,v) 0.
We now show that loss gaps quantitatively constrain the gate's allocation.
Proposition 3.3 (Loss gap implies sparsity). Let (li(x, k) be the per-expert losses and (x) E arg mine (((x y) be any minimizer. Define the mixture-of-losses E(z, v) Σ K (x(x) la y) and the loss gap
min (2.w) (La(2,,) v) lu (2, K 2,
A 9)
0,
K
4

---

Published as a conference paper at ICLR 2026
Then
(xr,)) 2 l (2. + (1- (1-TT;1 (x)) A A(x, v).
n..
Equivalently, for any A y) > 0.
Te (x) 2 1 ((x,0) v) - Ei (2.
A(x,))
This bound shows that to keep the mixture loss close to the best expert, the gate must assign sufficient weight to that expert. Furthermore, larger loss gaps make this requirement stronger. Such loss gaps are expected to arise as a consequence of diversity, which forces experts to attend to decorrelated subgraphs. We formalize this in the following assumption:
Assumption 3,4 (Diversity induces loss gaps). Let (xx, y) be in Definition 3.2
There exist
0, p [0, 1), and a measurable set S S X x y with D(S) 2 1 - such that
A(x,0) 2 V(z,y) E.
and on S the minimizer i (x.,) coincides with an expert aligned with the environment's causal mechanism. Semantic diversity (Definition) 3.1) promotes this condition by forcing experts to attend to decorrelated subgraphs.
This assumption formalizes the idea that, under semantic diversity, one causal expert should out- perform others by margin on most inputs. Intuitively, if one expert captures the causal subgraph, then by semantic diversity, the others focus on decorrelated (likely spurious) subgraphs, yielding non non-trivial loss gaps. This is validated empirically in Section 4 Hence, with sufficient diversity, loss gaps should arise on large fraction of inputs, and the MoE can match the best expert only by selecting it via sparse gating: diversity necessitates sparsity.
3.4 AN OOD RISK DECOMPOSITION COVERAGE AND SELECTION
We now show how these two properties jointly control OOD risk by decomposing risk into a cover- age term (controlled by diversity) and selection term (controlled by sparsity). First, we make the following assumption:
Assumption 3.5 (Mechanism coverage). Let (Dm : m E M) denote environments and h E arg mint, RD (h) the oracle predictor for environment Dmn For any OOD environment D' Dmn, there exists an expert it (m!) E [K] (depending only on the environment) such that
RD (h))(m)) S RDD. (Hime) + Ecov (m').
That is, for every unseen environment, at least one expert covers it by achieving risk within Ecov Ecov-(m)) of the oracle predictor. To minimize OOD risk, the MoE must then select this expert by allocating enough probability to it, yielding:
Theorem 3.6 (OOD risk: coverage + selection). Fix an OOD environment D' D For the environment-alignel expert it (mr), define
T m' (x, v)
max (L)(x,)) (4- v))- (m''
Then, under Assumptions 3.5 and 3.4
R (MoE)
< Ri (h
+
Ecov (m')
+ B(a.a))D-((1 The (x)) I' (x,
oracle risk
coverage diversity
selection penalty via sparsity
This decomposition expresses OOD risk three parts. The first is the oracle risk, the irreducible error in the target environment. The second is the coverage term, kept small by semantic diversity: only when experts are sufficiently diverse can the pool cover unseen test environments by ensuring at least one aligns with the causal mechanism. The third is the selection penalty, kept small by sparsity: loss gaps make the aligned expert identifiable, but only sparse gate reliably concentrate on it. Together, these conditions show that to jointly reduce OOD error, diversity needed to ensure a good expert exists, and sparsity is needed to ensure that it is selected.
5

---

Published as a conference paper at ICLR 2026
A
Expert
path
Deckward patht
Predictor
l
Subgraph Extractor
Diversty Regularization Loas
Gate
Expert
Lo
Gating Loss
Predictor
Task Lows
Figure 1: An illustration of the proposed MoE architecture. Each expert extracts distinet candidate causal subgraph and produces predictions, which are combined by the gating network into the final output.
3.5 IMPLEMENTATION
A detailed illustration of our implementation shown in Figure 1 Given an input graph æ (V, E.X), we first compute node embeddings with a shared GNN encoder. For each expert E [K], small expert-specific MLP takes the concatenated embeddings of edge endpoints and outputs
mask logit < for every edge e E E, which is later transformed into binary selection through the Gumbel-sigmoid straight-through estimator for differentiability (Maddison al.! 2017 Jang al. 2017). The resulting masked graph 2x1), representing the extracted causal subgraph by expert is then passed through expert-specific GNN and classifier head to produce logits 2, (z) Predictions are combined by lightweight MLP gate that consumes expert-derived statistics (c.g., confidence, entropy) and outputs weight vector (2x) A over experts. The final prediction is obtained as the weighted average of expert logits using и (x). The entire model is trained end-to-end with combination of task, regularization, diversity, and gating losses.
Task loss. For each expert i, the overall task loss is
(x y) Σ (T)
where (xx,y) is the per-sample cross-entropy for expert i. By weighting the loss with the gate probabilities, experts that contribute more to the final prediction receive stronger gradients, while those assigned little weight are suppressed, promoting specialization.
Regularization Loss. To control the size of extracted subgraphs, we regularize the average fraction of edges retained by each expert. For expert input let p; denote its observed keep-rate. We then penalize deviations from a target [0, 1] using
(xr
(( = (p, - 2) Teg
which discourages degenerate solutions where experts keep either too few or too many edges.
Diversity loss. To prevent experts from collapsing onto identical subgraphs, we compute
lai
K(K- Σ
max(0, lpb.) Teoer)
where Pi denotes the correlation between the masks of experts i and on input æ and the model is penalized when this correlation exceeds threshold Teor- This directly follows from the semantic diversity condition (Definition 3.1 and drives experts specialize distinct subgraphs, in turn ensuring coverage.
Gating loss. We train the gate with teacher-student objective. The teacher distribution is defined by normalizing the negative per-expert cross-entropy losses, giving higher weight to lower-loss ex- perts, while the student distribution is the gate output p (I:) This alignment teaches the gate

---

Published as a conference paper at ICLR 2026
which experts are competent which inputs. Since the task loss is gate- sate-weighted, non-selected experts are not heavily penalized, reinforeing specialization rather than foreing uniformity among experts. To further shape the gate, we add sparsity and balance regularizers, yielding the objective
Égasc KL(pllq) + Asparse Espars (p) Apal (p))
Here, Esarre penalizes high-entropy distributions to enforce instance-level sparsity, while boat en- courages even usage across the batch to avoid expert starvation, supporting coverage. These terms are complementary: sparsity sharpens routing for each input, while balance spreads usage globally. For stability, training begins with short warm-up phase of uniform routing all experts receive sufficient training signals before specialization. The gate can be fine-tuned after experts have been trained better align with the learned specializations.
Total loss. The final training objective is
C = lcE Avegfreg + Adiv + >
This loss function implements the components of the theoretical framework while incorporating additional regularizers that prevent degenerate solutions and expert starvation.
Theoretical computational complexity. Our MoE architecture is designed to introduce only mod- est computational overhead. As illustrated in Figure I all experts share a single subgraph extractor, and only lightweight per-expert output heads (MLPs) are replicated. Each expert consists of one GNN encoder followed by a small classifier (MIP), P), in contrast to existing OOD graph methods that require multiple full GNNs for classification (Gui et al. 2023 An et al., 2024).
Theoretically, if a GNN forward pass costs O(f(G) and each MLP forward pass costs (gg()), then the total complexity of our MoE is:
O(S(G) + K g() + K S(G) + K g() +
corresponding respectively to the shared extractor, per-expert heads, per-expert GNNs, per-expert classifiers, and the gating network. Since GNN computation dominates MLP computation, ((G
g(G) Goodffllow et al. 2016 ?), this simplifies to:
O(K f(G)).
indicating that the overhead grows linearly in the number of experts K and is dominated by the shared GNN encoder. In Appendix G we provide empirical analyses on the training and inference time of our method.
Assumption-ligh design without auxiliary invariance losses. Many prior OOD graph methods learn causal subgraphs through auxiliary objectives (e.g., adversarial discriminators in LECI (Gui et al. 2023), structural alignment in UIL (Sui et al. 2025)). While effective in single-model settings, these approaches are computationally expensive when replicated across experts and may be unstable (e.g., requiring finely tuned adversarial schedules that vary across experts). More critically, they rely on restrictive causal assumptions (e.g.. G. Y, invariant G.) that rarely hold heterogeneous real-world data (Zhang al. 2024).. By contrast, our design avoids such auxiliary losses entirely. As a result, it scales gracefully with the number of experts and requires neither environment labels nor restrictive causal assumptions. This key design choice makes method both highly scalable and assumption- light. We provide further discussion in Appendix D
4 EXPERIMENTS
We now evaluate our method empirically, guided by five research questions: RQ1: Does the method achieve strong performance on both synthetic and real-world datasets with structural shifts? RQ2: Does the method achieve strong performance across different causal assumptions? RQ3: Are the
key components of the framework necessary to obtain the reported improvements? RQ4: Does
enforcing semantic diversity induce larger loss gaps and promote expert specialization? RQ5: How sensitive is performance to hyperparameters?
Datasets. We evaluate on the GOOD benchmark (Gui et al. 2022), which provides training, OOD validation, and OOD test splits. Our study covers six datasets with structural shifts: HIV. Scaffold/Size (molecular), Length (sentiment). These span both synthetic and real-world domains under diverse structural
(synthetic motifs), Twigthr-Length (social), and SST2-
7

---

Published as a conference paper at ICLR 2026
Table 1: Results on graph classification datasets with structural shifts from the GOOD bench- mark. Values are classification accuracy (ROC-AUC for HIV) on OOD test sets, averaged over 5 runs with standard deviation in parentheses. Table sections correspond to domain generalization, data augmentation, and causal subgraph methods. We also report average performance and rank across datasets, Best results are in bold. indicates methods requiring environment labels.
HIV
Twitter t
Motif
SST2 7
Avrg Rank +
Method
Avrg
Scaffold
Size
Length
Basis
Size
Length
ERM
69.89 (1.95)
58.12 (2.5)
58.56 (1.0)
64.32 (104)
54.29 (544)
80.23 (0.8)
64,24
10.83
Coral
72.33 (2.1)
60.33 (3.5
57.33 (09)
65.39 (96)
52.39 (2.9)
79.23 (1.8)
64.50
10.33
IRM
71.11
60.67 (1.4)
57.79 (1.9)
62.64 (109)
54.14 (5.2)
80.37 (177)
64,45
10.33
VREx T
70,94 (3.0)
61.10 (3.0)
56.55 (0.7)
65.13 (5.0)
56.97 (6.3)
79.85 (1.6)
65.09
10.83
GDRO
67.13 (2.3)
56.91 (3.0)
56.73 (0.9)
62.63 (8.9)
52.01 (3.6)
81.33 (0.9)
62.79
12.50
DANN
67.69 (2.9)
62.05 (2.1)
56.09 (1.7)
52.65 (6:)
49.33 (5.4)
80.59 (0.9)
61.40
13.00
GM
66.06 (4.0)
66.24 (2.9)
56,97 (3.1)
67.33 (5.9)
61.43 (6:5)
81.96 (0.6)
66.67
7.83
AIA
71.23 (1.4)
62.33 (4.6)
57.13 (1.8)
74.18 (5.9)
56.07 (5.3)
80.91 (1.0)
66.98
7.67
GALA
74.51 (1.8)
64.89 (1.7)
60.79 (0.7)
79.11 (3:2)
72.13 (1.4)
82.42 (0.7)
72.31
2.67
LIRS
70.70 (2.3)
64.46 (2.9)
58.76 (1:4)
74.16 (3.0)
72.61 (6.9)
81.20 (0.7)
70.32
5.83
GSAT
70.76
61.76 (2.1)
57.13 (0.8
62.27 (0:)
54.12 (5:2)
80.62 (0:5)
64,44
10.67
CIGA
71.33 (1.1
63.09 (1.6)
58.01 (2.2)
38.01 (1:4)
55.69 (6.)
80.56 (1.7)
61.12
9.17
DIR
68.06 (5.5)
61.22 (0.8)
57.19 (0.9)
36.10 (2.5)
43.98 (3.1)
81.13 (0.7)
57.95
12.17
LECT
74.28 (1.7)
65.76 (1.4)
59.90 (0.2)
85.74 (3.0)
71.92 (1:4)
83.27 (0.3)
73,48
2.67
UIL
62.51 (1.7)
64.79 (0.8)
59.66 (0.9)
61.77 (4.8)
68.47 (3.1)
82.03 (0.4)
66.54
7.83
DISCO
71.55 (1.4)
66.98 (1.0)
61.13 (L.))
92.80 (1.4)
75.52 (2.9)
83.73 (1.4)
75.29
1.50
covariate shifts. We also evaluate on CFP-Motif (uui et al. three different causal assumptions: covariate, FIIF, and PIIF. Unless otherwise specified, results averaged over five seeds.
2023), which provides datasets with
Hyperparameters. All models adopt the Graph Isomorphism Network (GIN) Xu et al. 2019), the default GOOD backbone, with standard hyperparameters. Unless otherwise noted, we use eight experts within our MoE model. We tune mask keep-rate prior p. batch size, and learning rate over 10 trials on the OOD validation set.
Baselines. We compare against three groups of methods. (1) General domain generalization al-
gorithms: ERM. IRM (Arjovsky et al. 2019). Coral (Sun & Saenko 2016). V-REx (Krueger)
al. 202I). GroupDRO (GDRO for brevity) (Sagawa al. 2020), and DANN (Ganin et al.
2016 (2) Data augmentation methods that address instance heterogeneity: GraphMETRO (GM
for brevity) (Wu et al., 2024), AIA (Chen et al. (023),, and GALA (Sui al. 2023). (3) Causal subgraph algorithms: LIRS (Yao et al.) (225). GSAT (Miao et al.
2022). CIGA (Chen et al. 2022).
DIR (Wu et al. 2022), LECI (Gui et al. 2023), and UIL (Sui et al.
2025). All results are reproduced
using official repositories and hyperparameter settings.
4.1 RESULTS AND DISCUSSION
Does the method achieve strong performance on both synthetic and real-world datasets with structural shifts? On the synthetic datasets (HIV-Scaffold/Size and Motif-Basis/Size), our method achieves the best performance except on HIV-Scaffold. Notably, on Motif-Basis, a common sanity test for causal subgraph methods since labels are determined by the presence of specific motifs (Gui et al. 2022 2023), our method achieves 92.8% accuracy which is comparable to oracle- level performance (Gui al. 2023). and an 8.2% relative improvement over the next best method. On HiV-Scaffold, the severe class imbalance (over 95% majority class) makes the task especially challenging, and our method attains 71.55 ROC-AUC, ranking fourth overall. On the real-world datasets (Twitter and SST2), our method outperforms all baselines, demonstrating robustness under noisy, real-world distribution shifts. Across all datasets, our method achieves the highest average score (75.29) and lowest average rank (1.5). The closest competitors LECI (73.48 average, 2.67 rank) and GALA (72.31 average, 2.67 rank), confirming substantial improvements over existing baselines.
RQI:
8

---

Published as a conference paper at ICLR 2026
Table 2: Results on different causal assumption datasets from CFP-Motif. Values are are classifi- cation accuracy with standard deviation in parentheses. Baseline results are from Gui et al. (2033) Best results are in bold.
Covariate |
FIIF L
PII +
ERM
57.56 (9.59)
37.22 (3.70)
62.45 (9.21)
IRM
58.11 (5.14)
44.33 (1.52)
68.34 (10.40)
VREx
48.78 (7.81)
34.78 (1.34)
63.33 (6.35)
Coral
57.11 (8.35)
42.68 (7.09)
60.33 (8.85)
DANN
49.45 (8.05)
43.22 (6.64)
62.56 (10.39)
DIR
44.67 (0.00)
42,00 (6.77)
47.22 (8.79)
GSAT
68.22 (7.23)
51.56 (6.59)
61.22 (8.80)
CIGA
56.78 (2.99)
39.11 (7.70)
45.67 (7.52)
LECI
83.20 (5.89)
77.73 (3.85)
69.40 (7.54)
DISCO
90.83 (1.73)
84.17 (3.88)
77.19 (6.04)
Table Performance of the MoE framework on the GOOD benchmark under ablations of each loss term. Standard deviations are reported in parentheses.
HIV f
Twitter t Length
Motif f
SST2 f
Dropping
Scaffold
Size
Basis
Size
Length
Edic
65.95 (2.5)
65.23 (0.8)
60.10 (0.9)
91.13 (1.))
70.07 (28)
82.20 (1.3)
fgue
68.56 (0.9)
61.91 (1.2)
59.92 (0.)
89,60 (27)
74.36 (2.0)
81.97 (1.)
Eag
68.55 (3.9)
64.79 (2.2)
60.60
67.48
73.75 (2.5)
83.46 (1.5)
Standard (full loss)
71.55 (1.4)
66.98 (1.0)
61.13 (1.1)
92.80 (1.4)
75.52 (2.9)
83.73 (1.4)
RQ2: Does the method achieve strong performance across different causal assumptions? On the CFP-Motif dataset (Table 2). our method achieves the best performance under all three causal assumptions, with LECI as the strongest baseline. We attribute this improvement avoiding aux- iliary invariance losses that implicitly enforce a fixed causal assumption. By not committing to particular SCM, our method adapts more effectively the different causal assumptions present in CFP-Motif.
RQ3: Are the key components of the framework necessary for the reported improvements? Table 3 presents an ablation study on the GOOD benchmark, evaluating the contribution of each major component: semantic diversity, instance-level sparsity, and regularization. We systemati- cally remove the corresponding loss terms and observe substantial performance degradations across most datasets, with particularly large drops on Motif Baif-B HIV-Scaffold, and HIV-Size. These re- sults indicate that all three components essential for achieving the full performance gains of our framework.
RQ4: Does enforcing semantic diversity induce larger loss gaps and promote expert specializa- tion? To evaluate Assumption 3.4 (semantic diversity induces loss gaps), we compute the total loss gap within each batch and then average over the entire test set. We then compare models trained with and without the diversity objective. Jable that semantic diversity consistently increases the average per-batch loss gap across datasets. On Twitter, the gap rises from 0.13 to 0.19 46% increase); on SST2, from 0.07 to 0.22 (over 200% increase); and on Motif Basif, from 0.076 to 0.12 (a 58% increase). These results empirically validate our assumption that semantic diversity among experts induces larger loss gaps. We provide additional evidence in Appendix F.I which demon-
strates diverse expert specialization across causal subgraphs, and in Appendix F3 which presents visualizations of the causal subgraphs extracted by each expert.
RQ5: How sensitive performance to hyperparameters? We first evaluate the effect the edge keep-rate prior p Twitter, SST2, and Motif-Bass -asis by perturbing it around the tuned value (shown in brackets). Table shows that accuracy varies by at most 1-2% under shifts of +0.1 or #0.2, confirming robustness moderate deviations. For example, Twitter peaks at p 0.55 (61.13%), SST2 at p = 0.20 (83.73%) and Motif-Basis at p = 0.55 (92.800 80%). with nearby settings yielding comparable results. These results demonstrate that our framework is not overly sensitive the edge

---

Published as a conference paper at ICLR 2026
Table Sensitivity analysis of the edge keep- rate prior p on Twitter, SST2, and Motif-Basis. The value in brackets indicates the tuned p se- lected via validation for each dataset, with per- formance shown at this setting and at perturbed values.
Table 4: Total loss gaps per-batch under abla- tions of diversity on Twitter, SST2, and Motif- Basis. Standard deviations are reported in parentheses.
Dataset
Diversity
Loss Gap
P Twitter (0.55)
SST2 f (0.2)
Motif-Base Botif-Basis t (0.55)
w/
0.19 (0.011)
Twitter
w/o
0.13 (0.004)
w/
0.22 (0.12)
-0.2
58.75
90.13
SST2
w/o
0.07 (0.19)
-0.1
60.33
82.06
89.76
w/
0.12 (0.035)
Tuned
61.13
83.73
92.80
Motif- Basis
w/o
0.076 (0.024)
+0.1
59.51
83.15
92.03
+0.2
60.64
83.46
90.54
Table 6: Performance of the MoE framework with 4, and experts. reported in parentheses.
Standard deviations are
HIV 1
Twitter f Length
Motif 1
SST2 f Length
Experts
Avrg
Scaffold
Size
Basis
Size
1
67. 13 (22)
63.48 (0.9
58.68 (1.7)
89.3 (2.1)
65.31 (12)
81.11 (L.)
70.84
4
70.65
65.35 (1.3)
60.48 (0.7)
91.79 (1.5)
76.4 (21)
83.41
74.68
8
71,55 (1.4)
66.98 (1.0)
61.13 (1.1)
92.8 (1.4)
75.52 (29)
83.73 (144
75.29
keep-rate prior p. Furthermore, hyperparameters tuned over modest budget of only 10 trials. We then analyze the impact of the number of experts in Ta which compares 1, 4, and experts. Moving from 1 to 4 experts yields substantial gains across all datasets, demonstrating the importance of expert diversity and sparse gating. Increasing further to 8 experts provides smaller but generally positive improvements (except on Motif-Sizor Importantly, with only experts, the MoE framework achieves an average accuracy of 74.68%, outperforming all baselines in Table []] suggesting that our method is not highly sensitive the exact number of experts once diversity present.
5 CONCLUSION
In this work, we introduce a causal subgraph- MoE framework that explicitly addresses instance-level heterogeneity, enabling different experts capture distinct causal explanations within the same class. Our framework demonstrates that diversity among experts provides coverage of het- erogeneous causal mechanisms, while sparsity in the gating step enables effective selection, together reducing OOD eиror. We operationalize these principles in scalable, assumption-lipn architecture that requires neither environment labels nor restrictive causal assumptions. Empirically, the method achieves strong performance on the GOOD benchmark across both synthetic and real-world shifts, with ablations and visualizations confirming that experts specialize in distinct causal mechanisms. Looking ahead, broadening causal perspectives on OOD graph learning, through richer causal mech- anisms, more flexible expert designs, and closer pheory-practice integration, remains an important direction for building robust and generalizable graph learning systems. We hope this work estab- lishes MoE as a strong foundation for causal-base based OOD graph learning.
REPRODUCIBILITY STATEMENT
All mathematical proofs are provided in Appendix Implementation details are provided in Ap- pendix C
The source code and reproduction instructions are available here
REFERENCES
Weizhi An, Wenliang Zhong, Feng Jiang, Hehuan Ma, and Junzhou Huang. Causal subgraphs and information bottlenecks: Redefining ood robustness in graph neural networks. In European
10

---

Published as a conference paper at ICLR 2026
Conference Computer Vision, pp. 473-4779999 Springer, 2024.
Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint www.iv:1907002557 2019.
Xuanze Chen, Jiajun Zhou, Shanqing Yu, and Qi Xuan. Mixture of experts meets decoupled message passing: Towards general and adaptive node classification. In Companion Proceedings of the ACM on Web Conference 2025, pp. 907-910, 2025.
Yongqiang Chen, Yonggang Zhang, Yatao Bian, Han Yang, MA Kaili, Binghui Xie, Tongliang Liu, Bo Han, and James Cheng. Learning causally invariant representations for titto--distribution generalization on graphs. Advances Neural Information Processing Systems, 35:22131-22144 2022.
Yongqiang Chen, Yatao Bian, Kaiwen Zhou, Binghui Xie, Bo Han, and James Cheng. Does invariant graph learning via environment augmentation learn invariance? Advances in Neural Information Processing Systems, 36:71448800 36:71486-71519, 2023.
Shaohua Fan, Xiao Wang, Chuan Shi, Peng Cui, and Bai Wang, Generalizing graph neural networks on out-of-distiibttion graphs. IEEE transactions on pattern analysis and machine intelligence, 46(1):322-337, 2023.
Yaroslav Ganin, Bvgeniya Ustinova, Hana Ajakan, Paseal 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.
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT press, 2016.
Shurui Gui, Xiner Li, Limei Wang, and Shuiwang Ji. Good: A graph out-of-distribution benchmark. Advances in Neural Information Processing Systems, 35:2059-2000 2073, 2022.
Shurui Gui, Meng Liu, Xiner Li, Youzhi Luo, and Shuiwang Ji. Joint learning of label and envi- ronment causal independence for graph distr-iiittiiitionn generalization. Advances in Neural Information Processing Systems, 36:3945-3977 2023.
Fenyu Hu, Liping Wang, Qiang Liu, Shu Wu, Liang Wang, and Tieniu Tan. Graphdive: Graph classification by mixture of diverse experts. In Lud De Raedt (ed.), Proceedings of the Thirty-First International Joint Conference
Artificial Intelligence, IJCAI-22, pp. 2080-2086. International
Joint Conferences on Artificial Intelligence Organization, 2022. doi: 00.24933iicui222222/899 2022/289
URL httpp:://oi...oot org. 21..1...9933311 21t.9/////1...2222222...2..
Main Track.
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. Advances in neural information processing systems, 33:22118 11-22188221133 2020.
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.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparametrization with gumble-softmax. In International Conference on Learning Representations (ICLR 20/7). OpenReview. net, 2017.
Johan Ludwig William Valdemar Jensen. Sur les fonctions convexes et les inégalités entre les valeurs moyennes, Acta mathematica. 30(1):175-193, 1906.
Chaitanya K Joshi, Cristian Bodnar, Simon V Mathis, Taco Cohen, and Pietro Lio. On the expressive power of geometric graph neural networks. In International conference on machine learning, pp. 15330-15355. PMLR, 2023.
David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distriiatioo generalization via risk extrapo- lation (rex). In International conference machine learning, pp. 5815-5826. PMLR, 2021.
Bin Lu, Ze Zhao, Xiaoying Gan, Shiyu Liang, Luoyi Fu, Xinbing Wang, and Chenghu Zhou. Graph out-off-dstrrbutt generalization with controllable data augmentation. JEEE Transactions on Knowledge and Data Engineering, 36(11))63776663333 2024.

---

Published as a conference paper at ICLR 2026
C Maddison, A Mnih, and Y Teh. The concrete distribution: A continuous relaxation of discrete random variables. In Proceedings of the international conference on learning Representations. International Conference on Learning Representations, 2017.
Siqi Miao, Mia Liu, and Pan Li. Interpretable and generalizable graph learning via stochastic at- tention mechanism. In International conference on machine learning, pp. 15524-15543. PMLR, 2022.
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang, Distributionally robust neural networks. In International Conference on Learning Representations, 2020.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE transactions on neural networks, 20(1):61-80) 2008
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations, 2017.
Yongduo Sui, Qitian Wu, Jiancan Wu, Qing Cui, Longfei Li, Jun Zhou, Xiang Wang, and Xiangnan He. Unleashing the power of graph data augmentation on covariate distribution shift, Advances in Neural Information Processing Systems, 36:18109 1111--9991111,.11111111111111111 2023.
Yongduo Sui, Jie Sun, Shuyao Wang, Zemin Liu, Qing Cui, Longfei Li, and Xiang Wang. unified invariant learning framework for graph classification. In Proceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining V. I, pp. 1301-1312, 2025.
Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In European conference on computer vision, pp. 443-450. Springer, 2016.
Haotao Wang. Ziyu Jiang, Yuning You, Yan Han, Gaowen Liu, Jayanth Srinivasa. Ramana Kom- pella, Zhangyang Wang, et al. Graph mixture of experts: Learning on large-scale graphs with ex- plicit diversity modeling. Advances Neural Information Processing Systems, 36:50825-50835 2023a.
Zitai Wang, Qianqian Xu, Zhiyong Yang, Yuan He, Xiaochun Cao, and Qingming Huang. A unified generalization analysis of re-weighting and logit-adjustment for imbalanced learning. Advances in Neural Information Processing Systems, 36:48417-48430, 2023b.
Shirley Wu, Kaidi Cao, Bruno Ribeiro, James Zou, and Jure Leskovec. Graphmetro: Mitigating complex graph distribution shifts via mixture of aligned experts. Advances in Neural Information Processing Systems, 37:9358-9338 2024.
Yingxin Wu, Xiang Wang, An Zhang, Xiangnan He, and Tat-Seng Chua. Discovering invariant rationales for graph neural networks. In International Conference on Learning Representations, 2022.
Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learn- ing. Chemical science, 9(2):513 2018.
Keyulu Xu, Weihua Hu, Jure Leskovee, and Stefanic Jegelka. How powerful are graph neural networks? International Conference on Learning Representations, 2019.
Huaxiu Yao, Xinyu Yang, Xinyi Pan, Shengchao Liu, Pang Wei Koh, and Chelsea Finn. Improv- ing domain generalization with domain relations. In The Twelfth International Conference on Learning Representations, 2024.
Tianjun Yao, Yongqiang Chen, Kai Hu, Tongliang Liu, Kun Zhang, and Zhiqiang Shen. Learn- ing graph invariance by harnessing spuriosity. In The Thirteenth International Conference on Learning Representations, 2025.
Haotian Ye, Chuanlong Xie, Tianle Cai, Ruichen Li, Zhenguo Li, and Liwei Wang. Towards theoretical framework of outto--distribution generalization. generaaitoonnnnenraaiiaaiiooo Advances in Neural Information Pro- cessing Systems, 144-23519-235333 2021.
12

---

Published as a conference paper at ICLR 2026
Kexin Zhang, Shuhan Liu, Song Wang, Weili Shi, Chen Chen, Pan Li, Sheng Li, Jundong Li, and Kaize Ding. A survey of deep graph learning under distribution shifts: from graph out- of-distribution generalization adaptation. CoRR, 2024.
A APPENDIX: PROOFS
This appendix collects the proof's of all formal results stated in the main body, together with some natural extensions. Our goal is not to introduce new theory but to provide full technical details and clarify intermediate steps that were omitted for brevity in the main text.
A.1 SPARSITY REQUIREMENT IN AN EXTREME CASE
We begin with a simple but instructive lemma that is not included in the main body, It captures an extreme case of expert selection, where exactly one expert achieves a positive margin and all others non-positive (that is, only expert is correct). While not needed for the main theoretical results, we include here because highlights in the clearest possible terms why sparse gating is necessary: in such scenarios, correct predictions are only guaranteed if the gating mechanism concentrates sufficient mass on the correct expert.
Lemma A.1 (Sparsity requirement in an extreme case). Let per-expert margins be
mi m = (x)
max (x))c) oay
and let the mixture logits be
K
Emix (x) Σ
π(x) E A
If there exists an expert i with m.- (x) > 0 and my(a) S 0 for all j # it, then the mixture margin
Maxi (x) == 2mix (x) (x) max Xmix (x)
czV
strictly positive whenever
m my m2(2)
T4+ (x) > a,
a my mx (x) - minj my mz(x)
E (0,1).
Proof. We have
mmix (2) Σ T((x) (z)(x)
max Σ (x) (+))))(()) (x)(cc]
of
2 Σ (((x) 2(())al - Σ ne (w) max 2 [c) = (x)m,(x))
using max. Ex Oic Ee max, dic. With my (x) and my (x) for j + it,
Σ (0) 2 (x) (x) (1 म.. (x)) min my (x)
jf
Thus Mmix(x) (x) > whenever
Te (x) m2 (a:) (1 nz, (x)) min my(x) >
which rearranges Te- (x) as stated.
The previous lemma considered the binary-style case where only one expert correct and all others strictly incorrect. We now extend the analysis the multiclass setting by allowing non-causal experts to have bounded negative margins. This produces a similar threshold condition on the gating weight, ensuring that the correct expert dominates when sufficiently favored by the gate.
13

---

Published as a conference paper at ICLR 2026
Lemma A.2 (Multiclass sparse threshold under bounded negative margins), Let the one-vs-rest margin of expert i at input x be
m2(a) 2x())pl max (z)]cc] czy
Suppose there exists a (causal) expert it and constants m 0 and 7 2 0 such that
m4- (z) 2 m₃(x)) 2 -Y for all j + i²
Then the mixture margin satisfies
K
mmix (x)
Σ
(1(2)m(()) 2 (r-(x) (1 Tit (2)) 7.
and in particular Muxix (x) > 0 whenever
Tt. (x) >
?
+
Proof. By the max-sum inequality max, E; az S >; max, aso.
Mmix (x) Zmix (z)
max mmix (x)lc
cg
Σ (T)) - max of
(
(() m)()..
Using the margin bounds m.- (x) 2 m and my(x) 2 -y for 1 #
Σ [=( 2 ne. (x) m Σ [イ{r) (-7) m (2) m - (1 (1- (x)) g.
Hence
Mixx (x) 2 ne. (z) m (1- -1-55((6 (z)) 7,
which is strictly positive precisely when π.. (x) 7/ 7/ 7)
A.2
LOSS GAP IMPLIES SPARSITY
The margin- based results above illustrate the role of sparsity terms of logits and decision bound- aries. We now provide complementary perspective using losses directly. This proposition shows that a positive loss gap between the best expert and all others induces lower bound on the gating weight assigned to the best expert. This connects the concept of expert specialization to loss-based analysis.
Proposition A.3 (Loss gap implies sparsity). Let (Li(z, (L((,,))) be the per-expert losses and (x, y) E arg min (x y) be any minimizer. Define the mixture-of-losses f(x, y) = Σ (x) Ea(a,, u) and the loss gap
A(x,8) = mmin (2.)) ((,(,,))=
K> 2, la (2),
0.
K = 1.
Then
((x,y) 2 ly (z, + (1 - n.. (x)) A(x,
Equivalently, for any A(z, y) 0.
Te. (x) 2 1 [(. v) Lx (x.
A(z,y)
Proof. Since i" is minimizer, for all it we have (z(x,,)) Ez (x, v) 2 A(x, v). Therefore,
K
[(x,y) - l,- (x. Σ (x) (L, v) - (x, = Σ (z(x) (4, (2, ls (x. 2
jtt
Ex v) (1- m.- (x)) A (x,,))
which yields the stated inequality. Rearranging gives the equivalent lower bound म.. (1)).
14

---

Published as a conference paper at ICLR 2026
A.3 OOD RISK: COVERAGE SELECTION
Finally, we return to the main theorem on OOD risk. This result decomposes the risk of the MoE into three interpretable terms: the oracle risk (if the aligned expert were always selected), coverage penalty (arising when no expert perfectly matches the test distribution), and selection penalty (arising when the gating function fails to concentrate on the aligned expert). The proof makes explicit the role of Jensen's ss inequality and shows how the sparsity results above fit into the broader risk bound.
Theorem A.4 (OOD risk: coverage + selection). Fix an OOD environment D' = D For the enviromment-aligned expert it (m'', define
m (I. v) = max
(B((,)) - (mm))
jftr
Then, under Assumptions 3.5 and 3.4
R (MoE) S R (h
+
Ecov (m²)
+ B(z.))) Ti (-m (2)) T (x. v)
oracle risk
coverage via diversity
selection penalty sporsity
Proof. By convexity of cross-entropy in the logits,
Ro Ro (MoE) =
< ED(((e,V)),
(x, = Σ ( ((
Fix (x. y) and abbreviate = 1n (m'). Decompose
ez By definition of m' (I, y). each difference satisfies (,(x, y) ea- (z. v) S T m' (I. y). hence
E(r,y) n. (x)
Σ jz1
(x, Ex (e,(x)((((1,))) la- (2-(2))) 8)).
E(z, v) & l. (x,z) +
[m)(t))Fnr(e (, (z, + - n (x)) (x, v).
372
Taking expectations under D' yields
RD (MoE) < EDi li- (x, w)] + ED.[((1 Tiz (I)) I'm (2,
RD (h.r) ) + EDi(( Ttp (x)) I' (I, u)]
Finally, by Assumption 3.5 RO (bx.) S RD (hina)) ) + Ecor (m²), which proves the claim.
B DETAILED RELATED WORKS
B.1 CAUSAL SUBGRAPII-BAN OOD LEARNING.
LECI enforces label-environment independence by learning edge masks such that the label is inde- pendent of the environment and the environment independent of the causal subgraph (Gui al.) 2023). jointly enforces semantic and structural invariance, aligning graphs across environments using graphon distances (Sui et 2025). LIRS instead learns spurious features first and removes them from ERM-learned features, thereby capturing a broader set of invariant subgraphs (Yao et al.) 2025). DIR discovers label-causal subgraph rationales by intervening on the training distribution and selecting features invariant across the induced environments, filtering out spurious shortcuts (Wu et al.) 2022). CIGA learns causally invariant graph representations by extracting subgraphs that maximally preserve label-relevant intra-class information (Chen al.) 2022). GSAT applies an information boriiionn-bottenekkkkk driven stochastic attention to mask task-irrelevant nodes/edges, yielding faithful rationales (Miao et al. 2022). CSIB learns label-causal subgraphs under causal model while jointly optimizing invariant risk with a graph information bottleneck to balance invariance
versus compression for FUFF/IIF (An et al. 2024).
15

---

Published as a conference paper at ICLR 2026
B.2 INSTANCE HETEROGENEITY.
GSAT enforces adversarial consistency to discourage reliance on unstable features, thereby enabling more robust predictions across heterogeneous samples (Miao al. 2022). GALA explicitly models
multiple latent graph views and aggregates them to capture sample- happle-leee level causal diversity (Chen et al.)
2023 AIA generates new environments through adversarial augmentation while preserving stable features, simulating heterogeneity under covariate shift (Sui et al., 2023). Other approaches such as FLAG (Lu al., 2024), StableGNN Yaao et al.] 2024), and GraphMETRO (Wu al. 2024) similarly expand training coverage via data augmentation.
GSAT enforces adversarial consistency discourage reliance on unstable features, thereby enabling more robust predictions across heterogeneous samples (Miao et al. 2022). GALA explicitly models multiple latent graph views and aggregates them to capture sample- level causal diversity (Chen et al,
2023). AIA generates features, simulating heterogeneity under covariate shift (Sui et al. (023). FLAG enforces large-scale augmentation by applying instant feature-level adversarial noise during training while preserving graph topology (Lu et al. 2024). StableGNN pools subgraphs into high-level variables and adds a causal-variabes distinguishing penalty term to prioritize stable correlations under distribution shifts [Fan et al. 2023).
environments through adversarial augmentation while preserving stable
B.3 SEETURRE-OO-EPPRTTE FOR GNNs.
MoE architectures allocate computation across experts selected by gating function (Jacobs al., 1991 Shazeer et al. 2017). On graphs, GMoE leverages MoE layers to scale graph transformers efficiently (Wang et al., (023a). MixGNN employs expert routing improve efficiency on large- scale graph tasks (Chen et al. 2025). Other variants (Hu et al. 2022) similarly focus on distributing computation or scaling to large graphs. These methods demonstrate the potential of MoE for graphs
but do not target OOD generalization.
B.4 GRAPIMETRO..
In the OOD setting, GraphMETRO decomposes distributional heterogeneity and aligns referential representations across shifts by gated mixture-of-experte with shift- seeft-specialized experts (Wu et al.) 2024). However, this design requires the selection of shift types before training with no guarantee that the selected shift types are the ones that will be encountered at test time. Furthermore, these perturbations risk altering label semantics unknowingly. contrast, our approach uses MoE not for augmentation, but for causal subgraph identification, encouraging experts extract diverse causal subgraphs and thereby directly modeling instance heterogeneity.
B.5 OUR APPROACH.
Unlike prior methods, do not rely strong assumptions about the underlying SCM, as in many causal approaches, nor do we risk altering the label semantics through perturbations, as in augmentation--asse methods. Our framework instead learns identify candidate causal sub- graphs through expert-specific masks, with a sparse gating mechanism selecting among them on per-instance basis. This design encourages experts specialize in distinct causal explanations and allows the model to adaptively choose the most relevant one for each input, thereby capturing instance-level heterogeneity.
C IMPLEMENTATION DETAILS
C.I EXPERT SUBGRAPII EXTRACTION
Given an input graph x EV, we first compute node embeddings using a shared GNN encoder. Each expert k E [K] produces edge-level mask logits (k) via small expert-specific MLP
applied the concatenated embeddings of the edge endpoints. The logits are transformed into binary masks using Gumbel-sigmoid straight-through estimator:
m (k) 1
(00)
> 0.5
16

---

Published as a conference paper at ICLR 2026
where g is Gumbel noise, is the sigmoid, and 0.1 is the temperature. Hard binary masks m are used the forward pass to produce the masked graph G while gradients flow through the continuous relaxation during backpropagation. Node weights are induced from their incident edges, and isolated nodes are removed. Each expert has its own GNN encoder and classifier head. Unless otherwise specified, adopt the Graph Isomorphism Network (GIN) following the GOOD benchmark configuration: hidden dimension 300, depth 3, dropout 0.5.
C.2
TASK LOSS
For each expert compute the per-sample cross-entropy
c (a, v) = CE(O,, v).
The overall task loss is gate-weighted aggregation across experts:
(zn v) =
(I)
CE (z. (z,y),
k=1
ensuring that experts favored by the gate receive stronger gradients, thereby encouraging specializa-
tion.
C.3 MASK REGULARIZATION
To control the size of extracted subgraphs, we regularize the average fraction of edges retained by each expert. For expert k, the observed keep-rate on graph g is
(k)
m (k)
CEE
We penalize deviations from a target p E [0, 1] via
(k)
(0/ - p)2
reg
This discourages degenerate solutions where experts keep either too few or too many edges.
C.4
DIVERSITY LOSS
To prevent collapse of experts onto identical subgraphs, enforee the semantic diversity condition (Definition [.I.). Masks are standardized, correlations are computed, and we penalize high off- diagonal correlations:
Σ ma |C! - Tourr) A
Lav K(K -
where C/ is the correlation between standardized masks of experts i and j on graph g. This loss directly encourages experts to specialize on different subgraphs.
C.5 GATING MECHANISM
The gate does not operate directly on the raw graph but instead consumes diagnostic, label-free features derived from each expert's outputs. For expert and sample b, we construct feature vector Pb,6 E I0:
Gate inputs.
max max (c), marging, H(r.k), - llog Σ enl()) -KL(pE, p.raa))
Σ KL(pb& (pbs)) (( H (p na, ml]
J
where Ph, (26, are class probabilities from expert k, marging, is the difference between the top-1 and top-2 probabilities, H(-) is Shannon entropy, and the energy term is
17

---

Published as a conference paper at ICLR 2026
- log Ee exp(2x,x())). The KL terms capture stability under weak augmentations and pairwise disagreement between experts. yon and Spur are outputs of environment and spurious classifiers, while ne, and my denote the number of nodes and edges. The feature vector is passed through two-layer MLP with hidden dimension 64 and ReLU activations. Outputs are normalized with Entmax(a = 1.38), yielding sparse gate probabilities T6,1:K-
Gate loss.
The gate is trained with three complementary objectives:
(i) Stacher-student alignment. A teacher distribution is constructed from expert competence signals:
Tb (Z), 9b) Wia &LA,, Wea EFA.
Tb softmax
The student is the gate output = Entmaxx=1.35($5) where Sy are gate scores. The alignment loss is
Lalign -KL(gb po).
(ii) Balanced usage. To avoid collapse to single expert, we regularize the average gate distribution over batch:
Lbalance KT(u
p
Σ pb
where u is uniform over experts.
(iii) Per-sample sparsity. To encourage sparse routing, we penalize high-entropy gate outputs:
B
K
Lsparse
Σ H(p)),
H(p))
Σ
Pb, log Pb.k-
b=1
The final gate loss is
Lgaae Lulign + Lballnce Isparse-
C.6
FINAL OBJECTIVE
The overall training loss is
C lep + Areghreg + Adivediv + Lagto--
Here freg K Σx coverage, selection, and diversity principles while guarding against degenerate solutions
(k)
and Areg Adiv, Agate are loss weights. This formulation integrates the task, and expert
reg
starvation.
C.7 HYPERPARAMETER SETTINGS
For all experiments, we adopt common set of fixed hyperparameters. Each expert network is implemented a Graph Isomorphism Network (GIN) with hidden dimension 300, three layers, and a dropout rate of 0.5. The gating network is a two-layer MLP with hidden dimension 64 and ReLU activations, followed by Entmax transformation with 1.38 to produce sparse routing probabilities. The loss weights are set to 1.0 for the task cross-entropy, mask regularization, and diversity losses, and 0.1 for the gate loss. Optimization decay of 10-1 The learning rate is reduced by a factor of two if the validation performance does not improve by at least 0.001 for 10 consecutive epochs.
performed with Adam using a weight
In addition to these fixed settings, we tune small number of hyperparameters. The mask keep-rate prior p is sampled uniformly between 0.1 and 0.9. The learning rate is selected from (0.001, 0.0005, 0.0001). and the batch size is chosen from {32, 64, 128, 256, 512]. Hyperparam- eter tuning is conducted with 10 independent trials per dataset, each initialized with different random seed. The best configuration is selected based validation performance, retrained on the training set, and finally evaluated on the held-out test sets over five different seeds. This procedure ensures consistent model selection without any test leakage.
18

---

Published as a conference paper at ICLR 2026
D CAUSAL ASSUMPTIONS
Problem setup. Each input is a graph x - (V,E,)) with label y E V. Let Ge denote causal/stable substructures and G, denote sparious/eevirrnmnnttal substructures. We focus co- variate shift: the marginal over graphs changes across environments, while p(y x) is stable, i.e., Prrin (x) + Prea (x) but Prrin x) Pes x))
FIIF vs. PII Following (Gui et al. 2023), fully/partially informative invariant features formalize
how Gg interacts with G. and environments E. Under FIIF, G, is influenced by both G. and E, creating spurious dependence Y < Ge Gss under PIIF, collider induces G. ++ Y through Ge (e.g., Y + Ge + + Gy). These regimes explain why G, can correlate with even when noncausal.
WHAT PRIOR METHODS ASSUME
LECI Gui et al. 2023). LECI jointly enforces (i) label-enviromment causal independence: Ge 1 E and (i) label spurious independence: LY (operationalized via adversarial objectives), aiming to recover invariant subgraphs. In FIIF/PIIF regimes, these constraints are used to filter G2 and keep Ge
UIL (Sui et al., 2025). UIL posits stronger structural invariance: the stable (causal) part of graphs within each class shares class-specific graphon pattern that invariant across all (seen/unseen) environments; semantic invariance is then layered on top. In effect, Ge is assumed structurally invariant across environments.
DIR and related. DIR (Wu et cally constructing interveniionlllaugmentte environments and enforcing invariance of the predictive mechanism across them.
2022) also targets invariant rationales/causal attention, typi-
WHY THESE ASSUMPTIONS BREAK IN PRACTICE
Concrete examples.
Chemistry. Multiple active chemotypes for the same endpoint; scaffolldtimessorcee shifts alter G2 and even the prevalence of certain Ga's, violating structural uniqueness and G2 _ Y (Wu al., 2018 Hu et al. 2020).
Social/text graphs. Domain-specific syntax/community structures change across splits;
G, (e.g., degree/length) correlates with Y via FIFFPII pathways (Gui al. 2022
General reasons. (i) Instance heterogeneity: real tasks often admit multiple causal explanations within the same class (different G. ' per instance). In molecular property prediction, distinct chemo- types/functional groups can yield the same label (e.g., multiple acidic moieties), and scaffold splits explicitly emphasize cress-chemotype variation; thus any single class-graphor assumption can be violated. (ii) Y Y-G, correlation (FIIF/PIIF): even when G, is noncausal, it may correlate with Y via Ge. making G, VY and breaking LECI-style independence assumptions. Recent theory shows that environment augmentation cannot, in general, identify invariance without additional biases; moreover, G2 and Y can have arbitrary correlation, making environment inference/labeling funda- mentally hard (Chen et al. 2023).
OUR ASSUMPTION (MINIMAL AND ROBUST)
We assume only that each graph admits at least one causal subgraph Ge that governs Y. We do not assume (i) Ge is unique within a class, (ii) G, is structurally identical across environments (no class-graphon), or (iii) G, .LY Y or G. l E. This minimal assumption tolerates instance-level heterogeneity (different G per instance), admits FUFF/PIFF couplings, and aligns with practical datasets where multiple mechanisms yield the same label (e.g.. multiple binding motifs syntax patterns). Benefit: we avoid brittle structural/indeeendenne assumptions and instead learn to select among diverse causal hypotheses at the instance level.
19

---

Published as a conference paper at ICLR 2026
Table 7:
Performance of the MoE framework with 4, and experts. Standard deviations are
reported in parentheses.
HIV T
Twitter f Length
Motif f
SST2 f Length
Experts
Avrg
Scaffold
Size
Basis
Size
1 67.13 (2:2
63.48 (0.9)
58.68 (1.7)
89.3 (2.1)
65.31 (122)
81.11 a..
70.84
4
70,65
65,35 (1.3)
60.48 (0.7)
91.79 (1.5)
76.40 12.1)
83,41 (22)
74.68
8
71.55
66.98 (1.0)
61.13 (1.1)
92.8 (1.4)
75.52 (2.)
83.73
75.29
EMPIRICAL SUPPORT: SINGLE-EXPERY PERFORMANCE
Table 7 provides empirical evidence for the strength of this assumption-light design. the single- expert setting, the model reduces to the simplest form of causal subgraph extraction: single extrac- tor, regularized by feeding into a GNN trained like standard ERM. The GNN is optimized only with the task loss, while the extractor trained with the task loss plus the regularization term. This design makes no additional assumptions about the underlying causal SCM beyond the sparsity prior, representing a minimal instantiation of causal subgraph methods. Remarkably, despite its simplic- ity, this approach outperforms all but one causal baseline (LECI) in average performance across all datasets. This result suggests that the explicit structural assumptions encoded in prior causal meth- ods may in fact be too restrictive or fragile, and that a more assumption approach can provide stronger and more reliable generalization.
E DATASET DETAILS
We provide dataset-specific information for the four GOOD tasks used in the main paper. Each task introduces different type of structural or distributional shift, following the benchmark design in Gui al. (2022).
This task a molecular property prediction problem derived from the MoleculeNet HIV dataset. Each graph corresponds molecule, where nodes are atoms and edges are chemical bonds. The prediction task is binary classification: whether the molecule inhibits HIV replication. To evaluate OOD generalization, two types of environment splits are defined. The scaffold split par- titions molecules according to their core scaffolds, such that training and test sets contain molecules with distinct underlying structures. The size split creates a distribution shift by separating molecules based on the number of heavy atoms, exposing models to molecules of substantially different sizes at test time, These shifts test whether models can generalize beyond memorized molecular backbones and size ranges.
GOOD-HIV.
This synthetic dataset is designed to provide controlled graph-level classification tasks with interpretable shifts. Each graph is generated by attaching motifs (e.g., cycles, cliques, houses) onto random base graph. The label depends on the presence or type of motif. The bench- mark defines
GOOD-Motif.
kinds of shifts. In the motif-basis split, the set of motifs used for training differs from those used in evaluation, requiring extrapolation across structural primitives. In the size split, the base graphs differ in size between environments, requiring robustness distribution shifts in graph order and density. This dataset isolates structural shifts in a controlled synthetic setting, mak- ing useful for probing whether models can truly capture causal motif information.
This dataset consists of ego-networks from Twitter users. Each ego-network is represented as a graph where the central node corresponds the ego user, and edges represent social connections among the ego and their neighbors. The task is binary classification of user attributes. The primary distribution shifts are domain shifts across different user communities, which result in differences in graph sparsity, degree distributions, and local structural motifs. Since ego-networks are sampled from diverse domains, training and test sets differ significantly in their structural prop- erties, requiring models to generalize across heterogeneous social network subgraphs.
GOOD-Twitter:
GOOD-SST2. This dataset is based on the Stanford Sentiment Treebank (SST2). Each sentence parsed into dependency tree, which serves the input graph. The task is binary classification of
20

---

Published as a conference paper at ICLR 2026
sentence sentiment (positive negative). The distribution shifts are introduced by partitioning the data according to linguistic structures. Specifically, environments differ in the average tree depth and branching factors, leading to structural shifts in dependency graphs. These shifts test whether graph models can capture sentiment cues in syntactic structures when faced with substantial variation in parse tree topology across domains.
Overall, these four datasets cover both real-world domains and synthetic graphs, and they introduce diverse OOD challenges, including scaffold and size shifts, motif-basis changes, domain hetero- geneity, and structural variation. This variety makes them a comprehensive testbed for evaluating the robustness of graph OOD methods.
F INTERPRETABILITY RESULTS
Motif 1: Usage vs Accuracy (Per Expert)
Accursyythage
04
Experts
MACE BUsaae
Motif 2: Usage vs Accuracy (Per Expert)
1.2
AcuraryUUage
0.6
Experts
Motif Usage vs Accuracy (Per Expert)
01
Experts
Figure 2: Per-expert usage and accuracy across motifs (MOTIF-BASIS). For each motif, we report the average gating probability (usage) and the corresponding per-expert accuracy, averaged over all test instances containing that motif.
This section provides interpretable diagnostic showing how the gating network routes instances exhibiting different heterogeneous causal subgraphs to different experts, and whether routing aligns with expert performance.
We use the MOTIF-BASIS dataset from the GOOD benchmark, which contains three heterogeneous causal subgraphs (three distinct motifs) (Gui et al. 2022). After training, evaluate the test set and compute, for each motif m E {1,2, 3) and each expert k E (1....,K K} (here K = 8): (i) usage, defined as the average gating probability assigned to expert k. averaged over all test instances
21

---

Published as a conference paper at ICLR 2026
containing motif m; and (ii) accuracy, defined as the average classification accuracy of expert & on the same subset of instances. This yields x 8 (usage, accuracy) pairs summarizing routing behavior and expert specialization.
F.1 PER-MOTIF ROUTING PATTERNS.
Figure visualizes per-expert usage for each motif. We can observe that for each motif, the gate concentrates probability mass on a subset of experts rather than using all experts uniformly, Further- more, the subset of frequently-used experts changes across motifs, indicating that the gate responds to the underlying causal subgraph. Concretely, motif primarily uses experts 2, 3,8;; motif primarily uses {1, 4, 5, 6, 7); and motif primarily uses {3, 4, 5, 7.8). Several experts are nearly unused for specific motifs (e.g., expert 6 for motifs and 3), while becoming prominent for others (e.g., expert 6 for motif 2). This demonstrates that the experts are indeed trained to specialize on different causal subgraphs, empirically supporting our claim of semantic diversity.
F.2 ALIGNMENT BETWEEN ROUTING SELECTION AND EXPERT PERFORMANCE.
Figure plots, across all (m, k) pairs, expert usage expert accuracy, with points grouped by motif category. Routing aligns with expertise: experts that are more frequently selected for a motif tend achieve higher accuracy instances containing that motif. Quantitatively, Pearson correlation test between usage and accuracy over all (m, k) pairs yields r = 0.823 with < 10-?, indicating strong positive association between gate probability and expert correctness for the corresponding causal subgraph which supports our claim of correct expert selection.
F.3 VISUALIZATION OF PREDICTED SUBGRAPHS.
Figure 4 visualizes the subgraphs extracted by a four-expert MoE on three training-set graphs from Motif-basis containing the house motif. These visualizations show that different experts extract distinct subgraphs, demonstrating semantic diversity. Moreover, for each instance, only one or two experts correctly recover the house motif, highlighting expert specialization.
Overall, these results provides strong evidence for the interpretability of our method and visually support our claims diverse expert specialization and successful routing alignment.
Expert Usage Accuracy by Category
1.0
Categories
Category
Category
Category
Experts
Accuracy
Expert 1
Expert 2
0.4
Expert 3
Expert
Expert
Expert 6
0.2
Expert 7
Expert 8
0.0
0.00
0.05
0.10
0.15
0.20
0.25
Usage
Figure 3: Usage vs. accuracy by motif category. Each point corresponds one (motif, expert) pair. Higher usage aligns with higher accuracy, suggesting that the gate routes instances toward experts that specialize in the corresponding motif.
22

---

Published as a conference paper at ICLR 2026
Figure 4: Visualization of the subgraphs extracted by four-expert MoE on three training-set graphs from Motif-basis containing the house motif. Edges highlighted in green are extract by the expert.
Table 8: Effect of learned gating on Twitter and SST2. Uniform: uniform averaging of logits (no gating). Majority: majority voting (no gating). Top sparse selection of the top-1 expert predicted by the gate. Top 2: sparse selection of the top-2 experts predicted by the gate. Gate Weight (Soft): weighted averaging of logits using the gate weights (used in our method).
Learned Gating
Twitter
SST2
Uniform
59.92
82.72
Majority
58.39
82.99
Top
56.36
79.36
Top
59.11
82.13
Gate Weight (Soft)
61.13
83.73
G ADDITIONAL EXPERIMENTS
In this section, we provide additional experiments and discussions.
G.1 HIV DATASET
For HIV dataset, we apply class balanced sampling during batch sampling and logit adjustment during training (Wang al. 2023b). This follows the convention of LIRS and GALA who apply similar techniques to mitigate the class imbalance Yao et al. 2025 Chen et al. (0023). Without these techniques, our performance 66.98 to 65.35, respectively. this case, our method still achieves the best average performance and average rank across all evaluated baselines.
HIV-Scaffold and HIV-Size drops slightly from 71.55 70.65 and
G.2 GATING MECHANISM
We next assess the contribution of the gating mechanism. Our framework employs a learned gate that assigns input-dependent weights to experts, forming a weighted combination of their predictions. Table 8 compares this design to several alternatives. On both Twitter and SST2, uniform averaging already benefits from expert complementarity. but falls short the learned gate (59.92 vs. 61.13 on Twitter, 82.72 83.73 SST2). Majority voting performs similarly worse (58.39 and 82.99), suggesting that ignoring the confidence of individual experts limits robustness. Sparse selection of only the top-1 top-2 experts consistently underperforms (e.g., 56.36/79.36 and 59. 1/82.18/1))
23

---

Published as a conference paper at ICLR 2026
Table 9: Computational efficiency analysis by varying number of experts (1, 4, 8, 16. 32) and measuring training time, inference time, and OOD accuracy on Twitter, SST2, and Motif-Basis
Num Experts
Training Time
Inference Time
Parameters
Accuracy
Twitter
1
3m12s
4s
328,641
58.68
4
3m18s
4s
11,091,324
60.48
8
3m30s
4s
21,497,876
61.13
16
3m47s
5s
42,310,980
60.05
32
4m34s
6s
83,937,188
60.12
SST2
38m35s
26s
3,284,002
81.11
4
45m19s
42s
11,081,692
83.41
8
58m47s
50s
21,478,612
83.73
16
1h15m21s
50s
42,272.452
83.72
32
1h56m05s
1m20s
83,860, 132
83.92
Motif-Basis
13m45s
6s
2,361,796
83,9
4
15m38s
6s
8,083,168
91.79
8
17m07s
6s
15,711,664
92.8
16
22m34s
10s
30,968,656
92.47
32
31m16s
12s
61,482,640
92.33
showing that over-reliance on top-k experts all samples discards useful complementary informa- tion. By contrast, the learned gate yields the best performance on both datasets, highlighting that adaptive weighting crucial for fully exploiting expert specialization. Importantly, the gate not only aggregates predictions effectively but also mitigates expert collapse by encouraging specialization into complementary substructures.
G.3 EMPIRICAL COMPUTATIONAL EFFICIENCY ANALYSTS
In Section 3.5 it demonstrated that the computattional overhead grows linearly with respect to the number of experts. This is validated by the empirical results in Table 9 Notably, in prac- tice, both training and inference times increase sub-linearly with the number of experts while OOD accuracy remains largely stable with slight improvementsdeeraaationn depending on the dataset. Furthermore, the results indicate that eight experts provides good balance between computational efficiency and OOD accuracy.
G.4 TRAINING STABILITY
Across all experiments, we did not observe training instability or divergence. illustrate this, we report the training loss curves MOTIF-BASIS in 5iguhe loss decreases smoothly and consistently. We attribute the observed stability to two design choices:
which shows that each component of the
Uniform gating warm-up: we employ uniform gating in the first & 10 epochs to ensure that all experts receive sufficient training signal before the gating network begins selective routing. This warm-up ensures each expert learns reasonable initial representation, which stabilizes early optimization.
Module-specific loss: Different components of the loss optimize different modules of the model which reduces potential gradient conflicts. The gating module only receives the gating loss, the experts only receive the task loss (cross-entropy). the subgraph extractor is the only component that receives multiple losses (task, diversity, and regularization).
LARGE LANGUAGE MODELS STATEMENT
Large Language Models were used exclusively for editorial purposes, such as refining language and improving readability. All scientific contributions were developed solely by the authors.
24

---

Published as a conference paper at ICLR 2026
Training Curves on Motif-Basis
0.8
3
9
10
11
12
13 Epochs
14
15
16
10
19
20
21
23
Figure 5: Training losses on MOTIF-BASIS, We use uniform gating for the first ten epochs: thus, the diversity and gating losses are reported starting at epoch 10.
25


Reviews:
[{"review_id": "xJkX22WMXl", "reviewer": "Reviewer_4z1L", "summary": "The paper proposes a causal subgraph–based MoE framework for graph OOD generalization. The key idea is to consider instance-level different causal subgraph conceptual experts. The authors give a risk decomposition showing that diversity ensures at least one expert matches the unseen environment, and sparsity ensures the gate actually picks it. Empirically, on the GOOD benchmark, the method achieves the best average rank and clearly improves on strong OOD baselines such as LECI and GALA.", "questions": "N/A", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The paper correctly identifies that many causal-subgraph OOD methods assume one invariant causal subgraph per task/environment, but different instances often rely on different motifs.\n2. The idea of learning different causal subgraphs using MoE is fresh. I will explain from a different perspective, though: different MoE learns different spurious correlations. Given enough experts, there will always be one expert learning no spurious correlations with the diversity constraint.\n3. The experimental results are very strong with useful and complete ablations.\n4. The assumption is lighter, easier to extend.", "weaknesses": "1. Given the story, there is few interpretability results in terms of how MoE solves the instance-level heterogeneity problem.\n2. LECI provides several enhanced Motif datasets, CFP-Motif, which represent different causal assumptions. The authors may train the three splits together. If the MoE assumption stands, we should observe that, given different underlying causal assumptions, the gate activation should follow clear patterns. Without this experiment or other alternative, the claim will not be justified enough.\n3. The authors show 1, 4, 8 experts, but in MoE literature there is always a question: does diversity collapse or does routing become unstable when K=16/32 on small datasets?\n4. Lack of visualizations of learned subgraphs per expert.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wKNeMO9iiF", "reviewer": "Reviewer_GS29", "summary": "The paper introduces a novel MoE framework designed to extract diverse subgraphs and address dataset heterogeneity. The proposed approach theoretically demonstrates that both semantic diversity among experts and sparse expert selection contribute to reducing prediction error. Extensive experiments are conducted to validate the framework’s effectiveness across different datasets, highlighting its potential for improving generalization in heterogeneous graph learning scenarios.", "questions": "None, see weaknesses.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The paper proposes a novel MOE framework for extracting diverse subgraphs, aiming to address dataset heterogeneity.\n\n2. A theoretical analysis is provided to demonstrate that semantic diversity among experts, together with sparsity in expert selection, jointly contributes to error reduction.\n\n3. Extensive experiments are conducted to validate the effectiveness of the proposed approach.", "weaknesses": "1. It is necessary to clarify which components are genuinely novel, and whether the proposed method is merely a combination of existing MoE techniques.\n\n2. The total loss comprises multiple terms, so it would be helpful to explain how their relative influences are balanced. More detailed ablation studies are needed to further validate the design, and to examine whether each component is essential for the final performance.\n\n3. Considering the trade-off between performance and computational cost, how many experts are recommended to be used? A detailed analysis or estimation of computational efficiency would also strengthen the work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8yDWLWIR7G", "reviewer": "Reviewer_q7vv", "summary": "This work introduces a Mixture-of-Experts (MoE) approach for addressing out-of-distribution graph classification tasks. Intuitively, each expert would learn a (perhaps disjoint) subset of the environment distributions. By softly enforcing diversity and sparseness of the expert distribution, the authors ensure that the resulting method both covers the environment space and selects the correct expert for that environment with high probability. Experimental results highlight the effectiveness of the method for graph classification.\n\nIn summary, the work is nicely written, the theoretical analysis (although simple) provides a clear understanding of the method, empirical results are promising, and notation is clean. Although the manuscript presents minor issues (e.g., $l_{sparse}$ is only described in the appendix - with a different name), I am happy to provide a favorable review for the work.", "questions": "Please refer to the questions above.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The method (although unnamed) is clearly described. \n\n2. The theoretical analysis is to-the-point and provides an intuitive understanding for the chosen design.\n\n3. Empirical results highlight the method’s competitive performance.\n\n4. Ad-hoc terminology (e.g., semantic diversity, mechanism coverage) is rigorously defined.", "weaknesses": "1. I believe that the main disadvantage of the introduced design is the introduction of many hyperparameters for model regularization. In practice, how can they be selected? Although there are standard approaches in machine learning for hyperparameter selection, these are often expensive.\n\n2. Also, I am a bit doubtful about the use of the word OOD; usually, this means a difference between training and test datasets. It is unclear how it is represented in the experimental analysis in this section. Could the authors elaborate briefly on this?\n\n3. Diversity and sparsity constraints seem to solely lead to marginal improvements over the baseline method - which also presents competitive performance regardless. Are there datasets for which the imposed constraints lead to more significant performance boosts?\n\n4. Could the authors discuss the runtimes for the presented methods in Table 1? Also, when adding further constraints (e.g., diversity and sparsity), the optimization problem changes; does it significantly affect the convergence of the stochastic gradient algorithm?\n\n5. [Minor] $l_{sparse}$ in line 304 uses a different notation in the appendix. \n\n6. [Minor] I could not find the anonymous repository referred to in the manuscript. Could the authors share the code?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6Ls0TOGMH6", "reviewer": "Reviewer_nffD", "summary": "This paper proposes a Mixture-of-Experts (MoE) framework to address the out-of-distribution (OOD) generalization problem for graphs. The method explicitly models the instance-level causal heterogeneity prevalent in real-world data by forcing experts to extract diverse causal subgraph hypotheses and using a sparse gating network to adaptively select the best explanation for each instance.", "questions": "- Q1 (Regarding W1): The concepts of instance-level heterogeneity and causal diversity need further elaboration. What is causal diversity, and is it a concept specific to the instance level? A related issue is the assumption of instance-level heterogeneity, which should be explicitly stated. Does this assumption imply that different instances can have different causal graphs? If so, why would different instances from the same data source follow different causal graph structures, and what are the corresponding real-world scenarios?\n\n- Q2 (Regarding W2): My understanding of the method is as follows: different experts are used to extract different subgraph structures based on a diversity constraint. A selection mechanism then assigns the maximum predictive weight to the expert with the lowest prediction loss, and the subgraph extracted by this expert is considered the causal part. Is this understanding correct? If so, the generalization guarantee provided by Assumption 3.5 is, in essence, a strong constraint on the structure of the test graphs. Could you please discuss this?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- S1: The paper is well-motivated, with a clear and easy-to-follow logic.\n\n- S2: The paper provides a clear theoretical framework that decomposes the OOD risk into coverage and selection terms. It directly links the proposed diversity and sparsity mechanisms to the goal of reducing OOD error, effectively supporting the paper's core claims and contributions.\n\n- S3: The authors have conducted extensive experiments on standard benchmarks to demonstrate the effectiveness of their method. Furthermore, the ablation studies validate the necessity of each core component of the framework.", "weaknesses": "- W1: The core premises of the paper, namely instance-level heterogeneity and the resulting causal diversity, require more formal and clear definitions and introductions.\n\n- W2: Assumption 3.5 appears to be a critical condition for the method to achieve OOD generalization. It states that at least one expert can achieve satisfactory performance in the test environment. This assumption may be too strong, as it implicitly requires that the causal subgraph structures present in the test data must have been seen during training and learned by at least one expert.\n\n- W3: The method relies on training multiple expert networks, yet it lacks a corresponding analysis of computational efficiency. Although the paper claims the method is scalable, it does not provide a quantitative analysis of the relationship between computational cost (e.g., training time) and the number of experts, K.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)