Atom-jepa improves prediction of molecular and crystal properties
Atom-JEPA: Joint-Embedding Predictive Architecture for 3D Atomistic Systems
Machine LearningArtificial Intelligence
Summary
Predicting the properties of molecules and materials from their 3D structures is important but challenging. The authors present Atom-JEPA, a self-supervised learning method that trains on lots of unlabeled 3D data to understand atomic and substructure features jointly. This approach helps the model learn useful representations that improve its accuracy when later predicting different chemical and physical properties. Their results show better performance on various molecular and crystalline tasks compared to previous methods.
What this means in practice
- •For chemical engineers: Use pre-trained Atom-JEPA models to more accurately predict molecular properties relevant to drug design and materials engineering tasks.
- •For materials scientists: Employ Atom-JEPA for improved prediction of the physical properties of crystalline materials to aid in material discovery and optimization.
Authors
Kasper Helverskov Petersen, Rasmus Hannibal Tirsgaard, François R J Cornet, Mikkel Jordahn, Mikkel N. Schmidt
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