SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We now examine the objective to gain an intuition into the learning signal. Our first term (3) is an  $L_{2}$  regression to a target that, as we have seen, is a scaled vector between  $\mathbf{x}_t$  and  $\sqrt{\alpha_t} M_t(\mathbf{x}_0)$ . As the solution to an  $L_{2}$  regression problem is the conditional expectation of the target,  $s_t^\theta (\mathbf{X}_t)$  will learn to predict vectors pointing towards  $\mathbf{x}_0$  averaged over the possible correspondences between dimensions of  $\mathbf{x}_t$  and dimensions of  $\mathbf{x}_0$ . Thus, during sampling,  $s_t^\theta (\mathbf{X}_t)$  provides a suitable direction to adjust the current value  $\mathbf{X}_t$  taking into account the fact  $\mathbf{X}_t$  represents only a noisy subpart of a clean whole  $\mathbf{X}_0$ .
The second term (4) gives a learning signal for  $\overleftarrow{\lambda}_t^\theta$  and  $A_{t}^{\theta}$ . For  $A_{t}^{\theta}$ , we simply have a maximum likelihood objective, predicting the missing part of  $\mathbf{X}_t$  (i.e.  $\mathbf{x}_{t}^{\mathrm{add}}$ ) given the observed part of  $\mathbf{X}_t$  (i.e.  $\mathbf{Y}$ ). The signal for  $\overleftarrow{\lambda}_t^\theta$  comes from balancing two terms:  $-\overleftarrow{\lambda}_t^\theta (\mathbf{X}_t)$  and  $\overrightarrow{\lambda}_t(n_t)\log \overleftarrow{\lambda}_t^\theta (\mathbf{Y})$
which encourage the value of  $\overleftarrow{\lambda}_t^\theta$  to move in opposite directions. For a new test input  $\mathbf{Z}$ ,  $\overleftarrow{\lambda}_t^\theta (\mathbf{Z})$ 's value needs to trade off between the two terms by learning the relative probability between  $\mathbf{Z}$  being the entirety of a genuine sample from the forward process, corresponding to the  $\overleftarrow{\lambda}_t^\theta (\mathbf{X}_t)$  term in (4), or  $\mathbf{Z}$  being a substructure of a genuine sample, corresponding to the  $\overleftarrow{\lambda}_t^\theta (\mathbf{Y})$  term in (4). The optimum trade-off is found exactly at the time reversal  $\overleftarrow{\lambda}_t^*$  as we show in Appendix A.5.
We optimize  $\mathcal{L}(\theta)$  using stochastic gradient ascent, generating minibatches by first sampling  $t\sim$ $\mathcal{U}(0,T),\mathbf{X}_0\sim p_{\mathrm{data}}$  and then computing  $\mathbf{X}_t$  from the forward process. This can be done analytically for the  $\vec{\lambda}_t(n)$  functions used in our experiments. We first sample  $n_t$  by analytic integration of the dimension deletion Poisson process with time inhomogeneous rate  $\vec{\lambda}_t(n)$ . We then add Gaussian noise independently to each dimension under  $p_{t|0}(\mathbf{x}_t|\mathbf{X}_0,n_t,M_t)$  using a randomly drawn mask variable  $M_{t}$ . See Appendix B for further details on the efficient evaluation of our objective.
3.4 Parameterization
$s_t^\theta(\mathbf{X}_t)$ ,  $A_t^\theta(\mathbf{y}^{\mathrm{add}}, i|\mathbf{X}_t)$  and  $\overleftarrow{\lambda}_t^\theta(\mathbf{X}_t)$  will all be parameterized by neural networks. In practice, we have a single backbone network suited to the problem of interest e.g. a Transformer [23], an EGNN [24] or a UNet [25] onto which we add prediction heads for  $s_t^\theta(\mathbf{X}_t)$ ,  $A_t^\theta(\mathbf{y}^{\mathrm{add}}, i|\mathbf{X}_t)$  and  $\overleftarrow{\lambda}_t^\theta(\mathbf{X}_t)$ .  $s_t^\theta(\mathbf{X}_t)$  outputs a vector in  $\mathbb{R}^{n_d}$ .  $A_t^\theta(\mathbf{y}^{\mathrm{add}}, i|\mathbf{X}_t)$  outputs a distribution over  $i$  and mean and standard deviation statistics for a Gaussian distribution over  $\mathbf{y}^{\mathrm{add}}$ . Finally, having  $\overleftarrow{\lambda}_t^\theta(\mathbf{X}_t) \in \mathbb{R}_{\geq 0}$  be the raw output of a neural network can cause optimization issues due to the optimum  $\overleftarrow{\lambda}_t^*$  including a probability ratio which can take on very large values. Instead, we learn a component prediction network  $p_{0|t}^\theta(n_0|\mathbf{X}_t)$  that predicts the number of components in  $\mathbf{X}_0$  given  $\mathbf{X}_t$ . To convert this into  $\overleftarrow{\lambda}^\theta(\mathbf{X}_t)$ , we show in Proposition 3 how the optimum  $\overleftarrow{\lambda}_t^*(\mathbf{X}_t)$  is an analytic function of the true  $p_{0|t}(n_0|\mathbf{X}_t)$ . We then plug  $p_{0|t}^\theta(n_0|\mathbf{X}_t)$  into Proposition 3 to obtain an approximation of  $\overleftarrow{\lambda}_t^*(\mathbf{X}_t)$ .
Proposition 3. We have
$$
\stackrel {\leftarrow} {\lambda} _ {t} ^ {*} (\mathbf {X} _ {t}) = \vec {\lambda} _ {t} (n _ {t} + 1) \sum_ {n _ {0} = 1} ^ {N} \frac {p _ {t | 0} (n _ {t} + 1 | n _ {0})}{p _ {t | 0} (n _ {t} | n _ {0})} p _ {0 | t} (n _ {0} | \mathbf {X} _ {t}),
$$
where  $\mathbf{X}_t = (n_t, \mathbf{x}_t)$  and  $p_{t|0}(n_t + 1|n_0)$  and  $p_{t|0}(n_t|n_0)$  are both easily calculable distributions from the forward dimension deletion process.
3.5 Sampling
To sample the generative process, we numerically integrate the learned backward jump diffusion process using time-step  $\delta t$ . Intuitively, it is simply the standard continuous time diffusion sampling scheme [3] but at each timestep we check whether a jump has occurred and if it has, sample the new component and insert it at the chosen index as explained by Algorithm 1.
Algorithm 1: Sampling the Generative Process
t←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)$
$\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})$
$\mathbf{X}\leftarrow (n,\mathbf{x}),t\gets t - \delta t$
Our method jointly generates both dimensions and state values during the generative process whereas prior approaches [8, 11] are forced to first sample the number of dimensions and then run the diffusion process in this fixed dimension. When diffusion guidance is applied to these unconditional models [14, 26], users need to pick by hand the number of dimensions independent of the conditioning information even though the number of dimensions can be correlated with the conditioning parameter.
Instead of automatically learning when and how many dimensions to add during the generative process, previous work focusing on images [27, 28] hand pick dimension jump points such that the
resolution of images is increased during sampling and reaches a certain pre-defined desired resolution at the end of the generative process. Further, rather than using any equivalent of  $A_{t}^{\theta}$ , the values for new dimensions are simply filled in with Gaussian noise. These approaches mainly focus on efficiency rather than flexible generation as we do here.
The first term in our learning objective in Proposition 2 corresponds to learning the continuous part of our process (the diffusion) and the second corresponds to learning the discrete part of our process (the jumps). The first term can be seen as a trans-dimensional extension of standard denoising score matching [20] whilst the second bears similarity to the discrete space ELBO derived in [29].
Finally, jump diffusions also have a long history of use in Bayesian inference, where one aims to draw samples from a trans-dimensional target posterior distribution based on an unnormalized version of its density [30]: an ergodic jump diffusion is designed which admits the target as the invariant distribution [30-32]. The invariant distribution is not preserved when time-discretizing the process. However, it was shown in [33, 34] how general jump proposals could be built and how this process could be "Metropolized" to obtain a discrete-time Markov process admitting the correct invariant distribution, yielding the popular Reversible Jump Markov Chain Monte Carlo algorithm. Our setup differs significantly as we only have access to samples in the form of data, not an unnormalized target.
5.1 Molecules
We now show how our model provides significant benefits for diffusion guidance and interpolation tasks. We model the QM9 dataset [35, 36] of 100K varying size molecules. Following [8], we consider each molecule as a 3-dimensional point cloud of atoms, each atom having the features:  $(x,y,z)$  coordinates, a one-hot encoded atom type, and an integer charge value. Bonds are inferred from inter-atomic distances. We use an EGNN [24] backbone with three heads to predict  $s_t^\theta$ ,  $p_{0|t}^\theta(n_0|\mathbf{X}_t)$ , and  $A_t^\theta$ . We uniformly delete dimensions,  $K^{\mathrm{del}}(i|n) = 1/n$ , and since a point cloud is permutation invariant,  $A_t^\theta(\mathbf{y}^{\mathrm{add}}|\mathbf{X}_t)$  need only predict new dimension values. We set  $\vec{\lambda}_t$  to a constant except for  $t < 0.1T$ , where we set  $\vec{\lambda}_{t<0.1T} = 0$ . This ensures that all dimensions are added with enough generation time remaining for the diffusion process to finalize all state values.
We visualize sampling from our learned generative process in Figure 2; note how the process jointly creates a suitable number of atoms whilst adjusting their positions and identities. Before moving on to apply diffusion guidance which is the focus of our experiments, we first verify our unconditional sample quality in Table 2 and find we perform comparably to the results reported in [8] which use an FDDM. We ablate our choice of  $\vec{\lambda}_t$  by comparing with setting  $\vec{\lambda}_t$  to a constant for all  $t$  and with setting  $\vec{\lambda}_t = 0$  for  $t < 0.9T$  (rather than just for  $t < 0.1T$ ). We find that the constant  $\vec{\lambda}_t$  performs worse due to the occasional component being added late in the generation process without enough time for the diffusion process to finalize its value. We find the  $\vec{\lambda}_{t < 0.9T} = 0$  setting to have satisfactory sample quality however this choice of  $\vec{\lambda}_t$  introduces issues during diffusion guided generation as we see next. Finally, we ablate the parameterization of Proposition 3 by learning  $\vec{\lambda}_t^\theta (\mathbf{X}_t)\in \mathbb{R}$  directly as the output of a neural network head. We find that this reduces sample quality due to the more well-behaved nature of the target,  $p_{0|t}^{\theta}(n_0|\mathbf{X}_t)$  when using Proposition 3. We note pure autoregressive models perform significantly worse than diffusion based models as found in [8].
Figure 2: Visualization of the jump-diffusion backward generative process on molecules.
Table 2: Sample quality metrics for unconditional molecule generation. An atom is stable if it has the correct valency whilst a molecule is considered stable if all of its atoms are stable. Molecular validity is measured using RDKit [37]. All methods use 1000 simulation steps and draw 10000 samples.
Method	% Atom Stable (↑)	% Molecule Stable (↑)	% Valid (↑)
FDDM [8]	98.7	82.0	91.9
TDDM (ours)	98.3	87.2	92.3
TDDM, const x̂t	96.7	79.1	86.7
TDDM, x̂t<0.9T=0	97.7	82.6	89.4
TDDM w/o Prop. 3	97.0	66.9	87.1
Table 3: Conditional Molecule Generation for 10 conditioning tasks that each result in a different dimension distribution. We report dimension error as the average Hellinger distance between the generated and ground truth dimension distributions for that property as well as average sample quality metrics. Standard deviations are given across the 10 conditioning tasks. We report in bold values that are statistically indistinguishable from the best result at the  $5\%$  level using a two-sided Wilcoxon signed rank test across the 10 conditioning tasks.
Method	Dimension Error (↓)	% Atom Stable (↑)	% Molecule Stable (↑)	% Valid (↑)
FDDM	0.511±0.19	93.5±1.1	31.3±6.3	65.2±10.3
TDDM	0.134±0.076	93.5±2.6	59.1±11	74.8±9.3
TDDM, const x̂t	0.226±0.17	88.9±4.8	43.6±15	63.4±14
TDDM, x̂t<0.9T=0	0.390±0.38	95.0±2.1	61.7±17	77.8±13
TDDM w/o Prop. 3	0.219±0.12	93.8±3.2	55.0±19	73.8±13
5.1.1 Trans-Dimensional Diffusion Guidance
We now apply diffusion guidance to our unconditional model in order to generate molecules that contain a certain number of desired atom types, e.g. 3 carbons or 1 oxygen and 2 nitrogens. The distribution of molecule sizes changes depending on these conditions. We generate molecules conditioned on these properties by using the reconstruction guided sampling approach introduced in [9]. This method augments the score  $s_t^\theta(\mathbf{X}_t)$  such that it approximates  $\nabla_{\mathbf{x}_t} \log p_t(\mathbf{X}_t | y)$  rather than  $\nabla_{\mathbf{x}_t} \log p_t(\mathbf{X}_t)$  (where  $y$  is the conditioning information) by adding on a term approximating  $\nabla_{\mathbf{x}_t} \log p_t(y | \mathbf{X}_t)$  with  $p_t(y | \mathbf{X}_t) = \sum_{n_0} \int_{\mathbf{x}_0} p(y | \mathbf{X}_0) p_{0|t}(\mathbf{X}_0 | \mathbf{X}_t) \mathrm{d}\mathbf{x}_0$ . This guides  $\mathbf{x}_t$  such that it is consistent with  $y$ . Since  $\lambda_t^\theta(\mathbf{X}_t)$  has access to  $\mathbf{x}_t$ , it will cause  $n_t$  to automatically also be consistent with  $y$  without the user needing to input any information on how the conditioning information relates to the size of the datapoints. We give further details on diffusion guidance in Appendix C.
We show our results in Table 3. In order to perform guidance on the FDDM baseline, we implement the model from [8] in continuous time and initialize the dimension from the empirically observed dimension distribution in the dataset. This accounts for the case of an end user attempting to guide a unconditional model with access to no further information. We find that TDDM produces samples whose dimensions much more accurately reflect the true conditional distribution of dimensions given the conditioning information. The  $\vec{\lambda}_{t < 0.9T} = 0$  ablation on the other hand only marginally improves the dimension error over FDDM because all dimensions are added in the generative process at a time when  $\mathbf{X}_t$  is noisy and has little relation to the conditioning information. This highlights the necessity of allowing dimensions to be added throughout the generative process to gain the trans-dimensional diffusion guidance ability. The ablation with constant  $\vec{\lambda}_t$  has increased dimension error over TDDM as we find that when  $\vec{\lambda}_t > 0$  for all  $t$ ,  $\vec{\lambda}_t^\theta$  can become very large when  $t$  is close to 0 when the model has perceived a lack of dimensions. This occasionally results in too many dimensions being added hence an increased dimension error. Not using the Proposition 3 parameterization also increases dimension error due to the increased difficulty in learning  $\vec{\lambda}_t^\theta$ .
5.1.2 Trans-Dimensional Interpolation
Interpolations are a unique way of gaining insights into the effect of some conditioning parameter on a dataset of interest. To create an interpolation, a conditional generative model is first trained and then sampled with a sweep of the conditioning parameter but using fixed random noise [8]. The resulting series of synthetic datapoints share similar features due to the fixed random noise but vary in ways that are very informative as to the effect of the conditioning parameter. Attempting to interpolate with an FDDM is fundamentally limited because the entire interpolation occurs in the same dimension which is unrealistic when the conditioning parameter is heavily correlated with the dimension of the datapoint. We demonstrate this by following the setup of [8] who train a conditional FDDM conditioned on polarizability. Polariz
ability is the ability of a molecule's electron cloud to distort in response to an external electric field [38] with larger molecules tending to have higher polarizability. To enable us to perform a trans-dimensional interpolation, we also train a conditional version of our model conditioned on polarizability. An example interpolation with this model is shown in Figure 4. We find that indeed the size of the molecule increases with increasing polarizability, with some molecular substructures e.g. rings, being maintained across dimensions. We show how the dimension changes with polarizability during 3 interpolations in Figure 3. We find that these match the true dataset statistics much more accurately than interpolations using FDDM which first pick a dimension and carry out the entire interpolation in that fixed dimension.
Figure 3: Number of atoms versus polarizability for 3 interpolations with fixed random noise. The dataset mean and standard deviation for the number of atoms is also shown. FDDM interpolates entirely in a fixed dimensional space hence the number of atoms is fixed for all polarizabilities.
Figure 4: Sequence of generations for linearly increasing polarizability from  $39\mathrm{Bohr}^3$  to  $66\mathrm{Bohr}^3$  with fixed random noise. Note how molecular size generally increases with polarizability and how some molecular substructures are maintained between sequential generations of differing dimension. For example, between molecules 6 and 7, the single change is a nitrogen (blue) to a carbon (gray) and an extra hydrogen (white) is added to maintain the correct valency.
5.2 Video
We finally demonstrate our model on a video modeling task. Specifically we model the RoboDesk dataset [39], a video benchmark to measure the applicability of video models for planning and control problems. The videos are renderings of a robotic arm [40] performing a variety of different tasks including opening drawers and moving objects. We first train an unconditional model on videos of varying length and then perform planning by applying diffusion guidance to generate videos conditioned on an initial starting frame and a final goal frame [41]. The planning problem is then reduced to "filling in" the frames in between. Our trans-dimensional model automatically varies the number of in-filled frames during generation so that the final length of video matches the length of time the task should take, whereas the fixed dimension model relies on the unrealistic assumption that the length of time the task should take is known before generation.
We model videos at  $32 \times 32$  resolution and with varying length from 2 to 35 frames. For the network backbone, we use a UNet adapted for video [42]. In contrast to molecular point clouds, our data is no longer permutation invariant hence  $A_{t}^{\theta}(\mathbf{y}^{\mathrm{add}}, i|\mathbf{X}_{t})$  includes a prediction over the location to insert the new frame. Full experimental details are provided in Appendix D. We evaluate our approach on three planning tasks, holding stationary, sliding a door and pushing an object. An example generation conditioned on the first and last frame for the slide door task is shown in Figure 5, with
the model in-filling a plausible trajectory. We quantify our model's ability to generate videos of a length appropriate to the task in Table 4 finding on all three tasks we generate a more accurate length of video than FDDM which is forced to sample video lengths from the unconditional empirically observed length distribution in the training dataset.
Figure 5: A sample for the slide door task conditioned on the first and last frame (highlighted).
Table 4: Dimension prediction mean absolute error for three planning tasks with standard deviations estimated over 45 samples.
Method	Stationary (↓)	Slide Door (↓)	Push Object (↓)	Average (↓)
FDDM	14.16±1.41	13.39±1.34	17.06±1.47	14.87
TDDM	9.70±0.99	11.47±0.74	15.43±0.90	12.2
In this work, we highlighted the pitfalls of performing generative modeling on varying dimensional datasets when treating state values and dimensions completely separately. We instead proposed a trans-dimensional generative model that generates both state values and dimensions jointly during the generative process. We detailed how this process can be formalized with the time-reversal of a jump diffusion and derived a novel evidence lower bound training objective for learning the generative process from data. In our experiments, we found our trans-dimensional model to provide significantly better dimension generation performance for diffusion guidance and interpolations when conditioning on properties that are heavily correlated with the dimension of a datapoint. We believe our approach can further enable generative models to be applied in a wider variety of domains where previous restrictive fixed dimension assumptions have been unsuitable.
The authors are grateful to Martin Buttenschoen for helpful discussions. AC acknowledges support from the EPSRC CDT in Modern Statistics and Statistical Machine Learning (EP/S023151/1). AD acknowledges support of the UK Dstl and EPSRC grant EP/R013616/1. This is part of the collaboration between US DOD, UK MOD and UK EPSRC under the Multidisciplinary University Research Initiative. He also acknowledges support from the EPSRC grants CoSines (EP/R034710/1) and Bayes4Health (EP/R018561/1). WH and CW acknowledge the support of the Natural Sciences and Engineering Research Council of Canada (NSERC), the Canada CIFAR AI Chairs Program. This material is based upon work supported by the United States Air Force Research Laboratory (AFRL) under the Defense Advanced Research Projects Agency (DARPA) Data Driven Discovery Models (D3M) program (Contract No. FA8750-19-2-0222) and Learning with Less Labels (LwLL) program (Contract No.FA8750-19-C-0515). Additional support was provided by UBC's Composites Research Network (CRN), Data Science Institute (DSI) and Support for Teams to Advance Interdisciplinary Research (STAIR) Grants. This research was enabled in part by technical support and computational resources provided by WestGrid (https://www.westgrid.ca/) and Compute Canada (www.computecanada.ca). The authors would like to acknowledge the use of the University of Oxford Advanced Research Computing (ARC) facility in carrying out this work. http://dx.doi.org/10.5281/zenodo.22558
[1] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. International Conference on Machine Learning, 2015.
[2] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 2020.
[41] Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. International Conference on Machine Learning, 2022.
[42] William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. Advances in Neural Information Processing Systems, 2022.
[43] John L Kelley. General Topology. Courier Dover Publications, 2017.
[44] Stewart N Ethier and Thomas G Kurtz. Markov processes: Characterization and convergence. John Wiley & Sons, 2009.
[45] Giovanni Conforti and Christian Léonard. Time reversal of Markov processes with jumps under a finite entropy condition. Stochastic Processes and their Applications, 2022.
[46] Clement Vignac and Pascal Frossard. Top-n: Equivariant set and graph generation without exchangeability. International Conference on Learning Representations, 2022.
[47] Margaret Mitchell, Simone Wu, Andrew Zaldivar, Parker Barnes, Lucy Vasserman, Ben Hutchinson, Elena Spitzer, Inioluwa Deborah Raji, and Timnit Gebru. Model cards for model reporting. Proceedings of the Conference on Fairness, Accountability, and Transparency, 2019.