Mol-JEPA: A multimodal Joint Embedding Predictive Architecture for Molecules
Florian Rottach, Sebastian Schieferdecker, William Rudman, Randall Balestriero, Carsten Eickhoff
cs.LG, cs.AI
2026-08-24
Mol-JEPA from Boehringer Ingelheim and academic labs learns molecular world models by masking 14 modalities in latent space, reaching 0.402 average MAE versus Chemprop's 0.519.
Self-supervised pretraining for molecules usually invents extra views by masking atoms or bonds, then pulling those views together in embedding space. That trick works for photos. For chemistry it often does the wrong thing. Property cliffs are common: one functional-group swap can move binding or toxicity by an order of magnitude. Treating the mutated graph as a positive pair teaches the model that two chemically different objects should look the same.
Structure is also an incomplete description. Absorption, metabolism, cellular phenotype and off-target binding live in a biochemical environment that a SMILES string never records. Several studies already showed that pretrained molecular embeddings frequently fail to beat fingerprints plus a random forest. Contrastive models can fuse graphs with cell-painting images, but they need negatives, and extra modalities tend to collapse.
Mol-JEPA applies a Joint Embedding Predictive Architecture to molecules, with one design change: it never perturbs the graph. It masks whole modalities and trains a Transformer to predict the missing latents.
The pretraining corpus has about 4.69 million unique compounds from ChEMBL, TDC, PCBA-1328, nabla-squared DFT and Enamine REAL. Each molecule can carry up to 14 modalities:
Three learned encoders (atom, graph, vector) map these inputs to a shared width. A Bernoulli mask with rate r drops some present modalities, keeping at least one visible; missing slots get mask tokens. The predictor is a PyTorch TransformerEncoder with a CLS token for readout. Collapse control is SIGReg from LeJEPA, applied to target embeddings; applying it to predictions made training collapse. The tuned model has about 50 million parameters, mask rate 0.3, lambda 0.7. Training used 12 RTX-6000ADA GPUs, effective batch 1024, and converged in about 3 days (around 150 epochs) before the online probe overfit.
Downstream options: linear or nonlinear probes or TabICLv2 on the CLS token; a two-layer Transformer over per-modality embeddings; LoRA rank-16 on the predictor. Code and weights ship with the paper.
Evaluation uses OpenADMET ExpansionRx, ASAP-Polaris, a PXR blind set and Biogen ADME (23 endpoints). Splits are Taylor-Butina clusters on ECFP4 at Tanimoto 0.65, three cluster assignments. Baselines: random forest, LightGBM, TabICLv2, CheMeleon, Chemprop, CLAMP.
| Method | Mean MAE | Mean R² |
| Mol-JEPA Best | 0.402 | 0.42 |
| Mol-JEPA Transformer | 0.427 | 0.38 |
| TabICLv2 + AlvaDesc | 0.453 | 0.37 |
| CheMeleon | 0.471 | 0.34 |
| LightGBM + AlvaDesc | 0.468 | 0.33 |
| Chemprop | 0.519 | 0.28 |
| CLAMP | 0.576 | 0.27 |
| Random forest + ECFP4 | 0.559 | 0.13 |
Gains concentrate on smaller sets (ASAP, Biogen) and on molecules farther from the training set in Tanimoto space. In paired Wilcoxon tests Mol-JEPA wins about 38% of comparisons, TabICLv2 plus AlvaDesc about 47%, mostly on the larger ExpansionRx tables. On official temporal competition splits, TabICLv2 still wins more often on MAE and Wilcoxon; Mol-JEPA leads on R².
On the full corpus, adding the extra modalities on top of Graph plus ECFP4 cuts linear-probe MAE by 14% and nonlinear by 13%. On a 100k subset: treating experimental vectors as semi-supervised labels is worse than JEPA (linear MAE 0.467 vs 0.436); dropping them also hurts (0.449); width 256 to 512 helps a lot, 1024 saturates; a projection head hurts. Leave-one-out ranks the graph encoder first (the only trainable GNN backbone), then MOE and ECFP.
For ADMET and activity modeling this is a practical multimodal recipe: if structure augmentations are chemically invalid, treat other measurements and other models as the views. JEPA needs no negatives and tolerates sparse modalities, which matches how drug data actually looks. Better behavior in low-data and farther-from-train regimes is the setting most projects live in.
It is still incremental. TabICLv2 with classical descriptors is not blown out on large tables or competition splits, and Wilcoxon slightly favors it. Mol-JEPA's contribution is packing phenotypes, docking-like scores, quantum descriptors and assays into one latent prediction loop, not owning every endpoint. A linear probe is the honest first try.
The authors list four. The reconstruction loss ignores that some modalities cannot recover others. Out-of-distribution work remains hard. Training dynamics across sparsity, width and redundancy are still poorly understood. Most ablations ran on 100k molecules, not the 4.69M corpus.
A few extra caveats. Mol-JEPA Best cherry-picks the best probe and embedding variant, so it is not a single-head comparison. LoRA underperforms, which means the latents are not plug-and-play. Boltz-2 covers only about 102k molecules and 9 targets. CLOOME can hurt downstream. Experimental vectors are extremely sparse. On temporal leaderboard splits, gradient-boosted tables remain competitive. Calling this a molecular world model still overstates how far the model extrapolates into new chemical space.