Atom-JEPA: Joint-Embedding Predictive Architecture for 3D Atomistic Systems

Kasper Helverskov Petersen; Rasmus Hannibal Tirsgaard; François R J Cornet; Mikkel Jordahn; Mikkel N. Schmidt

cs.LGOct 6, 2026

Atom-JEPA: Joint-Embedding Predictive Architecture for 3D Atomistic Systems

Authors: Kasper Helverskov Petersen, Rasmus Hannibal Tirsgaard, François R J Cornet, Mikkel Jordahn, Mikkel N. Schmidt

Organizations: Department of Applied Mathematics and Computer Science Technical University of Denmark

Abstract

Large-scale self-supervised pretraining has reshaped modern machine learning, substantially advancing the ability of language and vision models to generalize across downstream tasks. While deep learning has driven considerable progress in modeling atomistic systems in recent years, self-supervised pretraining in this domain has not yet achieved comparable downstream generalization. To address this, we introduce Atom-JEPA, a self-supervised pretraining framework that learns latent representations from unlabeled 3D structures through complementary atom-level and substructure-level objectives inspired by joint-embedding predictive architectures. We pretrain Atom-JEPA on large-scale molecular and crystalline datasets and evaluate its transfer performance by fine-tuning on a diverse set of downstream property prediction tasks. Atom-JEPA achieves state-of-the-art performance on molecular ADMET and quantum-chemical property prediction tasks, and is highly competitive in predicting the physical properties of crystalline materials. These results demonstrate the potential of latent-space predictive pretraining to support broad downstream generalization from structural data alone. Code and pretrained model checkpoints are publicly available at https://github.com/khelverskovp/atom-jepa

Figures & tables

Figure 1: Overview of Atom-JEPA. Molecular and crystalline structures are partitioned into complementary context (red) and target (blue) subgraphs by sampling a $k$ -hop neighborhood around an anchor atom. A 3D equivariant graph neural network encodes the context, which is then used to predict target representations at two scales: individual atoms and mean-pooled substructures. Target representations are extracted from the full structure after using a target encoder updated as an exponential moving average of the context encoder. Pretraining minimizes prediction errors in latent space, with the two subgraphs alternating as context and target. The pretrained context encoder is subsequently fine-tuned for downstream property prediction.

Figure 2: Average rank across 38 endpoints from three benchmarks (Biogen 4 endpoints, ChEMBL-MT 25 endpoints, ExpansionRx 9 endpoints). Top part weights each dataset equally, whereas bottom part weighs each endpoint equally. Bars join models whose paired rank difference could not be distinguished from zero by a stratified seed/fold bootstrap over endpoints (

10^{6}

replicates, resampled within dataset,

95\%

BCa interval). Tables with individual results can be found in Appendix E .

	Biogen (6)		ChEMBL MT (25)		ExpansionRx (9)
Model	Best / Indist.	Sum	Best / Indist.	Sum	Best / Indist.	Sum	Total (40)
Atom-JEPA $\times$ 10 conf	3 / 3	6	6 / 10	16	7 / 2	9	31
Atom-JEPA	0 / 6	6	0 / 15	15	0 / 9	9	30
CKERMT ( Xue et al., 2026 )	0 / 6	6	11 / 8	19	1 / 3	4	29
KERMT * ( Adrian et al., 2025 )	1 / 1 †	2	3 / 9	12	0 / 1	1	15 †
Chemprop ( Graff et al., 2026 )	1 / 3	4	0 / 12	12	0 / 1	1	17

Table 1: ADMET results. For each endpoint, the model with the lowest mean error across folds is designated as Best , and an unpaired Tukey’s test is used to identify additional models that are statistically indistinguishable ( Indist. ) from the best model at the

95\%

confidence level (see Appendix E.1 ). Sum denotes the total best or indistinguishable endpoints. ∗ reports only the mean error across splits; these values are treated as constant in the statistical tests. † did not report Biogen results for two targets and is therefore excluded from the statistical tests for these targets.

Task	$\alpha$	$\Delta\varepsilon$	$\varepsilon_{\text{HOMO}}$	$\varepsilon_{\text{LUMO}}$	$\mu$	$C_{\nu}$	$G$	$H$	$R^{2}$	$U$	$U_{0}$	ZPVE
Units	$ma_{0}^{3}$	meV	meV	meV	mD	$\frac{\text{mcal}}{\text{mol K}}$	meV	meV	$ma_{0}^{2}$	meV	meV	meV
No pretraining
MACE ( Batatia et al., 2022 )	38	42	22	19	15	21	5.5	4.7	210	4.1	4.1	1.23
Equiformer ${}_{\text{V}2}$ ( Liao et al., 2024 )	47	29.0	14.4	13.3	9.9	23	7.57	6.22	186	6.49	6.17	1.47
GotenNet L ( Aykent and Xia, 2025 )	28	19.8	13.4	12.2	6.7	19	4.98	3.30	24	3.41	3.37	1.08
Supervised pretraining

Table 2: QM9 results. The evaluation metric is MAE

(\downarrow)

. Bold: best result in each column; underline : within 10% of the best. Grey rows with

\dagger

use a different random train/validation/test split from 110k/10k/11k and are shown for reference only, excluded from highlighting. “–” denotes unreported targets.

Tasks	Phonons	Dielectric	Log GVRH	Log KVRH	Perovskites	MP Gap	MP E Form	MP Is Metal
Units	cm -1	–	$\log_{10}$ (GPa)	$\log_{10}$ (GPa)	meV	eV	meV/atom	F1
#Samples	1,265	4,764	10,987	10,987	18,928	106,113	132,752	106,113
No pretraining
CGCNN ( Xie and Grossman, 2018 )	$57.8{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 12.3}$	$.599{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.083}$	$.090{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.002}$	$.071{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.003}$	$45.2{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.7}$	$.297{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.004}$	$33.7{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.6}$	$\mathbf{.946}{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.008}$
ALIGNN ( Choudhary and DeCost, 2021 )	$29.5{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 2.1}$	$.345{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.087}$	$.072{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.001}$	$.057{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.003}$	$28.8{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.9}$	$.186{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.003}$	$21.5{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.5}$	$.902{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.002}$
coGN ( Ruff et al., 2024 )	$29.7{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 2.0}$	$.309{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.086}$	$.069{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.001}$	$.054{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.003}$	$26.9{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.8}$	$.156{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.002}$	$17.0{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.3}$	$.901{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.003}$

Table 3: Matbench results. The evaluation metric is MAE

(\downarrow)

for all tasks except MP Is Metal, which uses F1

(\uparrow)

. Values are means over five folds, with standard deviations where available. Pretrained models are ordered by increasing pretraining dataset size, shown in parentheses after each method. Bold: best result in each task. “–” denotes targets not reported by the original authors. * results from ( Ramlaoui et al., 2026 ) .

Figure 3: UMAP projections of Atom-JEPA representations for molecules annotated as belonging to one chemical family from the 32,252 molecules in the manually annotated three-star subset of ChEBI ( Malik et al., 2026 ) , colored by chemical family, and 50,000 randomly sampled crystals from the Alexandria pretraining dataset, colored by crystal system.

Task (Units)	#Samples	No pretraining	Atom-JEPA	$\Delta$ (rel.)	Impr..	$p$
ADMET datasets
Biogen	3,521	$.3825$	$\mathbf{.3458}$	$9.6\%$	6/6	$.011$
ExpansionRX	7,608	$.3673$	$\mathbf{.3361}$	$8.5\%$	9/9	$<.001$
ChEMBL-MT	114,112	$\mathbf{.4543}$	.4567	$-.5\%$	12/25	.57
Matbench
Phonons (cm -1 )	1,265	$26.9{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 4.0}$	$\mathbf{23.3}{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 4.5}$	$13.5\%$	5/5	$.015$

Table 4: Effect of Atom-JEPA pretraining. Identical encoder architecture is compared with and without pretraining on identical official folds; values are means and standard deviations across the five folds or endpoints for ADMET tasks.

\Delta

(rel.) is the mean fold-wise or endpoint relative improvement. “Impr.” counts folds or endpoints on which pretraining wins;

p

is a two-sided paired

t

-test over the five Matbench folds or mean MAE over the folds in the endpoints for ADMET tasks.

Figure 4: Effect of pretraining objectives and dataset. Frozen-feature MLP probes across 12 QM9 targets. The horizontal axis shows relative MAE reduction to the baseline of each panel. Positive values indicate improvement. (a) Joint (Atom+substructure) and substructure-only objectives relative to atom-only pretraining on QM9. (b) The joint objective relative to atom-only pretraining on Uni-Mol. (c) Uni-Mol pretraining relative to QM9 pretraining, with the objective held fixed. (d) The joint objective relative to atom-only pretraining on Uni-Mol on Biogen ADME validation objective.

Appendix figures & tables40 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Samples	Tasks	Description
Large-scale pretraining datasets
Uni-Mol ( Zhou et al., 2023 )	18,837,320	–	3D molecular conformers
Alexandria PBE ( Cavignac et al., 2026 )	1,696,344	–	DFT-relaxed inorganic crystals
Downstream evaluation datasets
ChEMBL-MT ( Adrian et al., 2025 )	114,000	25	Multi-task ADMET
ExpansionRx ( MacDermott-Opeskin and Castellanos, 2026 )	7,608	10	Multi-task ADMET

Appendix

Table 5: Summary of datasets. Uni-Mol and Alexandria are used for large-scale pretraining. The remaining datasets are used for downstream evaluation. QM9 is additionally used for pretraining in the separate ablation experiments.

Evaluation dataset	Unique molecules	Found in Uni-Mol	Overlap (%)
			Exact	Parent	Broad parent	Murcko scaffold
ChEMBL-MT	23,213	11,964	52	52	73	90
CL Microsome Human	968	632	65	67	73	88
CL Microsome Mouse	136	84	62	66	69	92
CL Microsome Rat	341	226	66	66	69	87
CL Total Dog	90	54	60	71	94	97

Appendix

Table 6: Uni-Mol pretraining-set overlap in downstream evaluation data. Dataset rows are unions over their evaluation test membership. "Unique molecules" and "Found in Uni-Mol" use the Exact definition and with an observed label for that endpoint. Exact denotes identical RDKit canonical isomeric molecular graphs and retains stereochemistry, isotopes, formal charge, and disconnected fragments. Parent compares cleaned fragment parents, ignoring salts and solvents while retaining stereochemistry. Broad parent also normalizes charge and isotopes, removes stereochemistry, and uses a heteroatom-tautomer hash. Murcko scaffold is the percentage of unique Bemis-Murcko scaffolds observed in Uni-Mol and measures structural relatedness. Percentage rounded to nearest percentage.

Evaluation dataset	Unique molecules	Found in Uni-Mol	Overlap (%)
			Exact	Parent	Broad parent	Murcko scaffold
QM9	128,822	4,742	4	4	5	9

Appendix

Table 7: Full QM9 overlap with Uni-Mol pretraining data. Exact denotes an identical canonical isomeric molecular graph. Parent removes salts and solvents while retaining stereochemistry. Broad parent additionally normalizes charge, isotopes, stereochemistry, and common heteroatom tautomers. Murcko scaffold measures structural relatedness rather than molecule identity and excludes acyclic molecules from its denominator. All QM9 properties share this molecular inventory. The graph-based profiles exclude 1,819 PyG QM9 entries whose explicit-hydrogen SMILES cannot be sanitized by RDKit; Unique molecules therefore denotes unique valid graph identities.

Evaluation dataset	Unique structures	Found in Alexandria	Overlap (%)
			Strict	Relaxed	Prototype	Composition
Matbench	151,731	59,032	39	44	76	59
Phonons	1,265	1,190	94	96	100	99
Dielectric	4,764	3,399	71	80	93	92
Log GVRH	10,987	8,697	79	83	99	93
Log KVRH	10,987	8,697	79	83	99	93

Appendix

Table 8: Matbench overlap with Alexandria PBE pretraining data. Strict requires the same elemental composition and a tight primitive-cell match without lattice rescaling. Relaxed uses lattice rescaling and moderate geometric tolerances to accommodate relaxation differences. Prototype permits element substitution but requires the same canonical anonymous formula, Pearson symbol, space group, and Wyckoff sequence. Composition requires only the same reduced elemental composition. The latter two columns measure structural or chemical relatedness, not identical materials.

Panel of Figure 5 — Figure 5: UMAP visualization of crystal representations by crystal system.

Panel of Figure 6 — Figure 6: UMAP visualization of crystal representations by space group.

Neighbor selection	Same space group
Embedding neighbors	69.7%
Composition-only neighbors	14.4%
Random crystals with matching atom count	12.4%
Unmatched random crystals	4.9%

Appendix

Table 9: Structural neighborhood agreement.

Predictor	Log-volume-per-atom	Coordination
Composition + atom count	.942	.794
Embeddings	.854	.814
Both together	.973	.894

Appendix

Table 10: Test-set coefficient of determination

R^{2}

for geometric quantities.

Predictor	Energy above hull	Formation energy	Indirect band gap
Composition + size	.072	.821	.493
Composition + size + geometry	.079	.861	.507
Embeddings	.109	.712	.446
Embeddings + composition + size + geometry	.136	.908	.542

Appendix

Table 11: Test-set coefficient of determination

R^{2}

for physical properties.

Figure 7: Uni-Mol representations colored by molecular weight, formal charge, ring count, and aromatic heavy-atom fraction.

Figure 8: Uni-Mol representations with six functional groups highlighted individually.

Panel of Figure 9 — Figure 9: Uni-Mol representations with the 6 most frequent molecular scaffolds highlighted, shown combined and individually (remaining points in grey).

Input	Ring count	Aromatic fraction	Radius of gyration
Composition + size	.706	.642	.615
Embeddings	.909	.848	.964
Combined	.958	.856	.969

Appendix

Table 12: Uni-Mol structural probe results.

Figure 10: UMAP visualization of ChEBI representations, with exclusive chemical families highlighted.

Figure 11: Visualization of ChEBI representations, with chemical families highlighted individually and the rest in grey.

Input	Mean average precision	Mean AUROC
Composition + size	.186	.874
Embeddings	.547	.964
Combined	.554	.966

Appendix

Table 13: ChEBI family probe results.

Figure 12: The three most frequent size/composition profiles highlighted in the randomly initialized and the Atom-JEPA pretrained QM9 representation, respectively

Figure 13: Pretrained and randomly initialized QM9 representations colored by within-property percentile rank for all twelve targets.

Model	$\alpha$	$\Delta\varepsilon$	$\varepsilon_{\text{HOMO}}$	$\varepsilon_{\text{LUMO}}$	$\mu$	$C_{v}$	$G$	$H$	$\langle R^{2}\rangle$	$U$	$U_{0}$	ZPVE
	$a_{0}^{3}$	eV	eV	eV	D	$\frac{\text{cal}}{\text{mol K}}$	eV	eV	$a_{0}^{2}$	eV	eV	eV
Training-set mean	6.343	1.077	.434	1.055	1.165	3.195	7.511	8.304	203.364	8.244	8.170	.720
All layers
Atom-JEPA	.091	.095	.061	.064	.032	.031	.013	.012	.419	.012	.012	.001
Random initialization	.371	.279	.171	.221	.039	.108	.091	.089	1.850	.089	.091	.004
Last layer

Appendix

Table 14: QM9 frozen-backbone MLP probe test MAE results using all-layer or last-layer features for a pretrained encoder vs random initialization, with a training-set mean predictor as a baseline (atom-referenced for

U

U_{0}

H

G

Model	Biogen	ChEMBL-MT	ExpansionRx	Mean rank across endpoints	Group
Residual correlation $\bar{r}$ to peers (lower is more distinct)
Atom-JEPA $\times$ 10 conf	$.723$	$\mathbf{.791}$	$.730$	$\mathbf{1.98}$ [1.66, 2.34]	a
Chemprop	$\mathbf{.719}$	$.803$	$\mathbf{.681}$	$2.53$ [2.02, 2.88]	ab
Mol-JEPA $\cdot$ modalities	$.738$	$.801$	$.749$	$2.78$ [2.42, 3.14]	b
Morgan + RDKit $\cdot$ LGBM	$.724$	$.814$	$.761$	$2.98$ [2.59, 3.31]	b
CKERMT	$.779$	$.829$	$.787$	$4.73$ [4.48, 4.86]	c

Appendix

Table 15: Model distinctiveness. How reproducible each model’s errors are from the other four model families, measured two ways on the same per-molecule test residuals. Residual correlation is the averaged pairwise correlation between other models, averaged over test data-points, Uniqueness is the

1-R^{2}_{\mathrm{adj}}

from regressing a model’s residuals on all four peers’ residuals at once. Per-benchmark columns are means over each benchmark’s endpoints. Mean rank ranks the families within each benchmark and endpoint cell and averages with equal weight per benchmark, with rank 1 = most distinct. The brackets are a stratified cluster bootstrap over endpoints. Group shows the group of models that cannot separated by the paired 95% bootstrap interval on their mean-rank difference.

Model	Caco2	HIA	P-gp	Bioav.	Lipo.	Solub.
Unit	MAE $\downarrow$	AUROC $\uparrow$	AUROC $\uparrow$	AUROC $\uparrow$	MAE $\downarrow$	MAE $\downarrow$
GBDT models
Morgan	$.3530{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0370}$	$.8170{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0300}$	$.8650{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0230}$	$.5690{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0730}$	$.6060{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0170}$	$.9700{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0060}$
RDKit	$.3020{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0110}$	$.8900{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0500}$	$.8690{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0190}$	$.5800{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0430}$	$.5690{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0110}$	$.7630{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0110}$
Morgan + RDKit	$.2950{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0060}$	$.8940{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0400}$	$.8590{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0080}$	$.5960{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0550}$	$.5420{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0050}$	$.7760{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0280}$
Morgan + RDKit + Avalon + ErG	$.3290{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0150}$	$.9410{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0170}$	$.8690{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0160}$	$.5790{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0770}$	$.5400{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0110}$	$.8030{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0050}$

Appendix

Table 20: Results on the TDC-ADMET benchmark: absorption. Columns correspond, from left to right, to caco2_wang , hia_hou , pgp_broccatelli , bioavailability_ma , lipophilicity_astrazeneca , and solubility_aqsoldb . Numbers from Koleiev et al. (2026) .

Model	BBB	PPBR	VDss
Unit	AUROC $\uparrow$	MAE $\downarrow$	Spear $\uparrow$
GBDT models
Morgan	$.8300{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0110}$	$8.6780{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.1380}$	$.6070{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0230}$
RDKit	$.8770{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0150}$	$7.5460{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.1110}$	$.6220{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0090}$
Morgan + RDKit	$.8600{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0140}$	$7.5100{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0980}$	$.6250{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0260}$
Morgan + RDKit + Avalon + ErG	$.8600{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0260}$	$8.1450{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.3010}$	$.5750{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0230}$

Appendix

Table 21: Results on the TDC-ADMET benchmark: distribution. Columns correspond, from left to right, to bbb_martins , ppbr_az , and vdss_lombardo . Numbers from Koleiev et al. (2026) . RDKit + Morgan and RDKit + Morgan + Avalon + ErG MLP diverged during training on VDSS Lombardo

Model	2C9-I	2D6-I	3A4-I	2C9-S	2D6-S	3A4-S
Unit	AUPRC $\uparrow$	AUPRC $\uparrow$	AUPRC $\uparrow$	AUPRC $\uparrow$	AUPRC $\uparrow$	AUROC $\uparrow$
GBDT models
Morgan	$.6730{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0590}$	$.5300{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0100}$	$.7800{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0050}$	$.2810{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0000}$	$.4660{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0960}$	$.5990{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0190}$
RDKit	$.7110{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0090}$	$.5880{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0150}$	$.7840{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0100}$	$.2970{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0250}$	$.5420{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.1280}$	$.6170{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0280}$
Morgan + RDKit	$.7120{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0100}$	$.6000{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0110}$	$.8090{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0060}$	$.2810{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0090}$	$.5560{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0500}$	$.6310{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0090}$
Morgan + RDKit + Avalon + ErG	$.6930{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0120}$	$.5630{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0290}$	$.8100{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0110}$	$.2930{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0150}$	$.5140{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.1230}$	$.6210{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0300}$

Appendix

Table 22: Results on the TDC-ADMET benchmark: metabolism. Columns correspond, from left to right, to cyp2c9_veith , cyp2d6_veith , cyp3a4_veith , cyp2c9_substrate_carbonmangels , cyp2d6_substrate_carbonmangels , and cyp3a4_substrate_carbonmangels . Numbers from Koleiev et al. (2026) .

Model	Half-life	Cl-Mic	Cl-Hep
Unit	Spear $\uparrow$	Spear $\uparrow$	Spear $\uparrow$
GBDT models
Morgan	$.4130{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0430}$	$.4670{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0290}$	$.3690{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0250}$
RDKit	$.3520{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0350}$	$.5710{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0310}$	$.4030{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0480}$
Morgan + RDKit	$.4150{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0370}$	$.5510{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0190}$	$.4270{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0330}$
Morgan + RDKit + Avalon + ErG	$.3810{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0440}$	$.4620{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0680}$	$.3750{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0190}$

Appendix

Table 23: Results on the TDC-ADMET benchmark: excretion. Columns correspond, from left to right, to half_life_obach , clearance_microsome_az , and clearance_hepatocyte_az . Numbers from Koleiev et al. (2026) .

Model	hERG	AMES	DILI	LD50
Unit	AUROC $\uparrow$	AUROC $\uparrow$	AUROC $\uparrow$	MAE $\downarrow$
GBDT models
Morgan	$.8100{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0180}$	$.7690{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0170}$	$.8640{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0170}$	$.6750{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0170}$
RDKit	$.7910{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0220}$	$.8050{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0050}$	$.8920{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0230}$	$.6370{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0120}$
Morgan + RDKit	$.8120{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0170}$	$.8160{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0100}$	$.8870{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0180}$	$.6500{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0090}$
Morgan + RDKit + Avalon + ErG	$.7680{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0530}$	$.8050{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0130}$	$.8300{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0470}$	$.6400{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm.0130}$

Appendix

Table 24: Results on the TDC-ADMET benchmark: toxicity. Columns correspond, from left to right, to herg , ames , dili , and ld50_zhu . Numbers from Koleiev et al. (2026) .

Task	$\alpha$	$\Delta\varepsilon$	$\varepsilon_{\text{HOMO}}$	$\varepsilon_{\text{LUMO}}$	$\mu$	$C_{\nu}$	$G$	$H$	$R^{2}$	$U$	$U_{0}$	ZPVE
Units	$ma_{0}^{3}$	meV	meV	meV	mD	$\frac{\text{mcal}}{\text{mol K}}$	meV	meV	$ma_{0}^{2}$	meV	meV	meV
Pretrained on QM9 (130k)
atom+sub	115	130.6	86.1	88.0	56.0	40.8	22.3	22.9	396	22.6	22.1	1.8
atom-only	97	114.4	69.6	81.4	49.4	32.4	18.4	18.4	398	18.4	18.6	1.5
sub-only	104	131.6	90.0	80.2	82.3	42.7	23.0	22.7	420	22.5	22.4	1.9
Pretrained on Uni-Mol (19M)

Appendix

Table 25: All-layer MLP probes on frozen encoder features pretrained on the QM9 dataset and Unimol dataset respectively across combinations of the pretraining objectives. The evaluation metric is MAE

(\downarrow)

. Bold: best result within each pretraining group; bold and underline : best overall. “atom” and “sub” denote atom-level and substructure-level objectives, respectively.

Hyperparameter	Uni-Mol	Alexandria	QM9
Encoder
Maximum degree $L_{\max}$	2	2	6
Maximum order $M_{\max}$	2	2	2
Number of Transformer blocks	8	8	8
Embedding dimension $d_{\mathrm{embed}}$	$(2,256)$	$(2,256)$	$(6,256)$
Attention hidden dimension $d_{\mathrm{attn\text{-}hidden}}$	$(2,64)$	$(2,64)$	$(6,64)$

Appendix

Table 26: Hyperparameters used for pretraining on Uni-Mol, Alexandria, and QM9. We use EquiformerV3 for the context encoder, target encoder, and predictor. The context and target encoders use identical architectures. Predictor hyperparameters not explicitly listed are identical to those of the corresponding encoder.

Figure 14: Evolution of RankMe (top) and $\alpha$ -reQ (bottom) during pretraining for the Uni-Mol (left) and Alexandria (right) encoders.

Figure 15: Total training loss during pretraining for the Uni-Mol (left) and Alexandria (right) encoders.

Processor Hyperparameter	Interval	Prior
MLP
Learning rate	$(10^{-4},10^{-2})$	Log-uniform
Epochs	$(5,500)$	Log-uniform
Weight decay	$(10^{-4},10^{-1})$	Log-uniform
Head dropout	$(0,0.5)$	Uniform
Head hidden units	$(32,256)$	Log-uniform

Appendix

Table 27: Intervals and priors for the hyperparameters of the specific feature processor tuned over using Optuna.

Feature Hyperparameter	Interval	Prior
Morgan
Radius	$\{2,3\}$	Categorical
Number of bits	$\{1024,2048\}$	Categorical
Use Counts	$(\text{False, True)}$	Boolean
RDKit
–	–	–

Appendix

Table 28: Intervals and priors for the hyperparameters of the specific feature extractor tuned over using Optuna.

Dataset	$N$	Fragmented	Fragmented (%)	Extra atoms (%)
Caco2 Wang	906	9	1.0	4.6
HIA Hou	578	0	0.0	0.0
Pgb Broccatelli	1212	0	0.0	0.0
Bioavailability Ma	640	1	0.2	52.5
Lipophilicity AstraZeneca	4200	1	0.0	4.9
Solubility AqSolDB	9982	1098	11.0	30.2

Appendix

Table 29: Molecules in the TDC ADMET benchmark group that RDKit parses as more than one connected component (salts, hydrates, mixtures), over the de-duplicated train_val + test SMILES of each dataset. “Extra atoms” is the mean, over fragmented molecules only, of the share of atoms outside the largest fragment.

Figure 16: The validation error on the validation error (MAE) against the number of conformers used during training. The different curves represent different numbers of predictions from conformers averaged over during test time (test time augmentations).

Panel of Figure 17 — Figure 17: Histogram of the prediction error of validation datapoints versus the energy of the conformer. Left plot is the relative energy of the conformer from the mean, right plot is the rank of the conformer. For none of the plots a strong correlation between error and conformer energy was observed.

Hyperparameter	Biogen ADME	ChEMBL-MT	ExpansionRX
Prediction Head
Number of tasks	6	25	9
Trunk Architecture	Linear $(768,256)$ $\rightarrow$ SiLU $\rightarrow$ Dropout
Per Task Head Architecture	Linear $(256,64)$ $\rightarrow$ SiLU $\rightarrow$ Linear $(64,1)$
Shared Trunk dropout	.1
Pooling readout	Mean

Appendix

Table 30: Hyperparameters used for Atom-JEPA fine-tuning on Biogen ADME, ChEMBL-MT, ExpansionRX, and ADMET-TDC. All three start from the same Uni-Mol-pretrained context encoder; encoder architecture is unchanged from Table 26 . Every task gets its own MLP head on top of a shared trunk that takes the atom-pooled representation.

Hyperparameter	All targets
Optimization
Optimizer	AdamW
Adam $\epsilon$	$10^{-7}$
Maximum training epochs	1,000
Batch size	32
Encoder base learning rate	$10^{-4}$

Appendix

Table 31: Hyperparameters used for QM9 fine-tuning. All targets share the same optimization settings, with target-specific readout architectures for

\mu

and

R^{2}

. Encoder architecture hyperparameters are inherited from the pretrained model.

Hyperparameter	Small-to-medium tasks	Large tasks
Optimization
Optimizer	AdamW	AdamW
Adam coefficients $(\beta_{1},\beta_{2})$	$(.9,.95)$	$(.9,.95)$
Adam $\epsilon$	$10^{-8}$	$10^{-8}$
Number of epochs	100	12
Training batch size	4	32

Appendix

Table 32: Hyperparameters used for Matbench fine-tuning. Small-to-medium tasks comprise phonons, dielectric, Log GVRH, Log KVRH and Perovskites; large tasks comprise MP Gap, MP E Form, and MP Is Metal. Classification settings apply only to MP Is Metal. Encoder freezing applies only to pretrained runs; encoders trained from scratch are trainable from the first epoch.

Figure 18: Tukey HSD grid plot across the endpoints of Biogen Scaffold test data over 5 seeds. The analysis was conducted with a blocked ANOVA test followed by Tukey HSD using the Statsmodels Python package.

Figure 19: Tukey HSD grid plot across the endpoints of ExpansionRX temporal test data over 4 seeds. The analysis was conducted with a one-way ANOVA test followed by Tukey HSD using the Statsmodels Python package.

Figure 20: Tukey HSD grid plot across the endpoints of ChEMBL-MT test data over 4 seeds. The analysis was conducted with a blocked ANOVA test followed by Tukey HSD using the Statsmodels Python package.

Explore similar work

May 25, 2026cs.LG

UNATE: UNsupervised ATomic Embedding for crystal structures property prediction

Accurately predicting crystal properties is critical for accelerating materials discovery, but it is often limited by scarce labeled data and costly theoretical calculations. To alleviate this, we propose UNATE (Unsupervised Atomic Embedding), a framework that leverages structural information extracted from unlabeled crystal structures. UNATE integrates an unsupervised denoising autoencoder with self-supervised contrastive learning to learn robust atomic representations, which are then used as input features for downstream property prediction. Experimental results show that replacing raw atomic numbers with UNATE-pretrained node embeddings yields a 2.7% improvement over the full-data baseline. Notably, the benefits become more pronounced in scenarios with limited labeled data, reaching improvements of up to 10% when only 25% of the labeled data is used.

Laura Solà-Garcia, Àlex Solé, Javier Ruiz-Hidalgo

Signal Theory and Comms. Dept. UPC-Univ. Polit`ecnica de Catalunya Barcelona, Spain

Crystal Structure Prediction Molecular Property Prediction Unsupervised

Ask PDF

Aug 23, 2026cs.LG

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

Despite recent advances in molecular foundation models, several limitations remain, such as chemically invalid augmentations, modality collapse, and incomplete representation of biochemical environments. To address these challenges, we present \textbf{Mol-JEPA}, a scalable framework for learning molecular world models. Rather than relying on suboptimal molecular perturbations, our model uses modality masking to exploit information from molecular structures, cellular phenotypes, binding affinities, ADMET profiles, quantum chemistry simulations and other drug discovery data. Across various benchmarks, we show that the representations learned by Mol-JEPA deliver strong performance, demonstrating the value of incorporating biochemical context through latent space prediction.

Florian Rottach, Sebastian Schieferdecker, William Rudman +2

University of Tübingen · 2Boehringer Ingelheim · 3The University of Texas at Austin +1

Molecular Representation Learning Molecules Drug Discovery Latent Prediction S-Jepa

Ask PDF

Sep 15, 2026cs.LG

Procedural Pretraining for Molecular Property Prediction

Molecular property prediction is often limited by the small size of labeled downstream datasets, motivating pretraining on large corpora of unlabeled molecules. In this work, we ask whether useful inductive biases can instead be learned from abstract, procedurally generated data before a model sees any molecular data. We introduce a three-stage training pipeline consisting of procedural pretraining, molecular pretraining on SMILES, and downstream fine-tuning, and evaluate several procedural tasks spanning sequence structure, cellular automata, and graph reasoning. We find that procedural pretraining can improve molecular property prediction even after subsequent molecular pretraining: on Lipophilicity, \textsc{Reverse} reduces test error by 4.8%. For context, the magnitude of this improvement is roughly 90% of the performance difference between our 250K-molecule baseline and the publicly released MoLFormer checkpoint pretrained on approximately 100M molecules. Our analysis shows that the benefit is strongest under downstream data scarcity, depends on the structure of the procedural data rather than only surface-level statistics, and does not increase monotonically with additional procedural training. Instead, transfer typically peaks at an intermediate procedural budget and deteriorates as the model approaches convergence on the procedural task. We further find that, for several tasks, much of the transferable information is localized in the attention layers, while feed-forward layers can contribute to over-specialization. These results show that procedural data can provide transferable structure for molecular learning and offer a complementary route to improving performance when labeled molecular data are limited.

Moritz Friedemann, Zachary Shinnick, Philip Torr +1

University of Oxford · Adelaide University · Slater Labs

Molecular Property Prediction Pretraining Procedural

Ask PDF