Zatom-2 / README.md

Update README.md

48faf77 verified about 18 hours ago

5.62 kB

license: mit
language:
  - en
tags:
  - chemistry
  - biology
  - generative-model
  - predictive-model
  - representation-learning
  - transformer
  - molecule
  - material
  - protein
  - property
  - energy
  - forces
  - mlip

Zatom-2

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, 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 (recommended) or pip:

# 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:

# 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:

# 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:

# 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):

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:

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:

@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}
}