[論文レビュー] Optimal Transport Graph Neural Networks

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

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

ひとこと要約

この論文では、ノード埋め込みと学習可能なプロトタイプ点群の間の Wasserstein 距離を用いてグラフ埋め込みを計算する GNN、OT-GNN を提案する。これにより、ユニバーサル近似が可能になり、分子性質予測が向上する。最適輸送とパラメトリックなプロトタイプ、ノイズ対比正則化子を組み合わせることで、モデルは集約崩壊を回避し、標準的な 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 距離として計算される。
• プロトタイプが単一の点に退化するのを防ぐために、エンドツーエンド学習でノイズ対比正則化子を用いてモデルを訓練する。
• 最終的なグラフ埋め込みは、すべてのプロトタイプとの Wasserstein 距離を連結することで形成され、豊富で幾何学的感度の高い表現が得られる。
• 理論的裏付けとして、OT-GNN が定義する関数クラスが点群上でユニバーサル近似可能であることを証明している。これは、標準的な和ベースの集約とは異なり、そのような性質を有さない。

実験結果

リサーチクエスチョン

• RQ1 学習可能なプロトタイプを用いた最適輸送ベースのグラフ表現学習は、分子性質予測において標準的な GNN 集約を上回るか？
• RQ2 OT と組み合わせたパラメトリックなプロトタイプの使用は、和や平均集約と比較して GNN の表現能力を向上させるか？
• RQ3 エンドツーエンド学習における OT ベースの GNN の最適化課題、特にプロトタイプ崩壊をどのように緩和できるか？
• RQ4 学習されたプロトタイプは解釈可能な分子サブ構造や機能性基団に対応しているか？
• RQ5 得られる埋め込み空間は、標準的な GNN より滑らかでより頑健であるか？

主な発見

• OT-GNN は、GIN や GAT を含む最先端の GNN よりも、複数の分子性質予測ベンチマークで優れた性能を示し、一貫した実験的利得を確認した。
• ノイズ対比正則化子はプロトタイプ崩壊を効果的に防止し、最適輸送の幾何的構造を十分に活用できるようにした。
• 理論的分析により、OT-GNN が定義する関数クラスが点群上でユニバーサル近似可能であることが確認された。これは、標準的な GNN 集約手法とは対照的である。
• ベースラインと比較して、OT-GNN は滑らかなグラフ埋め込み空間を生成しており、一般化性能と頑健性の向上を示している。
• 解釈可能性分析の結果、学習されたプロトタイプは意味のある分子サブ構造に対応していることが示された。例えば、あるプロトタイプは環を含む分子（不溶性）に最も近づき、別のプロトタイプはより溶解性の高い分子を好む。
• ハイパーパramータサーチの結果、プロトタイプ数（M=10,20）やプロトタイプサイズ（N=10）に関わらず性能が安定しており、最適な隠れ次元（d=5,10）および正則化係数（nc_coef=0.1）で最良の性能が達成された。