[论文解读] Optimal Transport Graph Neural Networks

Benson Chen, Gary Bécigneul+3 more|arXiv (Cornell University)|Jun 8, 2020

Advanced Graph Neural Networks参考文献 61被引用 22

一句话总结

该论文提出OT-GNN，一种图神经网络，通过计算节点嵌入与可学习原型点云之间的Wasserstein距离来生成图嵌入，实现通用逼近并提升分子性质预测性能。通过结合最优传输与参数化原型及噪声对比正则化，该模型避免了聚合坍塌问题，生成的表示比标准GNN更加平滑且信息量更丰富。

ABSTRACT

Current graph neural network (GNN) architectures naively average or sum node embeddings into an aggregated graph representation -- potentially losing structural or semantic information. We here introduce OT-GNN, a model that computes graph embeddings using parametric prototypes that highlight key facets of different graph aspects. Towards this goal, we successfully combine optimal transport (OT) with parametric graph models. Graph representations are obtained from Wasserstein distances between the set of GNN node embeddings and ``prototype'' point clouds as free parameters. We theoretically prove that, unlike traditional sum aggregation, our function class on point clouds satisfies a fundamental universal approximation theorem. Empirically, we address an inherent collapse optimization issue by proposing a noise contrastive regularizer to steer the model towards truly exploiting the OT geometry. Finally, we outperform popular methods on several molecular property prediction tasks, while exhibiting smoother graph representations.

研究动机与目标

• 解决标准GNN通过简单聚合（如求和或平均）节点嵌入导致结构与语义信息丢失的局限性。
• 开发一种参数化图模型，利用最优传输（Wasserstein距离）将输入图与抽象、可学习的原型进行比较。
• 克服基于OT的模型在端到端训练中的优化挑战，特别是原型坍塌问题（即原型退化为单一点）。
• 为所提出的函数类相较于标准GNN聚合方法在表示能力上的提升提供理论依据。
• 在分子性质预测基准上实证验证其优越性能，同时生成更平滑、更具可解释性的图表示。

提出的方法

• 模型使用标准GNN（如GIN或GAT）将输入图编码为节点嵌入，形成嵌入空间中的点云。
• 引入可学习的原型点云作为自由参数，代表与下游任务相关的抽象结构或功能特征（如官能团）。
• 通过输入图的节点嵌入点云与每个原型点云之间的Wasserstein距离计算图级表示。
• 采用噪声对比正则化进行端到端训练，防止原型坍塌，从而保留OT的几何结构。
• 最终图嵌入通过拼接所有原型的Wasserstein距离得到，实现丰富且具备几何感知能力的表示。
• 理论上证明，由OT-GNN定义的函数类在点云上是通用逼近器，而标准GNN聚合方法不具备此性质。

实验结果

研究问题

• RQ1 基于最优传输与可学习原型的图表示学习是否能在分子性质预测中超越标准GNN的聚合方法？
• RQ2 结合OT与参数化原型是否能提升GNN的表示能力，相较于求和或平均聚合？
• RQ3 在基于OT的GNN端到端训练中，如何缓解优化挑战，特别是原型坍塌问题？
• RQ4 所学习的原型是否对应于可解释的分子亚结构或官能团？
• RQ5 所得嵌入空间是否比标准GNN更平滑且更具鲁棒性？

主要发现

• OT-GNN在多个分子性质预测基准上优于最先进GNN模型（包括GIN和GAT），展现出一致的性能提升。
• 噪声对比正则化有效防止原型坍塌，使模型能够充分利用最优传输的完整几何结构。
• 理论分析证实，OT-GNN定义的函数类在点云上是通用逼近器，而标准GNN聚合方法不具备此性质。
• 与基线模型相比，该模型生成的图嵌入空间更平滑，表明其泛化能力与鲁棒性得到提升。
• 可解释性分析表明，所学习的原型对应于有意义的分子亚结构：例如，一个原型最接近含环结构的分子（不溶性），而另一个则更倾向于可溶性分子。
• 超参数搜索表明，性能在不同原型数量（M=10,20）和原型大小（N=10）下保持稳定，最优性能在特定隐藏维数（d=5,10）和正则化系数（nc_coef=0.1）下实现。