All-atom Diffusion Transformers: Unified generative modelling of molecules and materials

Published: 01 May 2025, Last Modified: 18 Jun 2025ICML 2025 posterEveryoneRevisionsBibTeXCC BY 4.0
TL;DR: First broadly generalizable diffusion model for 3D molecular generation. State-of-the-art results for periodic crystals and non-periodic molecular systems through transfer learning.
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
Lay Summary: Molecules form the foundation of biological life and physical materials in our world. Designing new molecules has potential applications in clean energy, drug discovery, and other areas critical to human health and sustainability. However, current AI systems of generating molecular structures are customised to specific types of molecules. Our work, the All-atom Diffusion Transformer (ADiT), is the first general-purpose architecture for molecular modelling that can learn to generate a wide range molecular systems, from small organic molecules to inorganic periodic crystals.
Application-Driven Machine Learning: This submission is on Application-Driven Machine Learning.
Link To Code: https://github.com/facebookresearch/all-atom-diffusion-transformer
Primary Area: Applications->Chemistry, Physics, and Earth Sciences
Keywords: Latent Diffusion, Transformers, Molecules, Crystals, Diffusion Transformers, Material Science, Chemistry
Submission Number: 11487
Loading