File size: 4,033 Bytes
6d81eec
 
 
 
b408efc
 
 
6d81eec
 
 
 
1c275d6
 
6d81eec
 
 
dc319bd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1c275d6
dc319bd
 
 
1c275d6
dc319bd
 
1c275d6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
dc319bd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
---
datasets:
- Zatom-AI/qm9
- Zatom-AI/mp_20
- Zatom-AI/geom
- Zatom-AI/matbench
- Zatom-AI/qmof
- Zatom-AI/omol25
- Zatom-AI/mptrj
language:
- en
license: mit
pipeline_tag: other
tags:
- chemistry
- biology
- foundation-model
- generative-model
- predictive-model
- representation-learning
- transformer
- molecule
- material
- property
- energy
- forces
- mlip
---

<div align="center">

# Zatom-1

[![Paper](https://img.shields.io/badge/arXiv-2602.22251-b31b1b.svg)](https://arxiv.org/abs/2602.22251)

<a href="https://arxiv.org/abs/2602.22251"><img src="zatom_1.png" width="600"></a>

</div>

This repository contains the model weights for **Zatom-1**, the first end-to-end foundation model that unifies generative and predictive learning of 3D molecules and materials. Introduced in [Zatom-1: A Multimodal Flow Foundation Model for 3D Molecules and Materials](https://huggingface.co/papers/2602.22251), Zatom-1 is a Transformer trained with a multimodal flow matching objective that jointly models discrete atom types and continuous 3D geometries.

## GitHub repository

For the full implementation, training scripts, and configuration files, visit:
https://github.com/Zatom-AI/zatom

## Sample Usage

### Installation

To get started, clone the repository and install the dependencies:

```bash
# Clone project
git clone https://github.com/Zatom-AI/zatom
cd zatom

# Install requirements
pip install -e '.[cuda]'
```

### Evaluation

To generate evaluation metrics for molecule and material generation using the Zatom-1 weights:

```bash
python zatom/eval_fm.py \
    ckpt_path=checkpoints/zatom_1_joint_paper_weights.ckpt \
    model.sampling.num_samples=10000 \
    model.sampling.batch_size=1000 \
    name=eval_run \
    seed=42 \
    trainer=gpu
```

## Open-source resources

Zatom-1 builds upon the source code and data from the following projects:

- [all-atom-diffusion-transformer](https://github.com/facebookresearch/all-atom-diffusion-transformer)
- [flow_matching](https://github.com/facebookresearch/flow_matching)
- [flowmm](https://github.com/facebookresearch/flowmm)
- [jvp_flash_attention](https://github.com/amorehead/jvp_flash_attention)
- [lemat-genbench](https://github.com/LeMaterial/lemat-genbench)
- [lightning-hydra-template](https://github.com/ashleve/lightning-hydra-template)
- [PlatonicTransformers](https://github.com/niazoys/PlatonicTransformers)
- [ProteinWorkshop](https://github.com/a-r-j/ProteinWorkshop)
- [posebusters](https://github.com/maabuu/posebusters)
- [tabasco](https://github.com/carlosinator/tabasco)

We thank all their contributors and maintainers!

## Acknowledgements

This research used resources of the National Energy Research Scientific Computing Center, a DOE Office of Science User Facility supported by the Office of Science of the U.S. Department of Energy under Contract No. DE-AC02-05CH11231 using AI4Sci@NERSC award NERSC DDR-ERCAP0036206 awarded to AM. NBE would like to acknowledge support from the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, EXPRESS: 2025 Exploratory Research for Extreme-Scale Science program, and the Scientific Discovery through Advanced Computing (SciDAC) program, under Contract Number DE-AC02-05CH11231 at Berkeley Lab.

## Citation

If you use the code or data associated with this package or otherwise find this work useful, please cite:

```bibtex
@article{zatom_1_2026,
    title={Zatom-1: A Multimodal Flow Foundation Model for 3D Molecules and Materials},
    author={Alex Morehead* and Miruna Cretu* and Antonia Panescu* and Rishabh Anand* and Maurice Weiler* and Tynan Perez* and Samuel Blau and Steven Farrell and Wahid Bhimji and Anubhav Jain and Hrushikesh Sahasrabuddhe and Pietro Liò and Tommi Jaakkola and Rafael Gómez-Bombarelli and Rex Ying* and Ben Erichson* and Michael Mahoney*},
    year={2026},
    eprint={2602.22251},
    archivePrefix={arXiv},
    primaryClass={cs.LG},
    url={https://arxiv.org/abs/2602.22251},
    note={* denotes equal contribution}
}
```