[Paper Review] All-atom Diffusion Transformers: Unified generative modelling of molecules and materials

Chaitanya K. Joshi, Xiang Fu+5 more|ArXiv.org|Mar 5, 2025

Machine Learning in Materials Science3 citations

TL;DR

ADiT introduces a single latent diffusion Transformer framework to jointly generate both non-periodic molecules and periodic crystals, achieving state-of-the-art results and successful transfer learning between domains.

ABSTRACT

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

Motivation & Objective

• Motivate a unified diffusion approach across molecular and crystalline systems sharing underlying physics.
• Develop a shared latent representation via a VAE for both molecules and materials.
• Train a diffusion model in latent space to sample new atoms, decoded back to valid structures.
• Demonstrate transfer learning between periodic and non-periodic systems and show scalability.

Proposed method

• Represent atoms as sets with unified attributes (categorical atom type, 3D coordinates, fractional coordinates, lattice parameters).
• Use a Variational Autoencoder (VAE) to map inputs of both molecules and crystals into a shared latent space and reconstruct them.
• Train a Diffusion Transformer (DiT) in latent space to generate new latent embeddings, using classifier-free guidance during sampling.
• Decode latents with the VAE decoder to produce valid molecules (non-periodic) or crystals (periodic) via RDKit and PyMatGen, respectively.
• Adopt a two-stage training: first-stage VAE reconstruction, second-stage latent diffusion with a class-conditioned DiT and data-augmentation for roto-translation and periodic symmetries.
• Evaluate with chemistry-focused metrics (validity, stability, uniqueness, novelty) and Posebusters/DFT sanity checks; compare to dataset-specific baselines.

Experimental results

Research questions

• RQ1 Can a single unified diffusion model generate both molecules and crystals with high validity and realism?
• RQ2 Does joint training across molecular and crystal data improve generation quality via transfer learning?
• RQ3 How does model size and training scale affect performance and speed relative to equivariant baselines?
• RQ4 Is the approach extensible to more complex hybrids like MOFs and other periodic/non-periodic hybrids?
• RQ5 What are the efficiencies in training and inference compared to domain-specific diffusion models?

Key findings

• Jointly trained ADiT achieves state-of-the-art or competitive results for both QM9 molecules and MP20 crystals.
• For crystals, ADiT shows improved validity, stability, and novelty over several baselines in 10,000 samples.
• For molecules, ADiT matches or exceeds validity and uniqueness while meeting Posebusters checks for physical realism.
• Joint training yields higher validity and stability than QM9-only or MP20-only variants, indicating transfer learning benefits.
• ADiTs provide significant inference speedups over equivariant diffusion models, enabling generation of 10,000 samples on a single V100 GPU in under 20 minutes.
• Scaling the DiT denoiser size (DiT-S to DiT-L) yields predictable performance gains, suggesting further gains with more data and larger models.