从池化到子图GNN：更具表现力的图神经网络的通用框架
1. 基本信息
论文题目: From Relational Pooling to Subgraph GNNs: A Universal Framework for More Expressive Graph Neural Networks
作者:Cai Zhou, Xiyuan Wang, Muhan Zhang
机构:1. Department of Automation, Tsinghua University 2. Institute for Artificial Intelligence, Peking University
2. 摘要
本文从关系池化(Relational Pooling)框架出发,提出了一种显式为节点赋予标签的方法,以提高基于消息传递的神经网络区分非同构图的能力。该方法进一步扩展到更高维的WL层次,导出了一种新的k,l-WL算法和框架,其规范性更强,能够统一许多现有算法和GNN模型。文章还引入了子图概念,提出了局部化的k,l-WL框架,将一系列现有的子图GNN方法融入其中。从理论上分析了k,l-WL关于k和l的表达能力,并与传统k-WL进行了比较。最后,文章讨论了如何构建实用且高效的k,l-GNN模型,并在多个合成和真实数据集上进行了验证,证明了该框架能够普适地提高任意基础GNN模型的表达能力。
3. 介绍
图神经网络(GNN)在许多领域取得了巨大成功,但其区分非同构图的表示能力受到一维Weisfeiler-Lehman测试(1-WL)的限制。因此,设计表达能力更强的GNN成为研究热点。关系池化(RP)框架通过对基础模型的输出进行对称化来构建更强大的GNN,但计算开销巨大。为改进RP,本文提出了一种给节点明确赋予标签的方法,以提高基于消息传递的神经网络区分非同构图的能力。基于该思想,文中提出了k,l-WL算法和框架,其将k维WL测试应用于赋予l个标签的图上。k,l-WL框架不仅统一了关系池化和许多子图GNN方法,而且拥有比k-WL层次更强的表达能力。文章还讨论了如何构建实用的k,l-GNN模型,并在合成和真实数据集上验证了该框架能够普适地提升任意基础GNN模型的表达能力和性能。
本文的主要贡献有:
- 提出了ID-MPNN模型,改进了关系池化方法。
- 在ID-MPNN的基础上,提出通用的k,l-WL框架,统一了众多算法。
- 从理论上分析了k,l-WL的表达能力,建立了完整的层次结构。
- 讨论了k,l-GNN的实际问题,提出了多种技巧改进可扩展性。
- 在合成和真实数据集上验证了方法的有效性。
4. 方法
4.1 基于标签的消息传递:引入非对称性的增强
关系池化框架中,MLP和RNN等作为基模型具有强表达能力,但对图数据缺乏归纳偏置,且只能编码诱导子图,失去全局信息。为解决这一问题,文章提出了一种为节点明确赋予标签的消息传递方法。
给定输入图
,l-IDMPNN首先为l个节点赋予唯一标签
作为额外特征,然后在标记后的整个图上运行标准消息传递,最后对这些标记图的表示进行聚合以产生原图的表示。因此,IDMPNN既引入了表达能力更强的非对称性,又保留了消息传递的归纳偏置,并可编码全局图信息。
4.2 k,l-WL:高维度的增强
基于ID-MPNN,如果用更强大的k-WL代替消息传递,会产生什么效果?文章给出肯定回答:如果图中赋予l个标签,则运行k-WL(k≥3)会比1-WL(带l个标签)更强大。于是文章提出了k,l-WL算法:对赋予l个标签的图运行k维WL测试。
k,l-WL的关键区别在于初始化。k-WL根据k元组的同构类型确定初始颜色,使得初始颜色很有限。而k,l-WL在赋予l个标签后,根据k元组在带标签图中的同构类型确定初始颜色,使得初始颜色更丰富,增强了后续的表达能力。随后,k,l-WL的迭代更新过程与k-WL相同。
4.3 统一现有层次结构
k,l-WL框架统一了关系池化、原始k-WL和许多子图GNN。当l=0时,k,l-WL等价于k-WL;当k=1时,k,l-WL包含l-OSAN等算法;1,1-WL等价于一种子图WL方法,因此涵盖众多子图GNN。k,l-WL也包含一些超出子图WL范围的子图GNN,如I2-GNN。总之,k,l-WL建立了一个更完备的表达能力层次结构。
5. 实验发现
论文在合成和真实数据集上进行了实验。主要发现如下:
- 在两组合成数据上,l-IDMPNN和l-IDPPGN(l≥2)实现了完美的图同构区分,验证了理论分析。
- 在子结构计数任务上,IDMPNN的误差接近0,证明了模型能完全计数l节点内的子结构。
- 在3个真实分子数据集上,IDMPNN都取得了优于多种基线的结果,表明该框架能普适地提升性能。
6. 结论
本文提出k,l-WL框架,通过给部分节点赋予标签并运行k维WL,构建出更一般的层次结构,统一了众多关系池化和子图GNN方法。论文还讨论了k,l-GNN的实现问题,设计了实用的模型结构。通过理论分析和实验验证,证明了该框架能普适地提高任意基础GNN模型的表达能力。
