M-JEPA: Predictive Self-Supervised Learning for Molecular Graphs with Scaffold-Shift Evaluation on Tox21.

Area of Science:

• Computational chemistry
• Machine learning
• Drug discovery

Background:

• Self-supervised learning (SSL) enhances molecular property prediction with limited data.
• Contrastive SSL methods are sensitive to data augmentation and often focus solely on discrimination metrics.

Purpose of the Study:

• Introduce M-JEPA, a predictive SSL method for molecular graphs.
• Evaluate M-JEPA's performance against contrastive methods using a compute-matched protocol.
• Assess transfer learning capabilities on toxicology prediction tasks.

Main Methods:

• Developed M-JEPA using connected-subgraph masking and an exponential-moving-average (EMA) teacher.
• Employed a three-phase protocol: objective screening on ESOL, transfer learning on Tox21 with scaffold splits.
• Utilized compute-matched comparisons for fair evaluation.

Main Results:

• M-JEPA achieved significantly lower RMSE on ESOL prediction compared to InfoNCE (2.12 vs 3.75).
• M-JEPA demonstrated faster Phase-1 training time (24.09 vs 44.21 min).
• Fine-tuning M-JEPA improved Tox21 prediction (ROC-AUC from 0.561 to 0.609) and reduced calibration error (ECE from 0.279 to 0.064).

Conclusions:

• M-JEPA offers a more efficient and effective predictive SSL approach for molecular representation learning.
• The method shows robust transfer learning benefits, particularly in reducing prediction uncertainty.
• Attribution stability is generally high, though dependent on specific prediction tasks.