[論文レビュー] Learning Attributed Graph Representations with Communicative Message Passing Transformer

Jianwen Chen, Shuangjia Zheng+3 more|arXiv (Cornell University)|Jul 19, 2021

Computational Drug Discovery Methods参考文献 17被引用数 10

ひとこと要約

本稿では、トランスフォーマーに基づくアーキテクチャを用いて、ノードとエッジ間の双方向的メッセージパッシングを明示的にモデル化することで、分子グラフ表現を向上させるCoMPT（Communicative Message Passing Transformer）を提案する。トポロジカルなメッセージ拡散機構を導入することでメッセージ爆発と過剰平滑化を緩和し、SOTA（最先端）の性能を達成。グラフレベルのタスクでは平均して約4%の性能向上を達成し、1H-NMRのノードレベル予測では0.214 MAEを達成した。

ABSTRACT

Constructing appropriate representations of molecules lies at the core of numerous tasks such as material science, chemistry and drug designs. Recent researches abstract molecules as attributed graphs and employ graph neural networks (GNN) for molecular representation learning, which have made remarkable achievements in molecular graph modeling. Albeit powerful, current models either are based on local aggregation operations and thus miss higher-order graph properties or focus on only node information without fully using the edge information. For this sake, we propose a Communicative Message Passing Transformer (CoMPT) neural network to improve the molecular graph representation by reinforcing message interactions between nodes and edges based on the Transformer architecture. Unlike the previous transformer-style GNNs that treat molecules as fully connected graphs, we introduce a message diffusion mechanism to leverage the graph connectivity inductive bias and reduce the message enrichment explosion. Extensive experiments demonstrated that the proposed model obtained superior performances (around 4$\%$ on average) against state-of-the-art baselines on seven chemical property datasets (graph-level tasks) and two chemical shift datasets (node-level tasks). Further visualization studies also indicated a better representation capacity achieved by our model.

研究の動機と目的

• 既存のGNNが長距離依存関係を捉えることや、分子グラフにおけるエッジ情報の活用を十分に行えないという限界を解消すること。
• 深層GNNおよびトランスフォーマー風GNNにおける過剰平滑化とメッセージ豊かさの爆発問題を克服すること。
• ノードとエッジ特徴量の間の相互作用を明示的にモデル化する通信的メッセージパッシング機構を用いて、分子表現学習を向上させること。
• 分子のトポロジーを尊重する拡散機構を統合することで構造的インダクティブバイアスを導入し、一般化性能と表現力の両方を向上させること。

提案手法

• ノードとエッジの埋め込みが、ノードとエッジ間のクロスアテンションメカニズムを用いて反復的に更新されるトランスフォーマーに基づくアーキテクチャを採用する。
• 分子グラフのトポロジカルな隣接行列を用いて、メッセージ伝搬を制約し、爆発的なメッセージ蓄積を低減するためのメッセージ拡散機構を導入する。
• 構造的文脈を保持し、表現品質を向上させるために、学習可能なノード位置およびエッジ方向埋め込みを用いる。
• ノードとエッジが双方向的に情報を交換する通信的メッセージパッシングを実行することで、特徴量の表現力を強化する。
• 学習可能なリードアウト機構とノード単位の回帰ヘッドを備えることで、グラフレベルおよびノードレベルの予測タスクを両立する。
• タスクに応じて交差エントロピー損失または平均絶対誤差（MAE）損失を用いて、エンドツーエンドでモデルを訓練する。

実験結果

リサーチクエスチョン

• RQ1 ノード-エッジ相互作用を明示的にモデル化するトランスフォーマー基盤GNNが、標準的なGNNやバニラトランスフォーマーを上回る性能を示せるか？
• RQ2 トポロジカルなメッセージ拡散機構を統合することで、分子グラフタスクにおけるモデルの性能と安定性にどのような影響を与えるか？
• RQ3 CoMPTは深層アーキテクチャにおいて過剰平滑化をどの程度軽減し、識別可能な原子表現を保持できるか？
• RQ4 ノード位置、エッジ方向、メッセージ拡散が、それぞれ最終的な性能にどの程度寄与しているか？

主な発見

• CoMPTは、7つのグラフレベル分子性質予測データセットにおいて、最先端のベースラインと比較して平均して約4%の性能向上を達成した。
• ノードレベルタスクでは、1H-NMRで平均絶対誤差（MAE）が0.214 ppm、13C-NMRで1.321 ppmを達成し、すべてのベースラインを上回った。
• アブレーションスタディの結果、メッセージ拡散機構を削除すると最も顕著な性能低下が見られ、これがメッセージ爆発の制御において極めて重要な役割を果たしていることが示された。
• t-SNE可視化により、CoMPTが過剰平滑化を緩和し、スケルトン原子と明確に区別できるサイドチェーン官能基が分離されることを確認した。
• ノード位置およびエッジ方向埋め込みの導入がさらに性能向上に寄与したため、構造的認識の重要性が裏付けられた。
• 複数のランダムシードで繰り返し実験を実施した結果、モデルは高いロバストネスを示し、標準偏差が非常に低かった（例：1H-NMR MAEで0.003）。