D3MES: Diffusion Transformer for 3D Molecule Generation

• D3MES is a diffusion-based generative model that leverages SE(3)-equivariant multihead self-attention to synthesize chemically valid 3D molecular conformations.
• The architecture omits hydrogens during encoding, re-attaches them via valence completion, and encodes heavy atom coordinates, one-hot features, and bond matrices.
• It employs a discrete-time diffusion process with class-conditional sampling, achieving state-of-the-art performance on benchmarks like GEOM-QM9 and GEOM-Drugs.

D3MES (Diffusion Transformer with Multihead Equivariant Self-Attention for 3D Molecule Generation) is a diffusion-based generative model designed to synthesize accurate and chemically valid three-dimensional (3D) molecular conformations. D3MES leverages SE(3)-equivariant multihead self-attention to model the symmetries inherent in molecular structures and incorporates a classifiable diffusion process capable of conditioning on discrete molecule types. The architecture is engineered for robust generation of molecular conformers across diverse chemical classes, with specific innovations for hydrogen atom treatment and class-conditional sampling, achieving state-of-the-art performance on key molecular generation benchmarks (Zhang et al., 13 Jan 2025).

1. Architectural Overview and Data Encoding

D3MES represents each molecule using three primary data channels: the 3D coordinates of non-hydrogen (heavy) atoms $x \in \mathbb{R}^{N \times 3}$, atom-type one-hot features $h \in \{0,1\}^{N \times C}$, and the bond adjacency matrix $b \in \{0,1\}^{N \times N}$. Hydrogen atoms are stripped during the preprocessing stage to simplify learning and avoid combinatorial explosion from permutation-invariant hydrogens. They are subsequently re-attached by valence completion after generative sampling.

The model employs an SE(3)-equivariant multihead attention layer, which aggregates information over local neighborhoods $N_i$ in the point cloud formed by heavy atom coordinates and features. This frontend is inspired by SE(3)-Transformers [Thomas et al., 2018], with careful construction of attention weights and value projections so that the updates are equivariant to rigid transformations. The aggregation step across heads $l$ and feature channels $k$ for each atom $i$ is given by:

$$f_{\mathrm{out},i}^l = W_V^l f_{\mathrm{in},i}^l + \sum_{k \geq 0} \sum_{j \in N_i \setminus i} \alpha_{ij} W_V^{lk}(x_j-x_i) f_{\mathrm{in},j}^k,$$

where $\alpha_{ij}$ are softmax-normalized attention weights dependent on both features and relative positions, and the $W^{lk}(x)$ operators employ radial and spherical harmonic filters to maintain SE(3)-equivariance.

After the coordinate channels are refined by the equivariant attention, all molecule channels are concatenated, patchified $h \in \{0,1\}^{N \times C}$0, embedded via a linear layer, and further supplemented with timestep, (optional) class embedding, and positional encodings. The backbone comprises $h \in \{0,1\}^{N \times C}$1 transformer blocks (with $h \in \{0,1\}^{N \times C}$2, hidden dimension $h \in \{0,1\}^{N \times C}$3, MLP expansion ratio $h \in \{0,1\}^{N \times C}$4, and 8 attention heads), following the standard transformer configuration.

2. Diffusion Process and Denoising Parameterization

The generative process in D3MES follows discrete-time forward noising and reverse denoising steps as in standard diffusion models. The forward process is defined as:

$h \in \{0,1\}^{N \times C}$5

with a prescribed noise schedule. At any time step $h \in \{0,1\}^{N \times C}$6, $h \in \{0,1\}^{N \times C}$7 is realized as:

$h \in \{0,1\}^{N \times C}$8

where $h \in \{0,1\}^{N \times C}$9.

For generative sampling, the D3MES network parameterizes the reverse kernel as:

$b \in \{0,1\}^{N \times N}$0

with mean $b \in \{0,1\}^{N \times N}$1 (following the common parameterization):

$b \in \{0,1\}^{N \times N}$2

where $b \in \{0,1\}^{N \times N}$3 is the predicted noise residual. Both $b \in \{0,1\}^{N \times N}$4 and $b \in \{0,1\}^{N \times N}$5 are produced at each transformer output location.

3. Class-Conditional Diffusion and Equivariance

D3MES incorporates discrete class conditioning via embedding the label $b \in \{0,1\}^{N \times N}$6 (e.g., cyclic/noncyclic class) into a vector $b \in \{0,1\}^{N \times N}$7, which is injected into all transformer blocks using AdaLN-Zero (adaptive layer normalization). The denoising model thus learns $b \in \{0,1\}^{N \times N}$8, supporting flexible and multi-class molecular generation.

An optional classification head $b \in \{0,1\}^{N \times N}$9 is attached to the clean reconstruction $N_i$0, with classification loss:

$N_i$1

which encourages $N_i$2 to encode discriminative class features, and enables accurate class-conditional conformer synthesis.

An equivariance regularizer may be added to train the denoiser to be consistent under global SE(3) transformations:

$N_i$3

thus enforcing geometric symmetry priors at the representation level.

4. Hydrogen Atom Treatment and Valence Completion

Generation and denoising in D3MES omit hydrogens from both the coordinates and feature tensors, a practice that simplifies conformational modeling due to hydrogen's indistinguishability and relatively predictable valency. After generative sampling is complete, hydrogens are re-attached to the heavy atom backbone by valence completion. For each heavy atom $N_i$4, the number of hydrogens $N_i$5 is estimated as:

$N_i$6

where $N_i$7 is the atom's maximal valence and $N_i$8 is its heavy-atom degree. Hydrogens are placed with standard bond geometries (e.g., tetrahedral for sp$N_i$9) ensuring chemical plausibility in the final structure.

5. Training Objectives and Optimization

The total training objective for D3MES is:

$l$0

where:

• $l$1 is the mean squared error (MSE) between actual and predicted noise,
• $l$2 is a variational lower bound loss based on the KL divergence between the true and predicted transitions,
• $l$3 applies class-conditional supervision,
• $l$4 enforces SE(3)-equivariance.

The model is trained for approximately 200k steps with batch size 64 using Adam optimizer at learning rate $l$5 and weight decay $l$6. The diffusion process employs $l$7 timesteps with a linear noise schedule $l$8.

6. Generation and Postprocessing Pipeline

The unconditional or class-conditional generative pipeline in D3MES proceeds as follows:

1. Initialize $l$9 for heavy-atom channels.
2. For $k$0: sample $k$1 from model predictions, optionally conditioned on $k$2.
3. Extract denoised $k$3 as the heavy atom backbone.
4. Re-attach hydrogens using valence completion and standard chemical geometry.
5. Output the full 3D molecular structure.

7. Experimental Benchmarks and Performance Metrics

D3MES was evaluated on GEOM-QM9 (130k, ≤9 heavy atoms), GEOM-Drugs (290k, ≤91 heavy atoms), and a cyclic/noncyclic classification subset of QM9 (Zhang et al., 13 Jan 2025). Key validation metrics included:

• Atom-stable (%): correct valence for each atom after bond inference.
• Mol-stable (%): all atoms stable in a molecule.
• Valid (%): structural validity by RDKit.
• Uniqueness (%): unique SMILES canonicalization.
• Classification accuracy (%): for class-conditional generation.

D3MES achieved the following (samples of 10,000):

• QM9: atom-stable 99.3%, mol-stable 92.2%, valid 97.8%, uniqueness 98.3%.
• Drugs: atom-stable 99.8%, mol-stable 94.7%, valid 99.98%, uniqueness 99.9%.
• QM9 class-conditional: noncyclic accuracy 97.4%, cyclic accuracy 96.6%.

These results consistently match or surpass preceding approaches such as EDM, GCDM, MDM, MiDi, and GeoLDM. D3MES thus provides a state-of-the-art diffusion framework with strong class-conditional, chemically valid, and unique 3D molecule generation capabilities (Zhang et al., 13 Jan 2025).