SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Christian Weilbach
Valentin De Bortoli<sup>3</sup>
Tom Rainforth<sup>1</sup>
$^{1}$ Department of Statistics, University of Oxford, UK
$^{2}$  Department of Computer Science, University of British Columbia, Vancouver, Canada
$^{3}$ CNRS ENS Ulm, Paris, France
{campbell, rainforth, doucet}@stats.ox.ac.uk
{wsgh, weilbach}@cs.ubc.ca
valentin.debortoli@gmail.com
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.
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].
Figure 1: Illustration of the jump diffusion generative process on videos and molecules. The generative process consists of two parts: a diffusion part which denoises the current set of frames/atoms and a jump part which adds on a suitable number of new frames/atoms such that the final generation is a clean synthetic datapoint of an appropriate size.
This approach to trans-dimensional generative modeling is fundamentally limited due to the complete separation of dimension generation and state value generation. This is exemplified in the common use case of conditional diffusion guidance. Here, an unconditional diffusion model is trained that end-users can then easily and cheaply condition on their task of interest through guiding the generative diffusion process [3, 12-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.
Standard continuous-time diffusion models [3, 15-17] define a forward diffusion process through a stochastic differential equation (SDE) where  $\mathbf{x}_0\sim p_{\mathrm{data}}$  and, for  $t > 0$
$$
\mathrm {d} \mathbf {x} _ {t} = \overrightarrow {\mathbf {b}} _ {t} (\mathbf {x} _ {t}) \mathrm {d} t + g _ {t} \mathrm {d} \mathbf {w} _ {t}, \tag {1}
$$
where  $\mathbf{x}_t\in \mathbb{R}^d$  is the current state,  $\overrightarrow{\mathbf{b}}_t:\mathbb{R}^d\to \mathbb{R}^d$  is the drift and  $g_{t}\in \mathbb{R}$  is the diffusion coefficient.  $\mathrm{dw}_t$  is a Brownian motion increment on  $\mathbb{R}^d$ . This SDE can be understood intuitively by noting that in each infinitesimal timestep, we move slightly in the direction of the drift  $\vec{\mathbf{b}}_t$  and inject a small amount of Gaussian noise governed by  $g_{t}$ . Let  $p_t(\mathbf{x}_t)$  denote the distribution of  $\mathbf{x}_t$  for the forward diffusion process (1) so that  $p_0(\mathbf{x}_0) = p_{\mathrm{data}}(\mathbf{x}_0)$ .  $\vec{\mathbf{b}}_t$  and  $g_{t}$  are set such that at time  $t = T$ ,  $p_T(\mathbf{x}_T)$  is close to  $p_{\mathrm{ref}}(\mathbf{x}_T) = \mathcal{N}(\mathbf{x}_T;0,I_d)$ ; e.g.  $\vec{\mathbf{b}}_t(\mathbf{x}_t) = -\frac{1}{2}\beta_t\mathbf{x}_t$ ,  $g_{t} = \sqrt{\beta_{t}}$  for  $\beta_{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(\mathbf{x}_T)$  to  $p_0(\mathbf{x}_0)$  and satisfies the following reverse time SDE
$$
\mathrm {d} \mathbf {x} _ {t} = \overleftarrow {\mathbf {b}} _ {t} (\mathbf {x} _ {t}) \mathrm {d} t + g _ {t} \mathrm {d} \hat {\mathbf {w}} _ {t},
$$
where  $\overleftarrow{\mathbf{b}}_t(\mathbf{x}_t) = \overrightarrow{\mathbf{b}}_t(\mathbf{x}_t) - g_t^2\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t)$ , dt is a negative infinitesimal time step and  $\mathrm{d}\hat{\mathbf{w}}_t$  is a Brownian motion increment when time flows backwards. Unfortunately, both the terminal
distribution,  $p_T(\mathbf{x}_T)$ , and the score,  $\nabla_{\mathbf{x}_t} \log p_t(\mathbf{x}_t)$ , are unknown in practice. A generative model is obtained by approximating  $p_T$  with  $p_{\mathrm{ref}}$  and learning an approximation  $s_t^\theta (\mathbf{x}_t)$  to  $\nabla_{\mathbf{x}_t} \log p_t(\mathbf{x}_t)$  typically using denoising score matching [20], i.e.
$$
\min  _ {\theta} \mathbb {E} _ {\mathcal {U} (t; 0, T) p _ {0, t} (\mathbf {x} _ {0}, \mathbf {x} _ {t})} [ \| s _ {t} ^ {\theta} (\mathbf {x} _ {t}) - \nabla_ {\mathbf {x} _ {t}} \log p _ {t | 0} (\mathbf {x} _ {t} | \mathbf {x} _ {0}) \| ^ {2} ]. \tag {2}
$$
For a flexible model class,  $s^{\theta}$ , we get  $s_t^\theta (\mathbf{x}_t) \approx \nabla_{\mathbf{x}_t}\log p_t(\mathbf{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  $\mathbf{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,  $\mathbf{x}$ , giving  $\mathbf{X} = (n,\mathbf{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,  $\mathbf{X} \in \mathcal{X} = \bigcup_{n=1}^{N} \{n\} \times \mathbb{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}(\mathbf{Y}|\mathbf{X})$ :  $\mathcal{X} \times \mathcal{X} \rightarrow \mathbb{R}_{\geq 0}$ . Letting  $\mathbf{Y} = (m,\mathbf{y})$ , the transition kernel satisfies  $\sum_{m}\int_{\mathbf{y}}K_{t}(m,\mathbf{y}|\mathbf{X})\mathrm{d}\mathbf{y} = 1$  and  $\int_{\mathbf{y}}K_{t}(m = n,\mathbf{y}|\mathbf{X})\mathrm{d}\mathbf{y} = 0$ . The rate at which jumps occur (jumps per unit time) is given by a rate function  $\lambda_t(\mathbf{X}) : \mathcal{X} \rightarrow \mathbb{R}_{\geq 0}$ . For an infinitesimal timestep  $\mathrm{d}t$ , the jump diffusion can be written as
$$
\textbf {J u m p} \quad \mathbf {X} _ {t} ^ {\prime} = \left\{ \begin{array}{l l} \mathbf {X} _ {t} & \text {w i t h p r o b a b i l i t y 1 - \lambda_ {t} (\mathbf {X} _ {t}) d t} \\ \mathbf {Y} \sim K _ {t} (\mathbf {Y} | \mathbf {X} _ {t}) & \text {w i t h p r o b a b i l i t y \lambda_ {t} (\mathbf {X} _ {t}) d t} \end{array} \right.
$$
$$
\text {D i f f u s i o n} \quad \mathbf {x} _ {t + \mathrm {d} t} = \mathbf {x} _ {t} ^ {\prime} + \mathbf {b} _ {t} \left(\mathbf {X} _ {t} ^ {\prime}\right) \mathrm {d} t + g _ {t} \mathrm {d} \mathbf {w} _ {t} \quad n _ {t + \mathrm {d} t} = n _ {t} ^ {\prime}
$$
with  $\mathbf{X}_t\triangleq (n_t,\mathbf{x}_t)$  and  $\mathbf{X}_{t + \mathrm{d}t}\triangleq (n_{t + \mathrm{d}t},\mathbf{x}_{t + \mathrm{d}t})$  and  $\mathrm{dw}_t$  being a Brownian motion increment on  $\mathbb{R}^{n_t^{\prime}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  $\vec{\mathbf{b}}_t(\mathbf{X}) = -\frac{1}{2}\beta_t\mathbf{x}$  and  $\vec{g}_t = \sqrt{\beta_t}$  with  $\beta_{t} \geq 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  $\vec{\lambda}_t(\mathbf{X})$ . To formalize the deletion, we need to introduce some more notation. We let  $K^{\mathrm{del}}(i|n)$  be a user-defined distribution over which component of the current state to delete. We also define  $\operatorname{del}:\mathcal{X}\times \mathbb{N}\to \mathcal{X}$  to be the deletion operator that deletes a specified component. Specifically,  $(n - 1,\mathbf{y}) = \operatorname{del}((n,\mathbf{x}),i)$  where  $\mathbf{y}\in \mathbb{R}^{(n - 1)d}$  has the same values as  $\mathbf{x}\in \mathbb{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  $\vec{K}_t(\mathbf{Y}|\mathbf{X}) = \sum_{i = 1}^n K^{\mathrm{del}}(i|n)\delta_{\mathrm{del}(\mathbf{X},i)}(\mathbf{Y})$ . We note that only one component is ever deleted at a time meaning  $\vec{K}_t(m,\mathbf{y}|\mathbf{X}) = 0$  for  $m\neq n - 1$ . Further, the choice of  $K^{\mathrm{del}}(i|n)$  will dictate the behaviour of the reverse generative process. If we set  $K^{\mathrm{del}}(i|n) = \mathbb{I}\{i = n\}$  then we only ever delete the final component and so in the reverse generative
Table 1: Summary of forward and parameterized backward processes
Direction	b_t	g_t	λ_t(X)	K_t(Y|X)
Forward	-1/2β_t x	√β_t	→_t(n)	∑i=1n K^del(i|n) δdel(x,i)(Y)
Backward	-1/2β_t x - β_t s_t^θ(X)	√β_t	←_t(X)	∫yadd ∑i=1n+1 A_t^θ(yadd, i|X) δins(X,yadd, i)(Y) dyadd
direction, datapoints are created additively, appending components onto the end of the current state. Alternatively, if we set  $K^{\mathrm{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 noisig process is simulated from  $t = 0$  to  $t = T$  and should be such that at time  $t = T$ , the marginal probability  $p_t(\mathbf{X})$  should be close to a reference measure  $p_{\mathrm{ref}}(\mathbf{X})$  that can be sampled from. We set  $p_{\mathrm{ref}}(\mathbf{X}) = \mathbb{I}\{n = 1\} \mathcal{N}(\mathbf{x};0,I_d)$  where  $\mathbb{I}\{n = 1\}$  is 1 when  $n = 1$  and 0 otherwise. To be close to  $p_{\mathrm{ref}}$ , for the jump part, we set  $\vec{\lambda}_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  $\vec{\lambda}_t$  to depend only on the current dimension  $\vec{\lambda}_t(\mathbf{X}) = \vec{\lambda}_t(n)$  with  $\vec{\lambda}_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  $\beta_t$  schedule [2, 3] so that we are close to  $\mathcal{N}(\mathbf{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  $\overleftarrow{\mathbf{b}}_t(\mathbf{X})$ , diffusion coefficient  $\overleftarrow{g}_t$ , rate  $\overleftarrow{\lambda}_t(\mathbf{X})$  and transition kernel  $\overleftarrow{K}_t(\mathbf{Y}|\mathbf{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  $\overleftarrow{K}_t(\mathbf{Y}|\mathbf{X})$ .  $\overleftarrow{K}_t(\mathbf{Y}|\mathbf{X})$  should undo the forward deletion operation. Since  $\overrightarrow{K}_t(\mathbf{Y}|\mathbf{X})$  chooses a component and then deletes it,  $\overleftarrow{K}_t(\mathbf{Y}|\mathbf{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  $\mathbf{y}^{\mathrm{add}} \in \mathbb{R}^d$ . The insertion operator is defined as  $\operatorname{ins}:\mathcal{X} \times \mathbb{R}^d \times \mathbb{N} \to \mathcal{X}$ . It takes in the current value  $\mathbf{X}$ , the new component  $\mathbf{y}^{\mathrm{add}}$  and an index  $i \in \{1,\dots,n+1\}$  and inserts  $\mathbf{y}^{\mathrm{add}}$  into  $\mathbf{X}$  at location  $i$  such that the resulting value  $\mathbf{Y} = \operatorname{ins}(\mathbf{X},\mathbf{y}^{\mathrm{add}},i)$  has  $\operatorname{del}(\mathbf{Y},i) = \mathbf{X}$ . We denote the joint conditional distribution over the newly added component and the index at which it is inserted as  $A_t(\mathbf{y}^{\mathrm{add}},i|\mathbf{X})$ . We therefore have  $\overleftarrow{K}_t(\mathbf{Y}|\mathbf{X}) = \int_{\mathbf{y}^{\mathrm{add}}} \sum_{i=1}^{n+1} A_t(\mathbf{y}^{\mathrm{add}},i|\mathbf{X}) \delta_{\operatorname{ins}(\mathbf{X},\mathbf{y}^{\mathrm{add}},i)}(\mathbf{Y}) \mathrm{d}\mathbf{y}^{\mathrm{add}}$ . Noting that only one component is ever added at a time, we have  $\overleftarrow{K}_t(m,\mathbf{y}|\mathbf{X}) = 0$  for  $m \neq n+1$ .
This backward process formalism can be seen as a unification of diffusion models with autoregressive models. The diffusion part  $\overleftarrow{\mathbf{b}}_t$  denotes the current set of components in parallel, whilst the autoregressive part  $A_{t}(\mathbf{y}^{\mathrm{add}},i|\mathbf{X})$  predicts a new component and its location.  $\overleftarrow{\lambda}_t(\mathbf{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  $\overleftarrow{\mathbf{b}}_t(\mathbf{X}),\overleftarrow{g}_t,\overleftarrow{\lambda}_t(\mathbf{X})$  and  $A_{t}(\mathbf{y}^{\mathrm{add}},i|\mathbf{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  $\vec{\mathbf{b}}_t$ , diffusion coefficient  $\vec{g}_t$ , rate  $\vec{\lambda}_t(n)$  and transition kernel  $\sum_{i=1}^{n} K^{\mathrm{del}}(i|n) \delta_{\mathrm{del}(\mathbf{X}, i)}(\mathbf{Y})$  is given by a jump diffusion process with drift  $\vec{\mathbf{b}}_t^*(\mathbf{X})$ , diffusion coefficient  $\vec{g}_t^*$ , rate  $\vec{\lambda}_t^*(\mathbf{X})$  and transition kernel
$\int_{\mathbf{y}^{\mathrm{add}}}\sum_{i = 1}^{n + 1}A_t^* (\mathbf{y}^{\mathrm{add}},i|\mathbf{X})\delta_{\mathrm{ins}(\mathbf{X},\mathbf{y}^{\mathrm{add}},i)}(\mathbf{Y})\mathrm{d}\mathbf{y}^{\mathrm{add}}$  as defined below
$$
\overleftarrow {\mathbf {b}} _ {t} ^ {*} (\mathbf {X}) = \overrightarrow {\mathbf {b}} _ {t} (\mathbf {X}) - \overrightarrow {g} _ {t} ^ {2} \nabla_ {\mathbf {x}} \log p _ {t} (\mathbf {X}), \overleftarrow {g} _ {t} ^ {*} = \overrightarrow {g} _ {t},
$$
$$
\overleftarrow {\lambda} _ {t} ^ {*} (\mathbf {X}) = \vec {\lambda} _ {t} (n + 1) \frac {\sum_ {i = 1} ^ {n + 1} K ^ {\mathrm {d e l}} (i | n + 1) \int_ {\mathbf {y} ^ {\mathrm {a d d}}} p _ {t} (\operatorname {i n s} (\mathbf {X} , \mathbf {y} ^ {\mathrm {a d d}} , i)) \mathrm {d} \mathbf {y} ^ {\mathrm {a d d}}}{p _ {t} (\mathbf {X})},
$$
$$
A _ {t} ^ {*} (\mathbf {y} ^ {\text {a d d}}, i | \mathbf {X}) \propto p _ {t} \left(\operatorname {i n s} \left(\mathbf {X}, \mathbf {y} ^ {\text {a d d}}, i\right)\right) K ^ {\text {d e l}} (i | n + 1).
$$
All proofs are given in Appendix A. The expressions for  $\overline{\mathbf{b}}_t^*$  and  $\overleftarrow{g}_t^*$  are the same as for a standard diffusion except for replacing  $\nabla_{\mathbf{x}}\log p_t(\mathbf{x})$  with  $\nabla_{\mathbf{x}}\log p_t(\mathbf{X}) = \nabla_{\mathbf{x}}\log p_t(\mathbf{x}|n)$  which is simply the score in the current dimension. The expression for  $\lambda_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  $\mathbf{X}$ . If this is higher than the raw probability at time  $t$  that the forward process is at  $\mathbf{X}$  (the denominator) then we should have high  $\overleftarrow{\lambda}_t^*$  because  $\mathbf{X}$  is likely the result of a deletion of a larger datapoint. Finally the optimum  $A_t^*(\mathbf{y}^{\mathrm{add}}, i|\mathbf{X})$  is simply the conditional distribution of  $\mathbf{y}^{\mathrm{add}}$  and  $i$  given  $\mathbf{X}$  when the joint distribution over  $\mathbf{y}^{\mathrm{add}}, i, \mathbf{X}$  is given by  $p_t(\mathrm{ins}(\mathbf{X}, \mathbf{y}^{\mathrm{add}}, i))K^{\mathrm{del}}(i|n + 1)$ .
3.3 Objective for Learning the Backward Process
The true  $\overleftarrow{\mathbf{b}}_t^*$ ,  $\overleftarrow{\lambda}_t^*$  and  $A_t^*$  are unknown so we need to learn approximations to them,  $\overleftarrow{\mathbf{b}}_t^\theta$ ,  $\overleftarrow{\lambda}_t^\theta$  and  $A_t^\theta$ . Following Proposition 1, we set  $\overleftarrow{\mathbf{b}}_t^\theta(\mathbf{X}) = \overrightarrow{\mathbf{b}}_t(\mathbf{X}) - \vec{g}_t^2 s_t^\theta(\mathbf{X})$  where  $s_t^\theta(\mathbf{X})$  approximates  $\nabla_{\mathbf{x}} \log p_t(\mathbf{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  $\mathbb{E}_{p_{\mathrm{data}}(\mathbf{x}_0)}[\log p_0^\theta (\mathbf{x}_0)]$  [21]. We derive here an equivalent loss to learn  $s_t^\theta ,\stackrel {\leftarrow}{\lambda}_t^\theta$  and  $A_{t}^{\theta}$  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^\theta (\mathbf{X}_t)$  will resemble denoising score matching (2) but instead involve the conditional score  $\nabla_{\mathbf{x}_t}\log p_{t|0}(\mathbf{X}_t|\mathbf{X}_0) = \nabla_{\mathbf{x}_t}\log p_{t|0}(\mathbf{x}_t|\mathbf{X}_0,n_t)$ . This is difficult to calculate directly due to a combinatorial sum over the different ways the components of  $\mathbf{X}_0$  can be deleted to get to  $\mathbf{X}_t$ . We avoid this problem by equivalently conditioning on a mask variable  $M_{t}\in \{0,1\}^{n_{0}}$  that is 0 for components of  $\mathbf{X}_0$  that have been deleted to get to  $\mathbf{X}_t$  and 1 for components that remain in  $\mathbf{X}_t$ . This makes our denoising score matching target easy to calculate:  $\nabla_{\mathbf{x}_t}\log p_{t|0}(\mathbf{x}_t|\mathbf{X}_0,n_t,M_t) = \frac{\sqrt{\alpha_tM_t(\mathbf{x}_0) - \mathbf{x}_t}}{1 - \alpha_t}$  where  $\alpha_{t} = \exp (-\int_{0}^{t}\beta (s)\mathrm{d}s)$  [3]. Here  $M_{t}(\mathbf{x}_{0})$  is the vector removing any components in  $\mathbf{x}_0$  for which  $M_{t}$  is 0, thus  $M_{t}(\mathbf{x}_{0})$  and  $\mathbf{x}_t$  have the same dimensionality. We now state our full objective.
Proposition 2. For the backward generative jump diffusion process starting at  $p_{\mathrm{ref}}(\mathbf{X}_T)$  and finishing at  $p_0^\theta (\mathbf{X}_0)$ , an evidence lower bound on the model log-likelihood  $\mathbb{E}_{\mathbf{x}_0}\sim p_{\mathrm{data}}[\log p_0^\theta (\mathbf{x}_0)]$  is given by
$$
\mathcal {L} (\theta) = - \frac {T}{2} \mathbb {E} \left[ g _ {t} ^ {2} \| s _ {t} ^ {\theta} (\mathbf {X} _ {t}) - \nabla_ {\mathbf {x} _ {t}} \log p _ {t | 0} (\mathbf {x} _ {t} | \mathbf {X} _ {0}, n _ {t}, M _ {t}) \| ^ {2} \right] + \tag {3}
$$
$$
T \mathbb {E} \left[ - \stackrel {\leftarrow} {\lambda} _ {t} ^ {\theta} (\mathbf {X} _ {t}) + \vec {\lambda} _ {t} (n _ {t}) \log \stackrel {\leftarrow} {\lambda} _ {t} ^ {\theta} (\mathbf {Y}) + \vec {\lambda} _ {t} (n _ {t}) \log A _ {t} ^ {\theta} (\mathbf {x} _ {t} ^ {\mathrm {a d d}}, i | \mathbf {Y}) \right] + C, \tag {4}
$$
where expectations are with respect to  $\mathcal{U}(t;0,T)p_{0,t}(\mathbf{X}_0,\mathbf{X}_t,M_t)K^{\mathrm{del}}(i|n_t)\delta_{\mathrm{del}(\mathbf{X}_{t,i})}(\mathbf{Y}),C$  is a constant term independent of  $\theta$  and  $\mathbf{X}_t = \operatorname {ins}(\mathbf{Y},\mathbf{x}_t^{\mathrm{add}},i)$ . This evidence lower bound is equal to the log-likelihood when  $\overleftarrow{\mathbf{b}}_t^\theta = \overleftarrow{\mathbf{b}}_t^*$ ,  $\overleftarrow{\lambda}_t^\theta = \overleftarrow{\lambda}_t^*$  and  $A_{t}^{\theta} = A_{t}^{*}$ .