PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Florian Rottach, Sebastian Schieferdecker, William Rudman, Randall Balestriero, Carsten Eickhoff

arXiv:2608.22642v1cs.LGcs.AI

TL;DR

Molecular foundation models face chemically invalid augmentations, modality collapse, and incomplete biochemical context. Mol-JEPA uses modality masking to predict latent representations across diverse molecular and biochemical modalities, achieving strong performance across downstream tasks, particularly on smaller datasets and out-of-distribution settings. The authors identify modality selection and out-of-distribution robustness as continuing challenges.

  • Problem

    Molecular representation learning is limited by chemically invalid augmentations, modality collapse, and missing biochemical context, while pretrained representations may yield only marginal gains over traditional baselines.

  • Method

    Mol-JEPA masks available molecular modalities and predicts their latent representations using diverse biological, chemical, and computational data.

  • Results

    Mol-JEPA achieves strong performance across downstream tasks and is particularly better on smaller datasets and in out-of-distribution settings.

  • Takeaways & Limitations

    Latent-space prediction provides a scalable framework for integrating multiple molecular modalities and biochemical context.

  • Takeaways & Limitations

    Out-of-distribution generalization remains a key challenge, and most ablations used a considerably smaller dataset than full pretraining.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Mol-JEPA addresses limitations in molecular representation learning by replacing chemically questionable augmentations with modality masking and by incorporating biochemical context. The framework combines diverse modalities and reports competitive benchmark performance.

  • Pretrained molecular representations often fail to consistently outperform traditional baselines or provide only marginal gains.
  • Molecular structure alone omits biological context from protein networks, metabolism, cellular states, and physiological environments.
  • Mol-JEPA masks entire modalities and predicts their latent representations across molecular, biological, chemical, and computational data.
  • The framework introduces self-supervised learning through meaningful augmentations and uses a scalable multimodal Joint Embedding Predictive Architecture.
  • The multimodal dataset contains nearly 5 million molecules spanning diverse drug-like chemical space and biological, chemical, and computational modalities.
  • Mol-JEPA demonstrates highly competitive performance on relevant benchmarks while evaluating prediction quality, out-of-distribution behavior, and modality contributions.

2 RELATED WORK

Prior molecular self-supervised and multimodal methods can introduce augmentation errors, require negative samples, neglect biochemical effects, or risk modality collapse. Mol-JEPA instead masks modalities and predicts all latent representations while integrating biochemical-environment information.

  • Self-supervised molecular representation learning: Self-supervised molecular methods commonly use masked prediction or contrastive learning, including graph augmentations that may create false positives.
  • Self-supervised molecular representation learning: CLIP-based molecular models fuse graphs with other modalities but require negative samples that can introduce false negatives.
  • Multimodal architectures: Many multimodal architectures represent molecular structure while neglecting modalities that reflect effects in biochemical environments.
  • Multimodal architectures: Existing fusion approaches using concatenation or attention-based aggregation can result in modality collapse.
  • Multimodal architectures: Mol-JEPA avoids chemically invalid augmentations by masking entire modalities, integrates biochemical-environment knowledge, and predicts all modalities to encourage their use.

3 METHOD

Mol-JEPA is pretrained on a curated multimodal dataset and learns to predict masked molecular modalities from available context. Its objective combines modality prediction with isotropic regularization, and evaluation uses realistic drug-discovery datasets and established baselines.

  • 3.1 PRETRAINING DATA: The pretraining corpus merges public datasets, pretrained-model embeddings, quantum-chemical calculations, and experimental measurements into 4.69 million unique compounds.The resulting dataset captures multidimensional chemical and biological behavior across drug-like molecules.
  • 3.1 PRETRAINING DATA: Mol-JEPA represents each molecule with 14 modalities spanning embeddings, descriptors, computed properties, and multidimensional behavioral labels.Three learnable encoders project modalities to a shared dimensionality.
  • 3.3 MODEL ARCHITECTURE: Given a masking ratio r, the model deactivates selected modalities while retaining at least one available modality and replaces missing inputs with masking tokens.A Transformer predictor reconstructs masked embeddings from the available context, encouraging integration across alternative molecular environments.
  • 3.4 PRETRAINING OBJECTIVES: The final objective combines masked-modality prediction and SIGReg losses across modalities, with their contribution balanced by λ.SIGReg regularizes latent variables toward an isotropic Gaussian distribution to reduce trivial representation collapse.
  • 3.6 BASELINES: Evaluation compares Mol-JEPA with Random Forest, LightGBM, and other established baselines on recent OpenADMET blind challenges and high-quality datasets.The evaluation is designed to address concerns about noisy labels, inconsistent chemical representations, and limited real-world relevance in traditional benchmarks.

4 RESULTS

Mol-JEPA outperforms baselines across benchmark families, with particular advantages on smaller datasets and out-of-distribution settings. Analyses further show modality-dependent training behavior, redundancy, and computational scaling.

  • Benchmark performance: Mol-JEPA outperforms baselines across multiple datasets, especially on smaller ASAP/Polaris and Biogen ADME benchmarks and at greater distances from training data.The authors attribute the out-of-distribution advantage to pretraining on a large and diverse set of drug-like compounds.
  • Statistical significance: 38% of comparisons favored Mol-JEPA, compared with 47% for TabICL with AlvaDesc descriptors, while Mol-JEPA achieved more significant wins on smaller ASAP and Biogen benchmarks.The results suggest particular benefits in low-data settings.
  • Modality scaling: Mol-JEPA training time increased from roughly 9 hours with two modalities to about 13 hours with four and approximately 27 hours with eight on 100k molecules for 300 epochs.The authors report approximately linear scaling with modality count, also depending on modality and input dimensionality.
  • Modality analysis: Removing one modality at inference caused only minor performance differences, suggesting that the learned representations contain redundant information across modalities.The authors associate this redundancy with the reconstruction objective.
  • Multimodal training dynamics: Effective rank increased consistently across modality-specific embeddings, whereas prediction and SIGReg losses showed heterogeneous and sometimes unstable behavior across modalities.Sparse modalities were generally harder to optimize, and including CLOOME could decrease downstream performance.

5 CONCLUSION

Mol-JEPA predicts latent representations across many molecular modalities and achieves strong downstream performance. The model performs particularly well on smaller datasets and out-of-distribution settings, while modality selection remains an open direction.

  • 5 CONCLUSION: Mol-JEPA enables stable learning from many molecular modalities by predicting their latent representations.The authors evaluate the model across downstream tasks and analyze modality effects on training dynamics and performance.
  • 5 CONCLUSION: Mol-JEPA shows particularly strong performance on smaller datasets and in out-of-distribution settings.The conclusion frames this result in the context of limited data availability in drug discovery.
  • 5 CONCLUSION: The selection of the most effective modality set remains an important direction for future research.The authors also describe the approach as potentially applicable across other scientific domains.

6 LIMITATIONS

The paper identifies limitations in modality reconstruction, out-of-distribution robustness, understanding multimodal learning dynamics, and the scale of ablation studies.

  • 6 LIMITATIONS: The reconstruction objective does not account for modalities that may be unable to recover information contained in other modalities.The authors propose modality-specific dependencies and advanced masking strategies as possible improvements.
  • 6 LIMITATIONS: Out-of-distribution generalization remains a key challenge when molecules, tasks, or data distributions differ substantially from pretraining.The authors identify robustness in such settings as a future research need.
  • 6 LIMITATIONS: It remains unclear how individual modalities and factors such as data sparsity, dimensionality, information content, and redundancy affect representation learning.The authors connect resolving these questions to more effective modality selection and multimodal training strategies.
  • 6 LIMITATIONS: Most ablation studies used a considerably smaller dataset than the full pretraining data because of computational constraints, potentially affecting some conclusions.This limitation applies specifically to the scale of the ablation experiments.

A DATASET DETAILS

The dataset details describe processing of experimental assay data, benchmark transformations, and cluster-based dataset splitting. These procedures address assay coverage, sparsity, endpoint aggregation, and structural overlap between splits.

  • ChEMBL processing: ChEMBL processing pivots bioassay activity data into molecule-level assay-measurement vectors while balancing assay coverage against data sparsity.Increasing the minimum measurements required per assay reduces the number of retained endpoints.
  • Therapeutics Data Commons processing: Therapeutics Data Commons measurements are aggregated at the molecule level into feature vectors, with several datasets producing 672-dimensional vectors.Available measurements vary substantially across datasets and endpoints.
  • Benchmark data processing: Benchmark datasets are used in the organizers’ provided form, with ExpansionRx endpoints additionally transformed using log10(x + 0.001).The transformation matches competition preprocessing and adds an epsilon of 0.001 before taking the logarithm.
  • Dataset splitting: Cluster-based splitting moves selected clusters into test data while retaining the remaining clusters for training across benchmark datasets.Figure 7 visualizes the three splits for the OpenADMET LogD dataset.
  • Dataset splitting: Butina clustering at a Tanimoto similarity threshold of 0.65 reduces structural overlap between training and test sets but does not strictly enforce a maximum cross-split similarity of 0.65.The threshold therefore reduces, rather than completely bounds, structural similarity across splits.

B IMPLEMENTATION DETAILS

Mol-JEPA tuning explores architectural, masking, loss-balancing, probing, and training choices, with performance sensitive to λ but not consistently related to masking ratio.

  • Hyperparameter tuning: Hyperparameter optimization evaluated 300 randomly sampled configurations on 100,000 molecules and selected the model with the best online probe performance.The search covered encoder dimensions, aggregation structure, training parameters, masking ratios, and λ.
  • Hyperparameter tuning: Small values of λ deteriorate performance, whereas masking ratio shows no consistent relationship with downstream performance.These findings come from the reported hyperparameter sensitivity analysis.
  • Probing model details: Mol-JEPA uses a 512-dimensional CLS embedding as a global molecular representation for linear or nonlinear predictors, with TabICLv2 as an expressive alternative.The probing setup also evaluates modality-specific predicted embeddings.
  • Computational requirements: Training uses 12 RTX-6000ADA GPUs, an aggregated batch size of 1024, and converges after 3 days at approximately 150 epochs before overfitting.Online probe performance degrades after convergence.

D RESULTS AND MODEL EVALUATION

Mol-JEPA performs strongly across benchmark evaluations, including R2-based comparisons and cluster splits, while public temporal splits reveal room for improvement against TabICLv2 on some metrics.

  • Benchmark results: Mol-JEPA achieves the strongest overall average performance across datasets under the R2 evaluation metric.Negative R2 values are clipped to zero in the comparison table.
  • Public competition splits: TabICLv2 with AlvaDesc achieves higher Wilcoxon and MAE win rates on public temporal splits, while Mol-JEPA delivers the strongest R2 performance.The public splits are from ExpansionRX, ASAP, and PXR.
  • Public competition splits: Mol-JEPA consistently surpasses other foundation-model approaches on both public and cluster-based splits.The authors identify larger pretraining datasets and additional modalities as areas for further improvement.

E PRETRAINING DATA SCALING

Increasing the pretraining dataset size steadily improves downstream performance across datasets for both linear and nonlinear probes.

  • Pretraining data scaling: Pretraining performance steadily improves as dataset size increases for both linear and nonlinear probes across all downstream datasets.Performance is reported as 1-MAE, where higher values indicate better performance.

F ABLATION STUDIES

Ablations show that latent prediction, target-space regularization, a 512-dimensional embedding, and multimodal inputs are important design choices for Mol-JEPA performance.

  • Ablation setup: Experiments on 100,000 randomly sampled data points assess sensitivity to Mol-JEPA parameter and architectural choices.The subset is used for computational reasons.
  • Experimental modalities: Experimental modalities moderately affect performance, partly because missing values, imbalanced distributions, and noisy measurements limit extractable information.Despite these constraints, removing experimental modalities consistently worsens downstream performance.
  • Experimental modalities: JEPA latent prediction outperforms direct semi-supervised modality-label prediction, reducing MAE from 0.467 to 0.436 for linear probing and from 0.417 to 0.393 for nonlinear probing.The comparison supports latent prediction as a more robust strategy for noisy and incomplete experimental data.
  • Regularization: Applying SIGReg to target embeddings yields stable training and the best downstream performance, whereas applying it to predicted embeddings causes training collapse.The result distinguishes target-space regularization from alternative placements.
  • Embedding dimension: Increasing embedding dimension from 256 to 512 reduces MAE from 0.483 to 0.436 for linear probing and from 0.465 to 0.393 for nonlinear probing.Increasing dimension to 1024 yields nearly identical results, indicating diminishing returns.
  • Multimodal benefits: Using all modalities reduces MAE by 14% for linear probing and 13% for nonlinear probing compared with using only molecular graphs and ECFP4 fingerprints.The controlled comparison uses the full pretraining dataset and CLS-based probes.
  • Representation similarity analysis: CKA analysis forms clusters among related modalities, with quantum-chemistry modalities showing high mutual similarity and cellular-profiling modalities forming a distinct cluster.The analysis links embedding similarity to shared physicochemical or biological relationships while retaining modality-specific information.
Loading 2608.22642v1…