ICLR 2023 高分论文 | 基于子图草图的图神经网络用于链路预测

科技

2023-08-28 04:08

本文介绍了一项关于链路预测方向的研究，该论文在 openreview 上分别得到了 10, 8, 8, 8 的评分。作者研究了基于子图的用于链路预测（LP, Link Prediction）方法，提出了一个新奇、高效的全图 GNN 模型。

论文标题：

Graph Neural Networks for Link Prediction with Subgraph Sketching

论文链接：

https://arxiv.org/abs/2209.15486

代码链接：

https://github.com/melifluos/subgraph-sketching

研究背景

链路预测问题有着许多实际的应用。比如，推荐系统，药物发现和知识图结构等。GNNs（Graph Neural Networks），尤其是 MPNNs（Message-Passing Neural Networks），在图和结点级的任务表现出色，但是在链路预测问题上表现不佳，甚至不如简单的启发式算法。

在链路预测问题上，MPNNs有一些问题。一是由于消息传递等价于 WL（Weisfeiler-Leman）图同构测试，普通的 MPNNs 没能力三角计数、计数公共的邻居或者计算简单的启发式。二是由于基于 GNN 的链路预测方法结合置换等变结构的结点表征，这使得结点 u 与其他的自同构结点有相同的表征，从而忽略了像结点间的距离等重要特征。

针对上述的问题，一些方法被提出来增强 GNN 的表达力，如增加不同结点 ID，结构计数等。然而，增加结点 ID 有着泛化、训练收敛的问题，结构计数需要预训练，开销大。受到链路预测启发式表现优异的灵感，SGNN（Subgraph GNNs）被提出去学习链路预测的启发式。

但是 SGNN 也面临着一些问题。一方面，构建子图开销大且难以并行；另一方面，由于需要构建子图，每一步推理的开销较大。

分析SGNN

SGNNs 可以分解为四个步骤，提取结点对的子图，用结构特征增强子图结点，通过 GNN 传播特征，使用读出（readout）函数链路预测。受到 SGNN 表现优异的启发，作者对上述部分进行了深入研究。

2.1 结构特征

由于 GNN 无法区分自同构结点，结构特征被构建增强 GNN 的表达能力。作者首先在 Cora、Citeseer 和 Pubmed 数据集上实验分析了 0-1 编码（ZO, Zero-One Encoding）、双半径结点标签编码（DRNL, Double-Radius Node Labeling）和距离编码（Distance Encoding）三种主流的结构特征。

由 Figure 2（a）可以看出结构特征对于链路预测任务是重要的。基于 DRNL 出色的表现，作者进一步分析了不同最大距离 r 下，DRNL 结构特征的重要性（如下图）。作者发现对预测表现贡献最大的是距离近的结构特征。

2.2 传播/GNN

通常在 SGNNs，结构特征与结点特征是一起传播的。如下图所示，作者实验发现如果用 MLP 处理结构特征而不与结点特征一起传播，模型可以表现更好。作者尝试不传播原始结点特征，效果并不如意，但是可以通过预传播原始结点特征恢复模型表现。

2.3 读出/ 池化函数

对于 SGNNs，读出函数可以选择池化子图中所有结点表征后过 MLP 或者使用边池化（通常用哈达玛积）。作者对比了两种读出函数，发现边池化效果要好于图池化，实验结果如 Figure 2（b）所示。

2.4 分析总结

经过上述的实验，作者认为：

结构特征带来很大的提升。
嵌入结构特征并在 SGNN 中传播，这种方式在效率和表现都不是最好的。
大多数重要的结构特征位于最短的距离。
对于读出函数，边等级的池化函数效果要好于图等级的池化函数。

方法介绍

通过上述的实验与分析，作者认为如果必要的结构特征信息可以被编码在节点中，那么就无需为每对关联生成子图。作者设计了一个逐结点子图的草图来估计结构特征的全图 GNN 模型，从而避免构建子图。

3.1 结构特征计数

对于链接，距离结点为，距离结点为的结点个数。为了确保计算开销，作者在大小为的接受域中计算该结点个数。为了减轻信息的丢失，作者计算了距离结点为，距离 v 结点长度大于的结点个数，。Figure 6 展示了在大小为 2 的接受域中，两个结构特征与，结点邻居的关系。

作者进而得到如下关系式：

其中，是与结点距离不超过的邻居个数。

通过上式，作者敏锐的捕捉到了邻居计数中计算冗余，并使用基于统计、哈希的 Hyperloglog 和 MinHash 算法估算和，而不用显示的创建子图。这便是逐结点子图的草图，也解释标题中草图的由来。

3.2 基于哈希高效链路预测

作者进一步提出了 ELPH 模型（ELPH, Efficient Link Prediction with Hashes）。它的特征传播满足标准 MPNN 的范式：

其中，是学习函数，是一个局部置换不变聚合函数，是原始结点特征，和分别是 MinHashing 和 Hyperloglog 单结点的草图。如上述公式，通过 MinHashing 和 Hyperloglog 单结点的草图信息，，被估计，然后将估计的结构特征融合入消息传递中，从而免去了为每个关联创建子图，大大节省了计算开销。

最后，，的关联概率被得到：

其中，，是接受域范围，是一个 MLP，是逐元素积。

3.3 利用预处理加速推理

由于显存容量是有限的，而现实里大规模的数据比比皆是，作者从大规模数据角度进一步增强 ELPH。由前述实验可知，嵌入结构特征并在 SGNN 中传播在效率和表现都不是最好的，而预传播结点特征表现很好。于是，作者提出了 ELPH 大尺度版本 BUDDY，使用了预传播结点特征与预计算结构特征来进一步加速模型推理。

模型表现

作者最后在多个尺度的数据集将 ELPH 和 BUDDY 与启发式算法、传统 GNNs、主流 SGNNs 等方法进行对比。如下图可见，ELPH 在一般的数据集上表现优异，而 BUDDY 几乎超过了所有对比方法。

文章总结

这是一篇非常高质量的图神经方面的论文。首先，针对 GNNs 在链路预测任务表现不佳的现象，作者分析了最先进的链路预测方法，确定每个部分的贡献并得到了有趣的结论。其次，作者敏锐的捕捉到基于子图的邻居计数中的计算冗余，并通过 Hyperloglog 和 MinHash 估计结构特征。然后，作者设计智能、高效的全图 GNN 模型，同时考虑了大尺度数据集的情况。最后，多个数据集上的表现与理论分析都体现了该模型的优异。

合作联系： [email protected]

更多阅读