SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Pedro O. Pinheiro, Joshua Rackers, Joseph Kleinhenz, Michael Maser, Omar Mahmood, Andrew Martin Watkins, Stephen Ra, Vishnu Sresht, Saeed Saremi
Prescient Design, Genentech
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 [1] 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 (i.e., 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.
Finding novel molecules with desired properties is an important problem in chemistry with applications to 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^{60}$  [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 to trillions molecules) and these approaches cannot propose new molecules conditioned on some desiderata.
Generative models for molecules have been proposed to 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 SELFFIES [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 a 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 to be very effective in molecular discriminative tasks—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 a score function is undefined on discrete distributions,
Figure 1: Voxelized molecules generated by our model and their corresponding molecular graphs. Left, samples from a model trained on QM9 dataset  $(32^{3}$  voxels). Right, samples from a model trained on GEOM-drugs  $(64^{3}$  voxels). In both cases, each voxel is a cubic grid with side length of .25Å. Each color represents a different atom (and a different channel on the voxel grid). Best seen in digital version. See appendix for more generated samples.
some workaround is 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-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 tradeoffs between model expressivity and built-in equivariance).<sup>1</sup>
In this work we introduce VoxMol, a new score-based method to generate 3D molecules. Similar to [33], and unlike most recent approaches, we represent atoms as continuous (Gaussian-like) densities and molecules as 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 score-based generative models on images, e.g. [34, 35, 36]—are very effective and scale very well with data.
We start by training a neural network to denoise noisy voxelized molecules. Noisy samples are created simply by adding Gaussian noise (with a 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, 38] and 1D amino acid sequences [39].
Compared to point-cloud diffusion models, VoxMol is simpler to 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 to large, drug-sized molecules. Figure 1 (and Figures 8, 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—usually 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 a small dataset (QM9 [40]), while outperforms it (by a large margin) on a challenging, more realistic drug-like molecules dataset (GEOM-drugs [41]).
<sup>1</sup>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 a higher degree of equivariance compared to convolutional architectures [32].
Voxel-based unconditional 3D molecule generation. Skalic et al. [42] and Ragoza et al. [33] map atomic densities on 3D regular grids and train VAEs [43] using 3D convolutional networks to generate voxelized molecules. To recover atomic coordinates from the generated voxel grids², [33] introduces a simple optimization-based 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 to train a pocket-conditional model. Similar to these works, our model also relies on discretization of 3D space. Like [33], we use a simple peak detection algorithm to extract atomic coordinates from the generated voxel grids. However, our method differs 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, 51] use normalizing flows [52]. Hoogeboom et al. [15] proposes E(3) Equivariant Diffusion Models (EDM), a diffusion [53]-based 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/scaffolds 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 score-based generative modeling for images. These models are known to perform well and scale nicely with data. Second, message passing on graphs operates 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 a discussion on the many-body representation hypothesis). Third, voxel representations have a 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 to those on GEOM-drugs [41]) with reasonable resolution  $(.1 - .2\AA)$  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 a problem in practice as we can achieve competitive 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)-equivariance 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.
volitional network have been proposed [72, 73, 74]<sup>3</sup>, but current models do not scale as well as standard convnets, and it would be a challenge to apply them to drug-sized molecules. Built-in rotation equivariance is a good property to have, however equivariance can also be learned with strong data augmentation/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. However, 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) is represented by a different grid channel. See appendix for more information on how we discretize molecules. This discretization process gives us a dataset with  $n$  voxelized molecules  $\{x_{i}\}_{i = 1}^{n}, x_{i} \in \mathbb{R}^{d}, d = c \times l^{3}$ , where  $l$  is the length of each grid edge and  $c$  is the number of atom channels in the dataset. Each voxel in the grid can take values between 0 (far from all atoms) and 1 (at the center of atoms). Throughout our experiments, we consider a fixed resolution of .25 Å (we found it to be a good trade-off between accuracy and computation). Therefore, voxel grids occupy a volume of  $(l / 4)^{3}$  cubic Ångströms.
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  $\sigma^2 I_d$ . Equivalently,  $Y = X + N$ , where  $X \sim p(x)$ ,  $N \sim \mathcal{N}(0, \sigma^2 I_d)$ . Therefore  $Y$  is sampled from:
$$
p (y) = \int_ {\mathbb {R} ^ {d}} \frac {1}{(2 \pi \sigma^ {2}) ^ {d / 2}} \exp \left(- \frac {\| y - x \| ^ {2}}{2 \sigma^ {2}}\right) p (x) d x.
$$
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 least-square estimator of  $X$  is the Bayes estimator, i.e.,  $\hat{x}(y) = \mathbb{E}[X|Y = y]$ . Built on this result, Miyasawa [81] showed that, if the noisng process is Gaussian (as in our case), then the least-square estimator  $\hat{x}(y)$  can be obtained purely from the (unnormized) smoothed density  $p(y)$ :
$$
\hat {x} (y) = y + \sigma^ {2} g (y), \tag {1}
$$
where  $g(y) = \nabla_y\log p(y)$  is the score function [82] of  $p(y)$ . This interesting equation tells us that, if we know  $p(y)$  up to a 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  $\hat{x} (y)$ , we can compute the score function of  $p(y)$  via (1).
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(y)$  and the least-square estimator  $\hat{x}(y)$  from a dataset of voxelized molecules  $\{x_i\}_{i=1}^n$ , sampled from unknown  $p(x)$ . We leverage the (learned) estimator and score function to generate voxelized molecules in two steps: (i) sample  $y_k \sim p(y)$  with Langevin MCMC [83], and (ii) generate clean samples with the least-square estimator. The intuition is that it is much easier to sample from the smooth density than the original distribution. See Saremi and Hyvarinen [1] for more details.
3.2 Denoising voxelized molecules
We parametrize the Bayes estimator of  $X$  using a neural network with parameters  $\theta$  denoted by  $\hat{x}_{\theta}:\mathbb{R}^{d}\to \mathbb{R}^{d}$ . Since the Bayes estimator is the least-squares estimator, the learning becomes a least-squares denoising objective as follows:
$$
\mathcal {L} (\theta) = \mathbb {E} _ {x \sim p (x), y \sim \mathbb {N} \left(x, \sigma^ {2} I _ {d}\right)} | | x - \hat {x} _ {\theta} (y) | | ^ {2}. \tag {2}
$$
<sup>3</sup>We made an attempt at using an equivariant 3D convnets for denoising, but initial experiments were not successful. See appendix for details.
<sup>4</sup>We use the convention where we drop random variable subscripts from probability density function when the arguments are present:  $p(x) \coloneqq p_X(x)$  and  $p(y) \coloneqq p_Y(y)$ .
Figure 2: (a) A representation of our denoising training procedure. Each training sample (i.e., a voxelized molecule) is corrupted with isotropic Gaussian noise with a fixed noise level  $\sigma$ . The model is trained to recover clean voxel grids from the noisy version. To facilitate visualization, we threshold the grid values,  $\hat{x} = \mathbb{1}_{\geq 1}(\hat{x})$ . (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 of the denoising network<sup>5</sup>:
$$
g _ {\theta} (y) = \frac {1}{\sigma^ {2}} \left(\hat {x} _ {\theta} (y) - y\right). \tag {3}
$$
By minimizing the learning objective (2) we learn the optimal  $\hat{x}_{\theta}$  and by using (3) we can compute the score function  $g_{\theta}(y)\approx \nabla_y\log p(y)$ .
We model the denoising network  $\hat{x}_{\theta}$  with an encoder-decoder 3D convolutional network that maps every noised voxel on the grid to a clean version of it. Figure 2(a) shows a general overview of the denoising model. The noise level,  $\sigma$ , is kept constant during training and is a key hyperparameter of the model. Note that in the empirical Bayes formalism,  $\sigma$  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 to reconstruct voxelized molecules (within the noise levels considered in this paper). Our hypothesis is that this is 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  $g_{\theta}$  and the estimator  $\hat{x}_{\theta}$  to sample. We follow the walk-jump sampling scheme [1, 37, 38, 39] to generate voxelized molecules  $x_{k}$ :
(i) (walk step) For sampling noisy voxels from  $p(y)$ , we consider Langevin MCMC algorithms that are based on discretizing the underdamped Langevin diffusion [84]:
$$
d v _ {t} = - \gamma v _ {t} d t - u g _ {\theta} \left(y _ {t}\right) d t + (\sqrt {2 \gamma u}) d B _ {t} \tag {4}
$$
$$
d y _ {t} = v _ {t} d t,
$$
where  $B_{t}$  is the standard Brownian motion in  $\mathbb{R}^d$ ,  $\gamma$  and  $u$  are hyperparameters to tune (friction and inverse mass, respectively). We use the discretization algorithm proposed by Sachs et al. [85] to generate samples  $y_{k}$ , which requires a discretization step  $\delta$ . See appendix for a description of the algorithm.
(ii) (jump step) At an arbitrary time step  $k$ , clean samples can be generated by estimating  $X$  from  $y_{k}$  with the denoising network, i.e., computing  $x_{k} = \hat{x}_{\theta}(y_{k})$ .
This approach allows us to approximately sample molecules from  $p(x)$  without the need to compute (or approximate)  $\nabla_{x}\log p(x)$ . In fact, we 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, 86]. Figure 2(b) shows a schematic representation of the generation process. Following [37], we initialize the chains at by adding uniform noise to Gaussian noise (with the same  $\sigma$  used during training), i.e.,  $y_0 = N + U$ ,  $N\sim \mathcal{N}(0,\sigma^2 I_d)$ ,  $U\sim \mathcal{U}_d(0,1)$  (this was observed to mix faster in practice).
5 Alternatively, one can also parameterize the score function:  $g_{\theta}:\mathbb{R}^{d}\to \mathbb{R}^{d}$ , in which case  $\hat{x}_{\theta}(y) = y + \sigma^2 g_{\theta}(y)$ .
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 a 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 a walk-jump sampling chain, where generated molecules change gradually as we walk through the chain (the clean samples are shown every ten steps,  $\Delta k = 10$ ). This figure is 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/disappear/change as we move through the chain. Interestingly, this behavior happened on most chains we looked into explicitly.
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 0 all voxels with value less than .1, i.e.,  $x_{k} = \mathbb{1}_{\geq .1}(x_{k})$ . Then we run a simple peak detection to locate the voxel on 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 our pipeline to recover atomic coordinates and molecular graphs from generated voxelized molecules. See appendix for more details.
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 a feasible approach for unconditional generation (this is non-trivial) and (ii) it scales well with data, beating a established model on a large, drug-like dataset. In principle, VoxMol can be used for guided (or conditional) generation, an arguably more useful application for molecular sciences (see appendix for a 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