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 applying the same computation to every input creates a fundamental tradeoff: to resist perturbations, the model must limit how much it relies on shift-sensitive features, which leaves an irreducible floor on worst-case generalization error. Input-dependent routing---assigning different computation paths to different inputs---can break this ceiling, but brings new fragility: distribution shifts may misguide routing decisions, and perturbations cause the routing itself to fluctuate, 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 routing fluctuations before they reach downstream layers, and a regularized prediction head that bounds how much the output can amplify any residual upstream variation. Experiments on nine benchmarks spanning node, link, and graph tasks show that STEM-GNN achieves the best three-way balance, improving generalization under degree and homophily shifts and stability under feature masking and edge deletion, while maintaining competitive performance on clean graphs. The code and data are available at.