Navigating the Design Space of Equivariant Diffusion-Based Generative Models for De Novo 3D Molecule Generation
Authors: Tuan Le
Julian Cremer
Frank Noe
Djork-Arné Clevert
Kristof T Schütt
Venue: ICLR 2024 poster (2024)
Keywords: Generative Modelling
Molecule Design
Denoising Diffusion Probabilistic Models
Ablation Study
Equivariant Graph Neural Network
3D Molecule Generation
Diffusion Model
Primary area: generative models
TL;DR: We explore the design space of diffusion models for 3D molecular modelling. Based on our findings we provide an equivariant diffusion model that outperforms the current state-of-the-art models in the domain by a large margin.

Abstract: Deep generative diffusion models are a promising avenue for 3D de novo molecular design in materials science and drug discovery.
However, their utility is still limited by suboptimal performance on large molecular structures and limited training data.
To address this gap, we explore the design space of E(3)-equivariant diffusion models, focusing on previously unexplored areas.  
Our extensive comparative analysis evaluates the interplay between continuous and discrete state spaces.
From this investigation, we present the EQGAT-diff model, which consistently outperforms established models for the QM9 and GEOM-Drugs datasets.  
Significantly, EQGAT-diff takes continuous atom positions, while chemical elements and bond types are categorical and uses time-dependent loss weighting, substantially increasing training convergence, the quality of generated samples, and inference time. We also showcase that including chemically motivated additional features like hybridization states in the diffusion process enhances the validity of generated molecules.  
To further strengthen the applicability of diffusion models to limited training data, we investigate the transferability of EQGAT-diff trained on the large PubChem3D dataset with implicit hydrogen atoms to target different data distributions. Fine-tuning EQGAT-diff for just a few iterations shows an efficient distribution shift, further improving performance throughout data sets.   
Finally, we test our model on the Crossdocked data set for structure-based de novo ligand generation, underlining the importance of our findings showing state-of-the-art performance on Vina docking scores.

Published as a conference paper at ICLR 2024
NAVIGATING THE DESIGN SPACE OF EQUIVARIANT TAFUSION-BASHER GENERATIVE MODELS FOR DE Novo 3D MOLECULE GENERATION
Tuan 1.2, Julian Cremer
Frank No62.3, Djork-Arné Clevert Kristof Schütt!
1Pfizer Research & Development, FFeeie Universität Berlin 3 "Microsoft Research, *University Pompeu Fabra
ABSTRACT
Deep generative diffusion models a promising avenue for 3D de novo molec- ular design in materials science and drug discovery. However, their utility is still limited by suboptimal performance on large molecular structures and limited train- ing data. To address this gap, we explore the design space of E(3)-equivarian dif- fusion models, focusing on previously unexplored areas. Our extensive compara- tive analysis evaluates the interplay between continuous and discrete state spaces, From this investigation, we present the EQGAT-diff model, which consistently outperforms established models for the QM9 and GEOM-Drugs datasets. Signifi- cantly, EQGAT-diff takes continuous atom positions, while chemical elements and bond types are categorical and uses time-dependent loss weighting, substantially increasing training convergence, the quality of generated samples, and inference time. We also showcase that including chemically motivated additional features like hybridization states the diffusion process enhances the validity of gener- ated molecules. To further strengthen the applicability of diffusion models lim- ited training data, we investigate the transferability of EQGAT-diff trained on the large PubChem3D dataset with implicit hydrogen atoms target different data distributions. Fine-tuning EQGAT-diff for just few iterations shows an efficient distribution shift, further improving performance throughout data sets. Finally. we test our model on the Crossdocked data set for cruucture-bssed de novo ligand generation, underlining the importance of our findings showing state-of-the-art performance on Vina docking scores.
INTRODUCTION
The enormous success of machine learning (ML) in computer vision and natural language process- ing recent years has led to the adaptation of ML in many research areas in the natural sciences, such physics, chemistry, and biology, with promising results. Specifically, modern drug discov- ery widely utilizes ML to efficiently screen the vast chemical space for de novo molecule design in the carly-stage drug discovery pipeline. An important aspect is the structure-basse aware design of novel molecules in 3D space (Schneuing et al. 2023 Guan et al. 2023 Stärk et al., 2022 Corso al., 2023). However, incorporating the 3D geometries of molecules for rational and structure-based drug design is challenging, and the development of ML models in this domain is anything but easy, as these models need function with just limited amount of data learn physical rules in 3D space accurately. Fortunately, applying geometric deep learning to molecule generation has gained attention in the scientific community recent years, paving the way for novative approaches. These result in diffusion models quickly becoming state-of-the-arr in this area due to their ability to effectively learn complex data distributions (Hoogeboom al., 2022 [gashov et al. 2022; (Schneuing et al., 2023; Vignac et al. 2023) Guan et 2023). While this has enabled researchers to develop generative models for molecular design that can sample novel molecules in 3D space, several drawbacks and open questions remain prevalent for practitioners. Molecule gener-
target-
"Equal contribution. Correspondence {tuan. le, juli molliancceemrr]pp..... cremer] zom

---

Published as a conference paper at ICLR 2024
ative models are required to both generate realistic molecules in 3D space and preserve fundamental chemical rules, i.e., correct bonding and valencies. Various design decisions have to be taken into account that heavily impact the performance and complexity of those models. Hence, there is high need to better understand the design space of diffusion models for molecular modeling. Moreover, the availability of molecular data is not as abundant, confronting MI models with relatively narrow and specific data distributions. That is, ML models are usually trained explicitly for each data set, which is unfavorable regarding the efficient use of training data and computing resources.
This work introduces the E(3)eqquvvriiant graph attention denoising neural network EQGAT-diff: We systematically explore the design space of 3D equivariant diffusion models. including various parameterizations, loss weightings, data, and input feature modalities. Beyond that, we explore an efficient pre-training scheme on molecular data with implicit hydrogens. This enables a data- and time-efficient training and fine-tuning procedure leading to higher molecule stability. Our contribu- tions are the following:
We propose EQGAT-diff fast and accurate 3D molecular diffusion model that employs E(3)-equivariant graph attention. Our proposed model achieves SOTA results in shorter training time and with less trainable parameters than previous architectures.
We systematically explore various design choices for 3D molecular diffusion models and provide a thorough ablation study across the popular benchmark sets QM9 and GEOM- Drugs We propose a time-dependent loss weighting as a crucial component for fast training convergence, better inference speed, and sample quality.
We demonstrate the transferability of an EQGAT-diff model pre-trained on the PubChem3D dataset to smaller but complex molecular datasets. After short fine-tuning the target distribution, we show that the model outperforms models trained from scratch on the target data by only training on subsets.
We extend the diffusion process by modeling chemically motivated additional features and show a further significant increase in performance.
In summary, we found the following ingredients be crucial: =(3)-equivariant graph attention, dime-depeneent loss weighting, unconditional pretraining on large databases comprising 3D con- formers like PubChem3D, and adding chemical features like aromaticity and hybridization state as feature input to the denoising diffusion model.
2 RELATED WORK
Denoising diffusion probabilistic models (DDPM) (Sohl-Dickstein et al.) 2015 Ho et al. 2020 Kingma et al. 2021: Song et al., 2021b) have achieved great success various generation tasks due to their remarkable ability to model complicated distributions in the image and text processing community (Popov et al. 2021 Kong et al. [022]] Salimans Ho. 2022 Rombach et al. 2022 Karras et al. 2022 Li et al. 2022 Kingma & Gao, 2023). Deep generative modeling in the life sciences has become a promising research area, c.g., conditional conformer generation based on the 2D molecular graph, in which (Mansimov et al. 2019; Simm & Hernandez-Lobate 2020) leverage the idea of variational autoencoders, while recent work by (Xu et al., 2022 Ving et al. 2022) use DDPMs, to predict the 3D coordinates with the help of 3D equivariant graph neural networks. In the de novo setting, another line of research focuses on directly generating the atomic coordinates and elements, using either autoregressive models (Gebauer et al., 2019 2022 Luo & Jil 2022).
where atomic elements are generated one by one sequentially, or neural learning algorithms based on continuous normalizing flows (Satorras et al. 202T) that are computationally expensive due to the integration of an ordinary differential equation, leading to limited performance and scalability on large molecular systems. Diffusion models offer efficient training by progressively applying Gaussian noise to transform complex data distribution to approximately tractable Gaussian prior, intending to learn the reverse process. [Hoogeboom et al. (2022) introduced E(3) equivariant dif- fusion model (EDM) for de novo molecule design that simultaneously learns atomic elements next to the coordinates while treating chemical elements as continuous variables to utilize the formal- ism of DDPM. Follow-up works leverage EDM and develop diffusion models for linker design (Igashov et al. 2022) or ligand-protein complex modeling (Schneuing et al. (2023). Another line of work leverages the formalism of stochastic differential equations (SDRs) (Song al. (0211) and
2

---

Published as a conference paper at ICLR 2024
Schroedinger Bridges with extension to manifolds (De Bortoli et al. [021 2022) to generating 3D
conformer of fixed molecule into a protein pocket (Corso et al. 2023). while (Wu et al. 2022)
modifies the forward diffusion process to incorporate physical priors.
BACKGROUND
Problem Formulation and Notation We investigate the generation of molecular structures in de novo setting, where atomic coordinates, chemical elements, and the bond topology are sampled. molecular structure is given by where the vertices V (UI,..., UN) refer to the N atoms. Each vertex is tuple v; = (Ti, h.) comprised of the atomic coordinate in 3D space Ti and chemical element h4. The latter is one-hot encoded for K elements, i.c., h, (0, o..... 1,0), The edges E = (ety) (ery)== describe the connectivity of the molecule, where each edge feature can take five distinct values, namely the existence of no bond between atom and j. Additionally, we exclude self-loops in our data representation. We write node features as matrices X IR- and H E (0, 1] while the bond topology is given by E E {0, 1 x5, We aim to develop a probabilistic model that is invariant to the permutation of atoms of the same chemical element and roto-translation of coordinates in 3D space. That means, regardless of how atom indices in the node-feature matrix H are shuffled and coordinates X roto- translated, the probability for molecular structure N' remains unchanged.
single- double-, triple- or aromatic bond
3.1 DENOISING DIFFUSION PROBABILISTIC MODELS
Discrete-time diffusion models (Sohl-Dickstein et al. 2015 Ho et al. 2020) are latent variable gen-
erative models characterized by a forward and reverse Markov process over steps. Given a sample from the data distribution To ~ ((zu), the forward process ( (T)) = II ( trans- forms it into sequence of increasingly noisy latent variables F1.T = (I1, 22, ................................ 2r) and z; E X.
The learnable reverse Markov process po(Fu:)) p(rr) II. Po (2t 12x)) is trained to gradually denoise the latent variables approaching the data distribution. (Sohl-Dickstein et al. (2015) initially proposed a diffusion process for binary and continuous data, while the latter consists of Gaussian transition kernels. The learning process for discrete data has been introduced by Hoogeboom et al. (02T)) and Austin et al. (2021). leveraging categorical transition kernels in the form of doubly stochastic matrices. Crucially, both forward processes define tractable distributions determined by a noise schedule { such that the reverse generative model can be trained efficiently. As molecular data consists of atoms, bonds, and 31) coordinates, recent work leverages combination of Gaussian and categorical diffusion for 3D molecular generation (Peng et al., 2023 Vignac et al.) 2023 Guan al. 2023).. subtle property of tractable transition kernels is that the distribution of noisy state conditioned on data sample is also tractable, and for continuous or discrete data follows multivariate normal or categorical distribution
((x Xo) = (xxiVoxxx (1 - T) and ((ca/oo) = C(cala,co + +1 - az Ge).
(1)
where ai (1 Bx) ) E (0, 1), and (1 (l)))) determine a variance-preserving (VP) noise sched- uler Song et al. (2021b). The vector with T 1 = 1 determines the prior distribution of the cate- gorical diffusion, är + O. Possible prior distributions are the uniform distribution over K Kcclasses or the empirical distribution of categories in a dataset. In this work, we perturb atomic coordinates X, chemical elements H, and edge features E independently, using Gaussian and categorical diffu- sion. To conserve the edge-symmetry between atoms é and j. we only perturb the upper- apper-triangulrr elements of E. Diffusion models are trained by maximizing the variational lower-bound of the data log-likelihood (Sohl-Dickstenn et al. 2015) Kingma et al. 2021 Austin et al. 202T) decomposed as
DE (g(ær) (p(rr) denote the reconstruction, and prior loss. These two loss terms are commonly neglected during op- timization, while the diffusion loss L, = - -D [9(x,-1 [2,, xo) po( [xx]] has closed-fore ex- pression since q [2,, Iu)) is either a multivariate normal or categorical distribution, enabling ef- ficient KL divergence minimization by predicting the corresponding distribution parameters. These are defined as a function of x, and Zo. implying that the diffusion model is tasked to predict the clean data sample do optimize I., (Ho et al. 2020 Austin et al. 2021
log p(x) 2 Lo + I.prirr Σ Lt,, where Lu = log p(zo (1)) and Terrror
3

---

Published as a conference paper at ICLR 2024
4 EQGAT-DIFF
An essential requirement to obtain a data-efficient model is to reflect the permutational symmetry of atoms of the same chemical element and the otto-translational symmetries of 3D molecular struc- tures, In machine learning force fields, it has been shown that rotationally invariant features alone do not accurately represent the 3D molecular structure and hence require higher-order equivariant features (Schütt et al. 2021 Batzner et al. 2022 Thölke & Fabritiis 2022. Batatia et al.) (022)).
In short, a function f X + y mapping from input space X output space y equivariant to the group G iff f(g.z) = (x), where g. denotes the action of the group element g E G on an object x; y E X. As graph neural networks operate on graphs and map nodes into a feature space through shared transformations among all nodes, permutation equivariance is naturally preserved Bronstein et al. (202T). contrast, point clouds are embedded in 3D space, so we additionally con- sider the rotation, reflection, and translation group in RS often abbreviated as E(3). For the atomic coordinates, we require that f(XQ + t) matrix and t E Rd translation vector added now-wise. Group equivariance of a function f in the context of a diffusion model for molecular data is a requirement to preserve the group invariance for a probability density, as shown by Köhler et al. (220)) and [Xu and[Xu et al. (2022). To better address the challenge of molecular modeling, we propose a modified version of the EQGAT architecture Le et al. (2022), coined EQGAT-diff, which leverages attention-based feature aggregation of neigh- boring nodes. EQGAT-diff employs rotation equivariant vector features that can be interpreted as learnable vector bundles, which the denoising networks of EDM[Hoogebooo
[(X) + t, where Q E ((3) is a rotation or reflection
et al. (2022) and MiDi Vignac et al. (2023) are lacking. Point clouds are modeled as fully connected graphs, so message passing computes all pairwise interactions. Equivariant vector features are obtained through a tensor product of scalar features with normalized relative positions x(ji,)) Tx--x.I (xj x,) as similarly
proposed in the works of Jing et al. (202I) and |Shüttt al. (2021). We iteratively update hidden edge features within the EQGAT-diff architecture to handle the edge prediction between two atoms. To achieve this, we modify the message function of EQGAT as
= MLP((h)) (1) :h ( 60
( (1) :.. I)
())
where; ; denotes concatenation of E(3) invariant embeddings and MLP is 2-layer multi-layer per- (1) (1)
ceptron. The message embedding m = (a jr j Gj : d 2 E R is further split into sub-embeddings that serve as filter to aggregate node information from all other source nodes j.
(1)
exp(a
h (11)) h(")
w Who by
and e
(+1)
woooeef + +
E, exp(a!))
(1)
v(t2)
(0)
Σ
X jinn bb/ + +1 & () o ov, ((t)
N
(41) (1)
1
Σ
(1) ( ji Xjnn
x x;
N
where 1 = (1, and a is the SiLU activation function. The embeddings are further updated and normalized with details explained the Appendix I
EXPLORING THE DESIGN SPACE OF 3D MOLECULAR DIFFUSION MODELS
The design space of diffusion models has many degrees of freedom concerning. among others, the data representation, training objective, forward inference process, and the denoising neural network. In de novo 3D molecular generation, [Hooeeooom et al. (022) (EDM) utilized the e- parameterization and proposed to model chemical elements as well as atomic positions continuously. Vignac al. (2023) proposed MiDi, which generates the molecular graph and 3D structure simulta- neously. This model uses the to-parameterization and employs the framework developed by Austin et al. (02T) to model not only chemical elements but also formal charges and bond types in dis- crete state space. Both parameterizations optimize the same objective, i.c., aiming to minimize the KL divergence (g(2, i Ix, 2o) Ipo (z,-1 |x,)]. Ho et al. (2020) found that optimizing the dif- fusion model in noise-space on images results in improved generation performance than predicting
4

---

Published as a conference paper at ICLR 2024
Table 1: Comparison of EQGAT-diff on QM9 and GEOM-Drugs trained with wa or (x(t) loss- weighting. We report the mean values over five runs of selected evaluation metrics with the margin of error for the 95% confidence level given as subscripts. The best results are bold.
GEOM-Drugs
Weighting
Mol. Stability
Validity
Connect. Comp.
Mol. Stability 1
Validity 1
Connect, Comp. 1
97 97.38.u 33 n.2
99. auo
20.19
71. 71.41,n.22
86.57 a0..3
98.68.0.:1
28.96
99.94
91.60
84.0220.19
95.08.0.12
the original image from a noised version. While noise prediction might benefit the image domain, this does not necessarily generalize to 3D molecular data. In fact, MiDi outperforms EDM across all standard benchmark metries and datasets. However, whether the improved performance stems from the parameterization, the employment of categorical diffusion for discrete features, using bond types and other chemical features has still been unclear, leaving researchers and practitioners guessing which kind of diffusion model to deploy in their respective tasks.
In this section, we explore the design space of de novo molecular diffusion models in these three aspects while consistently using EQGAT-diff as the denoising neural network to isolate the effect of each change for better comparison. The diffusion models are evaluated on the QM9 dataset RRaakkishhann et al. 2014) containing molecules with up to 9 heavy atoms, and the GEOM-Drugs dataset (Axelrod & Gómez- Boanzzz-Bmbarlllii (2022) containing up to 15 heavy atoms. We utilize the data splits from Vignac et al. (2023) and benchmark all models on full molecular 3D graphs that include explicit hydrogens.
5.1 TRAINING DETAILS
We either employ noise prediction (e-parameterization) data prediction (2o parameterizaaton) to train EQGAT-diff such that the group equivariant network fo(z;) receives a noisy molecule T = (x, H, E,) and either outputs the applied noise = (éx, ÉH, (g.) ) or prediction of the clean data d = (X,, Ho: Ea) of coordinates, chemical elements well as bonds. We draw random batch of molecules and uniformly sample steps t E 2(1, 2") and optimize the diffusion loss L, for each sample. While we use the mean squared error loss for the e-model, the zo xo-model is optimized using loss functions la depending on the data modality d. Here, la is mean squared error for continuous and the cross- eross-yntropy loss for categorical data. This leads to a composite loss
L = (t)]le - (0 t)12 and Li. = w(t La(Io, (xt, t); Am),
(2)
weighting, which we adopt from Vignac et al. (2023) and set to >₂ 3, Xy 0.4, A 2. For noise learning, we adopt an atom-type feature scaling of 0.25 as in Hoooeboom et al. 2022)). Notably, w(t) is a loss weighting commonly set to 1 across all time steps, which has been previously found work best (Ho et al. 2020). In contrast to this result, we find this term to be crucial for molecular design, as discussed in Sec. A.4 Following Vignac et al.
where Am denotes a modality-depender
2023), we also employ an adaptive noise schedule (see Appendix A.I.I).
5.2 METRICS
Following (Hoogeboom et al. 2022), we measure validity using the success rate of RDKit sani- tization over 10,000 molecules (pre-selecting connected components only) with the caveat that the RDKit sanitization might add implicit hydrogens to the system to satisfy the chemical con- straints. Therefore, checking atomic and molecular stability for the correct valencies using pre- defined lookup table that complements the validation is essential. Further, we propose to include diversity/similarity measures. We evaluate the diversity of sampled molecules using the average Tanimoto distance and measure the similarity with the training dataset via Kullback-Leibler diver- gence and the Tanimoto distance. Lastly, following Vignac et al. (2023), we use the atom and bond total variations (AtomsTV and BondsTV) that measure the distance between the marginal distri- bution of atom types and bond types for the generated set and the test set, respectively. Moreover, employ the Wasserstein distance between valencies, bond lengths, and bond angles, with the latter two being 3D metrics to evaluate conformer accuracy. For more details, we refer to Vignac et al. (023) and Appendix A.2
5

---

Published as a conference paper at ICLR 2024
Kingma et al. (2021) have shown that the intermediate KL-divergence loss Le in the variational lower bound (VLB) for a Gaussian diffusion can be simplified to
L. w (1)) - 20(e,, 1)1 - - vo.
-1) - SNR( (Te, 1))))
where (L - a diffusion models for molecules are commonly set to 1, i.e, Wa = 1 in EDM or MiDi (Hoogeboom al. 2022 Vignac al. (023).
1-a, refers to the signal-to-noise ratio. However, the weighting coefficients in
We hypothesize that denoising requires high accuracy for timesteps close to the data distribution generate valid molecules, while errors close the noise distribution are neglectable. Such loss weighting has been proposed by (Salimans & Ho (2022) as 'truncated SNR', which we modify for our use case. Specifically, we perform experiments with the loss weighting
(t) max(0.05, min(1.5, SNR()))),
which matches our hypothesis about learning with higher weightings approaching the data distri- bution (see A and Fig. () We clip the maximum value of 1.5 to enforce larger weightings to enhance learning compared to uniform weighting, followed by an abrupt exponential decay. We train EQGAT-diff using Gaussian diffusion on atomic coordinates and categorical diffusion for chemical elements, formal charges, and bond features following the parameterization proposed by Vignac et al. 2023), predicting a clean data sample žu given noisy version H₁. As shown in Table | train- ing EQGAT-diff on GEOM-Drugh with w,(t) results in better generative model that can sample molecules preserving chemistry rules, measured in increased molecule stability of 91.60%, com- pared to the EQGAT-diff which was trained with Was only achieving 87.59%. As the w²,() loss weighting achieved better evaluation metrics and significantly faster training convergence on the QM9 and GEOM-Drugh datasets, we choose as default for the following experiments conducted in this work. We provide further empirical evidence in Appendix A.3
Table 2: Overall performance of EQGAT-diff on QM9 and GEOM- for discrete and continuous diffusion as well as noise (e) and data learning (zo). Discrete or continuous diffusion is denoted as 'disc' and 'cont', respectively, given as subscripts, E- and zo-parameterriztioo as superscripts. We report mean values over five sampling runs with 95% confidence intervals as subscripts. The best results are bold.
Dataset
QM9
GEOM-Drags
Model
EQGATS...
EQGAT
EQGATT:
EQGAT
EQGATES
EQGAT:
McL. Stab. f
98.68
96. 96.45
96.18-0.34
91.60-00144
90.46-00.0
85.19
Atom, Stub1
99.92200
99.79_mm
20.11
99.72-nutt
99.73.cmm
99.32
Validity
96.77 26.70gn
97 97.04-2-17
84.02
80196
79.13 16
Connect. Comp.
g0.94_om
99.82 90.82_am
99.71
95.008 20
93.30 03.30
94. 10
Novelty
64.03g0a1
0..5
73.40
99.87 99.87_
99.83_c.o
99. 09.82_.0.
Uniqueness
100.00
100.0
100.00
100.00
100.00
100.00
Diversity
91.72-xoz
91.51=pa0
-0.00
89.00-20us
88.87=0.0
KL Divergence
21.36
91.41 2
20:00 88
87.17 2.2.5
87.35 20:00
87. 2..00...
Train Similtriiyy
0.076
0.076
0.075
0,113
0.114
0.1 14 2
AtomsTV [10
2.0
3.400010
2.00
BoodsTV [10 L1
1.22000
1.8:nm
1.2:0.00
2.4-000
2.400.0
244
ValeneyW; [10-9]4
0.6
1.9amm
0.9.c.oo
1.2.0.10
1.9.c.
1.503.00
Besll enghoeW, [10
0.2g0
0.5
10
0.3gr.m
0.7g
BoodAngleeWW
1.86g00
0.52gg.00
0.92gamm
0.95g...0
1.77
5.3 DIFFUSION PARAMETERIZATION
20 AND DISCRETE VS CONTINUOUS
Diffusion models for continuous data are commonly implemented using the e-parameterization Ho al. (2020), which is connected to denoising score matching models proposed by Song Ermon (2019). ) Diffusion models have quickly adapted this setting for 3D molecular design (Hoogeboom cet al. 2022: Igashov et al. 2022; Schneuing et al., 2023). However, no comparative study of zo and e-parameterization in this domain has been performed yet. To close this gap, we benchmark the 6- vs. the Zy parameterizatioo on data modalities subject to a Gaussian diffusion. That is, we treat all node features (including atomic elements, charges, and coordinates) as well as the bond features as continuous variables and optimize our diffusion model using either the e- or Zo parameterrizatonnn with the loss functions defined in Eq. (2).

---

Published as a conference paper at ICLR 2024
100
40
90
50-
EQUATE
00-
BOCATER
Pre-tmvaed
60
70-
10
AwentV
ValaggW
Rigiaare Pre toulae
20
25
75
100
25
1.00
Sabset 1%6
Figure 1: Selected evaluation metrics for EQGAT-diff trained on GEOM-Drugs subsets (25, 50, 75%) from scratch or fine-tuned. We also report the results of the pre-trained, not fine-tuned model (0%)
In the following, we abbreviate EQGAT-diff with EQGAT to keep the notation clear, depicting the diffusion type subscripted and the parameterization superscripted. Table 2 shows that the T0- parameterization (EQGAT achieves higher molecule stability on QM9 and GEOM-Drugs than the e-parameterization (EQGAT cont). The performance gap is pronounced on the GEOM-Drugs dataset, which covers a broader range of larger and more complex molecules. On this more de- manding benchmark, EQGAT cont outperforms the e-model with 90.46% molccule stability against 85.19%. The lower molecule stability for the e-model is due to the molecular graph not being ac- curately denoised during the sampling. Thus, the final edge features do preserve the valency constraints of the chemical elements.
Next, compare how the choice of categorical or Gaussian diffusion for modeling the chemical elements, charges, and edge features affects the generation performance. Recall that the noising pro- cess in the categorical diffusion perturbs the one-hot encoding of discrete features by jumping from one class to another, or staying on the same class. Alternatively, noise from a multivariate normal distribution is added the (scaled) one-hot encodings, as described Eq. (T). For both settings, the diffusion models (EQGATIO disc and EQGAT comt) are tasked with predicting the original data point To. as there no e-parameterization when employing categorical diffusion. The previous ablation has shown that data prediction is superior to noise prediction when dealing with molecular data in a continuous setting. We discover that EQGAT disc outperforms EQGAT cont in all evaluation met- rics the QM9 and GEOM -Dugs dataset shown in Table 2 Hence, employing the categorical diffusion for discrete state-space in the Zo parameterizzaionn is the preferred choice.
TRANSFERABILITY OF MOLECULAR DIFFUSION MODELS
In many molecular design scenarios, only limited amount of training data is available for desired target distribution,
1.0
e.g., in structure-basce drug design. However, 3D generative molecular diffuusion models require a lot of training data yield a high ratio of valid and novel molecules. This sec- tion investigates how well diffusion model pre-trained on a general large set of molecules transfers to a target distri- bution specified by small training set of complex molecu- lar structures. We use the PubChem3D dataset [Bolton et al. (011)) for pre-training, which consists of roughly 95.7 mil- lion compounds from the PubChem database. It includes all molecules with chemical elements H. N, O, Si, P, S, CI, Br, and with less than 50 non-hydrogen atoms and max- imum 15 rotatable bonds. The 3D structures have been computed using OpenEye's OMEGA software (Hawkins & Nicholls 2012). We train EQGAT-diff on PubChem3D on four Nvidia A100 GPUs for epoch (~ 24 hours). In- terestingly, we found that by reducing the size of molecular graphs using only implicit hydrogens, we could reduce the
Molecale stabiliy
0.4
EOGATS
0.2
EOGATE2
EQGATS
20
30
40
50
60
90
Number
Figure 2: Comparing EQGAT
Neg with EQGAT Aisnt and EQGAT cont regarding molecule stability of 600 generated molecules with an increas- ing number of atoms. Standard devi- ations plotted in shaded areas.
7

---

Published as a conference paper at ICLR 2024
Table 3: Comparison of EQGAT disc models trained for 800 epochs on GEOM-Drugs. The super- scripts 'ft' and "af" abbreviate fine-tuned and additional-features confidence level is given as subscripts. We also compare EDM and the current SOTA, MiDi. Train- ing details for MiDi are given in Appendix A.6
The margin of error for the 95%
The best results are in bold.
Datasel
GEOM-Drugs
Model
EQGAT
EQGAT
EQGAT
EQGAT
EQGAT
EDM
MiDi
Mol. Slob. T
93. 11
93.9220
94.51 1s
95.01:0.37
40.3
89.7.000
Atrm, Stuh 7
99. 19.7910.31
99.1
99.33
99.84.000
97.8
99.7..0.1
Validity 1
85.86
88.04.0.17
87.89
88.42...2
87.8
70.5-n.41
Connect. Coma..
yn 25
96.36 25
20:00
41.4
08.76_.066
Novelly 1
99.82 10 19.03
99.84 20.04
99.82 io
99.82 10
100.00
100.00 .0..00_...
Divensity f
89.05 19.05
88.98 10:02
88.9 .0.0
KL Divergence f
87.6620.31
87.58
$8.38
87.62.n.19
Similarity
0.114
0.113:c.
0.114
2.114
AtrenTV [10 3
3 3.0220.06
3.02
2.88
2.91 10.10
21.2
5.11
BondTV [30 []0""]]
2.44
2040
2.42a000
2000
4.8
2:44000
ValencyW; [10-2]
L.18
L.20
2.85
0.90
28.5
2,48-n0s2
Bondl enggtsW: [10-9)4
0.56
0.10 10.00
0.50
0.20
0.2
0.2.0030
BondAnglesW,
p.83
0.79 0.7920.0
2.65
0.62.000
6.23
1,73g088
pre-training time significantly without sacrificing performance in fine-tuning. For a comparison to keeping explicit hydrogens in the pre pre-training. see Appendix [A.3] During fine-tuning, the diffusion model is tasked to adapt to the distribution of another dataset, now including explicit hydrogens.
To evaluate the effectiveness of pre-training. we fine-tune subsets of (25, 50, 75%) of the QM9 and GEOM-Drugs datasets. Our results suggest that using a pre-trained model and subsequent fine- tuning shows consistently superior performance across datasets, partly by a large margin (see Fig. [.. We demonstrate the importance of pre-training by evaluating molecule stability, validity, and the number of connected components of a fine-tuned model compared to training from scratch on the full data and its 25, 50, 75% subsets. As reference point (0%), show the pre-trained model without fine-tuning evaluated on the aforementioned metrics. Interestingly, the fine-tuned model shares similar (best) scores with EQGAT disc trained from scratch on 100% of the data when looking at atom type variation and valency as well as angle distance metrics using a hold-out test set as a reference. These metrics capture how well the model learns the underlying data distribution.
We find that the fine-tuned model effectively learns a distribution shift on GEOM-Drugs by only being trained on small subsets of the data. We list more detailed evaluation metrics and the evalu- ation on QM9 in Appendix A.3 Comparing the fine-tuned model EQGATTo.ft with EQGAT dise> disc*
disc and EQGAT cont. respectively. shown in Fig. 2 we can also observe that the fine-tuning leads to significantly more stable predictions for larger molecules. We suspect that these findings might also apply to learning building blocks large databases like the Enamine REAL Space to bias the gen- crative model towards, e.g., higher synthesizability while ensuring an efficient distribution shift on the target distribution.
INSERTING CHEMICAL DOMAIN KNOWLEDGE
In the previous sections, we examined and outlined the importance of design choices when em- ploying diffusion models for 3D molecular generation. Taking these results, we select the best two models- with and without fine-tuning: EQGAT disc and EQGAT=0 disc and train them to full con- vergence, comparing with EDM and MiDi. We demonstrate in Tab. 3 that EQGATTUl.f and even more so EQGAT disc and EQGAT=00.f.1 outperform MiDi on all evaluation metrics by a large disc
margin, while, most notably, our models converge significantly faster and are twice as fast compu- tationally (see Appendix A.6). Given the demonstrated ability of diffusion models to learn the data distribution of complex molecular structures, we insert more chemical domain knowledge into the diffusion model, going beyond bonding. We additionally utilize aromaticity, ring correspondence, and hybridization states to provide a more comprehensive description of the molecular structure. The new additional features are independently perturbed using the categorical transition kernels (see Eq. (I) and subsequently denoised by our model. We observe that these additional chemical fea- tures again improve the performance of our models (EQGAT disc and EQGAT compared to our previous models as well as EDM and MiDi.
disc
8

---

Published as a conference paper at ICLR 2024
8 STRUCTURE-BASED DE Novo LIGAND DESIGN
We train EQGAT-diff on the Crossdocked dataset Francoeur et al. (020)) for de novo structure- ligand design. Following (Guan et al., 2023 and (Schneuing et al. 2023), we consider the protein pocket as a condition to generate novel ligands. Here, the pocket is seen as a fixed 3D context, while the ligand's s coordinates, atom and bond types get diffused and denoised. In Tab. we report the validity. number of connected components as well as the Wasserstein distances of bond lengths and angles between generated set to the training set, respectively. We observe that the finetuned model with timestep loss weighting significantly outperforms the models that are trained from scratch on all metrics. For the models trained from scratch, using timestep weighting shows better performance than no loss weighting. These results further underline the relevance of our findings allowing for an effective transfer of our model to structure- molecule generation.
Table 4: Comparison of EQGAT-diff models trained on the Crossdocked dataset for pocket- conditioned de novo ligand generation. EQGAT dise and EQGAT disc are compared with and with- out loss weighting, each trained for 300 epochs. Mean values are reported over five runs of selected evaluation metrics with the margin of error for the 95% confidence level given as subscripts and best results in bold.
Model
Validity 1
Connect, Comp. f
BondLLengths W1 [10-"14
BondAAgles W1 4
+5.51
95.15+0.14
0.20+n.0
4.37+0.20
EQGATall. (w
19.62
1o.31
0.12gao
1.12
EQG (w,,()
05.65.o.12
99.66.0.30
0.11ga.
g.55
Based on these results. we sample ligands from EQGAT disc for docking. Following Luo et al. 22221) Peng et al. (2022), we draw 100 valid ligands per protein pocket and evaluate them using Vina (Hassan et al. 2017) as an empirical proxy of the ligand binding affinity. As shown Tab.
EQGATRR.. outperforms both TargetDiff Guan et al. (023)) and DiftSBDD (CrrBBDDS Schneuing et al. (2023) on the docking score and across all other metrics while generating more diverse ligands.
Table 5: Docking performance comparison between EQGAT TargetDiff and DiffSBDD trained on the Crossdocked dataset for pocket-conditioned de novo ligand generation. Best results in bold.
disc
Model
Vina (
Vina (Top-10%) L
QED
SA
Lipinski f
Diversity
EQGATH (vo.(t)
-7.423
_9.571_a.4
0.522
0.697 10:00
4.66_
0.742
TifgeDDfff
-7.318
-9.669 ja.00
0.483
0.584
4.594 _osa
0.718
DinSBDD-eood
4.950=2.0
-9.120=2.1e
0.469
0.57820.13
4.562-0.20
0.728
9 CONCLUSIONS
In this work, we have introduced EQGAT-diff, a framework for fast and accurate end-to-end differ- entiable de novo molecule generation in 3D space, jointly predicting geometry, topology, chemical composition and optionally other chemical features like the hybridization. The findings presented here are underpinned by comprehensive ablation studies, which address previously scientific blank spot by thoroughly exploring the design space of 3D equivariant diffusion models. We have specif- ically designed an equivariant diffusion model that combines Gaussian and discrete state space diffusion. Crucially, we have incorporated timestep-dependentt loss weighting that significantly enhances the performance and training time of EQGAT-diff and, furthermore, showcased the trans- ferability of our model being pre-trained on PubChem3D on small datasets. Our proposed models have significantly surpassed the current state-of-the-ar 3D diffusion models, particularly in gener- ating larger and more complex molecules, as evidenced by their high molecule stability and validity, which evaluate that chemistry rules are preserved. Most notably, we also showcased that our frame- work seamlessly transfers to larget-conditione de novo ligand design superior docking scores while ensuring high diversity in samples. Given these achievements, we anticipate our findings will open avenues for MI-ddriven de novo structure-based drug discovery.

---

Published as a conference paper at ICLR 2024
CODE AVAILABILITY
Our source code and implementation will be released under httpss://githei com/ pfizer- openapocccuur -diff
ACKNOWLEDGGMENTS
This study was partially funded by the European Union's Horizon 2020 research and innovation program under the Marie Sklodowska- Innovative Training Network European Industrial Doc- torate grant agreement No. 956832 "Advanced machine learning for Innovative Drug Discovery."
REFERENCES
Jacob Austin, Daniel Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete spacess In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL
hittps://openreviee.
Simon Axelrod and Rafael
molecular conforma- tions for property prediction and molecular generation. Sci. Data, 9(): Apr 2022. ISSN 2052-4463. doi: 10.
Geom,
URL https:///ed http:///..oorrg///// 1038/
841597-022-01288-
Christoph Bannwarth, Sebastian Ehlert, and Stefan Grimme. Gfn2-xttb accurate and broadly parametrized self-consistent tight-binding quantum chemical method with multipole electrostatics and density-dependent dispersion contributions. J. Chem. Theory Comput., 15(3)1165-11671, doi: (0.102//acsjctt..801111 URL https:://oi..crt 1021/acs jacascssctttt jctc 8b01176
Mar 2019. ISSN 1549-9618.
llyes Batatia, David Peter Kovacs, Gregor N. Simm, Christoph Ortner, and Gabor Csanyi. MACE: Higher order equivariant message passing neural networks for fast and accurate force fields. Ad- vances in Neural Information Processing Systems, 2022. URL https: //tp::s/////peeeiieeee net/
Simon Batzner, Albert Musaclian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Ko- mbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neu- ral networks for data-efficient and accurate interatomic potentials. Nat. Commun. 13(1), may 2022. doi: 10. 111111///4447--02-222000 3F34447--02222933--
URL pattps://ee http:://dc....o..ee mreeiorrg11... forg/10.1038
Evan E. Bolton, Jie Chen, Sunghwan Kim, Lianyi Han, Siqian He, Wenyao Shi, Vahan Simonyan, Yan Sun, Paul A. Thiessen, Jiyao Wang, Bo Yu, Jian Zhang. and Stephen H. Bryant. Pub- Chem3D: a new resource for scientists. Journal of Cheminformatics, 3(1):32, September 2011. (188/1758-2946---11
ISSN 1758-2946. doi:
URL https://Gcii /ott:::://cii..r//10... 1186/
1758 1758-29445 2946-3-32
Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Velivckovi'e. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. ArXiv, abs/2104.13477 2021. URL https:
Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi S. Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. In The Eleventh International Confer- ence on Learning Representations, 2023. URL https:://ppereeeeeeee net.ieewweeeeeteeeeeeeeeeeeeeeee ftrtt//ormmicc-
kKF8 K-mBbS
Valentin De Bortoli, James Thorton, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling. In Advances in Neural Information Processing Systems, volume 34, pp. 17695-17709. Curran Associates, Inc.,
2021. URL ottpaa/eeeeeeeeeeeeeeeeeeeeeeeee httppaappeeelleeeeppee////////// htw.a//9eee.fiaaaaaaoccoooorrraa
10

---

Published as a conference paper at ICLR 2024
Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thomton,, Yee Whye Teh, and Amaud Doucet. Riemannian Score-Basee Generative Modelling. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neu- ral Information Processing Systems, volume 35, pp. 2406-2422. Curran Associates, Inc., nwwripsuru.see Paper Cww.erenneeeeecceeeeeeeeeee pdf
2022. URL
Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat GANs on image synthesis. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neu- ral Information Processing Systems, 2021. URL AANuCvzaVE
net forum?id=
Matthias Fey and Jan Lenssen. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019. URL https:
70349949
Paul G. Francoeur, Tomohide Masuda, Jocelyn Sunseri, Andrew Jia, Richard B. Iovanisci, Ian Snyder, and David R. Koes. Three-dimensional convolutional neural networks and cross- docked data set for seuucture-bbsed drug design. Journal of Chemical Information and Mod- eling, 60(9))4200 42014 Sep 2020. ISSN 1549-9596. doi: 10. 1001022/acssccim.00001111 URL
https://dos www.orgggoo.. 1021 2o21/acc.... jaim. OcC0411
Niklas Gebauer, Michael Gastegger, and Kristof Schütt. Symmetry-adppted genera- tion of 3d point sets for the targeted discovery of molecules. In H. Wallach, H. Larochelle, A. Beygelzimer, d' d'Alché-Buc, E. Fox, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL htteppss//pooceeeingg...eeeee catesscccppppppp paperr files/paper/2019) file BLL14466222600002335557122222222
ccpppprrrrrpppppppppppppppapp... pdf
Niklas W. A. Gebauer, Michael Gastegger, Stefaan S. P. Hessmann, Klaus-Robert Müller, and Kristof T. Schütt. Inverse design of 3d molecular structures with conditional generative neural networks. Nature Communications, 13(1), feb 2022. doi: 10.10388641467--02---8838 022 28526
URL
Tttps://aol..rr/1111103833 1ttp://...138225441111
Jiaqi Guan, Wesley Wei Qian, Xingang Peng, Yufeng Su, Jian Peng, and Jianzhu Ma. 3d equivariant diffusion for target-aware molecule generation and affinity prediction. The Eleventh Interna- tional Conference Learning Representations, 2023. URL /ttpps://penreevee... net/
Forum?1d=k3gg3PX55.
Nafisa M Hassan, Amr A Alhossary, Yuguang Mu, and Chee- KeoeeKKeng Kwoh. Protein- ligand blind docking using QuickVina Vina-W W with inter-process spatio-temporal integration. Sei. Rep., 151)11 November 2017.
Paul C. D. Hawkins and Anthony Nicholls. Conformer generation with omega: Learning from the data set and the analysis of failures. Journal of Chemical Information and Modeling, 52 (11):2919-2936, 2012. doi: 10. 10.1021/c1300314k URL
c1300314K PMID: 23082786.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances Neu- ral Information Processing Systems, volume 33, pp. 6840-6851. Curran Associates, Inc.,
2020. URL lbttps://proceddinngg htwp:occeeddnngg..eeiieeeee neurips...... http://dsnnn..rrr........eeeeeee fttpr//eeeee///aaeeeeee//e/e///e
file/
Paper..... .pdf
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax flows and multinomial diffusion: Learning categorical distributions. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URLbttts: htlaattss//ppeeeiieeeee wtw.eeewweeeeeeeeeeeeeeeeeeeeeee http://www..ttrrrrrraaaaaaaauapa
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffu- sion for molecule generation in 3D. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba
11

---

Published as a conference paper at ICLR 2024
Szepesvari, Gang Niu, and Sivan Sabato (eds.). Proceedings of the 39th International Confer- ence on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 8867-8887. PMLR. 17-23 Jul 2022. URL http:://ppccceeaiggggeeeo///e//e mlr press/v162/ ottp
hoogeboom22a. html
Ilia Igashov, Hannes Stärk, Clément Vignac, Victor Garcia Satorras, Pascal Frossard, Max Welling, Michael M. Bronstein, and Bruno E. Correia. Equivariant 3d-conditional diffusion models for molecular linker design. ArXiv, abs/2210.05274, 2022. URL https://acxiv..rr//aa
2210.05274
Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. International Confer- ence on Learning Representations, 2021. URL htttss://pppnrreeeeestt ew. net. forum? did= TYLJDVSx6J4
Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi S. Jaakkola. Torsional diffusion for molecular conformer generation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https: Settps://openrest http///ppppepeeeeeeeeeeeeeeeeeep net forum? -id--idddww w6f]]r62r_____
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL
https://openrevieww http
Diederik Kingma and Ruiqi Gao. Understanding diffusion objectives as the ELBO with simple data augmentation. In seieth Conference Neural Information Processing Systems, 2023. htt.ttpps//ppeeneeeee....tt www.eeeewe//rrmiieee NrMMaadyyyy
Diederik P Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. On density estimation with dif- fusion models. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Ad- vances in Neural Information Processing Systems, 2021. URL https:://///// / http:///ppeeeeeeeeeeeeeeeee / Eorum?2d=22LdggxcL.
Jonas Köhler, Leon Klein, and Frank Noe. Equivariant flows: Exact likelihood generative learn- ing for symmetric densities. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learn- ing Research, pp. 1331153700 PMLR, 13-18 Jul 2020. TRL hwwp:ttts://// htttatass//prreediingg www.ingsmmlrr..... .mlr. press/ /enss/vlinn 19ol9///ooll222
html
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. In International Conference on Learning Representations, 2021. URL
net/ htt.://wwwweeeeeeeeeeeeeeeiieeee xF
Tuan Le, Frank Noe, and Djork-Amé Clevert. Representation learning on biomolecular structures using equivariant graph attention. In The First Learning on Graphs Conference, 2022. URL
https://openreviie wttt:oorumidd
Xiang Lisa Li, John Thickstun, Ishaan Guirajani, Percy Liang, and Tatsunori Hashimoto. Diffusion- LM improves controllable text generation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https: http:///ppeeeeiieeeeeeeeeeeeeeee wwweteeeoouurree forum? 3dd=3s9Irsssssssee
Shitong Luo, Jiaqi Guan, Jianzhu Ma, and Jian Peng. A 3d generative model for structure basedd drug design. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 6229-6239. Cur- ran Associates, Inc., 2021. URL http://ww.w.....ao.oooooo/ooe...
files/
Youzhi Luo and Shuiwang Ji. An autoregressive flow model for 3d molecular geometry generation from seratch. In International Conference on Learning Representations, 2022. URL https:
12

---

Published as a conference paper at ICLR 2024
Elman Mansimov, Omar Mahmood, Seokho Kang, and Kyunghyun Cho. Molecular geom- etry prediction using a deep generative graph neural network. Scientific Reports, 9(1,, dec 2019. doi: 10.1038/441598-019-555 2Fs41598
URL https://aci.org/111 10388
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic mod- els. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Confer- ence Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp- 8162-8171. PMLR, 18-24 Jul 2021. URL https::/prroeeenngg.. mir press/v13// http:pppoooeddnnng.......es/////
ni michol21a.html
Xingang Peng, Shitong Luo, Jiaqi Guan, Qi Xie, Jian Peng, and Jianzhu Ma. Pocket2Mol: Effi- cient molecular sampling based on 3D protein pockets. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.). Proceedings of the 39th Inter- national Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 17644-17655 PMLR, 17-23 Jul 2022. URL http:::/ .mll. mmlr.
press/v162 62 https//wwwl/pgg222...u pttp:/ppgg222bbbtmmme
Xingang Peng, Jiaqi Guan, Qiang Liu, and Jianzhu Ma. MolDiff: Addressing the atom-bond in- consistency problem in 3D molecule diffusion generation. In Andreas Krause, Emma Brun- skill, Kyunghyun Cho, Barbara Engelhardr, Sivan Sabato, and Jonathan Scarlett (eds.). Pro- ceedings of the 40th International Conference on Machine Learning. volume 202 of Proceed- ings of Machine Learning Research, pp. 27611-27629, PMLR, 23-29 Jul 2023. URL https: mlr. wtt.mm.p....////eeeeeeaaea htmml
Vadim Popov, Ivan Vovk. Vladimir Gogoryan, Tasnima Sadekova, and Mikhail Kudinov. Grad- tts: A diffusion probabilistic model for text-to-speech. In Marina Mcila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8599-8608 PMLR, 18-24 Jul 2021. URL https::/proceediing mmr. wtt.llpppree//t1///pppooo html
Raghunathan Ramakrishnan, Pavlo O.
Dral, Pavlo O. Dral, Matthias Rupp, and O. Anatole von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1, 2014. URL/tttss:/ htt.:tppps/pppp..mmmttiiisssses wwwanticscciittoe
org/Corpus ID:1333777 15367821
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF
Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10684-10695, June 2022. thecvf..e.ccccccc com/ wttp///otttttttceco /ttp://oteet//////////////////// R/om//// Saangge_Snthesiit With SiwieettDifffuiioo ModelsssioiMMMooeeeeeeoeeaaeeess CVPR_ SPFR_2022_papaee html
URL https:
High Resolution
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference on Learning Representations, 2022. URL. htt.tps:://peeeeviieeee
net/Forum?id=TTeee httefforrmmioo-TTllxxiiioee
Victor Garcia Satorras, Emiel Hoogeboom, Fabian Bernd Fuchs, Ingmar Posner, and Max Welling. H(n) equivariant normalizing flows. A. Beygelzimer, Dauphin, P. Liang, and Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https:// openreview www.nrevieennettEEee www. wwr.m Nid- RowVA
Arne Schneuing, Yuanqi Du, Charles Harris, Arian Jamasb, Ilia Igashov, Weitao Du, Tom Blun- dell, Pietro Lió, Carla Gomes, Max Welling, Michael Bronstein, and Brumo Correia. Structure- based drug design with equivariant diffusion models, 2023. URL https::/arxiv.co
abs/2210.13699
Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the pre- diction of tensorial properties and molecular spectra. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceed- ings of Machine Learning Research, pp. 9377-9388, PMLR, 18 18-24/ Jul 2021. URL hetps:
7/proceed ngs. ml www.ppresssvii. htt.tttt/cccuttt2/ htm
13

---

Published as a conference paper at ICLR 2024
Gregor Simm and Jose Miguel Hemnande--Lobtto A generative model for molecular distance ge- ometry. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Con- ference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp- 8949 8958. PMLR, 13-18 Jul 2020. URL hnttps://procee
ngs. press/v119)
51 mm mmm20a... ht.ll
Jascha Sohl- Sahl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Francis Bach and David Blei (eds.). Pro- ceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings Lille, France, 07-09 Jul 2015. PMLR. URL. of Machine Learning Research, pp. hwarrs/s/33/sollcccccttteeii1_ html
htttps:/pproceednngg
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna- net.
fional Conference Learning Representations, 2021a. URL
forum? d=teiid=sttii 1gi arCHLP
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data dis- tribution. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bue, E. Fox, and R. Gameet (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https::/prooceednggs. www.igs....r.sp....... www.inssuurrrppseeeeeeeeee cc/paper_ _frles/les//
Yang Song, Jascha Sohl-Dickster Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference on Learning Representations, 2021b. URL https: httptttt//peeeeeeeeeeeee /pw.ereeeereeeeeeeeeeeeeee net..
Eorum?1d=PxTIGI120
Hannes Stärk, Octavian Ganea, Lagnajit Pattanaik, Dr.Regina Barzilay, and Tommi Jaakkola. EquiBind: Geometric deep learning for drug binding structure prediction. In Kamalika Chaud- huri, Stefanie Jegelka, Le Song. Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Pro- ceedings of the 39th International Conference on Machine Learning. volume 162 of Proceed- ings of Machine Learning Research, pp. 20503-20521. PMLR, 17-23 Jul 2022. URL https: //proceedings.. mlr wttp:/iinnn............ee. sreessvvl6e stark22b.bb...... html
Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molec-
ular potentials. International Conference on Learning Representations, 2022. URL https:
9wrRB
Clément Vignac, Nagham Osman, Laura Toni, and Pascal Frossard. Midi: Mixed graph and 3d denoising diffusion for molecule generation. In Danai Koutra, Claudia Plant, Manuel Gomez Rodriguez, Elena Baralis, and Francesco Bonchi (eds.). Machine Learning and Knowledge Discovery in Databases: Research Track European Conference, ECML PKDD 2023, Turin, Italy, September 18-22, 2023, Proceedings, Part II, volume 14170 of Lecture Notes in Com- puter Science, pp. 560-576. Springer, 2023. doi: 10. 10000/977---001144415--0000 https: http://oi..ooo.
URL
031 10100331144113 43415-0 33
Lemeng Wu, Chengyue Gong, Xingchao Liu, Mao Ye, and qiang liu. Diffusion-based molecule generation with informative prior bridges. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL
net http://weeeeeeeeteeeeeee wwwumoiddiioooooo TJUNt ittitiiiKK
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Confer- ence on Learning Representations, 2022. https:// //tp://v/////ereeewee net/ forum?id
PZevXEMZvOC.
14

---

Published as a conference paper at ICLR 2024
A
APPENDIX
A. MODEL DETAILS
Before message passing, we create time embedding L. = 500 and concatenate those to the geometric-invariant (scalar) features, including atomic elements and charges, to pass the timestep in- formation into the network. After each round of message passing, we employ normalization layer for the position updates as proposed by Vignac et al. (2023), while sealar and vector features (h, v) normalized using a Layernorm followed by update block using gated equivariant transforma- tion as proposed in the original EQGAT architecture (Le et al., 2022). After L round of message passing and update blocks,
leverage the last layers' embeddings perform the final prediction Go = (X, H. E) shown in Figure For the case that additional (gcometric) invariant features modeled, including the atomic formal charges, aromaticity, or hybridization state, the hidden node matrix H includes them as output prediction by simple concatenation, i.c., predicting more output channels.
We implement EQGAT-diff using PyTorch Geometric (Fey & Lenssen 2019) and leverage the (sparse) coordinate (COO) format that stores the molecular data and respective edge indices of the fully connected graphs.
A.1.1 MODEL. TRAINING
We optimize EQGAT-diff under parameterization utilizing Gaussian diffusion for coordinates and categorical diffusion for discrete-valued data modalities, including chemical elements and bond types.
L - w.(:) Am |x Xo Xol1 + C,H,, Ho) + C, Eo)
where CE refers to the cross-entropy loss and (Ne, An, de) (3,0.4,) are weighting coefficients adapted from Vignac et al. (2023).
In all experiments, EQGAT-diff uses 256 scalar and vector features each and 128 edge features across 12 layers of fully connected message passing. This corresponds to 12. 2.3M trainable parameters.
We train for 200 epochs on QM9 and 400 epochs on GEOM-Drugs to achieve comparability across models while ensuring computational feasibility regarding many ablation experiments. We use fewer epochs for QM9 since the diffusion models quickly overfit such that the novelty of sampled molecules decreases. This is not the case with GEOM-Drugs.
We use the AMSGrad with learning rate 2 10-4, weight-decay of 10-12 and gradient clipping for values higher than ten throughout all experiments. The weights of the final model are obtained by an exponential moving average with decay factor of 0.999.
(
(4)
(
+
W.
W b
X(
concat
W + b
MLP
X
Figure 3: Prediction module that processes EQGAT-diff embeddings to obtain the predicted data modalities. The computational graph reads from top bottom.
15

---

Published as a conference paper at ICLR 2024
On the QM9 dataset, dataloader with a batch size of 800 following (Vignac et al. 2023 All models are trained on four Nvidia A100 GPUs.
batch size of 128; the GEogsDrugs dataset, adaptive
For training, we use an adaptive noise schedule proposed by Vignac et al. 2023
(t/T + 8)
a² = cos
1+5
The respective scaling hyperparameter y was set to V4 2.5, Vy 1.5, dataset. At the same time, for GEOM-Drugs we use coordinates, atom types, bond types, and charges, respectively. This noise scheduler accounts for the various variables of graph and 3D structure not being equally informative for the model and has been found by Vignac et al. (2023) to outperform the cosine schedule (Nichol & Dhariwal 2021 Hoogeboom al. 2022 significantly.
on the QM9 2 with Vr, N,, My and Mr denoting atom
A.1.2 MODEL SAMPLING
As mentioned in Sec. 5.2 the diffusion loss term L, wDw.DD(gaaaaaaa po) l is optimized by minimizing the KL-divergence. For the case of continuous data types, i.c., coordinates, the tractable reverse distribution (Sohl-Dickstein et al. 2015 Ho et al. 2020) is
Xx, Xo) = N
(xt- [t- (Xc, xo); 11-1);
(5)
with 14- (X,, Xo) the coordinate matrix is vectorized to have shape 3N.
/a X
1-ā--1 Var(1-ā. X₁ and E,-1 T-a. 3.I, where assume that 1-a,
1-ac
Sampling from that reverse distribution is obtained through the denoising network that predicts the clean coordinate matrix to parameterize (x, (x,) g(x;- |x,, xt) and sample via
vai-15t at
&o + Val Ci +
a1-1 Br ECM
Xt
(6)
where ECM
3NN Gaussian noise vector with zero mean.
3N
For discrete variables, we obtain a tractable reverse distribution that is categorical Austin et al.
202T
q(cr-1lco. c.) C(cr--lpo-- c,)):
with probability vector defined Pt-1 c.) = e,U,T Uc, where the entry [Ur]k denotes the transition probability jump from state to j and is defined
U,
U, (1 3.) 3. a,Ix + lxē
(8)
while the cumulative product after / timesteps starting from 1 can be simplified to
Ú = U,U2 U. ā + (1 ā. dxā
(9)
We recall that the one-hot encoding of each node or edge is perturbed independently during the forward process, such that the encoding c, E {0, 1 is obtained by sampling from the categorical distribution (co) C(c, ā.Co + (1 () described in Eq. (I).
Similar to (Austin et al. 2021] Vignac et al. 2023), we obtain the reverse process for discrete data types by marginalizing the network predictions (for each node in the graph)
po(Ce-1cc)) x Σ g(c,-1/ce, ex) Co.k
(10)
where e₁ is an one-hot-encoding with at index and 6 is the k-th entry in the softmaxed probability vector čp.
16

---

Published as a conference paper at ICLR 2024
A.2 METRICS
The Wasserstein distance between valencies is given as a weighted sum over the valency distribu- tions for each atom type
Valency W Σ
p(=) (Dra(z)), Dra()))
(11)
ZE atom types
with px (ar) being the marginal distribution of atom types in the training set and D (2) the marginal distribution valencies for atoms of type the generated set and D (x)) the same distribution in the test set. For the bond lengths metric, weighted sum of the distance between bond lengths for each bond type is used
BondLenghts
Σ
p(y) (Das( Dasalv)),
(12)
hond types
where (y) is the proportion of bond of types y in the training Dais (y) is the generated distribution bond lengths for the bond of type and Daia the same distribution computed the test set. Lastly, the distribution of bond angles for each atom type is weighted sum using the proportion of each atom type in the dataset, restricted to atoms with two or more neighbors, ensuring that angles can be defined
BondAngles (generated, target)
Σ
P(x) (Dimgle(())) D Dangles(z)) (a)):
(13)
E atom types
with (x) denoting the proportion of atoms of types in the training set, and Dangles (x) the distribution of geometric angles of the form / (Fk Ti, Ty T:), where is an atom of type and
and j are neighbors of (Vignac al. 2023
A.3 RESULTS AND DETAILS
We visualize the empirical distribution of the number of atoms and the chemical composition for the QM9, GEOM-Drugs, and PubChem3D datasets in Figure 4 For PubChem3D, we show the empirical distribution for the datasets with implicit and explicit hydrogens.
0.14
PubChem3D-If
PubChem3D-II
0.7
0.12
PubChem3D-noll
PubChem3D-nofi
GEOM-Drugs
0.6
GEOM-Drugs
0.10
QM9
QM9
Frequency
Frequency
0.5
0.08
0.4
0.06
0.3
0.04
0.2
0.1
0.00
0.0
20
40
60
100
H
Si P S As Hg
Number atoms
Figure 4: Empirical distributions over QM9, GEOM-Drugs, and PubChem3D with implicit and explicit hydrogens. a) Frequency for the number of atoms. b) Frequency for atomic elements.
A.4 TIME-DEPENDENT LOSs WEIGHTING
A.4.1 LOss WEIGHTING AND FINE-TUNING
In the study in Section A.4 conducted ablation analysis to evaluate the efficacy of loss weight- ing, comparing two weighting strategies denoted as w. (t) and was across different subsets (25, 50, 75, and 100%) of the QM9 and GEOM-Drugs datasets. In Fig. 5 the truncated loss weighting is
17

---

Published as a conference paper at ICLR 2024
10
Molscatle stability
0.9
1.0
0.8
0.7
10
w,(0
0.6
0.0
400
Epoch
Figure 5: Comparison of EQGAT-diff trained with wt(t) and wws respectively, GEOM-Druge a) Uniform (w,) versus modified SNR(() loss-weighting (w,(t)). b) Unweighted prediction errors for models trained with w, (t) loss weightisss increasing timesteps. c) Comparison between Wa and w. (t) regarding molecule stability convergence during training.
0.8
0.8
Malecule stability GaroM-33rgg
0.6-
0.6
100%. wj(t) +finetune
0.1
75%, w.(t) +finetue
0.2
50%, w.(t) +finetune
Validity
100%, wy(t)
200
75%, wy(t)
50% wu(t)
1.0
1.0
100%, Wo
75%, Wy
0.9
0.9
50% Wo
25%, W.
0M6
0.8
0.7
0.6
100
150
200
250
300
50
100
150
200
250
300
Epoch
Epoch
Figure 6: Comparison of different models and data subsets for training on GEOM-Drug and QM9, respectively. The dotted, solid lines depict the fine-tuned model using w. (f))-weighting The solid lines show the model using w. (fe-weigt and the dashed lines show the model trained without loss-weighting. While training, after every 20 epochs, 1000 sampled molecules evaluated molecule stability and validity.
depicted (left) besides the effect on the loss for lower timesteps illustrating the unweighted loss over time steps for a batch of 128 molecules, where the model trained with w. achieves lower prediction error for steps closer to 1 (middle) and the effect the molecule stability while training showing better performance and faster training convergence for molecule stability when using w, (t) (right).
As illustrated in Figure 6 applying loss weighting using (.(() consistently results in performance enhancements for the model. These enhancements are characterized by accelerated training conver- gence, leading to improved molecule stability and validity, even when the model operates on smaller subsets of the data. Notably, the case GEOM-Drugs, when trained with only 25% of the data and optimized with wt(t) (indicated by the yellow solid line), the model exhibits convergence be- havior similar to that of the model trained on 100% of the data with uniform weighting w., (indicated by the yellow dashed line). Furthermore, fine-tuning leads to superior performance (dotted, solid lines). After just 20 epochs of fine-tuning and only using 25% the data, the model already out- performs all its counterparts on molecule stability and validity even when they are trained 100% of the data and holds for both the GEOM-Drugs (first row) and QM9 (second row) datasets. Our findings, summarized in Table underscore the critical role of loss weighting using w, (t) in the training of diffusion models for molecular data and also highlight the importance of pre-training especially when the target distributions small and do not contain many data points.
18

---

Published as a conference paper at ICLR 2024
Table 6: Comparison of EQGAT-diff on QM9 and GEOM-Drugs trained on subsets of 25, 50 and 75% of the data. We report the mean values over five runs of Molecular Stability (Mol. Stability), Validity, and the number of Connected Components (Connect. Comp.) for training from scratch with and without modified SNR(() weighting and compare it with the performance of the fine-tuned model (SNR(()+fie-tunn) The best results are written in bold, and results with overlapping margins of errors are underlined. The margin of error for the 95% confidence level is given subscripts.
QM9
GEOM-Drugs
Subset
Mol. Stability
Validity
Connect. Comp.
Mol. Stability
Validity
Connect. Comp.
25%
96, 96.01+0.22
96.68+0.22
99.59=0.06
74. 7..12-0.29
51.32+0.35
68.88+0.26
50%
96.84.0.16
97. 97.45+0.15
10.75.... 03
85.20
61. 61.1920.
82.76.0.20
75%
16.19
96.8310.17
99. go.00
87.08_o.35
74.2710.20
88. 10.69
100%
47.39
97. 97.99
99. 70.
87.59.0.19
71.44g0.22
86.57
25%
97.34+0.10 18
97.777+0.00
99.81=0.0s
88.39
75.41+0.46
85.35+0..1
50%
98.3220.11
98.65.2..07
99.93
89.41
77.21
89.43
(t(())
75%
98.45
98.77 1.004
99. 9..33 10.93
91.88 2.220
82.77.1.:16
93.39
100%
98.6820.11 11
98.962.0.07
99.944.0
91.66_0.14
84.02go.11
95.08go.11
25%
99.00 18
99.244+0.10
99.96
90.82_0.07
83.01
90.77
50%
99.21 +0.00
99.41
99,96=0.01
91.24-0.62
83.83
94.66++0.77
at(+)++ fine-tune
75%
38.79 10
99. 99.1220.12
99. 00.955o.00 03
=2.97=0.15
86.51
95.92=0.14
100%
98.94_...77
99.28810.08
99.95_..02
93.19__0.07
86.83
96.31
0.0012
Samples
PubChem3D-noil
0.0010
GEOM-Drugs
Frequency
0. 0.0008
0.0006
0.0004
0.0002
0.0000
-5000
-1000
-3000
-2000
-1000
1000
Energy (Eh)
Figure 7: Comparison of the energy distributions calculated using XTB-GFN2 Bannwarth et al. (2019) for the GEOM-Drugs training dataset against the energies of sampled molecules. We also provide the energy distribution of PubChem3D (implicit hydrogens) to showcase the distribution shift; for quantum physics-based software, those molecules appear to be radicals, and hence, the energy distribution is shifted towards high energies. Nevertheless, the model effectively has to do this shift while fine-tuning.
A.5 PRE-TRAINING ON PUBCHEM3D
To emphasize more the capability of the model to learn the underlying data distribution, we follow (Hoogeboom al. 2022) and plot the distribution of energies for sampled molecules of model trained on GEOM- - Drugs against the energy distribution of the training dataset, as shown in Fig. 7 We observe that EQGAT-diff learns the training distribution well, showing a high overlap. Further- more, to highlight the shift in physical space the diffusion model has to perform while fine-tuning, we also report the energy distribution of PubChem3D with implicit hydrogens. All energies were calculated using the semi-empirical xTB-GFN2 software (Bannwarth et al. 2019).
We also pre-trained model on the PubChem3D dataset with explicit hydrogens. Interestingly, as shown in Tab. 7 we see decrease in performance for the model that is fine-tuned the pre- training with explicit hydrogens compared to the model using implicit hydrogens, even though pre- training with explicit hydrogens takes almost three times long. We suspect that when using explicit hydrogens in pre-training, the model overfits too much on the PubChem3D data distribution, having a more challenging time transferring to the GEOM-Drugh distribution.
We subsampled IM molecules from PubChem3D and GEOM- -Duugs and enumerated over bonds of selected pair atoms including carbon, hydrogen, nitrogen and oxygen atoms. We computed distances
19

---

Published as a conference paper at ICLR 2024
Table Comparison EQGAT-diff pre-trained with or without explicit hydrogens on PubChem3D and fine-tuned on GEOM-Drugh for 400 epochs. We report the mean values over five runs of selected evaluation metrics with the margin of error for the 95% confidence level given as subscripts. The best results are in bold.
Pretraining
Mol. Stab.
Validity 1
Connect. Comp. f
PubChem3D-noH PubChem3D-H
93.19_0.07
86.8340.20
96.31
92.70=0.09
85.46
94.78+0.19
250
50
200
150
20
50
100
4
10-
25
0
1.0
1.1
1.0
1.1
0.95
1.00
1.05
1.10
1.50
1.75
2.00
2.25
[A]
Distance H-N
Distance [A]
Distance C-S
10
PubChem3d
Geom Drugs
Frequency
1.
30
20
2-
0
1.2
1.4
1.6
1.0
1.2
1.4
1.6
1.25
1.50
1.75
1.50
1.75
2,00
2.25
C-O
(A]
Figure 8: Selected atom-pair distance distribution on PubChem3D and GEOM-Drugs.
and noticed that the hydrogen-oxygen distance distribution in PubChem3D seems to have a smaller variance than GEOM-Drugs in the last panel of Figure 8
A.6 EQGAT-DIFF VS MIDI
found EQGAT-diff outperforming MiDi by a large margin across both datasets, QM9 and GEOM-Drugs, and all metrics. In Fig. g we underpin this observation by comparing training curves of EQGAT=
dise and the MiDi model, observing that our model not only outperforms MiDi on molecule stability, validity and in the adaptation to the underlying data distribution, but also converges significantly faster. EQGATED disc converges to SOTA performance already after 150-200 epochs, while MiDi needs roughly 700 epochs weakly indicating convergence but to lower values.
Furthermore, EQGAT-diff' needs ~5 minutes per epoch using four Nvidia A100 GPUs, adaptive dataloading (taken from the MiDi code based on pyg wowdeorer..... Collater) with batch size of 200 per GPU. In contrast, MiDi takes 12 minutes, so EQGAT-diff is more than twice as fast.
For training MiDi, we used the official codebase on GitHub Tn and the given hyperparameter settings but trained on four Nvidia A100 GPUs (instead of two). As seen in TabBaand and shown here in Fig. (9, we could not reproduce the results reported in the paper. We also re-evaluated the checkpoint given on GitHub and again could not confirm the reported results.
A.7 EQGAT-DIFF WITH IMPLICIT HYDROGENS ON GEOM-DRUCH
We trained EQGAT-diff GEOM-Drugs with implicit hydrogens. To this end, pre-process the GEOM-Drugs dataset using the RDKit and remove hydrogens from molecule object mol using the Chem. RemoveHs function, with subsequent kekulization Chem. Kekulize, We list the evaluation results of models EQGAT disc and EQGAT cone in Table below. We discover that the Gaussian and categorical diffusion for the Zo parameterization achieves similar performance
20

---

Published as a conference paper at ICLR 2024
Molncule stablitty
0.0
0.6
Validity
0.6
0.4
0.4
EQGAT Usc
0.2
0.2
MIDI
200
400
100
100
200
400
600
800
0.150
12.5
AtomeTV
0.125
AnglesW1
1..0
0.100
7.5
0.050
0.025
200
400
600
800
200
4.00
600
500
Epoch
Epoch
Figure 9: Comparison between EQGAT-diff and MiDi for training on GEOM-Drugs. We com- pare both models regarding a) molecular stability, b) validity, c) AtomsTV, and d) AnglesW1 while training by sampling 1000 molecules every 20 epochs over 800 epochs of training. For molecular stability and validity, higher is better; for At and AnglesW1, lower better.
related validity and connected components. the same time, the Wasserstein- distance on the histograms for empirical bond angles is lower for the generated set from EQGAT disc to the
histogram the reference set.
Table Comparison of EQGAT-diff with implicit hydrogens on GEOM-Drugs for 400 epochs. We report the mean values over five runs of selected evaluation metrics with the margin of error for the 95% confidence level given as subscripts. The best results are bold.
Model
Validity
Connect. Comp. 1
BondLengths WI BondAngles W1 1
EQGAT disc
98.29
98.90
0.59 20.62
0.44
EQGATE comt
98.48 40.14
98.36
1.34
0.56
Table 9: Classifier- Classifierguiddance EQGAT-diff to shift the reverse sampling towards low high po- larizability values. We report the mean polarizability values of sampled molecules with standard deviations subscripts.
Guidance
Polarizability
Minimization
195. 195.1944.9
Maximization 400.21
A.8 CLASSIFIER GUIDANCE FOR CONDITIONAL MOLECULE DESIGN
For conditional molecule design, we can use a trained unconditional EQGAT-diff model together with classifier- classifier-guidanee, guidance, as proposed in (Dhariwal & Nichol] 2021). to steer the generation of sam- ples using the gradient of an external classifier/regressor during the reverse sampling trajectory from noise data. As a proof of concept, mizing for low/high polarizability showing promising results. For this, we trained polarizability regressor model and applied it to the reverse sampling of unconditional EQGAT-diff model test- ing for guiding towards low and high polarizability values, respectively. Afterwards we re-calculate the polarizability of all sampled molecules for both cases and compared the mean values. In Tab.
explored classífier-uiddane to generate molecules opti-
21

---

Published as a conference paper at ICLR 2024
9 we summarize the results. The mean value of the GEOM -rogg training dataset is 245.9 + 41.9. Hence, we see that we can successfully push the distribution of sampled molecules in the respective direction.
A.9 COMPARISON TO MOLDIFF
We compare against MolDiff Peng et al. (202)) by utilizing their evaluation pipeline that includes gost-processing steps on the generated molecules to potentially fix valency and aromaticity issues when parsing into the RDKit. Selecting the 5 x 10, 000 generated samples from our best performing model EQGAT we report mean validity and mean success rate in Table 10 As shown, our proposed pest-performing EQGAT-diff model achieves superior performance over MolDiff in generating chemically valid molecules but has lower diversity, which believe is caused by longer training time on our side. However, we believe the generative model should be able to faithfully sample molecules that satisfy valency constraints, it was also trained in such data, Suppose do not employ the post-pron post-processing scheme from MolDiff and determine the validity by parsing the generated molecule into RDKit's sanitization pipeline. that case, EQGAT-diff obtains mean validity of 0.916 and a mean success rate of 0.887. This shows that the post- post-processing applied in MolDiff substantially impacts model evaluation.
disc
Table 10: Evaluation metrics from EQGAT-diff against MolDiff.
Model
EQGAT-dift
MolDiff
Validity
0.998
0.947
Connectivity
0.968
0.908
Suec. Rate
0.966
0.860
Novelty 1.000
1,000
Uniqueness
1.000
1.000
Diversity 0.320
0.422
A.10 IMPROVED SAMPLING TIME
We experimented with the DDIM Song et al. (021a) sampling algorithm known for enhancing inference/sampling time in diffusion models trained via the standard DDPM procedure in image processing. The difference between DDIM and DDPM lies in the sampling algorithm, which we believe could also be applied in our molecular data setting. However, our best-performing scenario utilizes the Ho parameterization to preserve the correct data modalities for coordinates, atom, and bond features. Hence, applying DDIM directly to discrete-values data modalities is not straightfor- ward. We restricted DDIM to continuous coordinate updates, while discrete-valuee data modalities follow the approach outlined by (Austin et al., (02T) and explained in our Appendix in Eq. (00) Table I1 compares the evaluation performance of our base models when generating samples using DDIM or DDPM for varying numbers of reverse sampling steps 500, 250, 167. Given that all mod- els underwent training with T 500 discretized timesteps, we conducted DDIM sampling every or 3 steps of the reversed trajectory starting from index 500. Notably, we observed that employing DDIM did not enhance the quality of molecule generation with fewer sampling steps (250 166) compared to the 500 steps the models were trained on.
Another way enhance sampling time, is train the diffusion models with less discretized timesteps. We performed additional experiments and trained EQGAT disc. with T = 100 timesteps using the uniform and truncated SNR(() loss weighting. The rationale behind these experiments is to assess how the reduced number of timesteps affects performance while enabling faster inference time. We compare against the two corresponding models trained with serve that the model trained with truncated SNR(t) loss weighting over T = 100 timesteps performs better than the model trained with T 500 timesteps but uniform loss weighting as illustrated in Figure TO] This result clearly speaks for the usage of the proposed loss weighting while also achieving a diffusion model that has a faster sampling time using 100 reverse sampling steps only, i.e around 5x faster sampling time. Comparing the two models trained with truncated SNR(1) loss weighting, we notice that the model trained with T = 500 discretized steps still performs better than the identical model but trained with T 100 timesteps.
500 timesteps and ob-
22

---

Published as a conference paper at ICLR 2024
Table 11: Sampling results of trained models for DDPM and DDIM for 500 Molecules.
Model
Steps
Sampling
Runtime
Mol. Stability
Validity
AnglesW1
(NR())
500
DDPM
26min
0.9160
0,8100
0.83
Continuous-SNEn SNR(t)
500
DDPM
27min
0.8920
0.7600
0.90
Unifrre
500
DDPM
26min
0.8600
0.5960
1,36
Dnscrete-SNR(I)
250
DDIM
13min
0.6580
0.6260
2.26
Continuous SNR(t)
250
DDIM
13min
0,5680
0,3920
3,95
Discrete-Uniform
250
DDIM
13min
0.5400
0,3880
3.21
Continuous SNR())
250
DDPM
13min
0.5160
0.3400
4.74
Discrete-SNR(1)
250
DDPM
13min
0.4860
0.4600
4.60
Discrete-Unine Uniform
250
DDPM
14min
0.2620
0.1940
7.60
Discrete Discrete-SNR()
166
DDIM
9min
0.1980
0.2280
5.58
Discrete-Uniform
166
DDIM
9min
0.1240
0.1080
8.24
Continuous (SN()))
166
DDPM
Smin
0.1000
0.0380
13.68
Continuous-SNN (t)(())
166
DDIM
9min
0.0900
0.0520
14.11
Discrete-SNR(c)
166
DDPM
9min
0.0660
0.5200
12.21
Discrete Unieoo-
166
DDPM
9min
0.0200
0.0120
14.83
1.0
0.9
Molecule stability
0.8
0.7
0.6
Wt(t), T 500
0.5
Wg(t), 100
0.4
W.. T= 500
W.. 100
0.3
100
150
200
250
300
Epochs
Figure 10: Molecule stability learning curves for diffusion models trained with T 500 and 100 discrete timesteps. Again, we observe that the truncated SNR(t) loss weighting w,, (t) greatly improves performance.
23


Reviews:
[{"review_id": "bPn2Mqndqy", "reviewer": "Reviewer_XyXs", "summary": "This paper explores the design space of equivariant diffusion-based generative models for de novo 3D molecule generation, including various parameterizations, loss weightings and data modalities. The authors then introduce EQGAT-Diff, which can achieve SOTA results in shorter training time and with less trainable parameters.", "questions": "N/A", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This is an empirical paper. The experiments part is solid: various metics are computed, many ablation studies are performed and standard deviations are reported.  The explored design is helpful for future model design in the community of ML + chemistry.", "weaknesses": "- The novelty is limited. Most explored designs are easy to think about, like time-dependent loss weight, modeling discrete or continuous atom/bond types, parameterizing to match noise or x0, etc. The authors didn't come up with new designs. \n- It would be better to summarize useful designs / helpful findings clearly somewhere. \n- One important baseline MolDiff [1] is not compared with, although it has been cited. \n\n[1] Peng, X., Guan, J., Liu, Q., & Ma, J. MolDiff: Addressing the Atom-Bond Inconsistency Problem in 3D Molecule Diffusion Generation. ICML 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "GE9GtXRABU", "reviewer": "Reviewer_AUdU", "summary": "This paper describes EQGAT-diff, an improvement to the EQGAT architecture (a 3D small-molecule generator) that includes diffusion. The paper explores a couple of design choices in the formulation of the diffusion part of this updated method, such as the addition or not of the truncated SNR loss term, and the inclusion of a discrete loss term for the molecular graph construction. The paper shows that their improved model is better than MiDi, the method that introduced diffusion on the molecular graphs jointly with the coordinate reconstruction.", "questions": "The paper uses E3 symmetries which includes the reflection group and can therefore mix chiralities, in principle changing the input distribution of molecular graphs.  The authors seem to not address the possible distribution shifts due to chiral transformations.\n\nSmall molecules with 3D coordinates often have specific chirality.  Although this chirality is specified in the full set of 3D positions, the authors model is using E3 symmetries in its core so I wonder if it will mix chiral molecule in ways that would result in a distribution shift.  Have the authors checked the chirality of the generated molecular distributions by keeping say only one kind of stereo isomer and validating that their model would construct both? Would implementing the SE3 group allow the authors to also including chiral flags on the bonds, in the way that shown in the typical 2D representations of molecular graphs?\n\nIt was not clear to me why the inclusion of the eps-parameterization does not use a composite loss as did the x0-parameterization.  Is it not a priori clear a priori that the addition of a categorical loss for categorical data will outperform a mean squared error on categorical data?  I may be misunderstanding something simple here, so perhaps you can explain here and possibly reformulate the early parts of section 5.\n\nThe lack of the reproduction of the MiDi code benchmarks that the authors report in the supplement is interesting, though it probably would not change the results in the paper in a qualitative way.  Did the authors contact the authors of MiDi to resolve this issue and did they receive any feedback on that question? \n\nThe observation by the authors regarding deteriorating performance upon pretraining on explicit hydrogen on a larger molecular set is an interesting point that gets buried in the supplement. Is there a generalizable lesson and perhaps a better/different way to overcome this problem than dropping the hydrogens, or is this observation too close to the level of noise?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The method is a significant advance compared to the EQGAT method. The exploration of the design choices is reasonable and of some interest, even though the results are not surprising based on the existing literature or simple logic. \n\nThe introduction and method is well written and clear. Aspects of sections 5 and 6 could improve, in particular the figures.\n\nThe datasets that the authors use are standard in this line of study, and they have well know limitations, and it is surprising to me that the community has not yet moved to better data, which should be reasonably easy as QM calculations are not that hard to perform.\n\nThe overall significance of this work appears more specialized than a usual contribution to ICLR, however, the paper would be of some interest to the growing number of drug discovery researchers that apply ML to small molecules.", "weaknesses": "The novelty of the approach is more specialized than a typical ICLR as it is more of a collection of good ideas from a narrow field of study. The authors didn't discuss or try to draw ideas from recent efforts in generative diffusion models for images or for protein structure prediction that may have strong analogies and lessons to the study of small molecules as well.\n\nThe fonts in all the figure labels are unreadably (unpublishably?) small. If space is the main concern, then there certainly exist small chunks of text (e.g. the last paragraph of section 4) and some figure panels (e.g. most from Fig 2) that could move to the supplement, and some parts that could be rewritten concisely (e.g. intro to sec 5, which could also be clarified as it was not clear what \"these three aspects\" meant in paragraph 2, top of page 5).\n\n\nThe dataset is a weakness for this paper and most other papers in this area.  I disagree with the premise of the authors that the availability of molecular data is not as abundant, though I agree that the community has been stuck on trivial and no longer meaningful datasets (it feels similar to ML in vision prior to imagenet). The authors mention the catalogue of Enamine Real in the supplement which scales to the 10s of billions of molecular graphs (albeit limited in complexity, but there are also computational enumerations of valid molecules that can be much bigger), the published patent literature probably scales to the multiple 10s of millions.  QM methods of relatively reasonable accuracy exist that are not exceedingly expensive and datasets of 10s of millions of published 3D structures with simpler methods also exist.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "xe8VBjHC4h", "reviewer": "Reviewer_sxkU", "summary": "In this work, the authors investigate the design space of equivariant diffusion generative model for de novo molecular design. The work applies EQGAT-diff, a modification from EQ-GAT as the score network. It mainly explores three aspects of equiavriant diffusion model: 1) fixed vs. time-dependent loss weighting 2) $\\sigma$ vs. $x_0$ parametrization 3) discrete vs. continuous diffusion. Experiments on GEOM-QM9 and GEOM-DRUGS show a recipe for better de novo molecular design with equivariant diffusion model. By navigating through the design spaces, the model also achieves SOTA performance than previous baselines.", "questions": "1. Will different settings (like discrete vs continuous) lead to significantly different training time till convergence as well as inference time?\n2. In section 6, the authors use OMEGA to generate ground truth. However, the conformers from OMEGA may be inaccurate. Did the authors consider pretrain on DFT datasets like PCQM4Mv2 (>3M data)? Though it's much smaller than PubChem3D, the model may benefit from higher data quality. \n3. Following the previous question, Fig. 2 shows that with sufficient training data, the advantage of pretrained model is trivial. In fact, in some metrics, the pretrained model is doing slightly worse than training from scratch. What could be the reason? Could it be related to OMEGA generated pretraining data?\n4. One question about the big scope of this work is what the benefits of direct de novo 3D molecule generative model vs. two-stage model where first generate SMILES/graphs and then predict the conformers? I can see the advantages for structure-conditioned molecule design. However, this paper mostly investigate unconditional 3D molecule generation. \n5. Torsional diffusion (https://arxiv.org/abs/2206.01729) introduces a GEOM-XL dataset including large molecules (>100 atoms). It may also be interesting to see how the proposed method transfer to such OOD datasets.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The work is well-motivated. It is valuable to investigate the design space of equivariant diffusion model and provide a general recipe. \n2. The proposed EQGAT-diff achieves superior performance through navigating through the design space. \n3. The paper is well-written and easy to follow.", "weaknesses": "1. The work proposes to navigate through the design space of equivariant diffusion model for de novo 3D molecule design. However, there are still other design choices that are omitted, e.g., different diffusion kernel and backbones. I know it's infeasible to exhaust the whole design space in one paper. I want to put a note here about the possible extensions.\n2. It seems why EQ-GAT is chosen to model score network is missing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "hihUDkaiby", "reviewer": "Reviewer_WVdk", "summary": "The manuscript introduces the EQGAT-diff model, an adaptation of the previously proposed EQGAT network, to model denoising diffusion generation. While there isn't a significant technical innovation, the model does exhibit enhanced empirical performance based on the provided metrics. The primary focus of the model is the unconditional semi de novo generation of 3D molecular structures. Notably, the relevance of generating valence-correct molecular structures without specific conditions is not clear. The authors highlight potential applications in structure-based drug design, yet the model's current format doesn't seem optimized for this purpose. Additionally, there is no mention of code availability. The presentation and writing of the article require improvement. In its current state, the manuscript lacks a meaningful technical contribution and seems to address an ill-defined task. I suggest that the authors refine their model for structure-based drug design and consider resubmitting for a future conference.", "questions": "It would be beneficial for the authors to compare their 3D molecule generation model with a 2D molecule generation, specifically using OpenEye’s OMEGA, which is cited as the ground truth for conformation generation. Would the EQGAT-diff model maintain its efficiency advantage in such a comparison?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 1, "strengths": "## Empirical Performance\nThe model demonstrates superior performance based on the evaluation metrics presented.", "weaknesses": "## Task Relevance\nThe de novo generation of 3D molecular structures, as currently presented, appears to lack inherent value.\n\n## Lack of Technical Novelty\nThe work does not introduce new technical advancements.\n\n## Code Unavailability\nAbsence of open-source code hinders replication and comparison efforts.\n\n## Writing Quality\nThe referencing is inconsistent, with many citations missing journal or conference details. There are incorrect uses of quotation marks, and figure clarity is compromised, diminishing the manuscript's overall professionalism.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (poster)