ICLR 2022:Equivariant Subgraph Aggregation Networks

摘要

消息传递神经网络(MPNN)是在图结构数据上进行深度学习的领先架构，在很大程度上是因为它们的简单性和可扩展性。不幸的是，这些架构的表达能力有限。本文提出一种新的框架，等变子图聚合网络(ESAN)来解决这个问题。主要观察是，虽然两个图可能无法被MPNN区分，但它们往往包含可区分的子图。因此，本文建议将每个图表示为由一些预定义策略派生的一组子图，并使用合适的等变架构来处理它。本文开发了图同构的一维Weisfeiler-Leman (1-WL)测试的新变体，并证明了ESAN在这些新的WL变体方面的表达能力的下界。进一步证明，所提出方法提高了MPNN和更具表现力的架构的表达能力。此外，本文提供了理论结果，描述了子图选择策略和等变神经架构等设计选择如何影响架构的表达能力。为处理增加的计算成本，本文提出一种子图采样方案，可以看作是所提出框架的随机版本。在真实和合成数据集上的全面实验表明，所提出框架提高了流行GNN架构的表达能力和整体性能。

引言

由于其可扩展性和简单性，消息传递神经网络(mpnn)是用于在图结构数据上进行深度学习的领先的图神经网络(GNN)架构。然而，Morris等人(2019);Xu等人(2019)表明，这些架构最多与Weisfeiler- lehman (WL)图同构测试一样具有表现力。因此，mpnn无法区分非常简单的图(见图1)。鉴于这种局限性，自然会出现一个问题:有可能提高mpnn的表达能力吗?

最近的几项工作提出了更强大的架构。其中一种主要方法涉及等价于k-WL测试的高阶gnn，在表现力和时空复杂度之间提供了一种权衡。不幸的是，实现三阶网络(提供了3-WL的表达能力)已经很困难。替代方法使用富含节点和边结构编码的标准mpnn(例如，基于环或团计数)或将图提升为简单或细胞复合物，将消息传递机制扩展到这些高阶结构。这两种方法都需要预计算阶段，尽管这在实践中是合理的，但在最坏的情况下可能会非常昂贵。

为了设计简单、直观和更灵活的可证明表达性的图架构，本文开发了一种新的框架，称为等变子图聚合网络(ESAN)，以增强现有GNN的表达能力。我们的解决方案源于这样的观察:虽然两个图可能无法被MPNNs区分，但可以很容易地找到可区分的子图。更一般地说，不是像MPNNs和WL测试中那样编码节点颜色的多重集，而是选择编码子图的包(多重集)，并表明这种编码可以带来更好的表达能力。根据这种观察，建议将每个图表示为根据某种预定义策略选择的子图包，例如，可以通过从原始图中删除一条边得到的所有图。图1说明了这个想法。

子图包是高度结构化的对象，其对称性既来自于每个组成图的结构，也来自于整体的多重集。本文提出一种等变架构，专门为捕捉该物体的对称群而定制。首先，将一组图的对称群表述为集合和图的对称群的直积。然后，构建一个神经网络，其中包含对该组等变的层。受Maron等人(2020)的启发，这些层采用两个基本图编码器作为子程序:第一个编码器实现孪生网络，独立处理每个子图;第二个模块通过处理子图的聚合，充当信息共享模块。在经过几个这样的层处理后，集合学习模块将获得的子图表示聚合为原始图的不变表示，用于下游任务。

该方法的一个重要组成部分是子图选择策略:一个将图映射到子图包的函数，然后由等变神经网络处理。本文探讨了四种简单但功能强大的子图选择策略:节点删除子图、边删除子图和两种自网络变体。为了减轻可能的计算负担，还引入了该方法的一个有效的随机版本，根据上述策略对子图进行随机采样。

对所提出的方法进行了彻底的理论分析。证明了该架构可以实现著名的WL测试的新的、可证明更强大的变体，能够根据基本图编码器(例如MPNNs的WL)编码子图的多重集。此外，研究了架构的表达能力如何依赖于不同的主要设计选择，如底层的基本图编码器或子图选择策略。证明了该框架可以只用一个1-WL图编码器分离3-WL不可区分的图，并且可以增强PPGN等更强架构的表达能力。

然后，我们使用几个现有的gnn作为基本编码器，在广泛的合成和真实数据集上给出了经验结果。首先，我们使用Abboud等人(2020)引入的合成数据集研究了我们的方法的表达能力，并表明它达到了完美的准确性。然后，我们在流行的图形基准测试上评估ESAN，并显示它们始终优于其基础GNN架构，并且与最先进的方法表现更好或相当。

本文的主要贡献包括:(1)将图的学习概括为子图包的学习;(2)生成和处理子图包的等变框架;(3)对所提出的框架、子图选择策略及其表达能力进行了全面的理论分析;(4)在真实数据集和合成数据集上的实验结果表明，所提方法具有明显的性能提升。相信该方法是开发简单且可证明功能强大的GNN架构的一个进步，并希望它将激励未来的理论和实践研究工作。

Equivariant Subgraph Aggregation Networks (ESAN)

本文的主要贡献包括:(1)将图的学习概括为子图包的学习;(2)生成和处理子图包的等变框架;(3)对所提出的框架、子图选择策略及其表达能力进行了全面的理论分析;(4)在真实数据集和合成数据集上的实验结果表明，所提方法具有明显的性能提升。相信该方法是开发简单且可证明功能强大的GNN架构的一个进步，并希望它将激励未来的理论和实践研究工作。

在本节中，我们将介绍ESAN框架。它包括(1)处理子图包的神经网络架构(DSS-GNN和DS-GNN)和(2)子图选择策略。该方法的主要思想是将图G表示为一个包(多集)$S_G=G_1，\dots ，G_m$，并基于该子集对图进行预测，即$F(G):=F(S_G)$。该方法涉及两个关键问题:(1)应该使用哪种架构来处理图包，即如何定义$F(S_G)$?(2)采用何种子图选择策略，即如何定义$S_G$?

图袋编码器架构

为了解决第一个问题，我们从图袋的自然对称群开始。我们首先基于这一组设计了一个称为DSS-GNN的等变架构，然后提出了一个特别有趣的变体(DS-GNN)，该变体通过禁用DSS-GNN中的一个组件获得。

假设节点数n和子图数m都是固定的，G的子图包 S G = { { G 1 ， … ， G m } } S_G = \{\{G_1，…， G_m\}\} SG​={{G1​，…，Gm​}}可以表示为张量$(A,X)\in R^{n\times n\times m}\times R^{n\times d\times m}$。这里，A表示m个邻接矩阵，X表示m个节点特征矩阵。由于每个子图中的节点的顺序，以及这些子图在$S_G$中的顺序是任意的，因此我们的架构必须对这两个顺序的更改是等变的。由于每个子图中的节点顺序以及$S_G$中这些子图的顺序是任意的，因此我们的体系结构必须对这两种顺序的变化都是等价的。子图对称或节点置换由对称群$S_n$表示，$S_n$通过$(\sigma \dot{A}{)}_{ij}=A_{{\sigma }^{-1}(i){\sigma }^{-1}(j)}$，$(\sigma \dot{X}{)}_{il}=X_{{\sigma }^{-1}(i)l}$作用于图。集合对称性是$S_G$中子图的排列，由对称群$S_m$表示，$S_m$通过$(\tau \cdot A{)}_{ijk}=A_{ij{\tau }^{-1}(k)}$和$(\tau \cdot X{)}_{ilk}=X_{il{\tau }^{-1}(k)}$作用于图的集合。这两组可以用直接积的方法组合成一个单独的组$H=S_m\times S_n$，作用于$(a,X)$的方式如下:

换句话说，$\tau$置换集合中的子图，$\sigma$置换子图中的节点，如图2所示。重要的是，由于所有子图都来自同一个图，因此我们可以对它们的节点进行一致的排序，即所有子图中的第$i$个节点代表原始图中的第i个节点这个设置可以看作是DSS框架应用于图袋的一个特殊实例。

H-等变层. 本文的目标是提出一种H-等变层架构，可以处理子图包，说明其自然对称性(节点和子图排列的乘积)。这种等变体系结构的主要构建块是H-等变层。不幸的是，描述所有等变函数的集合可能是一项艰巨的任务，大多数工作局限于描述线性等变层的空间。不幸的是，描述所有等变函数的集合可能是一项艰巨的任务，大多数工作局限于描述线性等变层的空间。在我们的例子中，Maron等人(2020)用置换对称群P(在我们的例子中是$S_n$)描述了对称元素集(如图)的线性等变层空间。他们的研究表明，每个这样的层都由独立地对每个集合元素应用线性P-等变层的孪生组件和对集合中所有元素的聚合应用不同线性P-等变层的信息共享组件组成。

受线性表征的激励，采用具有孪生和信息共享组件的相同层结构，并使用任何等变GNN层(如MPNN)对这些层进行参数化。这些层$L:R^{n\times n\times m}\times R^{n\times d\times m}\to R^{n\times n\times m}\times R^{n\times d^{\prime }\times m}$将子图包映射到子图包，如下所示:

第一个组件是一个等变特性编码器,$E_{subgraphs}$实现作为几个h - 等变层的组成,图层(方程1)。子图编码器的目的是为所有子图中的所有节点学习有用的节点特性。在图编码器的顶部,我们应用一个子图读出图层$R_{subgraph}$:R,通过聚合图节点(和/或边缘)数据独立生成每个子图的不变特征向量,例如将节点特性求和。最后一层,E$sets$:R,是一个通用设置编码器,例如DeepSets或PointNet。图3(左面板)说明了DSS-GNN架构。

直观地说，DSS-GNN将子图(同时考虑到其他子图)编码为一组$S_n$不变表示($E_{subgraphs}\ast R_{subgraphs}$)，然后将产生的子图表示集编码为原始图$G$的单个$H$不变表示$E_{sets}$集。

DS-GNN是通过禁用信息共享组件，即通过设置L获得的DSS-GNN的变体
在$E_{subgraphs}$中，每一层2 = 0。因此，子图是独立编码的，并且子图实际上是一个暹罗网络。DS-GNN可以解释为一个孪生网络，独立编码每个子图，然后是一个编码子图表示集合的编码器。DS-GNN对于由环积获得的较大对称群是不变的。

在第3节中，我们将表明DS-GNN和DSS-GNN都是强大的架构，并且在某些情况下DSS-GNN具有优越的表达能力。

子图选择策略

第二个问题，如何选择子图?，对我们建筑的表现力有直接的影响。在这里，我们讨论子图选择策略，并提出简单但功能强大的方案，将在第3节中分析。

设G为所有节点数不超过n的图的集合，P(G)为其幂集，即所有子集S ⊆ G的集合。子图选择策略是一个函数π：G → P(G)，它将每个图G映射到其子图集合的一个子集π(G)。我们要求π对于图中节点的排列是不变的，即π(G) = π(σ · G)，其中σ ∈ Sn是一个节点排列，σ · G是应用σ后得到的图。我们将使用以下符号：π(G) := SπG。与之前一样，SπG可以用张量形式表示，其中我们以任意顺序堆叠子图，同时确保子图中节点的顺序是一致的。

在本文中，我们探讨了四种简单的子图选择策略，它们被证明在复杂性(子图的数量)和由此产生的表达能力之间取得了很好的平衡:节点删除子图(ND)，边删除子图(ED)和自我网络(EGO, EGO+)，如下所述。在节点删除策略中，图被映射到包含所有子图的集合，这些子图可以通过删除单个节点从原始图中获得。类似地，边删除策略通过删除一条边来定义。自我网络策略EGO将每个图映射到一组特定深度的自我网络，图中的每个节点都有一个自我网络(节点的k-自我网络是其具有诱导连通性的k跳邻域)。我们还考虑了egonetworks策略的一个变体，其中根节点持有一个识别特征(EGO+)。对于DS-GNN，为了至少保证所使用的基本图编码器的表达能力，对于任意策略π，可以将原始图G添加到$S_G^{\pi }$中，得到增强版本$\hat{\pi }$。例如，$\widehat{EGO}$是一个新策略，它为每个图输出上面描述的标准EGO策略和原始图。

随机抽样 对于较大的图，在某些策略中使用整个子图包的代价可能很高。为抵消这一问题，有时采用随机子图子采样方案，在每个训练epoch中，独立且均匀随机地选择子图SπG⊂SπG。在实践中,我们使用|πG | / |年代πG |∈{0.05,0.2,0.5}。除了能够在更大的图上运行外，子采样还允许更大的批处理大小，从而使每个epoch的运行速度更快。在推理过程中，我们将这些子图的子集(我们使用l = 5)和多数投票方案结合起来进行预测。虽然由此产生的模型不再是不变的，但该方法的这种有效变体增加了基本编码器的表达能力，并在经验上表现良好。

实验分析

在本节中，我们将通过可证明地分离非同构图的能力来研究我们的架构的表达能力。我们首先介绍面向ESAN的WL模拟，然后研究不同的设计选择如何影响我们架构的表达能力。

本文提出两种受WL同构测试启发的颜色细化程序。这些变体称为DSS-WL和DS-WL变体，受到具有1-WL等效基图编码器(例如MPNN)的DSS-GNN和DS-GNN的启发。我们的目标是形式化我们的直觉，即通过正确编码其包含的子图(子集)，图可能更容易分离。所设计的变体旨在提取和利用这种判别式信息源。

DSS-WL测试的输入是两个输入图G1, G2和一个子图选择策略π。该算法应用π在输入图上生成子图包。如果提供初始颜色，则为每个子图中的每个节点分配其在原始图中的原始标签。否则，为每个子图中的每个节点分配一个初始的、恒定的颜色，与包无关。

结论

本文提出了ESAN，一种用于学习图结构数据的新框架，基于将图分解为一组子图，并以有关其对称性结构的方式处理该集合。ESAN提高了几种现有图学习架构的表达能力，并在多个图分类基准上表现良好。该方法的主要限制是其相对于标准mpnn的计算复杂性增加。本文提出了一种更有效的随机版本来缓解这一限制，并在实践中表现良好。在后续研究中，我们计划研究以下方向:(1)是否有可能学习到有用的子图选择策略?(2)子图(例如，子图的图)上的高阶结构可以用于更好的表达能力吗?(3)对不同策略/聚合函数的随机版本和理论分析。