[論文レビュー] Motif-aware Attribute Masking for Molecular Graph Pre-training

Eric Inae, Gang Liu+1 more|arXiv (Cornell University)|Sep 8, 2023

Machine Learning in Materials ScienceMaterials Science被引用数 3

ひとこと要約

本稿では、化学的に意味のあるモチーフ内に存在するノード特徴量をすべてマスクすることで、モチーフ間の知識伝達を強化する、分子グラフ事前学習のためのモチーフに配慮した属性マスキング戦略であるMoAMaを提案する。モデルがクロスモチーフ情報を使って特徴量を再構築するよう強制することで、8つの分子性質予測データセットにおいて、最先端のベースラインと比較して平均1.3%の予測性能向上を達成した。

ABSTRACT

Attribute reconstruction is used to predict node or edge features in the pre-training of graph neural networks. Given a large number of molecules, they learn to capture structural knowledge, which is transferable for various downstream property prediction tasks and vital in chemistry, biomedicine, and material science. Previous strategies that randomly select nodes to do attribute masking leverage the information of local neighbors However, the over-reliance of these neighbors inhibits the model's ability to learn from higher-level substructures. For example, the model would learn little from predicting three carbon atoms in a benzene ring based on the other three but could learn more from the inter-connections between the functional groups, or called chemical motifs. In this work, we propose and investigate motif-aware attribute masking strategies to capture inter-motif structures by leveraging the information of atoms in neighboring motifs. Once each graph is decomposed into disjoint motifs, the features for every node within a sample motif are masked. The graph decoder then predicts the masked features of each node within the motif for reconstruction. We evaluate our approach on eight molecular property prediction datasets and demonstrate its advantages.

研究の動機と目的

• グラフニューラルネットワークの事前学習において、ランダムな属性マスキングの限界を解消する。これは局所的な隣接ノード特徴量に過度に依存し、モチーフ間の構造的知識を捉えられないためである。
• 事前学習中に化学的モチーフ間で知識を明示的に伝達することで、分子性質予測の性能を向上させる。
• モチーフ内特徴量伝搬に依存することを減らし、GNNにおけるメッセージパッシングのボトルネックを緩和する。
• モチーフ内構造パターンとモチーフ間の依存関係の両方を活用する事前学習戦略を設計する。
• 多様な分子性質予測ベンチマークにおいて、モチーフに配慮したマスキングの有効性を評価する。

提案手法

• 各分子グラフを事前に定義されたモチーフ語彙を用いて互いに素な化学的モチーフに分解する。
• 各モチーフ内のすべてのノード特徴量を学習可能なマスクトークンでマスクし、モデルがクロスモチーフ情報を使ってそれらを再構築するよう強制する。
• グラフデコーダーが、隣接モチーフからの文脈とモチーフ内部構造に基づいてマスクされたノード特徴量を予測する。
• マスクされた属性をモチーフ間で再構築するためのマスクオートエンコーダーの目的関数を用いて、モデルを事前学習する。
• MRRベースの指標を用いてモチーフ間の知識伝達を定量化し、低いスコアがより強いモチーフ間影響を示す。
• 主にAUCを指標として用い、8つの分子性質予測データセットでアプローチを評価する。

Figure 1: Our MoAMa masks every node in a motif to pre-train GNNs. The full masking of a motif forces the GNNs to learn to (1) pass feature information across motifs and (2) pass local structural information within the motif. Compared to the traditional random attribute masking strategies, the motif

実験結果

リサーチクエスチョン

• RQ1 ランダムマスキングと比較して、モチーフに配慮した属性マスキングは、分子グラフ事前学習におけるモチーフ間の知識伝達を向上させることができるか？
• RQ2 モチーフ内に存在するすべての特徴量を明示的にマスキングすることで、局所的隣接ノード特徴量への依存が軽減され、メッセージパッシングのボトルネックが緩和されるか？
• RQ3 MoAMaの事前学習戦略は、下流の分子性質予測性能をどの程度向上させるか？
• RQ4 さまざまなモデルにおいて、モチーフ間の知識伝達と予測性能の相関関係はどのように現れるか？
• RQ5 Tanimoto類似度やドメイン固有のヒューリスティクスに依存せずに、最先端の事前学習手法と同等の結果を達成できるか？

主な発見

• MoAMaは8つの分子性質予測データセットにおいて平均テストAUCが75.3を達成し、最良のベースラインを1.3%上回った。
• モデルは最高のモチーフ間影響比（ノードレベルで0.80、グラフレベルで0.51）を示し、モチーフ間での知識伝達が強いことを示した。
• MoAMaは最低のMRRスコア（ノードで0.59、グラフで0.55、モチーフで0.41）を記録し、モチーフ内特徴量への依存が低減され、モチーフ間の知識伝達が強化されていることを確認した。
• 予測AUCとモチーフ間影響比の間に中程度の正の相関（r = 0.52–0.55）が観察され、モチーフ間学習と性能向上の間の関連性が裏付けられた。
• AUCとMRRスコアの間に中程度の負の相関（r = -0.39 から -0.44）が確認され、モチーフ内依存の低減が下流の精度向上に寄与していることがさらに裏付けられた。
• 本手法は、モチーフ内構造とモチーフ間の依存関係を明示的に符号化することで、訓練の崩壊を効果的に緩和し、モデルの表現力の向上を実現した。

Figure 2: Inter-motif knowledge transfer score by motif count. A higher $\text{MRR}_{\text{inter}}^{(n)}$ score denotes greater inter-motif knowledge transfer.