最优传输图神经网络

当前的图神经网络（GNN）架构通常通过简单地对节点嵌入进行平均或求和来聚合得到图的表示，这一过程可能损失图结构或语义信息。本文提出一种新型模型——OT-GNN，该模型利用可学习的参数化原型（parametric prototypes），突出不同图特征的关键方面，从而实现更丰富的图嵌入。为实现这一目标，我们成功地将最优传输（Optimal Transport, OT）理论与参数化图模型相结合。图表示通过计算GNN节点嵌入集合与作为自由参数的“原型”点云之间的Wasserstein距离来获得。理论上，我们证明了与传统求和聚合方式不同，本模型在点云上的函数类满足一个基本的通用逼近定理。在实践层面，我们提出一种噪声对比正则化（noise contrastive regularizer），有效缓解了模型优化过程中固有的坍缩问题，引导模型真正利用最优传输的几何结构。实验结果表明，OT-GNN在多个分子性质预测任务中优于现有主流方法，同时展现出更平滑的图表示。