[论文解读] Probabilistically Rewired Message-Passing Neural Networks

Chendi Qian, Andrei Manolache+5 more|arXiv (Cornell University)|Oct 3, 2023

Advanced Graph Neural Networks被引用 4

一句话总结

该论文提出概率重布线消息传递神经网络（PR-MPNNs），一种可微分、原理严谨的框架，通过精确、可微分的k-子集采样，基于与任务相关的先验知识学习在图中添加或删除边。该方法增强了模型的表达能力，缓解了过压缩（over-squashing）和欠传播（under-reaching）问题，并在分子和引用图基准测试中实现了最先进或具有竞争力的性能，优于标准MPNN和图Transformer在关键设置下的表现。

ABSTRACT

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.

研究动机与目标

• 为解决消息传递GNN的局限性，如因固定图结构导致的过压缩和表达能力有限。
• 开发一种原理严谨、可微分的方法，用于学习与预测任务相适应的图结构修改。
• 克服启发式重布线方法的不足，这些方法不适应数据分布，且无理由地引入随机边。
• 理论上分析概率重布线如何增强模型表达能力，并识别其优于随机重布线的条件。
• 通过实证验证，PR-MPNN在真实世界数据集上实现优于或与图Transformer和最先进MPNN相当的性能。

提出的方法

• PR-MPNNs使用上游模型学习候选边的概率分布，该分布参数化为以恰好k条边为条件的k-子集分布。
• 该框架从该学习到的概率分布中采样多个k条边的邻接矩阵，然后通过聚合（例如减法）形成最终的、可微分的图结构。
• 通过k-子集分布的精确、可微分边际计算反向传播梯度，从而降低训练过程中的偏差和方差。
• 该方法利用近期在精确且可微分的k-子集采样方面的进展，实现边选择机制的端到端训练。
• 最终的图结构被用作下游MPNN的输入以进行预测，从而实现结构与表征学习的联合优化。
• 该方法设计为可扩展且可微分，支持通过随机边选择过程进行有效的反向传播。

实验结果

研究问题

• RQ1 与固定图或随机重布线图相比，概率重布线在何种条件下能增强消息传递GNN的表达能力？
• RQ2 可微分的k-子集采样机制如何实现边选择策略的有效且低方差训练？
• RQ3 PR-MPNN在多大程度上缓解了图神经网络中的过压缩和欠传播问题？
• RQ4 PR-MPNN能否在真实世界图基准上实现与图Transformer或最先进MPNN相当或更优的性能？
• RQ5 所学习的边选择策略如何适应底层数据分布和预测任务？

主要发现

• 在LRGB基准的Peptides-struct数据集上，PR-MPNN取得了最佳的平均绝对误差（MAE）得分0.2477 ± 0.0005，优于DRew-GIN（0.2799 ± 0.0016）及其他基线模型。
• 在异配性（heterophilic）和归纳性（transductive）的WebKB数据集上，PR-MPNN在Cornell（0.659 ± 0.040）、Texas（0.827 ± 0.032）和Wisconsin（0.750 ± 0.015）上均达到最高准确率，超越所有MPNN基线模型，并与图Transformer表现相当。
• PR-MPNN通过学习添加长程边以改善信息流动，显著缓解了过压缩和欠传播问题，这一效果在需要全局结构的任务中表现得尤为明显。
• 在间隙属性预测任务中，PR-MPNN相较基础GIN模型实现了1.98倍的平均绝对误差（MAE）改进（1.70 ± 0.01 vs. 3.37 ± 0.03）。
• 在ZPVE属性上，PR-MPNN相较基础GIN模型实现了14.13倍的改进（4.73 ± 0.08 vs. 66.87 ± 1.45），充分证明了其在缓解信息压缩方面的强大能力。
• 理论分析证实，在学习边分布及其与任务相关结构对齐的精确条件下，PR-MPNN的表达能力可超过标准MPNN。