Trans-Dimensional Generative Modeling via Jump Diffusion Models
Authors: Andrew Campbell
William Harvey
Christian Dietrich Weilbach
Valentin De Bortoli
Tom Rainforth
Arnaud Doucet
Venue: NeurIPS 2023 spotlight (2023)
Keywords: diffusion
score-based
score
markov chain
jump diffusion
poisson
TL;DR: We propose a new model class for generating varying dimensional data using jump diffusions.

Abstract: We propose a new class of generative model that naturally handles data of varying dimensionality by jointly modeling the state and dimension of each datapoint. The generative process is formulated as a jump diffusion process that makes jumps between different dimensional spaces. We first define a dimension destroying forward noising process, before deriving the dimension creating time-reversed generative process along with a novel evidence lower bound training objective for learning to approximate it.
Simulating our learned approximation to the time-reversed generative process then provides an effective way of sampling data of varying dimensionality by jointly generating state values and dimensions. 
We demonstrate our approach on molecular and video datasets of varying dimensionality, reporting better compatibility with test-time diffusion guidance imputation tasks and improved interpolation capabilities versus fixed dimensional models that generate state values and dimensions separately.

Trans-Dimensional Generative Modeling via Jump Diffusion Models
Andrew Camphell'
William Harvey²
Christian Weilbach²
Valentin De Bortoli
Tom Rainforth
Arnaud Doucet
Department Statistics, University of Oxford, UK
Department of Computer Science, University of British Columbia, Vancouver, Canada SCNS ENS Ulm, Paris, France
{campbell, rainforth, doucet)@stats, ac.
{wsgh, weilbach)@cs ubc. ca
valentin. dew.eroooooooolailllcccmommmo com
Abstract
We propose new class of generative models that naturally handle data of varying dimensionality by jointly modeling the state and dimension of cach datapoint. The generative process formulated as jump diffusion process that makes jumps be- tween different dimensional spaces. We first define dimension destroying forward noising process, before deriving the dimension creating time-reversed generative process along with novel evidence lower bound training objective for learning to approximate Simulating our learned approximation to the time-reversed generative process then provides effective way of sampling data of varying dimensionality by jointly generating state values and dimensions. We demonstrate approach molecular and video datasets of varying dimensionality, report- ing better compatibility with test-time diffusion guidance imputation tasks and improved interpolation capabilities versus fixed dimensional models that generate state values and dimensions separately.
1 Introduction
Generative models based on diffusion processes 1 have become widely used solving range of problems including text-to-image generation 4 audio synthesis [6] and protein design [7]. These models define a forward noising diffusion process that corrupts data to noise and then learn corresponding time-reversed backward generative process that generates novel datapoints from noise.
In many applications, for example generating novel molecules or videos 10,, the dimension of the data can also vary. For example, molecule can contain a varying number of atoms and video can contain a varying number of frames. When defining generative model over these data-types. it is therefore necessary to model the number of dimensions along with the raw values of each of its dimensions (the state). Previous approaches to modeling such data have relied on first sampling the number of dimensions from the empirical distribution obtained from the training data, and then sampling data using a fixed dimension diffusion model (FDDM) conditioned on this number of dimensions [8]. For conditional modeling, where the number of dimensions may depend on the observations, this approach does not apply and we are forced to first train an auxiliary model that predicts the number of dimensions given the observations ||
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

diffuse
j jump
diffum fuse
Figure 1: Illustration of the jump diffusion generative process on videos and molecules. The generative process consists two parts: a diffusion part which denoises the current set of frames/atoms and jump part which adds on suitable number of new frames/atoms such that the final generation is clean synthetic datapoint of an appropriate size.
This approach to trans-dimensional generative modeling is fundamentally limited due to the complete separation of dimension generation and state value generation. This is exemplified in the common use case of conditional diffusion guidance. Here, an unconditional diffusion model is trained that end-users can then easily and cheaply condition on their task of interest through guiding the generative diffusion process [3 12 [4| without needing to perform any further training or fine-tuning of the model on their task of interest. Since the diffusion occurs in fixed dimensional space, there is no way for the guidance to appropriately guide the dimension of the generated datapoint. This can lead incorrect generations for datasets where the dimension greatly affects the nature the datapoint created, e.g. small molecules have completely different properties to large molecules.
To generate data of varying dimensionality, we propose a jump diffusion based generative model that jointly generates both the dimension and the state, Our model can be seen as a unification of diffusion models which generate all dimensions in parallel with autoregressive type models which generate dimensions sequentially. We derive the model through constructing a forward noising process that adds noise and removes dimensions and backward generative process that denoises and adds dimensions, see Figure [ We derive the optimum backward generative process as the time-reversal of the forward noising process and derive novel learning objective learn this backward process from data. We demonstrate the advantages of our method on molecular and video datasets finding our method achieves superior guided generation performance and produces more representative data interpolations across dimensions.
2 Background
Standard continuous-time diffusion models 3 17 define forward diffusion process through stochastic differential equation (SDE) where Xg Paata and, for t > 0,
dx; b b. g:dw.,
where x, E Rd the current state, b.: Rd - Idd is the drift and ge E R is the diffusion coefficient. dw, is Brownian motion increment on Rd. This SDE can be understood intuitively by noting that in each infinitesimal timestep, we move slightly the direction the drift b, and inject small amount of Gaussian noised governed by gr. Let pt
(xx) denote the distribution of for the forward
diffusion process 1 so that po (Xo) = Painta (xo). b and St are set such that at time f T. pT(T))
close to (xT) A(xT;; 0, La): e.g. (x(x))
3,x,, g: 3 for 3,
12 3].
The time-reversal of the forward diffusion (I is also diffusion 1 19 which runs backwards in
time from (x(r) to Pu (xx) and satisfies the following reverse time SDE
dx, (x(xx) g:dw
where F((x,) (x,) Vx logg:((x) (x,), ), di is a negative infinitesimal time step and dw, is Brownian motion inerement when time flows backwards. Unfortunatcly, both the terminal

---

distribution, Pr (xr), and the score, Vx, (x,), are unknown in practice. A generative model obtained by approximating with Prof and learning an approximation (x,)) to Vx, log Pe typically using denoising score matching 20 i.c.
min (t.. (xx) Vx, log Prjo(X; [xo)]
(2)
For flexible model class, s" we get s. (x,) zo Vx, (x)) at the minimizing parameter.
3 Trans-Dimensional Generative Model
Instead of working with fixed dimension datapoints, we will instead assume our datapoints consist of variable number of components. A datapoint X consists of n components each of dimension d. For ease of notation, each datapoint will explicitly store both the number of components, n, and the state values, giving X (n, Since each datapoint have variable number of components from n = 1 to N, our overall space that our datapoints live in is the union all these possibilities, N
X X U n= [m] x Rnd For example, for a varying size point cloud dataset, components would refer to points in the cloud, each containing (ar 2) coordinates giving 3 and the maximum possible number of points in the cloud is N.
Broadly speaking, models. We will first define a forward noising process that both corrupts state values with Gaussian noise and progressively deletes dimensions. We then learn approximation to the time-reversal giving a backward generative process that simultaneously denoises whilst also progressively adding dimensions back until synthetic datapoint of appropriate dimensionality has been constructed.
approach will follow the same framework previous diffusion generative
3.1 Forward Process
Our forward and backward processes will be defined through jump diffusions. A jump diffusion process has two components, the diffusion part and the jump part. Between jumps, the process evolves according to standard SDE. When a jump occurs, the process transitions to a different dimensional space with the new value for the process being drawn from transition kernel K2(Y X):
X' x + Rzo. Letting Y = (m, y), the transition kernel satisfies Σn. Jy K2 (m, y|X)dy = 1 and Iy (m n. y X)dy= 0. The rate at which jumps occur (jumps per unit time) is given by a rate function (x) X + R20 For an infinitesimal timestep dt, the jump diffusion can be written as
X Y KY(XX.) with probability Ar Xx(X.)dt Xt X (x())) + g:dw, Re+ddt n/
with probability 1 - (Xx))
Jump
X
Diffusion
with X, - (n,, and X A (Me+de, X(+de) ) and dw, being a Brownian motion increment on Radd We provide a more formal definition in Appendix A
With the jump difffusion formalism in hand, We will use the diffusion part to corrupt existing state values with Gaussian noise and the jump part to destroy dimensions. For the diffusion part, we use the VP-SDE introduced in 2 3 with
construct our forward noising process.
b (x( = 3,x and g t = v3, with 3r 2 O.
When jump occurs in the forward process, one component of the current state will be deleted. For example, one point a point cloud or single frame video is deleted. The rate at which these deletions occur is set by user-defined forward rate X (X).. To formalize the deletion, we need to introduce some more notation. We let de (iln) of the current state to delete. We also define del X' x N + X' to be the deletion operator that deletes specified component. Specifically, (n 1,y) del((n, (nz, x). 2) where E [R 1)d has the same values as x E Rnd except for the d values corresponding to the ith component which have been removed. We now define the forward jump transition kernel as K (Y (Y Σn de (ill))) (ilm))aae(x (Y). We note that only one component is ever deleted time meaning (. Further, the choice of K (iln) will dictate the behaviour of the reverse generative process. we set K d (ilm)
user-defined distribution over which component
0 for m + 1.
lii n. then we only ever delete the final component and so in the reverse generative
3

---

Table 1: Summary of forward and parameterized backward processes
Direction
b,
(x)
K. X)
Forward - Pxx
v
(((n)
E Kdd (3)m)) (2))))
Backward
-30xx (
VB
XX(
Sycas 2im Al (y mn yst,, ( )dy
direction, datapoints created additively, appending components onto the end of the current state. Alternatively, if we set (ilm) = 1/n then components are deleted uniformly at random during forward corruption and in the reverse generative process, the model will need to pick the most suitable location for a new component from all possible positions.
The forward noising process is simulated from 0 to T and should be such that time t t= the marginal probability pr (X) should be close to a reference measure (re((X) that can be sampled from. We set Prur (X) = I[nn [x: 0, Ia) where I[n 1} is 1 when = 1 and 0 otherwise. To be close Ptef, for the jump part, set X, high enough such that time & T there is a high probability that all but one of the components the original datapoint have been deleted. For simplicity, we also set X to depend only on the current dimension Xx()) - (n) with Xx(n 1) 0 so that the forward process stops deleting components when there is only 1 left. In our experiments, we demonstrate the trade-off's between different rate schedules in time. For the diffusion part, we use the standard diffusion 3n schedule 2 3 so that we are close to N N(x: 0, La).
'
3.2 Backward Process
The backward generative process will simultaneously denoise and add dimensions back in order to construct the final datapoint. It will consist of backward drift b (X), diffusion coefficient & .. rate xx(X) and transition kernel K , (Y(X). We would like these quantities to be such that the backward process is the time-reversal of the forward process. In order find the time-reversal of the forward process, we must first introduce some notation to describe K (Y(X). (Y())) should undo the forward deletion operation. Since (Y(Y) chooses component and then deletes it, K((Y) X) will need to generate the state values for a new component, decide where the component should be placed and then insert it at this location. Our
component will be denoted yd Rd The insertion operator is defined as ins : X' x R2 x N + X. takes in the current value X, the new component ydd and an index i {1, ..... n + }} and inserts yaul into X location such that the resulting value Y = (ns(X, y i) has del( Y, = X. We denote the joint conditional distribution over the newly added component and the index at which it is inserted (y i i). We therefore have (Y = Syw En+1 Ar(yauld and mms) (Y dyyddd Noting that only one component is ever added at a time, we have (m = 0 for + + 1.
This backward process formalism can be seen as a unification of diffusion models with autoregressive models. The diffusion part b, denoises the current set of components in parallel, whilst the autoregressive part A Ayyyad, i|X) predicts between these parts controlling when and how many new components are added during generation.
component and its location. xx( is the glue
We now give the optimum values for ((()). T. xx(()) and A, (yddd, i ;iX) such that the backward process is the rime-reversal of the forward process.
Proposition 1. The time reversal of forward jump diffusion process given by drift D., diffusion coefficient ÿ. t. rate X R((n) and transition kernel Σ (iln))aa is given by a jump diffusion process with drift (2(X), diffusion coefficient 'g rate X (X) and transition kernel

---

syas Σx= A7 ) Ins (X, yau .)) (Y) )dyadd as defined below
Di(x) = =b,(x) V log -i=9. g'.
x((x)= = Xx(n+ 1) i= K (ipr n + 1) Syne Pr (X,, you, dy)dd
(x(
Ay , & pr (ins((XX,
y yould,
(i|m + 1).
All proofs are given in Appendix A The expressions for b; and :i are the same for a standard diffusion except for replacing Vx log Pr (x) with Vx lo = Vx pe which is simply the
score in the current dimension. The expression for X; : can be understood intuitively by noting that the numerator in the probability ratio is the probability that at time t, given deletion occurs, the forward process will arrive at X. If this is higher than the raw probability time that the forward process is at X (the denominator) then we should have high X because X is likely the result of a delction of a larger datapoint. Finally the optimum A (y i iXX) is simply the conditional distribution of you and i given X when the joint distribution over yoll i, X is given by pr (iisssX, y
y i)) (i/m (ill + +1.
3.3 Objective for Learning the Backward Process
The true bi: ; and A; are unknown so we need to learn approximations to them, 5 X and Af. Following Proposition 1 we set (X) (X) 772 (X) where sf (X) approximates
Vx logo p. ()). The forward and parameterized backward processes are summarized in Table
Standard diffusion models are trained using a denoising score matching loss which can be derived from maximizing an evidence lower bound on the model probability for Ep (x,) (log (xx)] (21).
We derive here an equivalent loss to learn s!. X and A7 for our jump diffusion process by leveraging the results of [I7] and (22). Before presenting this loss, we first introduce some notation. Our
objective for of (X,) will resemble denoising score matching (2) but instead involve the conditional score Xx, log (X,
Xx, Po (xx [Xo, nt). This difficult to calculate directly due to a combinatorial sum over the different ways the components of Xo can be deleted to get to X, We avoid this problem by equivalently conditioning on mask variable M, f0, 1 that 0 for components of Xo that have been deleted to get to X, and I for components that remain in X., This makes
our denoising score matching target easy to calculate: Xx, log prio (x n4, M,) y /or M, (Xo)
1-oz
where a exp(- 3(s)ds) [3]. Here Mc(xx) is the vector removing any components in x for which M2 is 0, thus M, (Xo) and x, have the same dimensionality. We now state our full objective. Proposition 2. For the backward generative jump diffusion process starting at Pret (Xr) and finishing at (X)), an evidence lower bound on the model log-likelihood Exo Dlg p (xo) is given by
(())=
T
F E (2 (X.) Xx, log Pt o (xe|Xo n4, M))
2
TE X((())) (m:) log X
+ X.(ne) log A (x m|x)]
+ +C,
where expectations are with respect to ld(t: 0, :o) X,, Kede (ilm,) () constant term independent of and X, = ins() Y, x! i). This evidence lower bound is equal to the log- -ogliiilloood when b = b; X X ; and A AZ.
(Y), C is
We now examine the objective to gain an intuition into the learning signal. Our first term (3) is an Le regression to target that,
As the solution to an La regression problem is the conditional expectation of the target, sf (X,) will learn to predict vectors pointing towards Xy averaged over the possible correspondences between dimensions of x, and dimensions of Xo. Thus, during sampling, (X.) provides a suitable direction to adjust the current value X, taking into account the fact X, represents only a noisy subpart of clean whole Xo-
have seen, is scaled vector between x2 and
The second term (4) gives learning signal for X & and AÇ For A4, we simply have maximum likelihood objective, predicting the missing part of X, (i.e. x given the observed part of X, (i.c. Y). The signal for X comes from balancing two terms: - A (X,) and X (((r.) log x (Y)
5

---

which encourage the value of X to move in opposite directions. For test input Z, X ( value needs to trade off between the two terms by learning the relative probability between Z being the entirety of genuine sample from the forward process, corresponding to the (xx)) term in 4 or being substructure a genuine sample, corresponding to the X (Y) term in The optimum trade-off is found exactly at the time reversal x show in Appendix A.5
(Z)'s
We optimize L(0) using stochastic gradient ascent, generating minibatches by first sampling & ~ (1(0, T), X ~ Pástta and then computing X, from the forward process. This can be done analytically for the (n) functions used in our experiments. We first sample n, by analytic integration of the dimension deletion Poisson process with time inhomogeneous rate X (n).. We then add Gaussian noise independently to each dimension under Prjo Xo. n4, M,) using a randomly drawn mask variable M, See Appendix B for further details the efficient evaluation of our objective.
3.4 Parameterization
(Xt), AQ (y 2 x(x,) and (X.) will all be parameterized by neural networks. In practice, we have a single backbone network suited to the problem of interest e.g. Transformer [23], an EGNN
24] or a UNet (25) onto which we add prediction heads for (X.), (y :i(x,) and X (X,.. of (Xx) outputs vector in gadd (y i[X,) outputs a distribution over i and mean and standard
deviation statistics for Gaussian distribution over y Finally, having X (X,, E Rzo be the
raw output of a neural network can cause optimization issues due to the optimum T including probability ratio which can take on very large values. Instead, we learn a component prediction network Phj. (no Xx) that predicts the number of components in Xo given X, To convert this into
(x,))) is an analytic function of the true Poje(nio) X.). We then plug Poje(nn X.) into Proposition 3 to obtain an approximation of X (X.).
(X,)), we show in Proposition 3 how the optimum
Proposition 3. We have
N
Peio[n + 1/mo)
Ti(K,)= Xe(ne +
(ejo mo) Poja (no(
where X, (nt, (1) and Pno(n, + 1(ma) and Pap(na no) are both easily calculable distributions from the forward dimension deletion process.
3.5 Sampling
Algorithm 1: Sampling the Generative Process
t t
To sample the generative process, we numerically integrate the learned backward jump diffusion process us- ing time-step St. Intuitively, is sim- ply the standard continuous time diffu- sion sampling scheme 3 but each timestep we check whether a jump has occurred and if has, sample the new component and insert it at the chosen index as explained by Algorithm [
X ~ Perl ()) = I(n = (x: Ln) while 1 > d
if (x))t with a 4(0. 1) then Sample x i A (x i end
ins(X, x!"
x ((()) ge Vôte with NN(0, Ina) X (n, x- S1
end
Related Work
Our method jointly generates both dimensions and state values during the generative process whereas prior approaches 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 images 227 28 hand pick dimension jump points such that the

---

resolution of images is increased during sampling and reaches a certain pre-defined desired resolution the end of the generative process. Further, rather than using any equivalent of A4, 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 the discrete space ELBO derived in [2]
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 unnormalized version of its density [30 an ergodic jump diffusion is designed which admits the target as the invariant The invariant distribution is not preserved when dime-discretizing the process.
distribution 30 32
how general jump proposals could be built and how this process could be "Metropolized" 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 samples in the form of data, not an unnormalized target.
However, was shown in [33 34
5 Experiments
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: (z, y, coordinates, 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 sf p ( (,)) and Ag- We uniformly delete dimensions, Kedel = 1/n, and since point cloud is permutation invariant, Af(y (y X.) need only
only predict new dimension values. We set X to a constant except for 0.12' where we set X i< IT 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, sample quality in Table and find we perform comparably to the results reported in g which use an FDDM. We ablate our choice of X, by comparing with setting Xx to a constant for all l and with setting X, = 0 for 1 < 0.9T (rather than just for < 0.17). We find that the constant X, performs worse due to the occasional component being added late in the generation process without enough time for the diffuuson process to finalize its value. We find the Xt<o.97 = 0 setting to have satisfactory sample quality however this choice of Xx introduces issues during diffusion guided generation as we see next. Finally, we ablate the parameterization of Proposition 3 by learning X (X,) IR directly as the output of neural network head. We find that this reduces sample quality
first verify our unconditional
due to the more well-behaved nature of the target, Po (x)) when using Proposition 3 We note pure autoregressive models perform significantly worse than diffusion based models found in g
Figure 2: Visualization of the cump-diffusion backward generative process on molecules.
7

---

Table 2: Sample quality metrics for unconditional molecule generation. An atom is stable it has the correct valency whilst molecule is considered stable if all of its atoms are stable. Molecular validity All methods use 1000 simulation steps and draw 10000 samples.
measured using iDKit 37
% Atom Stable (t)
% Molecule Stable (t)
Method
% Valid (f)
FDDM 8
98.7
82.0
91.9
TDDM (ours)
98.3
87.2
92.3
TDDM, const Xx
96.7
79.1
86.7
TDDM, X x
97.7
82.6
89.4
TDDM w/o Prop.
97.0
66.9
87.1
Table 3: Conditional Molecule Generation for 10 conditioning tasks that each result in 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 average sample quality metrics. Standard deviations are given across the 10 conditioning tasks. We report bold values that statistically indistinguishable from the best result at the 5% level using two-sided Wilcoxon signed rank test across the 10 conditioning tasks.
Dimension Error (4)
% Atom
% Molecule Stable (1)
Method
Stable (t)
% Valid (t)
FDDM
0.511
93.5+1..1
31.3+6.3
65.2+10.3
TDDM
0.13440.076
93.5+2.6
59. 1911
74.8+9.3
TDDM, const X
0.226++0.17
88.944.8
43.6+15
63.4+14
TDDM, X 0
0
0.390##0.3
95. 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 unconditional model in order to generate molecules that contain a certain number of desired atom types, e.g. 3 carbons or 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 (X,) such that it approximates Vx x, log Pt (X, ly) rather than Vx, log pr (X,) (where y is the conditioning information) by adding on a term approximating )xx This guides x, such that it is consistent with y. Since (X, has access x,, it will cause n, to automatically also be consistent with without the user needing to input any information on how the conditioning information relates to the size of the datapoints. We give further details diffusion guidance in Appendix C
Vx, log (y X2) with (5(y) (x) Eno Sx, p(y (x) Poje (Xo X,
We show our results in Table In order to perform guidance on the FDDM baseline, we implement the model from [ in continuous time and initialize the dimension from the empirically observed dimension distribution in the dataset. This accounts for the case of end user attempting to guide 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 conditioning information. The X r<00.97 = 0 ablation on the other hand only marginally improves the dimension error over FDDM because all dimensions are added in the generative process time when X; is noisy and has little relation to the conditioning information. This highlights the necessity of allowing dimensions to be added throughout the generative process gain the trans-dimensional diffusion guidance ability. The ablation with constant X has increased dimension error TDDM as we find that when X, > for all t, X can become very large when is close to 0 when the model has perceived lack of dimensions. This occasionally results in too many dimensions being added hence an increased dimension error. Not using the Proposition parameteriizaiinn also increases dimension error due to the increased difficulty in learning N
8

---

5.1.2
Interpolation
Interpolations are unique way of gaining in- sights into the effect of some conditioning pa- rameter on dataset of interest. To create interpolation, a conditional generative model is first trained and then sampled with sweep of the conditioning parameter but using fixed ran-
30
Dataset
Number Atoms
25
FDDM
TDDM
15
dom noise [8) The resulting series of synthetic datapoints share similar features due to the fixed random noise but vary in ways that are very in- formative as to the effect of the conditioning parameter. Attempting
10
330
Polarizability
interpolate with FDDM fundamentally limited because the en- tire interpolation occurs in the same dimension which is unrealistic when the conditioning pa- rameter is heavily correlated with the dimension of the datapoint. We demonstrate this by fol- lowing the setup of [8] who train a conditional FDDM conditioned on polarizability. Polariz-
Figure 3: Number of atoms versus polarizability for 3 interpolations with fixed random noise. The dataset mean and standard deviation for the num- ber of atoms is also shown. FDDM interpolates entirely in a fixed dimensional space hence the number of atoms is fixed for all polarizabilities.
ability is the ability of molccule's electron cloud to distort in response an external electric field (38) with larger molecules tending to have higher polarizability. To enable us to perform dimenssonl interpolation, we also train conditional version of our model conditioned on polarizability. An example interpolation with this model is shown in Figure E We find that indeed size molecule increases with increasing polarizability, with some molecular substructures e.g. rings, being maintained across dimensions. We show how the dimension changes with polariz- ability during 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 4: Sequence of generations for linearly increasing polarizability from 39 Bohr to 66 Bohr 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 nitrogen (blue) carbon (gray) and an extra hydrogen (white) is added to maintain the correct valency.
5.2
Video
We finally demonstrate our model on video modeling task. Specifically we model the RoboDesk dataset [39]. video benchmark to measure the applicability of video models for planning and control problems. The videos are renderings of robotic arm [40] performing variety of different tasks including opening drawers and moving objects. We first train unconditional model on videos of varying length and then perform planning by applying diffusion guidance to generate videos conditioned on initial starting frame and a final goal frame [1.. The planning problem is then reduced to "filling in" the frames in between. Our trans-ding dinenssina model automatically varies the number of in-filled frames during generation 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.
that the final length of video matches the length of
We model videos at 32 32 resolution and with varying length from 2 to 35 frames. For the network backbone, we use UNet adapted for video [42]. In contrast to molecular point clouds, our data is no longer permutation invariant hence A yd (.) includes prediction over the location to insert the new frame. Full experimental details are provided in Appendix [ We evaluate our approach on three planning tasks, holding stationary, sliding door and pushing an object. An example generation conditioned on the first and last frame for the slide door task shown in Figure 5 with

---

the model in-filling plausible trajectory. We quantify our model's ability to generate videos of length appropriate the task in fable4 on all three tasks we generate more accurate length of video than FDDM which is forced to sample video lengths from the unconditional empirically observed length distribution 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 (L)
Slide Door (4)
Push Object (L)
Average (L)
FDDM
14.16+1.11
13.39
17.06=1.47
14.87
TDDM
9.70=0.99
11.47
15.43=0.90
12.2
Discussion
In this work, we highlighted the pitfalls performing generative modeling on varying dimensional datasets when treating state values and dimensions completely separately. We instead proposed 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 jump diffusion and derived a novel evidence lower bound training objective for learning the generative found
process from data. In our experiments, better dimension generation performance for diffusion guidance and interpolations when conditioning on properties that heavily correlated with the dimension of datapoint. We believe our approach can further enable generative models to be applied in wider variety of domains where previous restrictive fixed dimension assumptions have been unsuitable.
dimenssonnn model provide significantly
Acknowledgements
The authors are grateful Martin Buttenschoen for helpful discussions. AC acknowledges support from the EPSRC CDT in Modern Statistics and Statistical Machine Learning (SP/S) AD acknowledges support of the UK Dstl and EPSRC grant EP/R0136 1611 This is part of the col- laboration 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/) 10/1) and Bayes4Health (EP/RO18561/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 Discov- ery Models (D3M) program (Contract No. 148750-19-2-02223 and Learning with Less Labels (LwLL) program (Contract C..FA8750-19-C-001
Additional support was provided by UBC's Composites Research Network (CRN), Data Science Institute (DSI) and Support for Teams to Ad- vance Interdisciplinary Research (STAIR) Grants. This research was enabled in part by technical support and computational resources provided by WestGrid (https: //wp:: 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. 22558
http:/ //xx doi org/ 10.
References
[1] Jascha Sohl-Dich Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised 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.
10

---

[3] Yang Song, Jascha Sohl- Dick-Dickstein, Diederik P Kingma, Abhishek Kumar. Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. International Conference on Learning Representations, 2021.
[4] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, Jonathan Ho, and Mohammad Fleet, David aand Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 2022.
[5] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06/25 2022.
[6] Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. International Conference on Learning Representations, 2021.
[7] Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian Trippe, Jason Yim, Helen Eisenach, Woody Ahem, Andrew Borst, Robert J Ragotte, Lukas F Milles, et al. Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models. bioRxiv, 2022.
[8] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. International Conference on Machine Learning, 2022.
[9] Jonathan Ho, Tim Salimans, Alexey Gritsenko. William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. Advances in Neural Information Processing Systems, 2022.
[10] Jonathan Ho. William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik Kingma, Ben Poole, Mohammad Norouzi, David Fleet, and Tim Saliman. Imagen video: High definition video generation with diffusion models. arXiv preprinl arXiv:2210.02303. 2022.
[11] Ilia Igashov, Hannes Stärk, Clément Vignac, Victor Garcia Satorras, Pascal Frossard, Max Welling. Michael Bronstein, and Bruno Correia. Equivariant 3d-conditional diffusion models for molecular linker design. arXiv preprint arXiv:2210.05274, 2022.
[12] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat GANs image synthesis.
Advances in Neural Information Processing Systems, 2021.
[13] Katherine Crowson.
Clip guided diffusion. wtwparrrrwwrere-iuuddddiieieieee guided-diffusion, 2021.
Web
Demo,
https://huggingfucecco httpss/htngiinfiiccco/ppcsssssst http:////...ttttteeeeeeeeeeeeeee
[14] Hengtong Zhang and Tingyang Xu. Towards controllable diffusion models via reward-guided exploration. arXiv preprint 20XV02230..0712,, 2023.
[15] Chin-Wei Huang, Jae Hyun Lim, and Aaron C Courville. A variational perspective on diffusion- based generative models and score matching. Advances in Neural Information Processing Systems, 2021.
[16] Tero Karras, Mika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of caitusion-baeel generative models. Advances Neural Information Processing Systems, 2022.
[17] Joe Benton, Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. From denoising diffusions to denoising Markov models. arXiv preprint arXiv:2211.03553 2022.
[18] Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 1982.
[19] Uirich G Haussmann and Etienne Pardoux. Time reversal of diffusions. The Annals of Probability, 1986.
[20] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computation, 2011.
11

---

[21] Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in Neural Information Processing Systems, 2021.
[22] Patrick Cheridito, Damir Filipovié, and Mare Yor. Equivalent and absolutely continuous measure changes for jump-diffusion processes. Annals of applied probability, 2005.
[23] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
[24] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. International Conference on Machine Learning, 2021.
[25] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. Medical Image Computing and Computer-Assisted Interven- sion, 2015.
[26] Tomer Weiss, Luca Cosmo, Eduardo Mayo Yanes, Sabyasachi Chakraborty, Alex M Bronstein, and Renana Gershoni-Poranne. Guided diffusion for inverse molecular design. ChemrXiv, 2023.
[27] Bowen Jing, Gabriele Corso, Renato Berlinghieri, and Tommi Jaakkola. Subspace diffusion generative models. European Conference on Computer Vision, 2022.
[28] Han Zhang. Ruili Feng, Zhantao Yang, Lianghua Huang, Yu Liu. Yifei Zhang, Yujum Shen, Deli Zhao, Jingren Zhou, and Fan Cheng. cimensionality-waryinn diffusion process. arXiv preprint arXiv:2211.16032, 2022.
[29] Andrew Campbell, Joe Benton, Valentin De Bortoli, Tom Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 2022.
[30] Ulf Grenander and Michael I Miller. Representations of knowledge in complex systems. Journal of the Royal Statistical Society: Series (Methodologicall, 1994.
[31] David Phillips and Adrian M Smith. Bayesian model comparison via jump diffusions. Markov Chain Monte Carlo in Practice, 1995.
[32] Michael Miller, Ulf Grenander, Joseph A O" Sulllivn,, and Donald L Snyder. Automatic target recognition organized via jump-diffusion algorithms. IEEE Transactions on Image Processing, 1997.
[33] Peter Green. Reversible jump Markov chain Monte Carlo computation and Bayesian model determination. Biometrika, 1995.
[34] Peter Green. Trans- dimensiimensoonal Markov chain Monte Carlo. Highly Structured Stochastic Systems, 2003.
[35] Lars Ruddigkeit, Ruud Van Deursen, Lorenz C Blum, and Jean-Louis Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17. Journal of chemical information and modeling, 2012.
[36] Raghunathan Ramakrishnan, Pavlo Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 2014.
[37] Rdkit: Open-source cheminformatics; http:://ww.rddte rdkit org. Accessed 2023.
[38] Eric V Anslyn and Dennis A Dougherty. Modern physical organic chemistry. University Science Books, 2005.
[39] Stephen Tian, Chelsea Finn, and Jiajun Wu. A control-centric benchmark for video prediction. arXiv preprint arXiv:2304. /3723, 2023.
[40] Harini Kannan, Danijar Hafner, Chelsea Finn, and Dumitru Erhan. Robodesk: A multi-task re- inforcement learning benchmark. https: // / github. com/ ppogle-researcth robooesskooooosko 2021.
12

---

[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 Leurning Representations, 2022.
[47] Margaret Mitchell, Simone Wu, Andrew Zaldivar, Parker Barns,, Lucy Vasserman, Ben Hutchin- son, Elena Spitzer, Inioluwa Deborah Raji, and Timnit Gebru. Model cards for model reporting. Proceedings of the Conference on Fairness, Accountabiliry, and Transparency, 2019.
13

---

Appendix
This appendix is organized as follows. In Section A present proofs for all of our propositions. Section presents a rigorous definition of our forward process using more specific notation. This is then used in (.2c.oon/.2.Ipp prove the time reversal for our jump diffusions. We also present an intuitive proof of the time reversal using notation from the main text in Section A.2.2 In Section A.3 we prove CropositionnJusing the notation from the main text. We prove Proposition 3 in Section A.4 and we analyse the optimum of our objective directly without using stochastic process theory in
Section A.5 In Section [ we give more details our objective and in Section we detail how apply diffusion guidance to our model. We give the full details for our experiments in Section D and finally, in Section E. discuss the broader impacts of our work.
A Proofs
A.1
Notation and Setup
We here introduce a more rigorous notation for defining our trans-dimensional notation that will be used in rigorous proof for the time-reversal of our jump diffusion. First, while it makes sense from a methodological and experimental point of view to present our setting as a transdimensional one, slightly change the point of view in order to derive theoretical results. We extend the space Rd to R2 Rd U {oo} using the one-point compactification of the space. We refer details on this space. The point will be understood as mask. For instance, let 1 42, 2g RRd
44 for
Then X (z1, #2,
(ied) with N 3 corresponds to vector for which all components observed whereas X' (2i oc, E (IRd)N corresponds a vector for which only the components on the first and third dimension are observed. The second dimension is masked in that case. Doing so, we will consider diffusion models on the space X = (iR))N with d, N E N. In the case of a video diffusion model, N can be seen as the max number of frames. We will always consider that this space is equipped with its Borelian sigma-field X' and all probability measures will be defined on X.
We denote dim: X + 13. which is given for any X = (xr) E X by
dim(X) (I:))
In other words, dim(X) is binary vector identifying the "dimension" of the vector X, i.e. which frames are observed. Going back to our example X (21, E (ied) and X' (211, oc, 22)) (R&)), we have that dim(X)
11, and dim(X') (1,0,1) 1). For any vector u. E {0, 11 we denote |u| = Ei us, i.e. the active dimensions of (or equivalently the non- masked frames). For any E and D E {0, we denote Xp = {X) with X X, D = and X = oo if D. = 0.
We denote (R R) the set of functions which are differentiable and bounded. Similarly, we denote Cj [Ra R) the set of functions which are & differentiable and compactly supported. The set C (Rd R) denotes the functions which are & differentiable and vanish when ||>||| + +oo. We note that f E C(Rd), if f E C(Rd) and f - ((o) E CIRd) and that f E C* (fRd)) for any & E N if the restriction of f to Rd is in O* (Rd) and f E C(R.).
A.1.1 Transdimensional infinitesimal generator
To introduce rigorously the transdimensional diffusion model defined in Section 3.1 we will introduce its infinitesimal generator. The infinitesimal generator of stochastic process can be roughly defined as its "probabilistic derivative". More precisely, assume that a stochastic process (X))e2 admits a transition semigroup (Pt)ızu, i.e. for any t 0, A X' and X E X we have P(X, A X,
I) P (z, A), then the infinitesimal generator is defined as A(f) = lim, (P (P,(f) - f)/t, for every for which this quantity is well-defined.
Here, we start by introducing the infinitesimal generator interest and give some intuition about its form. Then, prove a time-reversal formula for this infinitesimal generator.
14

---

We consider b: R + RR, a: {0,1) - RR. For any f (X) and X E X we define
A()(X) = ( vx.((X)) +
(X:)
(5)
- -DD.CD E DD. AM a(Do..... ...,, ,DM) E s(XD,,, 1)DD (dim(X)),
where M E N. Do {1} (A,) E N such that Σj=0 A N and for any j {0, .O......4. M 111. D is the subset of {0, (0,1) 1)....... such that D Dr and only DDD Dj+1. where is the pointwise multiplication operator, and [Djl [Dj+al + Aj. The condition Dj DJ+1= - =Dj+1 means that the non-masked dimensions in Dj+1 are also non-masked dimensions in D, The condition [Dj| = ID + A, means that order to go from D, to Dj+1. one needs to mask exactly A, dimensions.
Therefore, sequence {4, ر=D NM such that Σn- A N can be interpreted sequence of drops in dimension. At the core level, we have that |Dwl = N - |DM| 1, we have that at the end of the process, only dimension is considered.
ンドー M-1 A For instance if
We choose such that EDICo ... EDWCD oeDo....., DM) = 1. Therefore,
DM) corresponds to the probability to choose the dimension path Do + -...- Dm-
The part X r (b(X,), ((X) + Ar, fXX is more classical and corresponds to the continuous part of the diffusion process. We refer to (44) for a thorough introduction on infinitesimal generators.
For simplicity, we omit the schedule coefficients in (5).
A.1.2 Justification of the form of the infinitesimal generator
For any dimension path P = Do + ... + D (recall that Do = (1)A)) we define the jump kernel
JP as follows. For any E X, we have JX (x. = Ei (dim ( (dY). This operator corresponds to the deletion operator introduced in Section 3.1 Hence, for any dimension path P = Do .... ................................ + Das we can define the associated infinitesimal generator: for any f e (2) and X E X we define
A (f) = EM ((bla)) vx.((X)) + - (())88 -XX)
We can define the following jump kernel
J= == LDEoooo -ow 1eeo0-n------...... DM)JJ.
This corresponds to averaging the jump kernel over the different possible dimension paths. We have that for any / E (x) and X X
A()(X) Ei Va, + [ -
In other words, A EDECCDD ... ... -
DM)4
In what follows, we assume that there exists Markov process (X)) with infinitesimal generator A. In order to sample from (X.)ı20. one choice is to first sample the dimension path P according to the probability a. Second sample from the Markov process associated with the infinitesimal generator A². We can approximately sample from this process using the Lie-Trotter- Kato formula 44 Corollary 6.7, p.33].
Denote (Pt)ı20 the semigroup associated with AP, (Q1)ıo the semigroup associated with the continuous part of AP and (Jz) the semigroup associated with the jump part of AP, More precisely, we have that, (Q1)20 associated with Acont such that for any f C2 and X e
Acont (x) = EN (xx))), VI, + ((,,,(()))
In addition, we have that, (Q)) is associated with APump such that for any f E C2 and X x
A Aump (f(x) = {x(f(M) f(X)))
15

---

First, note that Acont corresponds to the infinitesimal generator of classical diffusion on the components which are set to sc. Hence, we can approximately sample from (Qr) by sampling according to the Fuler-Maruyama discretization of the associated diffusion, i.e. by setting
X, = Xu + tb(Xu) + Vtz,
where Z is a Gaussian random variable.
Similarly, order to sample from (Jz) one should sample from the jump process defined as follows. On the interval (0, we have X, Xg. At time т, define X; ~ (Xxo:) and repeat the procedure. In this case 7 is defined as an exponential random variable with parameter 1. For t > O small enough the probability that > т is of order t. Therefore, kernel, with probability t. Combining this approximation and 7 we get approximate samplers for (Q.)12 and (Jo) Under mild assumptions, the Lie-Trotter-Kato formula ensures that for any t 120 0
sample from i.e. the deletion
Pr
lim
12-+000
This justifies sampling according to Algorithm T the case of the forward process).
A.2 Proof of Proposition 1
For the proof of Proposition 1 we first provide rigorous proof using the notation introduced in A. We then follow this with second proof that aims to be more intuitive using the notation used in the main paper.
A.2.1 Time-reversal for the transdimensional infinitesimal generator and Proof of Proposition
We are now going to derive the formula for the time-reversal of the transdimensional infinitesimal generator A, (5). This corresponds to rigorous proof of Proposition [] We refer to Section A.222 for a more intuitive, albeit less-rigorous, proof. We start by introducing the kernel IKP given for any dimension path Do
DM. for any i {0,................. M 13. Y Di+1 and X by
KP (P,YY)) A) = - So (dim(Y)) JAnD,
M-1
(( No, , (d d dXDAAD.I Pe( Pdd
Note that this kernel is the same as the one considered in Proposition I It is well-defined under the following assumption.
Assumption 1. For any L > 0 and D 10 we have that X, conditioned to dimm..) D admits a density w.rt. the D Djd-dimensional Lebesgue measure, denoted pe (-Imim (X,) D).
The following result will be key to establish the time-reversal formula.
Lemma 1. Assume AT Let B,B& E X. Let be a dimension path Do > Da with M e N. Then, we have
E(LA(X,) J (x,, B)] = E|I,(X,) (X,,,
Proof. Let A. We have
E 1 (X,)) (X,,, B)] = E Bet---ElAAa (X)
= -EM- M-1 JAND. Pe (XD, |dim Xx) = P.) (X,) = ( JdXp.
= -EM- JAND, pr (Xo,,/dimmX (dim/(,)))) X.) = Pd) = Da)1a(XD...) JdX d.D.A....
= -M- SenD... 1s(XD...)
(AADD lA(Do (p (dim (.).(dim(X.) D.
EM-1 SenD... 18(N..t))
KP
= (ddi)(() = D A)g
EM' B(SD..,
which concludes the proof.
16

---

Lemma 1 shows that IP verifies the fixx equation associated with JP. The flux equation is the discrete state-space equivalent of the classical time-reversal formula for continuous state-space. We refer to 45 for a rigorous treatment of time-reversal with jumps under entropic conditions.
We also going to consider the following assumption which ensures that the integration by part formula is valid.
Assumption 2. For any > 0 and E NL....,N,, Xt admits a smooth density w.r.t the Nd- dimensional Lebesgue measure denoted Pt and we have that for any f.h E (B((R))))) for any E [0, and i E (1.....N}
Eldga ([[gg((X.))
Va, = (X. (V,, logPo (x,)()
The second assumption ensures that we can apply the backward Kolmogorov evolution equation. Assumption 3. For any 9 (x) and & >> 0, we have that for any u 0. I and X X, Oug(u, X) = A(g)(u, X 0, where for any u E [0, t] and X E X, g(u, X) = (Xg(X,) X, = X].
We refer 19 for conditions under A2 and AB are valid in the setting of diffusion processes.
Proposition 4. Assume Al Anadd A3 Assume that there exists Markov process (2) solution
of the martingale problem associated with (6) Let >> O and consider (Yt) = (XT-
Then (E))) is solution to the martingale problem associated with R. where for any f E (X). t E () and x E X we have
R(N((, = E (- Va, (X), Va, ((X)) + [X, dY).
Proof. Let f.g E (2)).. In what follows, we show that for any s, E [0, with 2
B((f(Y.) - ((Y)))((())))) .) = Big(Y,)) S R(f)(u, (Y) Y.)du]. More precisely, we show that for any s, E [0, 7 with t 2 s
B[(S(X.) ((x))))(X)))
S Ru))u,, X.)du).
Let s,t E [0,1], with t > s. Next, we denote for any E [0,t] and X E X. g(u (X,) X, X]. Using A3, we have that for any u E [0, and X E X, Ong(u, + (ug) X) = 0, i.. g satisfies the backward Kolmogorov equation. For any u E [0, and X EX. we have
A(Sg)(u, X) Oagln, (MXX) D4 X3)((x))
V X))
+ + +
(x)) +JX,ff
=.ggu
+ x))(())+ -40
Σ (1 (X,)
EN 8gdd (X.) (X),, (x),,,,,
V,,g(u,
M . ((K(X)) (x)) + (A.,f( X)ox
EN dgdd (X,) (V. ((X),, (u, + J(X, - ((X,9)
(8)
Using 2 we have that for any u E [0,t] f and i E (I....,NN
Eldra
V Vu, x.))
= -Bppp (X))))( X. (Ae, (Xx) + (V., logp. Ve....)))) (9)
In addition, we have that for any X E X and u 10,t] I(X, J(X,9)/(X) = -
()((() - ((X))XXXXXXXXX dY). Using Lemma we get
[XJ(X,, fg) J(X,, fg X..) = EEg((, X. (X,,, f)].
(10)
17

---

Therefore, using 8). 9 and
TTP. we have
E [lA(fg) x.)l
X.)).
Finally, we have
[((((x)) - f(X.))ggXXX Big(t, X,) (()) Xx,)
= (lf
= -EI,, (()) Xu) ).d -Xi)(X:) 1/ R())u, X,)dul,
which concludes the proof.
L
A.2.2 Intuitive Proof of Proposition
!
We recall Proposition I
Proposition 1. The time reversal of a forward jump diffussin process given by drift b diffusion coefficient Jt. rate X (((n) and transition kernel Σx=1 Kdel (aax)) is given by a jump diffusion process with drift (i(X), diffusion coefficient Ti. rate (i( and transition kernel Sym 1+1 Ay (iss)sssX you, ydyado as defined below
b;(x)=BAA - Ji Logpt( (X), D = J.
xi(x) = Xe(n+ Eht Kudl (i|n n + 1) Syss Pe (ins yiu dy)dy
(X)
Ay i|X) x Pe (ins (X, yadd, i)) de + 1).
Diffusion part.
Using standard diffusion models arguments such as 6,(x) (x) g²? (X.n..
18
45
we get
Jump part. We use the flux equation from 45 which intuitively relates the probability flow going
in the forward direction with the probability flow going the backward direction with equality being achieved at the time reversal.
Pr (X) Xi (Y PY
(X
X
m+ A; (y:uld, z ( you 13 ddyali
Pr
= pr (Y)
I del (i]n + 1) 1) (X).
(11)
To find (X), we sum and integrate both sides over m and y, with Y (m,
Σm= I md Pt (X) ; fyom E+ (y : yx ddyd >M Sy md Pt (Y) X (Y) E Kill(iln + 1 1) (X)dy.
Now the fact that 6 1)) (X) is for any F n + 1,
pc (x()) (x))=((x= X,(n 1) 1) SyERMIN Pe (Y) Σ (i)n + 1) .3) X)y
= xx,(n 2 dedel (i/n + 1) Jye Pt (Y (Y) do i) (X)dy
Now letting Y = =ins(X, ydd yau,, ,)),
pi (x) (x)xx(( Xe(n + 1)
f=1 nl K diel 1) Syn Pe Pinis( X, yilld ddydddd
N((Xx== Xe(n 2013 d (i|n+1) Sysea pt X. Jay pc (X)
To find (y 1 {)), we start from II and set Y = ins ndd j) to get
pr (x( (x() (z add j|x) (Y (n 1) 1).
By inspection, we see immediately that
(z .j x) x pa (ins X, nuتس K (j|n2 + 1). With re-labeling of 2 and j achieve the desired form
Ai (y .i(X) pe(ins(X X, ydd i (iln + 1).
18

---

A.3 Proof of Proposition 2
In this section we prove Proposition using the notation from the main paper by following the framework of [7]. We operate on a state space
Un-ifm] x Rnd On this space the gra- dient operator V C(),R) C(X, X) is defined as v/(X) v(na) s(X) where (nd) is the standard gradient operator defined as C(RRdd R) C(RTM _Rdd) with respect to E Ind. We will write integration with respect to a probability measure defined on X' as an explicit sum over the number of components and integral over Ind with respect probability density de- fined on Rnd i.e. [X))x(((d Ex SEERmA )(())p(x(())))
where, for A RRdd We will write p(X) as shorthand for p(n)p(x|n).
J(m.A) M(dX)= =
Following, 17, we start by augmenting our space with time variable so that operators become time inhomegeneous on the extended space. We write this as X (X, t) where X lives in the extended space S = X' x I20 In the proof, we use the infinitesimal generators for the the forward and backward processes. An infinitesimal generator is defined
E p ( [f(Y) f(X)
A lim r-40
t
and can be understood as probabilistic version of a derivative. For our process on the augmented space S, our generators decompose A A, where A operates only the spatial components of X i.e. X 17).
We now define the spatial infinitesimal generators for our forward and backward process. We will change our treatment of the time variable compared to the main text. Both our forward and backward processes will run from
0 to / = T. with the true time reversal of X following the forward process satisfying (Y)) = (XT-r)ie(o.n) Further, we will write g a as ge and T = -T- as we do not learn y and this is the optimal relation from the time reversal. We define
(x))(
Em SYERT.. (XIx))
well
-x(()) where A = (V-V) is the Laplace operator and 5 a dirac delta on X i.e. ΣN Iy SyEame Sx(Y)dy= 1 and Em= Syerms f(
(x))( = ((( TYX Sycrma
= f(X).
Verifying Assumption
The first step the proof is to verify Assumption in 17 Letting
(X) = (X), we assume we can write (xp:(X) = (X)) in the form Mv + cv = 0 for
some function e: S + R, where M is the generator of another auxiliary process on S and K; is the adjoint operator which satisfies (R; h) = (f. R,h) i.e. ΣN IXERnI N(X);; (f) Σ. fNERII (
(1)))x
We now find Kg. We start by substituting the form for Kr,
Ei f ((X) - SXERnA f(X) (X)(((((()))))) Vh)(X) (B
Jy ha..hhhnn
5x
We first focus on the RHS terms corresponding to the diffusion part of the process
EN-I IxERRnd (x)((e vn) + (
-E IXERR f(X) b: Vh)(X) + +07- vh(X)dx
-NN IxcESen f fxx) 6 +
>N IxcESSen - (X)V (fb,)(()++ MX (X(X)
E IXCEREN h(X(-v (000)/(VV(()))()))-0 (fb)) - Vf
= N
RE EXEER ( 6((x) VX() (((()) + 1.6 V Vf(X)
19

---

where we apply integration by parts twice to arrive at the third line and once to arrive at the fourth line. We now focus on the RHS term corresponding to the jump part of the process
En SXERRO ((x) (() Em SyERma (Y)
5x(Y))dy)dk dx
I hi h(Y) ( SXER 1(x) (x)) - (x(xx)xx
E h( (E N
SyErma (V) (x
where the last line we have relabelled to Y and Y to X. Putting both re-arranged forms for the RHS together, we obtain
En ExERna (())x=)
En SXER -XX(- (((X) - V/(X) (X( + V V -
V-V(X)+
Em SyERma ((Y) -X) 6x(X))dy)dx.
We therefore have
(x())()) = --XX)V D1(x) VS()) + + As(x)
Σm SyeRma ((Y) (X(O)( - (x.Y)....
Now we re-write (x), = KC; (x() in the form Mr + cv = We start by re-arranging
(x). = (X)((X) = 0 -= (X) (X))
Substituting in our form for Kū we obtain
0 0= 1 (X)V (x_(()) (x), (X) + Ax)
+ I LyERma
(xIx)) 5x(dd)
(12)
We define our auxiliary process to have generator M = 8, + Me with
M.()(x) = bb (X)
EN Sye (KY) (Y[X) 6x (Y))dy
which is jump diffusion process with drift b.M, diffusion coefficient It, rate X. and transition kernel K.M. Then we have b.M = -b
XM (X) Em Iailyye Syerma X X
(13)
and
(Y|)) & (
(14)
Then we have (2) can be rewritten as
0 (x) (X)V 5 5x,((X) + bi Vu + 300
m (X) + (X EN I SyErma (x)
(x) Em Sy md 1/ (K) (YIX) (x))dd
which is in the form ( (X)u 0f let
x - (x) - 5/ (X) I Syerme (Y)
Verifying Assumption 2. Now that we have verified Assumption 1, the second step the proof
Assumption 2 from [17]. We assume there is a bounded measurable function S + (0, oo) such
20

---

that aMf = [(fa) - fC for all functions f : X + R such that f D(M) and fa D(C).
Substituting in M and we get
(x) Dq (X) Vf(x)
+ A + XM Em SyERma ((Y) - 5x(Y))dy)
(X) + (X) v(fa:)(X) + ((fo))
+Xx(X) Em=l Jye yERmd (Y)
- -x(Y))dy
-((X) (X) + B. (X) Var (X) 1 (X)
+ +XX) Em= SyERma ad (Y) (Y) (Y - 5x(Y))dy)
(15)
Since does not depend on time, (x) = 0 and 0 (faz) = f(X) (X) thus the O terms on the RHS also cancel out. Comparing terms on the LHS and RHS relating to the diffusion part of the process we obtain
(X)
Vf(X)) + (x(
(. (()a)((X) + - Var(X) - {X)) (X).
Therefore, we get
(x) -xx()) (6), V/(X)) + (xx(X)))A/(X-- = +
((()) (x)
298 (2V (X) ((X)sa + (X)A((((
((X). VX, / OE (X).
Simplifying the above expression, we get
-x,(X) V/(X)) = ox. (X) by (X) VX(X + V (X)
(-a((X) b (X)) (X( (x) by (X) Va, Vf(X).
This is true for any implying
-a (X) b (X) = a ( Vx,
This implies that ag (X) satisfies
Vlog oe (X) - 9i
(b,(x) (x)
(16)
Comparing terms from the LHS and RHS of (15
relating to the jump part of the process we obtain
an (X) (X) Em SyERmo (K, (Y 8x (Y))dy =
(x) Σn Di LyERmO J(Y)a, (K, (Y - 5x(Y))dy
((x) X(X)x,(X) SYERmO On (X (x()d)dy.
Hence, we have
(x) .M
Sygrma J(Y) (Y o (X)A
X,(X) EN SyERma f(Y) Y (Y|x)dy
-((X) x(XX EN
E SYERoE ( (Y[x)dy.
Recalling the definitions of XM (X) and (YIX),
13 and 14
we get
(X) EM SyERme f(X) (Y)
K Y
(x))))( MN SyERme
T T- (Y
& T-I (x|)dy
X.(X) EN LyERma ((Y)au (Y),
-s(x)x EN SyErme ar
21

---

This equality is satisfied if ar (x)) follows the following relation
To (X) for + m
(17)
X.(X)R,(IXX
We only require this relation to be satisfied for * m because both (.(Y|X) and K (X are O for = m.. We note at this point, as in 17 that if we have or ( = 1/po(X) (X) and X f-t
and Kg - (X equal to the true time-reversal then both (16), and 17 are satisfied. However,
or (X) = 1/p.(2 (X) is not the only to satisfy these equations. (6 and 17 can be thought of as enforcing certain parameterization of the generative process in terms of Ox2
Concluding the proof. Now for the final part of the proof, we substitute our value for into the ZISM loss from [17 which is equal to the negative of the evidence lower bound on (X.) [log P xx) up constant independent of 0. Defining 3x (X,), we have
(X+)
(X.) + É. log Br (X.))d.. TISM (3) = So (x,) (x,))
We split the spatial infintesimal generator of the forward process into the generator corresponding to the diffusion and the generator corresponding to the jump part, L L + &t with
Ct (f)(x) B.( VX(X) + {{())
and
((()(X) = E f(Y) (()) 8x(Y))dy. By comparison with the approach to find the adjoint R2 we also have Li = Lequtt + C.a with
Ca (f)()) -f())v B.(x) vx(X) (x( + 2 f(xx).
In addition, we get
()) = - S (X)) -x {x(Y))
Finally, IISM becomes
TIsM (B) -000 (x.)
24 8, (x.)
(x(,)) Et + 2 log (xx))dtt L log Br (X,)) JJ"d Ep,(x.) B((X.)
((x))
= I (B) + (3),
where we have named the two terms corresponding to the diffusion and jump part of the process as ISM IISM respectively. For the diffusion part of the loss, use the denoising form of the objective proven in Appendix E of [17] which is equivalent to IiSM up constant independent of 0
ISM (B) = -1º" (x.x(x)
(priol (o)))))))) _diff (at (X. Xo
-c.
+ const.
To simplify this expression, we first re-arrange (ai (h) for some general function h + R.
Liift (h) Ld (log hi) V.VA
1 AA b, V log h - 199 log
h
2 V. h)
9 log htl2..
Setting Pt (-))a(-), our diffusion part of the loss becomes
7 (B) = 1..00 L X
log (X + Vlogar(X))
+ const
We then directly parameterize V log or X,, - (X.,)
lo Zriff (B) = 1 ( 1
(P(( const.
22

---

We now focus on the expectation within the integral to re-write it in an easy to calculate form
E (X.. V/ log Pajo (X, 2y (x =
||| (X)) 1 2s (X,) )" V log Po. P (X., X,)] + const
Now we note that we can re-write V log Po. (Xo. X,) using Mr where M, a mask variable Me (0, 10 that is 0 for components of Xo that have been deleted to get to X, and for components that remain in X
v log p Xo, X1)
(Xo, X+)
X.,
Σ Mi Po.e (Xu, X,, M,)
XX,
-MM, po,e(Xn,X, X,)
Vpo,c Me)
Me po (x)
p(ne.Me Xx)
V VP mns, Xo. Mc)
Ems, pix, mi M,)
piM, X x
Px Xu, M,)
V Prjo (x.,).X |m,, Xo, M,)
Substituting this back in get
(Xo..... V log Pejo Xo) - 56 (x:)11
= -Epo Ep,, (X.. [| (X, 12 25 ( Ept M. Xx, x.) V log Xo, M,)l + const
1X.. M.) V log Pr(o (x, Xo. M.) of (x.) + const.
Therefore, the diffusion part of LIsM can written as
ISM (3) T (et(,,,T))p X..X.. Lo V log MM, Xo, Mi) - (X,)11 + const.
We now focus on the jump part of the loss TisM We first substitute in Lt and Ct-
IISM So" Ep,(x,) SyERRm X,YY)
Y (.
(X, YY - dx+)))
xx(X,) EA SyERms K (Y(X,,) log (YYYdyy X.(X,) log 3r (X.))dI.
(18)
Noting that (,(X,) = (/ar(Xi)) we get
5 XX for n, + m
(19)
B(Y) X.(II)
or swapping labels for X, and Y.
(2. for n, # m
(20)
(X.)
(x
Substituting
19 into the second line and
20 into the first line of 18 and using the fact that
K. (X: Y) = 0 for n4 = m, we obtain
It -00 (x,) )l Emm fyERmd X Te (X)R R (X.
(X
EmE SyEkmd X
BYY
(X;)dy
((()
+ X((X,) EME SyErmS amd K Y X. (1 - log T (Y)
-log Ko (X,lY) + log XX,) + log R, (Y),)))
X.(X,) log (X.))dL
Hence, we have
I/SM .o¹ Ep((X.) (x))(()) Em=nim SyERms Kg (Y|x) Xx (x,)))
3((),)
Xx(X,) EN
Emm= Syerme R.(Y[x.)(( TQ_ -log Kg (d)xydd const.
23

---

This can be rewritten
(X,) IssM So (x.) 1 (X,) + (X, R (VXx)) 1- log x4_1 - llog K 7 (X, Y) const.
Therefore, have
Issm =T (x. x.) 1 (X,) (X,, log x (Y) (x(,, log (X, const.
Finally, using the definition of the forward and backward kernels, i.e.
,.(Y(,)) =
E (4))))00a))xx))))) and Kg xX,(() Iem Ei= (x inss) x"., (X, )dx we get
TIsM- (xutt,,T' Im=1 SYER E Kddll (i deel (Y) (T) Xx(X,)log (x)
-X((X, log K x (X, + const
We get
IIMM
(x)) (Xx) log 5 (Y) (xx,))x) log (X, Y). + const.
Therefore, we have
TISM (x
(X,, (X, log (
log 12- ,YY + const.
Putting expressions for Z and IisM together obtain
TISM E Poj n,, Xo. M1) ot [1]]
TE, X (X,) log 5 (Y) XX,,, log AG 1 Y) + const.
We get that - -IIM gives us our evidence lower bound on Pe... log po((() up to a constant that does not depend on 6. In the main text have used time notation such that the backward process runs backwards from f = T to t = 0. To align with the notation of time used in the main text we change 7 t to on subscripts for 77 and AG_ We also will use the fact that (X,),) depends only on the number of components X,, (X, = (x().).
[() F || (X.) Vx, log Pujo (x, |n,, Xo, M))
TE| xI (nx) log X (Y) X.(re) log A (x=*M, i Y) const.
Tightness of the lower bound Now that we have derived the ELBO as in Proposition we show that the maximizers of the ELBO are tight, i.c that they close the variational gap. We do this by proving the general ELBO presented in [17 has this property and therefore ours, which is special case of this general ELBO, also has that the optimum parameters close the variational gap.
To state our proposition, we recall the setting of [17]. The forward noising process is denoted (Y)) and associated with an infinitesimal generator L its extension (1, Y.)x20 is associated with the infinitesimal generator i.c. C 0 C. We also define the score-matching operator given for any f for which it is defined by
( C())/ - L(log())).
We recall that according
17
Equation (8)) and under 117 Assumption Assumption2), we have
log ( (Y,) Z Ellog (Yx)) Sh- ((v/B)/(v/3)) + C(log 3)dt),
with 02 Pr- for any t E J0. T). We define the variational gap Gap follows
Gap E(log Pr (Yu) log - 1-" L(v/B) + C(log 3)dt).
In addition, using Itô Formula, we have that log vy (Yr) - log up Yo) = - L(v)dt. Assuming that ( Yo)l] < +00, we get
El| log vr( Yr) log no
Gap = B - Clog v) + (l//B)/v//)) + f(lg 3) = L (c/3)dt).
In particular, using [17] Proposition 1]. we get that Gap 2 0 and Gap 0 if and only if B x In addition, the ELBO maximized if and only if 3 x v. see [17 Equation 10] and the remark that follows. Therefore, we have that: if we maximize the ELBO then the ELBO is tight. Combining this with the fact that the BL maximized the time reversal 17.. then we have that when our jump diffusion parameters match the time reversal, variational gap is O.

---

Other approaches. Another way to derive the ELBO is to follow the steps of [1]] directly, since 17 general framework extending this approach. The key formulae derive the result and the ELBO is 1) Teynman-Kac formula 2) a Girsanov formula. In the case of jump diffusions (with jump in TRd)= a Girsanov formula has been established by [22]. Extending this result to one-point compactification space would allow us to prove directly Proposition without having to rely on the general framework of [17]
A.4 Proof of Proposition
We start by recalling the form for the time reversal given in Proposition
!
x (n + E K 1) fyre Pr (ins (X, yoll, i) dy) /ld/ /pt (X).
We then introduce marginalization over Xo
;((x) Xi(n -
1) 1
(dde 1) sya Eno Sxx Po.e (Xo, ix, yuthd dx) /pe (X)
Xe(n - 1> - (il n + 1) Iym Eno Jxn
(x()a)
Pejo (ins X, you ()) ) cxxddyau
Pe(X)
-xin+ 1) .+ K (i|n + 1) yom Eno Sx0
Pija (x) Pt 1o (ins X yddd e) X dxx
Poj X
= =Xe(n + 1) m+ Kedel (i/m + 1) I Iyo Eno Jxa
Pojs Poy (xo X.no)
Prio X. Pr X..n)
x
PHo(n + ((oo) (x(,, yolld
XX: n + 1)dxndy=ii
where (n + 1,,1, (X, yuidd 1)) = ins( ix, y process does not depend on X6. only no. have Pr (n pnio(1(no) and Prjo(n + 1(X,) Pajo(n + 1/no). Using this result, we get
mdd 1). Now using the fact the forward component deletion
Ti(x) = =xe(n + +) (poo(+1imm))
Pejo((lna) Poje (no X)x
Entl (1]m++1 fyst (X.. 1j. Poje X, do) Sx1
(21)
We now focus on the probability ratio within the integral over Xg- We will show that this ratio 1. We start with the numerator, introducing marginalization over possible mask variables between Xo and (n + 1, z). denoted (n++)) with M M(+1) having n + 1 ones and no (n + 1) zeros.
E d (i/n + 1) Sym Prjo (x(
y ydd .)) n 1)dy
EMA (il 1) E M(++1 Sysu Pejo PMjo(M(n+1 (n+1) z(X, yiild, i) 10, n2 1)dy
EMOHH) EnS K (iin + ( Xo, 1| 1) Sysu Piio (zl,, i) i)lXo, 12 + 1, (n-11) Jdy
Now, for our forward process we have
(x) X)) n1 + M(++1)) = I1 (
aa)
where z is shorthand for z(X. add i), z(i is the vector in Rd for the jth component of z and M(n-1) 1) (X,)² is the vector in Rd corresponding to the component in X. corresponding the jth one in the M (n+1) mask. Integrating out yould we have
Iou (z) ydt i) Xo, n21, M((+
I[,,, (a.M(1 (X.), (1~ar)
where (n+1]J denotes a mask variable obtained by setting the ith one of M(++1) to zero. Hence, have
ΣT1 Kidel + 1) fyse (XX, youht X0, n + 1)dy
= EMinHD Ent1 dell (iln + (M 7 +
I Va (X)) (1- (1))))la).
(22)
We now re-write the denominator from 21 introducing marginalization over mask variables, (n)
Pxj(x||Xo,1) n) = E M0o (M Xo. n)) (x)])((((MMM) ( n)).
(23)
25

---

We use the following recursion for the probabilities assigned to mask variables
Pejo(M Xo n) Σ Minon EE I[M (V =
1))roo [xo n 1).
Substituting this into (23 gives
Pa n2) E Min) E EM
n+ I(M(++ 1 1)x
ii=
Prjo (n
Xo, n + 1 1) () Xo. ni)
Mi) Emantt
1 1) i=l I(M = M (m)) J. (ilm
m+
(M (x+1)
Xo. + 11 -V varM ,1 de)
Minnn E K (i/m + 11)::)(MMMMm Xo. 1)
- (M (1-a))la).
By comparing with we can see that
Xo n) E K (iin + 1) Jyz Prio Prjo(z((X, (z(X, yold, Xo: n + 1 1)dy
This shows that the probability ratio in 21 is 1. Therefore, we have
xi(x) Xe + 1) En, no Paja(n(no)
Pro[n+I|ng Pojs (ooo( (X) Sx0 Pojo no noxxx
pro(x+1|ma) x + 1) ma Pejo(m|no) Poja (no (X),
which concludes the proof.
(ej((n(mo) can be analytically calculated when xm) is simple enough form. When (n) does not depend on then the dimension deletion process simply becomes time inhomogeneous Poisson process. Therefore, we would have
Ptjo(1l = Si X _ds (mo-n)) exp(- exp(- So xds))
In our experiments we set Xe(n = 1) to stop the dimension deletion process when we reach a single component. If we have Xx(n) = X X²(m) for all n, m 1 then we can still use the time inhomogeneous Poisson process formula for > and find the probability for n (n 1mo) by requiring (ro((a(mo) to be a valid normalized distribution. Therefore, for the case
that (m) (m) for all n, m and Xe(n 1) O, we have
(La X dedo expl- fo X sss)
< n S no
(Mo-n)]
(G X _o m
m= (Ro m]] exp(- f X xss)) n=1
In cases where (n) depends on n not just for n = 1, (n)((()n)) can become more difficult to calculate analytically. However, since the probability distributions are all I-dimensional over it is very cheap to simply simulate the forward dimension deletion process many times and empirically estimate (1(o)) although we do not need to do this for our experiments.
A.5 The Objective is Maximized at the Time Reversal
In this section, we analyze the objective [(0) as a standalone object and determine the optimum values for =g 7 and A& directly. This is in order to gain intuition directly into the learning signal of C(0) without needing to refer to stochastic process theory.
The definition of (()) as in the main text is
(()) - - E || (X.) Vx; log (x, Xo, ne, 1))+
TE X + X (mx) log X (Y) - Xx(me) log Al [|Y)] + C.
with the expectations taken over 0, O, Po) (Xo, X,,
(Y).
26

---

Continuous optimum. We start by analysing the objective for sf. This part of [()) can be written as
2 Jo of B X. M,) [II
(X.
Vx, Pog Xo, ne, Md)
We now the fact that the function that minimizes L2 regression problem min Bp(ra)) [l)
y||a| is the conditional expectation of the target f (x) En(sl=]bul- Therefore the optimum value for sf (Xe) is
$f (X+) = Xo/x.) Vx, log (ejo Xo, ne. Mc]]
-EMe Eng ERO PMM, no: X xo (X, P:jox (xo no. ne. M2) dxx
EM, E EROER Pajo(xe n..
ao. (X+)
V Vx, (x Xu, no. ne, M, )dxo
= EM, Eng=1 IXOERTON
V Vx, Xx1 Xo, no. ne, dxo
p(nex.))
p(m(,x)) EM. Emo SxpERma V Vx.P(Xx Xo, no: ne, M2 )xx
p(n(xx)) Vx, Em.. Eno Exo p(x,, Xo. no. Re, Mi dxn
1
p(n)xx)) Vx, Xx,,(x;,,)) nz) Vx, (X,,)
Therefore, the optimum value for s (X.) is Xx, log)(X;) which is the value that gives b to be the time reversal of b as stated Proposition
The learning signal for X comes from these two terms in C())
Jump rate optimum.
7 + X (ms) log Xi(())
(24)
This expectation is maximized when for each test input Z and test time t, we have the following expression maximized
-Pe (Z)Z)) + Ei= Syw (ins Zz, yalld, 1))
Kdel (ilm, + 1)dy x Xe(n, + 1) log (Z))
because (Z) is the probability Z gets drawn as full sample from the forward process and k)) Σe=+ fym Pt yins (i)) bigger than Z gets drawn from the forward process and then component is deleted to get Z. Therefore the first probability is the probability that test input and test time / appear as the first term in (24) whereas the second probability is the probability that test input Z and test time appear as the second term in (24).
Idy is the probability that sample component
We now the fact that, for constants and
argmax, - ba + clog a = 6.
We therefore have the optimum (z( as
)i(z) = X, +1) 1) E.= Sys Priins( auy a)) de (1),, I)dy
(2)
which is the form for the time-reversal given in Proposition
Jump kernel optimum. Finally, analyse the part of L()) for learning Ag (x ; Y)
TE X (me) log Al (xx2") Y|
0 ((. XX (mee)) log Af Y)
- Ep((me) X ( Ep. (n.) Jou (i)ms Baacc (V) [log Al (xx" : Y)]ddt
We now re-write the joint probability distribution that the inner expectation is taken with respect to,
(x(x:|m))) (ilre) (Y) p( (e)p(( ; Y (x) (xx)
27

---

with
(()m))) = Eis Sxx Px(xx(n))) ( (i) ddll(x..)) (dxx,
and
p(x . Y) & pc (x:lmm))
and x² xasse
E: [R(n==14 1)d referring to the ne components of xs, that are not x3th i.c. X, =
ins
(xx
n4 x i). We then have
TE X (mm) log A (
-f² Bp.(ne) [ (C 18 (xt [log A (x: il d|]]]]
-So Epi(nn) [ (p (8)) log Af (xcptle d YY]]]d
-f" Bp.(me) X (Tee) Eayy m
[lg (x:dddd
Y)] + const
= = Bp.( (M))) ((y (xt=)) KL (px Y) A3 add
Y)]ldt + const.
Therefore, the optimum Al (x:"dd Y) which maximizes this part of (()) is
(x Y p(x;",, Y) ; x (x( (ilm)).
which is the same form as given in Proposition
B Training Objective
We estimate our objective L()) by taking minibatches from the expectation U(t;0, T) and then
U(I:0, (X,,, X,, (1im)) (Y). We first sample take samples from dataset Xu Pdaul Poo)) In order to sample Pejo M,, (X)) we need to both add noise, delete dimensions and sample a mask variable. Since the Gaussian noising process is isotropic, we add suitable amount of noise all dimensions of X, and then delete dimensions of that noised full dimensional value. More specifically, first sample X, (no, x,)
with x, (Xxr Vo, (1 mgg) for = exp (- 3(s)ds) using the analytic forward equations for the VP-SDE derived 3. Then we sample the number of dimensions to delete. This is simple to do when rate function is independent of n except for the case when which it is zero. We simply sample a Poisson random variable with mean parameter J X,ds and then clamp its value such that the maximum number of possible components that are deleted is no - 1. This gives the appropriate distribution over n. (tjo(((ng))
at
given Section A.4 To sample which dimensions are deleted, we can sample K Ki (mo) K K ( (i2|ni Ino 1)... K'del (ng= Im + 1) from which we can create the mask M, and apply it to X, to obtain X,, X, = (X,)).. When (de((ijn) 1/n this is especially simple do by simply randomly permuting the components of X,, and then removing the final mg n, components.
As is typically done in standard diffusion models, we parameterize sf in terms of noise prediction network that predicts e where x; = (x,M)(()) + /I- ate, e ~ NO, In.a). We then re-weight the score loss in time such that we have uniform weighting in time rather than the 'likelihood weighting' with o² 3 21.. Our objective to learn sf then becomes
(oo) (e,))(cs) (X,) - -
with X, (x, + arc (X,)
-1 vT (X,).
Further, by using the parameterization given in Proposition 3 we can directly supervise the value of Pa (no (n) by adding an extra term to our objective. We can treat the learning of P (ng X,) as standard prediction task where we aim to predict ny given access to X,, A standard objective for learning P/j, (no)X,) is then the cross entropy
(Xo. x,) Pog xxi)
max
8
28

---

Our augmented objective then becomes
[()) = -e
log X((+))
((7))+X((n))) (e(ms) ) log (x +) log poje xx)]
(25)
where the expectation is taken with respect to
(Xo)
nu ( (E; In,d)) (2
where X, (xx) + vI - age and y is loss weighting term for the cross entropy loss.
C Trans-Dimensional Diffusion Guidance
To guide an unconditionally trained model such that generates datapoints consistent with conditioning information,
the reconstruction guided sampling approach introduced in [ Our conditioning information will be the values for some of the components of Xn, 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 x Re.d and the components to be generated x IRaad Our
trained score function s (X,) approximates Vx, log Pt (xX)) whereas we would like the score to approximate Vx, log Ps Xx, xx(x)) In order do this, will need augment our unconditional score sf (x)) such that incorporates the conditioning information.
We first focus on the dimensions of the score vector corresponding to x². These can be calculated analytically from the forward process
Vxp logp(X, Xx; log (e] n4)
with Pejo (xf x6; n²) = (xf; Vaxxi (1 (x)) an Note that we assume a correspondence between xi and x8- For example, in video if condition the first and last frame, we assume that the first and last frame of the current noisy x, correspond to x and guide them towards their observed values. For molecules, the point cloud is permutation invariant and so we can simply assume the first n. components of x, correspond to xi and guide them to their observed values.
Now we analyse the dimensions of the score vector corresponding to x". We split the score as
Vx logg(X((xx)) = Vx; log p(x) X,) + V Vx log px(;))
p(x X2) is intractable to calculate directly and following
9
approximate it with
N aaa (Xt), I-ox Inga) where xX)(Xx) is point estimate of xo given from sf (X.) calculated
as
(x, x + (1 - a.)ss
Var
where again we have assumed a correspondence between x and xi. Our approximation for V xe log p(x]] Xx) is then
V xt log 1 x 2(1
at
|x
* (x))11
x
a1))
which can be calculated by differentiating through the score network sQ.
We approximate " (Xx(xx) and (y i|X, x)), with their unconditional forms X (X.) and Al (yadd, ,iXx),) We find this approximation still leads to valid generations because the guidance of the score network s/, results in X, containing the conditioning information which in turn leads to (X,) guiding the number of components in X, be consistent with the conditioning information too as verified in our experiments. Further, any errors in the approximation for A (ydddd i i|X,) are fixed by further applications of the guided score function, highlighting the benefits of our combined autoregressive and diffusion based approach.
29

---

D
Experiment Details
Our code is available at attpppe//gaaabooooooccaawwecceep
D.1 Molecules
D.1.1 Network Architecture
Backhone 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 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 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 final output projection layer to give the score network output sf (X,).
Component number prediction To obtain Poje (no) Xx), we take the embedding produced by the EGNN before the final output embedding layer and pass it through 8 transformer layers each consisting of 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 final projection layer to give softmax logits over the poj (no (x) distribution.
Autoregressive Distribution Our A i[X,) network has to predict the position and features for a new atom when is added to the molecule. Since the point cloud permutation invariant, we do not need to predict i and so we just need to parameterize A (y.) 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 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 (no for (no X.). We also input the distances from the nearest atom to all other atoms in the molecule
X.).
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 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, (ydid X.) outputs both mean and a standard deviation for 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 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 [6] with decay parameter of 0.9999. We train on the 100K molecules contained in the QM9 training split. We model hydrogens explicitly. Training model requires approximately 7 days on a single GPU which was done on an Academic cluster.
In 8 the atom type is encoded as one-hot vector and diffused as 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 guided performance when conditioning on certain atom-types being
30

---

present. We found our model is robust to this change and achieves similar sample quality to g as shown in Table 2
When deleting dimensions, we first shuffle the ordering of the nodes and then delete the final nnp T4 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 O. 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 We then add CoM free noise such that the noisy molecule also has CoM= 0. Our
that its CoM score model s is parameterized through noise prediction model ef which is trained to predict the CoM free noise that was added. Therefore, our score network learns suitable directions to maintain the process the CoM= 0 subspace. For the position prediction from A (y (x,) 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 CoM= 0 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 pc(X) as its stationary distribution rather than the process progressing toward po(X), (X). We use the same method to determine the corrector step size Ç 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-dimensiona. modeling by first simulating possible insertion using X and then simulating possible deletion using X t- We describe overall sampling algorithm in Algorithm
Algorithm 2: Sampling the Generative Process with Corrector Steps
Input: Number of corrector steps C
X Pret X = I(n (> Ia)
while t > 0 do
if u < (x))) with u ~ U(0, 1) then
Sample xald i in A (x ,ix)
X + (xs(Xxx"dd) .1)
end
x
+- x - (((x))1+ + de olo
ble with e (O, Ina)
for
= 1,.... C7 do
x x (5)
+ 2ÇE with e (O, Ind)
if u < ( (X)8t with 41(0 1) then Sample add (x x
X + ins(X, x .))
end
if u
< X (--se(n)8 with u U(0, 1) then
X + del(X, with deee
end
end
X < (n, ( t-
end
D.1.4 Evaluation
Unconditional For our unconditional sampling evaluation, start adding corrector steps when
I < 0.17 in the backward process and use corrector steps without the corrector steps on the number of components. We set 8 0.05 for t > 0.5l and 0.001 for t 0.5T such that the total number
31

---

Dataset
0.150
Samples
0.125
Density
0.100
0.075
0.050
0.025
0.000
5
10
15
20
25
30
Number of atoms
Figure 6: Distribution of the size of molecules in the QM9 dataset 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 Figure show more unconditional samples Figure 7 We find our model consistently generates realistic molecules and achieves size distribution similar to the training dataset even though this is not explicitly trained and arises from sampling our backward rate X Since we are numerically integrating a continuous time process and approximating the true time reversal rate X; some approximation error is expected. For this experiment, sampling all 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 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 given in T Molecules in the training dataset that meet the conditions in each task have a different distribution of sizes. The tasks were chosen that have 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, use 3 corrector steps throughout the backward process with ö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 Figure X 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 S and the size distribution for molecules in the training dataset that match the conditions of the task (green no hashing in Figure B). We find that indeed our model when guided by diffusion guidance can automatically produce a size distribution close 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 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 { who train new model conditioned on the polarizability of molecules the dataset. We train a conditional version of our model which can be achieved by simply adding in the polarizability as additional feature input to our backbone network and re-using all the same hyperparameters. We show more examples of interpolations in Figure 10
32

---

Figure 7: Unconditional samples from our model.
D.1.5 Ablations
For our main model, we set X X = to ensure that all dimensions are added with enough generation time remaining for the diffusion process finalize all state values. To verify this setting, we compare its performance with X Xico.037
and X X O. We show our results in Table setting generate reasonable sample quality but incur some extra dimension error due to the generative process sometimes observing a lack of dimensions near and adding too many dimensions. We observed the same effect in the paper for when setting X, to be constant for all
6 We find that the X X:c0.037
Table Further, the setting X X 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 X X to be a reasonable trade-off between these effects.
D.1.6 Uniqueness and Novelty Metrics
We here investigate sample diversity and novelty of our unconditional generative models. We measure uniqueness by computing the chemical graph corresponding to each generated sample and measure what proportion of the 10000 produced samples have unique chemical graph amongst this set of 10000 as is done in [8) We show our results in Table and find our TDDM method to have slightly lower levels of uniqueness when compared to the fixed dimension diffusion model baseline.
33

---

Table The 10 conditioning tasks used for evaluation. The number of each atom type required for the task is given in columns -5 whilst the average number of atoms in molecules that meet this condition the training dataset is given in the 6th column.
Mean Number of Atoms
Task
Carbon
Nitrogen
Oxygen
Fluorine
11.9
2
3
1
13.0
3
2
1
13.9
4
6
O
1
14.6
5
3
O
16.0
6
6
3
O
0
17.2
7
I
0
17.7
I
8
7
0
19.1
I
0
19.9
10
8
O
0
21.0
1.0
0.3
0.8
0.8-
0.6
0.6
0.3
0.4
0.1
0.2
0.1
Densty
0.0
0.0
20
30
10
20
30
20
30
10
20
30
10
20
0.4
0.3
0.3
0.3
0.3-
0.2
0.2
0.2
0.1
0.0
10
1.00
10
30
0.0
20
Number of atoms
Figure 8: Distribution of molecule sizes for each conditioning task. Tasks to right in the top row and tasks 6 - 10 are shown left to right in the bottom row. We show the unconditional size distribution from the dataset in blue vertical/horizontal hashing, the size distribution of our conditionally generated samples in orange diagonal hashing and finally the size distribution for molecules the training dataset that match the conditions of each task (the ground truth size distribution) in green no hashing.
shown left
Measuring novelty generative models trained the QM9 dataset is challenging because the QM9 dataset contains an exhaustive enumeration of all molecules that satisfy certain predetined constraints [46), [8) Therefore, if novel molecule is produced it means the generative model has failed to capture some of the physical properties of the dataset and indeed is found in [8] that during training, the model improved, novelty decreased. Novelty is therefore evaluating molecular diffusion models. For completeness, we include the novelty scores Table 7 as comparison to the results presented in [ Appendix C. We find that samples closer the statistics of the training dataset whilst still producing 'novel' samples at a consistent rate.
typically included in
D.2
Video
D.2.1 Dataset
We used the vp² benchmark, which consists of 35 000 videos, each 35 frames long. The videos
are evenly divided among seven tasks, namely: push (red, green, blue) button, open {slide, drawer), push {upright block, flat block} off table. The 5000 videos for each task were collected using scripted task-specific policy operating in the RoboDesk environ- ment 44). They sample action vector at every step during data generation by adding i.i.d. Gaussian
34

---

Janc
Figure 9: Samples generated by our model when conditional diffusion guidance applied. Each row represents one task with task at the top, down to task 10 at the bottom. For each task, 10 samples are shown each row.
noise each dimension of the action vector output by the scripted policy: For each task, they sample 2500 videos with noise standard deviation 0.1 and 2500 videos with standard deviation 0.2. We filter the lower-quality trajectories sampled with noise standard deviation 0.2, and 17 500 videos (2500) per task with noise standard deviation 0. We convert these videos to 32 x 32 resolution and then, that the data we train on has varying lengths, create each training example by sampling length / from uniform distribution over f2..... 35) and then taking a random I-frame subset of the video.
only the
D.2.2 Forward Process
The video domain differs from molecules in two important ways. The first is that videos cannot be reasonably treated as commutation-invariann set. This is because the order of the frames matters. Secondly, generating a full new component for the molecules with a single pass autoregressive network is feasible, however, a component for the videos is full frame which is challenging for single pass autoregressive network to generate. We design our forward process to overcome these challenges.
35

---

Figure 10: Interpolations showing a sequence of generations for linearly increasing polarizability from 39 Bohr to 66 Bohr² with fixed random noise. Each row shows an individual interpolation with Bohr increasing from left to right.
36

---

Table 6: Ablation of when to set the forward rate on the conditional molecule generation task. We report dimension eπrоr as the average Hellinger distance between the generated and ground truth conditional dimension distributions as well as average sample quality metrics. Metrics are reported after 620k training iterations.
Method
Dimension Error
% Atom stable
% Molecule Stable
% Valid
t<0.037 = 0
0.227 10.16
91.5133.7
56.519.
72.0111
1x00.TT
0.162+0.0m
92.4+2.8
53.9#12
72.7+9.6
1<0.3T
0.266#0.11
92.0+3.2
13.5
66.6+12
Table Uniqueness and novelty metrics on unconditional molecule generation. We produce 10000 samples for each method and measure validity using RDKit. Uniquenss is judged whether the chemical graph is unique amongst the 10000 produced samples. Amongst the valid and unique molecules, then find the percentage that have chemical graph not present in the training dataset.
Percentage of Valid and Unique Molecules that are Novel
Method
% Valid % Valid and Unique
FDDM [8]
91.9
90.7
65.7
TDDM (ours)
92.3
89.9
53.6
TDDM, const X
86.7
84.4
56.9
TDDM, À X = 0
89.4
86.1
51.3
TDDM w/o Prop. 3
87.1
85.9
63.3
We define our forward process to delete frames in random order. This means that during generation, frames can be generated in any order in the reverse process, enabling more conditioning tasks since we can always ensure that whichever frames we want to condition on are added first. Further, we use a non-isotropic noise schedule by adding noise just to the frame that is about to be deleted. Once it deleted, we then start noising the next randomly chosen frame. This is that, in the backward direction, when a new frame is added, it is simply Gaussian noise. Then the score network will fully denoise that new frame before the next new frame is added. We now specify exactly how our forward process is constructed.
We enable random-order deletion by applying an initial shuffling operation occurring at time I = 0. Before this operation, we represent the video x as an ordered sequence of frames, Xg [x1, X2,..., Xno] During shuffling, we sample random permutation of the integers 1.... no. Then the frames are kept in the same order, but annotated with an index variable so that
have Xg+ [6 T (1)), (x) (2)), ... (x()a), (nn))))
We will run the forward process from 0 to I 100N. We will set the forward rate such we delete down from ny ne at time (N + 1)100. This is achieved heuristically by setting
0
for t < (N - n, + 1)100,
Xe(ns)
=
8 for t 2 (N + 1)100.
We can see that at time (N n. + 1)100 we will quickly delete down from n, to n. 1 at which point X, (n)) will become thus stopping deletion until the process arrives at the next multiple of 100 in time. When we hit deletion event, delete the frame from X, that has the current highest index variable π (n). In other words
for n. = (x))(2). [2]. otherwise {
K del (1 de (iXx)
where we use x to refer the shuffle index variable for the ith current frame in Xt-
We now provide an example progression of the forward deletion process. Assume we have 7g 4,
N 5 and sample permutation such that T(1) = 3, (2)) 2, m3) = 4. and (44 1. Initially
37

---

the state is augmented to include the shuffle index. Then the forward process progresses from t = 0
500 with components being deleted in descending order of the shuffle index
X = [[x; (1) ,3,),
2), (x, (),,44 (x} 1)]
X 1000 + [(x; (2) ,3,, (x) (2)
2), (x, (3)
(x.
(4) 1)]
X200 + [(x]] ,,3, (x) 22) (x) (1)
1)]
[(x]²) (2), 2), (xi
(4) 1)]
X300+
[(x{") ,)]
X 100+
In this example, due the random permutation sampled, the final video frame remained after all others had been deleted. Note that the order of frames is preserved as we delete frames in the forward process although the spacing between them can change as we delete frames in the middle.
Between jumps, we use a noising process to add noise to frames. The noising process is non-isotropic in that adds noise to different frames at different rates such that the a frame is noised only in the
time window immediately preceding its deletion. For component i E [1,..., mel. we set the forward
(6) noising process such that Pijo( (x (8) |x0 M,)- (x x or (x (xxx)))) where xo is the clean frame
(8)
(0)
(0)
corresponding to
(i) as given by the mask Mt and (x()
(6) follows
for t < (N x; (2))
ox(xx() = 100
for l > (N - xx 12)100,
(N -
xx x! (2))
for (N x; [2)) t (N x! [2] + 1)100
where we again use x(1) [2] for the shuffle index of component i. This the VE-SDE from applied to each frame in turn. We note that we only add noise to the state values on not the shuffle index itself.
(3) The SDE parameters that result in the VE-SDE are b = 0 and Jt = V 21 2(N x; 10)
D.2.3 Sampling the Backward Process
When is not at a multiple of 100, the forward process is purely adding Gaussian noise, and so the reverse process is also purely operating on the continuous dimensions. We use the Heun sampler proposed by [66] update the continuous dimensions in this case, and also a variation of their discretisation of L specifically to update from e.g. / = 600 to = 500, we use their discretization of I as if the maximum value was 100 and then offset all values by 500.
To invert the dimension deletion process, we can use Proposition to derive our reverse dimension generation process. We re-write our parameterized T using Proposition as
(poo(n + 1mm
{(()))) = Xi(ne + 1) (no (x))
Pejo moo
At each time multiple of 100 in the backward process, we will have an opportunity to add a component.
At this time point, we estimate the expectation with single sample no ~Ve 1rjee X,). (n) If no > ne
then ( = oo. The new component will then be added at which point À 0 (X,) becomes 0 for the remainder of this block of time due to n, becoming n4, + 1. If ng n, then X (X,) 0 and no new component is added. 5 2 (X,) will continue to be 0 for the remainder of the backward process an opportunity to add component is not used.
When new frame is added, we use A& add x|X,) to decide where the frame added and its initial value. Since when we delete a frame it fully noised, AQ (yddd ji{,)) can simply predict Gaussian noise for the new frame ydd However, (f (yaddd .iX.,) will still learn predict a suitable location to place the new frame such that backward process is the reversal of the forward.
We give an example simulation from the backward generative process in Figure II
38

---

While would possible to train single network that estimates the score and all quantities needed for modelling jumps, we chose to train two separate networks in order to factorize our exploration of the design space. These were the score network of, and the rate and index prediction network modeling Poje (no X,) and AQ (i)..). The rate and index prediction network similar the first half' of the score network, that uses all U-net blocks up to and including the middle one. We then frame input, Figure 11: An example simulation of the backward generative process conditioned the first and last frame. Note how the process first adds a new frame and then fully denoises it before adding the next frame. Since the first and last frame are very similar, the process produces short video.
D.2.4 Network Architecture
Our video diffusion network architecture is based on the U-net used by 42 which takes input the index of each frame within the video, and uses the differences between these indices to control the interactions between frames via attention mechanism. Since, during generation, we do not know the final position of each frame within the Xg, we instead pass in its position within the ordered sequence X,.
One further difference that, since we are perform non- non-isotropic diffusion, the standard deviation of the added noise will differ between frames. We adapt to this by performing preconditioning, and inputting the timestep embedding, separately for each frame (6) based on (x(((() instead of basing them the global diffusion timestep Our timestep embedding and pre- and post-conditioning of network inputs/outputs are as suggested by [16], other than being done on per-frame basis. The architecture from 42 with these changes applied then gives
score network s.
flatten the 512 x 4 4 hidden state for each frame after this block such that, for we obtain n, x 8192 hidden state. These are fed through ID convolution with kernel size 2 and zero- vero-padding of size 1 on each end, reducing the hidden state to (n, 1) x 128, which is in turn fed
39

---

through ReLU activation function. This hidden state is then fed into three separate heads. One head maps the parameters of A. (x) via a ID convolution of kernel size 3. The output of size (ne + 1) is fed through softmax to provide the categorical distribution AQ (e) The second head averages the hidden state over the "frame" dimension, producing a 128-dimensional vector. This is fed through a single linear layer and a softmax to parameterize Poje X,). Finally, the third head
consists of a ID convolution of kernel size with 35 output channels. The (74 + 1) x 35 output is fed through softmax parameterize distributions over the number of frames that were deleted from Xa which came before the first in x,, the number of frames from Xn which were deleted between each pair of frames in X,, and the number deleted after the last frame in X,- We do not use this head at inference-time but found that including it improved the performance of the other heads by helping the network learn better representations.
For a final performance improvement, we note that under our forward process there is only ever one "noised" frame in xg. while there sometimes many clean frames. Since the cost of running architecture scales with the number of frames, running it on many clean frames may significantly increase the cost while providing little improvement to performance. We therefore only feed into the architecture the "noised" frame, the two closest "clean" frames before it, and the two closest "clean" frames after it. See our released source code for the full implementation of this architecture.
D.2.5 Training
To sample during training, we adapt the log- tog-normal distribution suggested by [6 in the context of isotropic diffusion over a single image. To apply it to our non-isotropic video diffusion, we first sample which frames have been deleted, which exist with no noise, and which have had noise added, by sampling the timestep from uniform distribution and simulating our proposed forward process. We then simply change the noise standard deviation for the noisy frame, replacing with a sample from the log-normal distribution. The normal distribution underlying our log- normal has mean -0.6 and standard deviation 1.8. This can be interpreted as sampling the timestep from a mixture of log-normal distributions, N Ei L.N(t 100i; 0.66, 1.8²). Here, the mixture index can be interpreted controlling the number of deleted frames.
We use the same loss weighting as [16] but, similarly to our use of preconditioning, compute the weighting separately for each frame x; as function of a, (x(()) noise.
) to account for the non-isotropic
D.2.6 Perceptual Quality Metrics
We now verify that our reverse process does not have any degradation in quality during the generation as more dimensions are added. We generate 10000 videos and throw away the 278 that were sampled to have only two frames. We then compute the FID score for individual frames in each of the remaining 9722 videos. We group together the scores for all the first frames be generated in the reverse process and then for the second frame to be generated and so on. We show our results in Table We find that when a frame is inserted has no apparent effect on perceptual quality and conclude that there is no overall degradation in quality as our sampling process progresses. We note that the absolute value of these FID scores may not meaningful due to the RoboDesk dataset being far out of distribution for the Inception network used to calculate FID scores. We can visually confirm good sample quality from Figure 5
Table 8: FID for video frames grouped by when they were inserted during sampling.
1st
2nd
3rd
3rd last
2nd last
last
34.2
34.9
34.7
34.2
34.1
34.4
E Broader Impacts
In this work, we presented a general method for performing generative modeling on datasets of varying dimensionality. We have not focused on applications and instead present generic method.
40

---

Along with other generic methods for generative modeling, we must consider the potential negative social impacts that these models can cause when inappropriately used. As generative modeling capabilities increase, it becomes simpler to generate fake content which can be used to spread misinformation. addition this, generative models are becoming embedded into larger systems that then have real effects on society. There will be biases present within the generations created by the model which in turn can reinforce these biases when the model's outputs are used within wider systems. In order to mitigate these harms, applications of generative models to real world problems must accompanied with studies into their biases and potential ways they can misused. Further, public releases of models must be accompanied with model cards 47 explaining the biases, limitations and intended uses of the model.
41


Reviews:
[{"review_id": "iGGic94Veb", "reviewer": "Reviewer_EZPD", "summary": "This paper proposes jump diffusion, which is a novel diffusion model to handle data with varying dimensions. The proposed method is derived from a special forward process that contains a jump part that changes the dimension of the generated samples. The corresponding backward process and the learning objective are derived, with two more components that need to be learned other than the standard diffusion drift. Numerical issues are properly handled to make the mathematical model work in practice. Experiments over various application scenarios are conducted to demonstrate the flexibility and versatility of the proposed framework.", "questions": "I have one question on the proposed diffusion model and I would appreciate answers from the authors. \n\nIn Algorithm Box 1, what is the initial distribution? Is it a random standard Gaussian noise with only 1 dimension? If so, how do we decide which dimension is the initial dimension? For example, if we have x = (3, x_1, x_2, x_3), what is the initial random noise? (1, N(0, I), 0, 0), (1, 0, N(0, I), 0), or (1, 0, 0,  N(0, I) ? Or are they uniformly sampled? \n\nMoreover, if the initial number of dimension is always 1, how do we guarantee that in setting the forward process? I feel like there must be some constraints on the scheduling of rate function $\\lambda(t)$.\n\nMaybe I missed some parts of the derivation in the paper.", "limitations": "The limitations and broader impacts are properly discussed.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method targets at a very important problem. It is a great contribution to invent theoretically sound diffusion models to deal with data with varying dimensions. This kind of problem frequently happens in real-world applications.\n\n2. The motivation is strong, the idea is reasonable, and the mathematical derivations makes the idea a solid framework. \n\n3. Applications on molecule and video generation showcases the potential of the proposed framework in reality.", "weaknesses": "1. Although I do think the proposed framework is interesting, generation with varying dimensionality can be addressed easily by \n\n(1) learning a distribution over the dimension numbers and sample dimension number from the learned distribution (2) sampling initial random noises according to the sampled dimension number at the beginning of the generation and pad the unused dimensions with 0. \n\nThis simple method has little modification over existing diffusion models, and the only extra efforts one need to do is to add paddings during training.\n\nIt would be very nice if I can hear comments from the authors on this simple baseline. I do appreciate the theoretical contribution of this paper, and it would be better if the authors can compare with this simple baseline .  \n\n2.In L84, should the summation be taken over m<n? Because there is only deletion in the forward process.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their engagement with our proposed methodology and thoughtful questions. We are grateful that our work is considered to be a solid framework to tackle a very important problem. We answer questions and respond to comments below.\n\n> *It would be very nice if I can hear comments from the authors on this simple baseline.*\n\nWe focussed in this work on training flexible unconditional models that jointly model all the relevant aspects of the data: the state and the dimension. This is in contrast to the mentioned baseline that models the state and dimension separately with two separate models. We showcase the benefit of the joint modelling approach by using the task of diffusion guidance where an unconditional model is first trained and then, at test time, different researchers with different goals can condition this model on their task of interest using diffusion guidance, for example Weiss et al.  2023 and Crowson 2021. This has the advantage that, should the unconditional model be powerful enough, any conditional generation task can be easily accomplished with a limited computational budget because no re-training is required, the only thing needed is a guidance model. In the case of the baseline that you mention, complete knowledge of the final generation task is needed before training because the dimension prediction model and score model both need to be conditioned on this task information. Therefore this type of approach is not applicable to our intended application where the end user does not have the resources for a complete re-training of the models for the tasks they are interested in.\n\nMotivated by this deficiency of models that treat dimensions and state values separately, we developed a generative model that jointly models dimension and state meaning it is powerful enough to be used as an unconditional model that can then be guided at test time for different generation tasks. In this case where the end user is guiding the model on their task of interest, the fixed dimensional model fails to capture the correct dimension information whereas our method can produce more accurate dimensions statistics for each desired task as we show in Table 3 in the paper.\n\n> *In L84, should the summation be taken over m<n? Because there is only deletion in the forward process.*\n\nThank you for pointing this out, you are correct that for our forward process $K_t( m , \\mathbf{y} | \\mathbf{X})$ will be zero for $m \\geq n$ due to the forward process only deleting dimensions. For the backward process we would instead have $K_t(m , \\mathbf{y} | \\mathbf{X})$ being $0$ for $m \\leq n$. On L84 we introduce jump diffusions for the first time and so try to keep the formulation general enough to cover both the forward and backward cases. However, given that this introduction is in the forward process section this could lead to confusion and we will make this clearer in a revision to the paper.\n\n> *In Algorithm Box 1, what is the initial distribution?*\n\nThe initial distribution is indeed a single dimension of standard Gaussian noise. However, we do not need to pre-determine which dimension in the final generated datapoint it corresponds to. We treat it genuinely as only a single dimension i.e. the initial sample is $(1, \\mathcal{N}(0, I))$ with no baked in knowledge of the length of the final datapoint to create nor the initial dimension’s place in that final datapoint. Then, when a new dimension is added it could be added to the left or to the right becoming $(2, z, \\mathcal{N}(0, I) )$ or $(2, \\mathcal{N}(0, I), z)$ where $z$ is the new value to be added. This process is repeated until we reach the terminal time $t=0$ of the reverse generative process. Therefore, that initial dimension could correspond to any of the dimensions in the final generated datapoint depending on where new dimensions are inserted. Furthermore, the final number of dimensions also varies depending on how many jumps occurred during generation. In your example e.g. $(1, 0, \\mathcal{N}(0, I), 0)$ this assumes the final dimension is 3 but we make no such assumption in our method.\n\n> *Moreover, if the initial number of dimension is always 1, how do we guarantee that in setting the forward process? I feel like there must be some constraints on the scheduling of rate function*\n\nThis is a good point and indeed this constraint influences the choice of forward process rate function, $\\lambda_t$. We set it such that, with high probability, the number of dimensions at the end of the forward process is $1$. This is the same idea as in fixed dimension diffusion models where the noising process is such that the corrupted data at the end of the noising process is approximately distributed according to $\\mathcal{N}( \\mathbf{x}; 0, I)$. Here, we have a noising process such that the corrupted data at the end of the noising process is approximately distributed according to $\\mathbb{I} \\\\{n = 1\\\\} \\mathcal{N}(\\mathbf{x}; 0, I)$. In practice we achieve this by setting $\\lambda_t$ large enough such that there is a high probability of enough jumps occurring to remove all but one of the dimensions in the datapoint. Further, we set $\\lambda_t$ to $0$ when the number of dimensions is equal to $1$ so that we can’t delete down to $0$ dimensions.\n\n## References\n\nTomer Weiss, Luca Cosmo, Eduardo Mayo Yanes, Sabyasachi Chakraborty, Alex M Bronstein,and Renana Gershoni-Poranne. *Guided diffusion for inverse molecular design.* ChemrXiv, 2023\n\nKatherine Crowson. *Clip guided diffusion.* Web Demo, 2021\n"}, {"review_id": "sA2MvPJfZP", "reviewer": "Reviewer_FTvW", "summary": "This paper addresses the problem of modelling data of various dimensions. This is achieved by generalising diffusion models as jump diffusion processes, allowing the content and dimension of data to be jointly modelled. The forward process gradually corrupts the data with gaussian noise while also gradually deleting dimensions until a single normally distributed dimension remains. The reverse process learns both the score function, when to add a dimension, and what data to place in the newly created dimension. Experimental results on molecule and video generation tasks show that the approach can well represent the distribution of data dimensions, while outperforming/being competitive in terms of sample quality. ", "questions": "- Do the authors have any rationale/evidence that fixing $\\lambda$ at $t<0.1$ is a reasonable choice? For instance, graphing metrics over a variety of limits.\n- Similarly, is there rationale why predicting new dimensions with a gaussian is sufficient; or if for more complex data this is problematic, is there a more expressive extension?", "limitations": "Limitations are well discussed throughout the paper. The previously discussed weaknesses could be mentioned as well.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The proposed approach of modelling the generative process as a jump diffusion process is interesting and is a sensible solution. \n- The paper is very clear, easy to read, and to my understanding, the method is technically sound. I particularly like the intuitive explanation of the loss in Equations 3-4; and predicting the original data dimension (line 191) is a good idea to address the optimisation issues.\n- Experimental results demonstrate the effectiveness of the proposed approach (Table 2), performing comparability to, or outperforming the baseline which sample dimension prior to sampling. I also like the evaluation of the impact of setting $\\lambda=0$ (dimension deletion/insertion rate) towards the end of the diffusion process.\n- It is shown that reconstruction guidance can be used to generate molecules with specified features (Table 3) much better than dimension independent approaches.\n- The problem of jointly modelling content and dimension of data is an important one to address. In my opinion, the proposed solution is compelling and I believe it will be very useful to others.\n", "weaknesses": "- Predicting the content of the newly created dimensions with a gaussian distribution is limiting, potentially creating a discrepancy between train/test time and requiring more diffusion time to correct.\n- Fixing $\\lambda=0$ towards the end of the diffusion process to ensure added dimensions have time to diffuse is a weakness of the proposed approach. Particularly for more complex data, there is no guarantee that the remaining time will be sufficient.\n- Finally, it would have been nice to see a perceptual quality metric for the video generation task, this would offset the above weaknesses.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their review and positive comments on the paper. We are especially happy to hear that our proposed method is considered to be a compelling solution to an important problem. You have raised important points regarding the new dimension distribution and the time needed to diffuse at the end of the process. Since these are related issues, we answer both here.\n\nWe first emphasize that our framework can be seen as an integration of autoregressive models with diffusion models. In the limit of no diffusion and only jumps then we arrive at a pure autoregressive model whereas if we include diffusion but initialize each new dimension with $\\mathcal{N}(0, I)$ and create all dimensions at the start of the reverse generative process then we arrive at a pure diffusion model. We show in the paper how there is a fertile middle ground where we can derive benefits from both types of model with the diffusion part giving good sample quality and the autoregressive part giving trans-dimensional generation capabilities.\n\nThere is a lot of flexibility in the choice for the new dimension creation distribution which we refer to as the autoregressive distribution. We chose to parameterize it with a Gaussian in the paper for simplicity but it can be parameterized with any likelihood model because the training signal for learning the autoregressive distribution is simply a maximum likelihood objective for predicting the missing part of the data given the observed part. Therefore, depending on the task, we could use more expressive alternatives such as normalizing flows or a G-SchNet architecture for molecules (Gebauer et al. 2019) which predicts new atom positions using discrete probabilities over binned distances.\n\nIn our experiments we found that a reasonably simple and effective approach is to have the network predict mean and standard deviation statistics for a Gaussian distribution and then refine this with the diffusion part of the process. In our preliminary experiments we found that $t=0.1$ to $t=0$ is sufficient time to clean up any errors and produce high quality samples whilst retaining the trans-dimensional nature of the model and avoiding instabilities near $t=0$.\n\nTo demonstrate this effect, we have run a sweep over when to set $\\lambda_t$ to $0$ on the molecule task, see Table 2 in the additional results pdf. We have found that the setting $\\lambda_{t < 0.03T} = 0$ 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 (L276) when setting $\\lambda_t$ to be constant for all $t$. Further, the setting $\\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. Hence, we believe the setting $\\lambda_{t<0.1T}$ to be reasonable in our case. Note these models have not trained for as long as the ones in the paper due to time constraints during the rebuttal period.\n\nWe agree that the specific balance could be different for different datasets and model architectures and ultimately this is a case of hyperparameter selection. Further, if it is found that a small amount of diffusion time is insufficient to generate high quality samples, then the expressivity of the autoregressive part can be arbitrarily increased, the only downsides being an increased computational cost and code complexity.\n\nThank you for the suggestion of perceptual metrics on the video dataset. To investigate the relative quality of frames that are added near the start of the diffusion process versus those that are added near the end, we calculated the FID for individual frames, grouped by when they are added in the generative process. Our results can be found in Table 3 in the additional results pdf. We find no systematic trend meaning there is no degradation in quality for frames added near the end of the generative process. Note that the absolute value of these FIDs may not be meaningful due to the RoboDesk dataset being far out of distribution for the Inception network used to calculate FID scores. We can visually confirm good sample quality from e.g. Figure 5 in the paper. \n\n## References\nNiklas W. A. Gebauer, Michael Gastegger, Kristof T. Schütt. *Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules.* Advances in Neural Information Processing Systems, 2019.\n"}, {"review_id": "dMeDDVlPOF", "reviewer": "Reviewer_fiGg", "summary": "This paper focuses on varying dimensional datasets and proposes a novel generative model to solve the varying dimensional problems. The proposed model is theoretically valid and has an interesting and novel contribution to extending the traditional score-based generative model by generating both state values and dimensions jointly during the generative process, which is idea-simple but effective. Experiments on molecule generation and video generation both show this model's effectiveness.", "questions": "1. I am confused about how to properly define $K^{{del}}(i|n)$ in $\\overrightarrow{K}_t(\\mathbf{Y}|\\mathbf{X})$? Any clarification about this issue will be helpful.", "limitations": "Yes, the authors have adequately addressed the limitations and potential negative societal impact of their work.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The idea of jointly modeling the state and dimension, in particular, the idea of using the intensity function to model the jump distribution, is interesting and novel.\n\n2. The theoretical contribution is valid and the experiment is thorough in supporting the proposed model.", "weaknesses": "None.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their review and we appreciate the praise for our method’s novelty and thorough experiments. We address your questions below.\n\n> *I am confused about how to properly define $K^\\text{del}(i | n)$ in $\\overrightarrow{K}_t(\\mathbf{Y} | \\mathbf{X})$ Any clarification about this issue will be helpful.*\n\n$K^\\text{del}(i | n)$ will dictate the ordering in which dimensions are deleted in the forward noising process. For example if $K^\\text{del}(i | n) = \\mathbb{I} \\\\{ i = n \\\\}$ then it is always the final dimension of the datapoint that is deleted when a jump occurs. Alternatively, you can set $K^\\text{del}(i|n) = 1/n$ so that the dimension to be deleted is chosen uniformly at random when a jump occurs.\n\nThe choice of $K^\\text{del}( i | n)$ impacts the reverse generative process. If $K^\\text{del}(i | n) = \\mathbb{I} \\\\{ i = n \\\\}$ then datapoints are constructed in an additive way, each new dimension is simply appended onto the end of the current datapoint. This is similar to how a text based autoregressive model builds sentences by generating new words and appending them onto the end of the current sentence. Instead, if $K^\\text{del}(i|n) = 1/n$ then during the reverse generative process, the generative model first picks a suitable place to add a new dimension and then inserts one at that chosen point.\n\nThe choice of $K^\\text{del}( i | n)$ depends on the dataset and problem being tackled. In our case, we choose $K^\\text{del}(i|n) = 1/n$ because for molecules there is no natural notion of a ‘final dimension’ due to permutation invariance of the point cloud. Further, for our videos we want to condition on the first and last frame thus an autoregressive appending style of generation would be unsuitable. We present our framework in a general manner so it can be easily applied to other problems where other choices of $K^\\text{del}( i | n)$ could be more suitable.\n\nWe will make this choice clearer in an update to the paper, thank you for raising this point of confusion.\n"}, {"review_id": "xHk5eUsJX9", "reviewer": "Reviewer_tNRG", "summary": "This paper proposes a new diffusion model based on jump diffusion processes. Compared with previous discrete and continuous formulations, the model introduces the usage of a transition kernel, which models the jump process in a semantically meaningful manner. The method absorbs standard constructions like diffusion guidance and produces good results on molecule and robot arm video tasks.", "questions": "Nothing more than the weaknesses section.", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The proposed methodology is elegant and theoretically sound.\n* Many constructions like diffusion guidance extend naturally.\n* The proposed methodology is widely applicable, as many natural data types include a continuous portion (modeling the space) as well as a discrete but unknown component.", "weaknesses": "* For the molecule task, the metrics are based on molecular validity properties. This seems like it wouldn't account for overfitting from the model, which seems more relevant for a generative modeling task. In particular, can you report a diversity/non-memorization metric?\n* The robot arm example, while sufficient for the paper, is rather toy.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their review and positive feedback. We are very pleased to hear that our approach is considered to be elegant and theoretically sound with wide applicability. We address the specific comments from the review here.\n\n> *For the molecule task, the metrics are based on molecular validity properties. This seems like it wouldn't account for overfitting from the model, which seems more relevant for a generative modeling task. In particular, can you report a diversity/non-memorization metric?*\n\nThank you for this suggestion, we have investigated sample diversity and novelty and will include these results in an update to the paper. On the molecule task, we can measure uniqueness by computing the chemical graph corresponding to each generated sample and measure what proportion of the 10000 produced samples have a unique chemical graph amongst this set of 10000 as is done by Hoogeboom et al. 2022. We show our results in Table 1 in the additional results pdf and find our method without any ablations has only slightly lower levels of uniqueness when compared to the fixed dimension diffusion model baseline. Measuring novelty on generative models trained on the QM9 dataset is challenging because the QM9 dataset contains an exhaustive enumeration of all molecules that satisfy certain predefined constraints as noted by Vignac et al. 2022 and Hoogeboom et al. 2022. Therefore, if a novel molecule is produced it means the generative model has failed to capture some of the physical properties of the dataset and indeed Hoogeboom et al. 2022 found that during training, as the model improved, novelty decreased. Novelty is therefore not typically included in evaluating molecular diffusion models. For completeness, we include the novelty scores in Table 1 as a comparison to the results presented in Hoogeboom et al. 2022 Appendix C. We find that our samples are closer to the statistics of the training dataset whilst still producing ‘novel’ samples at a consistent rate.\n\n\n## References\nEmiel Hoogeboom, Vıctor Garcia Satorras, Clément Vignac, and Max Welling. *Equivariant diffusion for molecule generation in 3d.* International Conference on Machine Learning, 2022.\n\nClement Vignac and Pascal Frossard. *Top-n: Equivariant set and graph generation without exchangeability.* International Conference on Learning Representations, 2022\n"}]

Author rebuttal:
We would like to thank all the reviewers for their analysis of our paper and very helpful reviews. We were pleased that reviewers considered our work to be a novel and theoretically sound method for tackling the important problem of modeling data with varying dimensionality.

We address comments made by each reviewer in individual responses and attach an additional results pdf here.


Meta-review: Authors  propose in this paper a new diffusion model based on jump diffusion processes that handles varying dimensions. All reviews agreed on the novelty and the soundness of the work. The proposed methods was validated on molecules, videos and robotics applications.
Decision: Accept (spotlight)