JEPA · 2026-10-07

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

Kasper Helverskov Petersen, Rasmus Hannibal Tirsgaard, Fran\c{c}ois R J Cornet, Mikkel Jordahn, Mikkel N. Schmidt

AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.

TL;DR

Atom-JEPA is a self-supervised pretraining framework for 3D atomistic systems that learns latent representations from unlabeled molecular and crystalline structures using complementary atom-level and substructure-level latent-space prediction objectives inspired by JEPA. The authors pretrain on large molecular (Uni-Mol) and crystal (Alexandria) datasets and fine-tune the pretrained context encoder for downstream property prediction.

Why it matters

The authors address limited downstream generalization in 3D atomistic self-supervised pretraining by using a JEPA-style latent prediction approach with an SE(3)-equivariant GNN. They report state-of-the-art performance on molecular ADMET and quantum-chemical property prediction tasks and competitive results on crystalline material properties, and release code and pretrained checkpoints.

Method

Use two JEPA encoders (context updated by gradients; target updated by EMA) with an SE(3)-equivariant GNN; predict target latent representations from context latent representations in two prediction directions.
Train with two latent-space objectives: an atom-level loss predicting each target atom representation and a substructure-level loss predicting mean-pooled target representations.
Create context/target views by partitioning a k-hop EgoNet around an anchor atom (k sampled from a predefined set) and taking the complement as the other view.

Limitation

The authors restrict training data to structures that are close to local minima of the potential energy surface; including off-equilibrium structures without information about relative stability (e.g., total energy) could make learning ambiguous.

Abstract (from arXiv)

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

Publication record

Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.

First public release: 2026-10-06 · Indexed by Pub-AI: 2026-10-07 UTC

preprint