Generalizing GNNs with Tokenized Mixture of Experts

Deployed graph neural networks (GNNs) operate as frozen snapshots, yet must simultaneously fit clean data, generalize under distribution shifts, and remain stable against input perturbations---three goals that are difficult to satisfy at once with a single fixed model. We first show theoretically that a single fixed inference rule can create a stability--generalization tradeoff: making the model insensitive to perturbations can also suppress task-relevant signals needed to fit and generalize. Input-dependent routing---assigning different computation paths to different inputs---can relax this tension, but brings new fragility: distribution shifts may mislead routing decisions, and perturbations can destabilize routing, compounding downstream errors. We formalize these effects through two risk decompositions that separate (i) how well the available paths cover diverse test conditions from how accurately the router selects among them, and (ii) how sensitive each fixed path is from how much routing fluctuation amplifies that sensitivity. Guided by these analyses, we propose STEM-GNN: Stable TokEnized Mixture-of-Experts GNN, a pretrain-then-finetune framework that couples a mixture-of-experts encoder providing diverse computation paths to cover heterogeneous test conditions, a vector-quantized token interface that maps encoder outputs to a discrete codebook to absorb small representation drift induced by input perturbations and routing fluctuations before downstream layers, and a Lipschitz-regularized prediction head that bounds how much the output can amplify residual upstream variation. Across eight node, link, and graph benchmarks, STEM-GNN maintains strong clean performance; on representative node benchmarks, it improves the three-way balance under degree/homophily shifts and feature/edge perturbations. The code and data are available at https://github.com/GXG-CS/STEM-GNN.