All-atom Diffusion Transformers: Unified generative modelling of molecules and materials
This work addresses the need for efficient and generalizable generative models in chemistry, offering a unified approach for both molecules and materials, though it is incremental as it builds on existing diffusion and transformer methods.
The paper tackles the problem of generative modeling for 3D atomic systems by introducing ADiT, a unified latent diffusion framework that jointly generates molecules and materials, achieving state-of-the-art results on datasets like MP20, QM9, and GEOM-DRUGS with significant speedups in training and inference.
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