File size: 5,616 Bytes
abde3bc
 
7438c3c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
abde3bc
7438c3c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
48faf77
7438c3c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f20ba73
7438c3c
 
 
 
 
 
 
 
 
 
 
 
b3ad457
7438c3c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
---
license: mit
language:
- en
tags:
- chemistry
- biology
- generative-model
- predictive-model
- representation-learning
- transformer
- molecule
- material
- protein
- property
- energy
- forces
- mlip
---

<div align="center">

# Zatom-2

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

<a href="https://arxiv.org/abs/2610.11454"><img src="zatom_2.png" width="600" alt="Zatom-2 Architecture and Domain Capabilities"></a>

</div>

This repository contains the model weights for **Zatom-2**, a unified atomistic model for generative modeling and representation learning across small molecules, periodic materials, and proteins. Introduced in [Zatom-2: Multitask Pretraining on Atomistic Data for Generative Modeling across Domains](https://huggingface.co/papers/2610.11454), Zatom-2 features a multiscale Transformer architecture coupled with conditional flow matching that supports force conditioning and foundational pretraining tasks—including 3D unconditional generation, structure prediction, and prediction of molecular and material energies and interatomic forces.

## GitHub repository

For the full implementation, training scripts, configuration files, and evaluation pipelines, visit:
- **Repository:** https://github.com/Zatom-AI/nucleus
- **Project Page:** https://zatom-ai.github.io/nucleus

## Sample Usage

### Installation

To get started, clone the repository and install the dependencies with [uv](https://docs.astral.sh/uv/) (recommended) or `pip`:

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

# Install using uv (recommended)
uv sync --locked --extra all
source .venv/bin/activate

# Or install from PyPI
pip install zatom-nucleus zatom2
# (or pip install "zatom-nucleus[all]" to include Zevals evaluation suites)
```

### Downloading Checkpoints

Download the primary Zatom-2 model weights (large joint-properties model) using the `nucleus` CLI:

```bash
# Discover available checkpoints
nucleus list-available

# Install the primary Zatom-2 weights (~/.nucleus/checkpoints)
nucleus install zatom2
```

### Generation

Generate 3D molecules, periodic materials, or proteins using `zatom2 design`:

```bash
# Generate a molecule
zatom2 design inputs=models/zatom2/docs/examples/demo.json out_dir=logs/molecule

# Generate a periodic material
zatom2 design inputs=models/zatom2/docs/examples/material.json out_dir=logs/material

# Generate a protein (downloading the SCOPe transfer checkpoint)
nucleus install zatom2-table3-scope-from-joint-properties-force-conditioned
zatom2 design \
    ckpt_path=zatom2-table3-scope-from-joint-properties-force-conditioned \
    inputs=models/zatom2/docs/examples/protein.json \
    out_dir=logs/protein
```

### Evaluation

To evaluate generated samples against benchmark distributions:

```bash
# Evaluate molecule generation on OMol25
zatom2 evaluate-generation \
    ckpt_path=zatom2 \
    out_dir=logs/evaluation/omol25 \
    evaluation=omol25_unconditional

# Evaluate periodic material generation on OMat24
zatom2 evaluate-generation \
    ckpt_path=zatom2 \
    out_dir=logs/evaluation/omat24 \
    evaluation=omat24_unconditional
```

To evaluate interatomic energy and force predictions (MLIP):

```bash
zatom2 evaluate-mlip ckpt_path=zatom2 out_dir=logs/evaluation/mlip
```

## Open-source resources

Zatom-2 builds upon the source code, data, and tools from the following projects:

- [AtomWorks](https://github.com/RosettaCommons/atomworks)
- [cuequivariance](https://github.com/NVIDIA/cuequivariance)
- [Foundry / RFdiffusion3](https://github.com/RosettaCommons/foundry)
- [lemat-genbench](https://github.com/LeMaterial/lemat-genbench)
- [lightning-hydra-template](https://github.com/ashleve/lightning-hydra-template)
- [Mol*](https://github.com/molstar/molstar)
- [posebusters](https://github.com/maabuu/posebusters)
- [ProteinMPNN](https://github.com/dauparas/ProteinMPNN)

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 (DOE) under Contract No. DE-AC02-05CH11231, using the AI4Sci@NERSC award NERSC DDRERCAP0036206 awarded to AM. NBE would also like to acknowledge that this work was supported in part by the U.S. Department of Energy's Genesis Mission and the Office of Science, Office of Advanced Scientific Computing Research's ModCon under Contract No. DE-AC02-05CH11231 at Lawrence Berkeley National Laboratory. Additionally, MC's PhD is funded by the EPSRC Centre of Doctoral Training in Automated Chemical Synthesis Enabled by Digital Molecular Technologies (SynTech CDT).

## Citation

If you use the code, weights, or data associated with this work, please cite:

```bibtex
@article{cretu2026zatom2,
    title={Zatom-2: Multitask Pretraining on Atomistic Data for Generative Modeling across Domains},
    author={Miruna Cretu* and Alex Abrudan$\dagger$ and Antonia Panescu$\dagger$ and Tynan Perez$\dagger$ and Rishabh Anand$\dagger$ and N. Benjamin Erichson and Michael W. Mahoney and Samuel Blau and Joseph Jacobson and Rafael G{\'o}mez-Bombarelli and Rex Ying and Tuomas Knowles and Pietro Li{\`o} and Alex Morehead*},
    journal={arXiv preprint arXiv:2610.11454},
    year={2026},
    eprint={2610.11454},
    archivePrefix={arXiv},
    primaryClass={cs.LG},
    url={https://arxiv.org/abs/2610.11454},
    note={* denotes equal contribution; $\dagger$ denotes equal core computational contribution}
}
```