Zatom-2

An atomistic foundation model for molecular, material, and protein generation. Install the shared Nucleus monorepo, then follow the inference CLI guide, paper reproduction guide, and data preparation guide.

The project page, preprint PDF, and LaTeX source are included. arXiv URL placeholder — submission pending.

Model Architecture

The Zatom-2 model is wired by Zatom2, with the simplified Hydra component using ZatomInitializer and Zatom2DiffusionModule. The current production configuration trains a flow-matching velocity model. At generation time, FlowMatchingEDMAdapter converts the RFdiffusion3-style EDM sampler state and noise schedule to flow state and time before each model call.

flowchart TD
    features["Input features f<br/>atom coordinates X_L and flow time t<br/>task, motif, material, and property metadata"] --> wrapper["Zatom2.forward"]

    wrapper --> init["ZatomInitializer<br/>static, geometry/time-agnostic conditioning"]
    init --> tok["Token path<br/>sequence/task/material + force statistics<br/>AllScAIP charge/spin embeddings -> S_I<br/>RPE + token bonds + ref coords -> Z_II"]
    init --> atom["Atom path<br/>atom/guidance features + S_I[token] -> C_L<br/>motif/ref coords + Z_II[token pairs] -> P_LL"]
    tok --> init_out["Initializer outputs<br/>S_I, Z_II, C_L, P_LL"]
    atom --> init_out

    wrapper --> train["Training branch<br/>single flow-matching model call"]
    wrapper --> infer["Generation branch<br/>iterative EDM-schedule rollout"]
    wrapper --> property["Property-inference branch<br/>single clean-endpoint model call"]
    init_out --> train
    init_out --> infer
    init_out --> property

    train --> seg{"self-embedding<br/>guidance?"}
    seg -->|optional| clean["Clean backbone pass<br/>restored sequence/element targets"]
    clean --> condition["Pool token or atom features<br/>into a system condition"]
    condition --> model_call
    seg -->|no| model_call

    infer --> cfg{"classifier-free<br/>guidance?"}
    cfg -->|yes| ref_init["Strip CFG features<br/>run reference initializer"]
    cfg -->|no| sampler
    ref_init --> sampler["ZatomConditionalDiffusionSampler<br/>RFdiffusion3-style EDM schedule"]
    sampler --> adapter["FlowMatchingEDMAdapter<br/>EDM sigma/state -> flow time/state"]
    adapter --> model_call["Zatom2DiffusionModule<br/>flow-matching velocity model"]
    property --> model_call

    model_call --> embed["Flow geometry/time conditioning<br/>scaled coordinates -> atom/token embeddings"]
    embed --> attn_graph["Attention graph construction<br/>atom local sparse indices<br/>optional token local indices"]
    attn_graph --> recycle{{"Recycle loop"}}
    recycle --> rec["Recycle feedback<br/>previous X_L only<br/>C-alpha distogram into Z_II"]
    rec --> backbone["TransformerBackbone<br/>atom encoder -> residual cross-attention Downcast<br/>token trunk -> residual cross-attention Upcast<br/>atom decoder"]
    backbone --> heads["Output heads<br/>flow velocity -> Euclidean endpoint X_L<br/>AdaLN-Zero sequence and atom-type logits<br/>fractional coords; system-downcast lattice<br/>optional energy, forces, and latent equivariance"]
    heads -->|next recycle| recycle
    heads --> output["Generated structure and optional properties"]

Relationship to RFdiffusion3 and Emyx

Zatom-2 is closest to RFdiffusion3 at the wrapper and sampler level, and closest to Emyx at the high-level atom-token-token-atom backbone layout. The main architectural difference is that Zatom-2 generalizes the representation and output heads for molecules and materials, while keeping a compact feature initializer and a shared TransformerBackbone.

References: RFdiffusion3 bioRxiv preprint, RFdiffusion3 Foundry source, and Emyx arXiv preprint.