JEPA · 2026-10-20
M-JEPA: Predictive Self-Supervised Learning for Molecular Graphs with Scaffold-Shift Evaluation on Tox21
Karthik Iyer, Nasser Sabar
AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.
TL;DR
M-JEPA adapts Joint Embedding Predictive Architectures to molecular graphs by masking chemically coherent connected subgraphs and training a context encoder to regress embeddings of an EMA teacher. The authors introduce a leakage-resistant, compute-matched three-phase evaluation and test on scaffold-shifted Tox21, reporting discrimination (ROC-AUC/PR-AUC) plus calibration (Brier/ECE) and motif-level integrated-gradient attribution.
Why it matters
The paper argues that molecular SSL should account for view construction sensitivity and realistic distribution shift. M-JEPA emphasizes chemically coherent masking and evaluates under Bemis–Murcko scaffold splits with joint discrimination–reliability analysis and attribution diagnostics, addressing cases where ranking improvements do not guarantee well-calibrated probabilities.
Method
- Pretrain a molecular JEPA: connected-subgraph masking; student context encoder predicts EMA teacher target embeddings via embedding-space regression (no explicit negative sampling).
- Use a leakage-resistant, compute-matched three-phase protocol: (1) 80-run ESOL objective screening under fixed compute, (2) JEPA-only refinement of masking/augmentation, (3) locked selection evaluated on Tox21 with Bemis–Murcko scaffold splits.
- Report both discrimination and calibration under shift (ROC-AUC/PR-AUC plus Brier score and ECE with reliability diagrams) and provide motif-level integrated-gradient attribution diagnostics.
Limitation
The authors’ key limitation is baseline breadth: they focus on a carefully compute-matched InfoNCE baseline rather than an exhaustive comparison to all recent molecular SSL objectives, and additional ablations (e.g., EMA removal, mask-ratio sensitivity) would further isolate which components drive gains.
Abstract (from arXiv)
Self-supervised molecular representation learning is increasingly used to improve transfer on small, label-limited property prediction tasks, yet commonly used contrastive objectives can be sensitive to view construction and are often evaluated with protocols that report only discrimination. We adapt Joint Embedding Predictive Architectures (JEPAs) to molecular graphs by masking chemically coherent connected subgraphs and training a context encoder to regress the masked target embedding produced by an exponential-moving-average (EMA) teacher. To evaluate predictive self-supervision under realistic distribution shift, we introduce a leakage-resistant, compute-matched three-phase protocol: (i) an 80-run objective screening on ESOL with fixed compute budgets, (ii) JEPA-only refinement of masking and augmentation choices, and (iii) per-assay Tox21 evaluation under Bemis-Murcko scaffold splits with joint discrimination-reliability analysis and attribution diagnostics. Under matched compute, M-JEPA achieves substantially lower ESOL proxy RMSE than a compute-matched InfoNCE baseline (mean RMSE 2.12 vs 3.75; mean paired dRMSE 1.61; 95% bootstrap CI [1.52, 1.72]; Wilcoxon p=2.4e-4) and typically completes runs faster (~15-25 min vs ~35-45 min on a single GPU under identical hardware), albeit with higher runtime variance. On Tox21, improvements in ROC-AUC/PR-AUC are assay-dependent and do not uniformly imply better calibration; reliability diagrams and Brier/ECE metrics reveal endpoints where fine-tuning improves ranking but leaves probabilities miscalibrated.
Related papers
- Atom-JEPA: Joint-Embedding Predictive Architecture for 3D Atomistic Systems
- AC-JEPA: Anchored Cross-Modal Joint Embedding Predictive Architecture for Visual-to-Symbolic Translation
- Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture
- Latent JEPA: Abstract Future Prediction for Latent Reasoning in Chemistry
- CI-JEPA: A Counterfactual Analysis of Latent Representations in Joint-Embedding Predictive Architectures for Self-Supervised Learning
Publication record
Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.
First public release: 2026-10-18 · Indexed by Pub-AI: 2026-10-20 UTC
M-JEPA: Predictive Self-Supervised Learning for Molecular Graphs with Scaffold-Shift Evaluation on Tox21
Preprint posted: 2026-10-18 · Original record ↗
- arxiv
- 2610.25163
Explicit source versions (1)
- v1 · 2026-10-18
- Preprint available
Source date: 2026-10-18 · First detected: 2026-10-25 UTC · Historical baseline
Source evidence (1)
- Official evidence ↗ · atom:published (official metadata)
Abstract · arxiv
Self-supervised molecular representation learning is increasingly used to improve transfer on small, label-limited property prediction tasks, yet commonly used contrastive objectives can be sensitive to view construction and are often evaluated with protocols that report only discrimination. We adapt Joint Embedding Predictive Architectures (JEPAs) to molecular graphs by masking chemically coherent connected subgraphs and training a context encoder to regress the masked target embedding produced by an exponential-moving-average (EMA) teacher. To evaluate predictive self-supervision under realistic distribution shift, we introduce a leakage-resistant, compute-matched three-phase protocol: (i) an 80-run objective screening on ESOL with fixed compute budgets, (ii) JEPA-only refinement of masking and augmentation choices, and (iii) per-assay Tox21 evaluation under Bemis-Murcko scaffold splits with joint discrimination-reliability analysis and attribution diagnostics. Under matched compute, M-JEPA achieves substantially lower ESOL proxy RMSE than a compute-matched InfoNCE baseline (mean RMSE 2.12 vs 3.75; mean paired dRMSE 1.61; 95% bootstrap CI [1.52, 1.72]; Wilcoxon p=2.4e-4) and typically completes runs faster (~15-25 min vs ~35-45 min on a single GPU under identical hardware), albeit with higher runtime variance. On Tox21, improvements in ROC-AUC/PR-AUC are assay-dependent and do not uniformly imply better calibration; reliability diagrams and Brier/ECE metrics reveal endpoints where fine-tuning improves ranking but leaves probabilities miscalibrated.