All-Atom Diffusion Model

• All-atom diffusion models are probabilistic frameworks that use Markov processes to gradually transform noise into realistic atomic structures.
• They employ neural network-parameterized denoising steps and innovative architectures (e.g., MiAD, OXtal, ADiT) to capture atomic positions, types, and lattice features.
• These models enable scalable, accurate generation of crystals, molecules, and proteins, validated by metrics such as S.U.N. rates and packing similarity.

An all-atom diffusion model is a probabilistic, generative framework that models atomic systems—molecules, materials, or crystals—at full atomic resolution using Markov processes that gradually transform noise into realistic atomic structures. Such models capture the joint distribution over atomic coordinates, types, and auxiliary features (e.g., lattice, cell, or chemical graph) via forward noising and reverse denoising steps parameterized by neural networks. This paradigm has powered state-of-the-art advances in crystal generation, molecular modeling, protein structure restoration, and atomistic simulations on timescales inaccessible to direct molecular dynamics. Representative models include MiAD (Mirage Atom Diffusion) (Okhotin et al., 18 Nov 2025), OXtal (Jin et al., 7 Dec 2025), ADiT (Joshi et al., 5 Mar 2025), AGeDi (Rønne et al., 24 Jul 2025), DiAMoNDBack (Jones et al., 2023), and GPP-NEB-based atomic kinetics (Saxena et al., 24 Jun 2025).

1. Mathematical Foundations and Core Diffusion Processes

All-atom diffusion models are built upon discrete or continuous-time Markov chains that operate directly on atomic variables:

• Atomic positions: $\mathbf{X}_t \in \mathbb{R}^{N\times3}$ or fractional coordinates ($F_t \in [0,1)^{N\times3}$), typically corrupted by additive Gaussian noise in the forward process: $$q(x_t|x_{t-1}) = \mathcal{N}(x_t; \sqrt{1-\beta_t}x_{t-1}, \beta_t I)$$, or by SDEs in the continuous-time limit: $dx_t = f(x_t,t)dt + g(t)dW_t$.
• Atom types: $A_t \in \{0,1,\dots,T\}^N$ or similar, modeled by categorical noise kernels (discrete D3PM) or continuous-time Markov processes toward masking states.
• Lattice/cell parameters: $L \in \mathbb{R}^{3\times3}$ for periodic materials/crystals; often denoised with equivariant Gaussian DDPMs (Okhotin et al., 18 Nov 2025), or omitted via lattice-free protocols (Jin et al., 7 Dec 2025).

Reverse (generative) steps are parameterized by neural modules that predict means or scores for extracting structure from noise. Loss objectives combine denoising score matching for coordinates, cross-entropy or KL for types, and, in crystal models, masking-based losses to handle variable atom count (mirage atoms in MiAD).

In molecular settings, the joint noising of atomic species and geometry enables a two-phase inference regime: high-noise exploration selects atomic identities, low-noise relaxation refines geometries toward low-energy conformations (EDM) (Rothchild et al., 2023).

2. Model Variants and Architectural Strategies

Recent models introduce architectural innovations tailored to unique challenges in atomistic modeling:

• Mirage Infusion (MiAD): Permits the dynamic creation and annihilation of atoms during sampling via a learned 'mirage' state (type 0), enabling size-adaptive generation of crystals. This adoption of categorical diffusion for atom types substantially surpasses fixed-size models by increasing S.U.N. rates by up to $\times$2.5 (Okhotin et al., 18 Nov 2025).
• Voxel-based Grand Canonical DPMs: Use continuous occupancy fields on 3D grids rather than point clouds, sidestepping fixed-N constraints and capturing creation/annihilation of atoms naturally. A 3D U-Net denoises the voxel field, which is decoded via 3D CNN heads into atomic positions (Lei et al., 28 Aug 2024).
• Latent Diffusion and Transformers (ADiT): Autoencodes atoms, coordinates, and cell features into shared latent vectors, enabling a unified transformer-based flow-matching denoiser for both molecules and crystals (Joshi et al., 5 Mar 2025).
• Equivariant GNNs and Neural Backbones: Architectures enforce permutation, $O(3)$, lattice, and periodic equivariance via message-passing networks, canonical frame alignment, or data augmentation, as in CSPNet (MiAD) and Pairformer (OXtal) (Okhotin et al., 18 Nov 2025, Jin et al., 7 Dec 2025).
• Autoregressive Conditional DDPMs: In protein backmapping (DiAMoNDBack), all-atom residues are restored residue-by-residue, using local environment conditioning and canonical alignment, yielding realistic, diverse, and accurate side-chain configurations (Jones et al., 2023).

3. Sampling, Generation, and Inference Protocols

All-atom diffusion models employ iterative denoising schemes to sample atomic structures:

• Size-adaptive sampling: Mirage atom methods initialize all atoms in the maximum allowed slots as potentially non-existent; mirage infusion lets denoising steps switch atoms between existence/non-existence states based on learned categorical probability vectors (Okhotin et al., 18 Nov 2025).
• Voxel decoding: Grand canonical DPMs decode occupancy grids to atomic positions by classifying occupied voxels and interpolating sub-voxel offsets (Lei et al., 28 Aug 2024).
• Latent diffusion: Models such as ADiT diffuse latent codes, then decode the final (clean) latent into physical atomic attributes using autoencoders (Joshi et al., 5 Mar 2025).
• Autoregressive backmapping: DiAMoNDBack initializes each residue in its local reference frame, samples with a conditional DDPM, and aligns the output back into the growing global chain (Jones et al., 2023).
• Physical/time-resolved diffusion kinetics: The GPP-NEB model offers explicit integration schemes for atomic concentrations and hopping rates over real timescales, with inputs derived from statistical mechanics and transition-state theory (Saxena et al., 24 Jun 2025).

Sampling often incorporates classifier-free guidance, property masking, or explicit conditioning for compositional, symmetry, or packing constraints (Rønne et al., 24 Jul 2025, Lei et al., 28 Aug 2024).

4. Evaluation Metrics and Benchmarks

Quantitative assessment in all-atom diffusion is domain-specific:

• S.U.N. metrics: Stable, Unique, Novel rates denote the fraction of generated crystals that are thermodynamically stable ($E_{\text{hull}}<0$), unique by structure matcher, and not found in the training set. MiAD achieves 8.2% S.U.N. rate on MP-20, exceeding ADiT joint training (5.3%) and other baselines (Okhotin et al., 18 Nov 2025, Joshi et al., 5 Mar 2025).
• Packing similarity: For crystals, the COMPACK algorithm scores how well sampled periodic clusters align to experimental lattices under RMSD thresholds. OXtal attains 87% packing similarity and 96% conformer recovery on rigid sets (Jin et al., 7 Dec 2025).
• Bond/clash scores, diversity, dihedral recovery: For protein backmapping, DiAMoNDBack reports bond length accuracy (99.2%), clash avoidance, and generative diversity vs VAE and deterministic baselines (Jones et al., 2023).
• Physical kinetics: For atomistic long-term diffusion, Arrhenius rates, equilibrium solute/vacancy isotherms (Langmuir–McLean), and enrichment times are validated against experimental and analytic results (Saxena et al., 24 Jun 2025).
• Precision–recall in generated atomistic datasets: AGeDi quantifies model coverage and fidelity against data manifolds; reported AUCs approach or exceed 0.88 on materials datasets (Rønne et al., 24 Jul 2025).

5. Practical Implementations and Computational Efficiency

Implementation details vary by domain but share recurring themes:

• Scalable training: MiAD and ADiT scale to $10^7$–$10^8$ parameters, leveraging equivariant architectures or transformers with SE(3)-invariant data augmentation (Okhotin et al., 18 Nov 2025, Joshi et al., 5 Mar 2025, Jin et al., 7 Dec 2025).
• Sampling speed: ADiT enables generation of 10,000 samples in under 20 minutes on a single V100, outperforming E(n)-equivariant baselines by more than an order of magnitude (Joshi et al., 5 Mar 2025).
• Efficient lattice-free protocols: OXtal's S$^4$ shell cropping enables learning long-range packing without direct lattice parametrization, sidestepping the combinatorial complexity of cell prediction (Jin et al., 7 Dec 2025).
• Rapid time-resolved simulations: The GPP-NEB model integrates atomic diffusion master equations with on-the-fly transition barrier lookup, achieving real-time scale bridging (seconds to years in hours of compute) (Saxena et al., 24 Jun 2025).
• Generalizability and extensibility: Models such as AGeDi offer modular software stacks (Python/PyTorch), plugging in backbone GNNs, noise schedules, and supporting property guidance (Rønne et al., 24 Jul 2025).

6. Applications, Limitations, and Future Prospects

All-atom diffusion models impact a range of problems:

• De novo crystal generation: MiAD, OXtal, ADiT, and AGeDi enable direct sampling of crystals, including variable composition, symmetry, and atom counts, with empirically validated rates for thermodynamic and structural novelty (Okhotin et al., 18 Nov 2025, Jin et al., 7 Dec 2025, Joshi et al., 5 Mar 2025, Rønne et al., 24 Jul 2025).
• Molecular modeling: Diffusion models now compete with or outperform classical force fields in generating low-energy molecular geometries, with potential to accelerate quantum calculations (DFT) and sample Boltzmann distributions without explicit Hamiltonian labeling (Rothchild et al., 2023).
• Protein structure restoration: DiAMoNDBack demonstrates accurate, diverse recovery of all-atom protein configurations from coarse-grained traces, vital for MD refinement and structural ensemble generation (Jones et al., 2023).
• Atomistic kinetics: GPP-NEB models capture chemo-thermo-mechanical coupling in alloy diffusion, validated for mass transport, segregation, and time-resolved property prediction (Saxena et al., 24 Jun 2025).

Limitations include handling extreme stoichiometry, highly disordered or amorphous regions, explicit symmetry/lattice output, and coupling to continuum fields. Extensions are actively pursued: symmetry-conditioned diffusion, broader chemical representations, integration with physical ranking (free energy), solvent and kinetic pathway conditioning, and active learning for reactive or defective systems (Jin et al., 7 Dec 2025, Lei et al., 28 Aug 2024, Rønne et al., 24 Jul 2025).

A plausible implication is that all-atom diffusion frameworks will subsume classical search-and-rank or expensive ab initio methods, rendering atomistic structure and property modeling tractable at unprecedented scale and flexibility.