SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} T \mathbb {E} [ \vec {\lambda} _ {t} (n _ {t}) \log A _ {t} ^ {\theta} (\mathbf {x} _ {t} ^ {\mathrm {a d d}}, i | \mathbf {Y}) ] \\ = \int_ {0} ^ {T} \mathbb {E} _ {p _ {t} (n _ {t})} [ \vec {\lambda} _ {t} (n _ {t}) \mathbb {E} _ {\tilde {p} (\mathbf {Y} | n _ {t}) p (\mathbf {x} _ {t} ^ {\mathrm {a d d}}, i | \mathbf {Y}) \delta_ {\mathbf {y}} (\mathbf {x} _ {t} ^ {\mathrm {b a s e}})} [ \log A _ {t} ^ {\theta} (\mathbf {x} _ {t} ^ {\mathrm {a d d}}, i | \mathbf {Y}) ] ] \mathrm {d} t \\ = \int_ {0} ^ {T} \mathbb {E} _ {p _ {t} (n _ {t})} [ \vec {\lambda} _ {t} (n _ {t}) \mathbb {E} _ {\tilde {p} (\mathbf {Y} | n _ {t}) p (\mathbf {x} _ {t} ^ {\mathrm {a d d}}, i | \mathbf {Y}) \delta_ {\mathbf {y}} (\mathbf {x} _ {t} ^ {\mathrm {b a s e}})} [ \log A _ {t} ^ {\theta} (\mathbf {x} _ {t} ^ {\mathrm {a d d}}, i | \mathbf {Y}) ] ] \mathrm {d} t \\ - \int_ {0} ^ {T} \mathbb {E} _ {p _ {t} (n _ {t})} [ \vec {\lambda} _ {t} (n _ {t}) \mathbb {E} _ {\tilde {p} (\mathbf {Y} | n _ {t}) p (\mathbf {x} _ {t} ^ {\mathrm {a d d}}, i | \mathbf {Y}) \delta_ {\mathbf {y}} (\mathbf {x} _ {t} ^ {\mathrm {b a s e}})} [ \log p (\mathbf {x} _ {t} ^ {\mathrm {a d d}}, i | \mathbf {Y}) ] ] \mathrm {d} t + \text {c o n s t} \\ = \int_ {0} ^ {T} \mathbb {E} _ {p _ {t} (n _ {t})} [ \vec {\lambda} _ {t} (n _ {t}) \mathbb {E} _ {\tilde {p} (\mathbf {Y} | n _ {t}) \delta_ {\mathbf {y}} (\mathbf {x} _ {t} ^ {\text {b a s e}})} [ - \mathrm {K L} (p (\mathbf {x} _ {t} ^ {\text {a d d}}, i | \mathbf {Y}) | | A _ {t} ^ {\theta} (\mathbf {x} _ {t} ^ {\text {a d d}}, i | \mathbf {Y}) ] ] \mathrm {d} t + \text {c o n s t .} \\ \end{array}
$$
Therefore, the optimum  $A_{t}^{\theta}(\mathbf{x}_{t}^{\mathrm{add}},i|\mathbf{Y})$  which maximizes this part of  $\mathcal{L}(\theta)$  is
$$
A _ {t} ^ {*} \left(\mathbf {x} _ {t} ^ {\text {a d d}}, i | \mathbf {Y}\right) = p \left(\mathbf {x} _ {t} ^ {\text {a d d}}, i | \mathbf {Y}\right) \propto p _ {t} \left(\mathbf {X} _ {t}\right) K ^ {\text {d e l}} (i | n _ {t}).
$$
which is the same form as given in Proposition 1.
B Training Objective
We estimate our objective  $\mathcal{L}(\theta)$  by taking minibatches from the expectation  $\mathcal{U}(t;0,T)p_{0,t}(\mathbf{X}_0,\mathbf{X}_t,M_t)K^{\mathrm{del}}(i|n_t)\delta_{\mathrm{del}(\mathbf{X}_{t,i})}(\mathbf{Y})$ . We first sample  $t\sim \mathcal{U}(t;0,T)$  and then take samples from our dataset  $\mathbf{X}_0\sim p_{\mathrm{data}}(\mathbf{X}_0)$ . In order to sample  $p_{t|0}(\mathbf{X}_t,M_t|\mathbf{X}_0)$  we need to both add noise, delete dimensions and sample a mask variable. Since the Gaussian noise process is isotropic, we can add a suitable amount of noise to all dimensions of  $\mathbf{X}_0$  and then delete dimensions of that noised full dimensional value. More specifically, we first sample  $\tilde{\mathbf{X}}_t = (n_0,\tilde{\mathbf{x}}_t)$  with  $\tilde{\mathbf{x}}_t\sim \mathcal{N}(\tilde{\mathbf{x}}_t;\sqrt{\alpha_t}\mathbf{x}_0,(1 - \alpha_t)I_{n_0d})$  for  $\alpha_{t} = \exp \left(-\int_{0}^{t}\beta (s)\mathrm{d}s\right)$  using the analytic forward equations for the VP-SDE derived in [3]. Then we sample the number of dimensions to delete. This is simple to do when our rate function is independent of  $n$  except for the case when  $n = 1$  at which it is zero. We simply sample a Poisson random variable with mean parameter  $\int_0^t\vec{\lambda}_s\mathrm{d}s$  and then clamp its value such that the maximum number of possible components that are deleted is  $n_0 - 1$ . This gives the appropriate distribution over  $n$ ,  $p_{t|0}(n|n_0)$  as given in Section A.4. To sample which dimensions are deleted, we can sample  $K^{\mathrm{del}}(i_1|n_0)K^{\mathrm{del}}(i_2|n_0 - 1)\ldots K^{\mathrm{del}}(i_{n_0 - n_t}|n_t + 1)$  from which we can create the mask  $M_{t}$  and apply it to  $\tilde{\mathbf{X}}_t$  to obtain  $\mathbf{X}_t$ ,  $\mathbf{X}_t = M_t(\tilde{\mathbf{X}}_t)$ . When  $K^{\mathrm{del}}(i|n) = 1 / n$  this is especially simple to do by simply randomly permuting the components of  $\tilde{\mathbf{X}}_t$ , and then removing the final  $n_0 - n_t$  components.
As is typically done in standard diffusion models, we parameterize  $s_t^\theta$  in terms of a noise prediction network that predicts  $\epsilon$  where  $\mathbf{x}_t = \sqrt{\alpha_t} M_t(\mathbf{x}_0) + \sqrt{1 - \alpha_t} \epsilon$ ,  $\epsilon \sim \mathcal{N}(0, I_{n_d})$ . We then re-weight the score loss in time such that we have a uniform weighting in time rather than the 'likelihood weighting' with  $g_t^2$  [3, 21]. Our objective to learn  $s_t^\theta$  then becomes
$$
\left. - \mathbb {E} _ {\mathcal {U} (t; 0, T) p _ {\text {d a t a}} (\mathbf {X} _ {0}) p \left(M _ {t}, n _ {t} \mid \mathbf {X} _ {0}\right) \mathcal {N} \left(\epsilon ; 0, I _ {n _ {t} d}\right)} \left[ \left\| \epsilon_ {t} ^ {\theta} (\mathbf {X} _ {t}) - \epsilon \right\| ^ {2} \right] \right.
$$
with  $\mathbf{x}_t = \sqrt{\alpha_t} M_t(\mathbf{x}_0) + \sqrt{1 - \alpha_t}\epsilon$ ,  $s_t^\theta (\mathbf{X}_t) = \frac{-1}{\sqrt{1 - \alpha_t}}\epsilon_t^\theta (\mathbf{X}_t)$ .
Further, by using the parameterization given in Proposition 3, we can directly supervise the value of  $p_{0|t}^{\theta}(n_0|\mathbf{X}_t)$  by adding an extra term to our objective. We can treat the learning of  $p_{0|t}^{\theta}(n_0|\mathbf{X}_t)$  as a standard prediction task where we aim to predict  $n_0$  given access to  $\mathbf{X}_t$ . A standard objective for learning  $p_{0|t}^{\theta}(n_0|\mathbf{X}_t)$  is then the cross entropy
$$
\max  _ {\theta} \mathbb {E} _ {p _ {0, t} (\mathbf {X} _ {0}, \mathbf {X} _ {t})} \left[ \log p _ {0 | t} ^ {\theta} (n _ {0} | \mathbf {X} _ {t}) \right]
$$
Our augmented objective then becomes
$$
\tilde {\mathcal {L}} (\theta) = T \mathbb {E} \left[ - \frac {1}{2} \| \epsilon_ {t} ^ {\theta} \left(\mathbf {X} _ {t}\right) - \epsilon \| ^ {2} - \overleftarrow {\lambda} _ {t} ^ {\theta} \left(\mathbf {X} _ {t}\right) + \overrightarrow {\lambda} _ {t} (n _ {t}) \log \overleftarrow {\lambda} _ {t} ^ {\theta} (\mathbf {Y}) + \overrightarrow {\lambda} _ {t} (n _ {t}) \log A _ {t} ^ {\theta} \left(\mathbf {x} _ {t} ^ {\mathrm {a d d}}, i | \mathbf {Y}\right) + \gamma \log p _ {0 | t} ^ {\theta} \left(n _ {0} | \mathbf {X} _ {t}\right) \right] \tag {25}
$$
where the expectation is taken with respect to
$$
\mathcal {U} (t; 0, T) p _ {\mathrm {d a t a}} (\mathbf {X} _ {0}) p \left(M _ {t}, n _ {t} \mid \mathbf {X} _ {0}\right) \mathcal {N} \left(\epsilon ; 0, I _ {n _ {t} d}\right) K ^ {\mathrm {d e l}} \left(i \mid n _ {t}\right) \delta_ {\mathrm {d e l} \left(\mathbf {X} _ {t}, i\right)} (\mathbf {Y})
$$
where  $\mathbf{x}_t = \sqrt{\alpha_t} M_t(\mathbf{x}_0) + \sqrt{1 - \alpha_t}\epsilon$  and  $\gamma$  is a loss weighting term for the cross entropy loss.
C Trans-Dimensional Diffusion Guidance
To guide an unconditionally trained model such that it generates datapoints consistent with conditioning information, we use the reconstruction guided sampling approach introduced in [9]. Our conditioning information will be the values for some of the components of  $\mathbf{X}_0$ , and thus the guidance should guide the generative process such that the rest of the components of the generated datapoint are consistent with those observed components. Following the notation of [9], we denote the observed components as  $\mathbf{x}^a \in \mathbb{R}^{n_a d}$  and the components to be generated as  $\mathbf{x}^b \in \mathbb{R}^{n_b d}$ . Our trained score function  $s_t^\theta (\mathbf{X}_t)$  approximates  $\nabla_{\mathbf{x}_t}\log p_t(\mathbf{X}_t)$  whereas we would like the score to approximate  $\nabla_{\mathbf{x}_t}\log p_t(\mathbf{X}_t|\mathbf{x}_0^a)$ . In order to do this, we will need to augment our unconditional score  $s_t^\theta (\mathbf{X}_t)$  such that it incorporates the conditioning information.
We first focus on the dimensions of the score vector corresponding to  $\mathbf{x}^a$ . These can be calculated analytically from the forward process
$$
\nabla_ {\mathbf {x} _ {t} ^ {a}} \log p (\mathbf {X} _ {t} | \mathbf {x} _ {0} ^ {a}) = \nabla_ {\mathbf {x} _ {t} ^ {a}} \log p _ {t | 0} (\mathbf {x} _ {t} ^ {a} | \mathbf {x} _ {0} ^ {a}, n _ {t})
$$
with  $p_{t|0}(\mathbf{x}_t^a | \mathbf{x}_0^a, n_t) = \mathcal{N}(\mathbf{x}_t^a; \sqrt{\alpha_t} \mathbf{x}_0^a, (1 - \alpha_t) I_{n_ad})$ . Note that we assume a correspondence between  $\mathbf{x}_t^a$  and  $\mathbf{x}_0^a$ . For example, in video if we condition on the first and last frame, we assume that the first and last frame of the current noisy  $\mathbf{x}_t$  correspond to  $\mathbf{x}_0^a$  and guide them towards their observed values. For molecules, the point cloud is permutation invariant and so we can simply assume the first  $n_a$  components of  $\mathbf{x}_t$  correspond to  $\mathbf{x}_0^a$  and guide them to their observed values.
Now we analyse the dimensions of the score vector corresponding to  $\mathbf{x}^b$ . We split the score as
$$
\nabla_ {\mathbf {x} _ {t} ^ {b}} \log p (\mathbf {X} _ {t} | \mathbf {x} _ {0} ^ {a}) = \nabla_ {\mathbf {x} _ {t} ^ {b}} \log p (\mathbf {x} _ {0} ^ {a} | \mathbf {X} _ {t}) + \nabla_ {\mathbf {x} _ {t} ^ {b}} \log p _ {t} (\mathbf {X} _ {t})
$$
$p(\mathbf{x}_0^a | \mathbf{X}_t)$  is intractable to calculate directly and so, following [9], we approximate it with  $\mathcal{N}(\mathbf{x}_0^a; \hat{\mathbf{x}}_0^{\theta a}(\mathbf{X}_t), \frac{1 - \alpha_t}{\alpha_t} I_{n_d})$  where  $\hat{\mathbf{x}}_0^{\theta a}(\mathbf{X}_t)$  is a point estimate of  $\mathbf{x}_0^a$  given from  $s_t^\theta (\mathbf{X}_t)$  calculated as
$$
\hat {\mathbf {x}} _ {0} ^ {\theta a} (\mathbf {X} _ {t}) = \frac {\mathbf {x} _ {t} ^ {a} + (1 - \alpha_ {t}) s _ {t} ^ {\theta} (\mathbf {X} _ {t}) ^ {a}}{\sqrt {\alpha_ {t}}}
$$
where again we have assumed a correspondence between  $\mathbf{x}_t^a$  and  $\mathbf{x}_0^a$ . Our approximation for  $\nabla_{\mathbf{x}_t^b}\log p(\mathbf{x}_0^a|\mathbf{X}_t)$  is then
$$
\nabla_ {\mathbf {x} _ {t} ^ {b}} \log p (\mathbf {x} _ {0} ^ {a} | \mathbf {X} _ {t}) \approx - \nabla_ {\mathbf {x} _ {t} ^ {b}} \frac {\alpha_ {t}}{2 (1 - \alpha_ {t})} \| \mathbf {x} _ {0} ^ {a} - \hat {\mathbf {x}} _ {0} ^ {\theta a} (\mathbf {X} _ {t}) \| ^ {2}
$$
which can be calculated by differentiating through the score network  $s_t^\theta$ .
We approximate  $\overleftarrow{\lambda}_t^*(\mathbf{X}_t|\mathbf{x}_0^a)$  and  $A_t^*(\mathbf{y}^{\mathrm{add}},i|\mathbf{X}_t,\mathbf{x}_0^a)$ , with their unconditional forms  $\overleftarrow{\lambda}_t^\theta(\mathbf{X}_t)$  and  $A_t^\theta(\mathbf{y}^{\mathrm{add}},i|\mathbf{X}_t)$ . We find this approximation still leads to valid generations because the guidance of the score network  $s_t^\theta$ , results in  $\mathbf{X}_t$  containing the conditioning information which in turn leads to  $\overleftarrow{\lambda}_t^\theta(\mathbf{X}_t)$  guiding the number of components in  $\mathbf{X}_t$  to be consistent with the conditioning information too as verified in our experiments. Further, any errors in the approximation for  $A_t^\theta(\mathbf{y}^{\mathrm{add}},i|\mathbf{X}_t)$  are fixed by further applications of the guided score function, highlighting the benefits of our combined autoregressive and diffusion based approach.
Our code is available at https://github.com/andreew-cr/jump-diffusion
D.1 Molecules
D.1.1 Network Architecture
Backbone For our backbone network architecture, we used the EGNN used in [8]. This is a specially designed graph neural network applied to the point cloud treating it as a fully connected graph. A special equivariant update is used, operating only on distances between atoms. We refer to [8] for the specific details on the architecture. We used the same size network as used in [8]'s QM9 experiments, specifically there are 9 layers, with a hidden node feature size of 256. The output of the EGNN is fed into a final output projection layer to give the score network output  $s_t^\theta(\mathbf{X}_t)$ .
Component number prediction To obtain  $p_{0|t}^{\theta}(n_0|\mathbf{X}_t)$ , we take the embedding produced by the EGNN before the final output embedding layer and pass it through 8 transformer layers each consisting of a self-attention block and an MLP block applied channel wise. Our transformer model dimension is 128 and so we project the EGNN embedding output down to 128 before entering into the transformer layers. We then take the output of the transformer and take the average embedding over all nodes. This embedding is then passed through a final projection layer to give softmax logits over the  $p_{0|t}^{\theta}(n_0|\mathbf{X}_t)$  distribution.
Autoregressive Distribution Our  $A_{t}^{\theta}(\mathbf{y}^{\mathrm{add}}, i|\mathbf{X}_{t})$  network has to predict the position and features for a new atom when it is added to the molecule. Since the point cloud is permutation invariant, we do not need to predict  $i$  and so we just need to parameterize  $A_{t}^{\theta}(\mathbf{y}^{\mathrm{add}}|\mathbf{X}_{t})$ . We found the network to perform the best if the network first predicts the nearest atom to the new atom and then a vector from that atom to the location of the new atom. To achieve this, we first predict softmax logits for a distribution over the nearest atom by applying a projection to the embedding output from the previously described transformer block. During training, the output of this distribution can be directly supervised by a cross entropy loss. Given the nearest atom, we then need to predict the position and features of the new atom to add. We do this by passing in the embedding generated by the EGNN and original point cloud features into a new transformer block of the same size as that used for  $p_{0|t}^{\theta}(n_0|\mathbf{X}_t)$ . We also input the distances from the nearest atom to all other atoms in the molecule currently as an additional feature. To obtain the position of the new atom, we will take a weighted sum of all the vectors between the nearest atom and other atoms in the molecule. This is to make it easy for the network to create new atoms 'in plane' with existing atoms which is useful for e.g. completing rings that have to remain in the same plane. To calculate the weights for the vectors, we apply an output projection to the output of the transformer block. The new atom features (atom type and charge) are generated by a separate output projection from the transformer block. For the position and features,  $A_{t}^{\theta}(\mathbf{y}^{\mathrm{add}}|\mathbf{X}_{t})$  outputs both a mean and a standard deviation for a Gaussian distribution. For the position distribution, we set the standard deviation to be isotropic to remain equivariant to rotations. In total our model has around 7.3 million parameters.
D.1.2 Training
We train our model for 1.3 million iterations at a batch size of 64. We use the Adam optimizer with learning rate 0.00003. We also keep a running exponential moving average of the network weights that is used during sampling as is standard for training diffusion models [2, 3, 16] with a decay parameter of 0.9999. We train on the 100K molecules contained in the QM9 training split. We model hydrogens explicitly. Training a model requires approximately 7 days on a single GPU which was done on an Academic cluster.
In [8] the atom type is encoded as a one-hot vector and diffused as a continuous variable along with the positions and charge values for all atoms. They found that multiplying the one-hot vectors by 0.25 to boost performance by allowing the atom-type to be decided later on in the diffusion process. We instead multiply the one-hot vectors by 4 so that atom-type is decided early on in the diffusion process which improves our guided performance when conditioning on certain atom-types being
present. We found our model is robust to this change and achieves similar sample quality to [8] as shown in Table 2.
When deleting dimensions, we first shuffle the ordering of the nodes and then delete the final  $n_0 - n_t$  nodes. The cross entropy loss weighting in (25) is set to 1.
Following [8] we train our model to operate within the center of mass (CoM) zero subspace of possible molecule positions. The means, throughout the forward and backward process, the average position of an atom is 0. In our transdimensional framework, this is achieved by first deleting any atoms required under the forward component deletion process. We then move the molecule such that its CoM is 0. We then add CoM free noise such that the noisy molecule also has  $\mathrm{CoM} = 0$ . Our score model  $s_t^\theta$  is parameterized through a noise prediction model  $\epsilon_t^\theta$  which is trained to predict the CoM free noise that was added. Therefore, our score network learns suitable directions to maintain the process on the  $\mathrm{CoM} = 0$  subspace. For the position prediction from  $A_t^\theta (\mathbf{y}^{\mathrm{add}}|\mathbf{X}_t)$  we train it to predict the new atom position from the current molecules reference frame. When the new atom is added, we then update all atom positions such that  $\mathrm{CoM} = 0$  is maintained.
D.1.3 Sampling
During sampling we found that adding corrector steps [3] improved sample quality. Intuitively, corrector steps form a process that has  $p_t(\mathbf{X})$  as its stationary distribution rather than the process progressing toward  $p_0(\mathbf{X})$ . We use the same method to determine the corrector step size  $\zeta$  as in [3]. For the conditional generation tasks, we also found it useful to include corrector steps for the component generation process. As shown in [29], corrector steps in discrete spaces can be achieved by simulating with a rate that is the addition of the forward and backward rates. We achieve this in the context of trans-dimensional modeling by first simulating a possible insertion using  $\overleftrightarrow{\lambda}_t^\theta$  and then simulating a possible deletion using  $\overrightarrow{\lambda}_t$ . We describe our overall sampling algorithm in Algorithm 2.
Algorithm 2: Sampling the Generative Process with Corrector Steps
Input: Number of corrector steps  $C$ $t \gets T$ $\mathbf{X} \sim p_{\mathrm{ref}}(\mathbf{X}) = \mathbb{I}\{n = 1\} \mathcal{N}(\mathbf{x};0,I_d)$   
while  $t > 0$  do  
if  $u < \overleftarrow{\lambda}_t^\theta (\mathbf{X})\delta t$  with  $u \sim \mathcal{U}(0,1)$  then  
Sample  $\mathbf{x}^{\mathrm{add}},i \sim A_t^\theta (\mathbf{x}^{\mathrm{add}},i|\mathbf{X})$ $\mathbf{X} \gets \mathrm{ins}(\mathbf{X},\mathbf{x}^{\mathrm{add}},i)$   
end  
 $\mathbf{x} \gets \mathbf{x} - \overleftarrow{\mathbf{b}}_t^\theta (\mathbf{X})\delta t + g_t\sqrt{\delta t}\epsilon$  with  $\epsilon \sim \mathcal{N}(0,I_{nd})$   
for  $c = [1,\dots,C]$  do  
 $\mathbf{x} \gets \mathbf{x} + \zeta s_{t - \delta t}^\theta (\mathbf{X}) + \sqrt{2\zeta}\epsilon$  with  $\epsilon \sim \mathcal{N}(0,I_{nd})$   
if  $u < \overleftarrow{\lambda}_{t - \delta t}^\theta (\mathbf{X})\delta t$  with  $u \sim \mathcal{U}(0,1)$  then  
Sample  $\mathbf{x}^{\mathrm{add}},i \sim A_{t - \delta t}^\theta (\mathbf{x}^{\mathrm{add}},i|\mathbf{X})$ $\mathbf{X} \gets \mathrm{ins}(\mathbf{X},\mathbf{x}^{\mathrm{add}},i)$   
end  
if  $u < \overrightarrow{\lambda}_{t - \delta t}(n)\delta t$  with  $u \sim \mathcal{U}(0,1)$  then  
 $\mathbf{X} \gets \mathrm{del}(\mathbf{X},i)$  with  $i \sim K^{\mathrm{del}}(i|n)$   
end  
end  
 $\mathbf{X} \gets (n,\mathbf{x}),t \gets t - \delta t$   
end
D.1.4 Evaluation
Unconditional For our unconditional sampling evaluation, we start adding corrector steps when  $t < 0.1T$  in the backward process and use 5 corrector steps without the corrector steps on the number of components. We set  $\delta = 0.05$  for  $t > 0.5T$  and  $\delta = 0.001$  for  $t < 0.5T$  such that the total number
Figure 6: Distribution of the size of molecules in the QM9 dataset as measured through the number of atoms versus the distribution of the size of molecules generated by our unconditional model.
of network evaluations is 1000. We show the distribution of sizes of molecules generated by our model in Figure 6 and show more unconditional samples in Figure 7. We find our model consistently generates realistic molecules and achieves a size distribution similar to the training dataset even though this is not explicitly trained and arises from sampling our backward rate  $\stackrel{\leftarrow}{\lambda}_t^\theta$ . Since we are numerically integrating a continuous time process and approximating the true time reversal rate  $\stackrel{\leftarrow}{\lambda}_t^*$ , some approximation error is expected. For this experiment, sampling all of our models and ablations takes approximately 2 GPU days on Nvidia 1080Ti GPUs.
Conditional For evaluating applying conditional diffusion guidance to our model, we choose 10 conditioning tasks that each result in a different distribution of target dimensions. The task is to produce molecules that include at least a certain number of target atom types. We then guide the first set of atoms generated by the model to have these desired atom types. The tasks chosen are given in Table 5. Molecules in the training dataset that meet the conditions in each task have a different distribution of sizes. The tasks were chosen so that we have an approximately linearly increasing mean number of atoms for molecules that meet the condition. We also require that there are at least 100 examples of molecules that meet the condition within the training dataset.
For sampling when using conditional diffusion guidance, we use 3 corrector steps throughout the backward process with  $\delta t = 0.001$ . For these conditional tasks, we include the corrector steps on the number of components. We show the distribution of dimensions for each task from the training dataset and from our generated samples in Figure 8. Our metrics are calculated by first drawing 1000 samples for each conditioning task and then finding the Hellinger distance between the size distribution generated by our method (orange diagonal hashing in Figure 8) and the size distribution for molecules in the training dataset that match the conditions of the task (green no hashing in Figure 8). We find that indeed our model when guided by diffusion guidance can automatically produce a size distribution close to the ground truth size distribution found in the dataset for that conditioning value. We show samples generated by our conditionally guided model in Figure 9. We can see that our model can generate realistic molecules that include the required atom types and are of a suitable size. For this experiment, sampling all of our models and ablations takes approximately 13 GPU days on Nvidia 1080Ti GPUs.
Interpolations For our interpolations experiments, we follow the set up of [8] who train a new model conditioned on the polarizability of molecules in the dataset. We train a conditional version of our model which can be achieved by simply adding in the polarizability as an additional feature input to our backbone network and re-using all the same hyperparameters. We show more examples of interpolations in Figure 10.
Figure 7: Unconditional samples from our model.
D.1.5 Ablations
For our main model, we set  $\vec{\lambda}_{t < 0.1T} = 0$  to ensure that all dimensions are added with enough generation time remaining for the diffusion process to finalize all state values. To verify this setting, we compare its performance with  $\vec{\lambda}_{t < 0.03T} = 0$  and  $\vec{\lambda}_{t < 0.3T} = 0$ . We show our results in Table 6. We find that the  $\vec{\lambda}_{t < 0.03T} = 0$  setting to generate reasonable sample quality but incur some extra dimension error due to the generative process sometimes observing a lack of dimensions near  $t = 0$  and adding too many dimensions. We observed the same effect in the paper for when setting  $\vec{\lambda}_{t}$  to be constant for all  $t$  in Table 3. Further, the setting  $\vec{\lambda}_{t < 0.3T} = 0$  also results in increased dimension error due to there being less opportunity for the guidance model to supervise the number of dimensions. We find that  $\vec{\lambda}_{t < 0.1T} = 0$  to be a reasonable trade-off between these effects.
D.1.6 Uniqueness and Novelty Metrics