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

Karthik Iyer; Nasser Sabar

cs.LGOct 18, 2026

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

Authors: Karthik Iyer, Nasser Sabar

Organizations: Department of Computer Science and Information Technology, La Trobe University, Melbourne, VIC 3086, Australia

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.

Figures & tables

Figure 1: Overview of predictive self-supervised learning on molecular graphs. A chemically coherent substructure is masked from a molecular graph, and a context encoder is trained to predict the latent representation of the masked target region produced by an exponential-moving-average (EMA) teacher. The resulting representations support downstream molecular property prediction with subgraph-level attribution diagnostics and calibrated probability estimates under scaffold-split evaluation.

Figure 2: Leakage-resistant, staged evaluation pipeline separating objective-family choice (Phase 1), configuration selection (Phase 2), and downstream scaffold-shift evaluation (Phase 3).

Group	Parameter	Value(s) in Phase-1	Notes
Input & data	Unlabelled corpus	ZINC (canonicalised)	Cached graph dataset used for pretraining.
	Labelled proxy task	ESOL (regression)	Target: measured log solubility (“measured log solubility in mols per litre”).
	Labelled sample size	5,000	Used only for proxy evaluation stage.
	Unlabelled sampling cap	100,000 sampled; 250,000 max graphs/run	Constrains Phase-1 compute while keeping the objective comparison meaningful.
Architecture	Backbone (gnn_type)	{GINE, D-MPNN, SchNet3D}	Same candidate sets both objectives.
	Depth / width	3 layers; 256 hidden	Held constantly to isolate objective effects.

Table 1: Phase-1 sweep configuration (paired, compute-capped ESOL proxy).

Objective family	Runs	Mean RMSE	Std	Median	Best	Worst
JEPA	40	2.117	0.121	2.113	1.886	2.386
Contrastive	40	3.746	0.195	3.717	3.401	4.268

Table 2: ESOL validation RMSE distributions (proxy), Phase-1 sweeps (n=80).

Figure 3: Phase-1 objective comparison under matched compute (80 runs). M-JEPA achieves substantially lower ESOL proxy RMSE than the compute-matched contrastive baseline.

Figure 4: Phase-2 sweep structure (parallel-coordinates view of the explored hyperparameter space and outcomes).

Figure 5: ESOL proxy RMSE by augmentation profile and backbone (Phase-2 sweep). Heavy geometric augmentation degrades performance relative to minimal/no augmentation.

Figure 6: Augmentation and masking effects across Phase-2 runs, highlighting regime-level performance differences.

Figure 7: Reliability diagrams for Tox21 assays under scaffold split, comparing baseline (frozen) and hybrid fine-tuning settings.

Figure 8: Assay-level discrimination–reliability trade-off under scaffold shift, plotted as ΔPR-AUC versus ΔECE from baseline to hybrid fine-tuning.

Figure 9: Representative integrated-gradient attribution patterns across Tox21 assays under scaffold shift.

Component	Specification
Vast Compute Instance ID	29371882
Vast Host ID/Machine ID	79986 / 14726
GPU Architecture	NVIDIA GPU (CUDA-capable)
CUDA Capability	CUDA 12.6 (runtime), CUDA Toolkit 12.9
GPU Compute Performance	325.6 TFLOPS (FP16/TF32 aggregate, provider-reported)
DLPerf Score	169.4 DLPerf

Table 5: *

Category	Package	Version / Notes
Deep Learning	PyTorch	2.8.* (CUDA build, cu12x)
	TensorFlow	Installed (version per pip resolution)
Graph Learning	torch-geometric	2.5.3
	torch-scatter	2.1.2 (wheel or source-built depending on CUDA compatibility)
Scientific Computing	NumPy	<2.0
	SciPy	Installed via conda

Table 6: *

Backbone	Seed	n (JEPA)	mean RMSE (JEPA)	N (Contrastive)	mean RMSE (Contrastive)	Δ RMSE (C − J)
dmpnn	1	4	2.0987	2	3.7115	1.6128
dmpnn	2	3	2.2148	2	3.6197	1.4049
dmpnn	4	4	2.1084	3	3.8470	1.7386
dmpnn	5	1	2.2266	2	3.6330	1.4065
gine	2	1	1.8857	4	3.9393	2.0536
gine	3	4	2.2435	2	3.6389	1.3954

Table 7: *

Group	Parameter	Value(s) explored in Phase 2	Notes
Input & data	Unlabelled corpus	sampled to 250,000 graphs	Fixed compute cap per run.
	Labelled proxy	ESOL (regression), 5,000 labelled molecules	Proxy only; used for selection.
	Unlabelled sampling cap	250,000 sampled; 250,000 max graphs/run	Constrains Phase-2 compute while keeping the objective comparison meaningful.
	Graph caching	cache-datasets=1	Keeps data pipeline constant across trials.
Architecture	Backbone (gnn_type)	GINE, D-MPNN, SchNet3D	Backbone is part of the selection target.
	Hidden width	{256, 384}	Larger width is tested to check capacity sensitivity.

Table 8: *

Model	Data	Work	ROC-AUC	Split regime	Task type
SCFP (CNN)	T21C	(Hirohara et al., 2018)	0.877	Cross validation	Multi-task *
DeepTox (DNN)	T21C	(Mayr et al., 2016)	0.754 ± 0.144	Cross validation	Multi-task *
DeepTox (DNN)	T21C	(Mayr et al., 2016)	0.723 ± 0.136	Cross validation	Single task *
Weave	T21C	(Kearnes et al., 2016)	0.88 ± 0.003	Cross validation	Multi-task
FP2Vec	MN	(Jeon & Kim, 2019)	0.876 ± 0.006	Random	Multi-task
FP2Vec	MN	(Jeon & Kim, 2019)	0.825 ± 0.011	Random	Single task

Table 9: *

Group	Parameter (CLI / config name)	Value(s)	what it controls
Input & data	Dataset	Tox21	12 binary toxicity assays; strong class imbalance makes PR-AUC and calibration metrics informative.
	–csv	/srv/mjepa/data/tox21/data.csv	Source table for labels/SMILES; used to generate graph batches.
	–tasks	all – not shown for brevity	Multi-assay evaluation: per-assay metrics reported.
	Split strategy	Scaffold (implied by pipeline)	Scaffold splitting reduces chemical leakage; typically lowers AUC vs random split.
Representation / encoder	–encoder-checkpoint	/data/mjepa/cache/pretrain/encoder.pt	Starting molecular encoder weights used for hybrid fine-tuning.
	–encoder-source / –evaluation-mode	hybrid	Enables staged hybrid schedule (freeze → unfreeze) and hybrid evaluation bookkeeping.

Table 10: *

Assay	Overridden parameter(s)	Value(s)	Rationale
NR-Aromatase	pos_weight	10	Strong imbalance: explicit weighting stabilises minority gradients.
NR-AR	pos_weight, focal_gamma	8, 1.5	Moderate imbalance; reduced γ avoids over-penalising easy negatives.
SR-p53	pos_weight, focal_gamma	10, 1.5	Noisy assay: softer focal loss improves calibration stability.
SR-MMP	pos_weight	5.33	Higher positive fraction; weaker weighting sufficient.

Table 11: *

Metric	Mean ± SD	Notes
RMSE	1.63 ± 0.04	Regression error on external labelled dataset; fresh heads trained without Tox21-specific optimisation.
MAE	1.26 ± 0.03	Consistent with RMSE magnitude; reported for completeness.
R²	0.37 ± 0.03	Moderate explanatory power; not directly comparable to ESOL proxy RMSE due to dataset and protocol differences.
Validation RMSE	2.01 ± 0.03	Validation split performance; not used for model selection.
Number of seeds	3	Independent head initialisations.
Train / Val / Test size	901 / 113 / 113	Explicit split strategy.

Table 12: *

Molecule	Spearman ρ	Max \|ΔIG\|	Mean \|ΔIG\|	Median \|ΔIG\|
1	0.472	0.350	0.0564	0.0154
2	0.946	0.125	0.0092	0.0030
3	0.912	0.129	0.0167	0.0120

Table 13: *

Metric	Mean Δ	Median Δ	Improved assays	95% bootstrap CI (mean Δ)	Wilcoxon p (two-sided)
ROC-AUC	+0.048	+0.030	12/12	[0.025, 0.078]	4.88×10 ˆ -4
PR-AUC	+0.026	+0.016	11/12	[0.013, 0.042]	9.77×10 ˆ -4
Brier	−0.068	−0.029	12/12	[−0.114, −0.028]	4.88×10 ˆ -4
ECE	−0.216	−0.229	10/12	[−0.300, −0.129]	4.88×10 ˆ -3

Table 14: *

Assay	n	Median ρ	IQR (Q25–Q75)	Frac. low (ρ<0.50)	Median Top‑3 overlap
Overall	62	0.90	0.61–1.00	0.19	1.00
NR‑AR	4	1.00	0.94–1.00	0.00	1.00
NR‑AR‑LBD	6	0.80	0.65–0.96	0.17	0.50
NR‑AhR	5	0.80	0.79–1.00	0.20	1.00
NR‑Aromatase	6	0.75	0.45–0.90	0.33	0.75
NR‑ER	5	0.99	0.79–1.00	0.00	1.00

Table 15: *

Graph SSL

Graph self-supervised learning trains graph representations from structure and features without relying on task labels.

See all 63 papers

Self-Supervised Pretraining of Molecular Graph Encoders with LeJEPA
Sep 2, 2026Michał Kulczykowski, Rafał ŁabędzkiSelf-Supervised Pre-Training Molecular Representation Learning
SGH Warsaw School of Economics
Mol-JEPA: A multimodal Joint Embedding Predictive Architecture for Molecules
Aug 23, 2026Florian Rottach, Sebastian Schieferdecker, William Rudman +2Joint-Embedding Predictive Architecture Multimodal Pretraining
University of Tübingen · 2Boehringer Ingelheim · 3The University of Texas at Austin +1
Hierarchical Molecular Representation Learning via Fragment-Based Self-Supervised Embedding Prediction
Feb 23, 2026Jiele Wu, Haozhe Ma, Zhihan Guo +2Molecular Representation Learning Graph Representation Learning
School of Computing, National University of Singapore, Singapore · Computer Science and Engineering, The Chinese University of Hong Kong, China

Sep 2, 2026J/K move · Enter open · S save

Self-Supervised Pretraining of Molecular Graph Encoders with LeJEPA

Michał Kulczykowski, Rafał Łabędzki

SGH Warsaw School of Economics

Self-Supervised Pre-Training Molecular Representation Learning Molecular Property Prediction Joint-Embedding Predictive Architecture

Open paper AskPDF arXiv

Self-supervised pretraining has transformed language and vision, but its value for molecular graph neural networks remains contested. We ask whether pretraining on a large unlabelled corpus improves molecular property prediction. We adapt LeJEPA, a predictor-free joint-embedding predictive architecture regularised by Sketched Isotropic Gaussian Regularisation (SIGReg), to molecular graphs, evaluating GPS and Chemprop-style D-MPNN encoders on the Wong et al. [1] antibiotic-activity dataset and ogbg-molhiv using a multi-seed, bootstrap-based protocol. Pretraining improves learned representations but does not robustly improve finetuning. A frozen probe on pretrained embeddings exceeds random initialisation on both tasks (ogbg-molhiv ROC-AUC 0.788 vs 0.665; +0.123), reaching the published self-supervised band, but this does not translate into finetuning gains. On the antibiotic scaffold split, a canonical partition is significant (delta AUPRC +0.041, p = 0.010), but the effect vanishes across five partitions (pooled +0.013, p = 0.095). Finetuning is null on the random split, ogbg-molhiv, and D-MPNN. The representational edge is nevertheless recoverable. Embeddings saturate at ~16-32 effective dimensions, whereas Morgan fingerprints improve to 1024 bits. At matched dimensionality, fingerprints lead validation (0.799 vs 0.782 at 128 dimensions) but trail shifted test scaffolds (0.759 vs 0.788). Truncating embeddings and combining them with a 1024-bit Morgan fingerprint raises ogbg-molhiv ROC-AUC from 0.805 to 0.832 (delta +0.027; 95% CI [+0.003, +0.054]; p = 0.014); an untrained encoder gains nothing (delta -0.003). Thus, pretraining supplies complementary information best realised through feature-level combination, while finetuning gains are weak and partition-dependent.

Abstract

Graph SSL

Explore similar work

Self-Supervised Pretraining of Molecular Graph Encoders with LeJEPA

Mol-JEPA: A multimodal Joint Embedding Predictive Architecture for Molecules

Hierarchical Molecular Representation Learning via Fragment-Based Self-Supervised Embedding Prediction