Self-Supervised Pretraining of Molecular Graph Encoders with LeJEPA

技術要約：LeJEPAによる分子グラフエンコーダの自己教師あり事前学習

問題提起
自己教師あり事前学習は言語や視覚モデルに革命をもたらしたが、その分子グラフグラフニューラルネットワーク（GNN）に対する有用性は議論の分かれるところである。化学分野では、ラベル付きデータは希少かつ高価である一方、合成可能な分子の空間は膨大である。他の領域で成功している従来の「事前学習・後微調整（pretrain-then-finetune）」パラダイムは、分子特性予測において混合した結果をもたらしている。既存の自己教師あり目的関数（対照学習、生成、および予測的アプローチ）は、手作業で作られた記述子を用いた高度にチューニングされた教師ありベースラインに対して、堅牢なダウンストリームの利点を提供することに苦戦することが多い。本研究では、特定のアーキテクチャであるLeJEPAを用いて、大規模なラベルなしコーパス上で分子グラフエンコーダを事前学習することが、ダウンストリームの特性予測を改善するかどうか、また、その恩恵がどのように現れるのかを調査する。

手法
著者らは、LeJEPA（Latent-Euclidean Joint Embedding Predictive Architecture）を分子グラフに適応させた。従来のJEPAが非対称アーキテクチャ（モメンタム・ティーチャー、ストップ勾配、および予測ネットワーク）に依存しているのに対し、LeJEPAは**Sketched Isotropic Gaussian Regularisation (SIGReg)**によって正則化された、予測器フリーのデザインを採用している。

• アーキテクチャ: エンコーダは分子グラフ $G$ を $K$ 次元の埋め込み $z$ にマッピングする。2つのバックボーンアーキテクチャが評価された：
  1. GPS Graph Transformer: ローカルなGINEメッセージパッシングとグローバルなマルチヘッド自己注意（multi-head self-attention）を組み合わせ、ランダムウォーク位置エンコーディングを利用したハイブリッドモデル。
  2. ChempropスタイルのD-MPNN: 有向メッセージパッシングネットワーク。
• 事前学習の目的関数: モデルはChEMBLから抽出された約290万個の生物活性精製分子を用いて訓練される。目的関数は、以下の2つの項の加重和を最小化する：
  1. 予測損失 ($L_{pred}$): 各分子に対して、複数のビュー（軽微な破損を伴う2つのグローバルビューと、BFS部分グラフによる8つのローカルビュー）が生成される。エンコーダは、ローカルビューの埋め込みからグローバルビューの埋め込みのセントロイドを予測しなければならない。極めて重要な点として、モデルが構造を無視して自明な解に陥るのを防ぐため、物理化学的記述子（RDKit）はグローバルビューにのみ付加されている。
  2. SIGReg損失 ($L_{SIGReg}$): この正則化項は、バッチの埋め込み分布を等方性ガウス分布 $N(0, I)$ に向かわせることで表現の崩壊を防ぐ。1024のランダムな方向における埋め込みの1次元投影の経験的特性関数を標準正規分布と比較するために、Epps–Pulley統計量を使用する。
• 評価プロトコル: 事前学習されたエンコーダは、厳格なマルチシード、ブートストラップベースのプロトコルの下で、以下の2つのタスクに対して評価された：
  1. 抗生物質活性: Staphylococcus aureus に対する39,312個の化合物データセット（Wongら）を用い、ランダム分割およびスキャフォールド分割（Bemis-Murcko）の両方で評価。
  2. ogbg-molhiv: HIV複製抑制に関する標準的なベンチマークであり、固定されたスキャフォールド分割を使用。
  • プローブ: 評価には「凍結プローブ（frozen probes）」（固定された埋め込みにヘッドを適合させる）と、フル「ファインチューニング（finetuning）」（エンコーダの重みを更新する）が含まれる。比較対象は、ランダム初期化および手作業によるベースライン（RDKit記述子、Morganフィンガープリント）である。

主な結果

1. 表現の質 vs. ファインチューニング:

   • 凍結プローブ: 事前学習は、有意に優れた表現をもたらす。ogbg-molhivにおいて、凍結されたLeJEPAの埋め込みは、ランダム初期化されたエンコーダのROC-AUC 0.665に対し、0.788を達成し（+0.123のリフト）、公開されている自己教師ありGNNの範囲に達した。抗生物質タスクにおいて、凍結された事前学習済み埋め込み（AUPRC 0.159）は、ランダム初期化（0.082）およびRDKit記述子単独（0.112）の両方を上回った。
   • ファインチューニング: この表現上の優位性は、ダウンストリームのファインチューニングの利点には頑健に変換されない。
     • 抗生物質のランダム分割において、ファインチューニングはスクラッチからの学習に対して利点を示さなかった。
     • 抗生物質のスキャフォールド分割において、単一の標準的なパーティションでは顕著な利点（+0.041 AUPRC）が見られたが、この効果は5つの異なるスキャフォールド・パーティション間で再現した際に消失した（プールされた $\Delta$ = +0.013, $p=0.095$）。
     • ogbg-molhivおよびD-MPNNバックボーンにおいても、ファインチューニングによる利点はゼロであった。

2. フィンガープリントとの相補性:

   • 著者らは、事前学習された埋め込みは、その情報的サブスペースが孤立している場合に限り、手作業で作られたMorganフィンガープリントと相補的であることを発見した。
   • 埋め込みは16〜32の有効次元で飽和するが、フィンガープリントは最大1024ビットまで改善する。
   • 次元を一致させた場合（128次元）、フィンガープリントは検証分割において埋め込みを上回ったが、シフトしたテストスキャフォールドにおいては下回った。一方で、埋め込みはスキャフォールドのシフトに対して優れた堅牢性を示した。
   • 特徴量レベルの結合: 事前学習された埋め込みをPCAにより32次元に削減し、それを1024ビットのMorganフィンガープリントと結合することで、ogbg-molhivのテストROC-AUCは0.803から0.830へと上昇した（$\Delta$ = +0.027, $p=0.014$）。
   • 決定的なことに、この利点は事前学習された埋め込みに特有のものであり、未学習（ランダム）の埋め込みを結合しても利点は得られなかった。これにより、この利点が単に次元を追加したことによるものではなく、学習された構造に由来することが確認された。

3. ベースラインの性能:

   • 抗生物質タスクにおいて、217個のRDKit記述子を用いた単純な決定木アンサンブル（グラフネットワークを使用しない）は、Chempropモデルや公開されているベースラインと同等またはそれを上回る性能を示した。これは、この特定のタスクにおいてはグラフネットワークが必ずしも必要ではないことを示唆している。

意義と主張
本論文は、LeJEPAの事前学習が、物理化学的構造を捉え、新しいケモタイプに対して良好に汎化する測定可能な優れた分子表現を生成すると結論付けている。しかし、著者らは、この利点は標準的なファインチューニングを通じては実現されない（ファインチューニングは、ランダム初期化や手作業による記述子よりも事前学習された重みを活用することに失敗することが多い）と主張している。

代わりに、その価値は特徴量レベルの結合を通じて実現される。事前学習された埋め込みは、手作業によるフィンガープリントが欠いているスキャフォールドの変化（分布シフト）に対する堅牢性を提供するが、それは埋め込みがその情報的サブスペースに削減され、他の表現と結合されている場合に限られる。著者らは、本研究を、表現自体は優れているものの、そのダウンストリームでの有用性は、標準的な手法による置き換えではなく、どのように統合されるかに条件付けられているというメカニズム的研究として位置づけている。また、「利点」はしばしば検証セットでは見えず、分布シフトの下で初めて明らかになるため、標準的なハイパーパラメータチューニングを通じて選択することが困難であることを強調している。

認められた限界
著者らは、ファインチューニングのエビデンスが本研究の中で最も弱い部分であることを指摘しており、単一の有意なスキャフォールド分割の結果が複数のパーティション間で再現されなかったことを挙げている。また、コンパクトなエンコーダが、より大規模な教師ありモデルと比較してベンチマークにおける絶対的な最高スコアに達していないこと、および事前学習コーパスがダウンストリームの化学空間を完全にはカバーしておらず、潜在的な転移の天井を制限していることも認めている。