All-atom Diffusion Transformer

Generative models for 3D atomic systems are usually built one substrate at a time. A diffusion model for drug-like molecules assumes a non-periodic point cloud with bonds; a model for inorganic crystals assumes a periodic unit cell with lattice parameters. The underlying physics is identical, but the architectures, the training pipelines, and the released checkpoints are not, so nothing learned about molecular geometry transfers to materials or back again.

The All-atom Diffusion Transformer (ADiT), from FAIR Chemistry at Meta AI with the University of Cambridge and MIT, collapses that split into a single model. An autoencoder maps a unified all-atom representation of any system — molecule or crystal — into a shared latent space, and a diffusion transformer is trained to generate new latent embeddings that the decoder turns back into 3D structures. Only the molecular half falls within biology's remit, but the two halves are not separable: the same weights produce both, and the paper's central empirical result is that training on both makes each better.

ADiT's second departure is architectural. Where most 3D generative chemistry models build rotational equivariance into the network, ADiT uses ordinary Transformers for both the autoencoder and the denoiser, with minimal inductive bias. That choice costs nothing in sample quality and buys a large amount of throughput and scalability. The work was published at ICML 2025.

#Key Features

• One latent space for two substrates: Molecules and periodic crystals are encoded into the same embedding space, so a single checkpoint samples either without a substrate-specific decoder or generative process.
• Transfer between periodic and non-periodic systems: Jointly trained models beat their molecule-only and crystal-only counterparts on validity and stability, evidence that geometric knowledge moves across the two domains.
• No equivariant machinery: Standard Transformers replace equivariant networks throughout, avoiding the per-step cost that limits how far equivariant diffusion models can be scaled.
• Single-stage latent diffusion: Atom types and coordinates are diffused together in a continuous latent space rather than through coupled categorical and continuous processes, which is where the large gains in compositional validity come from.
• Predictable scaling: Denoiser capacity from 32M to 450M parameters improves training loss and generation validity along a clean trend.

#Technical Details

Stage one is a Transformer autoencoder with KL regularization that embeds each atom of a molecule or crystal into a shared latent; standard Transformers proved both better reconstructors and better latent spaces for diffusion than an Equiformer-V2 alternative. Stage two is a DiT denoiser over those latents, evaluated at three sizes — DiT-S (32M), DiT-B (130M, used for the headline results), and DiT-L (450M). Training data spans QM9, MP20 (27K crystals), and GEOM-DRUGS (430,000 molecules of up to 180 atoms), with a further extension to 14,000 metal-organic frameworks from QMOF.

Across 10,000 sampled molecules, the jointly trained model reached 97.43% validity and 96.92% uniqueness on QM9, and 95.3% validity with 100% uniqueness on GEOM-DRUGS, passing 85.3% of PoseBusters checks against 87.5% for SemlaFlow and 59.7% for EQGAT-diff — notable because ADiT does not explicitly predict bonds. On crystals it reached a 5–6% DFT-verified stable-unique-novel rate against 4–5% for prior methods, with the smallest variant trained on MP20 alone hitting 6.5%. Sampling 10,000 structures takes under 20 minutes on a single V100, where equivariant baselines take up to 2.5 hours. The GitHub code is released under CC BY-NC 4.0, restricting commercial use; the Hugging Face checkpoints carry an MIT license and are an independent reproduction by the first author that the authors note may not exactly match the manuscript's reported metrics.

#Applications

For drug discovery, ADiT generates 3D conformations of drug-like molecules directly, supplying structure-based virtual screening and lead generation pipelines with physically plausible candidates rather than SMILES strings that still require conformer generation. Its materials capability serves catalysis and battery research on the same weights. Generation is currently unconditional, so practical inverse design — conditioning on a measured property, scaffolding a motif, or infilling a fragment — requires extending the model, which the authors identify as the natural next step.

#Impact

ADiT is a concrete demonstration that the equivariance-by-construction consensus in 3D generative chemistry is optional: a general architecture trained on more diverse data matches or beats specialized ones while running an order of magnitude faster. That argument echoes the trajectory of biomolecular structure prediction, where AlphaFold 3 showed standard Transformers with Gaussian diffusion handle systems of thousands of atoms. The demonstrated transfer between molecules, crystals, and frameworks suggests a route toward genuinely general foundation models for atomic systems, though the training sets used here remain small by foundation-model standards, and the non-commercial code license limits industrial reuse of the reference implementation.