arXiv preprint ·
M-JEPA: Predictive Self-Supervised Learning for Molecular Graphs with Scaffold-Shift Evaluation on Tox21
- Published
- Source
- arXiv preprint, posted before peer review
- DOI
- 10.1021/acs.jcim.6c00828
- Original
- Read it at arXiv
- Topics
Abstract
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.
The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.
Read it on ReadSearch
Each opens after you create an account, and brings you back to this paper.
More in Machine Learning
Stable Signer V3: 4D Deformable Primitive Splatting Is All You Need for Sign Language Production
When Better Opportunities Backfire: Service Worker Responses to Parallel AI Deployment in Customer Acquisition
Means-Ends Analysis Helps LLMs Plan in Continuous-State Domains
Safety Offset Learning: When Benign Fine-Tuning Misaligns LLMs
What Can't Jev Decide? Probing the Limits of One-Pass Decision Models
Get the few papers worth your time in Machine Learning every morning.
ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.