Zatom-1: Towards a Multimodal Foundation Model for 3D Molecules and Materials
Authors: Alex Morehead, Miruna Cretu, Antonia Panescu, Rishabh Anand, Maurice Weiler, Tynan Perez, Samuel Blau, Steven Farrell, Wahid Bhimji, Anubhav Jain, Hrushikesh Sahasrabuddhe, Pietro Lio, Tommi Jaakkola, Rafael Gomez-Bombarelli, Rex Ying, N. Benjamin Erichson, Michael W. Mahoney
Organizations: LBNL · ICSI · University of Cambridge · Yale University · MIT · UC Berkeley
General-purpose 3D modeling in chemistry encompasses molecules and materials, requiring both generative and predictive capabilities. However, most existing AI approaches are optimized for a single domain (molecules or materials) and a single task (generation or prediction), which limits representation sharing and transfer. We introduce Zatom-1, a cross-domain, general-purpose model architecture that unifies generative and predictive learning of 3D molecules and materials. Zatom-1 is a deliberately simplified Transformer trained with a multimodal flow matching objective that jointly models discrete atom types and continuous 3D geometries. This approach supports scalable pretraining with predictable gains as model capacity increases, while enabling fast and stable sampling. We use cross-domain generative pretraining as a universal initialization for downstream multi-task prediction of properties, energies, and forces. Empirically, Zatom-1 outperforms or competes with specialized baselines on both multi-task generative and predictive benchmarks in data-controlled settings, while improving generative inference speed by more than an order of magnitude. Our experiments demonstrate positive predictive transfer between data domains from joint generative pretraining: modeling materials during generative pretraining improves molecular property prediction accuracy. Open-source code and model weights are freely available at https://github.com/Zatom-AI/zatom.
Figures & tables
Figure 1: The Zatom-1 architecture for 3D molecules and materials. A unified Transformer trunk processes (noisy) atomic, structural, and conditional inputs. The model features two training stages: (1) multimodal flow pretraining using the final trunk layer L ’s representations to create new molecules or materials; and (2) multi-task finetuning of additional downstream task backbones and heads that use a specified trunk layer K ’s representations to predict properties, energies, and forces.
Figure 2
Model
# Params
Valid% ↑
Unique% ↑
Novel% ↑
Energy-based
Stability
Metastability
Ef (Std) ↓
Ehull (Std) ↓
RMSD (Std) ↓
Stable% ↑
⋆ SUN% ↑
Metastable% ↑
⋆ MSUN% ↑
With relaxation
Crystalite
67M
97.2
95.8
53.2
−0.89
0.09
0.13
12.7
1.5 (38)
51.6
22.6
MatterGen
47M
95.7
95.1
70.5
−0.70±0.79
0.18 ± 0.18
0.39±0.50
2.0
0.2 (6)
33.4
15.0
PLaID++
7B
96.0
77.8
24.2
−0.50±0.44
0.09 ± 0.16
0.13 ± 0.29
12.4
1.0 (26)
60.7
7.6
WyFormer
8B
93.4
93.0
66.4
−0.43±0.95
0.50±0.51
0.81±0.98
0.5
0.1 (3)
15.7
1.9
Table 1: Crystal generation results on LeMat-GenBench with MP20 training. Core model evaluation metrics are calculated for 2,500 sampled crystals using an MLIP ensemble with LeMat-Bulk as the reference set. Best values are shown in bold , second-best values are underlined . ⋆ denotes aggregate metrics that, summed together, constitute the primary objective to optimize for this benchmark (i.e., SUN + MSUN). Notably, jointly trained Zatom-1 (80M) achieves its results using ∼ 400 GPU training hours, 3x faster than jointly trained ADiT (180M). Further, jointly trained Zatom-1-WD achieves state-of-the-art SUN (and MSUN) rates for materials discovery in both relaxed (and unrelaxed) method categories [ 24 ] through aggressive weight decay ( 1e−2 ) and early stopping ( 1/5 into training, before molecule validation performance has fully converged).
Model
# Params
Validity (%) ↑
Unique (%) ↑
One-stage training
Equivariant Diffusion
20M
91.90
98.69
Symphony
2M
83.50
97.98
Two-stage training
GeoLDM
20M
93.80
98.82
QM9-only ADiT
180M
92.19
97.90
Table 2: Molecule generation results on QM9. We report (a) validity and uniqueness rates as well as (b) % pass rates on 7 sanity checks from PoseBusters for 10,000 sampled molecules. All models explicitly generate hydrogen atoms. Notably, jointly trained Zatom-1 (80M) achieves its results using ∼ 400 GPU training hours compared to jointly trained ADiT (180M), which requires ∼ 1,200 GPU pretraining + finetuning hours.
Figure 2: Zatom-1 is significantly faster than equivariant and latent diffusion models. We plot the number of integration steps for Zatom-1 , equivariant diffusion FlowMM/SemlaFlow, and latent diffusion ADiT/GeoLDM vs. time to generate 10,000 samples on a single NVIDIA GPU. Zatom-1 scales better with the number of integration steps compared to equivariant and latent diffusion models.
Metric (% pass) ↑
EQGAT-diff
SemlaFlow
ADiT
TABASCO
Zatom-1
GEOM-Drugs
Validity
94.6
93.9
95.3
97.6
93.6
-
Uniqueness
100.0
100.0
100.0
99.03
99.93
-
Atoms connected
84.4
92.3
93.0
99.9
99.6
-
Bond angles
86.9
94.8
92.3
99.2
99.6
-
Bond lengths
87.0
94.6
92.5
99.4
99.7
-
Aromatic ring flat
87.0
94.9
95.4
100.0
100.0
-
Table 3: Molecule generation results on GEOM-DRUGS. Validity, uniqueness, and % pass rates on PoseBusters for 10,000 sampled molecules. Zatom-1 matches or exceeds state-of-the-art equivariant/latent diffusion baselines and generates the largest number of PoseBusters-valid samples.
Model / Metrics ↓
# Params
α ( a03 )
Δε (meV)
εHOMO (meV)
εLUMO (meV)
μ (D)
Cν (cal/mol K)
GATOM (meV)
HATOM (meV)
R2 ( a02 )
UATOM (meV)
U0ATOM (meV)
ZPVE (meV)
Optimized single-task learning
DimeNet++ †
5M
.044
33.0 / 199
25
20
.030
.023
8.00
7.00
.331
6.00
6.00
1.21
EGNN †
1M
.071
48.0 / —-
29
25
.029
.031
12.0
12.0
.106
12.0
11.0
1.55
PaiNN
1M
.045
46.0 / —-
28
20
.012
.024
7.35
5.98
.066
5.83
5.85
1.28
TorchMD-NET
7M
.059
36.0 / —-
20 (16 w/ NN)
18 (13 w/ NN)
.011
.026
7.62
6.16
.033
6.38
6.15
1.84
SphereNet
2M
.046
32.0 / —-
23
18
.026
.021
8.00
6.00
.292
7.00
6.00
1.12
Table 4: Property prediction for 3D molecules (QM9) / materials (Matbench), with zero-shot multi-task material predictions. Results are reported as test mean absolute errors (for Zatom-1 , with standard deviations over three runs with different random seeds). † denotes using different data partitions. / represents a task for which both (QM9) molecule and (Matbench) material property labels are available. NN denotes pretraining using a Noisy Nodes self-supervised objective [ 118 ] . The best (or second-best) results for each optimized (i.e., hyperparameter-tuned) or unoptimized (i.e., non-hyperparameter-tuned) category are in bold (or underlined ). Notably, Zatom-1 (80M + 20M ) and AIM-MTL are the only multi-task learning (MTL) methods represented in these results, which they achieved with only ∼ 100 GPU hours of additional finetuning for molecule property prediction, signifying their computational efficiency and expressivity. In contrast, single-task learning (STL) baselines such as EquiformerV2 (11M) traditionally require ∼ 150 GPU training hours for each task , not including their extensive hyperparameter optimization per task. Further, for the majority of QM9’s properties, jointly pretrained Zatom-1 offers the first example of positive transfer learning between molecules and materials for prediction tasks driven by cross-domain generative pretraining.
Appendix figures & tables17 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 4: PCA plots of latent embeddings from Zatom-1 ’s Transformer trunk for 100 data points from the MP20 and QM9 datasets, as well as 100 molecules/materials each generated by Zatom-1 . Each point denotes an atom (including hydrogens for molecules), with its shape representing its chemical system type and its primary color indicating whether it comes from real or generated data. (a) Embeddings taken from an untrained trunk are distributed in an overly clustered and chemically-segmented manner. (b-c) Embeddings taken from a trunk pretrained on either QM9 or MP20 form less well-structured molecule and material clusters. (d) Embeddings taken from a jointly pretrained trunk, however, show structured yet overlapping clusters, suggesting positive representation transfer between molecules and materials has occurred .
Figure 5: PCA plots of latent embeddings for carbon, nitrogen, oxygen, and fluorine atoms from Zatom-1 ’s Transformer trunk for 100 data points from the MP20 and QM9 datasets, as well as 100 molecules/materials each generated by Zatom-1 . Each point denotes an atom, with its shape representing its chemical system type, its size defined by whether it comes from real or generated data, and its color indicating its element type. (a) Embeddings taken from an untrained trunk are distributed in an overly clustered and chemically-segmented manner, with a quantitatively high crystal-molecule shift between atom embeddings of the same element and source type. (b-c) Embeddings taken from a trunk pretrained on either QM9 or MP20 form less well-structured molecule and material clusters, with a quantitatively higher intra-element spread and crystal-molecule shift. (d) A jointly pretrained model, however, produces trunk embeddings that achieve the best quantitative balance across all clustering and transfer learning metrics. Moreover, in such a model, principal component 1 largely relates to whether a system is a crystal (within range -2 – 4) or a molecule, while principal component 2, generally speaking, relates to an atom’s type. Joint pretraining induces a latent space that shows distinct clusters for different atom types, with oxygen (and, to some extent, fluorine) atoms having similar representations in both molecules and crystals. Notably, the overlap in oxygen atom representations suggests that Zatom-1 ’s latent space captures shared chemical properties across non-periodic and periodic systems, enabling effective knowledge transfer during joint pretraining.
Figure 6: Generated crystals from Zatom-1 trained jointly on MP20 materials and QM9 molecules.
Figure 7: Generated molecules from Zatom-1 trained jointly on MP20 materials and QM9 molecules.
Figure 8: Generated metal-organic frameworks from Zatom-1 trained on QMOF.
Figure 9: Tuning inference hyperparameters for best performance. Best generative modeling results for crystals and molecules, balancing sample validity and uniqueness, are generally achieved with a limited number of integration steps (e.g., T=100 ), a high white noise scale γg (e.g., 50) for continuous, non-periodic atom positions, and a low noise scale (e.g., 0.01) for continuous, periodic modalities such as fractional coordinates.
Hyperparam.
Zatom-1-WD (80M)
Zatom-1 (80M)
Zatom-1-L (160M)
Zatom-1-XL (300M)
# Params.
77,350,264
77,350,264
162,615,544
293,887,096
Hidden size
512
512
768
1024
# Transformer blocks
16
16
16
16
# Attn. heads
8
8
8
8
Train t distribution
Beta(1.8,1)
Beta(1.8,1)
Beta(1.8,1)
Beta(1.8,1)
λdiscrete
0.1
0.1
0.1
0.1
Appendix
Table 5: Hyperparameters across different model versions when jointly pretraining on QM9 & MP20.
Hyperparam.
GEOM-Drugs & MP20
QMOF
OMol25
# Params.
77,350,264
77,350,264
77,350,264
Hidden size
512
512
512
# Transformer blocks
16
16
16
# Attn. heads
8
8
8
Train t distribution
Beta(1.8,1)
Beta(1.8,1)
Beta(1.8,1)
λdiscrete
0.1
0.1
0.1
Appendix
Table 6: Hyperparameters when jointly pretraining Zatom-1 on GEOM-Drugs, QMOF, or OMol25.
Hyperparam.
Mol. Properties (80M + 20M ) →
Mol. / Mat. Properties (80M + 20M ) →
Mol. / Mat. Energy & Forces (80M + 220M )
# Dedicated params.
18,104,595
18,104,595
217,255,140
Aux. hidden size
512
512
1024
# Aux. transformer blocks
4
4
8
# Attn. heads
8
8
8
Train t distribution
max(Beta(1.8,1), 0.98)
max(Beta(1.8,1), 0.98)
max(Beta(1.8,1), 1.0)
λforces
N/A
N/A
5.0
Appendix
Table 7: Finetuning hyperparameters across different downstream prediction tasks for Zatom-1 (80M), optionally trained in sequential order.
Linear
G -linear
Attention
G -attention
Compute FLOPS
O(ND2)
O(N∣G∣2C2)
O(N2D)
O(N2∣G∣C)
Memory
O(ND)
O(N∣G∣C)
O(ND)
O(N∣G∣C)
Parameters
D2
∣G∣⋅C2
—
—
Appendix
Table 8: Compute, memory, and parameter complexity of TFT and TFP. TFT and TFP operate on tensors of shapes (N,∣G∣,C) and (N,D) , respectively. Memory costs follow immediately from these shapes. Linear layers in TFT come with the usual quadratic number of FLOPS and parameters in their D channels. TFP is internally leveraging nn.Linear layers as well, however, after flattening features to shape (N,∣G∣⋅C) , implying a computational cost quadratic in ∣G∣ as well. Equivariance imposes a symmetry constraint on the matrices (weight sharing), resulting in a ∣G∣ times-enhanced parameter efficiency. TFT’s attention layers have the usual computational complexity linear in channels D , while those of TFP are linear in ∣G∣C . Neither introduces additional parameters.
Figure 10: QM9-only pretraining of Zatom-1 vs. Platom-1 . These plots show the training dynamics of Zatom-1 compared to the analogous tetrahedral group equivariant Platom-1 model described in Section E.1 . Due to its equivariance, the Platonic model does not explicitly need to learn symmetries, resulting in significantly faster convergence while achieving improved evaluation metrics. These findings translate to the test results in Table 9 .
Test%
QMOF ADiT
QMOF Zatom-1 *
Has carbon ↑
100.0
100.0
Has hydrogen ↑
99.6
100.0
Has atomic overlap ↓
8.3
8.8
Has overcoord. C ↓
23.6
1.1
Has overcoord. N ↓
1.5
0.0
Has overcoord. H ↓
1.0
2.9
Appendix
Table 10: Metal-organic framework generation results on QMOF. We report sanity checks from MOFChecker [ 48 ] for 1,000 sampled MOFs, comparing methods trained only on QMOF for at least 10,000 epochs. ( ↑ / ↓ indicate higher/lower is better, respectively, and * denotes a model that has not converged after 50,000 training epochs)
Metric (% pass) ↑
OMol25 Zatom-1
Validity
30.4
Uniqueness
93.5
Atoms connected
17.0
Bond angles
78.9
Bond lengths
54.6
Aromatic ring flat
100.0
Appendix
Table 11: Molecule generation results on OMol25. Validity, uniqueness, and % pass rates on PoseBusters for 16,000 sampled molecules. Zatom-1 , pretrained on the (non-equilibrium) molecules of OMol25, takes a first step toward non-equilibrium 3D molecule generation.
Model / Metrics ↓
# Params
α ( a03 )
Δε (meV)
εHOMO (meV)
εLUMO (meV)
μ (D)
Cν (cal/mol K)
GATOM (meV)
HATOM (meV)
R2 ( a02 )
UATOM (meV)
U0ATOM (meV)
ZPVE (meV)
Optimized single-task learning
DimeNet++ †
5M
.044
33.0 / 199
25
20
.030
.023
8.00
7.00
.331
6.00
6.00
1.21
EGNN †
1M
.071
48.0 / —-
29
25
.029
.031
12.0
12.0
.106
12.0
11.0
1.55
PaiNN
1M
.045
46.0 / —-
28
20
.012
.024
7.35
5.98
.066
5.83
5.85
1.28
TorchMD-NET
7M
.059
36.0 / —-
20 (16 w/ NN)
18 (13 (w/ NN)
.011
.026
7.62
6.16
.033
6.38
6.15
1.84
SphereNet
2M
.046
32.0 / —-
23
18
.026
.021
8.00
6.00
.292
7.00
6.00
1.12
Appendix
Table 12: Property prediction for 3D molecules (QM9) / materials (Matbench) with joint molecule and materials finetuning. Results are reported as test mean absolute errors (for Zatom-1 , with standard deviations over three runs with different random seeds). † denotes using different data partitions. / represents a task for which both (QM9) molecule and (Matbench) material property labels are available. NN denotes pretraining using a Noisy Nodes self-supervised objective [ 118 ] . The best (or second-best) results for each optimized (i.e., hyperparameter-tuned) or unoptimized (i.e., non-hyperparameter-tuned) category are in bold (or underlined ). Notably, finetuning jointly for multi-task molecule and ( Δε ) materials property prediction improves or marginally trades off Zatom-1 ’s molecule property prediction performance for expanded applicability of its property predictions to both molecules and materials simultaneously.
Model
# Params
Energy MAE/Atom (meV) ↓
Energy MAE (meV) ↓
Force MAE (meV/Å) ↓
Material Interatomic Potential Prediction on MPtrj
Non-pretrained Orb-v1
∼ 30M
-
350.00
122.00
Denoising-pretrained Orb-v1
∼ 30M
-
210.00
32.00
Jointly pretrained Zatom-1
80M + 220M
112.81
2631.22
30.92
Molecule Interatomic Potential Prediction on OMol25-4M
eSEN-md-d.
∼ 51M
1.32
77.13
6.78
Appendix
Table 13: Energy and force prediction for 3D materials (MPtrj) / molecules (OMol25-4M). The results are reported as validation Energy (in meV or meV/Atom) and Force (in meV/Å) mean absolute error (MAE). All models directly predict forces, rather than using the gradient of their predicted energies with respect to input 3D coordinates for conservative force prediction, to considerably increase their computational efficiency. * denotes a model that was pretrained on OMol25 and which has not converged after 80 training epochs.
Fine-Tuning Strategy
Validity ↑
Uniqueness ↑
Novelty ↑
PoseBusters Valid ↑
Trunk Unfreezing
0.17
0.33
0.33
0.00
LoRA
0.48
1.00
1.00
0.61
Trunk Freezing (Baseline)
0.95
0.97
0.96
0.99
Appendix
Table 14: QM9 molecule generation results after 20 epochs of QM9 multi-task property prediction finetuning.
Metric
ADiT
TABASCO
Zatom-1 (Ours)
GEOM-Drugs
Validity ↑
0.95
0.98
0.94
–
Uniqueness ↑
1.00
0.99
1.00
–
Diversity ↑
0.91*
0.89
0.89
0.90
Novelty ↑
0.97*
0.99
0.98
0.00
QED ↑
–
0.64
0.64
0.67
Lipinski ↑
–
4.90
4.80
5.00
Appendix
Table 15: Comprehensive GEOM-Drugs molecule generation results. – denotes unavailable results, and * indicates results were estimated from 1,000 samples per Vonessen et al. [103] .
Unified atomistic modeling has the potential to accelerate discovery in chemistry, materials science, and biology by bridging data-rich chemical domains and data-scarce biological contexts. However, existing generative approaches to atomistic modeling remain highly specialized to scientific disciplines (chemistry vs. biology) or do not leverage both high-volume organic (molecule) and inorganic (material) data for general-purpose pretraining. To this end, we introduce Zatom-2, an atomistic generative model pretrained on approximately five million structures from the OMol25 and OMat24 electronic structure datasets. Zatom-2 features a multiscale Transformer architecture coupled with conditional flow matching that supports force conditioning and foundational pretraining tasks such as generation, structure prediction, and prediction of molecular and material energies and forces. Empirically, Zatom-2 achieves better molecular distribution fidelity than Zatom-1 and achieves strong performance on existing molecule and material generation benchmarks. Zatom-2 demonstrates the ability to control sample generation across low- and high-force regimes, and enhances protein generation in a low-data setting through joint generative-predictive pretraining and transfer learning, increasing protein backbone designability in a length extrapolation setting from 67.8% without pretraining to 74.8% after finetuning on 2,000 protein domains.
Miruna Cretu, Alex Abrudan, Antonia Panescu +11
University of Cambridge · Yale University · MIT +3
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
Kasper Helverskov Petersen, Rasmus Hannibal Tirsgaard, François R J Cornet +2
Department of Applied Mathematics and Computer Science Technical University of Denmark
Most machine learning models for molecular property prediction rely on a single molecular representation (either a sequence, a graph, or a 3D structure) and treat molecular geometry as static. We present MolFM-Lite, a multi-modal model that jointly encodes SELFIES sequences (1D), molecular graphs (2D), and conformer ensembles (3D) through cross-attention fusion, while conditioning predictions on experimental context via Feature-wise Linear Modulation (FiLM). Our main methodological contributions are: (1) a conformer ensemble attention mechanism that combines learnable attention with Boltzmann-weighted priors over multiple RDKit-generated conformers, capturing the thermodynamic distribution of molecular shapes; and (2) a cross-modal fusion layer where each modality can attend to others, enabling complementary information sharing. We evaluate on four MoleculeNet scaffold-split benchmarks using our model's own splits, and report all baselines re-evaluated under the same protocol. Comprehensive ablation studies across all four datasets confirm that each architectural component contributes independently, with tri-modal fusion providing 7-11% AUC improvement over single-modality baselines and conformer ensembles adding approximately 2% over single-conformer variants. Pre-training on ZINC250K (~250K molecules) using cross-modal contrastive and masked-atom objectives enables effective weight initialization at modest compute cost. We release all code, trained models, and data splits to support reproducibility.
Syed Omer Shah, Mohammed Maqsood Ahmed, Danish Mohiuddin Mohammed +2
Department of Computer Science and Engineering, University at Buffalo, Buffalo, NY, USA · Khoury College of Computer Sciences, Northeastern University, Boston, MA, USA · Department of Computer Science, Muffakham Jah College of Engineering and Technology, Hyderabad, India +1