Zatom-2 improves atom structure generation for chemistry and biology research

Zatom-2: Multitask Pretraining on Atomistic Data for Generative Modeling across Domains

Machine LearningArtificial Intelligence

Summary

Designing new molecules and materials by predicting atom arrangements is key in chemistry and biology. Existing tools usually focus only on specific areas like molecules or materials, limiting their usefulness. The authors created Zatom-2, a model trained on millions of atomic structures from both molecules and materials to work well across these fields. It can generate realistic atom configurations and better design protein structures, especially in cases where data is scarce. This helps speed up designing new chemicals, materials, and proteins by bridging knowledge from different scientific areas.

What this means in practice

  • For computational chemists: Generate diverse and accurate molecular structures across chemistry and material science domains using a unified pretrained model.
  • For protein engineers: Enhance protein backbone design in low-data settings by applying Zatom-2’s joint generative-predictive pretraining and transfer learning.

Authors

Miruna Cretu, Alex Abrudan, Antonia Panescu, Tynan Perez, Rishabh Anand, N. Benjamin Erichson, Michael W. Mahoney, Samuel Blau, Joseph Jacobson, Rafael Gómez-Bombarelli, Rex Ying, Tuomas Knowles, Pietro Liò, Alex Morehead

Abstract

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.