cs.LGOct 8, 2026
SaveZatom-2: Multitask Pretraining on Atomistic Data for Generative Modeling across Domains
Organizations: University of Cambridge · Yale University · MIT · Lawrence Berkeley National Lab · International Computer Science Institute · UC Berkeley
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.
Figures & tables
Figure 1: Zatom-2 multitask pretraining across domains. Zatom-2 is jointly trained on OMol25 (4M) and OMat24 (1M) as a generative model. Its input initializer receives task-dependent conditioning, summarized by task strips: checkmarks indicate enabled inputs and dashes indicate disabled inputs. MLIP task supervision (force and energy prediction) uses clean coordinates at and known atom types, without bond or system force conditioning. The initializer provides atom and token conditioning through the dotted connections. Noisy coordinates and flow time enter the U-shaped atom encoder—token trunk—atom decoder. Atom features and token features feed task-specific readout heads. The recycling path converts the clean-coordinate estimate obtained from into a token-level distogram, updating token-pair conditioning for a second forward pass. In this work, we experiment with trunk layers.
| Datasets | Tasks | ||||||
| Model | Size | OMol25 | OMat24 | Generation | Structure prediction | Force & energy prediction | Force conditioning |
| Zatom-1 (released) | Large | – | – | – | – | ||
| Zatom-2 (I) | Small | – | – | – | – | ||
| Zatom-2 (II) | Small | – | – | – | |||
| Zatom-2 (III) | Small | – | – | ||||
| Zatom-2 (IV) | Base | – | |||||
Table 1: Pretraining configurations and generation quality. Panel (a) shows multiple training configurations for OMol25 and OMat24 mixed training. Panel (b) reports their performance, where we also compare to 10,000 dataset samples and Zatom-1. All models are trained for 80 epochs. Forces are calculated using UMA ( Wood et al., 2025 ) and metrics are reported in eV Å -1 . AUROC uses a conditioning mean force norm cutoff of 1 eV Å -1 . Structure prediction results are reported in Appendix Table 9 .
Figure 2: Force-conditioning consistency. Panel (a) compares the conditioned mean force norm with the mean force norm evaluated by UMA for unconditional OMol25 molecule generation. Panels (b) and (c) give the corresponding comparison for OMol25 structure prediction on the GEOM ORCA6 and ANI2x subsets, respectively. Red dotted lines mark the AUROC cutoff (1 eV Å -1 in (a)–(b), and 2.5 eV Å -1 in (c)).
| Model | Valid | Unique | PB-valid | Valid & unique | Valid & PB-valid |
|---|---|---|---|---|---|
| EQGAT-diff | 94.6 | 100.0 | 59.7 | 94.60 | 56.48 |
| SemlaFlow | 93.9 | 100.0 | 87.5 | 93.90 | 82.16 |
| ADiT | 95.3 | 100.0 | 85.3 | 95.30 | 81.29 |
| TABASCO | 97.6 | 99.03 | 91.6 | 96.65 | 89.40 |
| Zatom-1 | 93.6 | 99.93 | 94.1 | 93.53 | 88.08 |
| GEOM-Drugs (reference) | – | – | 94.0 | – | – |
Table 2: Generation benchmarks on GEOM-Drugs and MP20. Panel (a) reports unconditional GEOM-Drugs molecule generation metrics, and panel (b) reports MetaSUN yield for unrelaxed MP20 material generation.
Figure 3: Model and data scaling across training epochs. We evaluate Zatom-2 models trained either on 500k examples per dataset or on the full OMol25 and OMat24 datasets. We report (a) OMol25 molecular generation quality measured by AFD, (b) OMat24 structure prediction measured by top-1 match rate, and force prediction MAE on (c) OMol25 and (d) OMat24. Training on the full dataset with more model parameters consistently improves validation performance across tasks.
| Model / Zatom-2 initialization | Designability (%) | Per-sequence success (%) | Diversity (#clusters) | Novelty (%) |
|---|---|---|---|---|
| (a) In-distribution: 50–128 residues | ||||
| RFdiffusion3 | ||||
| None | ||||
| Zatom-2 (IV): Generation + structure (w/o FC) | ||||
| Zatom-2 (IV): Generation + structure | ||||
| Zatom-2 (V): Generation + structure + F/E pred. | ||||
Table 3: SCOPe-2k protein generation and transfer. (a): We evaluate 1,027 protein backbones for each configuration, 13 corresponding to each length between 50 and 128 (with means standard deviations over 3 random seeds). (b): The same protocol is repeated for 129–256 residue lengths. FC denotes force-conditioning, and F/E pred. refers to a model trained with force and energy prediction as auxiliary tasks (Zatom (V)). The Zatom (V)-finetuned model substantially outperforms all configurations on the length extrapolation task. Metrics are explained in Appendix D.9 .
Figure 4: Protein latent space organization and localized decoding responses. Left two: t-SNE views of final token features, colored by amino acid. Right two: increases in sequence and structure reconstruction losses after final-token perturbation; solid/dashed lines denote the target/other amino acids. Bands show 95% intervals from 2,000 protein bootstrap resamples. Pretrained Zatom-2 is notably more sensitive to local amino acid-specific perturbations than without pretraining.
Appendix figures & tables25 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: Mean force norm distributions across OMol25-4M subsets.
Figure 6: Mean force norm distributions across OMat24-1M subsets. Data subsets are grouped by generation protocol (rattled and then re-relaxed, rattled at various temperatures, or ab initio molecular dynamics trajectories; see Barroso-Luque et al. (2024) for more details). The legend shows subset medians in eV Å -1 .
Figure 7: Atom count distributions of OMol25-4M & OMat24-1M. Hydrogens are included.
| Dataset | Source data terms |
|---|---|
| OMol25 | Creative Commons Attribution 4.0 (CC BY 4.0) |
| OMat24 | CC BY 4.0 |
| QM9 | CC0 1.0 (original Figshare data deposit) |
| MP20 | MIT for the processed ADiT release used by Zatom-1; underlying Materials Project data are CC BY 4.0 |
| GEOM-Drugs | CC0 1.0 (Harvard Dataverse deposit) |
| SCOPe / ASTRAL | Provider states that all data are freely available to all users; no standard named license is specified |
Table 4: Dataset licenses and access terms. Dataset names link to the corresponding release or the provider’s usage statement.
| RFdiffusion3 | Emyx | Zatom-2 | |
|---|---|---|---|
| 1 | Build local atom neighborhoods (128 keys); use global token attention. | Build sparse atom/token graphs prioritizing sequence, bonds, ligand, and motif edges, then spatial neighbors. | Build local atom neighborhoods (128 keys, token-index/spatial neighbors); use global token attention. |
| 2 | Embed atom/token metadata; mix dense pair features with two triangle-free Pairformer blocks; pool atom pairs into token pairs. | Bottleneck-embed atom14 atom/token and global features; project relative positions and bonds to edge biases. | Bottleneck-embed atom1/atom14 metadata, domain, and charge/spin; residually-embed mean force norm; map bonds, relative positions, and any available fixed/reference geometry to head biases. |
| 3 | Add linear embeddings of ; condition on Fourier features of . | Add Fourier coordinate embeddings; average geometry-aware conditioning with sinusoidal time features. | Add linear atom and mean-pooled linear token embeddings of ; average each conditioning stream with Fourier features of . |
| 4 | Cache and before recycling. | Begin recycling before the atom encoder. | Begin recycling before the atom encoder. |
| 5 | Reset ; concatenate , noisy-input and recycled distograms; apply transitions and two triangle-free Pairformer blocks to obtain . | Reset to initial states plus projected, normalized ; set on sparse edges. | Reset ; set for all token pairs. No node-state feedback. |
| 6 | Reuse cached encoder/downcast outputs. | ; . | ; . |
Table 5: Aligned forward pass operations. Enc, Trunk, and Dec are attention stacks; Down/Up are gated cross-attention between input resolutions. By definition, no recycling feedback is used in the first pass.
| Hyperparameter | Setting |
|---|---|
| Approximate model size (S/base/L) | 70M / 140M / 270M parameters |
| Token-trunk blocks (S/base/L) | 9 / 18 / 36 |
| Atom encoder / decoder blocks | 3 / 3 |
| Atom / token / token-pair widths | 128 / 768 / 256 |
| Atom / token attention heads | 4 / 12 |
| Attention | Local atom attention (128 keys); global token attention |
Table 6: Key Zatom-2 hyperparameters. S/base/L vary the token trunk depth while retaining the same feature widths. Token budgets are per GPU before diffusion batch/view augmentation; a smaller token budget is paired with twice as much gradient accumulation.
| Model training configuration | Budget / endpoint | GPUs | GPU hours |
| Small, OMol25 G, no FC | 80 epochs | 16 | 1,830 |
| Small, OMol25 G, FC | 80 epochs | 16 | 1,820 |
| Small, OMol25 G+S, FC | 80 epochs | 64 | 2,280 |
| Base, OMol25+OMat24 G+S, FC | 80 epochs | 64 | 2,740 |
| Base, OMol25+OMat24 G+S+P, FC | 80 epochs | 64 | 2,850 |
| Large, OMol25+OMat24 G+S+P, FC | 80 epochs | 64 | 4,090 |
Table 7: Approximate 80GB A100 GPU hours by training configuration. G/S/P denote generation, structure prediction, and MLIP prediction; FC denotes force conditioning. Epoch budgets count completed epochs; MLIP finetuning runs end at raw epoch 119. All finetuning costs exclude pretraining the corresponding initial model weights.
| OMol25 4m / ct-scd-omol25 | OMat24 1m / ct-scd-amp20 | |
|---|---|---|
| 100 | ||
| 500 | ||
| 1,000 | ||
| 5,000 |
Table 8: AFD baselines for encoder dataset pairs used in this work (finite- floor): Each entry is the mean standard deviation of the AFD score between two non-overlapping ground-truth reference sets of size , over 5 replicates (no structure is shared within or across replicates at a given ).
| OMol25 | |||||
|---|---|---|---|---|---|
| Model | Mean RMSD | Median RMSD | Mean | Target MAE | AUROC |
| Zatom-2 (III) | 1.739 | 1.714 | 2.271 | 0.725 | 0.798 |
| Zatom-2 (IV) | 1.708 | 1.675 | 2.138 | 0.787 | 0.787 |
| Zatom-2 (V) | 1.762 | 1.729 | 1.953 | 0.637 | 0.801 |
| Zatom-2 (VI) | 1.723 | 1.743 | 1.211 | 1.214 | 0.818 |
Table 9: Molecule/material top-1 structure prediction after pretraining. All models use the 80-epoch pretrained checkpoints described in Table 1 . OMol25 RMSDs are reported in Å and forces are calculated with UMA. Compared to OMol25, OMat24 structure prediction is more sensitive to outliers, so we report median metrics for OMat24. OMol25 AUROC uses a target mean force norm cutoff of 1 eV Å -1 .
| OMol25 | |||||||||
|---|---|---|---|---|---|---|---|---|---|
| Generation | Structure prediction | ||||||||
| Model | AFD | Mean | Target MAE | AUROC | Mean RMSD | Median RMSD | Mean | Target MAE | AUROC |
| Zatom-2 (IV) + FT | 0.017 | 1.205 | 0.540 | 0.807 | 1.771 | 1.747 | 1.570 | 0.632 | 0.890 |
| Zatom-2 (V) + FT | 0.016 | 1.137 | 0.530 | 0.809 | 1.729 | 1.706 | 1.458 | 0.690 | 0.828 |
| Zatom-2 (V) + FT (extended) | 0.015 | 1.040 | 0.526 | 0.840 | 1.720 | 1.713 | 1.178 | 0.871 | 0.882 |
Table 10: MLIP accuracy and retention of generative capabilities. “+ FT” denotes 40 epochs of finetuning with 80% of training steps allocated to MLIP prediction and 10% for generation and structure prediction each; here, “extended” means adding 120 additional MLIP-heavy finetuning epochs. Model names follow Table 1 . Panel (a) reports generation and top-1 structure prediction results after finetuning, with forces calculated using UMA; corresponding pretrained results appear in Tables 1 and 9 . For OMat24, we observe that generation and structure prediction are sensitive to outliers and we report median values. Panels (b)–(c) report validation MAEs (Appendix D.8 ); external baselines are published values with partially differing training and evaluation protocols.
Figure 8: SCOPe-2k convergence From left to right, the panels show best-of-eight designability (CA scRMSD Å), the mean maximum exact, query-normalized TM-align score against the canonical SCOPe-2k training set among designable backbones (lower is more novel), and distinct joint sequence–structure clusters among designable backbones. Joint clusters require TM together with at least 10% sequence identity and 70% coverage. Each point uses 96 generated backbones and eight refolds per backbone.
| Model | Validity |
|---|---|
| Equivariant Diffusion | 91.90 |
| Symphony | 83.50 |
| GeoLDM | 93.80 |
| ADiT (QM9-only) | 92.19 |
| ADiT (joint QM9+MP20) | 94.45 |
| Zatom-1 (QM9-only, 80M) | 92.88 |
Table 11: QM9 generation results. Zatom-2 is trained jointly on QM9 and MP20 (Appendix B ). Published baselines are from Table 2 of Morehead et al. (2026) (arXiv v4).
| Model | Overall valid | Unique | Novel | MetaSUN |
|---|---|---|---|---|
| ADiT (joint QM9+MP20) | 90.6 | 87.8 | 26.0 | 1.0 |
| Crystal-GFN | 51.7 | 51.7 | 51.7 | 0.0 |
| Crystalformer | 69.9 | 69.4 | 31.8 | 3.1 |
| LLaMat2-CIF | 84.4 | 81.4 | 30.0 | 2.1 |
| LLaMat3-CIF | 15.4 | 15.2 | 10.5 | 0.2 |
| SymmCD | 73.4 | 73.0 | 47.0 | 2.4 |
Table 12: MP20 generation results (unrelaxed).
| Model | Atoms connected | Bond angles | Bond lengths | Aromatic ring flat | Double bond flat | Internal energy | No steric clash |
|---|---|---|---|---|---|---|---|
| EQGAT-diff | 84.4 | 86.9 | 87.0 | 87.0 | 87.0 | 86.8 | 82.9 |
| SemlaFlow | 92.3 | 94.8 | 94.6 | 94.9 | 94.2 | 94.8 | 92.0 |
| ADiT | 93.0 | 92.3 | 92.5 | 95.4 | 95.3 | 91.3 | 91.8 |
| TABASCO | 99.9 | 99.2 | 99.4 | 100.0 | 99.8 | 99.4 | 94.3 |
| Zatom-1 | 99.6 | 99.6 | 99.7 | 100.0 | 100.0 | 99.3 | 96.4 |
| GEOM-Drugs (reference) | – | – | – | – | – | – | – |
Table 13: Individual PoseBusters checks for GEOM-Drugs generation. Published baseline values are reported from Table 3 of Morehead et al. (2026) (arXiv v4). Results are percentages among RDKit-valid molecules.
| Model | COV-R mean | MAT-R mean | COV-P mean | MAT-P mean |
|---|---|---|---|---|
| GO-Flow | 94.82 | 0.7971 | 70.13 | 1.1068 |
| Zatom-2 | 86.09 | 0.8459 | 83.07 | 0.9083 |
Table 14: GEOM-Drugs conformer generation evaluation. Values are means across test molecules; COV is reported as a percentage and MAT as RMSD in Å. GO-Flow values are from Table 1 of Liu et al. (2026) .
| Model | MR (%) | RMSE |
|---|---|---|
| Crystalite | 66.09 | 0.0337 |
| Zatom-2 | 64.00 | 0.1156 |
Table 15: MP20 structure prediction evaluation. We report the Crystalite result from Hadži Veljković et al. (2026) .
Figure 9: Generated QM9 molecules. MolStar ball-and-stick views of nine Zatom-2 molecules. Carbon is gray, hydrogen white, nitrogen blue, and oxygen red. Bonds are perceived by MolStar from the saved structures (note that coordinates are unchanged).
Figure 10: Generated OMol25 molecules. MolStar ball-and-stick views spanning biomolecules, electrolytes, metal complexes, reactivity, and community datasets (labels identify the source subset). Notably, Zatom-2’s source category and chemical composition for each sample are not fixed during generation but rather are (re)produced. Note that Transition-1x and RGD1 supply metadata for known reaction path examples, but Zatom-2’s generated samples for these subsets are not verified transition states. Atoms are colored by element type, and bonds are perceived by MolStar.
Figure 11: Generated MP20 materials. MolStar views of unrelaxed Zatom-2 materials, with atoms colored by element. Each panel displays a periodic repetition of the saved cell; a gray outline marks one unit cell. Labels report the composition of a single saved cell (not a reduced formula) and its sample ID.
Figure 12: Generated OMat24 materials. MolStar views of nine force-conditioned Zatom-2 samples, with atoms colored by element. Each panel displays a periodic repetition of the generated cell, with one cell outlined in gray. Labels give the composition of the saved cell and an abbreviated sample ID. No relaxation, space group refinement, or symmetrization is applied.
Figure 13: Generated SCOPe-2k protein backbones. MolStar cartoons colored from the N-terminus (blue) to the C-terminus (red). Rows show helix-rich, mixed, and sheet-rich examples, respectively. Labels display backbone length, best-of-eight C scRMSD in Å, and abbreviated sample IDs. The coordinates rendered for each sample are the generated backbones, not their ESMFold2 refolded structures.
Figure 14: Base Zatom-2 system embeddings across OMol25 and OMat24. Left: Zatom-2 (IV), without energy/force prediction; right: Zatom-2 (V), with it. Each point represents the same validation example in both columns, mean-pooled over all atoms. Rows show results for the joint datasets (a,b), OMol25 subsets (c,d), and OMat24 subsets (e,f), with matching colors across checkpoints. Each panel has an independent PCA fit; axis percentages report variance. Note that axes are not aligned between models and that “(sub)” denotes subsampled OMat24 data.
Figure 15: Base Zatom-2 atom embeddings across OMol25 and OMat24. Atom-level counterpart of Figure 14 (a,b), separately for O, P, N, and Br, selected by shared structure coverage. Left: Zatom-2 (IV), without MLIP supervision; right: Zatom-2 (V), with it. Each point is one atom, colored by domain, with identical atoms in both columns. Titles give sampled atom and structure counts. Each panel fits centered PCA independently; axis percentages report variance, and axes are not aligned between models.
Figure 16: Base Zatom-2 atom embeddings within OMol25. Atom-level counterpart of Figure 14 (c,d) for H, C, N, and O, the four elements present in the most sampled molecules. Left: without MLIP supervision; right: with it. Colors indicate the same OMol25 subsets as in Figure 14 . Each point is an individual atom, with an independent PCA fit per panel.
Figure 17: Base Zatom-2 atom embeddings within OMat24. Atom-level counterpart of Figure 14 (e,f) for Hg, La, Ag, and K, selected by structure coverage in the sampled validation materials. Left: without MLIP supervision; right: with it. Each point is an atom, colored by the corresponding OMat24 subset; “(sub)” denotes a subsampled subset. Identical atoms appear in both columns, with independent PCA fits. MLIP supervision concentrates more variance in PC1.