3D molecule generation by denoising voxel grids
Authors: Pedro O. Pinheiro
Joshua Rackers
joseph Kleinhenz
Michael Maser
Omar Mahmood
Andrew Martin Watkins
Stephen Ra
Vishnu Sresht
Saeed Saremi
Venue: NeurIPS 2023 poster (2023)
Keywords: generative model
molecule generation
drug discovery
TL;DR: New score-based approach to generate 3D molecules represented as voxel grids. The model is competitive with current state of the art while been simpler to train and faster to sample drug-like molecules.

Abstract: We propose a new score-based approach to generate 3D molecules represented as atomic densities on regular grids.
First, we train a denoising neural network that learns to map from a smooth distribution of noisy molecules to the distribution of real molecules.
Then, we follow the _neural empirical Bayes_ framework [Saremi and Hyvarinen, 2019] and generate molecules in two steps: (i) sample noisy density grids from a smooth distribution via underdamped Langevin Markov chain Monte Carlo, and (ii) recover the "clean" molecule by denoising the noisy grid with a single step.
Our method, _VoxMol_, generates molecules in a fundamentally different way than the current state of the art (ie, diffusion models applied to atom point clouds). It differs in terms of the data representation, the noise model, the network architecture and the generative modeling algorithm.
Our experiments show that VoxMol captures the distribution of drug-like molecules better than state of the art, while being faster to generate samples.

3D molecule generation by denoising voxel grids
Pedro O. Pinheiro, Joshua Rackers, Joseph Kleinhenz, Michael Maser, Omar Mahmood, Andrew Martin Watkins, Stephen Ra, Vishnu Sresht, Saeed Saremi
Prescient Design, Genentech
Abstract
We propose score-based approach to generate 3D molecules represented atomic densities on regular grids. First, we train a denoising neural network that learns map from smooth distribution of noisy molecules to the distribution real molecules. Then, we follow the neural empirical Bayes framework [1] and generate molecules in steps: (i) sample noisy density grids from a smooth distribution via underdamped Langevin Markov chain Monte Carlo, and (ii) recover the "clean" molecule by denoising the noisy grid with single step. Our method, VoxMol, gener- ates molecules in a fundamentally different way than the current state of the art (i.e., diffusion models applied atom point clouds). differs terms of the data rep- resentation, the noise model, the network architecture and the generative modeling algorithm. Our experiments show that VoxMol captures the distribution of drug-like molecules better than state of the art, while being faster to generate samples.
Introduction
Finding novel molecules with desired properties is an problemm in chemistry with applications many scientific domains. In drug discovery in particular, standard computational approaches perform some sort of local search-by scoring and ranking molecules around a region of the molecular space (chosen based on some prior domain knowledge). The space of possible drug-like molecules is prohibitively large (it scales exponentially with the molecular size [2, 3], estimated to be around 10 [4]), therefore search in this space is very hard. Search-based approaches achieve some successes in practice, but have some severe limitations: we can only explore very small portions of the molecular space (on the order of billions trillions molecules) and these approaches cannot propose new molecules conditioned on some desiderata.
Generative models for molecules have been proposed overcome these limitations and explore the molecular space more efficiently
[5]. These approaches often consider one of the following types of molecule representations: (i) one-dimensional sequences such as SMILES [6] or SELFIES [7] (e.g., [8, 9, 10]) (ii) two-dimensional molecular graphs, where nodes represent atoms or molecular substructures and edges represent bonds between them (e.g., [11, 12, 13, 14]), or (iii) atoms as three-dimensional points in space. Molecules are entities laying on three-dimensional space, therefore 3D representations are arguably the most complete ones they contain information about atom types, their bonds and the molecular conformation.
Recent generative models consider molecules as set of points in 3D Euclidean space and apply diffusion models on them [15, 16, 17, 18, 19, 20). Point-cloud representations allow us to use equivariant graph neural networks [21, 22, 23, 24, 25] known be very effective in molecular discriminative tasks-ar as the diffusion model's ' denoising network. However, point-based diffusion approaches have some limitations when it comes to generative modeling. First, the number of atoms in the molecule (i.e., nodes on the 3D graph) to be diffused need to be known beforehand. Second, atom types and their coordinates have very different distributions (categorical and continuous variables, respectively) and are treated separately. Because score function is undefined on discrete distributions,
37th Conference on Neural Information Processing Systems (NeurIPS 2023).

---

Figure 1: Voxelized molecules generated by our model and their corresponding molecular graphs. Left, samples from model trained on QM9 dataset (323 voxels). Right, samples from model trained on GEOM-drugs (64 voxels). In both cases, each voxel is cubic grid with side length of .25A. Each color represents different atom (and different channel on the voxel grid). Best seen digital version. See appendix for more generated samples.
some workaround necessary. Finally, graph networks operate only on nodes and edges (single and pairwise iterations, respectively). Therefore, capturing long-range dependencies over multiple atoms (nodes) can become difficult as the number of atoms increases. This is related to the limitations of the message-paseee passing formalism in graph neural networks [26]. Higher-order message passing can alleviate this problem to a degree [27, 28], but they come at a significant computational cost and they have been limited to third-order models [29] (see next section for more discussions on the tradcoffs between model expressivity and built-in equivariance).
In this work we introduce VoxMol, a new score- boore-baed method to generate 3D molecules. Similar to [33], and unlike most recent approaches, we represent atoms as continuous (Gaussian-like) densities and molecules a discretization of 3D space on voxel (i.e. a discrete unit of volume) grids. Voxelized representations allow us to use the same type of denoising architectures used in computer vision. These neural networks- the workhorse behind the success of acore-based generative models on images, e.g. [34, 35, 36e-are very effective and scale very well with data.
We start by training neural network to denoise noisy voxelized molecules. Noisy samples are created simply by adding Gaussian noise (with fixed identity covariance matrix scaled by a large noise level) to each voxel in the molecular grid. This denoising network also parametrizes the score function of the smooth/noisy distribution. Note that in contrast to diffusion models, the noise process we use here does not displace atoms. Then, we leverage the (learned) denoising network and generate molecules in two steps [1]: (i) (walk) sample noisy density grids from the smooth distribution via Langevin Markov chain Monte Carlo (MCMC), and (ii) (jump) recover "clean" molecules by denoising the noisy grid. This sampling scheme, referred to as walk-jump sampling in [1], has been successfully applied before to 2D natural images (37 and 1D amino acid sequences [39].
Compared to point-cloud diffusion models, VoxMol is simpler train, it does not require knowing the number of atoms beforehand, and it does not treat features as different distributions (continuous, categorical and ordinal for coordinates, atom types and formal charge)- we only use the "raw" voxelized molecule. Moreover, due to its expressive network architecture, our method scales better large, drug-siz sized molecules. Figure/ (and Figures 9 on appendix) illustrates voxelized molecules and their corresponding molecular graphs generated by our model, trained on two different datasets. These samples show visually that our model learns valences of atoms and symmetries of molecules.
The main contributions of this work can be summarized as follows. We present VoxMol, a new score-based method for 3D molecule generation. The proposed method differs from current approaches-usuall diffusion models on point clouds-in terms of the data representation, the noise model, the network architecture, and the generative modeling algorithm. We show in experiments that VoxMol performs slightly worse than state of the art on small dataset (QM9 [40]). while outperforms it (by large margin) on a challenging, more realistic drug-liks like molecules dataset (GEOM-Goe -drugs [41]).
'On this topic and since our approach in 3D molecule generation is based in computer vision, we should also highlight here the great success of vision transformer models [30], where in contrast to convolutional neural networks [31] no notion of equivariance whatsoever is built into the model architecture. Quite remarkably, after training. vision transformers can achieve higher degree of equivariance compared to convolutional architectures [32].

---

2
can achieve competitive Related Work
Voxel-based unconditional 3D molecule generation. Skalic et al. [42] and Ragoza el al. [33] map atomic densities on 3D regular grids and train VAEs [43] using 3D convolutional networks to generated voxelized molccules. To recover atomic coordinates from the generated voxel grids", [33] introduces a simple optimization-
solution, while [42] trains another model that "translates" voxel structures into SMILES strings. Voxel representations are flexible and can trivially be applied to related problems with different data modalities. For instance, [44] proposes a GAN [45] on voxelized electron densities, while [46] leverages voxelized 3D pharmacophore features train pocket-conditional model. Similar to these works, our model also relies on discretization of 3D space. Like [33], we use a simple peak de- tection algorithm extract atomic coordinates from the generated voxel grids. However, our method dif- fers on the underlying generative modeling, architecture, datasets, input representations and evaluations.
Point cloud-based unconditional generation. Most recent models treat molecules as sets of points, where each node is associated with a particular atom type, its coordinates and potentially extra information like formal charge. Different modeling approaches have been proposed, e.g., [47, 48, 49] utilize autoregressive models to iteratively sample atoms, and [50, 51J use normalizing flows [52]. Hoogeboom et al. [15] proposes E(3) Equivariant Diffusion Models (EDM), diffusion [S3]-b approach that performs considerably better than previous models on this task. EDMs learn to denoise a diffusion process (operating on both continuous and categorical data) and generate molecules by iteratively applying the denoising network on an initial noise. Several works have been proposed on the top of EDM [54, 55, 20, 56]. For instance, Xu et al. [56] improves EDM by applying diffusion on a latent space instead of the atomic coordinates, while MiDi [20] shows that EDM results can be improved by jointly generating the 3D conformation and the connectivity graph of molecules (in this setting, the model has access to both the 3D structure and the 2D connectivity graph).
Conditional 3D molecule generation. A related body of work is concerned with conditional generation. In many cases, conditional generation is built on the top of unconditional generation methods. Some authors propose to predict the 3D structure of the molecules given a molecular graph (this is called the conformer generation task): VAEs [57, 58], normalizing flows [59], reinforcement learning [60], optimal transport [61], autoregressive models [62] and diffusion models [63, 16, 64] have been proposed to this task. Some work [65, 66] condition 3D generation on shape while other works condition molecule generation on other structures. For instance, [17, 18, 19, 67] adapt (unconditional) diffusion models to condition on protein pockets, while [68] adapt their previous work [33] to condition voxelized structures to protein targets. Finally, [46] proposes a hybrid conditional generation model by modeling fragments/seaffolds with point cloud representation and the 3D target structures and pharmacophores features [69] with voxel grids.
Comparison between voxel and point-cloud representations. Voxels have some advantages and disadvantages compared to point cloud representations. First, voxels are straightforward generalizations of 2D pixels to 3D space, therefore we can leverage similar machinery used in icore-based generative modeling for images. These models are known to perform well and scale nicely with data. Second, message passing on graphs operate on single and pairwise interactions while convolution filters (and potentially transformer layers applied to regular grids) can capture multiple local interactions by construction (see [70] for discussion on the many-body representation hypothesis). Third, voxel representations have higher memory footprint and lower random memory accesses than point cloud representations [71]. We note however, that developing models on drug-sized molecules (that is, molecules with size close those on GEOM-drugs [41]) with reasonable resolution (.1-22A) is possible on current GPU hardware. Fourth, recovering point coordinates from a discrete grid has no analytical solution, therefore voxel-based models require an extra step to retrieve atomic coordinates. We show empirically that this is not problem in practice results, even with a very simple peak detection algorithm.
Finally, graph networks are less expressive due to message passing formalism [26, 27]. but are a better fit for built-in SE(3) equiiarrire architectures (e.g. [21, 22, 23, 24, 25]). Rotation-equivariant 3D con-
"This is necessary if we want to extract the molecular graph. However, raw voxelized generations could be useful to other computational chemistry tasks.
3

---

volutional network have been proposed [72, 73, 74] but current models do not scale as well as standard convnets, and would be challenge to apply them to drug-sized molecules. Built-in rotation equiv- ariance is a good property to have, however equivariance can also be learned with strong data augmen- lation/larger datasets [75, 76, 32]. In fact, concurrently to this work, [77] also show that built-in SE(3)- equivariant architecture is not necessary to generate molecules. Our experiments show that an expressive denoiser scales up better, allowing VoxMol to outperform current state of the art on GEOM-drugs. How-
ever, we hope our results motivate exploration of more efficient SE(3)-equivariant convnet architectures.
3
Method
We follow previous work (e.g., [78, 33, 70, 79]) and represent atoms as continuous Gaussian- like atomic densities in 3D space, centered around their atomic coordinates. Molecules are generated by discretizing the 3D space around the atoms into voxel grids, where each atom type (element) repre- sented by different grid channel. See appendix for more information on how we discretize molecules. This discretization process gives us a dataset with n voxelized molecules { E IR where is the length of each grid edge and c is the number of atom channels in the dataset. Each voxel the grid can take values between (far from all atoms) and (at the center atoms). Throughout our experiments, we consider a fixed resolution of 25A (we found to be a good trade-off between accuracy and computation). Therefore, voxel grids occupy a volume of (1/4)² cubic Angströms.
, dd 13,
3.1 Background: neural empirical Bayes
Let p(x) be an unknown distribution of voxelized molecules and p(y) a smoother version of it obtained by convolving p(x) with an isotropic Gaussian kernel with a known covariance La Equivalently, Y = where X ppp(x, N N~N(00,0 Ta)). Therefore Y is sampled from:
Rd (2To²)d/2 exp (- |1
p(y)=
p( )dr.
2g²
This transformation will smooth the density of X while still preserving some of the structure information of the original voxel signals. Robbins [80] showed that if we observe Y = y. then the y]- Built this result, Miyasawa [81] showed that, if the noising process is Gaussian (as in our case), then the least-square estimator (ty) can be obtained purely from the (unnormalized) smoothed density ply):
least-square estimator of X is the Bayes estimator, i.e., 2(y) |[X
(1) where g(v) = y. log p(y) is the score function [82] of p(y). This interesting equation tells us that, if we know ply) up to normalizing constant (and therefore the score function associated with it), we can estimate the original signal x only by observing its noisy version y. Equivalently, if we have access to the estimator (yy). we can compute the score function of p(y) via (1).
a(y) - y+ gy))
Our generative model is based on the neural empirical Bayes (NEB) formalism [1]: we are interested in learning the score function of the smoothed density p(v) and the least-square estimator â(y) from a dataset of voxelized molecules (zx { sampled from unknown p(z). We leverage the (learned) estimator and score function to generate voxelized molecules in two steps: (i) sample yı ~p(() with Langevin MCMC [83], and (ii) generate clean samples with the least- estimator. The intuition is that it is much easier to sample from the smooth density than the original distribution. See Saremi
and Hyvärinen [1] for more details.
3.2 Denoising voxelized molecules
We parametrize the Bayes estimator of X using a neural network with parameters 0 denoted by to Rd gd. Since the Bayes estimator is the least-squares estimator, the learning becomes least-squares denoising objective as follows:
( =
(2)
We made an attempt using an equivariant 3D convnets for denoising, but initial experiments were not successful. See appendix details.
"We use the convention where we drop random variable subscripts from probability density function when the arguments are present: p(æ) px (I) and p(y): (y)).

---

+
Mi
(b)
Figure 2: (a) A representation of our denoising training procedure. Each training sample (i.e., voxelized molecule) corrupted with isotropic Gaussian noise with fixed noise level ø. The model is trained to recover clean voxel grids from the noisy version. To facilitate visualization, we threshold the grid values, 2. ( (b) Graphical model representation of the walk-jump sampling scheme. The dashed arrows represent the walk, a MCMC chain to draw noisy samples from p(y). The solid arrow represents the jump. Both walks and jumps leverage the trained denoising network.
Using (1), we have the following expression for the smoothed score function in terms the denoising network
go( = 02 (1) (v)) yy)
By minimizing the learning objective (2) we learn the optimal it and by using (3) can compute the score function (y(y)= R V Vylogg(y).
We model the denoising network 2t with an encoder-decoder 3D convolutional network that maps every noised voxel on the grid to clean version it. Figure 2(a) shows a general overview of the denoising model. The noise level, a,, is kept constant during training and is key hyperparameter of the model. Note that in the empirical Bayes formalism, a can be any (large) value.
Compared to diffusion models, this training scheme is simpler as the noise level is fixed during training. VoxMol does not require noise scheduling nor temporal embedding on the network layers. We observe empirically that single-step denoising is sufficient reconstruct voxelized molecules (within the noise levels considered in this paper). Our hypothesis that this due to the nature of the voxel signals, which contain much more "structure" than "texture" information, in comparison to natural images.
3.3 Sampling voxelized molecules
We use the learned score function ga and the estimator 2e to sample. We follow the walk-jump sampling scheme [1, 37, 38, 39] to generate voxelized molecules .yy
(walk step) For sampling noisy voxels from p(y), we consider Langevin MCMC algorithms that are based on discretizing the underdamped Langevin diffusion [84]:
dug = yu;dt ugo dt + (V2yy)dd
dy: = v,dt,
where B, the standard Brownian motion in RRd, 7 and u are hyperparameters to tune (friction and inverse mass, respectively). We use the discretization algorithm proposed by Sachs ef al. [85] to generate samples yk. which requires discretization step See appendix for description of the algorithm.
iii (jump step) At arbitrary time step k. clean samples can be generated by estimating X from y& with the denoising network, i.e., computing Ia = (yk)..
This approach allows us to approximately sample molecules from p(x) without the need to compute (or approximate) Volog p(x). In fact, do MCMC on the smooth density p(y). which is known to be easier to sample and mixes faster than the original density p(x) [1,38, 38, 86). Figure 2(b) shows schematic representation of the generation process. Following [37], initialize the chains by adding uniform noise to Gaussian noise (with the same a used during training), i.e., yo = + U. N N(0 L2). U (0,1) (this was observed mix faster in practice).
Alternativlly, can also parameterize the score function: go 1 + in which case ko(v) y yoo 0z? go
5

---

Alpu)
Figure 3: Illustration of walk-jump sampling chain. We do Langevin MCMC on the noisy distribution (walk) and estimate clean samples with the denoising network at arbitrary time (jump).
The noise level plays key role in this sampling framework. If the noise is low, denoising (jump step) becomes easier, with lower variance, while sampling a "less smooth" p(y) (walk step) becomes harder. If the noise is high, the opposite is true.
Figure 3 illustrates an example of walk-jump sampling chain, where generated molecules change grad- ually as we walk through the chain (the clean samples are shown every ten steps, Ak= :d:=100 10). This figure a demonstration of the fast-mixing properties of our sampling scheme in generating 3D molecules. For instance, some atoms (or other structures like rings) might appear/disappper/chann through the chain. Interestingly, this behavior happened on most chains looked into explicitly.
as we move
3.4 Recovering atomic coordinates from voxelized molecules
It is often useful to extract atomic coordinates from generated voxelized molecules (e.g., to validate atomic valences and bond types or compare with other models). We use a very simple algorithm (a simplified version of the approach used in [33]) to recover the set of atomic coordinates from generated voxel grids: first we set to voxels with value less than i.e. T& = 12 Then we run simple peak detection to locate the voxel the center of each Gaussian blob (corresponding to the center of each atom). Finally we run a simple gradient descent coordinate optimization algorithm to find the set of points that best create the generated voxelized molecule. Once we have obtained the optimized atomic coordinates, we follow previous work [33, 18, 17, 20] and use standard cheminformatics software to determine the molecule's ' atomic bonds. Figure 4 shows pipeline to recover atomic coordinates and molecular graphs from generated voxelized molecules. See appendix for more details.
Experiments
In this section, we evaluate the performance of our model on the task of unconditional 3D molecule generation. Our approach is the first of its kind and therefore the objective of our experiments is to show that (i) VoxMol is feasible approach for unconditional generation (this is non-trivial) and (ii) it scales well with data, beating established model large, drug-like like dataset. In principle, VoxMol can be used for guided (or conditional) generation, an arguably more useful application for molecular sciences (see appendix for discussion on how guidance can be used on generation).
We start with a description of our experimental setup, followed by results on two popular datasets for this problem. We then show ablation studies performed on different components of the model.
4.1 Experimental setup
The denoising network às used in both the walk and jump steps described above. Therefore, its parametrization is very important to the performance of this approach. We use 3D U-Net [87] architecture for our denoising network. We follow the same architecture recipe as DDPM [34], with two differences: we use 3D convnets instead of 2D and we use fewer channels on all layers. The model has levels of resolution and we use self-attention on the two lowest resolutions. We augment our dataset during training by applying random rotation and translation to every training sample. Our models are trained with noise level a = -9, unless stated otherwise. We train our models with batch size of 128 and 64 (for QM9 and GEOM-drugs, respectively) and we use Adam [88]
Architecture.

---

sample voxelized malecules
find peaks
atomic
Figure 4: Pipeline for recovering atomic coordinates from voxel grids: (i) VoxMol generates voxelized molecules, (ii) atomic coordinates are extracted from voxel grid with simple peak detection algorithm, (iii) we use cheminformatics software to add atomic bonds and extract SMILES strings, molecular graphs, etc.
(learning rate 2 x 10 weight decay 10 10-2) to optimize the weights. The weights are updated with exponential moving average with a decay of ,999. We use y= 1.0, u 1.0 and 6 .5 for all our MCMC samplings. See appendix for more details on the architecture, training and sampling.
Datasets. We consider two popular datasets for this task: QM9 [40] and GEOM-drugs [41]. QM9 contains small molecules with up to 9 heavy atoms (29 if we consider hydrogen atoms). GEOM-drugs contains multiple conformations for 430k drug- sized molecules and its molecules have 44 atoms on average (up to 181 atoms and over 99% are under 80 atoms). We use grids of dimension 32 and 64 for QM9 and GEOM-drugs respectively. These volumes are able to cover over 998 of all points on both datasets. All our models model hydrogens explicitly. For QM9, we consider all 5 chemical elements (C, H, O. and F) present on the dataset. For (EOOM-rrugh we consider elements (C, H. O. N, F, S. CI and Br). We ignore P. I and B elements they appear in less than 1 % of the molecules in the dataset. Finally, the input voxel grids are of dimension b x 32 and R 8 for QM9 and GEOM- M-duugs, re- spectively. We perform the same pre- pre-processing and dataset split [20] and end up with 100K/20K/13K molecules for QM9 and 1. LLMM/46K/1466 for GEOM-drugs (train, validation, test splits respectively).
compare our method with two state-of -oate-of-thh-arttte approaches: GSchNer [47], a point-clluu autoregressive model and FDDM [15], a point-cloud diffusion based model. We note that both methods rely on equivariant networks, while ours does not. Our results could potentially be improved by success- fully exploiting equivariant 3D convolutional networks. We also show results of VoxMolocack in our main results, where we assume we have access to real samples from the noisy distribution. Instead of per- forming MCMC to sample yk. we sample molecules from the validation set and add noise to them. This baseline assumes we would have serfectsampling of noisy samples (walk step) and let us assess the qual- ity of our model to recover clean samples. It serves as an upper bound for our model and allows us to dis- entangle the quality of the walk (sampling noisy samples) and jump (estimating clean molecules) steps.
Baselines.
All methods generate molecules as set of atom types and their coordinates (in the case of voxelized molecules, we use the post-processing described above to get the atomic coordinates). We follow previous work [33, 18, 17,20 20] and use standard cheminformaties software to determine the molecule's atomic bonds given the atomic coordinates 5 Using the same post-processing for all methods allows a more apples-to-apples comparison the models.
Metrics. Most metrics we use to benchmark our model come from [20] 7, We draw 10,000 samples from each method and measure performance with the following metrics: stable mol and stable atom, the percentage of stable molecules and atoms, respectively, as defined in [15]: validity, the percentage of generated molecules that passes RDKit [90]'s s sanitization filter; uniqueness, the proportion of valid molecules that have different canonical SMILES; valency W, the Wasserstein distance between the distribution of valencies in the generated and test set; atoms TV and bonds TV, the total variation
We follow previous work and use OpenBabel [89] to get the bond orders.
We use the official implementation from https: //g// gigthhu..... com/ www.wwc/iiggaa....
7

---

stable mol %
stable
valid
unique %
valency
bond
bond
bond
atom %
%
W14
TV4
TV
len W 14
ang W
data
98.7
99.8
98.9
99.9
.001
.003
.000
.000
.120
GSchNet
92.0
98.7
98.1
94.5
.049
.042
.041
.005
1.68
EDM
97.9
99.8
99.0
98.5
.011
.021
.002
.001
0.44
aox
84.2
98.2
98.1
77.2
.043
.171
.050
.007
3.80
(::16)
(4.3)
(2:4)
(2:1.7)
(2.00
(2.200)
(2.010)
(2.00
(4:7)
VoxMol
89.3
99.2
98.7
92.1
.023
.029
.009
.003
1,96
(#-5)
(4:1)
(4:1)
(4:33
(4:0022
(4:0099
(#:002)
(#:002)
(4:00)
VoxMoloracte 90.1
99.3
98.9
99.9
.024
.009
.002
.001
0.37
Table I: Results on QM9. We use 10,000 samples from each method. Our results are shown with mean/standard deviation across 3 runs.
between the distribution of atom types and bond types, respectively; bond length W1 and bond angle W1, the Wasserstein distance between the distribution of bond and lengths, respectively. Finally, we also report the strain energy metric proposed in [91]. This metric defined the difference between the internal energy of the generated molecule's 's pose and relaxed pose of the molecule. The relaxation and the energy are computed using the Universal Force Field (UFF) [92] within RDKit. See appendix for more details about the metrics.
4.2 Experimental results
Table and Table 2 show results on QM9 and GEOM-drugs respectively. We report results for models trained with and without data augmentation (VoxMol and VoxMolmara respectively) and generate 10,000 samples with multiple MCMC chains. Each chain is initialized with 1, 1,000 warm-up steps, as we observed empirically that slightly improves the quality of generated samples. Then, samples are generated after each 500 walk steps (each chain having a maximum of 1,000 steps after the warm-up steps). Results for our models are shown with mean/standard deviation among three runs. The row data on both tables are randomly sampled molecules from the training set.
On QM9, VoxMol performs similar to EDM in some metries while performing worse in others (specially stable molecule, uniqueness and angle lengths). On GEOM-drugs, more challenging and realistic drug-like dataset, the results very different: VoxMol outperforms EDM in eight out of nine metrics, often by a considerably large margin.
Figure 5(a,b) shows the cumulative distribution function (CDF) of strain energies for the generated molecules different models on QM9 and GEOM-drugs, respectively. The closer the CDF of gener- ated molecules from model is to that of data (samples from training set), the lower is the strain energy of generated molecules. The ground truth data has median strain energy of 43.87 and 54.95 kcal/mol for QM9 and GEOM-drugs. GEOM-drugsseeppectvvee respectively. On QM9, all models have median strain energy around the
same ballpark: 52.58, 66.32 and 56.54 kcal/mol for EDM, VoxMolno rot and VoxMol, respectively. On GEOM the molecules generated by VoxMol have considerably lower median strain energy than EDM: 951 .23 kcal/mol for EDM versus 286.06 and 171.57 for VoxMol and VoxMol.
We observe, as expected, that augmenting the training data with random rotations and translations improves the performance of the model. The improvement is bigger on QM9 (smaller dataset) than on GEOM-drugs. In particular, the augmentations help capture the distribution of bonds and angles between atoms and to generate more unique molecules. We note that, unlike EDM, our model does not require knowledge of the number of atoms beforehand (neither for training nor sampling). In fact, Figure 6 show that our model Icarns the approximate distribution of the number of atoms per molccule on both datasets. Implicitly learning this distribution can be particularly useful in applications related to in in-painting (e.g., pocket conditioning, linking, scaffold conditioning). Finally, our method generates drug-like molecules in fewer iterations and is faster than EDM on average (see Table 3). EDM sampling time scales quadratically with the number of atoms, while ours has constant time (but scales cubically with grid dimensions).
These results clearly show one of the main advantages our approach: a more expressive model scales better with data. Architecture inductive biases (such as built-in SE(3) equivariance) are helpful in the setting of small dataset and small molecules. However, on the large-scale regime, a more expressive model is more advantagcous in capturing the modes of the distribution want to model. Compared
8

---

QM9
GEOW-drugs
1.0
1.0
data
0.8
EDM
Cumulatiive prob.
0.8
VorMol
Cumulative
0.6
VorMol (oracio)
MocMMal (ru rol)
0.6
0.4
0.4
data
FDM
0.2
0.2
VouMal (no rot)
VisMfll (oracle)
0.0
0.0
10
10
Strain Enerqy (UFF, kcal/mol)
Strain Energy (UFF, kcalmol)
Figure The cumulative distribution function of strain energy of generated molecules on (a) QM9 and (b) GEOM-drugs. For each method, use 10,000 molecules.
stable
stable atom %+
valid
unique
valency
atom
bond
bond
bond
mol 1
%
%,
W11
TV1
TV
len W11
ang W1
data
99.9
99.9
99.8
100.
.001
.001
.025
.000
.050
EDM
40.3
97.8
87.8
99.9
.285
.212
.048
.002
6,42
not
44.4
96.6
89.7
99.9
.238
.025
.024
.004
2.14
(2.1)
(#.1)
(#:2)
(2:0)
(2.001)
(#:001)
(2.001)
(2:000)
(d.07)
VoxMol
75.0
98.1
93.4
99.1
.244
.033
.036
.002
0.64
(1:1.)
(4:3)
(4:55
(2:2)
(2.003)
(2.011)
(2.006)
(d,001)
(E.13)
VoxMoloncce
81.9
99.0
94.7
97.4
.253
.002
.024
.001
0.31
Table 2: Results GEOM-drugs. We use 10,000 samples from each method. Our results are shown with mean/standard deviation across 3 runs.
QM9
GEOM-drugs
0.14
0.040
0.12
VoxMol
0.035
VooMMol
0.08
0.025
0.08
0.020
0.04
0.010
0.005
0.00
0.000
#atoms/molacule
#atoms/moleccue
(s)
(b)
Figure 6: Empirical distribution of number of atoms per molecule on QM9 (left) and GEOM-drugs (right). We sample 10,000 molecules from train set and generate the same number of VoxMol samples.
to VoxMoloracle results, that VoxMol can still be vastly improved. We potentially close this gap by improving the quality of the denoising network (e.g., by improving the architecture, train on more data, efficient built-in -quuvvvanvv CNNs, etc).
4.3 Ablation studies
Noise level a. Unlike diffusion models, the noise level is considered fixed during training and sampling. It is an important hyperparameter as it poses a trade-off between the quality of the walk step (Langevin MCMC) and the jump step (empirical Bayes). The ideal noise level is the highest possible value such that the network can still learn how to denoise. We train models on QM9 with o [.6..7 1.2), while keeping all other hyperparameters the same. Figure 7(a,b.c) shows how noise level a influences the performance on the validation set. While most metrics get better as the noise level increases. others (like stable molecules and valency W1) get worse after value. We observe empirically that .9 is the sweet spot level that achieves better overall performance on the validation set of QM9.
Number of steps Ak. Table shows how VoxMol's performance on GEOM-drugs change with number of walk steps Ak on the Langevin MCMC sampling. this experiment, we use the same trained model and only change the number of steps during sampling. Results of EDM are also shown for comparison (it always requires 1,000 diffusion steps for generation). We see that some metrics barely change, while others improve Ak increases. The average time (in seconds) to generate

---

0.10
99
80
90
25
0.20
0.06
85
2015
0.04
65
80
0.05
13
(6]
Figure 7: Effect of noise level generation quality. Models trained QM9 with different
noise level. Each plot shows two metrics: (a) molecule stability and uniqueness, (b) atom and bond TV, (c) valency and angle lengths WI.
Ak
stable mol %
stable atom %6
valid
unique %,
valency
atom
bond
bond
bond
avg.
(n steps)
%,
W14
TV;
TV4
len Wia
ang W14
s/mol
50
78.9
98.7
87.8
.200
.073
.102
.002
1.18
0.90
100
78.6
98.6
95.5
94.3
.256
.050
.101
.002
1.62
1.64
200
77.9
98.4
94.4
98.6
.253
.037
.104
.002
1.02
3.17
500
76.7
98.2
99.2
.252
.043
.042
.002
0.56
7.55
1,000
75.5
98.4
93.4
99.8
.257
.029
.050
.002
0.79
14.9
EDM
40.3
97.8
87.8
99.9
.285
.212
.048
.002
6.42
9.35
Table 3: Effect of number of walk steps Ak on generation quality on GEOM-drugs (2,000 samples).
EDM results are shown for comparison.
molecule increases linearly with the number of steps, expected. We observe that even using 500 steps, our model is still faster than EDM on average, while achieving better performance these metrics. Remarkably, with only 50 steps, VoxMol already outperforms EDM in most metrics, while being an order of magnitude faster on average.
Atomic density radii We also assess how the performance of the model changes with respect to the size of atomic radii chosen during the voxelization step (while always keeping the resolution of the grid fixed at 22A).. See appendix for how this is done. We tried four different values for the radii (same for all elements): 25,, 5, .75 and 1.0. We observe o different versions of the model, with different hyperparameters that using a fixed radius of consistently outperform other values. Training does not converge with radius .25 and quality of generated samples degrades increase the radius. We also tried to use Van der Waals radii (where each atom type would have their own radius), but results were also not improved.
5 Conclusion
We introduce VoxMol, a novel score-based method for 3D molecule generation. This method generates molecules in fundamentally different way than the current state of the art (i.e., diffusion models applied to atoms). The noise model used also novel in the class of score-based generative models for molecules. We represent molecules on regular voxel grids and VoxMol trained predict "clean" molecules from its noised counterpart. The denoising model (which approximates the score function of the smoothed density) is used to sample voxelized molecules with walk-jump sampling strategy. Finally atomic coordinates are retrieved by extracting the peaks from the generated voxel grids. Our experiments show that VoxMol scales better with data and outperforms (by a large margin) a representative state the art point cloud- based diffusion model on GEOM-drugs, while being faster to generate samples.
Broader impact. Generating molecules conditioned on some desiderata can have huge impacts in many different domains, such as, drug discovery, biology, materials, agriculture, climate, etc. This work deals with unconditional 3D molecule generation (in pure algorithmic way): problem that initial stepping stone (out of many) to this long-term objective. We. as society, need
can be seen to find solutions to use these technologies in ways that are safe, ethical, accountable and exclusively beneficial to society. These important concerns and they need to be thought of the same time we design machine learning algorithms.
The authors would like to thank the whole Prescient Design team for helpful discussions and Genentech's 's HPC team for providing reliable environment train/analyse models.
A Acknovledgements.
10

---

References
[1] Saeed Saremi and Aapo Hyvarinen. Neural empirical Bayes. JMLR, 2019.
[2] Tobias Fink, Heinz Bruggesser, and Jean-Louis Reymond. Virtual exploration of the small-molecule chemical universe below 160 daltons. Angewandte Chemie International Edition, 2005.
[3] Jiankun Lyu, Sheng Wang, Trent E Balius, Isha Singh, Anat Levit, Yurii S Moroz, Matthew J O Mcara, Tao Che, Enkhjargal Algaa, Kateryna Tolmachova, et al. Ultra-large library docking for discovering new chemotypes. Nature, 2019.
[4] Regine Bohacek, Colin McMartin, and Wayne C Guida. The art and practice structure- -astd drug design: a molecular modeling perspective. Medicinal research reviews, 1996.
[5] Camille Bilodeau, Wengong Jin, Tommi Jaakkola, Regina Barzilay, and Klavs Jensen. Generative models for molecular discovery: Recent advances and challenges. Computational Molecular Science, 2022.
[6] David Weininger. Smiles, a chemical language and information system. 1. introduction to methodology and encoding rules. Journal of chemical information and computer sciences, 1988.
[7] Mario Krenn, Florian Häse, AkshatKumar Nigam, Pascal Friederich, and Alan Aspuru- Guzuk. Self- referencing embedded strings (selfies): A 100% robust molecular string representation. Machine Learning: Science and Technology, 2020.
[8] Marwin HS Segler, Thierry Kogej. Christian Tyrchan, and Mark P Waller. Generating focused molecule libraries for drug discovery with recurrent neural networks. ACS centrul science, 2018.
[9] Thomas Blaschke, Marcus Olivecrona, Ola Engkvist, Jürgen Bajorath, and Hongming Chen. Application of generative autoencoder in de novo molecular design. Molecular informatics, 2018.
[10] Gabriel Lima Guimaracs, Benjamin Sanchez-Longeling, Carlos Outciral, Pedro Luis Cunha Farias, and Alín Aspuru-Guzik Objective-reinforreed generative adversarial networks (organ) for sequence generation models. arXiv preprint arXiv: 7705. 10843, 2017.
[11] Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction tree variational autoencoder for molecular graph generation. In ICML, 2018.
[12] Yujia Li, Oriol Vinyals, Chris Dyer, Razvan Pascanu, and Peter Battaglia. Learning deep generative models of graphs. arXiv preprint / 0003.00334044 2018.
[13] Jiaxuan You, Rex Ying, Xiang Ren, William I.. Hamilton, and Jure Leskovec, Graphrnn: A deep generative model for graphs. In ICMI, 2018.
[14] Omar Mahmood, Elman Mansimov, Richard Bonneau, and Kyunghyun Cho. Masked graph modeling for molecule generation. Nature Communications, 2021.
[15] Emiel Hoogebbom, Victor Garcia Satorras, Clément Vignac, and Max Welling, Equivariant diffusion for molecule generation in 3D. In ICML. 2022.
[16] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermoo,, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In ICI.R 2022.
[17] Ilia Igashov, Hannes Stärk, Clement Vignac, Victor Garcia Satorras, Pascal Frossard, Max Welling, Michacl M Bronstein, and Bruno Corrcia. Equivariant 3d-conditional diffusion models for molecular linker design. In NeurlPS, Al for ScienceWorksing 2022.
[18] Arn Schneuing, Yuanqi Du Charles Harris, Arian Jamasb, Ilia Igashov, Weitao Du, Tom Blundell, Pietro Lió, Carla Gomes, Max Welling, et al. Structure-based drug design with equivariant diffusion models. arXiv preprint arXiv:2210. 2022.
[19] Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi S. Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. In ICLR, 2023.
[20] Clement Vignac, Nagham Osman, Laura Toni, and Pascal Frossard. Midi: Mixed graph and 3d denoising diffusion for moleculc generation. arXiv preprint 200.000000 2023.
[21] Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang. Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-an and translation ritiiiinnequvvriinnt neural networks for 3d point clouds. arXiv preprint mrXiv 08219, 2018

---

[22] Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael JL Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. ICLR, 2021.
[23] Victor Garcia Satorras, Emiel Hoogehoom, and Max Welling. E (n) equivariant graph neural networks. In ICML., 2021.
[24] Mario Geiger and Tess Smidt. e3nn: Euclidean neural networks. arXiv preprint arXiv:2207.04453, 2022.
[25] Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. ICLR, 2023.
[26] Keyulu Xu, Weihua Hu, Jure Leskovee, and Stefanic Jegelka. How powerful are graph neural networks? ICLR, 2019.
[27] Christopher Morris, Martin Ritzert, Matthias Fey. William L Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe. Weisfeiler and leman go neural: Higher-order graph neural networks. In AAAI, 2019.
[28] Haggai Maron, Heli Ben-Hamu, Hadar Serviansky, and Yaron Lipman. Provably powerful graph networks. Neurips, 32, 2019.
[29] Tlyes Batatia, David Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi. MACE: Higher order equivariant message passing neural networks for fast and accurate force fields. Neurips, 35:11423-11436 2022.
[30] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. 1929, 2020.
[31] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems, 25, 2012.
[32] Nate Gruver, Mare Finzi. Micah Goldblum, and Andrew Gordon Wilson. The Lie derivative for measuring learned equivariance. arXiv preprint arXiv:2210.02988 2022.
[33] Matthew Ragoza, Tomohide Masuda, and David Ryan Koes. Learning continuous representation of 3d
molecular structures with deep generative models. In Neurips, Structural Biology workshop, 2020.
[34] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In NeurlPS, 2020.
[35] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. NeurlPS, 2021.
[36] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björ Ommer. High-rcsolution image synthesis with latent diffusion models. In CVPR, 2022.
[37] Saeed Saremi and Rupesh Kumar Srivastava. Multimeasurement generative models. ICLR, 2022.
[38] Saeed Saremi, Rupesh Kumar Srivastava, and Francis Bach. Universal smoothed score functions for generative modeling. arXiv preprint 230.::20331116666 2023.
[39] Nathan C Frey, Dan Berenberg. Joseph Kleinhenz, Isidro Hotzel, Julien Lafrance- Vanasse, Ryan Lewis Kelly, Yan Wu, Arvind Rajpal, Stephen Ra, Richard Bonneau, Kyunghyun Cho, Andreas Loukas, Vladimir Gligorijevic, and Saced Saremi. Learning protein family manifolds with smoothed energy-based models. In ICLR, Workshop on Physics for Machine Learning, 2023.
[40] Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learning. Chemical science, 2018.
[41] Simon Axelrod and Rafael Komez-Bombaree Geom, energy-annotate molecular conformations for property prediction and molecular generation. Scientific Data, 2022.
[42] Miha Skalic, José Jiménez, Davide Sabbadin, and Gianni De Fabritiis. Shape-hased generative geneerartiv modeling for de novo drug design. Journal of chemical information and modeling. 2019.
[43] Diederik P Kingma and Max Welling. Auto-encoding variational Bayes. In ICLR, 2014.
[44] Lvwei Wang. Rong Bai, Xiaoxuan Shi, Wei Zhang. Yinuo Cui, Xiaoman Wang, Cheng Wang. Haoyu Chang, Yingsheng Zhang, Jielong Zhou, et al. A pocket-basee 3d molecule generative model fueled by experimental electron density. Scientific reports, 2022.
12

---

[45] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yosh Bengio. Generative adversarial nets. In NIPS, 2014.
[46] Fergus Imrie, Thomas Hadfield, Anthony R Bradley, and Charlotte M Deane. Deep generative design with 3d pharmacophoric constraints. Chemical science, 2021.
[47] Niklas Gebauer, Michael Gastegger, and Kristof Schüit.. Symmetry-adaptedd generation 3d point sets for the targeted discovery of molccules. In NeurlPS, 2019.
[48] Niklas Gebauer, Michael Gastegger, and Kristof Schlitt. Generating equilibrium molecules with deep neural networks. arXiv preprint arXiv: 2018.
[49] Youzhi Luo and Shuiwang Ji. An autoregressive flow model for 3d molecular geometry generation from scratch. In ICLR, 2022.
[50] Jonas Köhler, Leon Klein, and Frank Noć. Equivariant flows: exact likelihood generative learning for symmetric densities. In ICMI,, 2020.
[51] Victor Garcia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling. E (n) cquivariant normalizing flows. In NewrIPS, 2021.
[52] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In ICMI, 2015.
[53] Jascha Sohl-Dickstein, Fric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In ICMI., 2015.
[54] Lei Huang, Hengtong Zhang, Tingyang Xu, and Ka-Chun Wong. Mdm: Molecular diffusion model for
3d molecule generation. arXiv preprint arwiv:2209.0577777 2022.
[55] Lemeng Wu, Chengyue Gong, Xingchao I. Liu, Mao Ye, et al. Diffusion-based molecule generation with informative prior bridges. In NeurlPS, 2022.
[56] Minkai Xu, Alexander Powers, Ron Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models 3d molecule generation. ICML, 2023.
[57] Elman Mansimov, Omar Mahmood. Seokho Kang, and Kyunghyun Cho. Molecular geometry prediction using deep generative graph neural network. Scientific Reports, 2019.
[58] Gregor NC Simm and José Miguel Hernández-Lobato. A generative model for molecular distance geometry. JCML, 2020.
[59] Frank Noé. Simon Olsson. Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-hody systems with deep learning. Science, 2019.
[60] Gregor NC Simm, Robert Pinsler, Gábor Csányi, and José Miguel Hernández-Lobato Symmetry aware actor-critic for molecular design. arXiv preprint arXiv:2011.12747 2020.
[61] Octavian Ganea, Lagnajit Pattanaik, Connor Coley, Regina Barzilay, Klavs Jensen, William Green, and Tommi Jaakkola, Geomol: Torsional geometric generation molecular 3d conformer ensembles. In NeurlPS, 2021.
[62] Xingang Peng, Shitong Luo, Jiaqi Guan, Qi Xie, Jian Peng, and Jianzhu Ma. Pocket2mol: Efficient molecular sampling based on 3d protein pockets. In JCML, 2022.
[63] Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular conformation generation. In ICML, 2021.
[64] Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional diffusion for molecular conformer generation. arXiv preprint arXiv:2206.01722 01790012290000 2022.
[65] Siyu Long, Yi Zhou, Xinyu Dai, and Hao Zhou. Zero-shot drug design by sketching and generating. In NeurlPS, 2022.
[66] Keir Adams and Connor W Coley. Equivariant shape-conditioned generation of 3d molecules for ligand- based drug design. arXiv preprint arXiv:22
2022.
[67] 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. ICLR, 2023.
13

---

[68] Matthew Ragoza, Tomohide Masuda, and David Ryan Koes. Generating 3d molecules conditional on receptor binding sites with deep generative models. Chemical science, 2022.
[69] David Schaller, Dora Sribar, Theresa Noonan, Lihua Deng, Trung Ngoc Nguyen, Szymon Pach, David Machalz, Marcel Bermudez, and Gerhard Wolber. Next generation 3d pharmacophore modeling. Wiley Interdisciplinary Reviews: Computational Molecular Science, 2020.
[70] Raphacl JI. Townshend, Martin Vögele, Patricia Suriana, Alexander Derry, Alexander Powers, Yianni Laloudakis, Sidhika Balachandar, Bowen Jing. Brandon Anderson, Stephan Eismann, et al. Atom3d: Tasks on molecules in three dimensions. NeurIPS, 2020.
[71] Zhijian Liu, Haotian Tang, Yujun Lin, and Song Han. Point-voxel cnn for efficient 3d deep learning. In NeurIPS, 2019.
[72] Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco S Cohen. 3d steerable enns: Learning rotationally equivariant features in volumetric data, NeurlPS, 2018.
[73] Ivan Diaz, Mario Geiger, and Richard lain McKinley. An end-to-end se (3)-cqiivariant segmentation network. arXiv preprint ww..ii:3033.000000 00351,
[74] Jiehong Lin, Hongyang Li, Ke Chen, Jiangbo Lu, and Kui Jia. Sparse stecrable convolutions: An efficient learning of se (3) equivariant features for estimation and tracking of object poses in 3d space. Neur/PS, 2021.
[75] Karel Lenc and Andrea Vedaldi. Understanding image representations by measuring their equivariance and equivalence. In CVPR, 2015.
[76] Diane Bouchacourt, Mark Ibrahim, and Ari Morcos, Grounding inductive biases in natural images: invariance stems from variations in data. In NeurlPS, 2021.
[77] Daniel Flam- Shepherd and Alán Aspuru Guzik. Language models can generate molecules, materials, and protein binding sites directly in three dimensions xyz, cif, and pdb files. arXiv preprint 20X5523305.05770, 2023.
[78] Matthew Ragoza, Joshua Hochuli, Elisa Idrobo, Jocelyn Sunseri, and David Ryan Koes. Protein-ligand scoring with convolutional neural networks. Journal of chemical information and modeling, 2017.
[79] Michael Maser and SE Reisman. 3d computer vision models predict dft-level homo-lumo gap energies from force-fieldooptimized geometries. ChemRvix, 2021.
[80] Herbert Ellis Robbins. An empirical Bayes approach statistics. In Proc. 3rd Berkeley Symp. Math. Statist. Probab. 1956, 1956.
[81] Koichi Miyasawa. An empirical bayes estimator of the mean of a normal population. Bull, Inst. Internat. Statist, 1961.
[82] Aapo Hyvärinen. Estimation of non- normoomallzed statistical models by score matching. JMLR, 2005.
[83] Giorgio Parisi, Correlation functions and computer simulations. Nuclear Physics B, 1981.
[84] Xiang Cheng, Niladri S. Chatterji. Peter Bartlett, and Michael Jordan. Underdamped Langevin MCMC: non-asymptotic analysis. In COLT, 2018.
[85] Matthias Sachs, Benedict Leinkuhler, and Vincent Danos. Langevin dynamics with variable coefficients and nonconservative forces: from stationary states to numerical methods. Entropy, 2017.
[86] Saced Saremi, Ji Won Park, and Francis Bach. Chain of log-concave Markov chains. arXiv preprint arXiv:2305.19.0 19473, 2023.
[87] Olaf Ronncherger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In MICCAI, 2015.
[88] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. ICLR, 2019.
[89] Noel M O'Boyle, Michael Banck, Craig A James, Chris Morley, Tim Vandermeersch,, and Geoffrey R Hutchison. Open babel: An open chemical toolbox. Journal of cheminformatics, 2011.
[90] Greg Landrum. Rdkit: Open-source cheminformatics software, 2016.
14

---

[91] Charles Harris, Kieran Didi, Arian R Jamasb, Chaitanya K Joshi. Simon V Mathis, Pietro Lio, and Tom Blundell. Benchmarking generated poses: How rational is structure-based drug design with generative models? arXiv preprint mrxiiv:308.0041:1 2023.
[92] Anthony K Rappé, Carla J Casewit, KS Colwell, William A Goddard III. and W Mason Skiff. Uff. a full periodic table force field for molecular mechanics and molecular dynamics simulations. Journal of American chemical society, 1992.
[93] Lin Li, Chuan Li, and Emil Alexov. On the modeling of polar component of solvation energy using smooth gaussian-based dielectric function. Journal of Theoretical and Computational Chemistry, 2014.
[94] Gabriele Orlando, Daniele Raimondi, Ramon Duran- Buran-Romañaa, Yves Moreau, Joost Schymkowitz, and Frederic Rousseau. Pyuul provides an interface between biological structures and deep learning algorithms. Nature communications, 2022.
[95] Yuxin Wu and Kaiming He. Group normalization. In ECCV, 2018.
[96] Stefan Elfwing. Eiji Uchibe, Kenji Doya. Sigmoid- linear units for neural network function approximation in reinforcement learning. Neural Networks, 2018.
[97] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style. high-performance deep learning library. In NeurlPS, 2019.
15

---

A Extra implementation details
A.1 Voxel representation
Molecules in our datasets are converted into voxelized atomic densities. For each molecule, we consider a box around its center and divide it into discrete volume elements. We follow [93, 94] and first convert each atom (of each molecule) into 3D Gaussian-lite densities:
de
Va (dddra==xepp (.93-ra)2)
(5)
where Va is defined as the fraction of occupied volume by atom a of radius Ta at distance from its center. Although we could consider a different radius for each element, in this work we consider all atoms to have the same radius ra .5.. The occupancy of each voxel in the grid is computed by integrating the occupancy generated by every atom in molecule:
Na O,c,,,x=1-1 =1- II (1-1o Va, n=1
(6)
where N₂ is the number of atoms in the molecule, on is the n atom, Ci,j, are the coordinates (i.j.k) in the grid and Zn is the coordinates of the center of atom [93]. The occupancy takes the maximum value of I at the center of the atom and goes to 0 as it moves away from it. Every channel is considered independent of one another and they do not interact nor share volumetric contributions. We use the python package PUUUL [94] to generate the voxel grids from the raw molecules (.xyz .sdf format).
We use grids with 32 voxels on QM9 and 64 on HEOM-drugh and place the molecules on the center of the grid. These volumes are able to cover over 99% of all points in the datasets. We use all 5 chemical elements present on the dataset (C, H. O, and F), while for GEOM-drugs, we use 8 (C, H. O, N, F, S, CI and Br). We model hydrogen explicitly in all our experiments. Finally, the input voxel grids are of dimension Rxxx32 xxx32x32x*> 32x and R xxxxx64xx0 for QM9 and GEOM-drugs, respectively.
We augment the dataset during training by applying random rotation and translation to the molecules. For rotation, we sample three Euler angles uniformly (between [0,2 )) and rotate each training sample. For translation, we randomly shift the center of the molecules on each of the three dimensions by sampling an uniform shift between [0,.
A.2 Architecture
Our neural network architecture follows standard encoder-decoder convnet architecture. We use very similar architecture recipe to DDPM [34]. The model uses four levels of resolution: 32 to 43 for the QM9 dataset and 64 to 8 for the GEOM-drugs dataset. The input voxel is embedded into a 32 dimensional space with a grid projection layer (3D convnet with kernel size 3 x x3+ x 3). Each resolution (on both encoder and decoder) has two convolutional residual blocks. Each block contains a group normalization [95] layer, followed by SiLU [96] non-linearity and a 3D convnet (with kernel size 3 3 x 3). All convolutions have stride and we pad the feature maps with on each side. We use self-attention layers between the convolutional layers in the two lowest resolutions. We reduce (increase, respectively) the resolution of the encoder (decoder) with 2x2 x (stride 1) max-poolings sulimear-upsammlingg.... The model has skip connections at each resolution to concatenate the encoder feature map with the decoder feature map. We double the number of feature maps at each resolution, except the last resolution where we quadruple. VoxMol has approximately 11IM parameters. We also implemented a smaller version (with reduced number of channels per layer) with around 30M. These models achieve performance close to the base model and are faster to train and sample.
A.3 Built-in SE(3) equivariance experiments
In early experiments, we made an attempt at using a SEuiieequiaariant 3D U-Net using steerable convnets [72] for denoising, but initial experiments were not successful. The hypothesis is that a built-in SE(() equivariant version of our model, VoxMolagui would be advantageous over the non-equivariant version for the task of molecule generation. We start with the official implementation of [73] and tune several network hyperparameters (related to architecture, optimization and training) so that the network is able to achieve good denoising metries QM9. We then use the same procedure to
16

---

generate samples as described in the main paper (only switching the network from the non non-equivariant the equivariant version). We tried different sampling hy hyperparameters, but we were never able to achieve the same performance as non-equivariant Vox Table 4 compares the results of the model with and without built-in SE(3) equivariance.
QM9
stable mol %-
stable atom %
valid
unique
valency
atom
bond
bond
bond
%
%
W14
TV:
TV1
W11
ang W14
VoxMol
89.3
99.2 81.8
98.7
92.1
.023
.029
.009
.003
1.96
VoxMolagui 25.1
95.9
92.9
13.2
.104
.015
.015
5.31
Table 4: Results on QM9 of our model without (VoxMol) and with (VoxMolequi) built-in SE(3) equivariance.
There might be many reasons why this is the case: (i) the best reconstruction loss we found with the equivariant model higher than the non-equivariant (approx. 9.4 10 vs. 5.4 10-5 MSE on val. set), (ii) the equívariant model needs more capacity to be competitive with the non-equivariant one (currently it has over 9 fewer parameters), (iii) something in the sampling procedure needs to be different the equivariant version (unlikely).
We hypothesize that if an equivariant version of VoxMol achieves similar (or lower) reconstruction loss as the vanilla version, it will probably achieve competitive/hetter results in the task of molecule generation. Finally, our equivariant implementation is less efficient (around 50-60% slower) and consumes more memory than the original version. This poses, therefore, an extra challenge to scale up the size of the dataset and the size of the molecules (e.g., GEOM-drugs requires 64 voxel grid).
A.4 Training and sampling
The weights are optimized with batch size 128 and 64 (for QM9 and GFOM-drugs, respectively),
AdamW W optimizer (B1 -9, $2 .999), learning rate of 10-5 and weight decay of 10-2, The models are trained for 500 epochs on QM9 and around 24 epochs on GEOM-drugs. We discretize the underdamped Langevin MCMC (Equation 4) with the algorithm proposed by Sachs et al. [85] (this has been applied on images before [37]). Algorithm describes this process.
Algorithm 1: Walk-jump sampling [1] using the discretization of Langevin diffusion by [85]. Lines 6-1 13 correspond to the walk step and line 14 to the jump step.
Input 8 (step size), u (inverse mass), y (friction), K (steps taken)
2: Input Learned score function go(3) (3) = V ylogp(y) and noise level ø
3: Output Øx
4: Yo NO0, La) + (,,1,
5: no +-
6: for k= O....K - do
7:
M/+1 - 8k +
8:
g 90 (3k+1)
9:
UK+1 Un + ud
T
10:
e~N(0,,a)
11:
Uk+1 Uk+1 - exp( -70) ud
g expp -270))
g
2
12:
DE+1 1 v4+ 94+1 uk+1
13:
end for
14:
āK < yK o - (Uk)
We use u 1.0, w= == 1.0, .5 for all samplings and we generate multiple chains in parallel (200 chains for QM9 and 100 for GEOM-drugs) We follow [37] and initialize the chains by adding uniform noise to the initial Gaussian noise (with the same a used during training), i.e., yo N(0,0² 2 I,) (0,1)
(this was observed to mix faster in practice).
All experiments and analysis on this paper were done on 100 GPUs and with PyTorch [97]. The models on QM9 were trained with 2 GPUs and the models GEOM-drugs on 4 GPUs.
17

---

dsct
coords ref.
stable
stable
valid %
unique $b
valency
atom
bond
bond
bond
mol 1
atom %-
W11
TV,
TV,
len W 11
ang W
QM9
80.5
98.5
98.1
93.3
.051
.028
.005
.008
2.94
89.3
99.2
98.7
92.1
.023
.029
.009
.003
1.96
GEOM
73.9
99.0
94.7
98.6
236
.030
.038
.008
2.92
74.9
98.1
93.4
99.2
.254
.033
.036
.002
.63
Table 5: Effect of coordinate refinement QM9 and GEOM-drugs. We use 10,000 samples from each method.
A.5 Recovering atomic coordinates from voxel grid
Figure 4 shows our pipeline to recover atomic coordinates and molecular graphs from generated voxelized molecules. In the first step, we use the model to "jump" to the data manifold generating sample in the voxelized representation, Ik. We set all voxels with value less than .1, i.e., Jx (xx)) We then apply a simple peak finding algorithm to find the voxel coordinates corresponding to the peaks in the generated sample. Our peak finding algorithm uses maximum filter with 3 x 3 x 3 stencil to find local maxima. Note that this algorithm always returns points on the voxel grid and is therefore limited by the resolution of the discretization.
In order to further refine the atomic coordinates we take advantage of the fact that our voxelization procedure is differentiable to perform gradient based optimization of the coordinates. Specifically we use L-BFGS to optimize the atomic coordinates based on the L2 norm of the reconstruction error in the voxel representation. Note, unlike some previous work [33] we perform peak detection and refinement in single step and do not perform search over multiple possible numbers of atoms or atom identities.
Table 5 shows the effect of coordinate refinement molecule generation. We generate molecules on the same setting in the experimental section.
Once we have obtained the optimized atomic coordinates, we follow previous work [33, 18, 17, 20] and use standard cheminformatics software determine the molecule's atomic bonds.
A.6 Metrics
Most of the metrics used benchmark models come from [20] Below we describe the metries:
Atom stability: the percentage of generated atoms with the right valency. This metric is computed on the raw 3D sample (before any postprocessing), therefore it is more stringent than validity.
Molecule stability: the percentage of generated molecules where all its atoms stable.
Validity: The percentage of generated molecules that passes RDKit's sanitization filter.
Uniqueness:
The proportion of valid molecules (defined above) that has unique canonical SMILES (generated with RDKit) representation.
Atoms TV: The total variation between the distribution of bond types in the generated and test set. We consider atom types on QM9 and atom types on GEOM-drugh The histograms hatm and Aanm are generated by counting the number of each atom type on all molecules on both the generated and real sample set. The total variation is computed as:
Atoms (hatm Nastm) #aaulm types
Rram (x) -hatm (x))
Bonds TV: Similar to above, the histograms for real and generated samples are created by counting all bond types on all molecules. The total variation is computed as:
Bonds TV (ihhonnd,/honl
B (x)) -hond (1)
Ebood types
'We do not compare directly with [20], since this model is an extension of EDM and leverages more information (connectivity graph and formal charges) during training.
18

---

Valency W1: This is the weighted sum of the Wasserstein distance between the distribution of valencies for each atom type:
Valency W targe))
Σ
(ha)
Eatom types
where homa (1r) and hval (ar) the histogram of valencies for atom type ar for generated and holdout set samples, respectively.
Bond length W1: The weighted sum of Wasserstein distance between the distribution of bond lengths for each bond type:
Bond Len W (generated,,aargel =
Σ
Wbb)iihaais (bass (b)),
6Ehond types
where haist (6) and hdist (6) are the histogram of bond lengths for bond type 6, for generated and holdout set samples, respectively.
Bond angles W1: The weighted sum of Wasserstein distance between the distribution of bond angles (in degrees) for each atom type in the dataset:
Bond Ang W (generated,target)
Σ
(hanga (x))
@atttom types
where hang(1) (x) and hang (z) are the histogram of angles for atom type x for generated and holdout samples, respectively. See [20] for how angles are measured.
Strain energy: The strain energy for generated molecule is computed as the difference between the energy on the generated pose and the energy of a relaxed position. The relaxation and the energy are computed using UFF provided by RDKit. We use [91]'s implementation"
A.7 Guiding the generation process
Like diffusion models, our method also leverages (learned) score functions and relies on Langevin MCMC for sampling. Therefore, in theory diffusion models: by constraining the score function as we walk through the MCMC chain. In the case of diffusion models, the score function of all steps is constrained to guide the transition steps from noise to (conditioned) sample. In VoxMol, the constrained score function would affect the "walk steps" (the Langevin MCMC steps): it would restrict the region where the chain samples noisy molecules to p(ylc) (instead of p(y)). is the condition (e.g., gradient of classifier). The "jump step" (a forward pass of the denoising network over the noised molecules) is independent of the condition and remains unchangeable.
condition VoxMol similarly to how it is done in
Many of the innovations on conditioning diffusion models come from computer vision, where U-nets are usually used. Since VoxMol has the same architecture (albeit 3D instead of 2D), many the condition- ing techniques/tricks used in images may be more easily transferable. For example, we could in principle use the gradient of classifier (trained jointly) to guide the sampling (using the same trick as in Dhariwal and Nichol [35]) adapt gradient-free guidance ([34]). Pocket conditioning could also be possible, as in e.g.. [18, 67], but utilizing voxel representations instead of point clouds and neural empirical Bayes in- stead of diffusion models. In painting has also proven to work very well in 2D U-Nets, so it could poten- tially work with 3D U-Nets as well. These in-painting techniques could also be leveraged in the context of molecule generation on voxel grids, e.g., for linker generation, scaffold/fragment conditioning.
Coode taken from strain. py
com/ http:///thhuubcooccccc////////// http:hhhww.c.cccccccccccccoo///a
19

---

B
Generated samples
NH
NH
Figure 8: Random generated samples from VoxMol trained on QM9 (passing RDKit's sanitization). Molecular graphs are generated with RDKit.
20

---

Br
Figure 9: Random generated samples from VoxMol trained GEOM-drugs (passing RDKit's sanitization). Molecular graphs are generated with RDKit.
21


Reviews:
[{"review_id": "HGYKmTXbtm", "reviewer": "Reviewer_R82i", "summary": "This paper proposes doing diffusion on voxel space for molecule conformation generation. Because the voxel space is discrete, an efficient sampling method is proposed.\n", "questions": "I have some questions on the method section, and would like to confirm with the authors.\n- I am wondering what’s the advantage of using NEB for estimating p(x)? Based on lines 135-140, it seems to be a standard EBM with Langevine dynamics, puls an extra NEB module. Is this correct?\n- The authors mentioned that `Chains are initialized with Guassian noise with standard deviation`. I would like to double-check what does standard Guassian mean in the voxel grids?\n", "limitations": "N/A", "rating": 4, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper proposes doing diffusion on voxel space for molecule conformation generation. Because the voxel space is discrete, an efficient sampling method is proposed.\n", "weaknesses": "## Method:\n- It is not clear how the SE(3)-equivariant property is guaranteed in the model. This is critical because once the molecules are too large to fit in the voxel grids, we need to rotate it for a better position.\n- The objective function in Eq 2 is not clear. Concretely, it is not clear how the atom types are trained.\n\n## Experiments:\n- The authors also referred to MiDi in the paper. I’m wondering why the authors did not add MiDi? Is that because the performance of MiDi is better than VoxMol?\n- The MiDi paper introduces more domain-specific metrics, and authors should consider add them.\n- Also a closely-related baseline should be cited at least: https://arxiv.org/abs/2305.05708.\n\n## Minor:\n- `... capturing long-range dependencies over multiple atoms can become difficult …` It also depends on how the edges are constructured. For point clouds, the edges are often constructured by distance, thus this may not become an issue as long as the spatial distance is close.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback and suggestions. Below we address additional concerns from the reviewer.\n\n**Diffusion model:**\n \nWe would like to clarify that the proposed model is _not_ a diffusion model (this is a big point of the paper). Our model is based on neural empirical Bayes and walk-jump sampling (Saremi and Hyvarinen, 19). Although the model is also a generative model based on score functions, it is fundamentally different than diffusion models.\n\n**Advantage of using NEB for estimating p(x):**\n\nIn general , it is very difficult to use the score function (the gradient of the loglikelihood of P) to sample from P(x) (space of \"clean\" molecules). The intuition of NEB is that sampling the noisy molecules y from P(y) is much easier than sampling from highly sparse x in P(x). In NEB, the score function is used to sample noisy molecules y. A single chain of VoxMol sampling can generate multiple noisy molecules (with a fast mixing time in practice). We can then \"jump\" from noisy to clean molecules with the denoiser network. This is in contrast to diffusion models, where the noise is gradually reduced in the reverse diffusion process (thus one chain per single sample).\n\n**Comparison w. arXiv2305.05708:**\n \nWe note this paper appeared 6 days before the submission deadline. We will add this reference to the revised version of the manuscript. What we like about this paper is that—similar to our work—it shows that we can generate molecules with methods that are different from GNN on point clouds. They also show—like we did—that built-in SE(3)-equivariant architecture is not necessary to generate molecules; equivariance can be learned with data/augmentation, as argued in [A, B, C] and related works.\n\n**SE(3) equivariance:**\n\nVoxMol has built-in equivariance to translation and relies on data augmentation through rotation (as it has been done before) to learn rotation equivariance. Unlike arXiv2305.05708, 3D CNNs can in principle be adapted to have built-in rotation equivariance (Weiler et al, 18). Please see SE(3) equivariance paragraph in general rebuttal.\n\n**\"…capturing long-range dependencies over multiple atoms can become difficult…\":**\n \nWe agree that it depends on how the edges are constructed. However, we point that EDM considers interactions between all atoms and uses a fully-connected graph (paragraph after equation 11 on EDM paper). We will clarify that this is a limitation of EDM in particular and not GNNs in general.\n\n**Eq 2 is not clear. Concretely, it is not clear how the atom types are trained:**\n \nEach atom is represented by a different input channel (similar to how R, G and B colors are separated on images). The training is done by denoising every voxel location of every channel of the noisy voxel grids (similar to how image denoisers are applied to every pixel on the three color channels).\n\n**Chain initialization:**\n \nWe describe how we initialize the chain on Appendix A3: \"We follow [27] and initialize the chains by adding uniform noise to the initial Gaussian noise (with the same σ used during training), i.e., y0 = N(0,σ2Id)+Ud(0,1) (this was observed to mix faster in practice)\". y0 is a voxel grid of tensor dimensions `(n_channels, dim, dim, dim)`, where dim is the dimension of the voxel grid (32 on QM9 and 64 on GEOM-drugs) and `n_channel` is the number of atom channels considered. We will move this to the main text for clarity.\n\n[A] Understanding image representations by measuring their equivariance and equivalence, Lenc,Vedaldi, CVPR15\n\n[B] Grounding inductive biases in natural images: invariance stems from variations in data, Bouchacourt,et al,NeurIPS21\n\n[C] The Lie Derivative for Measuring Learned Equivariance. Gruveret al,22"}, {"review_id": "4Xxh1qeiPD", "reviewer": "Reviewer_j1qM", "summary": "This paper proposes a novel 3D molecule generation routine dubbed VoxMol. The highlight of it lies in its introduction of a connection between traditional molecular graph representation and 3D voxel representation. In VoxMol, the molecules are first translated into voxel representation. A denoising network is then trained to serve as a score approximation. In the sampling process, a novel walk-jump procedure is introduced to sample noisy p(y) and obtain the denoised x. ", "questions": "As the authors have stated in the paper, efforts have been made for equivariant 3D CNNs but were unsuccessful. I am curious about what detailed experiments have been conducted and how were the results. I believe by involving equivariance the work will make a stronger contribution.", "limitations": "The limitations are not discussed in the main text. I would recommend the authors add these discussions since I think the work is inspiring and will benefit more from future research.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. To my knowledge, the method is novel in terms of data representation and the training/sampling process. The paper showcases that instead of the traditional graph representation, molecule generation can also be performed on 3D voxels, shedding light on new research directions of connecting 3D CNNs to equivariant GNNs. The authors also tried to recover the molecular graph representation based on the denoised voxel, which I really appreciate.\n\n2. The evaluations are interesting. Though not always state-of-the-art, the authors have exhibited insightful investigations, e.g., how $\\sigma$ influences the tradeoff between validity and uniqueness.\n\n3. The overall presentation is clear and the method is very easy to follow. \n\n", "weaknesses": "1. I feel a little pity that equivariance is not properly injected into the model. I think it would make a much stronger impact if equivariance could be enforced even on this 3D voxel representation so that physical/chemical priors could be ensured.\n\n2. The method does not significantly beat EDM on these benchmarks. But I understand that as perhaps the first work to challenge the task of 3D molecule generation with 3D CNNs, it is difficult to beat all other GNN counterparts at once.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback and suggestions. Below we address additional concerns from the reviewer.\n\n**The method does not significantly beat EDM on these benchmarks:** \n\nTable 1 and 2 on the attached pdf compares our method with EDM using MiDi metrics (as suggested by reviewers). Although it underperforms EDM on QM9 (by a relatively small margin), Voxmol clearly beats EDM by a large margin on most metrics on the challenging (and more realistic!) GEOM-drugs dataset. These conclusions were the same with the original metrics, but the new metrics show the differences in performance more clearly. We will update the tables and plots of the experiments to reflect the new metrics.\n\n**Limitations:** \n\nWe compare the pros/cons of voxel/CNN vs point-cloud/GNN on the last paragraph of related works (L100-117). We also mention the limitation that it scales cubic with voxel grid dimension (although been constant time given a fixed grid dimension) on L257-259. We would make these limitations more clear by having an explicit \"Limitations\" paragraph at the end of the manuscript (and also updating with reviewers' remarks).\n\n**SE(3) equivariance:**\n\nPlease see SE(3) equivariance paragraph in general rebuttal.\n\n**SE(3)-Equivariant 3D CNN experiments:** \n\nWe start with the official implementation of [D] and tune several network hyperparameters (related to architecture, optimization and training) so that the network is able to achieve good denoising metrics on QM9.  We then use the same procedure to generate samples as described in the paper (only switching the network from the non-equivariant to the equivariant version). We tried different sampling hyperparameters, but we were never able to achieve the same performance as non-equivariant VoxMol (eg, the molecule stability (w MiDi metrics) dropped from .87 to .25 per cent). There might be many reasons why this is the case: (i) the best reconstruction loss we found with the equivariant model is higher than the non-equivariant (~9.4e-5 vs. 5.4e-5 MSE on val. set), (ii) the equivariant model needs more capacity to be competitive to the non-equivariant one (currently it has over 90x less parameters), (iii) something on the sampling procedure needs to be different on the equivariant version (unlikely). What makes it even more challenging is that the equivariant network is less efficient (around 50-60% slower) and we faced issues while scaling up to 64^3 grid size (required for GEOM-drugs). We will include a description of what we tried on this space on the appendix.\n\n[D] An end-to-end SE(3)-equivariant segmentation network, Diaz, Geiger, McKinley, 23\n"}, {"review_id": "Ynap7i7NtP", "reviewer": "Reviewer_bpDM", "summary": "This paper proposes VoxMol, a novel method for generating 3D molecules in the form of voxel grids. The proposed method adopts neural empirical Bayes as the basic probabilistic framework to develop generative models for 3D voxel grid representations of molecules. Experiments are conducted to show that the proposed method can successfully generate valid 3D molecules.", "questions": "In experiments, methods are evaluated by the average number of aromatic rings per generated molecule. Do aromatic rings are detected from 2D molecule graphs inferred from 3D structures? Are the generated aromatic rings are checked in 3D space to decide whether they respect chemical constraints (i.e., all atoms in an aromatic ring should be in the same plane)? Why this metric is used? What can be shown if a method can generate more aromatic rings?", "limitations": "Yes.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- This work proposes a novel method of generating 3D molecules in the form of 3D voxel grids by neural empirical Bayes framework. The exploration of generating 3D voxel grids representation of molecules is interesting and meaningful to the development of efficient drug discovery applications.\n- The writing of this paper is well-organized and easy to follow.", "weaknesses": "- In this paper, generating 3D molecules in the form of 3D voxel grids by VoxMol is argued to be advantageous compred with 3D molecular graphs because the number of atoms is required to know beforehand, the distributions of continuous atom coordinates and discrete atom features are not treated separately, and simpler and faster training and generation can be achieved. However, in my opinion, these are not strong advantages.  \n(1) Actually, the number of atoms can be sampled from a learned distribution in 3D graph generation model, and VoxMol is also implicitly learn the number of atoms to place in every channel of voxel grids. These two processes are not fundamentally different.  \n(2) It is not a big issue to handle discrete and continuous distributions separately as we can jointly train discrete and continuous generative models (e.g., discrete & contiuous diffusion models). Also, VoxMol need to implicitly learn a \"discrete distribution\" as for every 3D spatial location in the voxel grid, it needs to decide the atom type, i.e., which channel to place an atom.  \n(3) The simpler and faster training and generation are mainly due to the use of neural empirical Bayes method, which may also be applied to 3D molecular graph generation to achieve fast generation and training. Also, 3D molecular graph generation can be accelerated by some other methods like introducing fragment-based generation.  \nHence, it is not convincing that the proposed VoxMol method is really advantageous when compared to existing graph-based method.\n- The use of 3D voxel grids also introduce many disadvantages. 3D voxel grids lose SE(3)-invariance as they are not invariant to rotations or translations. The use of 3D voxel grids hamper the introduction of chemical domain knowledge to guide the generative models to produce chemically valid or drug-like molecules, such as adding constraints based on fragments. The effect of these disadvantages should be evaluated or discussed.\n- The experimental results in Table 1 do not demonstrate the advantages of the proposed VoxMol method as VoxMol does not outperform strong  baseline methods in several metrics.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback and suggestions. Below we address additional concerns from the reviewer.\n\n**Implicitly learning the number of atoms is not an advantage:**\n \nWe agree with the reviewer that VoxMol is implicitly learning the number of atoms (together with all other information about molecules). This might be arguable, but we assumed that learning representations more implicitly/end-to-end is simpler and can generalize better—therefore, we counted this as an advantage (R4 agrees: \"This comes with benefits, chief among which is the ability to generate without a known number of atoms.\"). In applications related to in-painting (eg, pocket conditioning, linking, scaffold conditioning), learning implicitly the number of atoms could be very helpful. We will tone down this contribution by mentioning instead that VoxMol learns implicitly the number of atoms (in contrast to EDM-based methods, that need to learn them explicitly).\n\n**Score functions on continuous-only vs continuous/discrete space:**\n \nBecause the score function is not defined on discrete spaces, it can be non-trivial to adapt the score-based modeling to categorical data. This is currently an active research area (eg [E-H]) and we still don't have a clear winner solution to the problem. It has also been pointed ([F]) that it can be more challenging to condition the generation of score-based models on discrete spaces. Therefore, as we agree with the reviewer that it is possible to learn score function on continuous-discrete space, we believe that training only on continuous space is simpler. \n\n**Simpler training because to NEB is simpler than diffusion:**\n \nWe agree with the reviewer. In fact, we mention in the text that VoxMol (based on NEB) is simpler to train than _point-cloud diffusion models_ (L54-55 and L154), but not a general advantage of voxels over point clouds.\nWe also point out that NEB requires a very large noise to work and denoising point clouds can be more challenging than denoising discrete grids (U-Nets works extremely well for this task and has been highly fine-tuned for this task with large image generation models).\n\n**Experimental results:** \n\nThe main conclusion from our experiments is that VoxMol underperforms EDM by a small margin on QM9 (Table 1) and outperforms EDM on GEOM-drugs (larger/more realistic dataset) by a larger margin (Table 2). This conclusion holds even further with the MiDi metrics suggested by the reviewers (see Table 1 and 2 on the attached pdf). \n\n**\"3D voxel grids lose SE(3)-invariance as they are not invariant to rotations or translations.\"... \"should be discussed\"** [we assume the reviewer means \"equivariance\" instead of \"invariance\". Please let us know if that is not the case]: \n\nOur method has built-in equivariance to translation and relies on data augmentation through rotation to learn rotation equivariance (as it has been done before). We mentioned in the manuscript that VoxMol (being based on CNNs) does not have rotation equivariance built-in in the architecture (L113-117 and L217-218). For clarity, we will include an explicit section mentioning the limitations of the model on the revised version of the paper. \n\n**Guide generation:**\n \nAdding constraints to guide VoxMol generation can in principle be done in the same way as in-painting or guidance is done with diffusion models on images (eg, Song et al. 20, Lugmayr et al. 22). Please see general rebuttal for discussion in guided/conditional sampling.\n\n**Measuring number of generated rings:**\n \nHoogeboom et al., 22 (appendix) points that predicting single bonds can easily increase validity. This metric was added to show that VoxMol achieves its performance without \"cheating\" and oversampling single bounds–in fact, it matches better the average number of rings per molecule from the ground-truth dataset. Fig 8 in the submission shows that EDM undersample aromatic bonds and oversamples single and double bonds, while VoxMol matches better the number of bonds. We recognize the rationale for this metric was not clearly explained in the text. We will move this discussion to the appendix of the paper and report the same metrics as in the MiDi paper, as suggested by the reviewers.\n\n[E] Argmax Flows and Multinomial Diffusion, Hoogeboom et al, Nuerips21\n\n[F] Continuous diffusion for categorical data, Dieleman et a., 22\n\n[G] Analog bits, Chen et al., ICLR23\n\n[H] Geometric Latent Diffusion Models for 3D Molecule Generation, Xu et al, ICML23\n"}, {"review_id": "LSOLtfo9gu", "reviewer": "Reviewer_zHHZ", "summary": "This paper describes a new method of unconditional small molecule generation by parameterizing small molecules as 3d voxel arrays.", "questions": "- Line 243 should be as noise *decreases*, correct?\n\n- We consistently see a large drop in the validity when the noise level goes from $\\sigma = 0.9$ to $\\sigma = 1$, why is such a small change in noise so detrimental? What are the characteristics of the generated molecules at these levels of noise? What happens if you train a model with $\\sigma = 2$?\n\n- What is the relevance of the number of aromatic rings per molecule? Especially given that this is one of the only metrics where the proposed model outperforms, I would like a good rationale for including this metric.\n\n- What is the relevance of the validity metric, if that is something that can be automatically detected? It seems in this case that generating unique molecules is more relevant, especially if there are no other metrics to optimize.\n\n- How can this model be extended to allow for conditional generation, given that most of the other methods being compared to have this ability?", "limitations": "Mainly, the authors should find better metrics to evaluate their work, such as those mentioned in the MiDi paper (at a minimum).", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- Novel characterization of the small molecule generation task and a new way to parameterize the data. This comes with benefits, chief among which is the ability to generate without a known number of atoms.\n\n- I appreciate the discussion in section 4.3, where the authors dive into some problems with the metrics that they have chosen.\n\n- The ablation studies are illuminating.\n\n- The paper is well-written and easy to follow.", "weaknesses": "- Unconditional generation of small molecules is a hard task to evaluate. There are no good metrics to optimize for, as the paper indicates. Nevertheless, MiDi (work that the authors cite and compare against) has a more representative set of metrics they compare against. I suggest the authors do the same.\n\n- Weaknesses of the one-step denoising procedure are not discussed. These include, for example, the inability to control generation through the use of auxiliary gradient signals, such as classifier-free guidance.\n\n- The authors don't mention the implications of having a fixed grid size for their generative model. This means that the molecules generated need to be within a specified volume, restricting the output space. They should contrast this to the fixed atom number of competing models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback and suggestions. Below we address additional concerns from the reviewer.\n\n**Metrics:**\n \nWe agree with the reviewers that unconditional molecule generation is hard to evaluate. We followed the reviewer's suggestion and evaluated VoxMol on the MiDi metrics (Table 1 and 2). VoxMol is slightly worse than EDM on the simple QM9 while performing much better than EDM on the more challenging/realistic GEOM-drugs dataset. We will update the tables and plots of the experiments to reflect the new metrics.\n\n**Limitations of one-step denoising:**\n \nNeural empirical Bayes (NEB) performs denoising in 1 step while diffusion models do in multiple (eg, 1000) steps.\nBecause it is a single-step denoising process (from a high noise level), it can be challenging to capture high-frequency components of the signal (eg, generated images are usually more blurred than diffusion model samples). We observed empirically that VoxMol is able to denoise voxelized molecules (they are signals with much \"lower frequencies\"  than natural images), and therefore NEB is a good fit for this problem–as we show in experiments. We briefly mention this on L154-158, but will make it more explicit in the text.\n\n**Conditional generation:**\n \nBoth NEB and diffusion models use a (learned) score function to sample with Langevin MCMC.. Therefore, we can in principle also guide the sampling by constraining the score function at each step of the MCMC chain. \nThe Langevin MCMC is used very differently in the two approaches. In DM, the chains transition from pure noise to the molecule (that is, one chain per sample). In NEB, the chain keeps sampling noisy molecules from p(y) (and we get clean molecules by denoising them with xhat). Please see the paragraph on guided/conditional sampling on the main rebuttal.\n\n**No mention about the implications of fixed grid size:**\n \nWith a fixed voxel grid, we are limited to a fixed volume in space (we mentioned this on L127, but will be more clear). With a 64^3 voxel grid (at resolution .25A), we can fit over 99% of the drug-like molecules from GEOM-drugs dataset. We wrote about the pros/cons of voxels vs point cloud (L100-117), but will be more specific on this tradeoff between max volume limitation/number of atoms.\n\n**Drop in validity from sigma=.9 to sigma=1:**\n \nAs we increase sigma, denoising gets more and more difficult to learn. At the same time, the smoother the space of noisy molecules is (ie, higher the sigma), the easier it gets to sample from this space with Langevin MCMC. In practice, we need to find the highest sigma such that the U-Net can still learn to denoise. sigma=0.9 is the best value we found on our hyperparameter search (on the validation set). We also found this surprising; it seems that higher sigma values work on natural images (Saremi and Hyvarinen 19, Saremi and Srivastava, 22). We hypothesize this has to do with the sparsity of voxelized molecules compared to natural images, but it needs more investigation. We mention this in L174-176, but will provide a better discussion.\n\n**Measuring number of generated rings:**\n \nHoogeboom et al., 22 (on appendix) points that predicting single bonds can easily increase validity. This metric was added to show that VoxMol achieves its performance without \"cheating\" and oversampling single bounds–in fact, it matches better the average number of rings from the ground truth dataset. Fig 8 shows that EDM undersample rings and oversamples single and double bonds, while VoxMol matches better the number of bonds. We recognize the rationale for this metric was not clearly explained in the text. We will move this discussion to the appendix of the paper and report the same metrics as in the MiDi paper, as suggested by the reviewers.\n\n**L243:** yes! Thanks for correcting.\n\n**Relevance of validity metrics:**\n \nWe used validity metrics as they correlate with other metrics we measured. Following the reviewer's suggestion, we will report the MiDi metrics (Table 1 and 2 on the attached pdf) on the updated version of the paper.\n"}]

Author rebuttal:
(R IDs: R1=R82i , R2=j1qM, R3=bpDM, R4= zHHZ)

We thank the reviewers for the detailed and helpful reviews. In particular, we thank the suggestion of evaluating on MiDi metrics [R1, R4]. The results on these metrics corroborate the findings of the submission: VoxMol performs slightly worse than EDM on QM9 (Table 1 on the attached pdf), while considerably outperforming EDM on harder GEOM-drugs (Table 2). The metrics shows more clearly that VoxMol (i) beats the current state of art on a (arguably) more realistic/useful dataset _by a large margin_, and (ii) seems to be more scalable. We also acknowledge that reviewers find that this work (iii) is novel in terms of data representation, training and/or sampling [R1, R2, R3, R4], (iv) has good evaluations/ablations [R2, R4], (v) is well written [R2, R3, R4], and (vi) can "shed light on new research directions of connecting 3D CNNs to equivariant GNNs" [R2].

PS: we fixed a bug on the rotation augmentation (randomly rotates the axis between [0, 2] instead of [0,2pi] radians), see `_random_rot_matrix()` function on `./dataset/dataset.py` (submitted source code). The bug fix improves the uniqueness score while other metrics mostly remain the same. Table 1,2 shows results with VoxMol trained with the correct rotation augmentation.

Next, we address the main concerns from reviewers.

**MiDi metrics [R1, R4]:**

First, we run the official MiDi evaluation on EDM's .xyz files (provided by authors) to reproduce results from the MiDi paper. Then, we run the same evaluation on VoxMol .xyz files. Table 1 and 2 on the attached pdf show results for QM9 and GEOM-drugs respectively. To summarize: 
* In QM9, VoxMol performs slightly worse than/similar to EDM in most metrics (except on bond angle W1 where the gap is large).
* In GEOM, VoxMol is considerably better (by larger margins in most cases) than EDM in all metrics except on bond length W1, where they perform relatively close. 

We will use these tables/metrics on the revised version of the manuscript. 

**Comparison w. MiDi [R1]:** 

MiDi is built on the top of EDM by leveraging molecular graphs (2D) as well as 3D atoms. Its contributions are therefore orthogonal to ours. In fact, a "VoxMol+MiDi" model could improve results the same way MiDi (3D+2D) improved over EDM (3D). Table1,2 show MiDi results for reference. We note that VoxMol outperforms MiDi (published on ICLR23 MLDD Workshop) on GEOM-drugs in most metrics _without_ using molecular graphs or formal charges information during training. 

**SE(3) equivariance [R1, R2, R3]:** 

Built-in rotation equivariance is a good property for a network to have, however equivariance can also be learned with strong data augmentation/larger datasets ([A, B, C] and related works). By choosing a 3D U-Net architecture, we lose built-in rotation equivariance but win on expressiveness of the denoising network (U-Nets have been highly optimized for denoising grids, eg, modern image generation). Our experiments show that an efficient denoiser can scale up better, allowing VoxMol to outperform current SOTA on GEOM-drugs despite not having built-in rotation equivariance.

3D CNNs can in theory be adapted to have built-in rotation equivariance (Weiler et al, 18). Successfully designing SE(3)-equivariant 3D CNNs for this problem can potentially improve results further. Currently, equivariant 3D CNNs do not scale as well as the standard ones (and it is a challenge to apply to large datasets). We hope our results motivate the community to explore more scalable SE(3)-equivariant 3D CNNs. We see these developments as architecture improvements and not as the main contributions of the paper.

**Guide/Conditional generation [R3, R4]:**

Like diffusion models (DM), our method also leverages (learned) score functions and relies on Langevin MCMC for sampling. Therefore, in theory we can condition VoxMol similarly to how it is done in diffusion models: by constraining the score function as we walk through the MCMC chain. In the case of DMs, the score function of all T steps are constrained to guide the transition steps from noise to a (conditioned) sample. In VoxMol, the constrained score function would affect the "walk steps" (the Langevin MCMC steps): it would restrict the region where the chain samples noisy molecules y to P(y|c) (instead of P(y)), c is the condition (eg, gradient of a classifier). The "jump step" (a forward pass of the denoising network over the noised molecules) is independent of the condition and remains unchangeable.

Many of the innovations on conditioning DMs come from computer vision, where U-nets are usually used. Since VoxMol has the same architecture (albeit 3D instead of 2D), many of the conditioning techniques/tricks used in images may be more easily transferable. Eg, we could in principle use the gradient of a classifier (trained jointly) to guide the sampling (using the same trick as in Dhariwal and Nichol 21) or adapt gradient-free guidance (Ho and Salisman, 22). Pocket conditioning could also be possible, as in eg, (Schneuing et al. 22, Guan et al. 23). In-painting (related to linker/scaffold/fragments cond.) has also proven to work very well in 2D U-Nets, so it could potentially work with 3D U-Nets as well.

We will add a section in the appendix on how we can condition the generation in different ways. We agree that conditional molecule generation is a more interesting problem in practice. We note that our approach is the first of its kind and the main focus of this submission is to show that (i) it is a feasible approach (this is non-trivial) and (ii) it scales well on unconditional generation, beating current SOTA on a large dataset.

[A] Understanding image representations by measuring their equivariance and equivalence, Lenc,Vedaldi, CVPR15

[B] Grounding inductive biases in natural images: invariance stems from variations in data, Bouchacourt,et al,NeurIPS21

[C] The Lie Derivative for Measuring Learned Equivariance. Gruveret al,22


Meta-review: This paper presents a new approach to molecular generation based on denoising voxel grids. By using a voxel-grid representation, this work takes a vision-based approach to molecular generation instead of graph-based or point clouds methods. The method has strong empirical performance. The paper is well written and includes source code.  While the reviews are mixed, the authors provide extensive discussion during rebuttals that I felt addressed their concerns. Finally, it is unclear how useful the task of unconditional generation is in practice, however, the authors have stated that they will include a discussion of different ways to condition the generation process, which will be a valuable contribution.
Decision: Accept (poster)