Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Generalizing GNNs with Tokenized Mixture of Experts

Xiaoguang Guo, Zehong Wang, Jiazheng Li, Shawn Spitzel, Qi Yang, Kaize Ding, Jundong Li, Chuxu Zhang

TL;DR

The paper tackles the impossible-triangle challenge of deploying frozen GNNs that must fit clean data, generalize under distribution shifts, and resist perturbations. It analyzes static inference vs instance-conditional computation (ICC) under a tri-objective framework, revealing a stability floor for static methods and actionable levers for ICC via coverage, selection, and stability decompositions. STEM-GNN implements robust ICC through a three-component design: (1) a MoE encoder to expand the deployed set of mechanisms, (2) vector-quantized tokens to stabilize the encoder-to-head interface, and (3) Lipschitz-regularized heads to bound output amplification. Across nine benchmarks, STEM-GNN achieves the best tri-objective balance, showing strong robustness to degree/homophily shifts and perturbations while remaining competitive on clean data, and demonstrates favorable transferability under diverse pretraining sources. This approach offers a practical path to robust graph generalization in fixed, production-like deployment settings.

Abstract

Deployed graph neural networks (GNNs) are frozen at deployment yet must fit clean data, generalize under distribution shifts, and remain stable to perturbations. We show that static inference induces a fundamental tradeoff: improving stability requires reducing reliance on shift-sensitive features, leaving an irreducible worst-case generalization floor. Instance-conditional routing can break this ceiling, but is fragile because shifts can mislead routing and perturbations can make routing fluctuate. We capture these effects via two decompositions separating coverage vs selection, and base sensitivity vs fluctuation amplification. Based on these insights, we propose STEM-GNN, a pretrain-then-finetune framework with a mixture-of-experts encoder for diverse computation paths, a vector-quantized token interface to stabilize encoder-to-head signals, and a Lipschitz-regularized head to bound output amplification. Across nine node, link, and graph benchmarks, STEM-GNN achieves a stronger three-way balance, improving robustness to degree/homophily shifts and to feature/edge corruptions while remaining competitive on clean graphs.

Generalizing GNNs with Tokenized Mixture of Experts

TL;DR

The paper tackles the impossible-triangle challenge of deploying frozen GNNs that must fit clean data, generalize under distribution shifts, and resist perturbations. It analyzes static inference vs instance-conditional computation (ICC) under a tri-objective framework, revealing a stability floor for static methods and actionable levers for ICC via coverage, selection, and stability decompositions. STEM-GNN implements robust ICC through a three-component design: (1) a MoE encoder to expand the deployed set of mechanisms, (2) vector-quantized tokens to stabilize the encoder-to-head interface, and (3) Lipschitz-regularized heads to bound output amplification. Across nine benchmarks, STEM-GNN achieves the best tri-objective balance, showing strong robustness to degree/homophily shifts and perturbations while remaining competitive on clean data, and demonstrates favorable transferability under diverse pretraining sources. This approach offers a practical path to robust graph generalization in fixed, production-like deployment settings.

Abstract

Deployed graph neural networks (GNNs) are frozen at deployment yet must fit clean data, generalize under distribution shifts, and remain stable to perturbations. We show that static inference induces a fundamental tradeoff: improving stability requires reducing reliance on shift-sensitive features, leaving an irreducible worst-case generalization floor. Instance-conditional routing can break this ceiling, but is fragile because shifts can mislead routing and perturbations can make routing fluctuate. We capture these effects via two decompositions separating coverage vs selection, and base sensitivity vs fluctuation amplification. Based on these insights, we propose STEM-GNN, a pretrain-then-finetune framework with a mixture-of-experts encoder for diverse computation paths, a vector-quantized token interface to stabilize encoder-to-head signals, and a Lipschitz-regularized head to bound output amplification. Across nine node, link, and graph benchmarks, STEM-GNN achieves a stronger three-way balance, improving robustness to degree/homophily shifts and to feature/edge corruptions while remaining competitive on clean graphs.
Paper Structure (46 sections, 10 theorems, 90 equations, 5 figures, 10 tables)

This paper contains 46 sections, 10 theorems, 90 equations, 5 figures, 10 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Problem Formulation and Analysis
  4. Problem Formulation
  5. $\mathcal{H}_1$: Static Inference under a Stability Budget
  6. $\mathcal{H}_2$: Instance-Conditional Computation under Frozen Deployment
  7. Methodology
  8. Coverage Expansion via MoE Encoder
  9. Stabilizing Intermediate Representations via VQ Tokenization
  10. Bounding Output Amplification via Lipschitz Regularization
  11. Experiments
  12. Experimental Setup
  13. Results on Clean Graphs
  14. Results on Noisy Graphs
  15. Out-Of-Distribution Generalization
  16. ...and 31 more sections

Key Result

theorem 1

Under the assumptions, we define: Here $\psi_{\mathrm{fit}}(\eta)$ is the monotone lower-bound function, relating the fitting risk to the reliance level $\eta$, and $e_1\in E_{\mathrm{test}}$ is the witness environment specified with its associated monotone lower bound $\psi_{e_1}(\eta)$. If $\eta_{\min}(\alpha)>\eta_{\max}(\varepsi

Figures (5)

  • Figure 1: An illustrative example: fixed GNN cannot adapt to diverse deployment scenarios. On a knowledge graph where task-relevant information lies in high-frequency components, a low-pass GNN removes key signals. On a citation graph dominated by high-frequency noise, the same low-pass GNN suppresses noise and improves stability. No single static computation rule resolves both cases simultaneously.
  • Figure 2: The STEM-GNN framework: the MoE message-passing encoder first expands mechanism coverage via input-dependent routing, then a VQ token interface discretizes intermediate representations with a shared codebook to stabilize the encoder-to-head pathway, and finally a Lipschitz-regularized prediction head bounds amplification of residual variations.
  • Figure 3: Ablation study on clean original graphs (mean $\pm$ std). We compare the full STEM-GNN with variants removing MoE (w/o MoE), Lipschitz regularization (w/o Lip), or VQ (w/o VQ) across all benchmarks, and report the average performance across datasets.
  • Figure 4: Stability under inference-time perturbation on two node benchmarks. Left: Cora. Right: PubMed. In each dataset, (a) shows feature masking with masking rate $\alpha$, and (b) shows structural corruption via edge deletion with deletion rate $p$.
  • Figure 5: Cross-domain transfer with increasing pre-training diversity. STEM-GNN is less sensitive to the source mixture than GFT on PubMed and WikiCS (least-squares fits shown).

Theorems & Definitions (29)

  • theorem 1: Witness-family tension induced by a stability budget
  • lemma 1: Margin-based invariance of the VQ interface
  • proof
  • proposition 1: Bounded output amplification under token switches
  • proof
  • definition 1: Static inference class $\mathcal{H}_1$
  • definition 2: Two-component decomposition induced by the graph
  • definition 3: A witness subfamily for $\mathcal{H}_1$ via a global reliance parameter
  • definition 4: Slice-optimal OOD risk for $\mathcal{H}_1$
  • Remark 1: General perturbations affecting both components
  • ...and 19 more