[论文解读] Stochastic Aggregation in Graph Neural Networks

Yuanqing Wang, Theofanis Karaletsos|arXiv (Cornell University)|Feb 25, 2021

Advanced Graph Neural Networks参考文献 24被引用 7

一句话总结

本文提出Stochastic Aggregation（STAG），一种统一框架，在图神经网络（GNNs）的消息传递过程中向边权重注入自适应噪声，从而提升模型表达能力并缓解过平滑问题。STAG在引文网络和分子图基准测试中均表现出色，基于变分推断的变体在计算开销极小的情况下实现了最先进性能。

ABSTRACT

Graph neural networks (GNNs) manifest pathologies including over-smoothing and limited discriminating power as a result of suboptimally expressive aggregating mechanisms. We herein present a unifying framework for stochastic aggregation (STAG) in GNNs, where noise is (adaptively) injected into the aggregation process from the neighborhood to form node embeddings. We provide theoretical arguments that STAG models, with little overhead, remedy both of the aforementioned problems. In addition to fixed-noise models, we also propose probabilistic versions of STAG models and a variational inference framework to learn the noise posterior. We conduct illustrative experiments clearly targeting oversmoothing and multiset aggregation limitations. Furthermore, STAG enhances general performance of GNNs demonstrated by competitive performance in common citation and molecule graph benchmark datasets.

研究动机与目标

• 解决由确定性聚合机制引起的GNN模型过平滑和表达能力有限的问题。
• 将现有的正则化技术（如Dropout、DropEdge和Graph DropConnect）统一到单一的随机框架下。
• 开发一种基于变分推断的STAG变体，其中噪声参数与模型参数联合学习。
• 在标准GNN基准数据集（包括引文网络和分子图）上实证验证STAG的有效性。
• 证明边权重中的随机性可提升泛化能力与模型鲁棒性，且计算成本无明显增加。

提出的方法

• STAG在每次消息传递步骤中，使用噪声分布对边权重进行随机采样，替代原有的确定性聚合机制。
• 通过使用连续分布（如Normal(1,1)）或离散分布（如Bernoulli）对边权重进行扰动，实现对邻域消息的随机重加权。
• 在训练过程中，通过蒙特卡洛采样估计梯度，支持无偏的端到端反向传播。
• 引入变分推断（VI）框架，将噪声分布的参数作为变分参数进行学习，可基于节点特征和图结构进行条件化。
• 在推理阶段，通过边缘权重分布的积分得到预测后验分布。
• 该框架在PyTorch和DGL中高效实现，仅需对消息传递操作进行一行代码修改。

实验结果

研究问题

• RQ1 与确定性聚合相比，GNN中对边权重进行随机扰动能否提升模型表达能力并减少过平滑？
• RQ2 STAG在性能和泛化能力方面与现有正则化方法（如Dropout、DropEdge和Graph DropConnect）相比如何？
• RQ3 基于变分推断的噪声参数学习框架能否在多种图结构化任务中提升GNN性能？
• RQ4 在未应用归一化的情况下，连续噪声分布（如Normal）是否在STAG中优于离散分布（如Bernoulli）？
• RQ5 STAG能否在不同类型的图（如引文网络和分子图）上实现一致的性能提升，展现良好的泛化能力？

主要发现

• 在Cora和Citeseer数据集上，使用连续噪声分布（如Normal(1,1)）的STAG优于确定性基线和离散噪声变体，且在多次运行中表现一致。
• 基于变分推断的STAG_VI变体在引文网络和分子图基准上，始终优于非自适应STAG和Graph DropConnect。
• 最具表达力的STAG_VI变体（学习边与特征相关的噪声参数）在ESOL和FreeSolv数据集上实现了最先进性能。
• 在V100 GPU上，STAG每轮前向传播仅增加5.9至9.3毫秒的训练时间，证明其轻量化特性。
• 当与连续噪声分布结合时，归一化操作会降低性能，表明STAG的噪声注入机制本身具有内在稳定性，无需此类校正。
• 该框架在不同图类型上泛化良好，在社交网络和分子图数据集上均表现出一致的性能提升。