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
[](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}
}
``` |