All-atom Diffusion Transformers: Unified generative modelling of molecules and materials
Authors: Chaitanya K. Joshi
Xiang Fu
Yi-Lun Liao
Vahe Gharakhanyan
Benjamin Kurt Miller
Anuroop Sriram
Zachary Ward Ulissi
Venue: ICML 2025 poster (2025)
Keywords: Latent Diffusion
Transformers
Molecules
Crystals
Diffusion Transformers
Material Science
Chemistry
Primary area: applications->chemistry_physics_and_earth_sciences
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

All-atom Diffusion Transformers:
Unified generative modelling of molecules and materials
Chaitanya K. Joshi 12 Xiang Fu'
Yi-Lun Liao 13 Vahe Gharakhanyan Zachary W. Ulissi '1
Benjamin Kurt Miller
Anuroop Sriram" *1
Abstract
1. Introduction
Generative modelling of the 3D structure of atomic systems has the potential to revolutionize inverse design of new molecules and materials. The current state-of- state-of-the-ar uses diffusion or flow matching models for tasks such as structure prediction (Abramson et al., 2024; Corso et al., 2023; Jiao et al., 2023) and conditional generation (Watson et al., 2023; Ingraham al., 2023; Zeni et al., 2025) for biomolecules and materials, as well for structure-based drug design (Schneuing al., 2024).
Diffusion models are the standard toolkit for gen- erative modelling of 3D atomic systems. How- ever, for different types of atomic systems - such as molecules and materials - the generative pro- cesses are usually highly specific the target system despite the underlying physics being the same. We introduce the All-atom Diffusion Trans- former (ADiT), a unified latent diffusion frame- work for jointly generating both periodic mate- rials and non-periodic molecular systems using the same model: (1) An autoencoder maps uni- fied, all-atom representations of molecules and materials a shared latent embedding space; and (2) A diffusion model is trained to gener- ate new latent embeddings that the autoencoder can decode to sample new molecules or materi- als. 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 re- sults par with molecule and -pystal-specific models. ADiT uses standard Transformers with minimal inductive biases for both the autoen- coder and diffusion model, resulting in signifi- cant speedups during training and inference com- pared to equivariant diffusion models. Scal- ing ADiT up to half a billion parameters pre- dictably improves performance, representing a step towards broadly generalizable foundation models for generative chemistry. Open source code: http://gggghbb.oooc/cccooreeeee dtom-diffusion-tranffomee
All atomic systems share the same underlying physical prin- ciples that determine their 3D structure and interactions. However, we currently do not have a unified formulation of diffusion models across different types of atomic systems such as small molecules, biomolecules, crystals, and their combinations. Most diffusion models are highly specific to each type of system, and involve multi-modal generative processes complex product manifolds of categorical and continuous data types. For example, de novo generation of small molecules is modelled as two independent diffusion processes for the atom types (categorical) and 3D coordi- nates (continuous) of set of atoms (Hoogeboom al., 2022). The denoiser model learns how atom types and 3D coordinates jointly evolve in order to sample new molecules but passes through unrealistic intermediate states during the denoising trajectory. Diffusion models for biomolecules treat groups of atoms as rigid bodies and add a third mani- fold (rotations) into the joint diffusion process (Yim et al., 2023b; Campbell al., 2024). For crystals and materials, the diffusion process needs to additionally handle periodic- ity and operates a joint manifold of atom types, fractional coordinates, lattice lengths, and lattice angles that together define the repeating unit cell (Xic et al., 2022; Miller et al., 2024).
In this paper, we pose the following question: How can we build unified diffusion models that can generate both periodic materials and non-periodic molecular systems?
Equul contribution 'Fundamental AI Research (FAIR) Meta UUiiversity Cambridge 'MIT. Correspondence to: Chaitanya K. Joshi <chaitanya @ ell.cam.c..uk>.
Our solution, the All-atom Diffusion Transformer (ADiT), illustrated in Figure 1, is a latent diffusion model based on two key ideas:
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

All-atom Diffusion Transformers
Stage Autoencoder for reconstruction
Input representation
Output representation
Atom types (B. N, 1) 3D coord. (B. N)
Atom types (B. N, 1) 3D coord. (B. N,
Latent representation d)
Frac. coord. (B. N, 3)
Encoder
Decoder
Frac. coord. (B. N. 3)
Cell lengths (B, 1. Cell angles (B, 1, 3)
lengths (B. 1.
Cell angles (B, 1, 3)
Stage 2: Latent diffusion generative model
Random Gaussian noise (B, d)
Sampled latent representation (B, N, d)
Class label periodicinon-periodic
Denoising with Diflusion Transformer
Figure Unified generative modelling of molecules and materials with All-atom Diffussin Transformers. ADiT performs generative modelling of chemical systems in two stages: (1) A Variational Autoencoder (VAE) learns shared latent space by reconstructing all-atom representations of both molecules (non-periodie and crystals (periodic); and (2) A Diffusion Transformer (DiT) samples new latents from the shared distribution using classifier-free guidance, which are decoded to valid molecules or crystals using the VAE. Our unified latent diffusion framework enables transfer learning and avoids the complexity of multiple diffusion processes on categorical-contiuuuu continuous product manifolds used by equivariant diffusion models.
1. All-atom unified latent representations: We treat both periodic and non-periodic atomic systems as sets of atoms in 3D space and develop unified representation with categorical and continuous attributes per atom. A Variational Autoencoder (VAE) (Kingma & Welling, 2014) embeds molecules and crystals into a shared latent space by training for all-atom reconstruction.
ADīTs are highly scalable, achieving significant speedups in both training and inference compared to equivariant dif- fusion models. By using standard Transformers with mini- mal inductive biases for both the autoencoder and diffusion model, ADIT' can generate 10,000 samples in under 20 min- utes on a single V100 GPU - an order of magnitude faster than baselines which take up to 2.5 hours on the same hard- ware. The practical efficiency of the DiT denoiser compared to equivariant networks allows lion parameters while keeping data scale fixed. Our scaling law analysis demonstrates that generative modelling perfor- mance improves predictably with model size, suggesting further gains are possible through continued scaling.
2. Latent diffusion using Transformers: We perform gen- erative modelling in the latent space of the VAE en- coder using a Diffusion Transformer (DiT) (Rombach et al., 2022; Peebles Xie, 2023). During inference, classifier-free guidance (Ho & Salimans, 2022) enables sampling new latents that can be reconstructed to valid molecules crystals using the VAE decoder.
scale ADiT to half bil-
All together, our work is the first to develop unified genera- tive models for both periodic and non-periodic atomic sys- tems, with state-of-the-at performance both molecules and crystals. ADiTs represent step towards broadly gener- alizable foundation models for generative chemistry.
ADIT's be trained jointly both periodic and non- periodic atomic systems, demonstrating broad generalizabil- ity. Training a single unified model on the QM9 molecular and MP20 materials datasets leads to statc-of-the-as perfor- mance in both domains, exceeding specialized equivariant diffusion models on physics- validations. DFT calcu- lations reveal that ADiTs generate stable, unique, and novel crystals at 5-6% S.U.N. rate, 25% improvement upon the 4-5% rates of previous methods. Joint training yields higher validity rates than QM9-only or MP20-only ADIT variants, demonstrating successful transfer learning between periodic and non non-periodic systems. ADiTs also match or exceed state-of-the-un equivariant models on the GEOM-DRUGS dataset of molecules with hundreds of atoms.
2. All-atom Diffusion Transformers
We use latent diffusion (Rombach et al., 2022) to unify generative modelling across periodic and non-periodic atomic systems. Our approach consists of two stages: (1) A Variational Autoencoder (VAE) (Kingma & Welling, 2014) learns a shared latent space by jointly reconstructing all-atom representations of both molecules and materials; and (2) A Diffusson Transformer (Peebles & Xie, 2023) generates new samples from this latent space
Overview.
2

---

All-atom Diffusion Transformers
which can be decoded into valid molecules crystals us- ing classifier-free guidance (Ho & Salimans, 2022). Latent diffusion shifts the complexity of handling categorical and continuous attributes into the autoencoder, enabling a sim- plified and scalable generative process in latent space. We discuss how our contributions are contextualized w.r.t. re- lated work in Appendix A.
via data augmentation. In Appendix D, we also ablated roto-translation equivariant VAEs based Equiformer-22 (Liao et al., 2024).
Algorithm 1: Pseudocode for VAE encoder
Input: 3D atomic system (fa,), (at). (fi), (11,22,1a))
Output: Latent reprenstations {++}
# Project inputs to dlmodel
1. h, Embeddingga.)
h, E RRoss
2.1. Stage I: Autoencoder for reconstruction
2. hi = h + Linear(Swishh) wwwsssssssiiiaaaiia Linear(x:)))
Unified representation of 3D atomic systems. All peri-
h; = h; + LLieer(SSwih((Lingg (f.))
odic and non-periodic atomic systems can be represented in a unified format as sets of atoms in 3D space (Duval et al., 2023). The key difference is that crystals require an additional periodic unit cell, while molecules have un- bounded coordinates. A crystal or molecule with N atoms is represented as a multi-modal object:
# Apply encoder network
4. {h.} = TransformerrEnoodr((())
# Down-project to mean pz and std oz
5.
Ns. = Linear(h:)
Mss C Rd
log T2 Lincar (h,)
Tzg E Rd
# Sample latents Z
2 Mz. t. Oe E e~ 10
2, E IRd
Atom types A- = {a
E Z
Algorithm 2: Pseudocode for VAE decoder D
3D coordinates X = {
E R
Input: Latent reprenstations {2:] Output:
Fractional coordinates F == [
E [0, 1)sx
3D atomic system (fa/), (2f). (M)) (4 (,))
Unit cell/lattice L == 42,
E Rx3
# Up-project latents to dmmodll
h2 = Linear(2;)
h, Rdmm
# Apply decoder network
The 3D coordinates X are in nanometers, and the frac- tional coordinates F are in the range [0, 1). The lattice
2. { = TransformerEEecodr(((,,,
# Predict outputs
matrix L represents a parallelepiped defining the shape of the repeating unit cell, and fractional coordinates are com-
3. df = argmax(Linear) (hx))
d ee
T Linear(h:)
of R3
puted as the inverse of the unit cell matrix multiplied by the 3D coordinates: F L L- We use Niggli reduction to uniquely determine the unit cell parameters for crystals (Grosse- Kunstleve et al., 2004). For non non-periodic molecules, we set the unit cell parameters and fractional coordinates to null values Φ
5. f Linear(h.)
f R
6. [4 1.144 = Linear 1 N Eim hs)
1' e R3
Reconstruction loss. We compute the loss for the predicted atom types A' via cross-entropy:
VAE architecture.
We use a Variational Autoencoder (VAE) to learn a shared latent representation of molecules and materials using a reconstruction objective. Given an input 3D atomic system (A, (A, F, L), an encoder maps each atom's 's attributes latent representation Z:
N
L
1 N
Σ
CrossEnt(a;, a()
For the predicted 3D coordinates X', we use the mean squared error (MSE) reconstruction loss after zero-centering both sets of coordinates:
z = E(A,XFF).
(1)
N
where 2 {=) E RdxN encodes information about the categorical atom type and continuous coordinates (unit cell parameters are encoded implicitly in the fractional co- ordinates). The decoder D reconstructs the input atomic system from the latent embedding:
34 = 24
:4
7=
x
N
N
1=1
1=1
N
1
Lx
Σ
|1,-
2/1
(4)
3N
We compute the reconstruction loss for the predicted frac- tional coordinates F" using MSE as well:
A, L'=
We describe the pseudocode for VAE encoder and de-
coder operations in Algorithms I and 2, respectively. For
CF
1 3N
Σ 1f. - 2111
N
(5)
the architecture of the encoder & and decoder D. we
used the standard Transformer (Vaswani et al. (2017), and learn symmetries
For the predicted lattice vectors L', we first convert to rotation-invariant lattice parameters: three side lengths of
torch. nn.
3

---

All-atom Diffusion Transformers
the unit cell L = (a, cf E RIx3, and three internal an- gles between them L. = (.) E [60", 1200|1x3 as described in Miller et al. (2024), We then compute the MSE reconstruction loss between the predicted and ground truth lattice parameters:
coordinates, constructing the molecule via RDKit. For peri- odic crystals, we combine the atom types, fractional coordi- nates, and lattice parameters to build the crystal structure us- ing PyMatGen. This split decoding strategy allows single unified model to share information between both domains while still respecting their distinct geometric constraints, enabling effective transfer learning between periodic and non-periodic systems.
Lr, =
3 (
+ (b- + (c-c)))
(6)
LL. (aa-a')² + (3 8) + (7 - -1))
2.2. Stage 2: Latent diffusion generative model
Note that in Lt,, we normalize the predicted and groundtruth lengths by the cube root of the number of atoms to account for the scaling of the unit cell with the number of atoms, following Xie et al. (2022). All angles are converted from degree to radians for numerical stability.
Diffusion formulation. We use Gaussian diffusion or flow matching as our generative framework, which iteratively de- noises latent samples from base distribution into samples from target distribution (Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Ho et al., 2020; Lipman et al., 2023). Our formulation uses linear interpolation between standard normal base distribution and the target distribution of VAE encoder latent representations of 3D atomic systems (we describe it in terms of flow matching. though both formu- lations are equivalent; see Gao et al. (2024)). Thus, the diffusion model is trained after training the first stage VAE.
The autoencoder is trained with a weighted reconstruction loss to balance the relative magnitudes of the various losses:
fax: = AAAA + Xxfxx + ArLF + XI, + XI. Lr
(8)
Depending on whether a training sample is periodic or non- periodic, we use different reconstruction loss weights:
Our model learns to generate a set of N latent represen- tations Z = {zi J. where each latent E Idd encodes information about one atom's 's type, coordinates and unit cell, which can be decoded to a valid 3D atomic system using the VAE decoder D. During training, given an input 3D atomic system (A, X, F, L), we first encode it to latent representation Z using the VAE encoder &. We denote Z as (1), a 'clean' training sample at time t = 1. We then sample a random initial latent (0) at time t = 0 from a di d-dimensional standard normal distribution (o 1)d, and perform zero-centering by subtracting the per-channel mean of Z Z We then use linear interpolation to construct a "noisy' interpolated sample Z(1) at randomly sampled time step t ~ L((0 1):
AA
Xx
XF
AL
ALa
Periodic
1.0
0.0
10.0
1.0
10.0
Non-periodie 1.0
10.0
0.0
0.0
0.0
Thus, the overall loss for periodic crystals trains the model to reconstruct the atom types, fractional coordinates and lattice parameters while ignoring the predicted 3D coordi- nates. Similarly, the overall loss for non-periodic molecules trains the model to reconstruct the atom types and 3D coor- dinates while ignoring the predicted fractional coordinates and lattice parameters.
Regularization. We use three regularization techniques to learn robust, informative latent representations: (1) A bottleneck architecture with latent dimension d significantly smaller than the encoder/decoder hidden dimension dmodel (e.g.. d 8 vs dmosel = 512). (2) A per-channel KL di- vergence penalty λKL ( (Z (2, 2) ) || (0,, 1)a) added to equation 8, following Rombach et al. (2022). (3) Denoising training with 10% of atoms having their types masked and coordinates perturbed by (0, 0.1) Gaussian noise. For non-equivariant encoders/decodees,, we also ran- domly rotate and translate each sample during training to learn symmetries via data augmentation.
z (1 Z (0) + t tz)
Thus, we can define groundtruth conditional vector field me ( ( along the path from the noisy latents Z at time step & to the clean latents Z
as:
(1) z(1)
u. (Z(0))(())
(10)
Samples from the base distribution can be transformed to samples from the target distribution by integrating the vector field we ( (t) Z (1) over time t.
Decoding latents to atomic systems. During infer- ence or sampling from the DiT, the desired output type peeriodiccnon-periodii determines how we process the de- coder outputs. The VAE decoder D generates four attributes for each system: (1) atom types, (2) 3D coordinates, (3) fractional coordinates, and (4) lattice parameters. For non- periodic molecules, we only utilize the atom types and 3D
The goal of conditional flow matching is to train a denoiser network FF to match this conditional vector field u, To do so, the denoiser takes as input the intermediate noisy latents (t) at time step and an additional class label e (described
'In practice, we set a minimum value for time step Imin 0.01.
4

---

All-atom Diffusion Transformers
subsequently) predict the final clean latents Z(t);;
from Gaussian noise O we iteratively denoise using the DiT model F for T steps. At each step, we perform Fuler integration of the vector field to gradually transform the noisy latents towards the target distribution. While we cur- rently use simple Fuler integration for efficiency, adaptive ODE solvers could potentially improve performance (Ma et al., 2024). Finally, we decode the denoised latents (1)) to a valid 3D atomic system using the VAE decoder D.
z = F(2 ,,,).
(11)
The denoiser is trained by minimizing an MSE loss be- tween the resulting predicted conditional vector field and the groundtruth conditional vector field:
Cim
Σ
N
(1)
2
(t)
(1)
(0)
2
(12)
N
1-t
1- t
i=1
N
Algorithm 3: Pseudocode for DiT sampling
1
1
Σ
|
(1)
200)
(1 - t)2
N
2
Input:
Class label c, num, integration steps T, cfg. scale 7
Output: Generated sample (A, X, F,F.II L)
# Sample initial noisy latents Z ((0) at =
In practice, we follow Yim et al. (2023a) and clip the value of at 0.9 to prevent numerical instability.
1.
0 = {2 (O, 1)d
2.
A A 1/T
# Step size
Denoiser architecture. As the denoiser network F, we use class-conditional Diffusion Transformer (DiT) (Peebles & Xie, 2023). The DiT largely follows standard Transformer architecture with the conditioning information incorporated via adaptive layer norm with zero- pero-initialization, which re- places all layer norm operations. For class conditioning, we use a binary embedding to denote whether the system being generated is periodic (crystal) or non pon-perio (molecule). This conditioning allows the model to learn domain specific features while sharing most parameters. During training, apply class label dropout with 10% probability to able classifier-free guidance during inference. We also in- corporate self-conditioning (Yim et al., 2023b) where the denoiser's prediction from the previous timestep is concate- nated to the current input with 50% dropout probability during training. While we currently only condition on the periodic/non--eriodde class label, the DiT architecture can incorporate additional conditioning signals like target prop- erties or geometric constraints to enable controlled genera- tion. This represents promising direction for future work in inverse design applications.
# Denoising loop
3.
for t in linspace(O.0., 1.0, T):
4.
Z.uml = ( ,f,,)
Condiiional prediction
Zimcond = F(Z (1,6) ..)) # &Unconditional prediction
# +Conditioning via Hassifier-free guidance
2'= (111->) Z + Zoomd
# Euler integration step
Z(+A0) 2 + Al. 2'-2'%
# Decode latents to 3D atomic system (Algorithm 2)
8. A, F, = D(Z
3. Experimental Setup
Datasets. For our main experiments, we train models on pe- riodic crystals from MP20 and non-periodic molecules from QM9, representing two distinct domains of atomic systems. MP20 (Xie et al., 2022) contains 45,231 metastable crystal structures from the Materials Project (Jain al., 2013), each with up 20 atoms in its unit cell and spanning 89 different element types. QM9 (Wu et al., 2018) 1) consists of 130,000 stable small organic molecules containing up to nine heavy atoms (C, N, O, F) along with hydrogens. We split the data following prior work (Xie et al., 2022: Hoogeboom et al., 2022) to ensure fair comparisons. We also include results the GEOM-DRUGS dataset of 430,000 large organic molecules up to 180 atoms (Axelrod & Gomez-Bombarelli, 2022), as well as the QMOF dataset of 14,000 metal-organic framework structures (Rosen et al., 2021).
Data augmentation. The DiT denoiser is trained with data augmentation to learn -rantsaaalati and periodic symmetries in the VAE's latent space. During training, each input system coordinates are randomly rotated and translated, and then converted to latents via the frozen VAE encoder & before being input to the DIT.
Sampling with classifier-free guidance. To generate new atomic systems from the trained diffusion model, we use classifier-free guidance (Ho& & Salimans, 2022) to steer the sampling process. At each denoising step, we compute both conditional prediction based on the periodic/non-perin periodic class label e and an unconditional prediction with null class label Φ. The final prediction is a weighted combination of these using guidance scale 7 allowing control over how strongly the generation follows the class conditioning. The full sampling procedure outlined in Algorithm 3. Starting
Training and hyperparameters. We sequentially train the first-stage VAE and then the second-stage DiT using AdamW optimizer with a constant learning rate 1e 4, 4, no weight decay, and batch size of 256. We use exponential moving average (EMA) of DiT weights over training with decay of 0.9999. Both models are trained to convergence for at most 5000 epochs up to 3 days on 8 V100 GPUs
For the first-stage VAE, we use a standard Transformer as both encoder E and decoder D with hidden dimension
5

---

All-atom Diffusion Transformers
Table 1: Crystal generation results on MP20. We report validity, stability, uniqueness, and novelty rates for 10,000 sampled crystals. ADiT shows improved performance over diffusion baselines across all metrics. We see significant gains for compositional validity due a single diffusion process in the latent space, as opposed to joint continuous and categorical diffusion for baselines. Joint training with both molecular and crystal data improves crystal generation performance over MP20-only models. (Stable: DFT Ehull <0.0, metastable: DFT Bhull <0.1, denotes results from MatterGen-MP for 1024 sampled crystals," denotes results replicated using the same DFT setup ADiT.)
Validity Rate (%)
Metastable rate (%)
Stable rate (%) f
M.S.U.N. rate (%) f
S.U.N. rate (%) f
Model
Structure
Composition
Overall
CDVAE
100.00
86.70
1.6
DiffCSP
100.00
83.25
5.0
3.3
MP20-only
UniMat
97.2
89.4
FlowMM
96.85
83.19
80.30
30.61
4.6
22.51
2.81
FlowLLM
99.94
90.84
90.81
66.91
13.91
26.31
4.71
MatterGen-MP
78
13*
21*
MP20-only ADiT
99.58
90.46
90.13
81.6
14.1
25.91
4.7
Jointly trained ADIT
99.74
92.14
91.92
81.0
15.4
28.2
5.3
dmodel = 512,8 attention heads, and 8 layers (51M param- eters total). The latent dimension is set to d = with KL regularization weight 'KL. = 1e 5 and 10% denoising per- turbation during training. For the second-stage DiT denoiser, report results primarily using DiT-B configurations: hid- den dimension dmodcc = 768, 12 attention heads, 12 layers, and 130M parameters total. We also evaluate smaller DīT-S (32M parameters) and larger DiT-L (450M) variants.
from data; (3) FlowLLM (Sriram et al., 2024), a two-stage framework which first finetunes the autoregressive Llama 2 language model on crystal structures (Touvron et al., 2023; Gruver et al., 2024), and then trains FlowMM with samples from the language model as the base distribution and MP20 as the target distribution.
For molecule generation on QM9, we compare to: (I) Equivariant Diffusion (Hoogeboom et al., 2022), a roto- translationally equivariant diffusion model operating on a multi-modal product manifold; (2) GeoLDM (Xu al., 2023), an alternative latent diffusion model using Equiv- ariant Diffusion in the latent space of roto-translationally equivariant autoencoder; (3) Symphony (Daigavane et al., 2024), an equivariant and autoregressive generative model that iteratively builds a molecule atom-by-atom.
Two key inference-tine byperparameters are the number of ODE integration steps T and the classifier-free guidance scale 7. We find T 500 or 1000 with 7 1.0 2.0 consistently works well for both molecules and crystals. Ad- ditional ablation studies comparing joint vs. dataset-specific training, architecture variants, regularization techniques, and inference settings are presented in Appendix D.
Evaluation metries. We evaluate the ability of ADiTs to sample valid and realistic molecules and crystals. Following prior work (Xie et al., 2022; Hoogeboom et al., 2022), we sample 10,000 crystals and molecules each and compute validity, stability, uniqueness and novelty rates using den- sity functional theory (DFT) for crystals as well as validity, uniqueness and Posebusters sanity checks (Buttenschoen et al., 2024) for molecules. Detailed descriptions of all evaluation metrics are provided in Appendix B.
4. Results
State-of-the-art crystals and molecule generation. Re- sults for crystal generation in Table 1 show that ADïTs generate high-quality crystals compared to baseline diffu- sion models, achieving improved performance across va- lidity, stability, uniqueness, and novelty metrics for 10,000 sampled crystals, with significant gains for compositional validity due to a single diffusion process in the VAE latent space rather than joint continuous and categorical diffusion. For molecule generation, ADiTs achieve state-of-tee of-the-art per- formance on validity and uniqueness metrics across 10,000 sampled molecules, as shown in Table 2(a), while Pose- busters sanity check metrics in Table 2(b) further confirm that ADiTs generate physically realistic molccular struc- tures, matching or exceeding baseline models across mea- sures like double bond flatness, reasonable internal energy and lack of steric clashes.
Baselines. We compare ADIT trained jointly both QM9 and MP20 to molecule-only and crystal-only ADiT vari- ants, as well as mate-offthe-aa baselines for both datasets. For crystal generation on MP20, we compare to: (1) three equivariant diffusion and flow matching- based models op- erating on multi-modal product manifolds: CDVAE (Xie al., 2022), DifICSP (Jiao et al., 2023), and FlowMM (Miller et al., 2024); (2) UniMat (Yang et al., 2024), non-equivariant diffusion model which learns symmetries
6

---

All-atom Diffusion Transformers
Table 2: Molecule generation results on QM9. We report (a) validity and uniqueness rates, as well as (b) % pass rates on 7 sanity checks from Posebusters for 10,000 sampled molecules. ADiTs match or improve performance w.r.. baselines, and sample physically realistic structures. Joint training with both molecular and crystal data improves molecular generation performance over QM9-only models. (* denotes models which explicitly generate hydrogen atoms.)
(a) Validity results
(b) PoseBusters results
Model
Validity (%) f
Unique (%) t
Test (% pass) f
Symphony
Eq. Diff.
ADiT
Equivariant Diffusion
97.50
96.71
Atoms connected
99.92
99.88
99.70
QM9-only
Equivariant Diffusion"
91.90
98.69
Bond angles
99.56
99.98
99.85
G
93.80
98.82
Symphony"
83,50
97.98
Bond lengths
98.72
100.00
99.41
QM9-only ADiT
96.02
97.76
Ring flat
100.00
100.00
100.00
QM9-only ADIT
92.19
97.90
Double bond flat
99.07
98.58
99.98
Internal energy
95.65
94.88
95.86
Jointly trained ADiT
97.43
96.92
No steric clash
98.16
99.79
99.79
Jointly trained ADIT
94.45
97.82
ADIT
ADIT5 (12M)
(130M)
Time to sample 10K molecules (mins)
ADIT- (130M)
402
GoOLDN
10K
1325
100
10:25
50
100
200
3.00
1.000
Number integration steps
Number integration steps
(a) Crystals MP20
(b) Molecules - QM9
Figure 2: ADiTs are significantly faster than equivariant diffusion models. We plot the number of integration steps for ADiTs and equivariant diffusion models vs. time to generate 10,000 samples on a single V100 GPU. ADiTs scale significantly better with the number of integration steps compared equivariant diffusion.
Joint training improves performance. Table 1 and Ta- ble 2 also show that jointly trained ADiTs (trained on both QM9 and MP20 together) exceed the performance of the MP20-only or QM9-only ADiTs for materials or molecules, respectively. Joint training improves validity and stability rates for both crystals and molecules, demonstrating effec- tive transfer learning between periodic and non-periodic atomic systems. These results validate that ADiT' effec- tively model diverse types of atomic systems within a single architecture.
tions between model size and performance metrics suggest further gains are possible from scaling both model size and data Alexandria (2M inorganic crystals), ZINC (250M molecules), and the Protein Data Bank (200K biomolecu- lar complexes) present promising opportunities for dataset scaling.
Speedup compared equivariant diffusion. ADiTs achieve significant inference speedup compared to equivari- ant diffusion under the same hardware conditions, as shown in Figure 2. When generating 10,000 samples on V100 GPU, ADiT's based on standard Transformers leads to better scaling with integration steps compared to FlowMM (Miller et al., 2024) for crystals and GeoLDM (Xu al., 2023) for molecules. both of which computationally intensive equivariant networks denoisers. It is significantly more practical to scale up Transformers than equivariant networks, as seen by the faster inference speed of ADiT-B compared to 100x smaller equivariant baselines.
Scaling up ADiT denoiser improves performance. In Figure 3, we see that generative modelling performance predictably improves as we scale the DiT denoiser from parameter counts of 32M (DiT-S) to 130M (DiT-B) all the way 450M (DiT-L),, even with current modest dataset size of 130K total samples. The diffusion training loss and validity rates consistently improve with larger model sizes, showing clear benefit from scale. Strong correla-

---

All-atom Diffusion Transformers
3.0
Pearson: -1.00
25
ADIT-8 (130M)
(450M)
1.08
Spearman: -1.00
Train Loss
2.0
1.06
Train Loss
1.04
Ep.
1.02
0.0
250
500
750
1000
1250
1500
1750
2000
3.0
3.5
4.0
4.5
5.0
5.5
6.0
6.5
Epoch
[Ng(Nume parameters M)
Pearson: 0.91
Spearman: 1.00
Crystal validity rate (%6
Ep. 2000; Crystal validity rate
0.91
0.4
0.90
(32M)
ADIT-8 (130M)
ADIT-L (450M)
500
1000
1250
1500
1750
2000
Epoch
lNg(Number of parameters in M)
Ep. 2000: Molecule validity rate (%)
Pearson: 0.94
Spearman: 1.00
Molecule validity rate (96
0.8
0.94
0.6
0.92
ADIT-S (32M)
(130M)
ADIT (450M)
0.86
1080
1750
Epoch
[Ng(Number of parameters in M)
Figure 3: Scaling up ADIT improves performance. We show the effect of increasing the number ADiT denoiser parameters on the training loss and generation validity rates. Left: training loss and validity rates vs. epochs. Right: Correlation plots for training loss and validity rates epoch 2,000 ADiT parameters (in Millions).
Extension to larger GEOM-DRUGS molecules. To
standard Transformer architecture with minimal molecular inductive biases and, unlike equivariant baselines, does not explicitly predict atomic bonds.
demonstrate the scalability of the ADIT architecture to larger systems, we experiment with the GEOM-DRUGS dataset of 430,000 molecules of up to 180 atoms. Our setup follows Vignac et al. (2023) and we compare to state-of-the-art equivariant diffusion (Le al.. 2024) and flow matching (Irwin al., 2025) baselines. In Table ADiT on par with equivariant models across validity and PoseBusters metrics. This is notable because ADIT is based on the
Additional results and visualizations. For additional results and ablation studies, including extensions metal- organic framework (MOF) generation, see Appendix C and Appendix D, respectively. In Appendix E, we further ana- lyze ADiT's joint latent space via PCA and visualize some sampled crystals, molecules, and MOFs.
8

---

All-atom Diffusion Transformers
Table 3: Molecule generation results on GEOM-DRUGS. Left: Validity, uniqueness and % pass rates on Posebusters for 10,000 sampled molecules (" PoseBusters results taken from Buttenschoen et al. (2025)). ADiT with minimal molecular inductive biases matches or exceeds state-of-the-art equivariant diffusion baselines, which explicitly predict atomic bonds, Right: We plot the number of integration steps for ADiTs and SemlaFlow vs. time to generate 10,000 molecules on single A100 GPU. ADiTs scale favorably with the number of integration steps compared to SemlaFlow, highly optimized equivariant diffusion model.
Metric (% pass) 1
EQGAT-diff
SemlaFlow
ADiT
920
Validity
94.6
93.9
95.3
Time to samelt 10K moleniies (mas))
40
ADTILL 1400
ADITS
Uniqueness
100.0
100.0
100.0
29
Semiat (49N)
Atoms connected
84.4
92.3
93.0
Bond angles
86.9
94.8
92.3
1925
150
330-
Bond lengths
87.0
94.6
92.5
Ring flat
87.0
94.9
95,4
220
Double bond flat
87.0
94.2
95.3
130
Internal energy
86.8
94.8
91.3
0-
No steric clash
82.9
92.0
91.8
10
300
350
500
750
100n
Number of integration steps
PoseBusters valid
59.7
87.5
85.3
5. Discussions
systems, could enable powerful transfer learning capabil- ities especially valuable for low-data domains. Finally, our current models only perform unconditional generation extending to conditional generation based on experimental properties, motif scaffolding, or molecular infilling would enable practical inverse design applications in drug discov- ery, materials science, and beyond.
Our work represents significant step towards a broadly applicable foundation model for generative chemistry. We have introduced unified latent diffusion framework for generating molecules and materials using a single model, and demonstrated the benefits of transfer learning from di- verse atomic systems. The All-atom Diffusion Transformer (ADiT) obtains state-of-the-ar ther performance in molecular generation using an architecture based primarily on standard Transformers with minimal inductive biases. This makes ADiT conceptually simpler and computationally more ef- ficient than previous domain-spednt -peciffc approaches based on equivariant diffusion.
Impact Statement
Our work represents a step towards foundation models for chemistry, which could have significant benefits and risks for society. Our architecture unifies generative modelling of molecules and materials, enabling transfer learning across diverse atomic systems. Future Al systems for inverse de- sign based on our architecture could accelerate the discovery of new materials and drugs, with potential applications in clean energy, drug discovery, and other areas critical hu- man health and sustainability. However, our methods could also be misused to design harmful molecular systems. We emphasize the importance of developing safeguards and guidelines for responsible deployment of these tools, and encourage the research community to carefully consider both positive and negative societal impacts as foundation models for chemistry continue to advance.
However, several limitations point to promising future di- rections. First, we currently use relatively small datasets for training, which may limit model generalization. Scaling to larger and more diverse datasets such as Alexandria and the Cambridge Structural Database for crystals, ZINC for small molecules, and the Protein Data Bank for biomolecu- lar complexes could significantly improve performance and enable learning of broadly applicable chemical principles. Second, while we demonstrate success on small molecules and crystals of to hundreds of atoms, we have not yet fully validated our approach on larger systems such as metal- organic frameworks or biomolecules containing thousands of atoms, though initial results for MOF generation are promising. Recent work biomolecular structure predic- tion with AlphaFold3 (Abramson al., 2024) demonstrates that simple Gaussian diffusion models with standard Trans- formers can effectively handle systems with thousands of atoms. Adapting ADIT' to larger scales, while maintaining its unified representation across periodic and non-periodis
References
Abramson, J., Adler, J., Dunger, J., Evans, R., Green, T.
Pritzel, A., Ronneberger, O., Willmore, L., Ballard, A. J.,
Bambrick, J., et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 2024.

---

All-atom Diffusion Transformers
Axelrod, S. and Gomez-Bombarelli, R. Geom, energy- annotated molecular conformations for property predic- tion and molecular generation. Scientific Data, 2022.
Deng, B., Zhong, P., Jun, K., Riebesell, J., Han, K., Bartel, C. J., and Ceder, G. Chgnet as pretrained universal neural network potential for charge-informed atomistic modelling. Nature Machine Intelligence, 2023.
Batatia, I., Benner, P., Chiang, Y., Elena, A. M., Kovács, D. P., Riebesell, J., Advincula, X. R., Asta, M., Avaylon, M., Baldwin, W. J., et al. A foundation model for atomistic materials chemistry, arXiv preprint arXiv:2401.00090 2023.
Dhariwal, P. and Nichol, A. Diffusion models beat gans image synthesis. Advances in neural information processing systems, 2021.
Duval, A., Mathis, V., Joshi, C. K., Schmidt, V., Miret, S., Malliaros, F. D., Cohen, T., Lio, P., Bengio, Y., and Bronstein, M. A hitchhiker's guide to geometric gnns for 3d atomic systems. arXiv preprint arXiv:2312.07511, 2023.
Betker, J., Goh, G., Jing, L., Brooks, T., Wang, J., Li, I.., Ouyang, L., Zhuang. J., Lee, J., Guo, Y., et al. Improving image generation with better captions, 2023.
Brooks, T., Peebles, B., Holmes, C., DePue, W., Guo, Y., Jing, L., Schnurr, D., Taylor, J., Luhman, T., Luhman, E., al. Video generation models world simulators. 2024.
Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F,, et al. Scaling rectified flow transformers for high-resot high-resolution image synthesis. In International Conference on Machine Learning, 2024.
Buttenschoen, M., Morris, G. M., and Deane, C. M. Pose- busters: Ai-based docking methods fail to generate physi- cally valid poses or generalise to novel sequences. Chem- ical Science, 2024.
hham-Shepherd, D. and Aspuru-Guzzie A. Language models generate molecules, materials, and protein binding sites directly in three dimensions as xyz, cif, and pdb files. arXiv preprint
Buttenschoen, M., Ziv, Y., Morris, G. M., and Deane, C. An evaluation of unconditional 3d molecular generation methods. In ICLR Workshop on Generative and Experi- mental Perspectives for Biomolecular Design, 2025.
2023.
Fu, X., Xie, T., Rosen, A. S., Jaakkola, T. S., and Smith, J. A. MOFDiff: Coarse-grained diffusion for metal-organic framework design. In ICLR, 2024.
Campbell, A., Yim, J., Barzilay, R., Rainforth. T., and Jaakkola, T. Generative flows on discrete state-spaces: Enabling multimodal flows with applications protein co-design. In Forty-first International Conference Machine Learning, 2024.
Gao, R., Hoogeboom, E., Heek, J., Bortoli, V. D., Murphy, K. P., and Salimans. T. Diffusion meets flow matching: Two sides of the same coin. 2024. URL https: /d ffusionflow g thub. 10/-
Chu, A. E.,, Kim, J., Cheng, L., El Nesr, G., Xu, M., Shuai, R. W., and Huang, P.-S. An all-atom protein generative model. Proceedings of the National Academy of Sciences, 2024.
Grosse-Kunstled R. W., Sauter, N. K., and Adams, P. D. Numerically stable algorithms for the computation of reduced unit cells. Acta Crystallogruphica Section A: Foundations of Crystallography, 2004.
Corso, G., Jing, B., Barzilay, R., Jaakkola, T., et al. Diffdock: Diffusion steps, twists, and turns for molecular docking. In International Conference on Learning Representations, 2023.
Gruver, N., Sriram, A., Madotto, A., Wilson, A. G., Zitnick, L., and Ulissi, Z. W. tine-tuned language models generate stable inorganic materials as text. In The Twelfth International Conference on Learning Representations, 2024.
Dai, X., Hou, J., Ma, C.-Y., Tsai, S., Wang, J., Wang, R., Zhang, P., Vandenhende, S., Wang, X., Dubey, A., et al. Emu: Enhancing image generation models ing photogenic needles in haystack. arXiv preprint arXiv:2309. 15807, 2023.
Harris, C., Didi, K., Jamasb, R., Joshi, C. K., Mathis, S. V., Lio, P., and Blundell, T. Benchmarking generated poses: How rational is gructure-based drug design with generative models? arXiv preprint arXiv:2308.074 2023.
Daigavane, A., Kim, S. E., Geiger, M., and Smidt, T. Sym- phony: Symmetry equivarraan point-centered spherical harmonics for 3d molecule generation. In ICLR, 2024.
Ho, and Salimans, T. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.1259% 2022.
Davies, D. W., Butler, K. T., Jackson, A. J.. Skelton, J. M., Morita, K., and Walsh, A. Smact: Semiconducting ma- terials by analogy and chemical theory. Journal of Open Source Software, 2019.
Ho, J., Jain, A., and Abbeel, Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 2020.
10

---

All-atom Diffusion Transformers
Hoogeboom, E., Satorras, V. G., Vignac, C., and Welling, M. Equivariant diffusion for molecule generation in 3d. In International conference machine learning, pp. 8867- 8887. PMLR, 2022.
Lipman, Y., Chen, R. T. Q.. Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In International Conference Learning Representations, 2023.
Ingraham, J. B., Baranov, M., Costello, Z., Barber, K. W., Wang. W., Ismail, A., Frappier, V., Lord, D. M., Ng- Thow-Hing, C., Van Vlack, E. R., et al. Illuminating protein space with programmable generative model. Nature, 2023.
Ma, N., Goldstein, M., Albergo, M. S., Boffi, N. M., Vanden- Eijnden, E., and Xie, S. Sit: Exploring flow and diffusion- based generative models with scalable interpolant trans- formers. In ECCV, 2024.
Martinkus, K., Ludwiczak, J., Liang, W.-C., Lafrance-
Vanasse, J., Hotzel, I., Rajpal, A., et al. Abdiffuser: full- atom generation of in-ritro functioning antibodies. Ad- vances in Neural Information Processing Systems, 2024.
Irwin, R., Tibo, A., Janet, J. P., and Olsson, S. Semlaflow efficient 3d molecular generation with latent attention and equivariant flow matching. In The 28th International Conference on Artificial Intelligence and Statistics, 2025.
Miller, B. K., Chen, T., Sriram, A., and Wood, B. M. Flowmm: Generating materials with riemannian flow matching. In Forty-first International Conference Machine Learning, 2024.
Jablonka, K. M. MOFChecker v0.9.6, 2023. URL https:
//github. com/ wtt.//////jaaaaat japppppapummmmmmmmmm wwtpa/////ccccc....ccmm
Jain, A., Ong. S. P., Hautier, G., Chen, W., Richards, W. D., Dacek, S., Cholia, S., Gunter, D., Skinner, D., Ceder, G., et al. Commentary: The materials project: A materials genome approach to accelerating materials innovation. APL materials, 2013.
O Pinheiro, P. O., Rackers, J., Kleinhenz, J., Maser, M., Mahmood, O., Watkins, A., Ra, S., Sresht, V., and Saremi, S. 3d molecule generation by denoising voxel grids. Ad- vances in Neural Information Processing Systems, 36: 69077-69097, 2023.
Jaao, R., Huang, W., Lin, P., Han, J., Chen, P., Lu, Y., and Liu, Y. Crystal structure prediction by joint equivari- ant diffusion. In seirth-seventh Conference on Neural Information Processing Systems, 2023.
Ong, S. P., Richards, W. D., Jain, A., Hautier, G., Kocher, M., Cholia, S., Gunter, D.,, Chevrier, L.., Persson, K. A., and Ceder, G. Python materials genomics (pymatgen): robust, open-source python library for materials analysis. Computational Materials Science, 2013.
Joshi, C. Transformers are graph neural networks. The Gradient, 2020.
Peebles, W. and Xie, S. Scalable diffusion models with transformers. In International Conference on Computer Vision, 2023.
Joshi, C. K., Bodnar, C., Mathis, S. V., Cohen, T., and Lio, P. On the expressive power of geometric graph neural networks. In International conference on machine learning, 2023.
Riebesell, J., Goodall, R. E., Jain, A., Benner, P., Persson, K. A., and Lee, A. A. Matbench discovery-an evalu- ation framework for machine learning crystal stability prediction. arXiv preprint S.XXv::338..19222 2023.
Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Židek, A., Potapenko, A., et al. Highly accurate protein structure prediction with alphafold. Nature, 2021.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In CVPR, 2022.
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. In International Conference on Learning Repre- sentations, 2014.
Rosen, A. S., Iyer, S. M., Ray, D., Yao, 7.,, Aspuru-Guzik, A., Gagliardi, L., Notestein, J. M., and Snurr, R. Q. Ma- chine learning the quantum- -hemmicl properties of metal- organic frameworks for accelerated materials discovery. Matter, 2021.
Le. T., Cremer, J., Noc, F., Clevert, D.-A., and Schütt, K. Navigating the design space of equivariant diffusion- based generative models for de novo 3d molecule genera- tion. In The Twelfth International Conference Learn- ing Representations, 2024.
Satorras, V. G., Hoogeboom, E., and Welling, M. E (n) equivariant graph neural networks. In International con- ference on machine learning. PMLR, 2021.
Liao, Y-L.,, Wood, B. M., Das, A., and Smidt, T. Equiformerv2: Improved equivariant transformer for scal- ing to higher-degree representations. In International Conference on Learning Representations, 2024.
Schneuing, A., Harris, C., Du, Y., Didi, K., Jamasb, A.,
Igashov, I., et al. maucture-based drug design with equiv- ariant diffusion models. Nature Computational Science, 2024.

---

All-atom Diffusion Transformers
Shoghi, N., Kolluru, A., Kitchin, J. R., Ulissi, Z. W., Zitnick, C. L., and Wood, B. M. From molecules to materials: Pre- training large generalizable models for atomic property prediction. In ICLR, 2024.
Xu, M., Powers, A. S., Dror, R. O., Ermon, S., and Leskovec, J. Geometric latent diffusion models for 3d molecule gen- eration. In International Conference Machine Learn- ing. 2023.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. In International conference on machine learning, 2015.
Yang. S., Cho, K., Merchant, A., Abbeel, P.,, Schuurmans, D., Mordatch, I., and Cubuk, E. D. Scalable diffusion for materials generation. In The Twelfth International Conference on Learning Representations, 2024.
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 2019.
Yim, J., Campbell, A., Foong, A. Y., Gastegger, M.. Jiménez-Luna, J., Lewis, S., Satorras, V. G., Veeling, B. S., Barzilay, R., Jaakkola, T., al. Fast protein back- bone generation with se (3) flow matching. arXiv preprint arXiv:2310.0529) 2023a.
Sriram, A., Miller, B. K., Chen, R. T. Q., and Wood, B. M. Flowllm: Flow matching for material generation with large language models base distributions. NeurIPS, 2024.
Yim, J., Trippe, B. L., De Bortoli, V., Mathieu, E., Doucet, A., Barzilay, R., and Jaakkola, T. Se (3) diffusion model with application to protein backbone generation. In In- ternational Conference on Machine Learning, pp. 40001- 40039. PMLR, 2023b.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine- tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
Zeni, C., Pinsler, R., Züggerr, D., Fowler, A., Horton, M., Fu, X., et al. Mattergen: generative model for inorganic materials design. Nature, 2025.
Vahdat, A., Kreis, K., and Kautz, Score-based generative modeling in latent space. Advances in neural information processing systems, 2021.
Zhang, L., Rao, A., and Agrawala, M. Adding conditional control to text-to-image diffusion models. In Proceedings of the IEEE/CVF International Conference Computer Vision, 2023.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. At- tention all you need. Advances in neural information processing systems, 2017.
Vignac, C., Osman, N., Toni, L., and Frossard, P. Midi: Mixed graph and 3d denoising diffusion for molecule generation. In ECML PKDD, 2023.
Wang, Y., Elhag, A. A., Jaitly, N., Susskind, J. M., and Bautista, M. A. Swallowing the bitter pill: Simplified scalable conformer generation. In International confer- ence on machine learning, 2024.
Watson, J. L., Juergens, D., Bennett, N. R., Trippe, B. L., Yim, J., Eisenach, H. E., Ahern, W., Borst, A. J., Ragotte, R. J., Milles, L. F., et al. De novo design of protein structure and function with rfdiffusion. Nature, 2023.
Wu, Z., Ramsundar, B., Feinberg, E. N., Gomes, J., Ge- niesse, C., Pappu, A. S., Leswing, K., and Pande, V. Moleculenet: benchmark for molecular machine learn- ing. Chemical science, 2018
Xie, T., Fu, X., Ganea, O.-E., Barzilay, R., and Jaakkola, T. S. Crystal diffusion variational autoencoder for peri- odic material generation. In International Conference on Learning Representations, 2022.
12

---

All-atom Diffusion Transformers
A. Related Work
Generative models for molecules and materials. Diffusion models have emerged as the state-of-the-art for generative modelling of atomic systems, with applications to molecules, crystals, and biomolecules. For molecule generation, Equivariant Diffusion (Hoogeboom et al., 2022) pioneered roto-translationally equivariant diffusion on the multi-modal product manifold of atom types and 3D positions, while GeoLDM (Xu et al., 2023) introduced latent diffusion in the space of an equivariant autoencoder. Schneuing et al. (2024) extended equivariant diffusion to generate molecules conditioned on binding protein partners for structure- based drug design, while Corso et al. (2023) explored similar architectures for protein-small molecule docking. For crystal generation, state-of-the-art approaches use equivariant diffusion on product manifolds of atom types, 3D/fractional coordinates, and lattice parameters. Notable examples include CDVAE (Xie et al., 2022), DifICSP (Jiao et al., 2023), and FM (Miller et al., 2024). MatterGen (Zeni et al., 2025) demonstrated conditional diffusion for inverse design based on target material properties and symmetry space groups. Language models have also been used for generating molecules and crystals as textual representations (Flam- -lem-hhe & Aspuru- Guzik, 2023; Gruver et al., 2024).
Our work stands out as the first to develop unified generative models capable of sampling both periodic crystals and non- periodic molecular systems jointly. The closest work to ADiT in terms of diffusion formulation is AlphaFold3 (Abramson et al., 2024), which applies standard Transformers and Gaussian diffusion to generate all-atom biomolccular complex. However, their formulation is specific to structure prediction for biomolecules and only diffuses 3D atomic coordinates in Cartesian space. In contrast, our latent diffusion formulation is sufficiently general work with both periodic and non- periodic systems, generating atom types, coordinates, as well as unit cell parameters unconditionally with classifier-fi free guidance. Our emphasis on joint representations of molecules and crystals also aligns with recent work foundation models for molecular dynamics (Shoghi et al., 2024; Batatia et al., 2023). Similarly, our unified latent diffusion framework can potentially be scaled up with larger and more diverse chemical datasets towards foundation models for generative chemistry.
Latent diffusion models. Latent diffusion models (Vahdat et al., 2021; Rombach et al., 2022) propose to do diffusion in the latent space of an autoencoder instead of the raw input space of high-dimensional continuous signals such as pixels, and have been extremely successful for generating images, audio, and videos (Esser et al., 2024; Betker et al., 2023; Brooks et al., 2024). Latent diffusion is a more computationally efficient alternative to standard diffusion as the autoencoder's latent space captures semantically meaningful features of the data, allowing for more efficient diffusion in a lower- dimensiinni space followed by reconstruction to the original data space. The original formulation was further improved by Diffusion Transformers (DiTs) (Peebles & Xie, 2023), which demonstrated that standard Transformers provide a highly scalable architecture for the denoiser network. Latent diffusion models can easily incorporate conditioning on additional information like class labels, text prompts, or infilling masks through classifier-based (Dhariwal & Nichol, 2021) and classifier-free guidance (Ho& & Salimans, 2022) as well as finetuning (Zhang et al., 2023; Dai et al., 2023).
Our work is the first to leverage latent diffusion for jointly generating the complex multi-modal product of categorical and continuous data types that constitute 3D atomic systems. This allows us to shift the complexity of handling atom types, coordinates, and unit cell parameters into an autoencoder while performing the generative process in latent space with DiTs, which is simpler and more scalable than alternative multi-modal equivariant diffusion models.
Equivariance and generative modelling. Geometric Graph Neural Networks (Duval et al., 2023), particularly roto- translationally equivariant networks, have been used denoisers in diffusion and flow matching approaches for generative modeling of 3D atomic systems. E(3) Eq3ivvauiaa Graph ConvNets (Satorras et al., 2021) are widely used as denoisers for molecule (Hoogeboom et al., 2022; Xu et al., 2023; Schneuing et al., 2024) and crystal generation (Jiao et al., 2023; Miller et al., 2024). More expressive architectures, like higher-order tensor networks (Liao et al., 2024) and Invariant Point Attention (Jumper et al., 2021). have been applicd to protein structure generation (Watson et al., 2023; Yim et al., 2023b) and protein-ligand docking (Corso et al., 2023).
However, equivariant networks are computationally expensive and harder to scale than standard Transformers in terms of data and model size. This is especially relevant for diffusion models, where denoisers typically process inputs as fully connected graphs to capture global structure (Joshi, 2020) and are iteratively run hundreds of times during inference. Recent work has challenges the necessity of 3D inductive biases and equivariance for generative structure prediction tasks, showing that standard Transformers can achieve strong performance on biomolecular complexes (Abramson et al., 2024) and small molecule conformations (Wang et al., 2024; O Pinheiro et al., 2023). Non-equivariant models have also shown promising results for protein structure generation (Chu et al., 2024; Martinkus et al., 2024). In the same vein, our work leverages
13

---

All-atom Diffusion Transformers
the simplicity and scalability of standard Transformers for generative modelling across both periodic and non-periodie 3D atomic systems, demonstrating that explicit equivariance and molecular inductive biases not a strict requirement for generating valid and realistic atomic structures at scale.
B. Evaluation Metrics
Crystal generation metrics. We follow the evaluation protocol established by Xie et al. (2022); Miller et al. (2024), where we sample 10,000 crystals and compute validity, stability, uniqueness, and novelty rates, defined as follows:
Structural validity: % of crystals with all pairwise distances >= 0.5 and crystal volume >= 0.1.
Compositional validity: % of crystal compositions with charge neutrality and electronegativity balance according to SMACT (Davies et al., 2019).
Overall validity: % of crystals which are both structurally and compositionally valid.
Stability: % of crystals with DFT energy above hull <0.0 eV/atom and no. of unique elements >= 2. (We also report
metastability as DFT energy above hull <0.1 eV/atom and no. of unique elements 2.)
Stable & unique: % of stable crystals which are unique, as defined by an all-to-all comparison using Structure Matcher²
from PyMatGen (Ong al., 2013).
Stable, unique & novel: % of stable, unique crystals which novel, as defined by an all-to-all comparison to all erystals in MP-20 using Structure Matcher.
To compute the stability, uniqueness, and novelty rates, we follow Miller al. (2024); Sriram et al. (2024): We first pre-relax the sampled crystals using fast ML potential, CHGnet (Deng et al., 2023), and then perform DFT relaxation. We then determine the DFT energy above hull for the relaxed structures against the Matbench Discovery convex hull (Riebesell et al., 2023). Note that there is a lower bound on the number of completed DFT calculations due to memory or timeout errors.
Molecule generation metries. We follow the evaluation protocol established by Hoogeboom et al. (2022); Daigavane et al. (2024), where we sample 10,000 molecules and compute validity and uniqueness rates as well success rates for 7 sanity checks from Posebusters (Buttenschoen al., 2024), as follows:
Validity: % of molecules with canonical SMILES string found by RDKit.
Uniqueness: % of unique SMILES among valid ones.
All-atoms connected: % of molecules where there exists path along bonds between all atoms.
Reasonable bond angles/lengths: % of molecules where all angles/lengths are within 0.75 of the lower and 1.25 of the
upper bounds determined by distance geometry.
Aromatic rings flatness: % of molecules where All-atoms in aromatic rings with or members are within 0.25A of the
closest shared plane molecule.
Double bond flatness: % of molecules where All-atoms of aliphatic carbon-carbon double bonds and their four neighbours are within 0.25 of the closest shared plane.
Reasonable internal energy: % of molecules where the calculated energy is no more than 100 times the average energy of an ensemble of 50 conformations generated for the input molecule.
No internal steric clash: % of molecules where the interatomic distance between pairs of non-covalently bound atoms is above 0.8 of the distance geometry lower bound.
The validity and uniqueness metrics focus on whether the chemical composition of generated molecules can be processed by RDKit, while the Posebusters sanity checks evaluate the physical realism of the generated 3D structures across multiple criteria, from geometric constraints like bond lengths to energetic considerations (Harris et al., 2023).
C. Additional Results
Extension to MOF generation. Having established that ADiT's benefit from transfer learning to generate high-quality crystals and molecules, we challenged our architecture to generate metal-orgame frameworks (MOFs), which represent a more complex class of hybrid materials with metal nodes connected by organic small molecule linkers. We trained ADIT' on an additional 14,000 MOFs of up 150 atoms from the QMOF database (Rosen et al., 2021) alongside QM9 and MP20.
Structure Matcher checks if two periodic structures are equivalent, even they are in different settings or have minor distortions.
14

---

All-atom Diffusion Transformers
using the same experimental settings and training for 10,000 epochs. We sampled 1,000 MOF's and evaluated their validity using 15 sanity checks from MOFChecker (Jablonka, 2023), including tests for: presence of metal/carhon/hyddogenn atoms, atomic overlaps, over/undervalent carbons and nitrogens, missing hydrogens, and excessive partial charges.
Table 4 shows that QMOF-only trained ADiT achieves 15% overall validity rate for MOF generation, which decreases to 10% with joint training on molecules, erystals, and MOFs simultaneously. However, the jointly trained model takes significant time to train and did not fully converge, suggesting that further improvements in MOF generation may be possible with larger models trained for longer. Comparing validity rates for joint vs. -ataset-specific ADiTs shows that the joint model benefits from transfer learning across molecules, crystals, and MOFs, achieving high validity rates earlier in training. Notably, the joint model achieves high validity rates of 91 % for crystals and 95% for molecules, matching best dataset-specific models and being able to additionally generate MOFs. While not directly comparable, models specialized for MOF such as MOFDiff (Fu et al., 2024) are trained on large synthetic dataset of 300,000 MOFs and achieve 30% validity rates based on MOFChecker after DFT relaxation of generated MOFs. Our results are reported without DFT relaxation.
Table 4: Metal-organic framework generation results. Left: We report sanity checks from MOFChecker for 1,000 sampled MOFs, trained on QMOF-only as well as jointly with QM9 and MP20. (+/4 indicate higher/lower better, respectively.) Right: Jointly trained and datasct-specific ADIT validity rates vs. epochs. The joint model benefits from transfer learning and requires fewer epochs per dataset to start generating valid samples.
Test (%)
QMOF ADIT
Joint ADiT
Has carbon f Has hydrogen f
100.0
100,0
100
99.6
100.0
Has atomic overlap & Has overcoord. C Has overcoord. J Has overcoord. H Has undercoord. C L Has undercoord. N
8.3
10.8
80-
23.6
34.3
Performance metric (%)
1.5
1.6
1.0
3.6
60.0
72.1
39.1
39.9
Has undercoord. rare carth Has metal f Has lone molecule +
0.4
0.8
Metric
QMOF steric clash rate
100.0
99.4
MP20 validity rate
72.9
83.2
20
QMB validity rate Model
Has high charge $ Has suspicious terminal oxo
0,9
2.5
Joint
2.6
5.8
Individual
Has undercoord. alkali I Has geom, exposed metal 1.
1.0
6.4
0
250
500
750
1000
1250
1500
1750
2000
7.0
9.6
Epoch
Validity rate (all passed) f
15.7
10.2
Histograms from DFT validation. In Figure 4, we show histograms of DFT energy above hull, formation energy, and number of unique elements per crystal for 10,000 generated crystals from ADiT, FlowMM, and FlowLLM compared to the MP20 training distribution. ADiT generates more thermodynamically stable crystals than prior models, as shown by the larger proportion of samples with DFT energy above hull below 0.0 e/atom. The distribution of DFT formation energies and number of unique elements per crystal from ADIT samples more closely matches the MP20 training data compared FlowMM and LLowLLM baselines, suggesting that ADIT better captures the underlying physical and chemical constraints of stable crystal structures. Note that we ran DFT calculations for all model samples under identical hardware and settings to ensure fair comparison.
Histogram of spacegroups. In Figure 5, we show the distribution of spacegroups for 10,000 generated crystals from ADIT, FlowMM, FloWLLM and the MP20 distribution. Diffusion-basel models (ADiT and FlowMM) tend to over sample crystals with P1 spacegroup, which represents the lowest symmetry group, likely due to their local, step-wise denoising process. In contrast, FlMw an autoregressive language model, tends to over sample spacegroups like Fm-3m, Pm-3m, and 14/mmm compared to the training data. While it would be straightforward to control the distribution of spacegroups generated by ADIT through classifier-free guidance conditioning, we leave this for future work since our current focus is on unconditional generation of diverse atomic systems.
SUN rate and scaling ADIT. In Table 5, we observe that the combined stability, uniqueness, and novelty (S.U.N.) rate for crystal generation decreases as we scale up the DiT denoiser from DIT-S (32M) to DÝT-I. (450M). While stability and
15

---

All-atom Diffusion Transformers
Moed FowwM
Adm
ASTT
1902-
FoodLM
FloeLux
Flovere
Count
2020
2020
-12
-0.3
0.0
0.4
0.5
2/6 (AV alom))
Formation energy Jek atom)
Namber of unique wlements per cryatal
(a) DFT energy above hull
(b) DFT formation energy
(c) Number of elements
Figure 4: Histograms from DFT validation of 10,000 generated crystals. ADiT is more likely to generate stable crystals with DFT energy above hull <0.0 eV/atom compared to prior models. Samples from ADIT most closely follow the distributions for DFT formation energy and number of unique elements per crystal from MP20.
4000
Modil
ADIT
25:00
FloXLLM
FlowMMM
3000
M9720 toat 545
2010
1500-
1000
P2.AR
Nymo
44/mm
8
10
180
International spaccegroup
Figure 5: Histogram of spacegroups for 10,000 generated crystals. Diffusion- ADIT and FlowMM tend to over sample crystals with P1 spacegroup compared to the MP20 training distribution. FlowLLM,, an autoregressive language, tends to over sample crystals with Fm-3m, Pm-3m, and I4/mmm spacegroups
uniqueness rates increase with model size, the S.U.N. rate decreases due to the larger model's greater capacity to memorize the small MP20 training dataset of 27K crystals. This suggests that larger models may be more prone to generating duplicate or depl-duplicate samples, which we plan to address by training on larger and more diverse datasets future work. For crystals, the Alexandria dataset of inorganic crystals and the Crystallography Open Database of organic crystals present promising opportunities for scaling up. Notably, ADiT-S trained on MP20-only achieves S.U.N. rate of 6.5%, representing significant improvement over previously published results from FlowMM (2.8%) and FlowLLM (4.7%). This demonstrates that even our smallest model variant substantially advances the state-of-ther the-tth-rrtt art for crystal generation.
Sensitivity of validity rate to number of samples and random seed. In Figure 6a, we plot the validity rates for crystal and molecule generation as we increase the number samples from 100 to 10,000 for 3 different random seeds. We observe that the validity rates generally converge and are stable across random seeds after sampling over 5,000 crystals or molecules.
Sensitivity of S.U.N. rate to number of samples. In Figure 6b, we plot the S.U.N. (stability, uniqueness, and novelty) rates for crystal generation as we increase the number of samples from 100 to 10,000 across 3 different random seeds, The S.U.N. rates converge after approximately 5,000 samples for diffusion- methods like ADiT and FlowMM. In contrast, autoregressive models like FlowLLM show higher variance in S.U.N. rates, likely due to more frequent generation of duplicate crystals during low-temperature sampling.
16

---

All-atom Diffusion Transformers
Table 5: Impact of scaling stability, uniqueness, and novelty rates for 10,000 generated crystals. We find that stability rate as well as stability & uniqueness rate increase as we increase the number of model parameters for ADiT from 32M to 450M. However, larger ADiT models have greater capacity to memorise the small MP20 training dataset of 27K crystals, resulting in decrease in the combined stability, uniqueness, novelty rate. ADiT-S trained on MP20-only achieves S.U.N. rate of 6.5%, representing a significant improvement over previously published state-of-the-ari models which attained S.U.N. rates up 4.7%.
Stability (Ebull <0.0)
Metatability (Ehull <0.1) M.S.U.N. (%) f
Model
(%)
S.U. (%) f
S.U.N. (%) f
M.S (%)
M.S.U. (%)
MP20-only ADiT-S (32M)
12.8
11.8
6.5
71.1
64.9
38.1
MP20-only ADiT-B (130M)
14.1
12.5
4.7
81.6
67.3
25.9
Joint ADIT-S (32M)
12.6
11.4
6.0
71.9
64.7
37.7
Joint ADiT-B (130M)
15.4
13.4
5.3
81.0
70.2
28.2
Joint ADīT-L (450M)
15.5
13.5
5.0
82.5
70.9
27.9
System Crystals
Model
ADIT-L
12
98
Molecules
ADIT-S
10
FlowLLM
Validity rate (96)
96
FlowMM
94
Stable. unique,
92
90
2
88
100
1000
2500
5000
10200
100
1000
2500
5000
7500
10000
Number of samples generated (ADIT)
Number of crystals sampled
(a) Validity rates are consistent across random seeds.
(b) S.U.N. rates converge after 5,000 samples.
Figure 6: Consistency of validity and S.U.N. rates as we increase number samples. plot the validity and S.U.N. rates vs. number of sampled crystals molecules. Error bars indicate 95% confidence interval across three different random seeds. Metrics are generally stable across seeds and converge after sampling over 5,000 crystals or molecules.
17

---

All-atom Diffusion Transformers
D. Ablation Study
Table 7 and Table 6 presents ablation studies as well as aggregated benchmarks for various configurations of ADiT's latent diffusion model and autoencoder, respectively. Key takeaways are highlighted below. Note that, unless otherwise stated, results in the main paper are reported for jointly trained ADiT-B which uses DiT-B denoiser, standard Transformer encoder and decoder, latent dimension d 8, and KL regularization weight λKI. = le 5.
Joint vs. dataset-specific training Joint training of the autoencoder to embed both molecules and crystals into shared latent space achieves similar better reconstruction performance compared to dataset-specific training, as shown in Table 6 (rows 3, 6, 10). The benefits of joint training are most evident in generative modelling performance samples from the joint model have higher validity rates for both crystals and molecules compared to dataset-specific models, demonstrating effective transfer learning between periodic and non-perion non-periodic atomic systems (Table 7, rows 12, 16, 20). These results provide strong evidence that ADÎT' can successfully unify the modelling of both periodic and non- non-periodic atomic systems within a single architecture, without compromising performance on either domain.
Denoiser architecture The DiT denoiser is a standard Transformer with key hyperparameters including the hidden dimension dynotee, number of attention heads, and number of layers. Scaling up the DiT denoiser from DìT-$ (32M parameters, dmodel- 384, 6 heads, 12 layers) to DiT-B (150M, duodell = 768, 12 heads, 12 layers) and DiT-L (450M, dmodel = 1024, 24 heads, 24 layers) consistently improves generative performance, as shown in Table (rows 12, 16, 20). We have additionally performed scaling law analysis for the training loss and validity rates in Figure 3, seeing strong correlations between model size and performance metrics. In Figure 6b, we further see that S.U.N rates for larger models are better than smaller models, further confirming the benefits of scaling up the DiT denoiser.
Autoencoder architecture For the architecture of the autoencoder's 's encoder and decoder, we explored both roto-translation cquivariant as well as non-equivariant VAEs. For the equivariant VAE variant, the encoder is Equiformer-V2 (Liao et al., 2024) and the decoder is an equivariant feedforward network adapted from output heads in the Equiformer- V2 codebase. We selected Equiformer- V2 as is theoretically expressive (Joshi et al., 2023) and has state-of-the- performance across diverse 3D atomic systems, As input to the Equiformer-V2 encoder, we use spherical harmonic embeddings of displacement vectors edge features and exclude the 3D coordinates in Algorithm 1, line 2, from the initial features {ha) as result. The initial features } are used the L = 0 scalar component of the initial spherical tensor features of Equiformer- V2. The rest of the pseudocode in Algorithms 1 and 2 remains the same.
As shown in Table (rows 1-4 and 5-8), the choice of autoencoder architecture has noticeable impact on reconstruction performance. Standard Transformers generally outperform Equiformer- V2 for both crystals and molecules, achieving higher match rates (% of test set samples where the reconstrueted structure matches the groundtruth, as determined by PyMatGen's 'ss More importantly, the latent space learned by standard Transformers proved more suitable for the latent diffusion process compared to Equiformer- V2's equivariant latent space, leading substantially better generative performance in terms of validity rates, particularly for crystals (Table 7, rows 1-4 and 5-8).
Autoencoder regularization As shown in Table 6 (rows 9-12), increasing the latent dimension and reducing the KI. regularization weight generally improved autoencoder reconstruction performance by lowering RMSD values which measure the average distance between the reconstructed and groundtruth structures. These improvements in reconstruction quality translated to better generative performance, with higher validity rates for both crystals and molecules at larger latent dimensions and lower KL weights (see Table 7, rows 9-12).
Sampling hyperparameters. Classifier-free guidance scale and number of integration steps are important hyperparameters for inference-time tuning. In Figure 7, we show a grid search over guidance scales 7 {1.0, 2.0, 3.0, 4.0, 6.0) and integration steps T E {10, 50, 100, 250, 500, 1000), finding that different combinations may be optimal for crystals vs. molecule generation. For each entry in Table 7, we have reported results for T and y which obtain the highest validity rates. T = 500 or 1000 with y = 1.0 or 2.0 tends to work well across both molecules and crystals.
18

---

All-atom Diffusion Transformers
Table 6: Autoencoder ablation study. We report match rate (computed with Structure-Matcher MoleculeMatcher from PyMatGen) and RMSD between the reconstructed and groundtruth structures for MP20 crystals and QM9 molecules.
Train
Autoencoder
hyperparameters
Crystals MP20
Molecules QM9
Set
Encoder
Latent
KL
Match Rate (%)
RMSD (Å) 4
Match Rate (%) f RMSD (Á) 4
MP20
Transformer
4
0.0001
85,50
0.0598
MP20
Equiformer-V V2
4
0.0001
81.70
0.1652
MP20
Transformer
8
0.0001
84.50
0,0502
MP20
Equiformer- V2
0.0001 8
88.90
0.0296
QM9
Transformer
0.0001
97.20
0.0747
QM9
Equiformer- V2
0.0001 4
96.20
0.0765
QM9
Transformer
0.0001 8
96.50
0.0823
QM9
Equiformer-V2
0.0001 8
96.20
0.0746
Joint
Transformer
0.0001
88.30
0.0471
96.60
0.0785
Joint
Transformer
0,00001
88.50
0.0468
98.50
0.0524
Joint
Transformer
8
0.0001
88.60
0.0269
96.60
0.0760
Joint
Transformer
8
0.00001
88.60
0.0239
97.00
0.0399
Table Latent diffusion model ablation study. We report validity rates for 10,000 generated crystals or molecules.
Autoencoder hyperparameters
Crystals - MP20
Molecules - QM9
Train Set
Diffusion Denoiser
Structure Valid (%)
Composition Valid (%) t
Overall Valid (%)
Validity
Validity
Encoder
Latent
KL
(%)
(%)
MP20
DIT-S
Transformer
4
0.0001
98.90
89.19
88.19
MP20
DiT-S
Equiformer- V2
0.0001
91.74
81.03
74.43
MP20
DIT-S
Transformer
0.0001
99.58
90.46
90.13
MP20
DIT-S
Equiformer-V V2
8
0.0001
99.26
86.09
85.50
QM9
DIT-S
Transformer
4
0.0001
95.94
92.19
QM9
DiT-S
Equiformer-V V2
0.0001
95.36
91.37
QM9
DIT-S
Transformer
0.0001
96.02
91.58
QM9
DIT-S
Equiformer- V2
0.0001
96.24
91.47
Joint
DiT-S
Transformer
4
0.0001
98.21
91.05
89.38
96.90
93.47
Joint
DIT-S
Transformer
4
0.00001
98.74
90.74
89.60
96.40
91.85
Joint
DiT-S
Transformer
0.0001
99.66
91.07
90.76
96.85
93.33
Joint
DiT-S
Transformer
0.00001
99.67
91.25
90.93
96.36
92.06
Joint
DiT-B
Transformer
0.0001
99.00
91.23
90.29
97.33
94.45
Joint
DiT-B
Transformer
4
0.00001
99.51
90.73
90.29
97.04
94.06
Joint
DiT-B
Transformer
0.0001
99.67
91.60
91.32
95.30
89.85
Joint
DiT-B
Transformer
0.00001
99.74
92.14
91.92
97.43
93.99
Joint
DIT-L
Transformer
0.0001
99.31
90.92
90.29
97.80
94.67
Joint
DIT-L
Transformer
4
0.00001
99.43
90.84
90.31
96.71
92.78
Joint
DiT-L
Transformer
8
0.0001
99.75
92.17
91.92
96.11
91.45
Joint
DIT-L
Transformer
0.00001
99.66
91.42
91.14
97.79
95.01
19

---

All-atom Diffusion Transformers
WAY
1,00
0,03
1,000
DOLS
33
302
522
3.5
000
CAN
cas
(a) Crystals MP20
(b) Molecules QM9
Figure 7: Tuning inference hyperparameters for best performance. Best generative modelling results for crystals and molecules are achieved with different classifier-free guidance scales and number integration steps T. = 500 or 1000 with y = 1.0 or tends to work well across both molecules and crystals.
E. Visualizations
PCA visualization of shared latent space. In Figure 8, we plot the first two PCA principal components of 100 random samples each from the MP20 and QM9 validation set, as well 100 generated crystals and 100 generated molecules sampled from ADiT. We observe that the joint latent space shows distinct clusters between molecules and crystals, with tighter clustering for molecules and more spread for crystals, reflecting the greater diversity of elements and local geometric environments in periodic crystal structures.
Next, we plot the same PCA but only keeping atoms of carbon, nitrogen, oxygen, and fluorine in Figure 9. These atoms appear in both QM9 molecules and MP20 crystals, allowing us to analyze how their representations compare across periodic and non-periodic systems. The visualization reveals clear patterns: principal component I primarily distinguishes between molecules (clustered between and and crystals, while principal component correlates with atom type. Most notably, oxygen atoms show similar latent representations whether they appear in molecules crystals, suggesting ADIT's latent space captures fundamental chemical properties that transfer across both domains. This shared representation of oxygen, a key element in both datasets, may help explain ADiT's successful joint learning and transfer between periodic and non non-periodic systems.
Generated crystals, molecules, and MOFs. In Figure 10, Figure 11, and Figure 12, we show samples generated crystals, molecules, and MOFs from ADiT, respectively. The generated crystals exhibit diverse spacegroups and compositions, while the generated molecules show a wide range chemical structures and conformations. These visualizations demonstrate that the jointly trained ADIT model successfully generates high-quality and chemically diverse atomic systems in both periodic and non-periodic domains.
20

---

All-atom Diffusion Transformers
6
System
Crystal
Molecule
Source
Dataset
Generated
2
Principle component
-6
-6
-2
2
4
6
Principle component
Figure 8: PCA plot of latent embeddings from ADIT's VAE for 100 data points from the MP20 and QM9 datasets, as well as 100 ADiT-generated crystals/molecules each. Fach point represents an atom, coloured by the system type and sized by whether comes from real data or generated latents. The joint latent space shows distinct clusters between molecules and crystals, with tighter clustering for molecules and more spread for crystals, reflecting the greater diversity of elements and local geometric environments in periodic erystal structures.
21

---

All-atom Diffusion Transformers
Atom type
C
N
O
Source
Dataset
Generated
System
Crystal
x Molecule
Principle component 2
0
22
-6
-6
-2
2
4
6
Principle component
Figure 9: PCA plot of latent embeddings for carbon, nitrogen, oxygen, and fluorine atoms from ADiT's VAE for 100 data points from the MP20 and QM9 datasets, as well as 100 ADiT-generated crystals/molecules each. Each point represents an atom, coloured by atom type and sized by whether it comes from real data or generated latents. Principle component visually correlates with whether a system is molecule (within range -2 - 2) or crystal. Principle component visually correlates with the atom type. The joint latent space shows distinet clusters for different atom types, with oxygen atoms having similar representations in both molecules and crystals. This overlap in oxygen atom representations suggests that ADiT's latent space captures shared chemical properties across periodic and non-periodic systems, enabling effective knowledge transfer during joint training.
22

---

All-atom Diffusion Transformers
4
1a Cla (PI)
Ce2 N₂ S (P-3ml)
Er, Brg (P63/mme)
Ba
Se
Pa
Fm
Bag Pa Fe Og (F3-3m)
K Tm Se (R-3m)
Pm Asa (Cm)
9
Bh
Pm Pba (Pm-3m)
Rb Nb F. (PI)
Sig Rh4 Zea (Pnma)
To
Ter As; 1 (P1)
Tm Fa (Imm2)
Yb, W₂ Og (P-1)
Figure 10: Generated crystals from ADIT trained jointly on MP20 crystals and QM9 molecules.
23

---

All-atom Diffusion Transformers
Figure 11: Generated molecules from ADiT trained jointly MP20 crystals and QM9 molecules.
24

---

All-atom Diffusion Transformers
Mn
Figure 12: Generated metal-organic frameworks from ADIT trained jointly QMOF150 metal-organic frameworks, MP20 crystals and QM9 molecules.
25


Reviews:
[{"review_id": "xYGcraaFH8", "reviewer": "Reviewer_WqZ3", "summary": "The authors note that current generative models for atomic systems - such as molecules and crystals - are are fragmented and overly specialized to the specific type of system they model. They propose all atom diffusion transformers (ADiT), which uses a two-step latent diffusion framework in which, first, mixed categorical and numerical data describing an atomic system are embedded into a shared latent space by training a VAE, and next, a diffusion transformer model is trained to model the latent distribution and generate new samples. Experimental results show that the proposed method effectively transfers knowledge between various types of atomic systems, outperforms baselines in various metrics, and scales in performance with increasing model capacity.\n\n#### Update after rebuttal period\nMy review score remains the same.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive comments and excellent summary of the work! Thanks for appreciating that the latent diffusion idea can be further applicable to multi-modality data in other scientific domains, too.\n\n> does not necessarily indicate that the model has learned the underlying physics \n\nADiT obtains very strong results on both MP20 and QM9, exceeding the current state-of-the-art. These systems are small but they are complex and challenging. **Doing well on them requires models to learn something about the underlying physics of atomic interactions,** and cannot be solved by simple co-occurence/statistical approaches.\n\nLet’s take MP20 as an example:\n- MP20 is restricted to up to 20 atoms in a unit cell because this is practically most relevant and largely covers the Materials Project distribution of known materials. Each atom's local environment is highly complex due to infinite periodic tiling of a unit cell.\n- The S.U.N. rate metric for crystals/MP20 is a very challenging and practically relevant metric for materials discovery. ADiT has **improved S.U.N. rate from 4-5% of previous works up to 6.5%** (see Table 4). ADiT significantly outperformed the recently published MatterGen (Nature, 2024). Previous works showing smaller improvements have been published at top conferences (most recent example: FlowLLM at NeurIPS, 2024). There is a large community of researchers who care about S.U.N. rate on MP20.\n\n> or that it can generalize well to very large-scale data\n\nWe agree, but the goal of this paper was to introduce the **first unified architecture** for generative modelling of 3D atomic systems, and to demonstrate **transfer learning** across periodic crystals and non-periodic molecular systems. We believe this will make a novel and valuable contribution to the ICML community.\n\nBelow, we presented results for the larger GEOM dataset of small molecules with up to 180 atoms. ADiT continues to show strong results, generating physically realistic molecules and outperforming baselines.\n\nAlso note that our best models are based on **pure Transformers** - which have shown to be extremely scalable to large-scale inputs - so there is nothing inherent to our methodology that would prevent scaling up to larger systems or more samples.\n\n---\n\n# **TO ALL REVIEWERS - PLEASE READ BELOW - RESULTS ON GEOM**\n\n---\n\nTo demonstrate the scalability of ADiT to larger systems, we have run experiments on GEOM, as suggested by Reviewers KRy3 and 7DhL. GEOM includes 430,000 unique molecules and consists of larger systems than QM9, up to 180 atoms. Our experiments and evaluation follows the EDM/MiDi paper, with additional PoseBusters physics-based tests. Model setup and hyperparameters used are exactly as described in the paper.\n\n| Metric  | EDM | MiDi | ADiT-S\n| --- | --- | --- | --- |\n| Validity | 87.8 | 77.8 | 93.4\n| Uniqueness | 99.9 | 100.0 | 100.0\n| Atoms connected | 41.4 | 90.0 | 94.9\n| Bond angles | 91.8 | - | 96.2\n| Bond lengths | 90.2 | - | 96.8\n| Ring flat | 99.0 | - | 100.0\n| Double bond flat | 98.2 | - | 99.9\n| Internal energy | 89.2 | - | 94.2\n| No steric clash | 85.2 | - | 93.5\n\nADiT-S (32M) outperforms or matches EDM and MiDi across all metrics. ADiT generates physically realistic molecules based on all PoseBusters criteria. Notably, molecules generated by ADiT are significantly more likely than EDM/MiDi to be connected as measured by the 'Atoms connected' score (measures % of molecules where there exists a path along bonds between any two atoms).\n\nWe have not trained the larger ADiT-B and ADiT-L models due to resource constraints in the short rebuttal period. We expect larger models to further improve performance based on the scaling trends we have seen on QM9 and MP20.", "claims_and_evidence": "The claims are supported.", "methods_and_evaluation_criteria": "The methods are valid.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The experiments appear well-designed, and the analysis sound.", "relation_to_broader_scientific_literature": "Generative models for atomic systems have thus far been very specialized for the specific application areas they model. This paper represents a unification of atomic system models based on the idea that the underlying physics in these systems should hold constant and generalize between seemingly different domains.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "### Strengths\n\n- The paper is well-written, and explains just enough of the physics and chemistry background needed for someone who is unfamiliar.\n- The implementation details are very thorough, and clearly describe the steps needed to reproduce the proposed method.\n- The proposed method is a promising approach not just for atomic systems, but also for other scientific domains that may have related data from heterogenous sources with mixed categorical and numerical traits.\n- Empirical results are extensive and back up the major claims made by the authors.\n\n### Weaknesses\n\n- The authors refer to this point in their discussion, but it bears repeating that transfer of performance between the two related atomic systems that the authors mainly test does not necessarily indicate that the model has learned the underlying physics or that it can generalize well to very large-scale data.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hNhY2T5PvW", "reviewer": "Reviewer_yi4C", "summary": "The paper introduced a unified framework called All-atom Diffusion Transformer (ADiT) for generating periodic (crystals) and non-periodic (molecules) atomic systems. ADiT employs a two-stage approach: A Variational Autoencoder that maps atomic systems into a shared latent space, and a Diffusion Transformer generates decoded latent samples into valid structures. Evaluations on QM9 and MP20 benchmarks showed the effectiveness of the proposed approach.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> trained on relatively small datasets - scalability on larger datasets\n\n---\n\n**[In our response to Reviewer WqZ3](https://openreview.net/forum?id=89QPmZjIhv&noteId=RrA8d9t9eq), we presented results for the larger GEOM dataset of small molecules with up to 180 atoms. ADiT continues to show strong results, generating physically realistic molecules and outperforming baselines.**\n\n---\n\nWe agree that systems from QM9 and MP20 are small, but they are complex and challenging. Doing well on them requires models to learn something about the underlying physics of atomic interactions, and cannot be solved by simple approaches.\n\nLet’s take MP20 as an example:\n- MP20 is restricted to up to 20 atoms in a unit cell because this is practically most relevant and largely covers the Materials Project distribution of known materials. Each atom's local environment is complex due to periodicity.\n- The S.U.N. rate metric for crystals/MP20 is a very challenging and practically relevant metric for materials discovery. ADiT has **improved S.U.N. rate from 4-5% of previous works up to 6.5%** (see Table 4). ADiT significantly outperformed the recently published MatterGen (Nature, 2024). Previous works showing smaller improvements have been published at top conferences (most recent example: FlowLLM at NeurIPS, 2024). There is a large community of researchers who care about S.U.N. rate on MP20.\n\nUltimately, the goal of this paper was to introduce the **first unified architecture** for generative modelling of 3D atomic systems, and to demonstrate **transfer learning** across periodic crystals and non-periodic molecular systems. We believe this will make a novel and valuable contribution to the ICML community.\n\nAlso note that our best models are based on **pure Transformers** - which have shown to be extremely scalable to large-scale inputs - so there is nothing inherent to our methodology that would prevent scaling up to larger systems or more samples.\n\n> lack of theoretical justification for why a shared latent space works\n\nThere are strong theoretical justifications of our choices:\n\n- **The underlying physics of atomic interactions is the same across all 3D atomic systems.** Interatomic distance between a carbon atom double bonded to an oxygen atom will be the same, whether they are part of a molecule or a crystal structure. Thus, a shared/unified latent space enables the model to learn shared principles of interatomic interactions.\n- Using a shared latent space **unifies system modalities** (both periodic and non-periodic atomic systems embedded together), as well as **unifies multi-modality data** (categorical, numerical, floating point embedded together). This makes it very easy to train a simple Gaussian diffusion model on the latent representations, instead of more complex equivariant diffusion formulations.\n- The advantage of jointly embedding periodic & non-periodic systems for ML force fields, aka. *universal interatomic potentials* was shown by recent works like JMP and MACE. More broadly, all of AI research is moving towards **joint training of large models on multiple datasets**, ie. learning unified latent representations.\n\n> trade-offs of equivariance? ability to respect physical symmetries?\n\nWhen developing ADiT, we ablated the impact of enforcing equivariance on the model -- **see Appendix D**. We tried to be pragmatic about whether or not to use equivariance.\n\nThe non-equivariant version **improved generative performance** and **physical realism** of the crystals/molecules, as judged by evaluation metrics explicitly focussed on physics-based tests s.a. PoseBusters suite for molecules, and DFT-based rate for crystals.\n\nOverall, we and others have found that equivariance is not strictly necessary for training a good generative model. Perhaps it can even be an advantage: if you start with two rotated versions of a noisy Gaussian and non-equivariant denoising leads to two different (but equally valid) molecules/crystals, that gives you a more diverse generative model without sacrificing on performance!\n\n> joint training on MOFs reduces validity?\n\nSimple answer: training did not converge during the time period that we were allocated compute resources on a large enough GPU cluster for this experiment.\n\nNote that joint training lead to very strong validity for both crystals (91%) and molecules (95%), while reducing validity for MOFs. When looking at training dynamics (right side plot in Table 3), the model first learns crystals, then molecules, and finally starts learning MOFs later in training. We believe training for longer is likely to close the gap and potentially improve beyond the MOF-only variant of ADiT.\n\nUltimately, our goal with including preliminary MOF results was to encourage the community to work further on MOFs - hybrid organic-inorganic materials for carbon capture and sustainability. We will release **open source code** and **datasets** which make it easy for others to build upon our initial experiments.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "No.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "The authors built their approach on several works in the literature, like Variational Autoencoder, Diffusion Transformer, and classifier-free guidance technique.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "**Strengths**:\n\n- Unified framework: I think the idea of using a joint generative model for molecules and crystals is novel, avoiding domain-specific methods. \n- Empirical performance: The paper showed strong empirical results outperforming domain-specific methods for crystals and molecule generations.  \n\n**Weaknesses**:\n\n- Dataset limitations: The method is trained on relatively small datasets (QM9 ~ 130K molecules; MP20 ~ 45K crystals). As already mentioned by the authors, scalability on larger datasets is unexplored.\n- Theoretical gaps: The paper lacks a theoretical justification for why a shared latent space works. More discussion on the latent space properties would be helpful.", "other_comments_or_suggestions": "No.", "questions_for_authors": "- Can you explain in more detail why joint training on MOFs reduces validity? \n\n- The comparison to the equivariant baseline focuses on speed, but are there trade-offs in terms of equivariance properties? For example, does this affect the model's ability to respect the physical symmetries of crystals/ molecules?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tmiWDhMMQ7", "reviewer": "Reviewer_KRy3", "summary": "The authors introduce All-atom Diffusion Transformer (ADiT), a framework aimed at unifying latent diffusion approach across different spatial molecular structure modalities. The proposed method focuses specifically on small molecules and crystals. ADiT leverages a combination of joint variational autoencoder and latent diffusion transformer models to generate these molecular structures. The latent diffusion model constructs atom-wise latent codes, which the decoder then maps into atom descriptions, positions, and, optionally, crystal lattice parameters. The authors evaluate their approach on unconditional small molecule and crystal generation to benchmark its effectiveness.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> relatively small structures - benchmark against GEOM\n\n---\n\n**[In our response to Reviewer WqZ3](https://openreview.net/forum?id=89QPmZjIhv&noteId=RrA8d9t9eq), we presented results for GEOM based on your suggestion. ADiT continues to show strong results, generating physically realistic molecules and outperforming baselines.**\n\n---\n\nWe agree that systems from QM9 and MP20 are small, but they are complex and challenging. Doing well on them requires models to learn something about the underlying physics of atomic interactions.\n\nLet’s take MP20 as an example:\n- MP20 is restricted to up to 20 atoms in a unit cell because this is practically most relevant and largely covers the Materials Project distribution of known materials. Each atom's local environment is complex due to periodicity.\n- The S.U.N. rate metric for crystals/MP20 is a very challenging and practically relevant metric for materials discovery. ADiT has **improved S.U.N. rate from 4-5% of previous works up to 6.5%** (see Table 4). ADiT significantly outperformed the recently published MatterGen (Nature, 2024). Previous works showing smaller improvements have been published at top conferences (most recent example: FlowLLM at NeurIPS, 2024). There is a large community of researchers who care about S.U.N. rate on MP20.\n\nUltimately, the goal of this paper was to introduce the **first unified architecture** for generative modelling of 3D atomic systems, and to demonstrate **transfer learning** across periodic crystals and non-periodic molecular systems. We believe this will make a novel and valuable contribution to the ICML community.\n\nAlso note that our best models are based on **pure Transformers** - which have shown to be extremely scalable to large-scale inputs - so there is nothing inherent to our methodology that would prevent scaling up to larger systems or more samples.\n\n> MPTS-52\n\nNot applicable as MPTS is meant for evaluating **structure prediction**, where TS = temporal split, which is different from our de novo generation task.\n\n> unjustified claims of primacy - Flam-Shephard\n\nTo the best of our knowledge, ADiT is the first **unified, jointly trained** generative model for periodic and non-periodic systems to demonstrate **transfer learning**. \n\nFlam-Shepherd trains 3 independent models on 3 different datasets with a different tokenisation strategy for each. They don't discuss how/whether their method can be applied for one unified model. Also, our results on MP20 are better than FlowLLM, another language model which outperforms Flam-Shepherd. \n\nWe will discuss their paper in Related Work, as well as all other references you shared.\n\n> latent diffusion for generating chemical structures has been explored prior\n\nOur novelty is about **how** we used latent diffusion: For unification of system types (periodic and non-periodic systems together), as well as unification of mulit-modal data (categorical, numerical, floating point). \n\nWe think this is new - **nobody has used latent diffusion in this unified manner** - and will be of interest to the ICML community.\n- [a] - Not directly related to 3D structure\n- [b] - They technically do latent diffusion, but the latents are still multi-modal (four latent scalars & one latent vector) - thus GeoLDM still uses equivariant diffusion, which is slow, and GeoLDM is not applicable beyond small molecules. Additionally, there are **several discrepancies** between the implementation reported in the paper and released on GitHub (most concerning: https://github.com/MinkaiXu/GeoLDM/issues/6). ADiT’s latent diffusion formulation is unified and faster, as well as outperforms GeoLDM.\n- [c] - Highly specific to proteins as latent representations are created by aggregating along the sequence\n\n(We will cite [a] and [c])\n\n> high number of lambda coefficients, lack of intuition\n\nThe coefficients are needed for each of the different data types that constitute a 3D atomic system. The coefficients are **not dataset dependent**. The intuition for choosing them is simple: **“balance the relative magnitudes of the various losses”**. This is stated in the paper and is standard practice for supervised learning tasks like the VAE reconstruction used here.\n\n> conditional downstream tasks\n\nThe focus of this paper on introducing the architecture and demonstrating transfer learning. Efforts for extensions to conditional tasks such as protein-pocket conditioned molecule generation (ie. SBDD) as well as property conditioned crystal generation (similar to MatterGen) are underway as independent papers. From a methods standpoint, adding conditioning is straightforward in Diffusion Transformers via classifier-free guidance tokens.\n\nWe don’t feel that property conditioned small molecule generation, as benchmarked in papers like EDM/GeoLDM, is very practically relevant. We feel that it is unlikely that generating molecules with specific HOMO-LUMO gaps or dipole moments are relevant for drug discovery.\n\n> Suggestions\n\nWe will implement both.", "claims_and_evidence": "The overall claims in the paper appear valid and reasonable. However, I recommend avoiding unjustified claims of primacy, particularly in lines 103-105: “our work is the first to develop unified generative models for both periodic and non-periodic atomic systems”. Preliminary approach [a] has proposed a unified language model for small molecule and crystal generation, as well as protein pocket generation. \n\na. Language models can generate molecules, materials, and protein binding sites directly in three dimensions as XYZ, CIF, and PDB files, Flam-Shepherd et al., 2023", "methods_and_evaluation_criteria": "The chosen evaluation metrics and baselines are appropriate and reasonable, and incorporating DTF and PoseBusters-derived metrics introduces a novelty into the evaluation of spatial small molecule structures generation. However, the choice of datasets raises concerns.\nThe QM9 and MP20 datasets contain relatively small structures (up to 9 heavy atoms in QM9 and up to 20 heavy atoms in MP20). In contrast, state-of-the-art generative models [a, b] for spatial small molecules unconditional generation typically benchmark against GEOM-DRUGS [c], which includes structures with up to 91 heavy atoms. For crystal generation, prior work [d] has used more complex datasets, such as MPTS-52 (up to 52 heavy atoms) and [e] used custom subsets with up to 30 heavy atoms, in addition to MP20.\nWhile using QM9 and MP20 facilitates direct comparison with a broader range of baselines, incorporating larger and more complex datasets would provide deeper insights into the model’s capabilities and generalizability. The integration of larger structures dataset would further enhance the value of the proposed approach and allow comparison with state-of-the-art approaches on harder tasks.\n\na. MolDiff: Addressing the atom-bond inconsistency problem in 3D molecule diffusion generation, Peng et al., 2023\n\nb. BindGPT: A Scalable Framework for 3D Molecular Design via Language Modeling and Reinforcement Learning, Zholus et al., 2024\n\nc. Geom, energy-annotated molecular conformations for property prediction and molecular generation, Axelrod & Gomez-Bombarelli, 2022\n\nd. FlowMM: Generating Materials with Riemannian Flow Matching, Miller et al., 2024\n\ne. Fine-Tuned Language Models Generate Stable Inorganic Materials as Text, Gruver et al., 2024", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "As discussed in “Methods and Evaluation Criteria”, the unconditional generation setups for molecules, crystals, and MOFs are reasonable and valid. The results from single-dataset and jointly trained models demonstrate the necessity of multidomain training. Additionally, the Ablation Study in the Appendix provides a justification for the chosen architectural design.\n\nHowever, unconditional generation primarily serves as a proof-of-concept to demonstrate model capability, rather than its applicability to real-world tasks. A key question remains: Can the proposed approach be effectively adapted for conditional downstream tasks? For example, conformation generation [a], pocket-conditional generation [b], text description-conditioned generation [c], and MOF property optimization [d]. \n\nWhile the authors acknowledge the downstream tasks as future work in the Discussion section, the absence of at least one conditional setup or case study significantly reduces the paper’s impact. Including an experiment on conditional generation would strengthen the evaluation and provide clearer evidence of the model’s broader applicability.\n\na. Torsional Diffusion for Molecular Conformer Generation, Jing et al., 2022\n\nb. 3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction, Guan et al., 2023\n\nc. Fine-Tuned Language Models Generate Stable Inorganic Materials as Text, Gruver et al., 2024\n\nd. MOFDiff: Coarse-grained Diffusion for Metal-Organic Framework Design, Fu et al., 2023", "relation_to_broader_scientific_literature": "The idea of applying the latent diffusion framework for generating chemical structures has been explored in prior work for 2D [a] and 3D molecules [b], as well as 3D proteins [c]. While the proposed modification to integrate crystal lattice parameters is relatively straightforward, I consider it to be a novel contribution of the paper.\n\na. 3M-Diffusion: Latent Multi-Modal Diffusion for Language-Guided Molecular Structure Generation, Zhu et al., 2024\n\nb. Geometric Latent Diffusion Models for 3D Molecule Generation, Xu et al., 2023\n\nc. A Latent Diffusion Model for Protein Structure Generation, Fu et al., 2023", "essential_references_not_discussed": "As mentioned in the Claims and Evidence section, preliminary work [a] on multidomain spatial structure generation was not cited or compared against in the paper. \n\na. Language models can generate molecules, materials, and protein binding sites directly in three dimensions as XYZ, CIF, and PDB files, Flam-Shepherd et al.,  2023", "other_strengths_and_weaknesses": "The paper is well-written and easy to follow. The metrics effectively cover different aspects of spatial structure generation. The proposed modification to the latent diffusion framework is straightforward and focuses on small molecules and crystals, which, while limited, still represents a novel contribution. Additionally, the inference time efficiency of the approach, especially in comparison to pure diffusion models, is a strength of the method.\n\nA minor weakness of the approach is the high number of lambda coefficients in the autoencoder loss. There is a lack of intuition on how to choose these coefficients and whether they are dataset-dependent.", "other_comments_or_suggestions": "1. I would recommend making Figure 2 smaller or moving it to the Appendix. Instead, it would be more useful to include examples of generated structures in the main text, as visual inspection of results would be beneficial for readers.\n\n2. In Table 2, Validity results, I assume that one of the QM9-only ADiT entries should be marked with an asterisk *.", "questions_for_authors": "My concerns are addressed in the previous sections. The key issues are the lack of downstream tasks beyond unconditional generation and the use of relatively small datasets for small molecule and crystal generation.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "w1Ar03sxVV", "reviewer": "Reviewer_7DhL", "summary": "This paper proposes a new model based on all-atom Diffusion with transformers, for generating both periodic crystals and non-periodic molecules. Experiments are performed on standard benchmarks for these appications.  The authors show how this model helps speedup standard equivariant diffusion models and how it can benefit from scaling.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> small systems - simple setting…include experiments on more challenging tasks\n\n---\n\n**[In our response to Reviewer WqZ3](https://openreview.net/forum?id=89QPmZjIhv&noteId=RrA8d9t9eq), we presented results for GEOM based on your suggestion. ADiT continues to show strong results, generating physically realistic molecules and outperforming baselines.**\n\n---\n\nWe agree that systems from QM9 and MP20 are small, but they are complex and challenging. Doing well on them requires models to learn something about the underlying physics of atomic interactions.\n\nLet’s take MP20 as an example:\n- MP20 is restricted to up to 20 atoms in a unit cell because this is practically most relevant and largely covers the Materials Project distribution of known materials. Each atom's local environment is complex due to periodicity.\n- The S.U.N. rate metric for crystals/MP20 is a very challenging and practically relevant metric for materials discovery. ADiT has **improved S.U.N. rate from 4-5% of previous works up to 6.5%** (see Table 4). ADiT significantly outperformed the recently published MatterGen (Nature, 2024). Previous works showing smaller improvements have been published at top conferences (most recent example: FlowLLM at NeurIPS, 2024). There is a large community of researchers who care about S.U.N. rate on MP20.\n\nUltimately, the goal of this paper was to introduce the **first unified architecture** for generative modelling of 3D atomic systems, and to demonstrate **transfer learning** across periodic crystals and non-periodic molecular systems. We believe this will make a novel and valuable contribution to the ICML community.\n\nAlso note that our best models are based on pure Transformers - which have shown to be extremely scalable to large-scale inputs - so there is nothing inherent to our methodology that would prevent scaling up to larger systems or more samples.\n\n> extremely small systems - up to 10 atoms\n\nJust to clarify the details here:\n- MP20: **20 atoms** in a periodically repeating unit cell. This means that each atom can **interact with a far greater number of atoms** than the size of the unit cell due to periodicity.\n- QM9: 9 heavy atoms, which doesn’t include hydrogen. When including hydrogen (which is how our models are trained), systems go up to **20-30 atoms on average**.\n\nWhich is why we wrote “tens of atoms”, not 10 atoms.\n\n> extend metrics for small molecules, following MiDi\n\nOur evaluation followed **[Symphony](https://arxiv.org/abs/2311.16199)**, the most recently accepted small molecule generation paper at ICLR 2024. We believe that **Symphony's evaluation metrics are an improvement upon MiDi**. \n\nHere are our justifications:\n- We have included the validity and uniqueness metrics from MiDi/EDM. We found ADiT to outperform EDM and GeoLDM, both of which do better than MiDi.\n- Instead of atom and molecule stability metrics from MiDi/EDM which based on simple valency heuristics, Symphony used the **PoseBusters** metrics suite which provides information about the physical realism of generated molecules. PoseBusters provides a more holistic set of metrics using force field relaxations and physics-based checks instead of relying on heuristics.\n- PoseBusters also includes metrics for bond distances, angles, rings, and geometries, which supersede MiDi's histogram metrics and are more interpretable (in our opinion).\n- PoseBusters is now widely known and used in academia and industry, e.g. AlphaFold3 used PoseBusters.\n\n> \"the paper proposes another approach to generate both material and molecules; as such the idea of learning a model across modalities is interesting and fairly new\"\n\nThe goal of this paper was to introduce the **first unified architecture** for generative modelling of 3D atomic systems, and to demonstrate **transfer learning** across modalities. \n\nWe disagree with your characterisation that our work is just “another” approach as we are not aware of any other generative architecture that jointly generates periodic and non-periodic systems. We believe this will make a novel and valuable contribution to the ICML community, and would kindly request you to reconsider your position on our novelty.\n\n> citing Flam-Shepherd et al. 2023 and Pinheiro et al 2023\n\nWe will definitely cite and discuss both papers in the Related Work - thanks for the pointers. We will also cite and discuss the MiDi paper and metrics in our revision.\n\n(Note that this year, ICML does not allow us to upload an updated PDF.)", "claims_and_evidence": "the claims are well supported.", "methods_and_evaluation_criteria": "using a transformer for building a unified model for molecule and materials makes a lot of sense; moreover, replacing equivariance-based constraints by data augmentation to leverage more expressive architecture is also a recent trend in this field.", "theoretical_claims": "there were no theoretical claims", "experimental_designs_or_analyses": "yes", "relation_to_broader_scientific_literature": "the paper proposes another approach to generate both material and molecules; as such the idea of learning a model across modalities is interesting and fairly new.", "essential_references_not_discussed": "given that the authors emphasize on the benefits of not having equivariance in their model, the authors could cite other work that also removed equivariance constraints in generative models for molecules / materials e.g. \"Language models can generate molecules, materials, and protein binding sites directly in three dimensions as xyz, cif, and pdb files\" Flam-Shepherd et al. 2023 or other work for unconditional molecule generation e.g. \"3D molecule generation by denoising voxel grids\" Pinheiro et al 2023. This section of related work without equivariance is missing in the related work, yet the authors did mention AlphaFold3 and Wang et al 2024 which are also relevant work on this front, not explicitly on generation.", "other_strengths_and_weaknesses": "overall, I found the paper well executed with a nice idea of unifying different generative tasks; in particular the choices on the ML front, follow many recent trends in generative models which make a lot of sense. yet the paper is only focused on extremely small systems (up to 10 atoms as mentioned in the conclusion); arguably this is a simple setting. I would recommend the authors to include some experiments on some more challenging tasks; in particular staying with small molecule generation, GEOM-drugs is a more challenging setting for which comparison to other baselines would be more convincing. Moreover, I recommend the authors to extend the metrics in the small molecule front e.g. following the framework of MiDi: Mixed Graph and 3D Denoising Diffusion for Molecule Generation for which several models are compared agains.", "other_comments_or_suggestions": "see above", "questions_for_authors": "please address the limitations on the scale of the datasets and metrics of small molecule generation.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This is an interesting paper addressing an important area of research. The reviewers are generally supportive to this paper, esp. after author rebuttals. Thus an accept is recommended.
Decision: Accept (poster)