SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We get that  $-\mathcal{I}_{\mathrm{ISM}}$  gives us our evidence lower bound on  $\mathbb{E}_{p_{\mathrm{data}}(\mathbf{X}_0)}[\log p_0^\theta (\mathbf{X}_0)]$  up to a constant that does not depend on  $\theta$ . 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  $\overleftarrow{\lambda}_{T - t}^{\theta}$  and  $A_{T - t}^{\theta}$ . We also will use the fact that  $\vec{\lambda}_t(\mathbf{X}_t)$  depends only on the number of components in  $\mathbf{X}_t$ ,  $\vec{\lambda}_t(\mathbf{X}_t) = \vec{\lambda}_t(n_t)$ .
$$
\begin{array}{l} \mathcal {L} (\theta) = - \frac {T}{2} \mathbb {E} [ g _ {t} ^ {2} \| s _ {t} ^ {\theta} (\mathbf {X} _ {t}) - \nabla_ {\mathbf {x} _ {t}} \log p _ {t | 0} (\mathbf {x} _ {t} | n _ {t}, \mathbf {X} _ {0}, M _ {t}) \| ^ {2} ] + \\ T \mathbb {E} \left[ - \overleftarrow {\lambda} _ {t} ^ {\theta} (\mathbf {X} _ {t}) + \overrightarrow {\lambda} _ {t} (n _ {t}) \log \overleftarrow {\lambda} _ {t} ^ {\theta} (\mathbf {Y}) + \overrightarrow {\lambda} _ {t} (n _ {t}) \log A _ {t} ^ {\theta} (\mathbf {x} ^ {\text {a d d}}, i | \mathbf {Y}) \right] + \text {c o n s t .} \\ \end{array}
$$
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  $(\mathbf{Y}_t)_{t\geq 0}$  and associated with an infinitesimal generator  $\hat{\mathcal{L}}$  its extension  $(t,\mathbf{Y}_t)_{t\geq 0}$  is associated with the infinitesimal generator  $\mathcal{L}$ , i.e.  $\mathcal{L} = \partial_t + \hat{\mathcal{L}}$ . We also define the score-matching operator  $\Phi$  given for any  $f$  for which it is defined by
$$
\Phi (f) = \mathcal {L} (f) / f - \mathcal {L} (\log (f)).
$$
We recall that according to [17, Equation (8)] and under [17, Assumption 1, Assumption2], we have
$$
\log p _ {T} (\mathbf {Y} _ {0}) \geq \mathbb {E} \left[ \log p _ {0} (\mathbf {Y} _ {T}) - \int_ {0} ^ {T} \mathcal {L} (v / \beta) / (v / \beta) + \mathcal {L} (\log \beta) d t \right],
$$
with  $v_{t} = p_{T - t}$  for any  $t\in [0,T]$ . We define the variational gap Gap as follows
$$
\operatorname {G a p} = \mathbb {E} \left[ \log p _ {T} \left(\mathbf {Y} _ {0}\right) - \log p _ {0} \left(\mathbf {Y} _ {T}\right) + \int_ {0} ^ {T} \mathcal {L} (v / \beta) / (v / \beta) + \mathcal {L} (\log \beta) \mathrm {d} t \right].
$$
In addition, using Ito Formula, we have that  $\log v_{T}(\mathbf{Y}_{T}) - \log v_{0}(\mathbf{Y}_{0}) = \int_{0}^{T}\mathcal{L}(v)\mathrm{d}t$ . Assuming that  $\mathbb{E}[|\log v_T(\mathbf{Y}_T) - \log v_0(\mathbf{Y}_0)|] < +\infty$ , we get
$$
\mathrm {G a p} = \mathbb {E} \left[ \int_ {0} ^ {T} - \mathcal {L} (\log v) + \mathcal {L} (v / \beta) / (v / \beta) + \mathcal {L} (\log \beta) \mathrm {d} t \right] = \mathbb {E} \left[ \int_ {0} ^ {T} \Phi (v / \beta) \mathrm {d} t \right].
$$
In particular, using [17, Proposition 1], we get that  $\mathrm{Gap} \geq 0$  and  $\mathrm{Gap} = 0$  if and only if  $\beta \propto v$ . In addition, the ELBO is maximized if and only if  $\beta \propto 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  $\mathbb{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].
We start by recalling the form for the time reversal given in Proposition 1
$$
\overleftarrow {\lambda} _ {t} ^ {*} (\mathbf {X}) = \overrightarrow {\lambda} _ {t} (n + 1) \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}).
$$
We then introduce a marginalization over  $\mathbf{X}_0$
$$
\begin{array}{l} \stackrel {\leftarrow} {\lambda} _ {t} ^ {*} (\mathbf {X}) = \vec {\lambda} _ {t} (n + 1) \sum_ {i = 1} ^ {n + 1} K ^ {\mathrm {d e l}} (i | n + 1) \int_ {\mathbf {y} ^ {\mathrm {a d d}}} \sum_ {n _ {0}} \int_ {\mathbf {x} _ {0}} p _ {0, t} (\mathbf {X} _ {0}, \mathrm {i n s} (\mathbf {X}, \mathbf {y} ^ {\mathrm {a d d}}, i)) \mathrm {d} \mathbf {x} _ {0} \mathrm {d} \mathbf {y} ^ {\mathrm {a d d}} / p _ {t} (\mathbf {X}) \\ = \vec {\lambda} _ {t} (n + 1) \sum_ {i = 1} ^ {n + 1} K ^ {\mathrm {d e l}} (i | n + 1) \int_ {\mathbf {y} ^ {\mathrm {a d d}}} \sum_ {n _ {0}} \int_ {\mathbf {x} _ {0}} \frac {p _ {0} (\mathbf {X} _ {0})}{p _ {t} (\mathbf {X})} p _ {t | 0} (\operatorname {i n s} (\mathbf {X}, \mathbf {y} ^ {\mathrm {a d d}}, i) | \mathbf {X} _ {0}) \mathrm {d} \mathbf {x} _ {0} \mathrm {d} \mathbf {y} ^ {\mathrm {a d d}} \\ = \vec {\lambda} _ {t} (n + 1) \sum_ {i = 1} ^ {n + 1} K ^ {\mathrm {d e l}} (i | n + 1) \int_ {\mathbf {y} ^ {\mathrm {a d d}}} \sum_ {n _ {0}} \int_ {\mathbf {x} _ {0}} \frac {p _ {0 | t} (\mathbf {X} _ {0} | \mathbf {X})}{p _ {t | 0} (\mathbf {X} | \mathbf {X} _ {0})} p _ {t | 0} (\operatorname {i n s} (\mathbf {X}, \mathbf {y} ^ {\mathrm {a d d}}, i) | \mathbf {X} _ {0}) \mathrm {d} \mathbf {x} _ {0} \mathrm {d} \mathbf {y} ^ {\mathrm {a d d}} \\ = \overrightarrow {\lambda} _ {t} (n + 1) \sum_ {i = 1} ^ {n + 1} K ^ {\mathrm {d e l}} (i | n + 1) \int_ {\mathbf {y} ^ {\mathrm {a d d}}} \sum_ {n _ {0}} \int_ {\mathbf {x} _ {0}} \frac {p _ {0 | t} (n _ {0} | \mathbf {X}) p _ {0 | t} (\mathbf {x} _ {0} | \mathbf {X} , n _ {0})}{p _ {t | 0} (n | \mathbf {X} _ {0}) p _ {t | 0} (\mathbf {x} | \mathbf {X} _ {0} , n)} \times \\ p _ {t | 0} (n + 1 | \mathbf {X} _ {0}) p _ {t | 0} (\mathbf {z} (\mathbf {X}, \mathbf {y} ^ {\mathrm {a d d}}, i) | \mathbf {X} _ {0}, n + 1) \mathrm {d} \mathbf {x} _ {0} \mathrm {d} \mathbf {y} ^ {\mathrm {a d d}} \\ \end{array}
$$
where  $(n + 1, \mathbf{z}(\mathbf{X}, \mathbf{y}^{\mathrm{add}}, i)) = \mathrm{ins}(\mathbf{X}, \mathbf{y}^{\mathrm{add}}, i)$ . Now using the fact the forward component deletion process does not depend on  $\mathbf{x}_0$ , only  $n_0$ , we have  $p_{t|0}(n|\mathbf{X}_0) = p_{t|0}(n|n_0)$  and  $p_{t|0}(n + 1|\mathbf{X}_0) = p_{t|0}(n + 1|n_0)$ . Using this result, we get
$$
\begin{array}{l} \stackrel {\leftarrow} {\lambda} _ {t} ^ {*} (\mathbf {X}) = \vec {\lambda} _ {t} (n + 1) \sum_ {n _ {0}} \left\{\frac {p _ {t | 0} (n + 1 | n _ {0})}{p _ {t | 0} (n | n _ {0})} p _ {0 | t} (n _ {0} | \mathbf {X}) \times \right. \\ \int_ {\mathbf {x} _ {0}} \frac {\sum_ {i = 1} ^ {n + 1} K ^ {\mathrm {d e l}} (i | n + 1) \int_ {\mathbf {y} ^ {\mathrm {a d d}}} p _ {t | 0} (\mathbf {z} (\mathbf {X} , \mathbf {y} ^ {\mathrm {a d d}}, i) | \mathbf {X} _ {0} , n + 1) \mathrm {d} \mathbf {y} ^ {\mathrm {a d d}}}{p _ {t | 0} (\mathbf {x} | \mathbf {X} _ {0} , n)} p _ {0 | t} (\mathbf {x} _ {0} | \mathbf {X}, n _ {0}) \mathrm {d} \mathbf {x} _ {0} \}. \tag {21} \\ \end{array}
$$
We now focus on the probability ratio within the integral over  $\mathbf{x}_0$ . We will show that this ratio is 1. We start with the numerator, introducing a marginalization over possible mask variables between  $\mathbf{X}_0$  and  $(n + 1,\mathbf{z})$ , denoted  $M^{(n + 1)}$  with  $M^{(n + 1)}$  having  $n + 1$  ones and  $n_0 - (n + 1)$  zeros.
$$
\begin{array}{l} \sum_ {i = 1} ^ {n + 1} K ^ {\mathrm {d e l}} (i | n + 1) \int_ {\mathbf {y} ^ {\mathrm {a d d}}} p _ {t | 0} (\mathbf {z} (\mathbf {X}, \mathbf {y} ^ {\mathrm {a d d}}, i) | \mathbf {X} _ {0}, n + 1) \mathrm {d} \mathbf {y} ^ {\mathrm {a d d}} \\ = \sum_ {i = 1} ^ {n + 1} K ^ {\mathrm {d e l}} (i | n + 1) \sum_ {M ^ {(n + 1)}} \int_ {\mathbf {y} ^ {\mathrm {a d d}}} p _ {t | 0} \left(M ^ {(n + 1)}, \mathbf {z} (\mathbf {X}, \mathbf {y} ^ {\mathrm {a d d}}, i) \mid \mathbf {X} _ {0}, n + 1\right) \mathrm {d} \mathbf {y} ^ {\mathrm {a d d}} \\ = \sum_ {M ^ {(n + 1)}} \sum_ {i = 1} ^ {n + 1} K ^ {\mathrm {d e l}} (i | n + 1) p _ {t | 0} \left(M ^ {(n + 1)} \mid \mathbf {X} _ {0}, n + 1\right) \int_ {\mathbf {y} ^ {\mathrm {a d d}}} p _ {t | 0} (\mathbf {z} (\mathbf {X}, \mathbf {y} ^ {\mathrm {a d d}}, i) \mid \mathbf {X} _ {0}, n + 1, M ^ {(n + 1)}) \mathrm {d} \mathbf {y} ^ {\mathrm {a d d}} \\ \end{array}
$$
Now, for our forward process we have
$$
p _ {t | 0} (\mathbf {z} (\mathbf {X}, \mathbf {y} ^ {\mathrm {a d d}}, i) | \mathbf {X} _ {0}, n + 1, M ^ {(n + 1)}) = \prod_ {j = 1} ^ {n + 1} \mathcal {N} (\mathbf {z} ^ {(j)}; \sqrt {\alpha_ {t}} M ^ {(n + 1)} (\mathbf {X} _ {0}) ^ {j}, (1 - \alpha_ {t}) I _ {d})
$$
where  $\mathbf{z}$  is shorthand for  $\mathbf{z}(\mathbf{X},\mathbf{y}^{\mathrm{add}},i)$ ,  $\mathbf{z}^{(j)}$  is the vector in  $\mathbb{R}^d$  for the  $j$ th component of  $\mathbf{z}$  and  $M^{(n + 1)}(\mathbf{X}_0)^j$  is the vector in  $\mathbb{R}^d$  corresponding to the component in  $\mathbf{X}_0$  corresponding to the  $j$ th one in the  $M^{(n + 1)}$  mask. Integrating out  $\mathbf{y}^{\mathrm{add}}$  we have
$$
\int_ {\mathbf {y} ^ {\mathrm {a d d}}} p _ {t | 0} (\mathbf {z} (\mathbf {X}, \mathbf {y} ^ {\mathrm {a d d}}, i) | \mathbf {X} _ {0}, n + 1, M ^ {(n + 1)}) \mathrm {d} \mathbf {y} ^ {\mathrm {a d d}} = \prod_ {j = 1} ^ {n} \mathcal {N} (\mathbf {x} ^ {(j)}; \sqrt {\alpha_ {t}} M ^ {(n + 1) \backslash i} (\mathbf {X} _ {0}) ^ {j}, (1 - \alpha_ {t}) I _ {d}),
$$
where  $M^{(n + 1)\backslash i}$  denotes a mask variable obtained by setting the  $i$ th one of  $M^{(n + 1)}$  to zero. Hence, we have
$$
\begin{array}{l} \sum_ {i = 1} ^ {n + 1} K ^ {\mathrm {d e l}} (i | n + 1) \int_ {\mathbf {y} ^ {\mathrm {a d d}}} p _ {t | 0} (\mathbf {z} (\mathbf {X}, \mathbf {y} ^ {\mathrm {a d d}}, i) | \mathbf {X} _ {0}, n + 1) \mathrm {d} \mathbf {y} ^ {\mathrm {a d d}} \\ = \sum_ {M ^ {(n + 1)}} \sum_ {i = 1} ^ {n + 1} K ^ {\mathrm {d e l}} (i | n + 1) p _ {t | 0} \left(M ^ {(n + 1)} \mid \mathbf {X} _ {0}, n + 1\right) \\ \prod_ {j = 1} ^ {n} \mathcal {N} \left(\mathbf {x} ^ {(j)}; \sqrt {\alpha_ {t}} M ^ {(n + 1) \backslash i} \left(\mathbf {X} _ {0}\right) ^ {j}, (1 - \alpha_ {t}) I _ {d}\right). \tag {22} \\ \end{array}
$$
We now re-write the denominator from (21) introducing a marginalization over mask variables,  $M^{(n)}$
$$
p _ {t | 0} (\mathbf {x} \mid \mathbf {X} _ {0}, n) = \sum_ {M ^ {(n)}} p _ {t | 0} (M ^ {(n)} \mid \mathbf {X} _ {0}, n) p _ {t | 0} (\mathbf {x} \mid M ^ {(n)}, \mathbf {X} _ {0}, n). \tag {23}
$$
We use the following recursion for the probabilities assigned to mask variables
$$
p _ {t | 0} (M ^ {(n)} | \mathbf {X} _ {0}, n) = \sum_ {M ^ {(n + 1)}} \sum_ {i = 1} ^ {n + 1} \mathbb {I} \{M ^ {(n + 1) \backslash i} = M ^ {(n)} \} K ^ {\mathrm {d e l}} (i | n + 1) p _ {t | 0} (M ^ {(n + 1)} | \mathbf {X} _ {0}, n + 1).
$$
Substituting this into (23) gives
$$
\begin{array}{l} p _ {t | 0} (\mathbf {x} | \mathbf {X} _ {0}, n) = \sum_ {M ^ {(n)}} \sum_ {M ^ {(n + 1)}} \sum_ {i = 1} ^ {n + 1} \mathbb {I} \{M ^ {(n + 1) \backslash i} = M ^ {(n)} \} K ^ {\mathrm {d e l}} (i | n + 1) \times \\ p _ {t | 0} (M ^ {(n + 1)} | \mathbf {X} _ {0}, n + 1) p _ {t | 0} (\mathbf {x} | M ^ {(n)}, \mathbf {X} _ {0}, n) \\ = \sum_ {M ^ {(n)}} \sum_ {M ^ {(n + 1)}} \sum_ {i = 1} ^ {n + 1} \mathbb {I} \{M ^ {(n + 1) \backslash i} = M ^ {(n)} \} K ^ {\mathrm {d e l}} (i | n + 1) \\ \times p _ {t | 0} (M ^ {(n + 1)} | \mathbf {X} _ {0}, n + 1) \prod_ {j = 1} ^ {n} \mathcal {N} (\mathbf {x} ^ {(j)}; \sqrt {\alpha_ {t}} M ^ {(n)} (\mathbf {X} _ {0}) ^ {j}, (1 - \alpha_ {t}) I _ {d}) \\ = \sum_ {M ^ {(n + 1)}} \sum_ {i = 1} ^ {n + 1} K ^ {\mathrm {d e l}} (i | n + 1) p _ {t | 0} \left(M ^ {(n + 1)} | \mathbf {X} _ {0}, n + 1\right) \\ \times \prod_ {j = 1} ^ {n} \mathcal {N} (\mathbf {x} ^ {(j)}; \sqrt {\alpha_ {t}} M ^ {(n + 1) \backslash i} (\mathbf {X} _ {0}) ^ {j}, (1 - \alpha_ {t}) I _ {d}). \\ \end{array}
$$
By comparing with (22), we can see that
$$
p _ {t | 0} (\mathbf {x} | \mathbf {X} _ {0}, n) = \sum_ {i = 1} ^ {n + 1} K ^ {\mathrm {d e l}} (i | n + 1) \int_ {\mathbf {y} ^ {\mathrm {a d d}}} p _ {t | 0} (\mathbf {z} (\mathbf {X}, \mathbf {y} ^ {\mathrm {a d d}}, i) | \mathbf {X} _ {0}, n + 1) \mathrm {d} \mathbf {y} ^ {\mathrm {a d d}}.
$$
This shows that the probability ratio in (21) is 1. Therefore, we have
$$
\begin{array}{l} \stackrel {\leftarrow} {\lambda} _ {t} ^ {*} (\mathbf {X}) = \vec {\lambda} _ {t} (n + 1) \sum_ {n _ {0}} \left\{\frac {p _ {t | 0} (n + 1 | n _ {0})}{p _ {t | 0} (n | n _ {0})} p _ {0 | t} (n _ {0} | \mathbf {X}) \int_ {\mathbf {x} _ {0}} p _ {0 | t} (\mathbf {x} _ {0} | \mathbf {X}, n _ {0}) \mathrm {d} \mathbf {x} _ {0} \right\} \\ = \overrightarrow {\lambda} _ {t} (n + 1) \sum_ {n _ {0}} \frac {p _ {t | 0} (n + 1 | n _ {0})}{p _ {t | 0} (n | n _ {0})} p _ {0 | t} (n _ {0} | \mathbf {X}), \\ \end{array}
$$
$p_{t|0}(n|n_0)$  can be analytically calculated when  $\vec{\lambda}_t(n)$  is of a simple enough form. When  $\vec{\lambda}_t(n)$  does not depend on  $n$  then the dimension deletion process simply becomes a time inhomogeneous Poisson process. Therefore, we would have
$$
p _ {t | 0} (n | n _ {0}) = \frac {\left(\int_ {0} ^ {t} \vec {\lambda} _ {s} \mathrm {d} s\right) ^ {n _ {0} - n}}{(n _ {0} - n) !} \exp \left(- \int_ {0} ^ {t} \vec {\lambda} _ {s} \mathrm {d} s\right).
$$
In our experiments we set  $\vec{\lambda}_t(n = 1) = 0$  to stop the dimension deletion process when we reach a single component. If we have  $\vec{\lambda}_t(n) = \vec{\lambda}_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  $\vec{\lambda}_t(n) = \vec{\lambda}_t(m)$  for all  $n, m > 1$  and  $\vec{\lambda}_t(n = 1) = 0$ , we have
$$
p _ {t | 0} (n | n _ {0}) = \left\{ \begin{array}{l l} \frac {\left(\int_ {0} ^ {t} \vec {\lambda} _ {s} d s\right) ^ {n _ {0} - n}}{\left(n _ {0} - n\right) !} \exp \left(- \int_ {0} ^ {t} \vec {\lambda} _ {s} d s\right) & 1 <   n \leq n _ {0} \\ 1 - \sum_ {m = 2} ^ {n _ {0}} \frac {\left(\int_ {0} ^ {t} \vec {\lambda} _ {s} d s\right) ^ {n _ {0} - m}}{\left(n _ {0} - m\right) !} \exp \left(- \int_ {0} ^ {t} \vec {\lambda} _ {s} d s\right) & n = 1 \end{array} \right.
$$
In cases where  $\vec{\lambda}_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  $\mathcal{L}(\theta)$  as a standalone object and determine the optimum values for  $s_t^\theta$ ,  $\overleftarrow{\lambda}_t^\theta$  and  $A_t^\theta$  directly. This is in order to gain intuition directly into the learning signal of  $\mathcal{L}(\theta)$  without needing to refer to stochastic process theory.
The definition of  $\mathcal{L}(\theta)$  as in the main text is
$$
\begin{array}{l} \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] + \\ T \mathbb {E} [ - \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}) ] + C. \\ \end{array}
$$
with the expectations taken over  $\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})$
Continuous optimum. We start by analysing the objective for  $s_t^\theta$ . This part of  $\mathcal{L}(\theta)$  can be written as
$$
- \frac {1}{2} \int_ {0} ^ {T} g _ {t} ^ {2} \mathbb {E} _ {p _ {0, t} (\mathbf {X} _ {0}, \mathbf {X} _ {t}, M _ {t})} [ \| 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} ] \mathrm {d} t
$$
We now use the fact that the function that minimizes an  $L_{2}$  regression problem  $\min_{f} \mathbb{E}_{p(x,y)}[\| f(x) - y\|^2]$  is the conditional expectation of the target  $f^*(x) = \mathbb{E}_{p(y|x)}[y]$ . Therefore the optimum value for  $s_t^\theta(\mathbf{X}_t)$  is
$$
\begin{array}{l} s _ {t} ^ {*} (\mathbf {X} _ {t}) = \mathbb {E} _ {p (M _ {t}, \mathbf {X} _ {0} | \mathbf {X} _ {t})} [ \nabla_ {\mathbf {X} _ {t}} \log p _ {t | 0} (\mathbf {x} _ {t} | \mathbf {X} _ {0}, n _ {t}, M _ {t}) ] \\ = \sum_ {M _ {t}} \sum_ {n _ {0} = 1} ^ {N} \int_ {\mathbf {x} _ {0} \in \mathbb {R} ^ {n _ {0} d}} p (M _ {t}, n _ {0}, \mathbf {x} _ {0} | \mathbf {X} _ {t}) \nabla_ {\mathbf {x} _ {t}} \log p _ {t | 0} (\mathbf {x} _ {t} | \mathbf {x} _ {0}, n _ {0}, n _ {t}, M _ {t}) \mathrm {d} \mathbf {x} _ {0} \\ = \sum_ {M _ {t}} \sum_ {n _ {0} = 1} ^ {N} \int_ {\mathbf {x} _ {0} \in \mathbb {R} ^ {n _ {0} d}} \frac {p (M _ {t} , n _ {0} , \mathbf {x} _ {0} | \mathbf {X} _ {t})}{p _ {t | 0} (\mathbf {x} _ {t} | \mathbf {x} _ {0} , n _ {0} , n _ {t} , M _ {t})} \nabla_ {\mathbf {x} _ {t}} p _ {t | 0} (\mathbf {x} _ {t} | \mathbf {x} _ {0}, n _ {0}, n _ {t}, M _ {t}) \mathrm {d} \mathbf {x} _ {0} \\ = \sum_ {M _ {t}} \sum_ {n _ {0} = 1} ^ {N} \int_ {\mathbf {x} _ {0} \in \mathbb {R} ^ {n _ {0} d}} \frac {p (\mathbf {x} _ {0} , n _ {0} , n _ {t} , M _ {t})}{p (n _ {t} , \mathbf {x} _ {t})} \nabla_ {\mathbf {x} _ {t}} p _ {t | 0} (\mathbf {x} _ {t} | \mathbf {x} _ {0}, n _ {0}, n _ {t}, M _ {t}) \mathrm {d} \mathbf {x} _ {0} \\ = \frac {1}{p \left(n _ {t} , \mathbf {x} _ {t}\right)} \sum_ {M _ {t}} \sum_ {n _ {0} = 1} ^ {N} \int_ {\mathbf {x} _ {0} \in \mathbb {R} ^ {n _ {0} d}} \nabla_ {\mathbf {x} _ {t}} p (\mathbf {x} _ {t}, \mathbf {x} _ {0}, n _ {0}, n _ {t}, M _ {t}) \mathrm {d} \mathbf {x} _ {0} \\ = \frac {1}{p \left(n _ {t} , \mathbf {x} _ {t}\right)} \nabla_ {\mathbf {x} _ {t}} \sum_ {M _ {t}} \sum_ {n _ {0} = 1} ^ {N} \int_ {\mathbf {x} _ {0} \in \mathbb {R} ^ {n _ {0} d}} p \left(\mathbf {x} _ {t}, \mathbf {x} _ {0}, n _ {0}, n _ {t}, M _ {t}\right) \mathrm {d} \mathbf {x} _ {0} \\ = \frac {1}{p (n _ {t} , \mathbf {x} _ {t})} \nabla_ {\mathbf {x} _ {t}} p (\mathbf {x} _ {t}, n _ {t}) = \nabla_ {\mathbf {x} _ {t}} \log p (\mathbf {X} _ {t}). \\ \end{array}
$$
Therefore, the optimum value for  $s_t^\theta(\mathbf{X}_t)$  is  $\nabla_{\mathbf{x}_t} \log p(\mathbf{X}_t)$  which is the value that gives  $\overline{\mathbf{b}}_t$  to be the time reversal of  $\overrightarrow{\mathbf{b}}_t$  as stated in Proposition 1.
Jump rate optimum. The learning signal for  $\overleftarrow{\lambda}_t^\theta$  comes from these two terms in  $\mathcal{L}(\theta)$
$$
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}) \right] \tag {24}
$$
This expectation is maximized when for each test input  $\mathbf{Z}$  and test time  $t$ , we have the following expression maximized
$$
- p _ {t} (\mathbf {Z}) \stackrel {\leftarrow} {\lambda} _ {t} ^ {\theta} (\mathbf {Z}) + \sum_ {i =} ^ {n _ {z} + 1} \int_ {\mathbf {y} ^ {\text {a d d}}} p _ {t} \left(\operatorname {i n s} \left(\mathbf {Z}, \mathbf {y} ^ {\text {a d d}}, i\right)\right) K ^ {\text {d e l}} (i | n _ {z} + 1) \mathrm {d} \mathbf {y} ^ {\text {a d d}} \times \overrightarrow {\lambda} _ {t} \left(n _ {z} + 1\right) \log \stackrel {\leftarrow} {\lambda} _ {t} ^ {\theta} (\mathbf {Z}),
$$
because  $p_t(\mathbf{Z})$  is the probability  $\mathbf{Z}$  gets drawn as a full sample from the forward process and  $\sum_{i=1}^{n_z+1} \int_{\mathbf{y}^{\mathrm{add}}} p_t(\mathrm{ins}(\mathbf{Z}, \mathbf{y}^{\mathrm{add}}, i)) K^{\mathrm{del}}(i|n_z+1) \, \mathrm{d}\mathbf{y}^{\mathrm{add}}$  is the probability that a sample one component bigger than  $\mathbf{Z}$  gets drawn from the forward process and then a component is deleted to get to  $\mathbf{Z}$ . Therefore the first probability is the probability that test input  $\mathbf{Z}$  and test time  $t$  appear as the first term in (24) whereas the second probability is the probability that test input  $\mathbf{Z}$  and test time  $t$  appear as the second term in (24).
We now use the fact that, for constants  $b$  and  $c$ ,
$$
\operatorname {a r g m a x} _ {a} - b a + c \log a = \frac {c}{b}.
$$
We therefore have the optimum  $\overleftarrow{\lambda}_t^\theta (\mathbf{Z})$  as
$$
\overleftarrow {\lambda} _ {t} ^ {*} (\mathbf {Z}) = \overrightarrow {\lambda} _ {t} (n _ {z} + 1) \frac {\sum_ {i =} ^ {n _ {z} + 1} \int_ {\mathbf {y} ^ {\mathrm {a d d}}} p _ {t} \left(\operatorname {i n s} \left(\mathbf {Z} , \mathbf {y} ^ {\mathrm {a d d}} , i\right)\right) K ^ {\mathrm {d e l}} (i | n _ {z} + 1) \mathrm {d} \mathbf {y} ^ {\mathrm {a d d}}}{p _ {t} (\mathbf {Z})}
$$
which is the form for the time-reversal given in Proposition (1).
Jump kernel optimum. Finally, we analyse the part of  $\mathcal{L}(\theta)$  for learning  $A_t^\theta (\mathbf{x}_t^{\mathrm{add}},i|\mathbf{Y})$
$$
\begin{array}{l} T \mathbb {E} \left[ \vec {\lambda} _ {t} (n _ {t}) \log A _ {t} ^ {\theta} (\mathbf {x} _ {t} ^ {\mathrm {a d d}}, i | \mathbf {Y}) \right] \\ = \int_ {0} ^ {T} \mathbb {E} _ {p _ {t} (\mathbf {X} _ {t}) K ^ {\mathrm {d e l}} (i | n _ {t}) \delta_ {\mathrm {d e l} (\mathbf {X} _ {t}, i)} (\mathbf {Y})} [ \vec {\lambda} _ {t} (n _ {t}) \log A _ {t} ^ {\theta} (\mathbf {x} _ {t} ^ {\mathrm {a d d}}, i | \mathbf {Y}) ] \mathrm {d} t \\ = \int_ {0} ^ {T} \mathbb {E} _ {p _ {t} (n _ {t})} [ \vec {\lambda} _ {t} (n _ {t}) \mathbb {E} _ {p _ {t} (\mathbf {x} _ {t} | n _ {t}) K ^ {\mathrm {d e l}} (i | n _ {t}) \delta_ {\mathrm {d e l} (\mathbf {X} _ {t}, i)} (\mathbf {Y})} [ \log A _ {t} ^ {\theta} (\mathbf {x} _ {t} ^ {\mathrm {a d d}}, i | \mathbf {Y}) ] ] \mathrm {d} t. \\ \end{array}
$$
We now re-write the joint probability distribution that the inner expectation is taken with respect to,
$$
p _ {t} (\mathbf {x} _ {t} | n _ {t}) K ^ {\mathrm {d e l}} (i | n _ {t}) \delta_ {\mathrm {d e l} (\mathbf {X} _ {t}, i)} (\mathbf {Y}) = \tilde {p} (\mathbf {Y} | n _ {t}) p (\mathbf {x} _ {t} ^ {\mathrm {a d d}}, i | \mathbf {Y}) \delta_ {\mathbf {y}} (\mathbf {x} _ {t} ^ {\mathrm {b a s e}}).
$$
$$
\tilde {p} (\mathbf {Y} | n _ {t}) = \sum_ {i = 1} ^ {n _ {t}} \int_ {\mathbf {x} _ {t}} p _ {t} (\mathbf {x} _ {t} | n _ {t}) K ^ {\mathrm {d e l}} (i | n _ {t}) \delta_ {\mathrm {d e l} (\mathbf {x} _ {t}, i)} (\mathbf {Y}) \mathrm {d} \mathbf {x} _ {t},
$$
$$
p \left(\mathbf {x} _ {t} ^ {\text {a d d}}, i | \mathbf {Y}\right) \propto p _ {t} \left(\mathbf {x} _ {t} \mid n _ {t}\right) K ^ {\operatorname {d e l}} (i \mid n _ {t}),
$$
and  $\mathbf{x}_t^{\mathrm{base}}\in \mathbb{R}^{(n_t - 1)d}$  referring to the  $n_t - 1$  components of  $\mathbf{x}_t$ , that are not  $\mathbf{x}_t^{\mathrm{add}}$  i.e.  $\mathbf{X}_t = \mathrm{ins}((\mathbf{x}_t^{\mathrm{base}},n_t - 1),\mathbf{x}_t^{\mathrm{add}},i)$ . We then have