Zatom-1: Unified 3D Molecule & Material Model

• Zatom-1 is a multimodal flow foundation model that unifies discrete atom types and continuous geometric variables for 3D molecules and materials.
• It employs a Transformer architecture with a multimodal flow matching objective, integrating both generative and predictive learning efficiently.
• Joint pretraining on datasets like QM9 and MP20 enables state-of-the-art performance in property, energy, and force predictions across various chemical environments.

Zatom-1 is a multimodal flow foundation model designed to unify generative and predictive learning across 3D molecules and materials. Distinct from existing models restricted to either molecules or materials and specializing in a single task (generation or prediction), Zatom-1 employs a single Transformer-based architecture with a multimodal flow matching objective. This construction allows the model to jointly encapsulate discrete atom types and continuous variables corresponding to geometries, enabling scalable pretraining and efficient inference. Zatom-1 supports joint generative pretraining, which further serves as a universal initialization for a range of downstream multi-task predictive tasks, including properties, energies, and forces. Empirically, Zatom-1 consistently matches or surpasses specialized baselines while demonstrating significant advances in generative inference efficiency. Modeling materials during joint pretraining has been shown to improve molecular property prediction accuracy, indicating positive predictive transfer (Morehead et al., 24 Feb 2026).

1. Model Architecture

Zatom-1's architecture consists of a Transformer "trunk" that provides a unified representation for both discrete and continuous modalities:

• Modalities Embedded:
  • Discrete atom types.
  • Four continuous geometric modalities: 3D Cartesian coordinates, fractional coordinates, cell lengths, cell angles.

Given a structure with $N$ atoms, notation is as follows:

• $$\bm{A} = \{a_i\}_{i=1}^N \in \mathbb{Z}^{1 \times N}$$ (atom types)
• $$\bm{X} = \{x_i\}_{i=1}^N \in \mathbb{R}^{3\times N}$$ (Cartesian coordinates)
• $\bm{F} = \{f_i\}_{i=1}^N \in [0,1)^{3\times N}$ (fractional coordinates, where $\bm{F} = \bm{L}^{-1}\bm{X}$)
• $\bm{L}_{\mathrm{len}} \in \mathbb{R}^{3\times1}$ (cell lengths)
• $\bm{L}_{\mathrm{ang}} \in \mathbb{R}^{3\times1}$ (cell angles)

For molecules, lattice inputs are masked; for periodic materials, Cartesian inputs are masked.

Each atom receives a joint embedding:

$$h_i = \operatorname{Embed}(a_i) + \operatorname{Embed}(c) + \operatorname{Embed}(t) + \sum_{v \in \{x_i, f_i, L_{\mathrm{len}}, L_{\mathrm{ang}}\}} W_v v$$

where $c$ is a binary type indicator (molecule vs. material), $t$ is the flow time, and $$\bm{A} = \{a_i\}_{i=1}^N \in \mathbb{Z}^{1 \times N}$$0 are learned linear projections. The sequence $$\bm{A} = \{a_i\}_{i=1}^N \in \mathbb{Z}^{1 \times N}$$1 is passed through an $$\bm{A} = \{a_i\}_{i=1}^N \in \mathbb{Z}^{1 \times N}$$2-layer Transformer, yielding deep latent representations $$\bm{A} = \{a_i\}_{i=1}^N \in \mathbb{Z}^{1 \times N}$$3.

A residual cross-attention decoder block attends from each trunk token back to its input embedding, producing modality-specific hidden states ($$\bm{A} = \{a_i\}_{i=1}^N \in \mathbb{Z}^{1 \times N}$$4), each mapped to their prediction heads—for atom types, a SiLU-activated linear layer; for geometries, a bias-free linear layer following layer normalization.

For downstream prediction (properties, energies, forces), trunk representations are taken from a selected layer $$\bm{A} = \{a_i\}_{i=1}^N \in \mathbb{Z}^{1 \times N}$$5, frozen, and processed by an $$\bm{A} = \{a_i\}_{i=1}^N \in \mathbb{Z}^{1 \times N}$$6-layer auxiliary Transformer for each target, with outputs read out via pooling or per-atom projections.

2. Multimodal Flow Matching Objective

Zatom-1 is pretrained using a multimodal conditional flow matching (CFM) loss—integrating both continuous (for geometry) and discrete (atom type) modalities.

Continuous CFM (Geometries):

• Noisy input at time $$\bm{A} = \{a_i\}_{i=1}^N \in \mathbb{Z}^{1 \times N}$$7: $$\bm{A} = \{a_i\}_{i=1}^N \in \mathbb{Z}^{1 \times N}$$8
• Network predicts vector field $$\bm{A} = \{a_i\}_{i=1}^N \in \mathbb{Z}^{1 \times N}$$9.
• Metric loss per modality:

$$\bm{X} = \{x_i\}_{i=1}^N \in \mathbb{R}^{3\times N}$$0

where $$\bm{X} = \{x_i\}_{i=1}^N \in \mathbb{R}^{3\times N}$$1 is the denoised prediction.

Discrete CFM (Atom Types):

• Noised type distribution: $$\bm{X} = \{x_i\}_{i=1}^N \in \mathbb{R}^{3\times N}$$2
• Loss:

$$\bm{X} = \{x_i\}_{i=1}^N \in \mathbb{R}^{3\times N}$$3

Combined loss:

$$\bm{X} = \{x_i\}_{i=1}^N \in \mathbb{R}^{3\times N}$$4

with $$\bm{X} = \{x_i\}_{i=1}^N \in \mathbb{R}^{3\times N}$$5; the loss is reweighted by $$\bm{X} = \{x_i\}_{i=1}^N \in \mathbb{R}^{3\times N}$$6 with $$\bm{X} = \{x_i\}_{i=1}^N \in \mathbb{R}^{3\times N}$$7 to stabilize near $$\bm{X} = \{x_i\}_{i=1}^N \in \mathbb{R}^{3\times N}$$8 (Morehead et al., 24 Feb 2026).

3. Generative Pretraining and Sampling Procedures

Zatom-1 is jointly pretrained on both QM9 (small molecules, $$\bm{X} = \{x_i\}_{i=1}^N \in \mathbb{R}^{3\times N}$$9130K) and MP20 (periodic crystals, $\bm{F} = \{f_i\}_{i=1}^N \in [0,1)^{3\times N}$045K) datasets using the Transformer trunk and identical flow matching objective. All coordinates are modeled directly in $\bm{F} = \{f_i\}_{i=1}^N \in [0,1)^{3\times N}$1 without autoencoders or explicit priors.

Sampling Algorithm:

1. Initialize discrete atom types $\bm{F} = \{f_i\}_{i=1}^N \in [0,1)^{3\times N}$2 and continuous vectors $\bm{F} = \{f_i\}_{i=1}^N \in [0,1)^{3\times N}$3.
2. Iterate time steps $\bm{F} = \{f_i\}_{i=1}^N \in [0,1)^{3\times N}$4 using Euler integration:
   • Discrete $\bm{F} = \{f_i\}_{i=1}^N \in [0,1)^{3\times N}$5 updated via discrete flow [Campbell et al.]
   • Continuous update:

   $\bm{F} = \{f_i\}_{i=1}^N \in [0,1)^{3\times N}$6

   with $\bm{F} = \{f_i\}_{i=1}^N \in [0,1)^{3\times N}$7, $\bm{F} = \{f_i\}_{i=1}^N \in [0,1)^{3\times N}$8, $\bm{F} = \{f_i\}_{i=1}^N \in [0,1)^{3\times N}$9.

3. Return decoded $\bm{F} = \bm{L}^{-1}\bm{X}$0.

Zatom-1 achieves rapid generation: 10,000 samples in under 4 minutes on a single NVIDIA A100 (using 50–100 steps), over an order of magnitude faster than a 500M-parameter ADiT latent diffusion baseline and 3$\bm{F} = \bm{L}^{-1}\bm{X}$1 faster than prior flow-matching crystal models.

4. Downstream Predictive Capabilities

Following pretraining, the model's Transformer trunk is kept frozen for downstream property, energy, and force prediction. The downstream protocol is:

• Extract embeddings $\bm{F} = \bm{L}^{-1}\bm{X}$2 from trunk layer $\bm{F} = \bm{L}^{-1}\bm{X}$3.

• Input to an $\bm{F} = \bm{L}^{-1}\bm{X}$4-layer auxiliary Transformer head (separate for each task), using cross-attention to the frozen trunk.

• Readouts are either mean- or per-atom-projected for final prediction.

Loss functions used per batch of size $\bm{F} = \bm{L}^{-1}\bm{X}$5 (with $\bm{F} = \bm{L}^{-1}\bm{X}$6 atoms) include:

• Properties: $\bm{F} = \bm{L}^{-1}\bm{X}$7

• Energy: $\bm{F} = \bm{L}^{-1}\bm{X}$8

• Forces: $\bm{F} = \bm{L}^{-1}\bm{X}$9

Set $\bm{L}_{\mathrm{len}} \in \mathbb{R}^{3\times1}$0.

5. Empirical Evaluations

Zatom-1 is evaluated for both generative and predictive tasks across established datasets. Key results:

• Crystal Generation (MP20/LeMat-GenBench, 2,500 samples; Zatom-1-XL):

  • Validity 94.9%, Uniqueness 90.1%, Novelty 4.8%
  • Formation energy $\bm{L}_{\mathrm{len}} \in \mathbb{R}^{3\times1}$1 eV/atom
  • Energy above hull $\bm{L}_{\mathrm{len}} \in \mathbb{R}^{3\times1}$2 eV/atom
  • RMSD of relaxation $\bm{L}_{\mathrm{len}} \in \mathbb{R}^{3\times1}$3 Å
  • Stable fraction 3.3%, SUN (Stable∧Unique∧Novel) ≈ 0.36%
  • Matches or exceeds specialized crystal models (DiffCSP, FlowMM) without any relaxation.
• Molecule Generation (QM9, 10,000 samples; Zatom-1-L):
  • Validity 95.26%, Uniqueness 96.84%
  • PoseBusters pass-rate ≈ 99.9% on all 7 sanity checks
  • Surpasses ADiT (180M) and all prior equivariant diffusion or latent-space methods.
• Large-Molecule Generation (GEOM-Drugs, 10,000 samples):
  • Validity 93.6%, Uniqueness 99.93%, PoseBusters valid 94.1%
  • Outperforms TABASCO, SemlaFlow by $\bm{L}_{\mathrm{len}} \in \mathbb{R}^{3\times1}$4.
• Predictive Finetuning (QM9 Multi-task, 80M + 20M parameter configuration):
  • $\bm{L}_{\mathrm{len}} \in \mathbb{R}^{3\times1}$5
  • State-of-the-art multi-task performance, outperforming EquiformerV2, DimeNet++ with approximately 100 GPU-hours.
• Transfer and Generalization:
  • Joint QM9+MP20 pretraining provides positive transfer—best property MAEs obtained with joint pretraining.
  • Zero-shot transfer: embeddings generalize from molecules to materials (Matbench results for $\bm{L}_{\mathrm{len}} \in \mathbb{R}^{3\times1}$6).
• Energy & Force Prediction (MLIPs):
  • On MPtrj, OMol25-4M: force MAE ≈ 31 meV/Å, energy MAE 2631 meV (slightly better than Orb-v1). This suggests Zatom-1 is a promising universal interatomic potential.

6. Design Parameters, Hyperparameters, and Scaling Behavior

The Zatom-1 implementation includes distinct model sizes:

• Zatom-1: 80M parameters (hidden 512, 16 blocks, 8 heads)
• Zatom-1-L: 160M parameters (hidden 768, 16 blocks, 8 heads)
• Zatom-1-XL: 300M parameters (hidden 1024, 16 blocks, 8 heads)

Datasets include: QM9 (130K), MP20 (45K), GEOM-Drugs (430K), QMOF (14K).

Pretraining (80M configuration): batch 256 ($\bm{L}_{\mathrm{len}} \in \mathbb{R}^{3\times1}$78 rotations, effective 32,768), AdamW, learning rate $\bm{L}_{\mathrm{len}} \in \mathbb{R}^{3\times1}$8, $\bm{L}_{\mathrm{len}} \in \mathbb{R}^{3\times1}$9, $\bm{L}_{\mathrm{ang}} \in \mathbb{R}^{3\times1}$0 sampling steps.

Scaling laws indicate both training loss and generative validity improve predictably with increasing model size (80M $\bm{L}_{\mathrm{ang}} \in \mathbb{R}^{3\times1}$1 160M $\bm{L}_{\mathrm{ang}} \in \mathbb{R}^{3\times1}$2 300M), with no observed saturation, suggesting further gains are possible at larger scale.

7. Significance and Outlook

Zatom-1 is the first model to simultaneously unify Transformer-based learning across molecules and crystals using multimodal flow matching, provide orders-of-magnitude speedup for 3D structure generation relative to prior diffusion/flow models, and deliver state-of-the-art performance on a suite of predictive tasks (properties, energies, forces) with minimal adaptation. Positive predictive transfer between domains and efficient scaling behavior mark its initial deployment as a general-purpose foundation model for 3D chemical environments (Morehead et al., 24 Feb 2026).