Atom-JEPA: Joint-Embedding Predictive Architecture for 3D Atomistic Systems

2026-10-06

AI Summary

This study addresses the limited downstream generalization of self-supervised pretraining for atomic systems by proposing Atom-JEPA, a framework that introduces the Joint-Embedding Predictive Architecture to the atomic domain for the first time. Leveraging large-scale molecular and crystal data, the method performs self-supervised pretraining through dual atomic- and substructure-level prediction objectives in latent space, thereby learning high-quality representations from unlabeled 3D structures. Experimental results demonstrate that the proposed model achieves state-of-the-art or highly competitive performance across a broad spectrum of downstream tasks, including ADMET profiling, quantum chemistry, and crystal property prediction. By effectively capturing transferable structural priors, Atom-JEPA significantly enhances cross-domain generalization capabilities for structure-based molecular and materials modeling.

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

Problem

• self-supervised pretraining
• atomistic systems
• downstream generalization
• 3D structures

Innovation

• Self-supervised pretraining
• Joint-embedding predictive architecture
• 3D atomistic systems
• Latent representation
• Downstream generalization