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

Abstract

Self-supervised molecular representation learning can improve transfer on label-limited property prediction tasks, but contrastive objectives are sensitive to view construction and are often evaluated using protocols that report discrimination alone. We introduce M-JEPA (Molecular Joint Embedding Predictive Architectures), a predictive self-supervised method for molecular graphs based on connected-subgraph masking and an exponential-moving-average (EMA) teacher, and evaluate it with a compute-matched three-phase protocol that screens objectives on ESOL and tests transfer on Tox21 under Bemis-Murcko scaffold splits. Under matched compute, M-JEPA achieves lower ESOL proxy RMSE than a compute-matched InfoNCE baseline (2.12 vs 3.75; paired ΔRMSE 1.61, 95% CI 1.52-1.72; Wilcoxon p = 2.4 × 10-4) and shorter Phase-1 wall-clock time under the present implementation (24.09 ± 2.53 vs 44.21 ± 4.03 min on a single GPU). On Tox21, hybrid fine-tuning from the M-JEPA checkpoint improves mean ROC-AUC from 0.561 to 0.609 and reduces mean ECE from 0.279 to 0.064 relative to a supervised-from-scratch baseline under the same scaffold split, with paired cross-assay tests supporting systematic rather than assay-specific benefits (Wilcoxon p ≤ 5 × 10-3 for all four metrics). Motif-level Integrated Gradients attributions are stable across fine-tuning in most cases (median Spearman ρ = 0.90 across 62 molecule-assay pairs), although attribution stability and discrimination gains are end point-dependent. Scope. The conclusions in this study are specific to Tox21 scaffold-shift transfer with predictive versus contrastive self-supervision under matched compute; generalization to other MoleculeNet end points is a natural next step and is left for future work.