Robustness and Generalization in Uncertainty-Aware Message Passing Neural Networks

Alesia Chernikova, Moritz Laber, Narayan G. Sabhahit, Tina Eliassi-Rad
Proceedings of The 29th International Conference on Artificial Intelligence and Statistics, PMLR 300:3826-3834, 2026.

Abstract

Existing theoretical guarantees for message passing neural networks (MPNNs) assume deterministic node features. We address a more realistic scenario where noise or finite measurement precision introduces uncertainties in node feature values. First, we quantify uncertainty by propagating the moments of node-feature distributions through the MPNN architecture. To propagate moments through activation functions, we use the Taylor expansion and the pseudo-Taylor polynomial expansion. We then use the resulting node embedding distributions to analytically derive probabilistic adversarial robustness certificates for node classification tasks against L2-bounded perturbations of node features. Second, we model node features as multivariate random variables and introduce Feature Convolution Distance $FCD_p$, a pseudometric based on the Wasserstein distance. $FCD_p$ corresponds to the discriminative power of MPNNs at the node level. We show that MPNNs are globally Lipschitz continuous functions with respect to the pseudometric $FCD_p$. Using the covering number of the resulting pseudometric space, which is a subset of the Wasserstein space, we derive generalization bounds for MPNNs with uncertainties in node features. Together, these two complementary approaches—moment propagation for adversarial robustness and $FCD_p$ on the subset of the Wasserstein space for generalization—establish a unified theoretical framework that comprehensively addresses MPNN reliability under node feature uncertainty.

Cite this Paper

BibTeX
@InProceedings{pmlr-v300-chernikova26a, title = { Robustness and Generalization in Uncertainty-Aware Message Passing Neural Networks }, author = {Chernikova, Alesia and Laber, Moritz and Sabhahit, Narayan G. and Eliassi-Rad, Tina}, booktitle = {Proceedings of The 29th International Conference on Artificial Intelligence and Statistics}, pages = {3826--3834}, year = {2026}, editor = {Khan, Emtiyaz and Li, Yingzhen and Solin, Arno and Ramdas, Aaditya}, volume = {300}, series = {Proceedings of Machine Learning Research}, month = {02--05 May}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v300/main/assets/chernikova26a/chernikova26a.pdf}, url = {https://proceedings.mlr.press/v300/chernikova26a.html}, abstract = { Existing theoretical guarantees for message passing neural networks (MPNNs) assume deterministic node features. We address a more realistic scenario where noise or finite measurement precision introduces uncertainties in node feature values. First, we quantify uncertainty by propagating the moments of node-feature distributions through the MPNN architecture. To propagate moments through activation functions, we use the Taylor expansion and the pseudo-Taylor polynomial expansion. We then use the resulting node embedding distributions to analytically derive probabilistic adversarial robustness certificates for node classification tasks against L2-bounded perturbations of node features. Second, we model node features as multivariate random variables and introduce Feature Convolution Distance $FCD_p$, a pseudometric based on the Wasserstein distance. $FCD_p$ corresponds to the discriminative power of MPNNs at the node level. We show that MPNNs are globally Lipschitz continuous functions with respect to the pseudometric $FCD_p$. Using the covering number of the resulting pseudometric space, which is a subset of the Wasserstein space, we derive generalization bounds for MPNNs with uncertainties in node features. Together, these two complementary approaches—moment propagation for adversarial robustness and $FCD_p$ on the subset of the Wasserstein space for generalization—establish a unified theoretical framework that comprehensively addresses MPNN reliability under node feature uncertainty. } }
Endnote
%0 Conference Paper %T Robustness and Generalization in Uncertainty-Aware Message Passing Neural Networks %A Alesia Chernikova %A Moritz Laber %A Narayan G. Sabhahit %A Tina Eliassi-Rad %B Proceedings of The 29th International Conference on Artificial Intelligence and Statistics %C Proceedings of Machine Learning Research %D 2026 %E Emtiyaz Khan %E Yingzhen Li %E Arno Solin %E Aaditya Ramdas %F pmlr-v300-chernikova26a %I PMLR %P 3826--3834 %U https://proceedings.mlr.press/v300/chernikova26a.html %V 300 %X Existing theoretical guarantees for message passing neural networks (MPNNs) assume deterministic node features. We address a more realistic scenario where noise or finite measurement precision introduces uncertainties in node feature values. First, we quantify uncertainty by propagating the moments of node-feature distributions through the MPNN architecture. To propagate moments through activation functions, we use the Taylor expansion and the pseudo-Taylor polynomial expansion. We then use the resulting node embedding distributions to analytically derive probabilistic adversarial robustness certificates for node classification tasks against L2-bounded perturbations of node features. Second, we model node features as multivariate random variables and introduce Feature Convolution Distance $FCD_p$, a pseudometric based on the Wasserstein distance. $FCD_p$ corresponds to the discriminative power of MPNNs at the node level. We show that MPNNs are globally Lipschitz continuous functions with respect to the pseudometric $FCD_p$. Using the covering number of the resulting pseudometric space, which is a subset of the Wasserstein space, we derive generalization bounds for MPNNs with uncertainties in node features. Together, these two complementary approaches—moment propagation for adversarial robustness and $FCD_p$ on the subset of the Wasserstein space for generalization—establish a unified theoretical framework that comprehensively addresses MPNN reliability under node feature uncertainty.
APA
Chernikova, A., Laber, M., Sabhahit, N.G. & Eliassi-Rad, T.. (2026). Robustness and Generalization in Uncertainty-Aware Message Passing Neural Networks . Proceedings of The 29th International Conference on Artificial Intelligence and Statistics, in Proceedings of Machine Learning Research 300:3826-3834 Available from https://proceedings.mlr.press/v300/chernikova26a.html.

Related Material