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

Abstract

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 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.

TL;DR: Mol-JEPA is a multimodal JEPA-based molecular foundation model that learns representations by predicting masked modalities (cellular effects, binding affinities, ADMET profiles, quantum chemistry) in latent space, achieving strong performance on drug discovery benchmarks.

Highlights

›No invalid augmentations - masks modalities instead of perturbing molecular structure

›Beyond molecular structure - integrates diverse modalities and scales efficiently to new ones

›No modality collapse - all modalities are predicted, ensuring none are ignored during training

Approach

Figure 1: Mol-JEPA. We model the biochemical environment by encoding molecular modalities from various datasets and backbone models. During training, we randomly mask modalities and predict them with a transformer module, taking the remaining modalities as input. Mode collapse is avoided by isotropic regularization via SIGReg.

Mol-JEPA is a multimodal joint embedding predictive architecture that predicts masked modalities from available ones. Given a masking ratio \(r\), some of the encoded modalities are deactivated, while ensuring that at least one modality is always available. The prediction module is implemented as a transformer encoder and optimized to predict the masked embeddings.

Training Objective. The Mol-JEPA pretraining loss combines the masked embedding prediction loss and the isotropic regularization loss across all modalities:

\[ \mathcal{L}_{\text{JEPA}} = \sum_{m \in \mathcal{M}} \left( \mathcal{L}^{m}_{\text{pred}} + \lambda\,\mathcal{L}^{m}_{\text{SIGReg}} \right), \qquad \lambda > 0 \]

The prediction loss \(\mathcal{L}^m_{\text{pred}}\) minimizes the MSE between masked and predicted modality embeddings. SIGReg conditions the latent variables to approximate an isotropic Gaussian distribution \(\mathcal{N}(0, I_K)\) using randomized sketching and the Epps–Pulley statistic:

\[ \mathcal{L}^m_{\text{SIGReg}} = \frac{1}{L} \sum_{l=1}^{L} T\!\left(\{a_l^\top s_j^m\}_{j=1}^{B}\right) \]

Our model integrates 14 modalities spanning embeddings from pretrained models (UMA, CLOOME, BioXMol, ChemGPT, Boltz-2, GNN), molecular descriptors (MOE, ECFP, xTB, ∇²DFT), and experimental measurements (ChEMBL bioactivity, PCBA, TDC ADMET). The pretraining dataset contains 4.69 million unique compounds.

Results

Mol-JEPA Evaluation Results

(a) Mol-JEPA demonstrates strong performance across tasks. The bars represent the aggregated mean absolute error of models within each benchmark family along with the standard deviation across the individual datasets.

(b) Mol-JEPA generalizes better. The lines show the average MAE for molecules across all datasets, binned by Tanimoto similarity to the train dataset. Mol-JEPA shows the best out-of-distribution performance.

Benchmark Comparison

Mean absolute error comparison across benchmark datasets for 3 Butina-cluster splits. We evaluate two variants of Mol-JEPA: the best performing model (Mol-JEPA Best) and the transformer probe on multimodal embeddings (Mol-JEPA Transformer). We compare against several baselines, including CLAMP (non-linear probe), CheMeleon (finetuned), Chemprop (finetuned), TabICLv2 (AlvaDesc), Random Forest (ECFP4), and LightGBM (AlvaDesc). The results are averaged over 3 random seeds. Mol-JEPA outperforms baselines on multiple datasets, with a clear advantage on smaller datasets (ASAP/Polaris, Biogen ADME) - especially important in drug discovery where large, high-quality labeled datasets are often prohibitively expensive. More results can be found in the paper.

Dataset	Mol-JEPA Best	Mol-JEPA Transformer	CLAMP	CheMeleon	Chemprop	TabICLv2	RF	LGBM
ExpansionRx
Caco-2 Permeability	0.34±.02	0.35±.02	0.40±.05	0.37±.05	0.39±.04	0.42±.02	0.43±.02	0.42±.02
Caco-2 Efflux	0.25±.02	0.26±.01	0.31±.03	0.27±.05	0.24±.03	0.24±.02	0.28±.04	0.24±.02
LogD	0.33±.02	0.46±.04	0.50±.05	0.35±.02	0.74±.06	0.36±.02	0.73±.07	0.43±.04
KSOL	0.45±.03	0.49±.04	0.53±.04	0.45±.06	0.60±.05	0.46±.06	0.62±.04	0.49±.03
HLM CLint	0.39±.02	0.43±.03	0.50±.02	0.40±.06	0.46±.01	0.39±.05	0.47±.02	0.39±.03
MLM CLint	0.40±.05	0.38±.07	0.36±.04	0.39±.06	0.42±.04	0.41±.09	0.44±.08	0.38±.07
MBPB	0.30±.02	0.30±.03	0.31±.01	0.29±.06	0.45±.10	0.26±.03	0.43±.01	0.34±.04
MGMB	0.30±.05	0.28±.06	0.35±.08	0.31±.13	0.41±.13	0.30±.09	0.39±.12	0.33±.10
MPPB	0.26±.04	0.28±.05	0.31±.03	0.31±.02	0.41±.01	0.26±.05	0.41±.06	0.29±.05
ASAP
MERS-CoV-2 Potency	0.59±.10	0.70±.08	0.96±.04	0.86±.09	0.52±.08	0.54±.09	0.58±.04	0.59±.13
SARS-CoV-2 Potency	0.62±.06	0.62±.06	1.02±.12	0.71±.13	0.73±.18	0.64±.17	0.77±.25	0.68±.16
LogD	0.68±.09	0.72±.11	1.09±.14	0.84±.19	0.64±.14	0.89±.07	0.99±.05	0.74±.14
KSOL	0.42±.19	0.53±.08	0.91±.19	0.60±.02	0.67±.19	0.65±.12	0.58±.06	0.70±.11
HLM	0.39±.11	0.39±.11	0.58±.03	0.50±.24	0.36±.09	0.42±.15	0.43±.07	0.43±.10
MLM	0.53±.01	0.53±.01	1.02±.37	0.67±.13	0.64±.08	0.59±.15	0.50±.06	0.51±.09
MDR1 Efflux	0.42±.13	0.44±.14	0.74±.14	0.51±.30	0.45±.22	0.48±.26	0.56±.36	0.58±.41
PXR
PXR Activity	0.55±.04	0.55±.06	0.75±.06	0.59±.13	0.79±.11	0.87±.20	0.81±.11	0.87±.17
Biogen ADME
Solubility	0.30±.02	0.30±.02	0.38±.02	0.35±.03	0.42±.04	0.33±.01	0.43±.01	0.33±.01
HLM CLint	0.33±.02	0.34±.02	0.45±.06	0.35±.01	0.51±.09	0.37±.03	0.55±.08	0.36±.02
RLM CLint	0.37±.04	0.38±.04	0.46±.04	0.39±.01	0.57±.07	0.39±.03	0.56±.04	0.39±.04
HPPB	0.33±.06	0.34±.07	0.45±.05	0.48±.08	0.56±.07	0.35±.04	0.67±.11	0.40±.07
RPPB	0.43±.07	0.45±.08	0.48±.18	0.55±.13	0.57±.18	0.49±.12	0.68±.15	0.56±.12
MDR1 Efflux	0.27±.03	0.29±.01	0.38±.05	0.30±.06	0.39±.02	0.31±.04	0.55±.08	0.32±.04
Average MAE	0.40±.05	0.43±.05	0.58±.08	0.47±.09	0.52±.09	0.45±.08	0.56±.08	0.47±.09

(a) Comparison win rates. The bars show how often a model wins across all pair-wise comparisons. While Mol-JEPA is best using MAE and R², it wins slightly less often when using Wilcoxon test for evaluation.

(b) Downstream probe performance. Besides full fine-tuning (FT), the Transformer probe on multimodal embeddings and TabICLv2 on CLS tokens are the best downstream predictors.

Modality Analysis

(a) Performance increases with more modalities. The lines represent the training time and epochs on a 100k subset dataset, annotated with the average downstream MAE.

(b) Leave one modality out analysis. We analyze the effect of removing modalities during training. The reported scores are the average downstream probe MAEs.

Training time increases approximately linearly with the number of modalities. In all experiments, we observe better performance when increasing the number of modalities. The graph modality contributes the most to downstream performance, likely because it is the only modality with a trainable GNN backbone. Molecular fingerprint modalities (MOE, ECFP4) rank among the most important, indicating that expert-designed features contain valuable information. Modalities capturing binding affinities and cellular interaction profiles also contribute, highlighting the value of incorporating biological context.

Multimodal Performance Gains

Multimodal performance gains. We re-train the model on the full dataset using only Graph and ECFP4 modalities. We find that the downstream performance is significantly better when using all modalities - a 14% reduction in MAE for linear probes and 13% for non-linear probes.

Citation (BibTeX)

@misc{rottach2026moljepamultimodaljointembedding,
      title={Mol-JEPA: A multimodal Joint Embedding Predictive Architecture for Molecules}, 
      author={Florian Rottach and Sebastian Schieferdecker and William Rudman and Randall Balestriero and Carsten Eickhoff},
      year={2026},
      eprint={2608.22642},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2608.22642}, 
}