Source-linked AI summary
All-atom Diffusion Transformers: Unified generative modelling of molecules and materials
Chaitanya K. Joshi, Xiang Fu, Yi-Lun Liao, Vahe Gharakhanyan, Benjamin Kurt Miller, Anuroop Sriram, Zachary W. Ulissi
TL;DR
Existing diffusion models for 3D atomic systems are usually specialized to molecules or materials despite shared underlying physics. ADiT unifies both domains through shared all-atom latent representations and Transformer latent diffusion, achieving state-of-the-art results with substantial speedups and predictable scaling. The approach is promising for broadly generalizable generative chemistry, but its validation remains limited by small training datasets and incomplete testing on very large systems.
Problem
Diffusion models lack a unified formulation for generating different atomic systems, including periodic materials and non-periodic molecules, despite their shared physical principles.
Method
ADiT jointly reconstructs molecules and materials in a shared all-atom VAE latent space, then uses a Diffusion Transformer with classifier-free guidance to generate decodable samples.
Results
ADiT achieves state-of-the-art molecular and crystal generation, including a 5-6% S.U.N. rate for crystals and 10,000 samples in under 20 minutes on one V100 GPU.
Takeaways & Limitations
A single standard-Transformer-based model can generate both molecules and materials while supporting transfer learning, computational efficiency, and scaling to half a billion parameters.
Takeaways & Limitations
Training uses relatively small datasets, and the approach has not been fully validated on larger systems such as metal-organic frameworks or biomolecules containing thousands of atoms.
Abstract
from arXiv · showhide
Diffusion models are the standard toolkit for generative modelling of 3D atomic systems. However, for different types of atomic systems -- such as molecules and materials -- the generative processes are usually highly specific to the target system despite the underlying physics being the same. We introduce the All-atom Diffusion Transformer (ADiT), a unified latent diffusion framework for jointly generating both periodic materials and non-periodic molecular systems using the same model: (1) An autoencoder maps a unified, all-atom representations of molecules and materials to a shared latent embedding space; and (2) A diffusion model is trained to generate new latent embeddings that the autoencoder can decode to sample new molecules or materials. Experiments on MP20, QM9 and GEOM-DRUGS datasets demonstrate that jointly trained ADiT generates realistic and valid molecules as well as materials, obtaining state-of-the-art results on par with molecule and crystal-specific models. ADiT uses standard Transformers with minimal inductive biases for both the autoencoder and diffusion model, resulting in significant speedups during training and inference compared to equivariant diffusion models. Scaling ADiT up to half a billion parameters predictably improves performance, representing a step towards broadly generalizable foundation models for generative chemistry. Open source code: https://github.com/facebookresearch/all-atom-diffusion-transformer
1 Introduction
ADiT addresses the lack of a unified diffusion formulation for periodic materials and non-periodic molecules by combining shared all-atom latent representations with Transformer-based latent diffusion. Joint training achieves strong molecular and crystal generation while improving efficiency and scalability.
- Motivation: ADiT targets unified generation of periodic materials and non-periodic molecules despite existing diffusion models being system-specific.The paper frames this as a question motivated by shared physical principles and differing categorical-continuous generative formulations.
- Method: A VAE reconstructs all-atom molecules and crystals into a shared latent space, treating both as sets of atoms with categorical and continuous attributes.This representation supports both non-periodic and periodic systems within one model.
- Method: A Diffusion Transformer samples new shared-space latents with classifier-free guidance, which the VAE decodes into valid molecules or crystals.Latent diffusion shifts categorical-continuous complexity into the autoencoder and simplifies the generative process.
- Results: 5-6% S.U.N. rate: ADiT generates stable, unique, and novel crystals, improving on previous methods’ 4-5% rates.The result comes from DFT calculations reported for jointly trained ADiTs.
- Results: 10,000 samples in under 20 minutes on one V100 GPU, versus baselines taking up to 2.5 hours, while scaling reaches half a billion parameters.Performance improves predictably with model size, and joint training outperforms single-domain variants on validity.
- Conclusion: ADiT provides unified, state-of-the-art generation across molecules and crystals and represents a step toward broadly generalizable generative-chemistry models.The framework also demonstrates transfer learning between periodic and non-periodic atomic systems.
2 All-atom Diffusion Transformers
ADiT unifies periodic crystals and non-periodic molecules as all-atom 3D systems, encoding both into a shared latent space and generating new samples with latent diffusion.
- Unified representation: ADiT represents molecules and crystals as atom sets with categorical and continuous attributes, adding periodic unit-cell information for crystals.Molecules use unbounded coordinates, whereas crystals additionally use a periodic unit cell, fractional coordinates, and lattice parameters.
- Stage 1: Autoencoder: A VAE jointly reconstructs all-atom systems to learn a shared latent representation containing atom types and continuous structural information.The encoder maps atomic attributes to latent representations, and the decoder reconstructs the input system from those embeddings.
- Stage 1: Autoencoder: The autoencoder balances atom-type, Cartesian-coordinate, fractional-coordinate, and lattice losses, using periodicity-specific objectives for crystals and molecules.Crystals reconstruct atom types, fractional coordinates, and lattice parameters; molecules reconstruct atom types and 3D coordinates.
- Decoding: The unified decoder uses atom types and 3D coordinates for molecules, while combining atom types, fractional coordinates, and lattice parameters for crystals.Molecules are constructed with RDKit and crystals with PyMatGen, allowing one model to share information across both domains.
- Stage 2: Latent diffusion: Latent diffusion uses a Transformer denoiser conditioned on periodic versus non-periodic class labels to generate embeddings that the VAE decoder converts into atomic systems.Classifier-free guidance combines conditional and unconditional predictions during iterative Euler denoising from Gaussian noise.
3 Experimental Setup
Experiments train and evaluate ADiT jointly on molecular and crystal datasets against specialized and general baselines using validity, stability, uniqueness, and novelty measures.
- Datasets: Main experiments use MP20 crystals and QM9 molecules, covering up to 20 crystal unit-cell atoms and up to nine heavy atoms in small organic molecules.MP20 contains 45,231 structures, while QM9 contains 130,000 molecules.
- Training: The VAE and DiT are trained sequentially with AdamW, batch size 256, and up to 5000 epochs on eight V100 GPUs.The VAE uses 51M parameters, while the primary DiT-B configuration uses 130M parameters.
- Inference: Inference uses ODE integration steps T of 500 or 1000 and classifier-free guidance scales γ of 1.0 or 2.0, which work consistently across both domains.Additional ablations examine joint versus dataset-specific training, architecture, regularization, and inference settings.
- Evaluation: ADiT is evaluated by sampling 10,000 crystals and molecules and measuring validity, stability, uniqueness, novelty, and molecule-specific PoseBusters checks.Crystal stability is assessed with density functional theory, while molecular evaluation includes validity and uniqueness metrics.
- Baselines: Crystal baselines include equivariant diffusion and flow-matching models, UniMat, and FlowLLM; molecule baselines include Equivariant Diffusion, GeoLDM, and Symphony.These comparisons span specialized models operating on multimodal product manifolds, latent diffusion, and autoregressive generation.
4 Results
ADiT generates high-quality molecules and crystals, with joint training improving validity and stability across domains. Its Transformer-based design also delivers favorable scaling, speed, and performance on larger molecules.
- Generation quality: ADiT achieves state-of-the-art crystal and molecule generation across validity, stability, uniqueness, novelty, and physical-realism checks.Crystal results improve across validity, stability, uniqueness, and novelty, while QM9 molecules perform strongly on validity, uniqueness, and PoseBusters checks.
- Joint training: Joint training improves validity and stability for both crystals and molecules compared with domain-specific ADiT variants.The jointly trained model exceeds MP20-only performance for materials and QM9-only performance for molecules.
- Scaling: Larger denoisers consistently improve generative performance from 32M to 450M parameters, even with approximately 130K training samples.Training loss and validity rates improve with model size, and correlations suggest further gains from scaling model and dataset size.
- Larger molecules: On GEOM-DRUGS molecules of up to 180 atoms, ADiT matches or exceeds equivariant baselines across validity and PoseBusters metrics.This result uses minimal molecular inductive biases and does not explicitly predict atomic bonds.
- Additional analyses: Additional ablations and visualizations cover metal-organic framework generation, transfer learning, and samples from crystals, molecules, and MOFs.These analyses are reported in the appendices.
5 Discussions
ADiT presents a unified, Transformer-based framework for molecule and material generation, with transfer learning across atomic-system types. The authors identify limited training data and incomplete validation on very large systems as key boundaries.
- Contribution: ADiT unifies molecule and material generation in one latent diffusion model and demonstrates benefits from transfer learning across diverse atomic systems.The framework is positioned as a step toward a broadly applicable foundation model for generative chemistry.
- Design implications: Standard Transformers with minimal inductive biases make ADiT conceptually simpler and computationally more efficient than domain-specific equivariant diffusion approaches.The discussion connects this design to the reported efficiency and broad applicability.
- Limitations: Relatively small training datasets may limit generalization, motivating scaling to larger and more diverse crystal, molecule, and biomolecular datasets.The authors name Alexandria, the Cambridge Structural Database, ZINC, and the Protein Data Bank as future data sources.
- Limitations: ADiT has not been fully validated on larger systems such as metal-organic frameworks or biomolecules containing thousands of atoms.Initial MOF results are described as promising, but the larger-system scope remains future work.
A Related Work
Prior work applies diffusion and flow matching to molecules, crystals, and biomolecules, often with system-specific equivariant formulations. ADiT extends latent diffusion with standard Transformers to jointly model periodic and non-periodic atomic systems.
- Generative models: Diffusion models are established for generating molecules, crystals, and biomolecules, including structure prediction, conditional generation, and structure-based drug design.Equivariant diffusion and latent diffusion are prominent approaches in this literature.
- Latent diffusion: Latent diffusion moves generation into an autoencoder’s lower-dimensional latent space, while Diffusion Transformers provide a scalable denoiser architecture.Classifier-free guidance can incorporate conditioning during latent sampling.
- ADiT’s position: ADiT is presented as the first latent-diffusion approach to jointly generate periodic crystals and non-periodic molecular systems with categorical and continuous atomic attributes.Its autoencoder handles atom types, coordinates, and unit-cell parameters before latent-space generation.
- Equivariance: Equivariant networks capture geometric structure but are computationally expensive and harder to scale, especially when denoisers run hundreds of times during inference.Recent work suggests standard Transformers can perform strongly without explicit equivariance in several 3D structure-generation settings.
- Equivariance: ADiT uses standard Transformers across periodic and non-periodic systems, demonstrating that explicit equivariance and molecular inductive biases are not strictly required for valid, realistic structures at scale.This positions simplicity and scalability as alternatives to specialized equivariant architectures.
B Evaluation Metrics
The evaluation measures generated crystals using structural, compositional, stability, uniqueness, and novelty criteria, while molecule evaluation combines chemical validity with 3D physical-realism checks.
- Crystal validity requires minimum pairwise distances and volume, while compositional validity requires charge neutrality and electronegativity balance.
- Crystal stability requires DFT energy above hull <0.0 eV/atom and at least two unique elements.
- Crystal stability, uniqueness, and novelty are assessed after CHGnet pre-relaxation and DFT relaxation against the Matbench Discovery convex hull.
- Molecule evaluation reports RDKit validity and uniqueness alongside seven Posebusters sanity checks for generated 3D structures.
- The Structure Matcher checks whether two periodic structures are equivalent despite different settings or minor distortions.
C Additional Results
Additional experiments evaluate ADiT on metal-organic frameworks, crystal distributions, sampling stability, and model scaling. Joint training transfers across domains, while larger models and diffusion sampling introduce distinct trade-offs.
- MOF generation: Joint MOF training used 14,000 QMOF structures alongside QM9 and MP20 and evaluated 1,000 generated MOFs with 15 MOFChecker sanity checks.
- MOF generation: 15% overall MOF validity for QMOF-only ADiT decreased to 10% with joint training, although the joint model also retained 91% crystal and 95% molecule validity.
- Crystal distributions: ADiT-generated crystals more often satisfy DFT energy above hull <0.0 eV/atom, while formation-energy and unique-element distributions more closely match MP20 than FlowMM or FlowLLM.
- Crystal distributions: Diffusion-based ADiT and FlowMM oversample the lowest-symmetry P1 spacegroup, whereas FlowLLM oversamples Fm-3m, Pm-3m, and I4/mmm.
- Scaling: ADiT-S achieves a 6.5% crystal S.U.N. rate, compared with 2.8% for FlowMM and 4.7% for FlowLLM.
- Sampling stability: Validity and S.U.N. rates generally stabilize across random seeds after approximately 5,000 samples, while autoregressive FlowLLM shows higher S.U.N. variance.
D Ablation Study
Ablations show that standard Transformer autoencoders and larger latent diffusion denoisers improve ADiT’s reconstruction and generation performance, with inference quality depending on guidance and integration-step settings.
- Joint training: Joint autoencoder training achieves similar or better reconstruction than dataset-specific training and improves generative validity for both crystals and molecules.
- Denoiser architecture: Scaling the DiT denoiser from 32M to 150M to 450M parameters consistently improves generative performance and correlates with better scaling-law metrics.
- Autoencoder architecture: Standard Transformer autoencoders outperform Equiformer-V2 in reconstruction and provide latent spaces more suitable for diffusion, especially for crystal validity.
- Autoencoder regularization: Larger latent dimensions and lower KL regularization weights reduce reconstruction RMSD and increase validity rates for crystals and molecules.
- Sampling hyperparameters: Different classifier-free guidance scales and integration steps are optimal for crystals versus molecules; T = 500 or 1000 with γ = 1.0 or 2.0 works broadly.
E Visualizations
Visualizations examine ADiT’s shared latent space and generated structures across crystals, molecules, and MOFs. The latent embeddings separate system types while generated examples display chemical and structural diversity.
- Shared latent space: PCA embeddings form distinct molecule and crystal clusters, with tighter molecular clustering and more dispersed crystal representations.
- Atom representations: For shared C, N, O, and F atoms, principal component 1 primarily separates molecules from crystals, while principal component 2 correlates with atom type.
- Generated systems: Generated crystals show diverse spacegroups and compositions, while generated molecules show varied chemical structures and conformations.
- Generated systems: ADiT also generates metal-organic frameworks when jointly trained on QMOF150, MP20, and QM9.