SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Algorithm 1: Walk-jump sampling [1] using the discretization of Langevin diffusion by [85]. Lines 6-13 correspond to the walk step and line 14 to the jump step.
1: Input  $\delta$  (step size),  $u$  (inverse mass),  $\gamma$  (friction),  $K$  (steps taken)  
2: Input Learned score function  $g_{\theta}(y) \approx \nabla_y \log p(y)$  and noise level  $\sigma$   
3: Output  $\widehat{x}_K$   
4:  $y_0 \sim \mathcal{N}(0, \sigma^2 I_d) + \mathcal{U}_d(0, 1)$   
5:  $v_0 \gets 0$   
6: for  $k = 0, \dots, K-1$  do  
7:  $y_{k+1} \gets y_k + \frac{\delta}{2} v_k$   
8:  $g \gets g_{\theta}(y_{k+1})$   
9:  $v_{k+1} \gets v_k + \frac{u\delta}{2} g$   
10:  $\varepsilon \sim \mathcal{N}(0, I_d)$   
11:  $v_{k+1} \gets \exp(-\gamma\delta) v_{k+1} + \frac{u\delta}{2} g + \sqrt{u(1 - \exp(-2\gamma\delta))}\varepsilon$   
12:  $y_{k+1} \gets y_{k+1} + \frac{\delta}{2} v_{k+1}$   
13: end for  
14:  $\hat{x}_K \gets y_K + \sigma^2 g_{\theta}(y_K)$
We use  $\gamma = 1.0$ ,  $u = 1.0$ ,  $\delta = .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  $\sigma$  used during training), i.e.,  $y_0 = \mathcal{N}(0, \sigma^2 I_d) + \mathcal{U}_d(0, 1)$  (this was observed to mix faster in practice).
All experiments and analysis on this paper were done on A100 GPUs and with PyTorch [97]. The models on QM9 were trained with 2 GPUs and the models on GEOM-drugs on 4 GPUs.
dset	coordinates ref.	stable mol %↑	stable atom %↑	valid %↑	unique %↑	valency W1↓	atom TV↓	bond TV↓	bond len W1↓	bond ang W1↓
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 on 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 a sample in the voxelized representation,  $x_{k}$ . We set to 0 all voxels with value less than .1, i.e.,  $x_{k} = \mathbb{1}_{\geq .1}(x_{k})$ . 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 a maximum filter with a  $3\times 3\times 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 a single step and do not perform search over multiple possible numbers of atoms or atom identities.
Table 5 shows the effect of coordinate refinement on 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 to determine the molecule's atomic bonds.
A.6 Metrics
Most of the metrics used to benchmark models come from  $[20]^8$ . Below we describe the metrics:
- 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 are stable.
- Validity: The percentage of generated molecules that passes RDKit's sanitization filter.
- Uniqueness:. The proportion of valid molecules (defined above) that has a 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 5 atom types on QM9 and 8 atom types on GEOM-drugs. The histograms  $\hat{h}_{\mathrm{atm}}$  and  $h_{\mathrm{atm}}$  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:
$$
\text {A t o m s} \mathrm {T V} (\hat {h} _ {\mathrm {a t m}}, h _ {\mathrm {a t m}}) = \sum_ {x \in \text {a t o m t y p e s}} | \hat {h} _ {\mathrm {a t m}} (x) - h _ {\mathrm {a t m}} (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:
$$
\text {B o n d s} \mathrm {T V} (\hat {h} _ {\text {b o n d}}, h _ {\text {b o n d}}) = \sum_ {x \in \text {b o n d t y p e s}} | \hat {h} _ {\text {b o n d}} (x) - h _ {\text {b o n d}} (x) |
$$
<sup>8</sup>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.
- Valency  $W_{1}$ : This is the weighted sum of the Wasserstein distance between the distribution of valencies for each atom type:
$$
\operatorname {V a l e n c y} \mathrm {W} _ {1} (\text {g e n e r a t e d}, \text {t a r g e t}) = \sum_ {x \in \text {a t o m t y p e s}} p (x) W _ {1} \left(\hat {h} _ {\text {v a l}} (x), h _ {\text {v a l}} (x)\right),
$$
where  $\hat{h}_{\mathrm{val}}(x)$  and  $h_{\mathrm{val}}(x)$  are the histogram of valencies for atom type  $x$  for generated and holdout set samples, respectively.
- Bond length  $W_{1}$ : The weighted sum of Wasserstein distance between the distribution of bond lengths for each bond type:
$$
\text {B o n d L e n W} _ {1} (\text {g e n e r a t e d}, \text {t a r g e t}) = \sum_ {b \in \text {b o n d t y p e s}} p (b) W _ {1} \left(\hat {h} _ {\text {d i s t}} (b), h _ {\text {d i s t}} (b)\right),
$$
where  $\hat{h}_{\mathrm{dist}}(b)$  and  $h_{\mathrm{dist}}(b)$  are the histogram of bond lengths for bond type  $b$ , for generated and holdout set samples, respectively.
- Bond angles  $W_{1}$ : The weighted sum of Wasserstein distance between the distribution of bond angles (in degrees) for each atom type in the dataset:
$$
\text {B o n d A n g W} _ {1} (\text {g e n e r a t e d}, \text {t a r g e t}) = \sum_ {x \in \text {a t o m t y p e s}} p (x) W _ {1} \left(\hat {h} _ {\text {a n g}} (x), h _ {\text {a n g}} (x)\right),
$$
where  $\hat{h}_{\mathrm{ang}}(x)$  and  $h_{\mathrm{ang}}(x)$  are the histogram of angles for atom type  $x$  for generated and holdout set samples, respectively. See [20] for how angles are measured.
- Strain energy: The strain energy for a 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<sup>9</sup>.
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 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 diffusion models, the score function of all steps is 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 (e.g., 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 diffusion models 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. For example, 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 [35]) or 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 instead of diffusion models. In-painting has also proven to work very well in 2D U-Nets, so it could potentially 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.
$^9$ Code taken from https://github.com/cch1999/posecheck/blob/main/posecheck/util/strain.py
B Generated samples
Figure 8: Random generated samples from VoxMol trained on QM9 (passing RDKit's sanitization). Molecular graphs are generated with RDKit.
Figure 9: Random generated samples from VoxMol trained on GEOM-drugs (passing RDKit's sanitization). Molecular graphs are generated with RDKit.