SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Architecture. The denoising network  $\hat{x}_{\theta}$  is used in both the walk and jump steps described above. Therefore, its parametrization is very important to the performance of this approach. We use a 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 4 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  $\sigma = .9$  , unless stated otherwise. We train our models with batch size of 128 and 64 (for QM9 and GEOM-drugs, respectively) and we use AdamW [88]
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 \times 10^{-5}$ , weight decay  $10^{-2}$ ) to optimize the weights. The weights are updated with exponential moving average with a decay of .999. We use  $\gamma = 1.0$ ,  $u = 1.0$  and  $\delta = .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 contain 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^3$  and  $64^3$  for QM9 and GEOM-drugs respectively. These volumes are able to cover over  $99.8\%$  of all points on both datasets. All our models model hydrogens explicitly. For QM9, we consider all 5 chemical elements (C, H, O, N and F) present on the dataset. For GEOM-drugs, we consider 8 elements (C, H, O, N, F, S, Cl and Br). We ignore P, I and B elements as they appear in less than  $1\%$  of the molecules in the dataset. Finally, the input voxel grids are of dimension  $\mathbb{R}^{5\times 32\times 32\times 32}$  and  $\mathbb{R}^{8\times 64\times 64\times 64}$  for QM9 and GEOM-drugs, respectively. We perform the same pre-processing and dataset split as [20] and end up with 100K/20K/13K molecules for QM9 and 1.1M/146K/146K for GEOM-drugs (train, validation, test splits respectively).
Baselines. We compare our method with two state-of-the-art approaches:  $GSchNet$  [47], a point-cloud autoregressive model and  $EDM$  [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 successfully exploiting equivariant 3D convolutional networks. We also show results of  $\mathrm{VoxMol}_{\mathrm{oracle}}$  in our main results, where we assume we have access to real samples from the noisy distribution. Instead of performing MCMC to sample  $y_{k}$ , we sample molecules from the validation set and add noise to them. This baseline assumes we would have perfect sampling of noisy samples (walk step) and let us assess the quality of our model to recover clean samples. It serves as an upper bound for our model and allows us to disentangle the quality of the walk (sampling noisy samples) and jump (estimating clean molecules) steps.
All methods generate molecules as a 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] and use standard cheminformatics software to determine the molecule's atomic bonds given the atomic coordinates  $^6$ . Using the same post-processing for all methods allows a more apples-to-apples comparison of 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 sanitization filter; uniqueness, the proportion of valid molecules that have different canonical SMILES; valency  $W_{1}$ , the Wasserstein distance between the distribution of valencies in the generated and test set; atoms  $TV$  and bonds  $TV$ , the total variation
<sup>6</sup>We follow previous work and use OpenBabel [89] to get the bond orders.
<sup>7</sup>We use the official implementation from https://github.com/cvignac/MiDi.
stable mol %↑	stable atom %↑	valid %↑	unique %↑	valency W1↓	atom TV↓	bond TV↓	bond len W1↓	bond ang W1↓
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
VoxMolno rot	84.2 (±1.6)	98.2 (±.3)	98.1 (±.4)	77.2 (±1.7)	.043 (±.0)	.171 (±.200)	.050 (±.010)	.007 (±.0)	3.80 (±.7)
VoxMol	89.3 (±.6)	99.2 (±.1)	98.7 (±.1)	92.1 (±.3)	.023 (±.002)	.029 (±.009)	.009 (±.002)	.003 (±.002)	1.96 (±.04)
VoxMoloracle	90.1	99.3	98.9	99.9	.024	.009	.002	.001	0.37
Table 1: 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  $W_{1}$  and bond angle  $W_{1}$ , the Wasserstein distance between the distribution of bond and lengths, respectively. Finally, we also report the strain energy metric proposed in [91]. This metric is defined as the difference between the internal energy of the generated molecule's pose and a 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 1 and Table 2 show results on QM9 and GEOM-drugs respectively. We report results for models trained with and without data augmentation (VoxMol and VoxMol $_{\text{no rot}}$ , respectively) and generate 10,000 samples with multiple MCMC chains. Each chain is initialized with 1,000 warm-up steps, as we observed empirically that it 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 metrics while performing worse in others (specially stable molecule, uniqueness and angle lengths). On GEOM-drugs, a more challenging and realistic drug-like dataset, the results are 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 of different models on QM9 and GEOM-drugs, respectively. The closer the CDF of generated molecules from a 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, respectively. On QM9, all models have median strain energy around the same ballpark: 52.58, 66.32 and 56.54 kcal/mol for EDM,  $\mathrm{VoxMol}_{\mathrm{no rot}}$  and  $\mathrm{VoxMol}$ , respectively. On GEOM-drugs, 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  $\mathrm{VoxMol}_{\mathrm{no rot}}$  and  $\mathrm{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 to 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 learns the approximate distribution of the number of atoms per molecule on both datasets. Implicitly learning this distribution can be particularly useful in applications related to 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 of 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 advantageous in capturing the modes of the distribution we want to model. Compared
Figure 5: The cumulative distribution function of strain energy of generated molecules on (a) QM9 and (b) GEOM-drugs. For each method, we use 10,000 molecules.
stable mol %↑	stable atom %↑	valid %↑	unique %↑	valency W1↓	atom TV↓	bond TV↓	bond len W1↓	bond 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
VoxMolno rot	44.4 (±.1)	96.6 (±.1)	89.7 (±.2)	99.9 (±.0)	.238 (±.001)	.025 (±.001)	.024 (±.001)	.004 (±.000)	2.14 (±.02)
VoxMol	75.0 (±.1.)	98.1 (±.3)	93.4 (±.5)	99.1 (±.2)	.254 (±.003)	.033 (±.041)	.036 (±.006)	.002 (±.001)	0.64 (±.13)
VoxMoloracle	81.9	99.0	94.7	97.4	.253	.002	.024	.001	0.31
Table 2: Results on GEOM-drugs. We use 10,000 samples from each method. Our results are shown with mean/standard deviation across 3 runs.
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 VoxMol<sub>oracle</sub> results, we see that VoxMol can still be vastly improved. We can 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 SE(3)-equivariance CNNs, etc).
4.3 Ablation studies
Noise level  $\sigma$ . 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  $\sigma$  in  $\{.6, .7, ..., 1.2\}$ , while keeping all other hyperparameters the same. Figure 7(a,b,c) shows how noise level  $\sigma$  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 a value. We observe empirically that  $\sigma = .9$  is the sweet spot level that achieves better overall performance on the validation set of QM9.
Number of steps  $\Delta k$ . Table 3 shows how VoxMol's performance on GEOM-drugs change with the number of walk steps  $\Delta k$  on the Langevin MCMC sampling. In 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 as  $\Delta k$  increases. The average time (in seconds) to generate a
Figure 7: Effect of noise level  $\sigma$  on generation quality. Models are trained on QM9 with a different noise level. Each plot shows two metrics: (a) molecule stability and uniqueness, (b) atom and bond TV, (c) valency and angle lengths W1.
Δk(n steps)	stable mol %↑	stable atom %↑	valid %↑	unique %↑	valency W1↓	atom TV↓	bond TV↓	bond len W1↓	bond ang W1↓	avg. t s/mol.↓
50	78.9	98.7	96.3	87.8	.250	.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	93.8	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  $\Delta k$  on generation quality on GEOM-drugs (2,000 samples). EDM results are shown for comparison.
molecule increases linearly with the number of steps, as expected. We observe that even using 500 steps, our model is still faster than EDM on average, while achieving better performance in 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 .25 Å). 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—throughout different versions of the model, with different hyperparameters—that using a fixed radius of .5 consistently outperform other values. Training does not converge with radius .25 and quality of generated samples degrades as we 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.
We introduce VoxMol, a novel score-based method for 3D molecule generation. This method generates molecules in a fundamentally different way than the current state of the art (i.e., diffusion models applied to atoms). The noise model used is also novel in the class of score-based generative models for molecules. We represent molecules on regular voxel grids and VoxMol is trained to 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 of 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 a pure algorithmic way): a problem that can be seen as an initial stepping stone (out of many) to this long-term objective. We, as a society, need to find solutions to use these technologies in ways that are safe, ethical, accountable and exclusively beneficial to society. These are important concerns and they need to be thought of at the same time we design machine learning algorithms.
Acknowledgements. The authors would like to thank the whole President Design team for helpful discussions and Genentech's HPC team for providing a reliable environment to train/analyse models.
[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 arXiv:2308.07413, 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 the 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-Romana, 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, and Kenji Doya. Sigmoid-weighted 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 NeurIPS, 2019.
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-like densities:
$$
V _ {a} (d, r _ {a}) = \exp \left(- \frac {d ^ {2}}{(. 9 3 \cdot r _ {a}) ^ {2}}\right), \tag {5}
$$
where  $V_{a}$  is defined as the fraction of occupied volume by atom  $a$  of radius  $r_a$  at distance  $d$  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  $r_a = .5\AA$ . The occupancy of each voxel in the grid is computed by integrating the occupancy generated by every atom in a molecule:
$$
\operatorname {O c c} _ {i, j, k} = 1 - \prod_ {n = 1} ^ {N _ {a}} \left(1 - V _ {a _ {n}} \left(\left| \left| C _ {i, j, k} - x _ {n} \right| \right|, r _ {a _ {n}}\right)\right), \tag {6}
$$
where  $N_{a}$  is the number of atoms in the molecule,  $a_{n}$  is the  $\mathfrak{n}^{th}$  atom,  $C_{i,j,k}$  are the coordinates (i,j,k) in the grid and  $x_{n}$  is the coordinates of the center of atom  $n$  [93]. The occupancy takes the maximum value of 1 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 PyUUL [94] to generate the voxel grids from the raw molecules (.xyz or .sdf format).
We use grids with  $32^{3}$  voxels on QM9 and  $64^{3}$  on GEOM-drugs 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, N and F), while for GEOM-drugs, we use 8 (C, H, O, N, F, S, Cl and Br). We model hydrogen explicitly in all our experiments. Finally, the input voxel grids are of dimension  $\mathbb{R}^{5\times 32\times 32\times 32}$  and  $\mathbb{R}^{8\times 64\times 64\times 64}$  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\pi)$ ) 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,.25]\mathring{\mathrm{A}}$ .
A.2 Architecture
Our neural network architecture follows standard encoder-decoder convnet architecture. We use a very similar architecture recipe to DDPM [34]. The model uses four levels of resolution:  $32^3$  to  $4^3$  for the QM9 dataset and  $64^3$  to  $8^3$  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 \times 3 \times 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 \times 3 \times 3$ ). All convolutions have stride 1 and we pad the feature maps with 1 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  $2 \times 2 \times 2$  (stride 1) max-poolings (bilinear-upsampling). 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 111M 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 SE(3)-equivariant 3D U-Net using steerable convnets [72] for denoising, but initial experiments were not successful. The hypothesis is that a built-in SE(3) equivariant version of our model,  $\mathrm{VoxMole_{equi}}$ , 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 metrics on QM9. We then use the same procedure to
generate samples as described in the main 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. Table 4 compares the results of the model with and without built-in SE(3) equivariance.
QM9	stable mol %↑	stable atom %↑	valid %↑	unique %↑	valency W1↓	atom TV↓	bond TV↓	bond len W1↓	bond ang W1↓
VoxMol	89.3	99.2	98.7	92.1	.023	.029	.009	.003	1.96
VoxMolequi	25.1	81.8	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 is higher than the non-equivariant (approx.  $9.4 \times 10^{-5}$  vs.  $5.4 \times 10^{-5}$  MSE on val. set), (ii) the equivariant model needs more capacity to be competitive with the non-equivariant one (currently it has over  $90 \times$  fewer parameters), (iii) something in the sampling procedure needs to be different on 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/better 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^{3}$  voxel grid).
A.4 Training and sampling
The weights are optimized with batch size 128 and 64 (for QM9 and GEOM-drugs, respectively), AdamW optimizer  $(\beta_{1} = .9, \beta_{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 1 describes this process.