Learning the Neighborhood: Contrast-Free Multimodal Self-Supervised Molecular Graph Pretraining

技術要約：Learning the Neighborhood (C-FREE)

問題提起
高品質な分子表現は、物性予測や創薬において極めて重要であるが、ラベル付きの大量のデータセットは依然として不足している。分子グラフに対する既存の自己教師あり学習（SSL）アプローチには、しばしば以下の特有の制限が存在する：

1. モダリティの制限： 多くの手法は2Dトポロジーまたは1Dシーケンス（SMILES）のみに依存しており、貴重な3D構造情報を十分に活用できていない。
2. 複雑さとコスト： 対照学習（Contrastive methods）は、注意深く設計された負例（negative samples）や大きなバッチサイズに依存するが、これらは不規則なグラフ構造（例：キラル異性体の区別）に対して定義することが困難である。生成的手法（Generative methods）は、入力空間における離散的な再構成や複雑なグラフトークン化を必要とすることが多い。潜在空間予測的手法（Latent predictive methods）は、高価な拡張（augmentation）やクラスタリング手順（例：METIS）に依存することが多く、表現の崩壊（representation collapse）に陥る可能性がある。
3. データの希少性： 精選されたデータセットは、特に低データ領域において利用できないことが多く、新しい化学ドメインへ効果的に転移できる堅牢な事前学習戦略が必要とされる。

手法：C-FREE
著者らは、2Dグラフと3Dコンフォーマーのアンサンブルを統合したマルチモーダル自己教師あり学習フレームワークであるC-FREE（Contrast-Free Representation learning on Ego-nets）を提案する。この手法は、潜在空間における予測学習戦略を採用しており、負例、位置エンコーディング、および入力空間の再構成を回避している。

• コアメカニズム（Ego-nets）： 本フレームワークは、固定半径のk-EgoNets（ここで $k \in \{3, 4\}$）を用いて分子グラフを互いに素な部分グラフに分割する。特定のアンカーノードに対して、$k$ホップの近傍が**コンテキスト（context）部分グラフを形成し、残りのエッジとノードがターゲット（target）**部分グラフを形成する。境界ノードは、エッジの互いに素性を保証するために複製される。
• マルチモーダル・エンコーディング：
  • 2Dエンコーダー： Graph Isomorphism Network with Edge features (GINE) がグラフのトポロジーを処理する。
  • 3Dエンコーダー： 等変ネットワーク（PaiNNまたはSchNet）が複数のコンフォーマーからの3D座標を処理する。
  • 融合（Fusion）： 両方のモダリティからのノードレベルの埋め込みを、学習可能な分類用トークン（hCLS）および分離用トークン（hSEP）と共に結合し、グローバルな依存関係を捉えるためのTransformerによって処理されるマルチモーダル・シーケンスを形成する。
• 予測目的関数： モデルは、コンテキスト部分グラフの埋め込みとターゲット部分グラフの埋め込みを整合させる。
  • コンテキストエンコーダーが入力ビューを処理する。
  • 予測ネットワーク（軽量なTransformerとそれに続くMLP）がコンテキストの埋め込みをターゲット空間へと写像する。
  • ターゲットエンコーダー（コンテキストエンコーダーの指数移動平均、EMA）がターゲットの埋め込みを生成する。
  • 目的関数は、予測されたコンテキストの埋め込みとターゲットの埋め込みの間の平均二乗L2距離を最小化する。
• ファインチューニング： ダウンストリームタスクにおいて、ターゲットエンコーダーは事前学習済みのバックボーンとして機能する。本フレームワークは2つの評価戦略をサポートしている：
  1. C-FREEMOL： 分子全体の埋め込みに対する線形プロービング（Linear probing）。
  2. C-FREESUB： DeepSetsを用いて部分グラフの埋め込みを集約する、Equivariant Subgraph Aggregation Networks (ESAN) をシミュレートした手法。

主な貢献

1. 新規な事前学習タスク： コストのかかる手作りの拡張や負例サンプリングを必要とせず、k-EgoNet部分グラフを利用した、広く適用可能な予測目的関数。
2. 堅牢性： 本フレームワークは、マルチモーダル設定（2D + 3D）と2Dのみの設定の両方で強力な性能を発揮し、コンフォーマーが存在しない場合でも効果的な転移を示す。
3. 単純さと表現力： 本スキームは、事前学習とファインチューニングの不一致や負例を排除している。理論的には、DeepSetsヘッドを使用する場合、C-FREEは1-Weisfeiler-Leman (1-WL) テストよりも表現力が高く、階層的な目的関数やクラスタリングを必要とせずにESANの表現力に匹敵する。
4. 最先端の性能： 本手法はMoleculeNetおよびその他のベンチマークにおいてトップクラスの結果を達成し、対照学習、生成的、およびその他のマルチモーダルSSL手法を凌駕している。

実験結果
モデルはGEOMデータセット（約33万個の3Dコンフォーマーを持つ分子）で事前学習され、複数のベンチマークで評価された：

• MoleculeNet (Frozen Backbone)： C-FREEは8つのデータセット全体で最高の平均性能を達成し、ベースラインと比較して8つのタスクのうち6つで上回った。マルチモーダル・バリアント（C-FREEMM）は一般に2Dのみのバリアントよりも優れた性能を示したが、2DのみのC-FREEであっても、平均して他のすべてのベースラインを上回った。
• MoleculeNet (Full Fine-Tuning)： GEOMで事前学習されたPaiNNエンコーダーを使用した場合、C-FREEは大幅に大きなデータセット（最大1900万分子）で学習された大規模な基盤モデル（例：Uni-Mol, GEM）と同等またはそれを上回る性能を示し、最高の平均ROC-AUCを達成した。
• QM9 および ZINC： QM9において、C-FREEはパラメータ数が大幅に少ない（910万）にもかかわらず、巨大なUni-Mol2（11億パラメータ）に対し4つのターゲットで上回った。ZINCにおいては、GraphJEPAを上回った。
• ラベル効率： Drugs-75Kデータセットにおいて、事前学習は低ラベル領域（1%、10%のラベル）で明確な利点をもたらし、ゼロから学習したモデルを一貫して上回った。
• アブレーション解析：
  • モダリティ： 3D情報は2D単独よりも強い信号を提供したが、両方を組み合わせることで最良の結果が得られた。
  • 予測器（Predictor）： 予測器を除去すると学習の崩壊（collapse）が発生した。Transformer予測器はMLPよりも優れた性能を示した。
  • EMA： EMAターゲットエンコーダーの初期減衰率を下げることで、表現の質が向上した。
  • EgoNet半径： 半径 $k=5$ が最高の性能を示した。これは、同程度のサイズのコンテキストとターゲットの部分グラフが最適であることを示唆しており、$k=1$ では十分な構造的信号を捉えられなかった。

意義と主張
論文は、C-FREEがマルチモーダル分子表現学習における最初の非対照的（non-contrastive）、かつ非生成的（non-generative）な予測フレームワークであると主張している。その意義は以下の点にある：

• 複雑性の排除： グラフSSLにおける一般的なボトルネックである、負例、大きなバッチサイズ、高価なクラスタリング、および位置エンコーディングの必要性を排除した。
• 3D幾何学の活用： 複雑な生成的目的関数を必要とせずに、分子の形状に依存する特性に不可欠な3Dコンフォーマーの多様性を効果的に統合している。
• 汎用性： 本手法は、自己教師あり事前学習が多様な化学ドメインや低データ設定に対して効果的に転移できる強力な初期化を提供できることを実証しており、2Dトポロジーが利用可能であっても3D情報に基づいた表現が重要であることを強調している。

著者らは、自らのモデルが一部の基盤モデルよりも小規模であるものの、その性能は、アーキテクチャの選択と事前学習パイプライン（具体的にはEgoNetに基づく予測タスク）が非常に効果的であることを示唆していると述べている。今後の課題として、より大きなデータセットへのスケーリング、SMILESのような追加のモダリティの組み込み、およびキラル認識エンコーダーの探索などが挙げられている。