Title: Trans-Dimensional Generative Modeling via Jump Diffusion Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: diffusion

score-based

score

markov chain

jump diffusion

poisson
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

Abstract

We propose a new class of generative models that naturally handle 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.

1 Introduction

Generative models based on diffusion processes [1-3] have become widely used in solving a range of problems including text-to-image generation [4, 5], audio synthesis [6] and protein design [7]. These models define a forward noising diffusion process that corrupts data to noise and then learn the corresponding time-reversed backward generative process that generates novel datapoints from noise.

In many applications, for example generating novel molecules [8] or videos [9, 10], the dimension of the data can also vary. For example, a molecule can contain a varying number of atoms and a video can contain a varying number of frames. When defining a 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 [11].

diffuse

diffuse

jump

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-14] without needing to perform any further training or fine-tuning of the model on their task of interest. Since the diffusion occurs in a fixed dimensional space, there is no way for the guidance to appropriately guide the dimension of the generated datapoint. This can lead to incorrect generations for datasets where the dimension greatly affects the nature of 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 a backward generative process that denoises and adds dimensions, see Figure 1. We derive the optimum backward generative process as the time-reversal of the forward noising process and derive a novel learning objective to 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, 15-17] define a forward diffusion process through a stochastic differential equation (SDE) where x 0 ∼ p data and, for t > 0 ,

where x t ∈ R d is the current state, - → b t : R d → R d is the drift and g t ∈ R is the diffusion coefficient. d w t is a Brownian motion increment on R d . This SDE can be understood intuitively by noting that in each infinitesimal timestep, we move slightly in the direction of the drift - → b t and inject a small amount of Gaussian noised governed by g t . Let p t ( x t ) denote the distribution of x t for the forward diffusion process (1) so that p 0 ( x 0 ) = p data ( x 0 ) . - → b t and g t are set such that at time t = T , p T ( x T ) is close to p ref ( x T ) = N ( x T ; 0 , I d ) ; e.g. - → b t ( x t ) = -1 2 β t x t , g t = √ β t for β t > 0 [2, 3].

The time-reversal of the forward diffusion (1) is also a diffusion [18, 19] which runs backwards in time from p T ( x T ) to p 0 ( x 0 ) and satisfies the following reverse time SDE

where ← -b t ( x t ) = - → b t ( x t ) -g 2 t ∇ x t log p t ( x t ) , d t is a negative infinitesimal time step and dˆ w t is a Brownian motion increment when time flows backwards. Unfortunately, both the terminal

distribution, p T ( x T ) , and the score, ∇ x t log p t ( x t ) , are unknown in practice. A generative model is obtained by approximating p T with p ref and learning an approximation s θ t ( x t ) to ∇ x t log p t ( x t ) typically using denoising score matching [20], i.e.

For a flexible model class, s θ , we get s θ t ( x t ) ≈ ∇ x t log p t ( x t ) at the minimizing parameter.

3 Trans-Dimensional Generative Model

Instead of working with fixed dimension datapoints, we will instead assume our datapoints consist of a 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, x , giving X = ( n, x ) . Since each datapoint can have a variable number of components from n = 1 to n = N , our overall space that our datapoints live in is the union of all these possibilities, X ∈ X = ⋃ N n =1 { n } × R nd . For example, for a varying size point cloud dataset, components would refer to points in the cloud, each containing ( x, y, z ) coordinates giving d = 3 and the maximum possible number of points in the cloud is N .

Broadly speaking, our approach will follow the same framework as previous diffusion generative models. We will first define a forward noising process that both corrupts state values with Gaussian noise and progressively deletes dimensions. We then learn an approximation to the time-reversal giving a backward generative process that simultaneously denoises whilst also progressively adding dimensions back until a synthetic datapoint of appropriate dimensionality has been constructed.

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 a standard SDE. When a jump occurs, the process transitions to a different dimensional space with the new value for the process being drawn from a transition kernel K t ( Y | X ) : X × X → R ≥ 0 . Letting Y = ( m, y ) , the transition kernel satisfies ∑ m ∫ y K t ( m, y | X )d y = 1 and ∫ y K t ( m = n, y | X )d y = 0 . The rate at which jumps occur (jumps per unit time) is given by a rate function λ t ( X ) : X → R ≥ 0 . For an infinitesimal timestep d t , the jump diffusion can be written as

with X t ≜ ( n t , x t ) and X t +d t ≜ ( n t +d t , x t +d t ) and d w t being a Brownian motion increment on R n ′ t d . We provide a more formal definition in Appendix A.

With the jump diffusion formalism in hand, we can now construct our forward noising process. 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 - → b t ( X ) = -1 β t x and - → g t = √ β t with β t ≥ 0 .

̸

When a jump occurs in the forward process, one component of the current state will be deleted. For example, one point in a point cloud or a single frame in a video is deleted. The rate at which these deletions occur is set by a user-defined forward rate → λ t ( X ) . To formalize the deletion, we need to introduce some more notation. We let K del ( i | n ) be a user-defined distribution over which component of the current state to delete. We also define del : X× N →X to be the deletion operator that deletes a specified component. Specifically, ( n -1 , y ) = del (( n, x ) , i ) where y ∈ R ( n -1) d has the same values as x ∈ R nd except for the d values corresponding to the i th component which have been removed. We can now define the forward jump transition kernel as - → K t ( Y | X ) = ∑ n i =1 K del ( i | n ) δ del ( X ,i ) ( Y ) . We note that only one component is ever deleted at a time meaning - → K t ( m, y | X ) = 0 for m = n -1 . Further, the choice of K del ( i | n ) will dictate the behaviour of the reverse generative process. If we set K del ( i | n ) = I { i = n } then we only ever delete the final component and so in the reverse generative

direction, datapoints are created additively, appending components onto the end of the current state. Alternatively, if we set K del ( i | n ) = 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 t = 0 to t = T and should be such that at time t = T , the marginal probability p t ( X ) should be close to a reference measure p ref ( X ) that can be sampled from. We set p ref ( X ) = I { n = 1 }N ( x ; 0 , I d ) where I { n = 1 } is 1 when n = 1 and 0 otherwise. To be close to p ref, for the jump part, we set → λ t high enough such that at time t = T there is a high probability that all but one of the components in the original datapoint have been deleted. For simplicity, we also set → λ t to depend only on the current dimension → λ t ( X ) = → λ t ( n ) with → λ t ( n = 1) = 0 so that the forward process stops deleting components when there is only 1 left. In our experiments, we demonstrate the trade-offs between different rate schedules in time. For the diffusion part, we use the standard diffusion β t schedule [2, 3] so that we are close to N ( x ; 0 , I d ) .

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 a backward drift ← -b t ( X ) , diffusion coefficient ← -g t , rate ← -λ t ( X ) and transition kernel ← -K t ( Y | X ) . We would like these quantities to be such that the backward process is the time-reversal of the forward process. In order to find the time-reversal of the forward process, we must first introduce some notation to describe ← -K t ( Y | X ) . ← -K t ( Y | X ) should undo the forward deletion operation. Since - → K t ( Y | X ) chooses a component and then deletes it, ← -K t ( 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 new component will be denoted y add ∈ R d . The insertion operator is defined as ins : X × R d × N → X . It takes in the current value X , the new component y add and an index i ∈ { 1 , . . . , n +1 } and inserts y add into X at location i such that the resulting value Y = ins ( X , y add , i ) has del ( Y , i ) = X . We denote the joint conditional distribution over the newly added component and the index at which it is inserted as A t ( y add , i | X ) . We therefore have ← -K t ( Y | X ) = ∫ y add ∑ n +1 i =1 A t ( y add , i | X ) δ ins ( X , y add ,i ) ( Y )d y add . Noting that only one component is ever added at a time, we have ← -K t ( m, y | X ) = 0 for m = n +1 .

̸

This backward process formalism can be seen as a unification of diffusion models with autoregressive models. The diffusion part ← -b t denoises the current set of components in parallel, whilst the autoregressive part A t ( y add , i | X ) predicts a new component and its location. ← -λ t ( X ) is the glue between these parts controlling when and how many new components are added during generation.

We now give the optimum values for ← -b t ( X ) , ← -g t , ← -λ t ( X ) and A t ( y add , i | X ) such that the backward process is the time-reversal of the forward process.

Proposition 1. The time reversal of a forward jump diffusion process given by drift - → b t , diffusion coefficient - → g t , rate → λ t ( n ) and transition kernel ∑ n i =1 K del ( i | n ) δ del ( X ,i ) ( Y ) is given by a jump diffusion process with drift ← -b ∗ t ( X ) , diffusion coefficient ← -g ∗ t , rate ← -λ ∗ t ( X ) and transition kernel

All proofs are given in Appendix A. The expressions for ← -b ∗ t and ← -g ∗ t are the same as for a standard diffusion except for replacing ∇ x log p t ( x ) with ∇ x log p t ( X ) = ∇ x log p t ( x | n ) which is simply the score in the current dimension. The expression for ← -λ ∗ t can be understood intuitively by noting that the numerator in the probability ratio is the probability that at time t , given a deletion occurs, the forward process will arrive at X . If this is higher than the raw probability at time t that the forward process is at X (the denominator) then we should have high ← -λ ∗ t because X is likely the result of a deletion of a larger datapoint. Finally the optimum A ∗ t ( y add , i | X ) is simply the conditional distribution of y add and i given X when the joint distribution over y add , i, X is given by p t ( ins ( X , y add , i )) K del ( i | n +1) .

3.3 Objective for Learning the Backward Process

The true ← -b ∗ t , ← -λ ∗ t and A ∗ t are unknown so we need to learn approximations to them, ← -b θ t , ← -λ θ t and A θ t . Following Proposition 1, we set ← -b θ t ( X ) = - → b t ( X ) - - → g 2 t s θ t ( X ) where s θ t ( X ) approximates ∇ x log p t ( X ) . The forward and parameterized backward processes are summarized in Table 1.

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 E p data ( x 0 ) [log p θ 0 ( x 0 )] [21]. We derive here an equivalent loss to learn s θ t , ← -λ θ t and A θ t for our jump diffusion process by leveraging the results of [17] and [22]. Before presenting this loss, we first introduce some notation. Our objective for s θ t ( X t ) will resemble denoising score matching (2) but instead involve the conditional score ∇ x t log p t | 0 ( X t | X 0 ) = ∇ x t log p t | 0 ( x t | X 0 , n t ) . This is difficult to calculate directly due to a combinatorial sum over the different ways the components of X 0 can be deleted to get to X t . Weavoid this problem by equivalently conditioning on a mask variable M t ∈ { 0 , 1 } n 0 that is 0 for components of X 0 that have been deleted to get to X t and 1 for components that remain in X t . This makes our denoising score matching target easy to calculate: ∇ x t log p t | 0 ( x t | X 0 , n t , M t ) = √ α t M t ( x 0 ) -x t 1 -α t where α t = exp ( -∫ t 0 β ( s )d s ) [3]. Here M t ( x 0 ) is the vector removing any components in x 0 for which M t is 0 , thus M t ( x 0 ) and x t have the same dimensionality. We now state our full objective.

Proposition 2. For the backward generative jump diffusion process starting at p ref ( X T ) and finishing at p θ 0 ( X 0 ) , an evidence lower bound on the model log-likelihood E x 0 ∼ p data [log p θ 0 ( x 0 )] is given by

where expectations are with respect to U ( t ; 0 , T ) p 0 ,t ( X 0 , X t , M t ) K del ( i | n t ) δ del ( X t ,i ) ( Y ) , C is a constant term independent of θ and X t = ins ( Y , x add t , i ) . This evidence lower bound is equal to the log-likelihood when ← -b θ t = ← -b ∗ t , ← -λ θ t = ← -λ ∗ t and A θ t = A ∗ t .

We now examine the objective to gain an intuition into the learning signal. Our first term (3) is an L 2 regression to a target that, as we have seen, is a scaled vector between x t and √ α t M t ( x 0 ) . As the solution to an L 2 regression problem is the conditional expectation of the target, s θ t ( X t ) will learn to predict vectors pointing towards x 0 averaged over the possible correspondences between dimensions of x t and dimensions of x 0 . Thus, during sampling, s θ t ( X t ) provides a suitable direction to adjust the current value X t taking into account the fact X t represents only a noisy subpart of a clean whole X 0 .

The second term (4) gives a learning signal for ← -λ θ t and A θ t . For A θ t , we simply have a maximum likelihood objective, predicting the missing part of X t (i.e. x add t ) given the observed part of X t (i.e. Y ). The signal for ← -λ θ t comes from balancing two terms: - ← -λ θ t ( X t ) and → λ t ( n t ) log ← -λ θ t ( Y )

which encourage the value of ← -λ θ t to move in opposite directions. For a new test input Z , ← -λ θ t ( Z ) 's value needs to trade off between the two terms by learning the relative probability between Z being the entirety of a genuine sample from the forward process, corresponding to the ← -λ θ t ( X t ) term in (4), or Z being a substructure of a genuine sample, corresponding to the ← -λ θ t ( Y ) term in (4). The optimum trade-off is found exactly at the time reversal ← -λ ∗ t as we show in Appendix A.5.

We optimize L ( θ ) using stochastic gradient ascent, generating minibatches by first sampling t ∼ U (0 , T ) , X 0 ∼ p data and then computing X t from the forward process. This can be done analytically for the → λ t ( n ) functions used in our experiments. We first sample n t by analytic integration of the dimension deletion Poisson process with time inhomogeneous rate → λ t ( n ) . We then add Gaussian noise independently to each dimension under p t | 0 ( x t | X 0 , n t , M t ) using a randomly drawn mask variable M t . See Appendix B for further details on the efficient evaluation of our objective.

3.4 Parameterization

s θ t ( X t ) , A θ t ( y add , i | X t ) and ← -λ θ t ( X t ) will all be parameterized by neural networks. In practice, we have a single backbone network suited to the problem of interest e.g. a Transformer [23], an EGNN [24] or a UNet [25] onto which we add prediction heads for s θ t ( X t ) , A θ t ( y add , i | X t ) and ← -λ θ t ( X t ) . s θ t ( X t ) outputs a vector in R n t d . A θ t ( y add , i | X t ) outputs a distribution over i and mean and standard deviation statistics for a Gaussian distribution over y add . Finally, having ← -λ θ t ( X t ) ∈ R ≥ 0 be the raw output of a neural network can cause optimization issues due to the optimum ← -λ ∗ t including a probability ratio which can take on very large values. Instead, we learn a component prediction network p θ 0 | t ( n 0 | X t ) that predicts the number of components in X 0 given X t . To convert this into ← -λ θ ( X t ) , we show in Proposition 3 how the optimum ← -λ ∗ t ( X t ) is an analytic function of the true p 0 | t ( n 0 | X t ) . We then plug p θ 0 | t ( n 0 | X t ) into Proposition 3 to obtain an approximation of ← -λ ∗ t ( X t ) .

Proposition 3. We have

where X t = ( n t , x t ) and p t | 0 ( n t +1 | n 0 ) and p t | 0 ( n t | n 0 ) are both easily calculable distributions from the forward dimension deletion process.

3.5 Sampling

To sample the generative process, we numerically integrate the learned backward jump diffusion process using time-step δ t . Intuitively, it is simply the standard continuous time diffusion sampling scheme [3] but at each timestep we check whether a jump has occurred and if it has, sample the new component and insert it at the chosen index as explained by Algorithm 1.

4 Related Work

Our method jointly generates both dimensions and state values during the generative process whereas prior approaches [8, 11] are forced to first sample the number of dimensions and then run the diffusion process in this fixed dimension. When diffusion guidance is applied to these unconditional models [14, 26], users need to pick by hand the number of dimensions independent of the conditioning information even though the number of dimensions can be correlated with the conditioning parameter.

Instead of automatically learning when and how many dimensions to add during the generative process, previous work focusing on images [27, 28] hand pick dimension jump points such that the

Algorithm 1: Sampling the Generative Process

resolution of images is increased during sampling and reaches a certain pre-defined desired resolution at the end of the generative process. Further, rather than using any equivalent of A θ t , the values for new dimensions are simply filled in with Gaussian noise. These approaches mainly focus on efficiency rather than flexible generation as we do here.

The first term in our learning objective in Proposition 2 corresponds to learning the continuous part of our process (the diffusion) and the second corresponds to learning the discrete part of our process (the jumps). The first term can be seen as a trans-dimensional extension of standard denoising score matching [20] whilst the second bears similarity to the discrete space ELBO derived in [29].

Finally, jump diffusions also have a long history of use in Bayesian inference, where one aims to draw samples from a trans-dimensional target posterior distribution based on an unnormalized version of its density [30]: an ergodic jump diffusion is designed which admits the target as the invariant distribution [30-32]. The invariant distribution is not preserved when time-discretizing the process. However, it was shown in [33, 34] how general jump proposals could be built and how this process could be 'Metropolized' to obtain a discrete-time Markov process admitting the correct invariant distribution, yielding the popular Reversible Jump Markov Chain Monte Carlo algorithm. Our setup differs significantly as we only have access to samples in the form of data, not an unnormalized target.

5 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: ( x, y, z ) coordinates, a one-hot encoded atom type, and an integer charge value. Bonds are inferred from inter-atomic distances. We use an EGNN [24] backbone with three heads to predict s θ t , p θ 0 | t ( n 0 | X t ) , and A θ t . We uniformly delete dimensions, K del ( i | n ) = 1 /n , and since a point cloud is permutation invariant, A θ t ( y add | X t ) need only predict new dimension values. We set → λ t to a constant except for t < 0 . 1 T , where we set → λ t< 0 . 1 T = 0 . This ensures that all dimensions are added with enough generation time remaining for the diffusion process to finalize all state values.

We visualize sampling from our learned generative process in Figure 2; note how the process jointly creates a suitable number of atoms whilst adjusting their positions and identities. Before moving on to apply diffusion guidance which is the focus of our experiments, we first verify our unconditional sample quality in Table 2 and find we perform comparably to the results reported in [8] which use an FDDM. We ablate our choice of → λ t by comparing with setting → λ t to a constant for all t and with setting → λ t = 0 for t < 0 . 9 T (rather than just for t < 0 . 1 T ). We find that the constant → λ t performs worse due to the occasional component being added late in the generation process without enough time for the diffusion process to finalize its value. We find the → λ t< 0 . 9 T = 0 setting to have satisfactory sample quality however this choice of → λ t introduces issues during diffusion guided generation as we see next. Finally, we ablate the parameterization of Proposition 3 by learning ← -λ θ t ( X t ) ∈ R directly as the output of a neural network head. We find that this reduces sample quality due to the more well-behaved nature of the target, p θ 0 | t ( n 0 | X t ) when using Proposition 3. We note pure autoregressive models perform significantly worse than diffusion based models as found in [8].

？

5.1.1 Trans-Dimensional Diffusion Guidance

We now apply diffusion guidance to our unconditional model in order to generate molecules that contain a certain number of desired atom types, e.g. 3 carbons or 1 oxygen and 2 nitrogens. The distribution of molecule sizes changes depending on these conditions. We generate molecules conditioned on these properties by using the reconstruction guided sampling approach introduced in [9]. This method augments the score s θ t ( X t ) such that it approximates ∇ x t log p t ( X t | y ) rather than ∇ x t log p t ( X t ) (where y is the conditioning information) by adding on a term approximating ∇ x t log p t ( y | X t ) with p t ( y | X t ) = ∑ n 0 ∫ x 0 p ( y | X 0 ) p 0 | t ( X 0 | X t )d x 0 . This guides x t such that it is consistent with y . Since ← -λ θ t ( X t ) has access to x t , it will cause n t to automatically also be consistent with y without the user needing to input any information on how the conditioning information relates to the size of the datapoints. We give further details on diffusion guidance in Appendix C.

We show our results in Table 3. In order to perform guidance on the FDDM baseline, we implement the model from [8] in continuous time and initialize the dimension from the empirically observed dimension distribution in the dataset. This accounts for the case of an end user attempting to guide a unconditional model with access to no further information. We find that TDDM produces samples whose dimensions much more accurately reflect the true conditional distribution of dimensions given the conditioning information. The → λ t< 0 . 9 T = 0 ablation on the other hand only marginally improves the dimension error over FDDM because all dimensions are added in the generative process at a time when X t is noisy and has little relation to the conditioning information. This highlights the necessity of allowing dimensions to be added throughout the generative process to gain the trans-dimensional diffusion guidance ability. The ablation with constant → λ t has increased dimension error over TDDM as we find that when → λ t > 0 for all t , ← -λ θ t can become very large when t is close to 0 when the model has perceived a lack of dimensions. This occasionally results in too many dimensions being added hence an increased dimension error. Not using the Proposition 3 parameterization also increases dimension error due to the increased difficulty in learning ← -λ θ t .

5.1.2 Trans-Dimensional Interpolation

Interpolations are a unique way of gaining insights into the effect of some conditioning parameter on a dataset of interest. To create an interpolation, a conditional generative model is first trained and then sampled with a sweep of the conditioning parameter but using fixed random noise [8]. The resulting series of synthetic datapoints share similar features due to the fixed random noise but vary in ways that are very informative as to the effect of the conditioning parameter. Attempting to interpolate with an FDDM is fundamentally limited because the entire interpolation occurs in the same dimension which is unrealistic when the conditioning parameter is heavily correlated with the dimension of the datapoint. We demonstrate this by following the setup of [8] who train a conditional FDDM conditioned on polarizability. Polariz-

50

60

70

80

90

Polarizability

5

10

15

20

25

30

Number of Atoms

Dataset

FDDM

TDDM

ability is the ability of a molecule's electron cloud to distort in response to an external electric field [38] with larger molecules tending to have higher polarizability. To enable us to perform a trans-dimensional interpolation, we also train a conditional version of our model conditioned on polarizability. An example interpolation with this model is shown in Figure 4. We find that indeed the size of the molecule increases with increasing polarizability, with some molecular substructures e.g. rings, being maintained across dimensions. We show how the dimension changes with polarizability during 3 interpolations in Figure 3. We find that these match the true dataset statistics much more accurately than interpolations using FDDM which first pick a dimension and carry out the entire interpolation in that fixed dimension.

+

x4.

r

5.2 Video

We finally demonstrate our model on a video modeling task. Specifically we model the RoboDesk dataset [39], a video benchmark to measure the applicability of video models for planning and control problems. The videos are renderings of a robotic arm [40] performing a variety of different tasks including opening drawers and moving objects. We first train an unconditional model on videos of varying length and then perform planning by applying diffusion guidance to generate videos conditioned on an initial starting frame and a final goal frame [41]. The planning problem is then reduced to 'filling in' the frames in between. Our trans-dimensional model automatically varies the number of in-filled frames during generation so that the final length of video matches the length of time the task should take, whereas the fixed dimension model relies on the unrealistic assumption that the length of time the task should take is known before generation.

We model videos at 32 × 32 resolution and with varying length from 2 to 35 frames. For the network backbone, we use a UNet adapted for video [42]. In contrast to molecular point clouds, our data is no longer permutation invariant hence A θ t ( y add , i | X t ) includes a prediction over the location to insert the new frame. Full experimental details are provided in Appendix D. We evaluate our approach on three planning tasks, holding stationary, sliding a door and pushing an object. An example generation conditioned on the first and last frame for the slide door task is shown in Figure 5, with

the model in-filling a plausible trajectory. We quantify our model's ability to generate videos of a length appropriate to the task in Table 4 finding on all three tasks we generate a more accurate length of video than FDDM which is forced to sample video lengths from the unconditional empirically observed length distribution in the training dataset.

6 Discussion

In this work, we highlighted the pitfalls of performing generative modeling on varying dimensional datasets when treating state values and dimensions completely separately. We instead proposed a trans-dimensional generative model that generates both state values and dimensions jointly during the generative process. We detailed how this process can be formalized with the time-reversal of a jump diffusion and derived a novel evidence lower bound training objective for learning the generative process from data. In our experiments, we found our trans-dimensional model to provide significantly better dimension generation performance for diffusion guidance and interpolations when conditioning on properties that are heavily correlated with the dimension of a datapoint. We believe our approach can further enable generative models to be applied in a wider variety of domains where previous restrictive fixed dimension assumptions have been unsuitable.

References

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. International Conference on Machine Learning , 2015.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 2020.

Yang Song, Jascha Sohl-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.

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 J aand Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems , 2022.

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 , 2022.

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.

Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J 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.

Emiel Hoogeboom, Vıctor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. International Conference on Machine Learning , 2022.

Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. Advances in Neural Information Processing Systems , 2022.

Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, and Tim Saliman. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02303 , 2022.

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.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat GANs on image synthesis. Advances in Neural Information Processing Systems , 2021.

Katherine Crowson. Clip guided diffusion. Web Demo, https://huggingface.co/spaces/EleutherAI/clip-guided-diffusion , 2021.

Hengtong Zhang and Tingyang Xu. Towards controllable diffusion models via reward-guided exploration. arXiv preprint arXiv:2304.07132 , 2023.

Chin-Wei Huang, Jae Hyun Lim, and Aaron C Courville. A variational perspective on diffusionbased generative models and score matching. Advances in Neural Information Processing Systems , 2021.

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. Advances in Neural Information Processing Systems , 2022.

Joe Benton, Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. From denoising diffusions to denoising Markov models. arXiv preprint arXiv:2211.03595 , 2022.

Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications , 1982.

Ulrich G Haussmann and Etienne Pardoux. Time reversal of diffusions. The Annals of Probability , 1986.

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computation , 2011.

Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in Neural Information Processing Systems , 2021.

Patrick Cheridito, Damir Filipovi´ c, and Marc Yor. Equivalent and absolutely continuous measure changes for jump-diffusion processes. Annals of applied probability , 2005.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems , 2017.

Vıctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. International Conference on Machine Learning , 2021.

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. Medical Image Computing and Computer-Assisted Intervention , 2015.

Tomer Weiss, Luca Cosmo, Eduardo Mayo Yanes, Sabyasachi Chakraborty, Alex M Bronstein, and Renana Gershoni-Poranne. Guided diffusion for inverse molecular design. ChemrXiv , 2023.

Bowen Jing, Gabriele Corso, Renato Berlinghieri, and Tommi Jaakkola. Subspace diffusion generative models. European Conference on Computer Vision , 2022.

Han Zhang, Ruili Feng, Zhantao Yang, Lianghua Huang, Yu Liu, Yifei Zhang, Yujun Shen, Deli Zhao, Jingren Zhou, and Fan Cheng. Dimensionality-varying diffusion process. arXiv preprint arXiv:2211.16032 , 2022.

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.

Ulf Grenander and Michael I Miller. Representations of knowledge in complex systems. Journal of the Royal Statistical Society: Series B (Methodological) , 1994.

David B Phillips and Adrian F M Smith. Bayesian model comparison via jump diffusions. Markov Chain Monte Carlo in Practice , 1995.

Michael I Miller, Ulf Grenander, Joseph A O'Sullivan, and Donald L Snyder. Automatic target recognition organized via jump-diffusion algorithms. IEEE Transactions on Image Processing , 1997.

Peter J Green. Reversible jump Markov chain Monte Carlo computation and Bayesian model determination. Biometrika , 1995.

Peter J Green. Trans-dimensional Markov chain Monte Carlo. Highly Structured Stochastic Systems , 2003.

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.

Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data , 2014.

Rdkit: Open-source cheminformatics; http://www.rdkit.org . Accessed 2023.

Eric V Anslyn and Dennis A Dougherty. Modern physical organic chemistry. University Science Books , 2005.

Stephen Tian, Chelsea Finn, and Jiajun Wu. A control-centric benchmark for video prediction. arXiv preprint arXiv:2304.13723 , 2023.

Harini Kannan, Danijar Hafner, Chelsea Finn, and Dumitru Erhan. Robodesk: A multi-task reinforcement learning benchmark. https://github.com/google-research/robodesk , 2021.

Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. International Conference on Machine Learning , 2022.

William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. Advances in Neural Information Processing Systems , 2022.

John L Kelley. General Topology . Courier Dover Publications, 2017.

Stewart N Ethier and Thomas G Kurtz. Markov processes: Characterization and convergence. John Wiley & Sons , 2009.

Giovanni Conforti and Christian Léonard. Time reversal of Markov processes with jumps under a finite entropy condition. Stochastic Processes and their Applications , 2022.

Clement Vignac and Pascal Frossard. Top-n: Equivariant set and graph generation without exchangeability. International Conference on Learning Representations , 2022.

Margaret Mitchell, Simone Wu, Andrew Zaldivar, Parker Barnes, Lucy Vasserman, Ben Hutchinson, Elena Spitzer, Inioluwa Deborah Raji, and Timnit Gebru. Model cards for model reporting. Proceedings of the Conference on Fairness, Accountability, and Transparency , 2019.

Appendix

This appendix is organized as follows. In Section A, we present proofs for all of our propositions. Section A.1 presents a rigorous definition of our forward process using a more specific notation. This is then used in Section A.2.1 to 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 Proposition 2 using 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 B we give more details on our objective and in Section C we detail how we apply diffusion guidance to our model. We give the full details for our experiments in Section D and finally, in Section E, we 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 a 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, we slightly change the point of view in order to derive our theoretical results. We extend the space R d to ˆ R d = R d ∪ {∞} using the one-point compactification of the space. We refer to [43] for details on this space. The point ∞ will be understood as a mask. For instance, let x 1 , x 2 , x 3 ∈ R d . Then X = ( x 1 , x 2 , x 3 ) ∈ ( ˆ R d ) N with N = 3 corresponds to a vector for which all components are observed whereas X ′ = ( x 1 , ∞ , x 3 ) ∈ ( ˆ R d ) N corresponds to 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 = ( ˆ R d ) N with d, N ∈ 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 →{ 0 , 1 } N which is given for any X = { x i } N ∈ X by

In other words, dim( X ) is a binary vector identifying the 'dimension' of the vector X , i.e. which frames are observed. Going back to our example X = ( x 1 , x 2 , x 3 ) ∈ ( ˆ R d ) N and X ′ = ( x 1 , ∞ , x 3 ) ∈ ( ˆ R d ) N , we have that dim( X ) = { 1 , 1 , 1 } and dim( X ′ ) = { 1 , 0 , 1 } . For any vector u ∈ { 0 , 1 } N we denote | u | = ∑ N i =1 u i , i.e. the active dimensions of u (or equivalently the nonmasked frames). For any X ∈ X and D ∈ { 0 , 1 } N , we denote X D = { X ′ i } N i =1 with X ′ i = X i if D i = 1 and X ′ i = ∞ if D i = 0 .

We denote C k b ( R d , R ) the set of functions which are k differentiable and bounded. Similarly, we denote C k b ( R d , R ) the set of functions which are k differentiable and compactly supported. The set C k 0 ( R d , R ) denotes the functions which are k differentiable and vanish when ∥ x ∥ → + ∞ . We note that f ∈ C( ˆ R d ) , if f ∈ C( R d ) and f -f ( ∞ ) ∈ C 0 ( R d ) and that f ∈ C k ( ˆ R d ) for any k ∈ N if the restriction of f to R d is in C k ( R d ) and f ∈ C( ˆ R d ) .

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 a stochastic process can be roughly defined as its 'probabilistic derivative'. More precisely, assume that a stochastic process ( X t ) t ≥ 0 admits a transition semigroup (P t ) t ≥ 0 , i.e. for any t ≥ 0 , A ∈ X and X ∈ X we have P ( X t ∈ A | X 0 = x ) = P t ( x, A ) , then the infinitesimal generator is defined as A ( f ) = lim t → 0 (P t ( f ) -f ) /t , for every f for which this quantity is well-defined.

Here, we start by introducing the infinitesimal generator of interest and give some intuition about its form. Then, we prove a time-reversal formula for this infinitesimal generator.

We consider b : R d → R d , α : { 0 , 1 } NM → R + . For any f ∈ C 2 ( X ) and X ∈ X we define

where M ∈ N , D 0 = { 1 } N , { ∆ j } M -1 j =0 ∈ N M such that ∑ M -1 j =0 ∆ j < N and for any j ∈ { 0 , . . . , M -1 } , D ∆ j j is the subset of { 0 , 1 } { 1 ,...,N } such that D j +1 ∈ D ∆ j j if and only if D j · D j +1 = D j +1 , where · is the pointwise multiplication operator, and | D j | = | D j +1 | + ∆ j . The condition D j · D j +1 = D j +1 means that the non-masked dimensions in D j +1 are also non-masked dimensions in D j . The condition | D j | = | D j +1 | +∆ j means that in order to go from D j to D j +1 , one needs to mask exactly ∆ j dimensions.

Therefore, a sequence { ∆ j } M -1 j =0 ∈ N M such that ∑ M -1 j =0 ∆ j < N can be interpreted as a sequence of drops in dimension. At the core level, we have that | D M | = N -∑ M -1 j =0 ∆ j . For instance if | D M | = 1 , we have that at the end of the process, only one dimension is considered.

We choose α such that ∑ D 1 ⊂ D ∆0 0 · · · ∑ D M ⊂ D ∆ M -1 M -1 α ( D 0 , . . . , D M ) = 1 . Therefore, α ( D 0 , . . . , D M ) corresponds to the probability to choose the dimension path D 0 →··· → D M .

The part X ↦→⟨ b ( X i ) , ∇ x i f ( X ) ⟩ + 1 2 ∆ x i f ( X ) 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 = D 0 →··· → D M (recall that D 0 = { 1 } N ), we define the jump kernel J P as follows. For any x ∈ X , we have J P ( X, d Y ) = ∑ M -1 i =0 δ D i (dim( X )) δ X D i +1 (d Y ) . This operator corresponds to the deletion operator introduced in Section 3.1 . Hence, for any dimension path P = D 0 →··· → D M , we can define the associated infinitesimal generator: for any f ∈ C 2 ( X ) and X ∈ X we define

We can define the following jump kernel

This corresponds to averaging the jump kernel over the different possible dimension paths. We have that for any f ∈ C 2 ( X ) and X ∈ X

In what follows, we assume that there exists a Markov process ( X t ) t ≥ 0 with infinitesimal generator A . In order to sample from ( X t ) t ≥ 0 , one choice is to first sample the dimension path P according to the probability α . Second sample from the Markov process associated with the infinitesimal generator A P . We can approximately sample from this process using the Lie-Trotter-Kato formula [44, Corollary 6.7, p.33].

Denote (P t ) t ≥ 0 the semigroup associated with A P , (Q t ) t ≥ 0 the semigroup associated with the continuous part of A P and (J t ) t ≥ 0 the semigroup associated with the jump part of A P . More precisely, we have that, (Q t ) t ≥ 0 is associated with A cont such that for any f ∈ C 2 ( X ) and X ∈ X

In addition, we have that, (Q t ) t ≥ 0 is associated with A P jump such that for any f ∈ C 2 ( X ) and X ∈ X

First, note that A cont corresponds to the infinitesimal generator of a classical diffusion on the components which are not set to ∞ . Hence, we can approximately sample from (Q t ) t ≥ 0 by sampling according to the Euler-Maruyama discretization of the associated diffusion, i.e. by setting

where Z is a Gaussian random variable.

Similarly, in order to sample from (J t ) t ≥ 0 , one should sample from the jump process defined as follows. On the interval [0 , τ ) , we have X t = X 0 . At time τ , we define X 1 ∼ J ( X 0 , · ) and repeat the procedure. In this case τ is defined as an exponential random variable with parameter 1 . For t > 0 small enough the probability that t > τ is of order t . Therefore, we sample from J , i.e. the deletion kernel, with probability t . Combining this approximation and (7), we get approximate samplers for (Q t ) t ≥ 0 and (J t ) t ≥ 0 . Under mild assumptions, the Lie-Trotter-Kato formula ensures that for any t ≥ 0

This justifies sampling according to Algorithm 1 (in the case of the forward process).

A.2 Proof of Proposition 1

For the proof of Proposition 1, we first provide a rigorous proof using the notation introduced in A.1. We then follow this with a 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 1

We are now going to derive the formula for the time-reversal of the transdimensional infinitesimal generator A , see (5). This corresponds to a rigorous proof of Proposition 1. We refer to Section A.2.2 for a more intuitive, albeit less-rigorous, proof. We start by introducing the kernel K P given for any dimension path D 0 →··· → D M , for any i ∈ { 0 , . . . , M -1 } , Y ∈ D i +1 and A ∈ X by

Note that this kernel is the same as the one considered in Proposition 1. It is well-defined under the following assumption.

Assumption 1. For any t > 0 and D ⊂ { 0 , 1 } N , we have that X t conditioned to dim( X t ) = D admits a density w.r.t. the | D | d -dimensional Lebesgue measure, denoted p t ( ·| dim( X t ) = D ) .

The following result will be key to establish the time-reversal formula.

Lemma 1. Assume A1 . Let A , B ∈ X . Let P be a dimension path D 0 →··· → D M with M ∈ N . Then, we have

Proof. Let A , B ∈ X . We have

which concludes the proof.

Lemma 1 shows that K P verifies the flux equation associated with J P . 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 are also going to consider the following assumption which ensures that the integration by part formula is valid.

Assumption 2. For any t > 0 and i ∈ { 1 , . . . , N } , X t admits a smooth density w.r.t. the Nd -dimensional Lebesgue measure denoted p t and we have that for any f, h ∈ C 2 b (( R d ) N ) for any u ∈ [0 , t ] and i ∈ { 1 , . . . , N }

The second assumption ensures that we can apply the backward Kolmogorov evolution equation.

2

Assumption 3. For any g ∈ C ( X ) and t > 0 , we have that for any u ∈ [0 , t ] and X ∈ X , ∂ u g ( u, X ) + A ( g )( u, X ) = 0 , where for any u ∈ [0 , t ] and X ∈ X , g ( u, X ) = E [ g ( X t ) | X u = X ] .

We refer to [19] for conditions under A2 and A3 are valid in the setting of diffusion processes.

Proposition 4. Assume A1 , A2 and A3 . Assume that there exists a Markov process ( X t ) t ≥ 0 solution of the martingale problem associated with (6) . Let T > 0 and consider ( Y t ) t ∈ [0 ,T ] = ( X T -t ) t ∈ [0 ,T ] . Then ( Y t ) t ∈ [0 ,T ] is solution to the martingale problem associated with R , where for any f ∈ C 2 ( X ) , t ∈ (0 , T ) and x ∈ X we have

Proof. Let f, g ∈ C 2 ( X ) . In what follows, we show that for any s, t ∈ [0 , T ] with t ≥ s

More precisely, we show that for any s, t ∈ [0 , T ] with t ≥ s

Let s, t ∈ [0 , T ] , with t ≥ s . Next, we denote for any u ∈ [0 , t ] and X ∈ X , g ( u, X ) = E [ g ( X t ) | X u = X ] . Using A3, we have that for any u ∈ [0 , t ] and X ∈ X , ∂ u g ( u, X ) + A ( g )( u, X ) = 0 , i.e. g satisfies the backward Kolmogorov equation. For any u ∈ [0 , t ] and X ∈ X , we have

Using A2, we have that for any u ∈ [0 , t ] and i ∈ { 1 , . . . , N }

In addition, we have that for any X ∈ X and u ∈ [0 , t ] , J ( X,fg ) -J ( X,g ) f ( X ) = ∫ X g ( u, Y )( f ( Y ) -f ( X )) J ( X, d Y ) . Using Lemma 1, we get

Therefore, using (8), (9) and (10), we have

Finally, we have

which concludes the proof.

A.2.2 Intuitive Proof of Proposition 1

We recall Proposition 1.

Proposition 1. The time reversal of a forward jump diffusion process given by drift - → b t , diffusion coefficient - → g t , rate → λ t ( n ) and transition kernel ∑ n i =1 K del ( i | n ) δ del ( X ,i ) ( Y ) is given by a jump diffusion process with drift ← -b ∗ t ( X ) , diffusion coefficient ← -g ∗ t , rate ← -λ ∗ t ( X ) and transition kernel ∫ y add ∑ n +1 i =1 A ∗ t ( y add , i | X ) δ ins ( X , y add ,i ) ( Y )d y add as defined below

Diffusion part. Using standard diffusion models arguments such as [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.

To find ← -λ ∗ t ( X ) , we sum and integrate both sides over m and y , with Y = ( m, y ) ,

̸

Now we use the fact that δ del ( Y ,i ) ( X ) is 0 for any m = n +1 ,

Now letting Y = ins ( X , y add , i ) ,

To find A ∗ t ( y add , i | X ) , we start from (11) and set Y = ins ( X , z add , j ) to get

By inspection, we see immediately that

With a re-labeling of z add and j we achieve the desired form

A.3 Proof of Proposition 2

In this section we prove Proposition 2 using the notation from the main paper by following the framework of [17]. We operate on a state space X = ⋃ N n =1 { n } × R nd . On this space the gradient operator ∇ : C ( X , R ) → C ( X , X ) is defined as ∇ f ( X ) = ∇ ( nd ) x f ( X ) where ∇ ( nd ) x is the standard gradient operator defined as C ( R nd , R ) → C ( R nd , R nd ) with respect to x ∈ R nd . 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 R nd with respect to a probability density defined on R nd i.e. ∫ X f ( X ) µ (d X ) = ∑ N n =1 ∫ x ∈ R nd f ( X ) p ( n ) p ( x | n )d x where, for A ⊂ R nd , ∫ ( n,A ) µ ( d X ) = ∫ x ∈ A p ( n ) p ( x | n )d x . We will write p ( X ) as shorthand for p ( n ) p ( x | n ) .

Following, [17], we start by augmenting our space with a 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 × R ≥ 0 . In the proof, we use the infinitesimal generators for the the forward and backward processes. An infinitesimal generator is defined as

and can be understood as a probabilistic version of a derivative. For our process on the augmented space S , our generators decompose as A = ∂ t + ˆ A t where ˆ A t operates only on 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 t = 0 to t = T , with the true time reversal of X following the forward process satisfying ( Y t ) t ∈ [0 ,T ] = ( X T -t ) t ∈ [0 ,T ] . Further, we will write - → g t as g t and ← -g t = g T -t as we do not learn g and this is the optimal relation from the time reversal. We define

Verifying Assumption 1. The first step in the proof is to verify Assumption 1 in [17]. Letting ν t ( X ) = p T -t ( X ) , we assume we can write ∂ t p t ( X ) = ˆ K ∗ t p t ( X ) in the form M ν + cν = 0 for some function c : S → R , where M is the generator of another auxiliary process on S and ˆ K ∗ t is the adjoint operator which satisfies ⟨ ˆ K ∗ t f, h ⟩ = ⟨ f, ˆ K t h ⟩ i.e.

We now find ˆ K ∗ t . We start by substituting in the form for ˆ K t ,

We first focus on the RHS terms corresponding to the diffusion part of the process

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

where on the last line we have relabelled X to Y and Y to X . Putting both re-arranged forms for the RHS together, we obtain

We therefore have

Now we re-write ∂ t p t ( X ) = ˆ K ∗ t p t ( x ) in the form M ν + cν = 0 . We start by re-arranging

Substituting in our form for ˆ K ∗ t we obtain

We define our auxiliary process to have generator M = ∂ t + ˆ M t with

which is a jump diffusion process with drift b t M , diffusion coefficient g t , rate λ t M and transition kernel K t M . Then if we have b t M = - ← -b θ T -t ,

and

Then we have (12) can be rewritten as

which is in the form M ( ν )( ¯ X ) + c ( ¯ X ) ν ( ¯ X ) = 0 if we let

Verifying Assumption 2. Now that we have verified Assumption 1, the second step in the proof is Assumption 2 from [17]. We assume there is a bounded measurable function α : S → (0 , ∞ ) such

that α M f = L ( fα ) -f L α for all functions f : X → R such that f ∈ D ( M ) and fα ∈ D ( L ) . Substituting in M and L we get

Since f does not depend on time, ∂ t f ( X ) = 0 and ∂ t ( fα t )( X ) = f ( X ) ∂ t α t ( X ) thus the ∂ t terms on the RHS also cancel out. Comparing terms on the LHS and RHS relating to the diffusion part of the process we obtain

Therefore, we get

Simplifying the above expression, we get

This is true for any f implying

This implies that α t ( X ) satisfies

Comparing terms from the LHS and RHS of (15) relating to the jump part of the process we obtain

Hence, we have

Recalling the definitions of λ t M ( X ) and K M ( Y | X ) , (13) and (14), we get

This equality is satisfied if α t ( X ) follows the following relation

̸

̸

We only require this relation to be satisfied for n = m because both - → K t ( Y | X ) and ← -K θ T -t ( X | Y ) are 0 for n = m . We note at this point, as in [17], that if we have α t ( X ) = 1 /p t ( X ) and ← -λ θ T -t and ← -K θ T -t ( X | Y ) equal to the true time-reversals, then both (16), and (17) are satisfied. However, α t ( X ) = 1 /p t ( X ) is not the only α t to satisfy these equations. (16) and (17) can be thought of as enforcing a certain parameterization of the generative process in terms of α t [17].

Concluding the proof. Nowfor the final part of the proof, we substitute our value for α into the I ISM loss from [17] which is equal to the negative of the evidence lower bound on E p data ( X 0 ) [log p θ 0 ( X 0 )] up to a constant independent of θ . Defining β t ( X t ) = 1 /α t ( X t ) , we have

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 diff t + ˆ L J t with

and

By comparison with the approach to find the adjoint ˆ K ∗ t , we also have ˆ L ∗ t = ˆ L diff ∗ t + ˆ L J ∗ t with

In addition, we get

Finally, I ISM becomes

where we have named the two terms corresponding to the diffusion and jump part of the process as I diff ISM , I J ISM respectively. For the diffusion part of the loss, we use the denoising form of the objective proven in Appendix E of [17] which is equivalent to I diff ISM up to a constant independent of θ

To simplify this expression, we first re-arrange ˆ L diff t ( h ) for some general function h : S → R .

Setting h = p t | 0 ( ·| X 0 ) α t ( · ) , our diffusion part of the loss becomes

We then directly parameterize ∇ log α t ( X t ) as -s θ t ( X t )

We now focus on the expectation within the integral to re-write it in an easy to calculate form

Now we note that we can re-write ∇ log p 0 ,t ( X 0 , X t ) using M t where M t is a mask variable M t ∈ { 0 , 1 } n 0 that is 0 for components of X 0 that have been deleted to get to X t and 1 for components that remain in X t .

Substituting this back in we get

Therefore, the diffusion part of I ISM can be written as

We now focus on the jump part of the loss I J ISM . We first substitute in ˆ L J t and ˆ L J ∗ t

Noting that β t ( X t ) = 1 /α t ( X t ) , we get

or swapping labels for X t and Y ,

̸

Substituting (19) into the second line and (20) into the first line of (18) and using the fact that - → K t ( X t | Y ) = 0 for n t = m , we obtain

Hence, we have

̸

This can be rewritten as

Therefore, we have

Finally, using the definition of the forward and backward kernels, i.e. - → K t ( Y | X t ) = ∑ n i =1 K del ( i | n ) δ del ( X ,i ) ( Y ) and ← -K θ T -t ( X t | Y ) = ∫ x add ∑ n i =1 A θ t ( x add , i | Y ) δ ins ( Y , x add ,i ) ( X t )d x add , we get

We get

Therefore, we have

Putting are expressions for I diff ISM and I J ISM together we obtain

We get that -I ISM gives us our evidence lower bound on E p data ( X 0 ) [log p θ 0 ( X 0 )] up to a constant that does not depend on θ . In the main text we have used a time notation such that the backward process runs backwards from t = T to t = 0 . To align with the notation of time used in the main text we change T -t to t on subscripts for ← -λ θ T -t and A θ T -t . We also will use the fact that → λ t ( X t ) depends only on the number of components in X t , → λ t ( X t ) = → λ t ( n t ) .

Tightness of the lower bound Now that we have derived the ELBO as in Proposition 2, we show that the maximizers of the ELBO are tight, i.e. 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 a 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 t ) t ≥ 0 and associated with an infinitesimal generator ˆ L its extension ( t, Y t ) t ≥ 0 is associated with the infinitesimal generator L , i.e. L = ∂ t + ˆ L . We also define the score-matching operator Φ given for any f for which it is defined by

We recall that according to [17, Equation (8)] and under [17, Assumption 1, Assumption2], we have

with v t = p T -t for any t ∈ [0 , T ] . We define the variational gap Gap as follows

In addition, using Itô Formula, we have that log v T ( Y T ) -log v 0 ( Y 0 ) = ∫ T 0 L ( v )d t . Assuming that E [ | log v T ( Y T ) -log v 0 ( Y 0 ) | ] < + ∞ , we get

In particular, using [17, Proposition 1], we get that Gap ≥ 0 and Gap = 0 if and only if β ∝ v . In addition, the ELBO is maximized if and only if β ∝ 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 ELBO is maximized at the time reversal [17], then we have that when our jump diffusion parameters match the time reversal, our variational gap is 0 .

Other approaches. Another way to derive the ELBO is to follow the steps of [15] directly, since [17] is a general framework extending this approach. The key formulae to derive the result and the ELBO is 1) a Feynman-Kac formula 2) a Girsanov formula. In the case of jump diffusions (with jump in R d ) a Girsanov formula has been established by [22]. Extending this result to one-point compactification space would allow us to prove directly Proposition 2 without having to rely on the general framework of [17].

A.4 Proof of Proposition 3

We start by recalling the form for the time reversal given in Proposition 1

We then introduce a marginalization over X 0

where ( n +1 , z ( X , y add , i )) = ins ( X , y add , i ) . Now using the fact the forward component deletion process does not depend on x 0 , only n 0 , we have p t | 0 ( n | X 0 ) = p t | 0 ( n | n 0 ) and p t | 0 ( n +1 | X 0 ) = p t | 0 ( n +1 | n 0 ) . Using this result, we get

We now focus on the probability ratio within the integral over x 0 . We will show that this ratio is 1 . We start with the numerator, introducing a marginalization over possible mask variables between X 0 and ( n +1 , z ) , denoted M ( n +1) with M ( n +1) having n +1 ones and n 0 -( n +1) zeros.

Now, for our forward process we have

where z is shorthand for z ( X , y add , i ) , z ( j ) is the vector in R d for the j th component of z and M ( n +1) ( X 0 ) j is the vector in R d corresponding to the component in X 0 corresponding to the j th one in the M ( n +1) mask. Integrating out y add we have

where M ( n +1) \ i denotes a mask variable obtained by setting the i th one of M ( n +1) to zero. Hence, we have

We now re-write the denominator from (21) introducing a marginalization over mask variables, M ( n )

We use the following recursion for the probabilities assigned to mask variables

Substituting this into (23) gives

By comparing with (22), we can see that

This shows that the probability ratio in (21) is 1. Therefore, we have

which concludes the proof.

p t | 0 ( n | n 0 ) can be analytically calculated when → λ t ( n ) is of a simple enough form. When → λ t ( n ) does not depend on n then the dimension deletion process simply becomes a time inhomogeneous Poisson process. Therefore, we would have

In our experiments we set → λ t ( n = 1) = 0 to stop the dimension deletion process when we reach a single component. If we have → λ t ( n ) = → λ t ( m ) for all n, m > 1 then we can still use the time inhomogeneous Poisson process formula for n > 1 and find the probability for n = 1 , p t | 0 ( n = 1 | n 0 ) by requiring p t | 0 ( n | n 0 ) to be a valid normalized distribution. Therefore, for the case that → λ t ( n ) = → λ t ( m ) for all n, m > 1 and → λ t ( n = 1) = 0 , we have

In cases where → λ t ( n ) depends on n not just for n = 1 , p t | 0 ( n | n 0 ) can become more difficult to calculate analytically. However, since the probability distributions are all 1-dimensional over n , it is very cheap to simply simulate the forward dimension deletion process many times and empirically estimate p t | 0 ( n | n 0 ) 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 L ( θ ) as a standalone object and determine the optimum values for s θ t , ← -λ θ t and A θ t directly. This is in order to gain intuition directly into the learning signal of L ( θ ) without needing to refer to stochastic process theory.

The definition of L ( θ ) as in the main text is

with the expectations taken over U ( t ; 0 , T ) p 0 ,t ( X 0 , X t , M t ) K del ( i | n t ) δ del ( X t ,i ) ( Y ) .

Continuous optimum. We start by analysing the objective for s θ t . This part of L ( θ ) can be written as

We now use the fact that the function that minimizes an L 2 regression problem min f E p ( x,y ) [ ∥ f ( x ) -y ∥ 2 ] is the conditional expectation of the target f ∗ ( x ) = E p ( y | x ) [ y ] . Therefore the optimum value for s θ t ( X t ) is

Therefore, the optimum value for s θ t ( X t ) is ∇ x t log p ( X t ) which is the value that gives ← -b t to be the time reversal of - → b t as stated in Proposition 1.

Jump rate optimum. The learning signal for ← -λ θ t comes from these two terms in L ( θ )

This expectation is maximized when for each test input Z and test time t , we have the following expression maximized

because p t ( Z ) is the probability Z gets drawn as a full sample from the forward process and ∑ n z +1 i = ∫ y add p t ( ins ( Z , y add , i )) K del ( i | n z +1)d y add is the probability that a sample one component bigger than Z gets drawn from the forward process and then a component is deleted to get to Z . Therefore the first probability is the probability that test input Z and test time t appear as the first term in (24) whereas the second probability is the probability that test input Z and test time t appear as the second term in (24).

We now use the fact that, for constants b and c ,

We therefore have the optimum ← -λ θ t ( Z ) as

which is the form for the time-reversal given in Proposition (1).

Jump kernel optimum. Finally, we analyse the part of L ( θ ) for learning A θ t ( x add t , i | Y ) ,

We now re-write the joint probability distribution that the inner expectation is taken with respect to,

with and

and x base t ∈ R ( n t -1) d referring to the n t -1 components of x t , that are not x add t i.e. X t = ins (( x base t , n t -1) , x add t , i ) . We then have

Therefore, the optimum A θ t ( x add t , i | Y ) which maximizes this part of L ( θ ) is

which is the same form as given in Proposition 1.

B Training Objective

We estimate our objective L ( θ ) by taking minibatches from the expectation U ( t ; 0 , T ) p 0 ,t ( X 0 , X t , M t ) K del ( i | n t ) δ del ( X t ,i ) ( Y ) . We first sample t ∼ U ( t ; 0 , T ) and then take samples from our dataset X 0 ∼ p data ( X 0 ) . In order to sample p t | 0 ( X t , M t | X 0 ) we need to both add noise, delete dimensions and sample a mask variable. Since the Gaussian noising process is isotropic, we can add a suitable amount of noise to all dimensions of X 0 and then delete dimensions of that noised full dimensional value. More specifically, we first sample ˜ X t = ( n 0 , ˜ x t ) with ˜ x t ∼ N (˜ x t ; √ α t x 0 , (1 -α t ) I n 0 d ) for α t = exp ( -∫ t 0 β ( s )d s ) using the analytic forward equations for the VP-SDE derived in [3]. Then we sample the number of dimensions to delete. This is simple to do when our rate function is independent of n except for the case when n = 1 at which it is zero. We simply sample a Poisson random variable with mean parameter ∫ t 0 → λ s d s and then clamp its value such that the maximum number of possible components that are deleted is n 0 -1 . This gives the appropriate distribution over n , p t | 0 ( n | n 0 ) as given in Section A.4. To sample which dimensions are deleted, we can sample K del ( i 1 | n 0 ) K del ( i 2 | n 0 -1) . . . K del ( i n 0 -n t | n t +1) from which we can create the mask M t and apply it to ˜ X t to obtain X t , X t = M t ( ˜ X t ) . When K del ( i | n ) = 1 /n this is especially simple to do by simply randomly permuting the components of ˜ X t , and then removing the final n 0 -n t components.

As is typically done in standard diffusion models, we parameterize s θ t in terms of a noise prediction network that predicts ϵ where x t = √ α t M t ( x 0 ) + √ 1 -α t ϵ , ϵ ∼ N (0 , I n t d ) . We then re-weight the score loss in time such that we have a uniform weighting in time rather than the 'likelihood weighting' with g 2 t [3, 21]. Our objective to learn s θ t then becomes

Further, by using the parameterization given in Proposition 3, we can directly supervise the value of p θ 0 | t ( n 0 | X t ) by adding an extra term to our objective. We can treat the learning of p θ 0 | t ( n 0 | X t ) as a standard prediction task where we aim to predict n 0 given access to X t . A standard objective for learning p θ 0 | t ( n 0 | X t ) is then the cross entropy

Our augmented objective then becomes

where the expectation is taken with respect to

where x t = √ α t M t ( x 0 ) + √ 1 -α t ϵ and γ is a loss weighting term for the cross entropy loss.

C Trans-Dimensional Diffusion Guidance

To guide an unconditionally trained model such that it generates datapoints consistent with conditioning information, we use the reconstruction guided sampling approach introduced in [9]. Our conditioning information will be the values for some of the components of X 0 , and thus the guidance should guide the generative process such that the rest of the components of the generated datapoint are consistent with those observed components. Following the notation of [9], we denote the observed components as x a ∈ R n a d and the components to be generated as x b ∈ R n b d . Our trained score function s θ t ( X t ) approximates ∇ x t log p t ( X t ) whereas we would like the score to approximate ∇ x t log p t ( X t | x a 0 ) . In order to do this, we will need to augment our unconditional score s θ t ( X t ) such that it incorporates the conditioning information.

We first focus on the dimensions of the score vector corresponding to x a . These can be calculated analytically from the forward process

with p t | 0 ( x a t | x a 0 , n t ) = N ( x a t ; √ α t x a 0 , (1 -α t ) I n a d ) . Note that we assume a correspondence between x a t and x a 0 . For example, in video if we condition on the first and last frame, we assume that the first and last frame of the current noisy x t correspond to x a 0 and guide them towards their observed values. For molecules, the point cloud is permutation invariant and so we can simply assume the first n a components of x t correspond to x a 0 and guide them to their observed values.

Now we analyse the dimensions of the score vector corresponding to x b . We split the score as

p ( x a 0 | X t ) is intractable to calculate directly and so, following [9], we approximate it with N ( x a 0 ; ˆ x θa 0 ( X t ) , 1 -α t α t I n a d ) where ˆ x θa 0 ( X t ) is a point estimate of x a 0 given from s θ t ( X t ) calculated as

where again we have assumed a correspondence between x a t and x a 0 . Our approximation for ∇ x b t log p ( x a 0 | X t ) is then

which can be calculated by differentiating through the score network s θ t .

We approximate ← -λ ∗ t ( X t | x a 0 ) and A ∗ t ( y add , i | X t , x a 0 ) , with their unconditional forms ← -λ θ t ( X t ) and A θ t ( y add , i | X t ) . We find this approximation still leads to valid generations because the guidance of the score network s θ t , results in X t containing the conditioning information which in turn leads to ← -λ θ t ( X t ) guiding the number of components in X t to be consistent with the conditioning information too as verified in our experiments. Further, any errors in the approximation for A θ t ( y add , i | X t ) are fixed by further applications of the guided score function, highlighting the benefits of our combined autoregressive and diffusion based approach.

D Experiment Details

Our code is available at https://github.com/andrew-cr/jump-diffusion

D.1 Molecules

D.1.1 Network Architecture

Backbone For our backbone network architecture, we used the EGNN used in [8]. This is a specially designed graph neural network applied to the point cloud treating it as a fully connected graph. A special equivariant update is used, operating only on distances between atoms. We refer to [8] for the specific details on the architecture. We used the same size network as used in [8]'s QM9 experiments, specifically there are 9 layers, with a hidden node feature size of 256 . The output of the EGNN is fed into a final output projection layer to give the score network output s θ t ( X t ) .

Component number prediction To obtain p θ 0 | t ( n 0 | X t ) , we take the embedding produced by the EGNN before the final output embedding layer and pass it through 8 transformer layers each consisting of a self-attention block and an MLP block applied channel wise. Our transformer model dimension is 128 and so we project the EGNN embedding output down to 128 before entering into the transformer layers. We then take the output of the transformer and take the average embedding over all nodes. This embedding is then passed through a final projection layer to give softmax logits over the p θ 0 | t ( n 0 | X t ) distribution.

Autoregressive Distribution Our A θ t ( y add , i | X t ) network has to predict the position and features for a new atom when it is added to the molecule. Since the point cloud is permutation invariant, we do not need to predict i and so we just need to parameterize A θ t ( y add | X t ) . We found the network to perform the best if the network first predicts the nearest atom to the new atom and then a vector from that atom to the location of the new atom. To achieve this, we first predict softmax logits for a distribution over the nearest atom by applying a projection to the embedding output from the previously described transformer block. During training, the output of this distribution can be directly supervised by a cross entropy loss. Given the nearest atom, we then need to predict the position and features of the new atom to add. We do this by passing in the embedding generated by the EGNN and original point cloud features into a new transformer block of the same size as that used for p θ 0 | t ( n 0 | X t ) . We also input the distances from the nearest atom to all other atoms in the molecule currently as an additional feature. To obtain the position of the new atom, we will take a weighted sum of all the vectors between the nearest atom and other atoms in the molecule. This is to make it easy for the network to create new atoms 'in plane' with existing atoms which is useful for e.g. completing rings that have to remain in the same plane. To calculate the weights for the vectors, we apply an output projection to the output of the transformer block. The new atom features (atom type and charge) are generated by a separate output projection from the transformer block. For the position and features, A θ t ( y add | X t ) outputs both a mean and a standard deviation for a Gaussian distribution. For the position distribution, we set the standard deviation to be isotropic to remain equivariant to rotations. In total our model has around 7.3 million parameters.

D.1.2 Training

We train our model for 1.3 million iterations at a batch size of 64 . We use the Adam optimizer with learning rate 0 . 00003 . We also keep a running exponential moving average of the network weights that is used during sampling as is standard for training diffusion models [2, 3, 16] with a decay parameter of 0 . 9999 . We train on the 100K molecules contained in the QM9 training split. We model hydrogens explicitly. Training a model requires approximately 7 days on a single GPU which was done on an Academic cluster.

In [8] the atom type is encoded as a one-hot vector and diffused as a continuous variable along with the positions and charge values for all atoms. They found that multiplying the one-hot vectors by 0 . 25 to boost performance by allowing the atom-type to be decided later on in the diffusion process. We instead multiply the one-hot vectors by 4 so that atom-type is decided early on in the diffusion process which improves our guided performance when conditioning on certain atom-types being

present. We found our model is robust to this change and achieves similar sample quality to [8] as shown in Table 2.

When deleting dimensions, we first shuffle the ordering of the nodes and then delete the final n 0 -n t nodes. The cross entropy loss weighting in (25) is set to 1 .

Following [8] we train our model to operate within the center of mass (CoM) zero subspace of possible molecule positions. The means, throughout the forward and backward process, the average position of an atom is 0 . In our transdimensional framework, this is achieved by first deleting any atoms required under the forward component deletion process. We then move the molecule such that its CoM is 0 . We then add CoM free noise such that the noisy molecule also has CoM = 0 . Our score model s θ t is parameterized through a noise prediction model ϵ θ t which is trained to predict the CoM free noise that was added. Therefore, our score network learns suitable directions to maintain the process on the CoM = 0 subspace. For the position prediction from A θ t ( y add | X t ) we train it to predict the new atom position from the current molecules reference frame. When the new atom is added, we then update all atom positions such that CoM = 0 is maintained.

D.1.3 Sampling

During sampling we found that adding corrector steps [3] improved sample quality. Intuitively, corrector steps form a process that has p t ( X ) as its stationary distribution rather than the process progressing toward p 0 ( 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-dimensional modeling by first simulating a possible insertion using ← -λ θ t and then simulating a possible deletion using → λ t . We describe our overall sampling algorithm in Algorithm 2.

D.1.4 Evaluation

Unconditional For our unconditional sampling evaluation, we start adding corrector steps when t < 0 . 1 T in the backward process and use 5 corrector steps without the corrector steps on the number of components. We set δ = 0 . 05 for t > 0 . 5 T and δ = 0 . 001 for t < 0 . 5 T such that the total number

0

5

10

15

20

25

30

Number of atoms

0.000

0.025

0.050

0.075

0.100

0.125

0.150

Density

Dataset

Samples

of network evaluations is 1000 . We show the distribution of sizes of molecules generated by our model in Figure 6 and show more unconditional samples in Figure 7. We find our model consistently generates realistic molecules and achieves a size distribution similar to the training dataset even though this is not explicitly trained and arises from sampling our backward rate ← -λ θ t . Since we are numerically integrating a continuous time process and approximating the true time reversal rate ← -λ ∗ t , some approximation error is expected. For this experiment, sampling all of our models and ablations takes approximately 2 GPU days on Nvidia 1080 Ti GPUs.

Conditional For evaluating applying conditional diffusion guidance to our model, we choose 10 conditioning tasks that each result in a different distribution of target dimensions. The task is to produce molecules that include at least a certain number of target atom types. We then guide the first set of atoms generated by the model to have these desired atom types. The tasks chosen are given in Table 5. Molecules in the training dataset that meet the conditions in each task have a different distribution of sizes. The tasks were chosen so that we have an approximately linearly increasing mean number of atoms for molecules that meet the condition. We also require that there are at least 100 examples of molecules that meet the condition within the training dataset.

For sampling when using conditional diffusion guidance, we use 3 corrector steps throughout the backward process with δ t = 0 . 001 . For these conditional tasks, we include the corrector steps on the number of components. We show the distribution of dimensions for each task from the training dataset and from our generated samples in Figure 8. Our metrics are calculated by first drawing 1000 samples for each conditioning task and then finding the Hellinger distance between the size distribution generated by our method (orange diagonal hashing in Figure 8) and the size distribution for molecules in the training dataset that match the conditions of the task (green no hashing in Figure 8). We find that indeed our model when guided by diffusion guidance can automatically produce a size distribution close to the ground truth size distribution found in the dataset for that conditioning value. We show samples generated by our conditionally guided model in Figure 9. We can see that our model can generate realistic molecules that include the required atom types and are of a suitable size. For this experiment, sampling all of our models and ablations takes approximately 13 GPU days on Nvidia 1080 Ti GPUs.

Interpolations For our interpolations experiments, we follow the set up of [8] who train a new model conditioned on the polarizability of molecules in the dataset. We train a conditional version of our model which can be achieved by simply adding in the polarizability as an additional feature input to our backbone network and re-using all the same hyperparameters. We show more examples of interpolations in Figure 10.

承

D.1.5 Ablations

For our main model, we set → λ t< 0 . 1 T = 0 to ensure that all dimensions are added with enough generation time remaining for the diffusion process to finalize all state values. To verify this setting, we compare its performance with → λ t< 0 . 03 T = 0 and → λ t< 0 . 3 T = 0 . We show our results in Table 6. We find that the → λ t< 0 . 03 T = 0 setting to generate reasonable sample quality but incur some extra dimension error due to the generative process sometimes observing a lack of dimensions near t = 0 and adding too many dimensions. We observed the same effect in the paper for when setting → λ t to be constant for all t in Table 3. Further, the setting → λ t< 0 . 3 T = 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 → λ t< 0 . 1 T = 0 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 a unique chemical graph amongst this set of 10000 as is done in [8]. We show our results in Table 7 and find our TDDM method to have slightly lower levels of uniqueness when compared to the fixed dimension diffusion model baseline.

0

10

20

30

0.0

0.1

0.2

0.3

0

10

20

30

0.0

0.2

0.4

0.6

0.8

1.0

0

10

20

30

0.0

0.2

0.4

0.6

0.8

0

10

20

30

0.0

0.2

0.4

0.6

0

10

20

30

0.0

0.1

0.2

0.3

0.4

0

10

20

30

0.0

0.1

0.2

0.3

0

10

20

30

0.0

0.1

0.2

0.3

0.4

0

10

20

30

0.0

0.1

0.2

0.3

0

10

20

30

0.0

0.1

0.2

0.3

0

10

20

30

0.0

0.1

0.2

0.3

Number of atoms

Density

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 [46], [8]. 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 it is found in [8] 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 7 as a comparison to the results presented in [8] 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.

D.2 Video

D.2.1 Dataset

We used the VP 2 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 a scripted task-specific policy operating in the RoboDesk environment [40]. They sample an action vector at every step during data generation by adding i.i.d. Gaussian

在

t

女

noise to 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 out the lower-quality trajectories sampled with noise standard deviation 0.2, and so use only the 17 500 videos (2500) per task with noise standard deviation 0.1. We convert these videos to 32 × 32 resolution and then, so that the data we train on has varying lengths, we create each training example by sampling a length l from a uniform distribution over { 2 , . . . , 35 } and then taking a random l -frame subset of the video.

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 a permutation-invariant 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 a full frame which is challenging for a single pass autoregressive network to generate. We design our forward process to overcome these challenges.

We define our forward process to delete frames in a 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 is deleted, we then start noising the next randomly chosen frame. This is so 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 t = 0 . Before this operation, we represent the video x as an ordered sequence of frames, x 0 = [ x 1 , x 2 , . . . , x n 0 ] . During shuffling, we sample a random permutation π of the integers 1 , . . . , n 0 . Then the frames are kept in the same order, but annotated with an index variable so that we have x 0 + = [( x (1) 0 + , π (1)) , ( x (2) 0 + , π (2)) , . . . , ( x ( n 0 ) 0 + , π ( n 0 ))] .

We will run the forward process from t = 0 to t = 100 N . We will set the forward rate such we delete down from n t to n t -1 at time ( N -n t +1)100 . This is achieved heuristically by setting

We can see that at time t = ( N -n t +1)100 we will quickly delete down from n t to n t -1 at which point → λ t ( n t ) will become 0 thus stopping deletion until the process arrives at the next multiple of 100 in time. When we hit a deletion event, we delete the frame from X t that has the current highest index variable π ( n ) . In other words

where we use x ( i ) t [2] to refer to the shuffle index variable for the i th current frame in x t .

We now provide an example progression of the forward deletion process. Assume we have n 0 = 4 , N = 5 and sample a permutation such that π (1) = 3 , π (2) = 2 , π (3) = 4 , and π (4) = 1 . Initially

the state is augmented to include the shuffle index. Then the forward process progresses from t = 0 to t = 500 with components being deleted in descending order of the shuffle index

In this example, due to 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 it 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 ∈ [1 , . . . , n t ] , we set the forward noising process such that p t | 0 ( x ( i ) t | x ( i ) 0 , M t ) = N ( x ( i ) t ; x ( i ) 0 , σ t ( x ( i ) t ) 2 ) where x ( i ) 0 is the clean frame corresponding to x ( i ) t as given by the mask M t and σ t ( x ( i ) t ) follows

where we again use x ( i ) t [2] for the shuffle index of component i . This is the VE-SDE from [3] applied to each frame in turn. We note that we only add noise to the state values on not the shuffle index itself.

The SDE parameters that result in the VE-SDE are - → b t = 0 and - → g t = √ 2 t -2( N -x ( i ) t [2])100 .

D.2.3 Sampling the Backward Process

When t 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 [16] to update the continuous dimensions in this case, and also a variation of their discretisation of t - specifically to update from e.g. t = 600 to t = 500 , we use their discretization of t as if the maximum value was 100 and then offset all values by 500 .

To invert the dimension deletion process, we can use Proposition 3 to derive our reverse dimension generation process. We re-write our parameterized ← -λ θ t using Proposition 3 as

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 a single sample n 0 ∼ p θ 0 | t ( n 0 | X t ) . If n 0 > n t then ← -λ θ t ( X t ) = ∞ . The new component will then be added at which point ← -λ θ t ( X t ) becomes 0 for the remainder of this block of time due to n t becoming n t +1 . If n 0 = n t then ← -λ θ t ( X t ) = 0 and no new component is added. ← -λ θ t ( X t ) will continue to be 0 for the remainder of the backward process once an opportunity to add a component is not used.

When a new frame is added, we use A θ t ( y add , i | X t ) to decide where the frame is added and its initial value. Since when we delete a frame it is fully noised, A θ t ( y add , i | X t ) can simply predict Gaussian noise for the new frame y add . However, A θ t ( y add , i | X t ) will still learn to predict a suitable location i 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 11.

D.2.4 Network Architecture

Our video diffusion network architecture is based on the U-net used by [42], which takes as input the index of each frame within the video, and uses the differences between these indices to control the interactions between frames via an attention mechanism. Since, during generation, we do not know the final position of each frame within the x 0 , we instead pass in its position within the ordered sequence x t .

One further difference is that, since we are perform 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 x ( i ) t based on σ t ( x ( i ) t ) instead of basing them on the global diffusion timestep t . Our timestep embedding and pre- and post-conditioning of network inputs/outputs are as suggested by [16], other than being done on a per-frame basis. The architecture from [42] with these changes applied then gives us our score network s θ t .

While it would be possible to train a 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 s θ t , and the rate and index prediction network modeling p θ 0 | t ( n 0 | X t ) and A θ t ( i | X t ) . The rate and index prediction network is similar to the first half of the score network, in that it uses all U-net blocks up to and including the middle one. We then flatten the 512 × 4 × 4 hidden state for each frame after this block such that, for an n t frame input, we obtain a n t × 8192 hidden state. These are fed through a 1D convolution with kernel size 2 and zero-padding of size 1 on each end, reducing the hidden state to ( n t +1) × 128 , which is in turn fed

through a ReLU activation function. This hidden state is then fed into three separate heads. One head maps it to the parameters of A θ t ( i | X t ) via a 1D convolution of kernel size 3. The output of size ( n t +1) is fed through a softmax to provide the categorical distribution A θ t ( i | X t ) . 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 p θ 0 | t ( n 0 | X t ) . Finally, the third head consists of a 1D convolution of kernel size 3 with 35 output channels. The ( n t +1) × 35 output is fed through a softmax to parameterize distributions over the number of frames that were deleted from X 0 which came before the first in x t , the number of frames from X 0 which were deleted between each pair of frames in x t , and the number deleted after the last frame in x t . 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 x t , while there are sometimes many clean frames. Since the cost of running our 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 t during training, we adapt the log-normal distribution suggested by [16] 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 a uniform distribution and simulating our proposed forward process. We then simply change the noise standard deviation for the noisy frame, replacing it 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, 1 N ∑ N -1 i =0 LN ( t -100 i ; -0 . 6 , 1 . 8 2 ) . Here, the mixture index i can be interpreted as 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 ( i ) t as a function of σ t ( x ( i ) t ) to account for the non-isotropic noise.

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 to be generated in the reverse process and then for the second frame to be generated and so on. We show our results in Table 8. 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 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 Figure 5.

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 a generic method.

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. In addition to 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 be accompanied with studies into their biases and potential ways they can be misused. Further, public releases of models must be accompanied with model cards [47] explaining the biases, limitations and intended uses of the model.

Peer reviews:
comments: 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.

STRENGTHS:
* The proposed methodology is elegant and theoretically sound.
* Many constructions like diffusion guidance extend naturally.
* 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?
* The robot arm example, while sufficient for the paper, is rather toy.

QUESTIONS:
Nothing more than the weaknesses section.

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

comments: 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.

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.

2. The theoretical contribution is valid and the experiment is thorough in supporting the proposed model.

WEAKNESSES:
None.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nLimitations are well discussed throughout the paper. The previously discussed weaknesses could be mentioned as well.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nI 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.\n\nLIMITATIONS:\nThe limitations and broader impacts are properly discussed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}