Message-passing graph neural networks (MPNNs) emerged as powerful tools for processing graph-structured input. However, they operate on a fixed input graph structure, ignoring potential noise and missing information. Furthermore, their local aggregation mechanism can lead to problems such as over-squashing and limited expressive power in capturing relevant graph structures. Existing solutions to these challenges have primarily relied on heuristic methods, often disregarding the underlying data distribution. Hence, devising principled approaches for learning to infer graph structures relevant to the given prediction task remains an open challenge. In this work, leveraging recent progress in exact and differentiable $k$-subset sampling, we devise probabilistically rewired MPNNs (PR-MPNNs), which learn to add relevant edges while omitting less beneficial ones. For the first time, our theoretical analysis explores how PR-MPNNs enhance expressive power, and we identify precise conditions under which they outperform purely randomized approaches. Empirically, we demonstrate that our approach effectively mitigates issues like over-squashing and under-reaching. In addition, on established real-world datasets, our method exhibits competitive or superior predictive performance compared to traditional MPNN models and recent graph transformer architectures.

翻译：消息传递图神经网络 (MPNN) 已成为处理图结构输入数据的强大工具。然而，它们通常在固定的输入图结构上运行，忽略了潜在的噪声和缺失信息。此外，其局部聚合机制可能导致过压缩及在捕获相关图结构方面表达能力受限等问题。针对这些挑战的现有解决方案主要依赖启发式方法，往往忽略了底层数据分布。因此，设计原则性方法来学习推断与给定预测任务相关的图结构，仍然是一个开放挑战。在本工作中，我们借助可微精确 $k$-子集采样的最新进展，设计了概率性重连的MPNN (PR-MPNN)，该网络能够学习添加相关边并去除收益较低的边。我们首次从理论上分析了PR-MPNN如何增强表达能力，并确定了其优于纯随机方法的精确条件。实验表明，我们的方法有效缓解了过压缩和欠覆盖等问题。此外，在已建立的真实世界数据集上，与传统的MPNN模型及最近的图变换器架构相比，我们的方法展现出具有竞争力或更优的预测性能。