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.

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