[论文解读] Wasserstein Embedding for Graph Learning

Soheil Kolouri, Navid Naderializadeh+2 more|arXiv (Cornell University)|Jun 16, 2020

Advanced Graph Neural Networks参考文献 64被引用 4

一句话总结

WEGL 通过使用节点嵌入之间的 Wasserstein 距离来构建固定大小的向量表示，提出了一种线性且计算高效的图嵌入框架，实现了图级别的预测，准确率达到当前最先进水平，且在图的数量上具有线性复杂度，而无需像以往方法那样进行二次方级别的成对比较。

ABSTRACT

We present Wasserstein Embedding for Graph Learning (WEGL), a novel and fast framework for embedding entire graphs in a vector space, in which various machine learning models are applicable for graph-level prediction tasks. We leverage new insights on defining similarity between graphs as a function of the similarity between their node embedding distributions. Specifically, we use the Wasserstein distance to measure the dissimilarity between node embeddings of different graphs. Unlike prior work, we avoid pairwise calculation of distances between graphs and reduce the computational complexity from quadratic to linear in the number of graphs. WEGL calculates Monge maps from a reference distribution to each node embedding and, based on these maps, creates a fixed-sized vector representation of the graph. We evaluate our new graph embedding approach on various benchmark graph-property prediction tasks, showing state-of-the-art classification performance while having superior computational efficiency. The code is available at https://github.com/navid-naderi/WEGL.

研究动机与目标

• 解决图核方法中因需要计算所有图对之间的成对相似性而导致的高计算成本问题。
• 克服图神经网络（GNNs）和基于核的方法在大规模图数据集上的可扩展性限制。
• 开发一种图嵌入方法，支持任意下游机器学习模型，同时通过 Wasserstein 距离保持度量保真度。
• 将基于 Wasserstein 的图相似性估计的计算复杂度从二次方降低到线性，以应对图的数量。
• 实现在大规模数据集上图属性预测任务的高效训练与推理。

提出的方法

• WEGL 从参考节点嵌入分布到每个图的节点嵌入计算 Monge 映射，以推导出传输计划。
• 它利用最优传输映射将每个图嵌入到希尔伯特空间中，从而生成固定大小的向量表示。
• 嵌入后图之间的欧氏距离近似于其 2-Wasserstein 距离，从而保持了度量结构。
• 该方法利用线性最优传输，避免了所有图对之间的成对距离计算，将复杂度从 O(M²) 降低到 O(M)。
• 节点嵌入用作计算 Wasserstein 距离的输入分布，参考分布则来自图之间池化或学习得到的分布。
• 由于其显式的希尔伯特空间嵌入，该框架支持任意下游分类器，包括 SVM、随机森林和梯度提升树。

实验结果

研究问题

• RQ1 我们能否在降低图嵌入计算复杂度的同时，实现最先进或接近最先进水平的图分类性能？
• RQ2 我们能否将图嵌入到希尔伯特空间中，使得嵌入后图之间的欧氏距离近似于其节点嵌入分布之间的 2-Wasserstein 距离？
• RQ3 基于线性复杂度的 Wasserstein 图相似性方法是否在可扩展性和性能上优于二次复杂度方法？
• RQ4 所提出的嵌入框架是否可与任意标准机器学习分类器配合使用，而不仅限于核方法？
• RQ5 在大规模图数据集上，WEGL 与 GNN 和图核基线方法相比，在效率和准确性方面表现如何？

主要发现

• WEGL 在多个基准数据集上实现了最先进或具有竞争力的分类准确率，包括 ogbg-molhiv 和 TUD 图分类任务。
• 在 IMDB-BINARY 数据集上，WEGL 搭配 GBDT 达到了 75.2% 的准确率，优于 GIN 和其他基线方法。
• 在 PROTEINS 数据集上，WEGL 搭配 GBDT 达到了 92.9% 的准确率，显著优于次佳方法。
• 在图数量庞大的数据集上，WEGL 的训练时间比 WWL 和 GIN 快几个数量级，这主要得益于其线性复杂度。
• WEGL 的推理速度优于基于 CPU 的 GIN，且远快于 WWL，即使 GIN 使用 GPU 加速也是如此。
• 该方法在多种数据集上均保持高性能，包括分子性质预测和社会网络分析任务，且计算开销极低。