3D self-supervised pre-training · 5×5 cross validation · two data sets

A 3D encoder, fine-tuned, against a frozen one read in context

Atom-JEPA¹ pre-trains an equivariant graph network on 19 million 3D conformers by predicting the latent representation of one part of a molecule from another, and then fine-tunes it on each data set. Its paper reports it best or tied on every ExpansionRx and Biogen endpoint against fingerprint models, Chemprop and KERMT. Run here with the authors' own code and weights on the same fifteen endpoints, folds and test sets as everything else on this site, that claim replicates: it beats fine-tuned CheMeleon on 38 of 45 endpoint × metric combinations and LightGBM on 45. Against Monroe + TabPFN 3², which the paper did not run, the two data sets disagree. On ExpansionRx fine-tuned Atom-JEPA is ahead on 6 and behind on 3 of the 9 endpoints on all three metrics (18 ahead and 9 behind of 27 endpoint × metric combinations); on Biogen it is behind on 4 of the 6 endpoints on all three metrics, and split on the other two (3 ahead and 13 behind of 18 endpoint × metric combinations).

The scoreboard

Tukey HSD per endpoint and metric, counted the way every page on this site counts: a method is on top if Tukey's correction cannot separate it from the best mean, and it is only best alone if nothing else is on top with it. The tables show MAE, in the units of each endpoint.

ExpansionRx⁷

Biogen ADME⁸

Each data set has a different leader. On ExpansionRx fine-tuned Atom-JEPA is on top of 19 of 27 endpoint × metric combinations and Monroe + TabPFN 3 of 9. On Biogen it is the other way round: Monroe on 17 of 18, Atom-JEPA on 5. Of the rest, only ChemProp + CheMeleon, fine-tuned reaches the top of anything, 3 in all. The frozen Atom-JEPA arm and LightGBM are on top of none.

The paper's claim replicates

The paper compares Atom-JEPA with fingerprint models, Chemprop, KERMT and CKERMT, and finds it best or statistically tied on all fifteen of these endpoints. Two of those comparators have a counterpart here. Paired by fold, fine-tuned Atom-JEPA is ahead of LightGBM on Morgan counts on 45 of 45 endpoint × metric combinations and behind on none, and ahead of fine-tuned CheMeleon on 38 and behind on two. A win or a loss is a paired t-test at p < 0.05 on one endpoint and one metric, uncorrected, so read the counts as a direction rather than a verdict on any single cell.

Against Monroe, the data sets disagree

On ExpansionRx fine-tuned Atom-JEPA has the lower MAE on LOG_HLM, LOG_MBPB, LOG_MGMB, LOG_MPPB, LogD and LogS, significantly on all six endpoints, and is significantly behind on LOG_Caco_AB, LOG_Caco_Efflux and LOG_MLM. On Biogen it is behind on LOG_HPPB, LOG_MDR1_ER, LOG_RLM, LOG_RPPB and LOG_SOL and not behind on LOG_HLM alone.

The obvious explanation is data volume: Biogen's folds are smaller, so perhaps fine-tuning needs more labels than an in-context head does. It does not hold. Across the fifteen endpoints the rank correlation between the number of molecules a fold trains on and fine-tuned Atom-JEPA's margin over Monroe is −0.04. Its two largest wins, LOG_MPPB and LOG_MBPB, train on 1,036 and 776 molecules per fold, fewer than three Biogen endpoints it loses.

Nor is it simply the endpoints that are hard for everything. On Caco-2 A→B, the worst of fine-tuned Atom-JEPA's losses, Monroe + TabPFN 3 reaches an R² of 0.471 ± 0.041 over the 25 folds. Every other arm is far below it: fine-tuned Atom-JEPA 0.103 ± 0.219, fine-tuned CheMeleon 0.160 ± 0.191, Atom-JEPA frozen 0.135 ± 0.103 and LightGBM −0.014 ± 0.129. Whatever Monroe's embedding carries about permeability, the 3D encoder does not, frozen or trained.

The representation costs, the fine-tuning pays

The frozen Atom-JEPA arm splits the step from Monroe to fine-tuned Atom-JEPA in two. The first step keeps TabPFN 3 and swaps Monroe's frozen embedding for Atom-JEPA's. The second keeps the encoder and fine-tunes it instead of freezing it. Each column is the change in mean R² over 25 folds, positive better.

As a frozen representation Atom-JEPA is worse than Monroe on 14 of 15 endpoints, by a median −0.090 R² on ExpansionRx and −0.094 on Biogen; paired by fold it is behind on 41 of 45 endpoint × metric combinations and ahead on none. Fine-tuning it then adds a median +0.126 R² on ExpansionRx and +0.074 on Biogen, improves on the frozen version on 14 of 15 endpoints, and is behind it on none of the 45 paired calls. On ExpansionRx that more than makes up the deficit; on Biogen it does not. What the paper's arm wins, it wins through fine-tuning, not through what pre-training put in the frozen encoder.

What it costs

Each fold fine-tunes the 39-million-parameter encoder for fifty epochs and scores ten conformers of every validation molecule after each one. The frozen arm is negligible next to it: the embedding is made once in about a minute per data set, and its 375 TabPFN 3 fits took about twenty minutes together. The fine-tuned folds ran on two machines, 42 on an RTX 5070 Ti and 8 on a DGX Spark, with the same code, the same conformers and the same seeds.

What this does and does not say

Reported the way “Even More Thoughts on ML Method Comparisons” and the protocol paper it points at⁹ argue comparisons should be: distributions over folds, simultaneous confidence intervals, paired tests on the folds, and no bolded maxima. Where two methods cannot be separated, both are reported as tied.

Code, per-fold metrics and every figure: studies/atom-jepa.

References

1. Petersen, K. H.; Tirsgaard, R. H.; Cornet, F. R. J.; Jordahn, M.; Schmidt, M. N. Atom-JEPA: Joint-Embedding Predictive Architecture for 3D Atomistic Systems. Preprint, 2026. The method this page puts to the test, run with the authors' code (github.com/khelverskovp/atom-jepa) and their released encoder (huggingface.co/atom-jepa/atom-jepa). arXiv:2610.08400
2. Banaszewski, B.; Fitzgibbon, A. W. Monroe: A Molecular Foundation Model for In-Context Probabilistic Inference. Preprint, 2026. The frozen encoder in the strongest arm of the source comparison. arXiv:2608.18982
3. Grinsztajn, L.; Flöge, K.; Key, O.; et al. TabPFN-3: Technical Report. Preprint, 2026. The head of the Monroe arm and of the frozen Atom-JEPA arm. arXiv:2605.13986
4. Burns, J. W.; Zalte, A. S.; Abreu, C. R. A.; et al. Deep Learning Foundation Models for Low-Data Regimes from Classical Molecular Descriptors. J. Chem. Inf. Model. 2026, articles ASAP. The pre-trained encoder of the other fine-tuned arm. 10.1021/acs.jcim.6c01546
5. Graff, D. E.; Morgan, N. K.; Burns, J. W.; et al. Chemprop v2: An Efficient, Modular Machine Learning Package for Chemical Property Prediction. J. Chem. Inf. Model. 2026, 66 (1), 28–33. 10.1021/acs.jcim.5c02332
6. Rogers, D.; Hahn, M. Extended-Connectivity Fingerprints. J. Chem. Inf. Model. 2010, 50 (5), 742–754. The LightGBM baseline's representation. 10.1021/ci100050t
7. OpenADMET; Expansion Therapeutics. OpenADMET–ExpansionRx Blind Challenge Data. Hugging Face, 2026, CC BY 4.0. 10.57967/hf/9687
8. Fang, C.; Wang, Y.; Grater, R.; et al. Prospective Validation of Machine Learning Algorithms for Absorption, Distribution, Metabolism, and Excretion Prediction: An Industrial Perspective. J. Chem. Inf. Model. 2023, 63 (11), 3263–3274. 10.1021/acs.jcim.3c00160
9. Ash, J. R.; Wognum, C.; Rodríguez-Pérez, R.; et al. Practically Significant Method Comparison Protocols for Machine Learning in Small Molecule Drug Discovery. J. Chem. Inf. Model. 2025, 65 (18), 9398–9411. 10.1021/acs.jcim.5c01609