# Where to Diffuse, How to Diffuse, and How to Get Back: Automated Learning
for Multivariate Diffusions

Raghav Singhal∗,1, Mark Goldstein∗1, Rajesh Ranganath1,2
Courant Institute of Mathematical Sciences1, New York University
Center for Data Science2, New York University

###### Abstract

Diffusion-based generative models (dbgms) perturb data to a target noise
distribution and reverse this inference diffusion process to generate samples.
The choice of inference diffusion affects both likelihoods and sample quality.
For example, extending the inference process with auxiliary variables leads to
improved sample quality. While there are many such multivariate diffusions to
explore, each new one requires significant model-specific analysis, hindering
rapid prototyping and evaluation. In this work, we study Multivariate
Diffusion Models (mdms). For any number of auxiliary variables, we provide a
recipe for maximizing a lower-bound on the mdms likelihood without requiring
any model-specific analysis. We then demonstrate how to parameterize the
diffusion for a specified target noise distribution; these two points together
enable optimizing the inference diffusion process. Optimizing the diffusion
expands easy experimentation from just a few well-known processes to an
automatic search over all linear diffusions. To demonstrate these ideas, we
introduce two new specific diffusions as well as learn a diffusion process on
the mnist, cifar10, and imagenet32 datasets. We show learned mdms match or
surpass bits-per-dims (bpds) relative to fixed choices of diffusions for a
given dataset and model architecture.

##  1 Introduction

Diffusion-based generative models (dbgms) perturb data to a target noise
distribution and reverse this process to generate samples. They have achieved
impressive performance in image generation, editing, translation (Dhariwal &
Nichol, 2021; Nichol & Dhariwal, 2021; Sasaki et al., 2021; Ho et al., 2022),
conditional text-to-image tasks (Nichol et al., 2021; Ramesh et al., 2022;
Saharia et al., 2022) and music and audio generation (Chen et al., 2020; Kong
et al., 2020; Mittal et al., 2021). They are often trained by maximizing a
lower bound on the log likelihood, featuring an inference process interpreted
as gradually “noising” the data (Sohl-Dickstein et al., 2015; Ho et al.,
2020). ††∗ Equal Contribution. Correspondence to {rsinghal,goldstein} at
nyu.edu.

The choice of this inference process affects both likelihoods and sample
quality. On different datasets and models, different inference processes work
better; there is no universal best choice of inference, and the choice matters
(Song et al., 2020b).

While some work has improved performance by designing score model
architectures (Ho et al., 2020; Kingma et al., 2021; Dhariwal & Nichol, 2021),
Dockhorn et al. (2021) instead introduce the critically-damped langevin
diffusion (cld), showing that significant improvements in sample generation
can be gained by carefully designing new processes. cld pairs each data
dimension with an auxiliary “velocity” variable and diffuses them jointly
using second-order Langevin dynamics.

A natural question: if introducing new diffusions results in dramatic
performance gains, why are there only a handful of diffusions (variance-
preserving stochastic differential equation (vpsde), variance exploding (ve),
cld, sub-vpsde) used in dbgms? For instance, are there other auxiliary
variable diffusions that would lead to improvements like cld? This avenue
seems promising as auxiliary variables have improved other generative models
and inferences, such as normalizing flows (Huang et al., 2020), neural
ordinary differential equations (odes) (Dupont et al., 2019), hierarchical
variational models (Ranganath et al., 2016), ladder variational autoencoder
(Sønderby et al., 2016), among others.

Despite its success, cld also provides evidence that each new process requires
significant model-specific analysis. Deriving the evidence lower bound (elbo)
and training algorithm for diffusions is challenging (Huang et al., 2021;
Kingma et al., 2021; Song et al., 2021) and is carried out in a case-by-case
manner for new diffusions (Campbell et al., 2022). Auxiliary variables
seemingly complicate this process further; computing conditionals of the
inference process necessitates solving matrix Lyupanov equations (section
3.3). Deriving the inference stationary distribution—which helps the model and
inference match—can be intractable. These challenges limit rapid prototyping
and evaluation of new inference processes.

Concretely, training a diffusion model requires:

  1. (R1):

Selecting an inference and model process pair such that the inference process
converges to the model prior

  2. (R2):

Deriving the elbo for this pair

  3. (R3):

Estimating the elbo and its gradients by deriving and computing the inference
process’ transition kernel

In this work, we introduce Multivariate Diffusion Models (mdms) and a method
for training and evaluating them. mdms are diffusion-based generative models
trained with auxiliary variables. We provide a recipe for training mdms beyond
specific instantiations–like vpsde and cld—to all linear inference processes
that have a stationary distribution, with any number of auxiliary variables.

First, we bring results from gradient-based MCMC (Ma et al., 2015) to
diffusion modeling to construct mdms that converge to a chosen model prior
(R1); this tightens the elbo. Secondly, for any number of auxiliary variables,
we derive the mdm elbo (R2). Finally, we show that the transition kernel of
linear mdms, necessary for the elbo, can be computed automatically and
generically, for higher-dimensional auxiliary systems (R3).

With these tools, we explore a variety of new inference processes for
diffusion-based generative models. We then note that the automatic transitions
and fixed stationary distributions facilitate directly learning the inference
to maximize the mdm elbo. Learning turns diffusion model training into a
search not only over score models but also inference processes, at no extra
derivational cost.

##### Methodological Contributions.

In summary, our methodological contributions are:

  1. 1.

Deriving elbos for training and evaluating multivariate diffusion models
(mdms) with auxiliary variables.

  2. 2.

Showing that the diffusion transition covariance does not need to be manually
derived for each new diffusion. We instead demonstrate that a matrix
factorization technique, previously unused in diffusion models, can
automatically compute the covariance analytically for any linear mdm.

  3. 3.

Using results from gradient-based Markov chain Monte Carlo (mcmc) to construct
mdms with a complete parameterization of inference processes whose stationary
distribution matches the model prior.

  4. 4.

Combining the above into an algorithm called Automatic Multivariate Diffusion
Training (amdt) that enables training without diffusion-specific derivations.
amdt enables training score models for any linear diffusion, including
optimizing the diffusion and score jointly.

To demonstrate these ideas, we develop mdms with two specific diffusions as
well as learned multivariate diffusions. The specific diffusions are
accelerated Langevin diffusion (alda) (introduced in Mou et al. (2019) as a
higher-order scheme for gradient-based mcmc) and an alteration, modified
accelerated Langevin diffusion (malda). Previously, using these diffusions for
generative modeling would require significant model-specific analysis.
Instead, amdt for these diffusions is derivation-free.

##### Empirical contributions.

We train mdms on the mnist, imagenet32 and cifar-10 datasets. In the
experiments, we show that:

  1. 1.

Training new and existing fixed diffusions, such as alda and malda, is easy
with the proposed algorithm amdt.

  2. 2.

Using amdt to learn the choice of diffusion for the mdm matches or surpasses
the performance of fixed choices of diffusion process; sometimes the learned
diffusion and vpsde do best; other times the learned diffusion and cld do
best.

  3. 3.

There are new and existing mdms, trained and evaluated with the mdm elbo, that
account for as much performance improvement over vpsde as a three-fold
increase in score model size for a fixed univariate diffusion.

These findings affirm that the choice of diffusion affects the optimization
problem, and that learning the choice bypasses the process of choosing
diffusions for each new dataset and score architecture. We additionally show
the utility of the mdm elbo by showing on a dataset that cld achieves better
bits-per-dims (bpds) than previously reported with the probability flow ode
(Dockhorn et al., 2021).

##  2 Setup

We present diffusions by starting with the generative model and then
describing its likelihood lower bound (Sohl-Dickstein et al., 2015; Huang et
al., 2021; Kingma et al., 2021). Diffusions sample from a model prior
$\boldsymbol{\mathbf{z}}_{0}\sim\pi_{\theta}$ and then evolve a continuous-
time stochastic process $\boldsymbol{\mathbf{z}}_{t}\in\mathbb{R}^{d}$:

$\displaystyle
d\boldsymbol{\mathbf{z}}=h_{\theta}(\boldsymbol{\mathbf{z}},t)dt+\beta_{\theta}(t)d\boldsymbol{\mathbf{B}}_{t},\quad
t\in[0,T]$ (1)

where $\boldsymbol{\mathbf{B}}_{t}$ is a $d$-dimensionsal Brownian motion. The
model is trained so that $\boldsymbol{\mathbf{z}}_{T}$ approximates the data
$\boldsymbol{\mathbf{x}}\sim q_{\text{data}}$.111Following Huang et al.
(2021); Dockhorn et al. (2021) we integrate all processes in forward time $0$
to $T$. It may be helpful to think of an additional variable
$\hat{\boldsymbol{\mathbf{x}}}_{t}\triangleq\boldsymbol{\mathbf{z}}_{T-t}$ so
that $\hat{\boldsymbol{\mathbf{x}}}_{0}$ approximates
$\boldsymbol{\mathbf{x}}\sim q_{\text{data}}$. Maximum likelihood training of
diffusion models is intractable (Huang et al., 2021; Song et al., 2021; Kingma
et al., 2021). Instead, they are trained using a variational lower bound on
$\log p_{\theta}({\mathbf{z}}_{T}=x)$. The bound requires an inference process
$q_{\phi}(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{x}}=x)$:222We use
$\boldsymbol{\mathbf{y}}$ as the inference variable over the same space as the
model’s $\boldsymbol{\mathbf{z}}$.

$\displaystyle
d\boldsymbol{\mathbf{y}}=f_{\phi}(\boldsymbol{\mathbf{y}},s)ds+g_{\phi}(s)d\widehat{\boldsymbol{\mathbf{B}}}_{s},\quad
s\in[0,T]$ (2)

where $\widehat{\boldsymbol{\mathbf{B}}}_{s}$ is another Brownian motion
independent of $\boldsymbol{\mathbf{B}}_{t}$. The inference process is usually
taken to be specified rather than learned, and chosen to be i.i.d. for each
$y_{tj}$ conditional on each $x_{j}$. This leads to the interpretation of the
$y_{tj}$ as noisy versions of features $x_{j}$ (Ho et al., 2020). While the
diffusion elbo is challenging to derive in general, Huang et al. (2021); Song
et al. (2021) show that when the model process takes the form:

$\displaystyle
d{\mathbf{z}}=\left[g_{\phi}^{2}(T-t)s_{\theta}({\mathbf{z}},T-t)-f_{\phi}({\mathbf{z}},T-t)\right]dt+g_{\phi}(T-t)d{\mathbf{B}}_{t},$
(3)

the elbo is:

$\displaystyle\begin{split}\log
p_{\theta}(x)\geq\mathcal{L}^{\text{ism}}(x)=\mathbb{E}_{q_{\phi}(\boldsymbol{\mathbf{y}}|x)}\Bigg{[}\log\pi_{\theta}(\boldsymbol{\mathbf{y}}_{T})+\int_{0}^{T}-\frac{1}{2}\|s_{\theta}\|^{2}_{g^{2}_{\phi}}-\nabla\cdot(g_{\phi}^{2}s_{\theta}-f_{\phi})ds\Bigg{]},\end{split}$
(4)

where $f_{\phi},g_{\phi},s_{\theta}$ are evaluated at
$(\boldsymbol{\mathbf{y}}_{s},s)$,
$\|{\mathbf{x}}\|^{2}_{\mathbf{A}}={\mathbf{x}}^{\top}{\mathbf{A}}{\mathbf{x}}$
and $g^{2}=gg^{\top}$. Equation 4 features the Implicit Score Matching (ism)
loss (Song et al., 2020a), and can be re-written as an elbo
$\mathcal{L}^{\text{dsm}}$ featuring Denoising Score Matching (dsm) (Vincent,
2011; Song et al., 2020b), see section E.1.

##  3 A recipe for Multivariate Diffusion Models

As has been shown in prior work (Song et al., 2021; Dockhorn et al., 2021),
the choice of diffusion matters. Drawing on principles from previous
generative models (section 6), we can consider a wide class of diffusion
inference processes by constructing them using auxiliary variables.

At first glance, training such diffusions can seem challenging. First, one
needs an elbo that includes auxiliary variables. This elbo will require
sampling from the transition kernel, and setting the model prior to the
specified inference stationary distribution. But doing such diffusion-specific
analysis manually is challenging and hinders rapid prototyping.

In this section we show how to address these challenges and introduce an
algorithm, amdt, to simplify and automate modeling with mdms. amdt can be used
to train new and existing diffusions, including those with auxiliary
variables, and including those that learn the inference process.

###  3.1 Multivariate Model and Inference

For the $j^{th}$ data coordinate at each time $t$, mdms pair
$\boldsymbol{\mathbf{z}}_{tj}\in\mathbb{R}$ with a vector of auxiliary
variables $\boldsymbol{\mathbf{v}}_{tj}\in\mathbb{R}^{K-1}$ into a joint
vector $\boldsymbol{\mathbf{u}}_{t}$ and diffuse in the extended space:

$\displaystyle\boldsymbol{\mathbf{u}}_{0}\sim\pi_{\theta},\quad\quad
d\boldsymbol{\mathbf{u}}=h_{\theta}(\boldsymbol{\mathbf{u}}_{t}=\begin{bmatrix}\boldsymbol{\mathbf{z}}_{t}\\\
\boldsymbol{\mathbf{v}}_{t}\end{bmatrix},t)dt+\beta_{\theta}(t)d\boldsymbol{\mathbf{B}}_{t}.$
(5)

mdms model the data $\boldsymbol{\mathbf{x}}$ with
$\boldsymbol{\mathbf{z}}_{T}$, a coordinate in
$\boldsymbol{\mathbf{u}}_{T}\sim p_{\theta}$. For the $j^{th}$ feature
$\boldsymbol{\mathbf{x}}_{j}$, each
$\boldsymbol{\mathbf{u}}_{tj}\in\mathbb{R}^{K}$ consists of a “data” dimension
$\boldsymbol{\mathbf{u}}_{tj}^{z}$ and auxiliary variable
$\boldsymbol{\mathbf{u}}_{tj}^{v}$. Therefore
$\boldsymbol{\mathbf{u}}\in\mathbb{R}^{dK}$. We extend the drift coefficient
$h_{\theta}$ from a function in
$\mathbb{R}^{d}\times\mathbb{R}_{+}\rightarrow\mathbb{R}^{d}$ to the extended
space $\mathbb{R}^{dK}\times\mathbb{R}_{+}\rightarrow\mathbb{R}^{dK}$. We
likewise extend the diffusion coefficient to a matrix $\beta_{\theta}$ acting
on Brownian motion $\boldsymbol{\mathbf{B}}_{t}\in\mathbb{R}^{dK}$.

Because the mdm model is over the extended space, the inference distribution
$\boldsymbol{\mathbf{y}}$ must be too. We then set
$q(\boldsymbol{\mathbf{y}}_{0}^{v}|\boldsymbol{\mathbf{y}}_{0}^{z}=x)$ to any
chosen initial distribution, e.g.
$\mathcal{N}(\boldsymbol{\mathbf{0}},\boldsymbol{\mathbf{I}})$ and discuss
this choice in section 4. Then $\boldsymbol{\mathbf{y}}_{s}$ evolves according
to the auxiliary variable inference process:

$\displaystyle
d\boldsymbol{\mathbf{y}}=f_{\phi}(\boldsymbol{\mathbf{y}},s)ds+g_{\phi}(s)d\widehat{\boldsymbol{\mathbf{B}}}_{s},$
(6)

where the inference drift and diffusion coefficients $f_{\phi},g_{\phi}$ are
now over the extended space
$\boldsymbol{\mathbf{y}}=[\boldsymbol{\mathbf{y}}^{z},\boldsymbol{\mathbf{y}}^{v}]$.
The function $f_{\phi}$ lets the $z$ and $v$ coordinates of
$\boldsymbol{\mathbf{y}}_{tj}$ interact in the inference process.

### Assumptions

This work demonstrates how to parameterize time-varying Itô processes, used
for diffusion modeling, to have a stationary distribution that matches the
given model prior. To take advantage of the automatic transition kernels also
presented, the inferences considered for modeling are linear time-varying
processes and take the form:

$\displaystyle
d{\mathbf{y}}={\mathbf{A}}_{\phi}(s){\mathbf{y}}ds+g_{\phi}(s)d{\mathbf{B}}_{s}$

where ${\mathbf{A}}_{\phi}(s):\mathbb{R}_{+}\rightarrow dK\times dK$ and
$g_{\phi}(s):\mathbb{R}_{+}\rightarrow dK\times dK$ are matrix-valued
functions.

###  3.2 elbo for mdms

We now show how to train mdms to optimize a lower bound on the log likelihood
of the data. Like in the univariate case, we use the parameterization in eq. 3
to obtain a tractable elbo.

######  Theorem 1.

The mdm log marginal likelihood of the data is lower-bounded by:

$\displaystyle\begin{split}\log&p_{\theta}(x)\geq\mathbb{E}_{q_{\phi}(\boldsymbol{\mathbf{y}}|x)}\Bigg{[}\underbrace{\log\pi_{\theta}(\boldsymbol{\mathbf{y}}_{T})}_{\ell_{T}}-\int_{0}^{T}\frac{1}{2}\|s_{\theta}\|^{2}_{g^{2}_{\phi}}+\nabla\cdot(g_{\phi}^{2}s_{\theta}-f_{\phi})ds-\underbrace{\log
q_{\phi}(\boldsymbol{\mathbf{y}}_{0}^{v}|x)}_{\ell_{q}}\Bigg{]}\quad(\mathcal{L}^{\text{mism}})\\\
&=\mathbb{E}_{q_{\phi}(\boldsymbol{\mathbf{y}}|x)}\Bigg{[}\ell_{T}+\int_{0}^{T}\frac{1}{2}\|s_{q}\|^{2}_{g^{2}_{\phi}}-\frac{1}{2}\|s_{\theta}-s_{q}\|^{2}_{g^{2}_{\phi}}+(\nabla\cdot
f_{\phi})ds-\ell_{q}\Bigg{]}\quad(\mathcal{L}^{\text{mdsm}}).\end{split}$ (7)

where divergences and gradients are taken with respect to
$\boldsymbol{\mathbf{y}}_{s}$ and
$s_{q}=\nabla_{\boldsymbol{\mathbf{y}}_{s}}\log
q(\boldsymbol{\mathbf{y}}_{s}|x)$.

###### Proof.

The proof for the $\mathcal{L}^{\text{mism}}$ bound is in appendix E. In
short, we introduce auxiliary variables, apply Theorem 1 of Huang et al.
(2021) (equivalently, Theorem 3 of Song et al. (2021) or appendix E of Kingma
et al. (2021)) to the joint space, and then apply an additional variational
bound to ${\mathbf{v}}_{0}$. The $\mathcal{L}^{\text{mdsm}}$ bound follows by
analogy of the ism-dsm connection from Huang et al. (2021); Song et al. (2021)
for univariate diffusion models. ∎

We call the first bound the mdm ism bound $\mathcal{L}^{\text{mism}}$ and the
second the mdm dsm bound $\mathcal{L}^{\text{mdsm}}$. In practice, estimates
of $\mathcal{L}^{\text{mism}}$ can be computationally prohibitive. We train
mdm’s by estimating the gradients of $\mathcal{L}^{\text{mdsm}}$. For
numerical stability, the integral in eq. 7 is computed on $[\epsilon,T]$
rather than $[0,T]$. One can regard this as a bound for a variable
$\boldsymbol{\mathbf{u}}_{\epsilon}$. To maintain a proper likelihood bound
for the data, one can choose a likelihood
$\boldsymbol{\mathbf{u}}_{0}|\boldsymbol{\mathbf{u}}_{\epsilon}$ and compose
bounds as we demonstrate in appendix H. We report the elbo with this
likelihood term, which plays the same role as the discretized Gaussian in
Nichol & Dhariwal (2021) and Tweedie’s formula in Song et al. (2021).

###  3.3 Ingredient 1: Computing the transition
$q_{\phi}(\boldsymbol{\mathbf{y}}_{s}|x)$

To estimate eq. 7 and its gradients, we need samples from
$q(\boldsymbol{\mathbf{y}}_{s}|x)$ and to compute $\nabla\log
q(\boldsymbol{\mathbf{y}}_{s}|x)$. While an intractable problem for mdms in
general, we provide two ingredients for tightening and optimizing these bounds
in a generic fashion for linear inference mdms.

We first show how to compute the mean and covariance of
$q(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})$ and then how to
adapt this to the mean and covariance of $q(\boldsymbol{\mathbf{y}}_{s}|x)$.
For linear mdms of the form:

$\displaystyle
d{\mathbf{y}}={\mathbf{A}}(s){\mathbf{y}}ds+g(s)d{\mathbf{B}}_{s},$

the transition kernel
$q(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})$ is Gaussian
(Särkkä & Solin, 2019). Let
$f(\boldsymbol{\mathbf{y}},s)=\boldsymbol{\mathbf{A}}(s)\boldsymbol{\mathbf{y}}$.
Then, the mean and covariance are solutions to the following odes:

$\displaystyle d\boldsymbol{\mathbf{m}}_{s|0}/ds$
$\displaystyle=\boldsymbol{\mathbf{A}}(s)\boldsymbol{\mathbf{m}}_{s|0}$
$\displaystyle d\boldsymbol{\mathbf{\Sigma}}_{s|0}/ds$
$\displaystyle=\boldsymbol{\mathbf{A}}(s)\boldsymbol{\mathbf{\Sigma}}_{s|0}+\boldsymbol{\mathbf{\Sigma}}_{s|0}\boldsymbol{\mathbf{A}}^{\top}(s)+g^{2}(s).$
(8)

The mean can be solved analytically:

$\displaystyle\boldsymbol{\mathbf{m}}_{s|0}=\exp\left[\int_{0}^{s}\boldsymbol{\mathbf{A}}(\nu)d\nu\right]\boldsymbol{\mathbf{y}}_{0}\underbrace{=\exp(s\boldsymbol{\mathbf{A}})\boldsymbol{\mathbf{y}}_{0}}_{\text{no
integration if $\boldsymbol{\mathbf{A}}(\nu)=\boldsymbol{\mathbf{A}}$}}.$ (9)

The covariance equation does not have as simple a solution because eq. 9 as
the unknown matrix $\boldsymbol{\mathbf{\Sigma}}_{s|0}$ is being multiplied
both from the left and the right.

Instead of solving eq. 8 for a specific diffusion manually, as done in
previous work (e.g. pages 50-54 of Dockhorn et al. (2021)), we show that a
matrix factorization technique (Särkkä & Solin (2019), sec. $6.3$) previously
unused in diffusion-based generative models can automatically compute
$\boldsymbol{\mathbf{\Sigma}}_{s|0}$ generically for any linear mdm. Define
$\boldsymbol{\mathbf{C}}_{s},{\boldsymbol{\mathbf{H}}_{s}}$ that evolve
according to:

$\displaystyle\begin{pmatrix}d\boldsymbol{\mathbf{C}}_{s}/ds\\\
d\boldsymbol{\mathbf{H}}_{s}/ds\end{pmatrix}=\begin{pmatrix}\boldsymbol{\mathbf{A}}(s)&g^{2}(s)\\\
\boldsymbol{\mathbf{0}}&-\boldsymbol{\mathbf{A}}^{\top}(s)\end{pmatrix}\begin{pmatrix}\boldsymbol{\mathbf{C}}_{s}\\\
\boldsymbol{\mathbf{H}}_{s}\end{pmatrix},$ (10)

then
$\boldsymbol{\mathbf{\Sigma}}_{s|0}=\boldsymbol{\mathbf{C}}_{s}{\boldsymbol{\mathbf{H}}_{s}}^{-1}$
for $\boldsymbol{\mathbf{C}}_{0}=\boldsymbol{\mathbf{\Sigma}}_{0}$ and
$\boldsymbol{\mathbf{H}}_{0}=\boldsymbol{\mathbf{I}}$ (Appendix C). These
equations can be solved in closed-form,

$\displaystyle\begin{pmatrix}\boldsymbol{\mathbf{C}}_{s}\\\
\boldsymbol{\mathbf{H}}_{s}\end{pmatrix}=\exp\Bigg{[}\begin{pmatrix}[\boldsymbol{\mathbf{A}}]_{s}&[g^{2}]_{s}\\\
\boldsymbol{\mathbf{0}}&-[\boldsymbol{\mathbf{A}}^{\top}]_{s}\end{pmatrix}\Bigg{]}\begin{pmatrix}\boldsymbol{\mathbf{\Sigma}}_{0}\\\
\boldsymbol{\mathbf{I}}\end{pmatrix}\underbrace{=\exp\Bigg{[}s\begin{pmatrix}\boldsymbol{\mathbf{A}}&g^{2}\\\
\boldsymbol{\mathbf{0}}&-\boldsymbol{\mathbf{A}}^{\top}\end{pmatrix}\Bigg{]}}_{\text{no
integration if
$\boldsymbol{\mathbf{A}}(\nu)=\boldsymbol{\mathbf{A}},g(\nu)=g$}}\begin{pmatrix}\boldsymbol{\mathbf{\Sigma_{0}}}\\\
\boldsymbol{\mathbf{I}}\end{pmatrix},$ (11)

where
$[\boldsymbol{\mathbf{A}}]_{s}=\int_{0}^{s}\boldsymbol{\mathbf{A}}(\nu)d\nu$.
To condition on $\boldsymbol{\mathbf{y}}_{0}=(x,v)$, we set
$\boldsymbol{\mathbf{\Sigma}}_{0}=\boldsymbol{\mathbf{0}}$.

##### Computing $q_{\phi}(\boldsymbol{\mathbf{y}}_{s}|x)$.

For the covariance $\boldsymbol{\mathbf{\Sigma}}_{s|0}$, to condition on $x$
instead of $\boldsymbol{\mathbf{y}}_{0}$, we set
$\boldsymbol{\mathbf{\Sigma_{0}}}$ to

$\displaystyle\boldsymbol{\mathbf{\Sigma_{0}}}=\begin{pmatrix}0&0\\\
0&\boldsymbol{\mathbf{\Sigma_{v_{0}}}}\end{pmatrix},$

To compute the mean, it is the same expression as for
$q(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})$, but with a
different initial condition:

$\displaystyle\boldsymbol{\mathbf{m}}_{s|0}=\exp\left[\int_{0}^{s}\boldsymbol{\mathbf{A}}(\nu)d\nu\right]\begin{pmatrix}x\\\
\mathbb{E}_{q}[\boldsymbol{\mathbf{y}}^{v}_{0}|x]\end{pmatrix}$ (12)

See appendix C for more details.

Table 1: Runtime Comparison: we compare the run time of sampling from the cld diffusion analytically versus using the automated algorithm. Method |  cifar-10 | mnist
---|---|---
Analytical | 0.027 | 0.0062
Automated | 0.029 | 0.007

##### A fast and simple algorithm.

We show in algorithm 2 (appendix G) that computing the transition kernel only
requires knowing $f,g$ and requires no diffusion-specific analysis. For $K-1$
auxiliary variables, $\boldsymbol{\mathbf{A}},g$ are $K\times K$. Like for
scalar diffusions, these parameters are shared across data coordinates. This
means matrix exponentials and inverses are done on $K\times K$ matrices, where
$K$ is only $2$ or $3$ in our experiments. In table 1, we compare the time to
sample a batch of size $256$ from the transition kernel for cifar10 and mnist.
The table shows the extra computational cost of the automated algorithm is
negligible.

Input: Data $\\{x_{i}\\}$, inference process matrices
${\mathbf{Q}}_{\phi},{\mathbf{D}}_{\phi}$, model prior $\pi_{\theta}$, initial
distribution $q_{\phi}({\mathbf{y}}^{v}_{0}\mid x)$, and score model
architecture $s_{\theta}$

Returns: Trained score model $s_{\theta}$

while $s_{\theta}$ not converged do

Sample $x\sim\sum_{i=1}^{N}\frac{1}{N}\delta_{x_{i}}$, $v_{0}\sim
q_{\phi}({\mathbf{y}}^{v}_{0}\mid x)$

Sample ${\mathbf{s}}\sim{\mathbf{U}}[0,T]$ and
${\mathbf{y}}_{s},{\mathbf{y}}_{T}\sim q_{\phi}({\mathbf{y}}_{s}\mid x)$ using
algorithm 2

Estimate the stochastic gradient of the mdm elbo,
$\nabla_{\theta}\mathcal{L}(\theta,\phi)$, using eq. 7

$\theta\leftarrow\theta+\alpha\nabla_{\theta}\mathcal{L}(\theta,\phi)$.

if learning inference then

$\phi\leftarrow\phi+\alpha\nabla_{\phi}\mathcal{L}(\theta,\phi)$

end if

end while

Output $s_{\theta}$

Algorithm 1 Automatic Multivariate Diffusion Training

###  3.4 Ingredient 2: mdm Parameterization

The mdm elbo (eq. 7) is tighter when the inference
$\boldsymbol{\mathbf{y}}_{T}$ tends toward the model’s prior $\pi_{\theta}$.
Here we construct inference processes with the model prior $\pi_{\theta}$ as a
specified stationary distribution $q_{\infty}$.

Ma et al. (2015) provide a complete recipe for constructing gradient-based
mcmc samplers; the recipe constructs non-linear time-homogeneous Itô processes
with a given stationary distribution, and show that the parameterization spans
all such Itô processes with that stationary distribution.

Diffusion models usually have time-varying drift and diffusion coefficients
(e.g. use of the $\beta(t)$ function). To build diffusion models that match
the model prior, we first extend Theorem 1 from Ma et al. (2015) to construct
non-linear Itô processes with time-varying drift and diffusion coefficients
with a given stationary distribution (Appendix B). Then, to keep transitions
tractable (per Section 3.3), we specialize this result to linear Itô
diffusions.

We directly state the result for linear time-varying diffusions with
stationary distributions. The parameterization requires a skew-symmetric
matrix $-\boldsymbol{\mathbf{Q}}(s)=\boldsymbol{\mathbf{Q}}(s)^{\top}$, a
positive semi-definite matrix $\boldsymbol{\mathbf{D}}(s)$, and a function
$\nabla H(\boldsymbol{\mathbf{y}})$ such that the desired stationary
distribution $q_{\infty}$ is proportional to
$\exp[-H(\boldsymbol{\mathbf{y}})]$. Linear Itô diffusions have Gaussian
stationary distributions (Särkkä & Solin, 2019) meaning that $\nabla H$ is
linear and can be expressed as
$\boldsymbol{\mathbf{S}}\boldsymbol{\mathbf{y}}$ for some matrix
$\boldsymbol{\mathbf{S}}$. For a matrix $\boldsymbol{\mathbf{A}}$, let
$\sqrt{\boldsymbol{\mathbf{A}}}$ refer to the matrix square root defined by
$\boldsymbol{\mathbf{a}}=\sqrt{\boldsymbol{\mathbf{A}}}\iff\boldsymbol{\mathbf{A}}=\boldsymbol{\mathbf{a}}\boldsymbol{\mathbf{a}}^{\top}$.
Then, the Itô diffusion:

$\displaystyle
d\boldsymbol{\mathbf{y}}=\underbrace{-\Big{[}\boldsymbol{\mathbf{Q}}(s)+\boldsymbol{\mathbf{D}}(s)\Big{]}\boldsymbol{\mathbf{S}}\boldsymbol{\mathbf{y}}}_{f(\boldsymbol{\mathbf{y}},s)}ds+\underbrace{\sqrt{2\boldsymbol{\mathbf{D}}(s)}}_{g(s)}d\widehat{\boldsymbol{\mathbf{B}}}_{s},$
(13)

has Gaussian stationary distribution
$\mathcal{N}(\boldsymbol{\mathbf{0}},\boldsymbol{\mathbf{S}}^{-1})$ where
$\boldsymbol{\mathbf{Q}}(s),\boldsymbol{\mathbf{D}}(s)$ and
$\boldsymbol{\mathbf{S}}$ are parameters. For a discussion of uniqueness and
convergence to the stationary distribution, see appendix B. In appendix B, we
show that existing diffusion processes such as vpsde and cld fall under this
parameterization, and discuss ways to parameterize skew-symmetric and positive
semi-definite matrices. We display the elbo in terms of
$\boldsymbol{\mathbf{Q}}/\boldsymbol{\mathbf{D}}$ in appendix F and an
algorithm in appendix G.

###  3.5 Learning the inference process

The choice of diffusion matters, and the elbos in eq. 7 have no requirement
for fixed $q_{\phi}$. We therefore learn the inference process jointly with
$s_{\theta}$. Under linear transitions (ingredient $1$), no algorithmic
details change as the diffusion changes during training. Under stationary
parameterization (ingredient $2$), we can learn without the stationary
distribution going awry. In the experiments, learning matches or surpasses
bpds of fixed diffusions for a given dataset and score architecture.

In $\mathcal{L}^{\text{mdsm}}$ or $\mathcal{L}^{\text{mism}}$,
$q_{\phi,\infty}$ may be set to equal $\pi_{\theta}$, but it is
${\mathbf{y}}_{T}\sim q_{\phi,T}$ for the chosen $T$ that is featured in the
elbo. Learning $q_{\phi}$ can choose $\boldsymbol{\mathbf{y}}_{T}$ to reduce
the cross-entropy:

$\displaystyle-\mathbb{E}_{q_{\phi}(\boldsymbol{\mathbf{y}}_{T}|x)}[\log\pi_{\theta}(\boldsymbol{\mathbf{y}}_{T})].$
(14)

Minimizing eq. 14 will tighten the elbo for any $s_{\theta}$. Next, $q_{\phi}$
is featured in the remaining terms; optimizing for $q_{\phi}$ will tighten and
improve the elbo alongside $s_{\theta}$. Finally, $q_{\phi}$ is featured in
the expectations and the $-\log q_{\phi}$ term:

$\displaystyle\log p_{\theta}(\boldsymbol{\mathbf{u}}_{T}^{z}=x)$
$\displaystyle\geq=\underbrace{\mathbb{E}_{q_{\phi}(\boldsymbol{\mathbf{y}}_{0}^{v}=v|x)}}\Big{[}(\mathcal{L}^{\text{dsm}}\text{
or }\mathcal{L}^{\text{ism}})\underbrace{-\log
q_{\phi}(\boldsymbol{\mathbf{y}}_{0}^{v}=v|x)}\Big{]}$ (15)

The $q_{\phi}(\boldsymbol{\mathbf{y}}_{0}^{v}|x)$ terms impose an optimality
condition that
$p_{\theta}(\boldsymbol{\mathbf{u}}_{T}^{v}|\boldsymbol{\mathbf{u}}_{T}^{z})=q_{\phi}(\boldsymbol{\mathbf{y}}_{0}^{v}|\boldsymbol{\mathbf{y}}_{0}^{z})$
(appendix D), When it is satisfied, no looseness in the elbo is due to the
initial time zero auxiliary variables.

To learn, $\boldsymbol{\mathbf{Q}},\boldsymbol{\mathbf{D}}$ need to be
specified with parameters $\phi$ that enable gradients. We keep
$\boldsymbol{\mathbf{S}}$ fixed at inverse covariance of $\pi_{\theta}$. The
choice of $\boldsymbol{\mathbf{Q}},\boldsymbol{\mathbf{D}}$ generates the
distribution parameters of the Gaussian distributions
$q_{\phi}(\boldsymbol{\mathbf{y}}_{s}|x)$ through the mean
$\boldsymbol{\mathbf{m}}_{s|0}$ and covariance $\boldsymbol{\Sigma}_{s|0}$.
Gaussian distributions permit gradient estimation with reparameterization or
score-function gradients (Kingma & Welling, 2013; Ranganath et al., 2014;
Rezende & Mohamed, 2015; Titsias & Lázaro-Gredilla, 2014). Reparameterization
is accomplished via:

$\displaystyle\boldsymbol{\mathbf{y}}_{s}={\mathbf{m}}_{s|0}+{\mathbf{L}}_{s|0}\epsilon$
(16)

where $\epsilon\sim{\mathcal{N}}(0,I_{dK})$ and ${\mathbf{L}}_{s|0}$ satisfies
${\mathbf{L}}_{s|0}{\mathbf{L}}_{s|0}^{\top}=\mathbf{\Sigma}_{s|0}$, derived
using coordinate-wise Cholesky decomposition. Gradients flow through eq. 16
from $\boldsymbol{\mathbf{y}}_{s}$ to $\boldsymbol{\mathbf{m}}_{s|0}$ and
$\boldsymbol{\Sigma}_{s|0}$ to
$\boldsymbol{\mathbf{Q}},\boldsymbol{\mathbf{D}}$ to parameters $\phi$.

Algorithm 1 displays Automatic Multivariate Diffusion Training (amdt). Amdt
provides a training method for diffusion-based generative models for either
fixed ${\mathbf{Q}},{\mathbf{D}}$ matrices or for learning the
${\mathbf{Q}}_{\phi},{\mathbf{D}}_{\phi}$ matrices, without requiring any
diffusion-specific analysis.

##  4 Insights into Multivariate Diffusions

##### Scalar versus Multivariate Processes.

Equation 13 clarifies what can change while preserving $q_{\infty}$. Recall
that $\boldsymbol{\mathbf{Q}}$ and $\boldsymbol{\mathbf{D}}$ are $K\times K$
for $K-1$ auxiliary variables. Because $0$ is the only $1\times 1$ skew-
symmetric matrix, scalar processes must set $\boldsymbol{\mathbf{Q}}=0$. With
$q_{\phi,\infty}=\mathcal{N}(0,\boldsymbol{\mathbf{I}})$, the process is:

$\displaystyle
d\boldsymbol{\mathbf{y}}=-\boldsymbol{\mathbf{D}}(s)\boldsymbol{\mathbf{y}}ds+\sqrt{2\boldsymbol{\mathbf{D}}(s)}d\widehat{\boldsymbol{\mathbf{B}}}_{s}.$
(17)

What is left is the vpsde process used widely in diffusion models where
$\boldsymbol{\mathbf{D}}(s)=\frac{1}{2}\beta(s)$ is $1\times 1$ (Song et al.,
2020b). This reveals that the vpsde process is the only scalar diffusion with
a stationary distribution.333There are processes such as sub-vpsde (Song et
al., 2020b) which are covered in the sense that they tend to members of this
parameterization as $T$ grows: sub-VP converges to vpsde. This also clarifies
the role of $\boldsymbol{\mathbf{Q}}$: it accounts for mixing between
dimensions in multivariate processes, as do non-diagonal entries in
$\boldsymbol{\mathbf{D}}$ for $K>1$.

#####  cld optimizes a lower bound on the data log-likelihood.

Taking the derivative of $\mathcal{L}^{\text{mdsm}}$ (eq. 7) with respect to
the score model parameters, we show that the objective for cld (Dockhorn et
al., 2021) formally maximizes a lower bound on $\log p_{\theta}(x)$ without
appealing to the probability flow ode.

#####  Does my model use auxiliary variables?

An example initial distribution is
$q(\boldsymbol{\mathbf{y}}_{0}^{v}|x)=\mathcal{N}(0,\boldsymbol{\mathbf{I}})$.
It is also common to set
$\pi_{\theta}=\mathcal{N}(0,\boldsymbol{\mathbf{I}})$. Because the optimum for
diffusions is $p_{\theta}=q$, the optimal model has main and auxiliary
dimensions independent at endpoints $0$ and $T$. Does this mean that the model
does not use auxiliary variables? In appendix A, we show that in this case the
model can still use auxiliary variables at intermediate times. A sufficient
condition is non-diagonal $\boldsymbol{\mathbf{Q}}+\boldsymbol{\mathbf{D}}$.

##  5 Experiments

We test the mdm framework with handcrafted and learned diffusions. The
handcrafted diffusions are (a) Alda, used in (Mou et al., 2019) for
accelerated gradient-based mcmc sampling (eq. 32) and (b) Malda: a modified
version of alda (eq. 33). Both have two auxiliary variables. We also learn
diffusions with $1$ and $2$ auxiliary variables. We compare with vpsde and
elbo-trained cld.

Table 2: Bpd upper-bounds on image generation for a fixed architecture. cifar-10: learning outperforms cld, and both outperform the standard choice of vpsde. mnist: learning matches vpsde while the fixed auxiliary diffusions are worse. imagenet32: all perform similarly. Learning matches or surpasses the best fixed diffusion, while bypassing the need to choose a diffusion. Model | $K$ |  cifar-10 |  imagenet32 | mnist
---|---|---|---|---
vpsde | 1 | $3.20$ | $3.70$ | $1.26$
Learned | 2 | $3.07$ | $3.71$ | $1.28$
Learned | 3 | $3.08$ | $3.72$ | $1.33$
cld | 2 | $3.11$ | $3.70$ | $1.35$
malda | 3 | $3.13$ | $3.72$ | $1.65$
alda | 3 | $29.43$ | $33.08$ | $124.60$
Table 3: Parameter Efficiency. The first two rows display diffusions from previous work: vpsde and cld, both using score models with 108 million parameters on cifar-10. We train the rest using a score model with 35.7 million parameters. The learned diffusion matches the performance of vpsde-large; changes in the inference can account for as much improvement as a 3x increase in score parameters. bpds are upper-bounds.  Model | $K$ | Parameters |  cifar-10
---|---|---|---
vpsde-large (Song et al., 2021) | 1 | 108M | $3.08$
cld-large (Dockhorn et al., 2021) | 2 | 108M | $3.31$
Learned | 2 | 35.7M | $3.07$
cld | 2 | 35.7M | $3.11$
vpsde | 1 | 35.7M | $3.20$

Following prior work, we train dbgms for image generation. We use the U-Net
from Ho et al. (2020). We input the auxiliary variables as extra channels,
which only increases the score model parameters in the input and output
convolutions (cld and Learned 2 have $7,000$ more parameters than vpsde on
cifar-10 and imagenet32 and only $865$ more for mnist). We use simple uniform
dequantization. We report estimates of $\mathcal{L}^{\text{mdsm}}$ (which
reduces to the standard $\mathcal{L}^{\text{dsm}}$ for $K=1$). We sample times
using the importance sampling distribution from Song et al. (2021) with
truncation set to $\epsilon=10^{-3}$. To ensure the truncated bound is proper,
we use a likelihood described in appendix H.

##### Results.

Table 2 shows that the inference process matters and displays. It displays
dbgms that we train and evaluate on cifar-10, imagenet32 and mnist. This
includes the existing vpsde and cld, the new malda and alda, and the new
learned inference processes. All are trained with the 35.7M parameter
architecture. For cifar-10, learning outperforms cld, and both outperform the
standard choice of vpsde. For mnist, learned diffusions match vpsde while the
three fixed auxiliary diffusions are worse. On imagenet32, all perform
similarly. The take-away is that learning matches or surpasses the best fixed
diffusion performance and bypasses the choice of diffusion for each new
dataset or score architecture. In Figure 1 we plot the generated samples from
cifar10.

Table 3’s first two rows display diffusion models from previous work: vpsde
(Song et al., 2021) and cld (Dockhorn et al., 2021) both with the 108 million
score model from Song et al. (2021) (labeled “large”). The rest are dbgms that
we train using the U-Net with 35.7 million parameters for cifar-10 and
imagenet32 and 1.1 million for mnist. Despite using significantly fewer
parameters, the learned diffusion achieves similar bpd compared to the larger
models, showing that changes in inference can account for as much improvement
as a three-fold increase in parameters. While the larger architecture requires
two GPUs for batch size 128 on cifar-10 on A100s, the smaller one only
requires one; exploring inference processes can make diffusions more
computationally accessible. Table 3 also demonstrates a tighter bound for cld
trained and evaluated with the mdm elbo ($\leq 3.11$) relative to existing
probability flow-based evaluations ($3.31$).

Figure 1: cifar10 samples generated from the “learned 2” and malda generative
models.

##  6 Related Work

Evidence Lower Bounds. Song et al. (2021); Huang et al. (2021) derive the ism
and dsm lower-bounds on the model log likelihood. Our work extends their
analysis to the multivariate diffusion setting to derive lower bounds on the
log marginal of the data in the presence of auxiliary variables.

##### Auxiliary variables.

Dupont et al. (2019) shows that augmented neural odes model a richer set of
functions and Huang et al. (2020) uses this principle for normalizing flows.
Hierarchical variational models and auto-encoders marginalize auxiliary
variables to build expressive distributions (Ranganath et al., 2016; Sønderby
et al., 2016; Maaløe et al., 2019; Vahdat & Kautz, 2020; Child, 2020). We
apply this principle to dbgms, including and extending cld (Dockhorn et al.,
2021).

##### Learning inference.

Learning $q_{\phi}$ with $p_{\theta}$ is motivated in previous work (Kingma &
Welling, 2013; Sohl-Dickstein et al., 2015; Kingma et al., 2021). Kingma et
al. (2021) learn the noise schedule for vpsde. For mdms, there are parameters
to learn beyond the noise schedule; $\boldsymbol{\mathbf{Q}}$ can be non-zero,
$\boldsymbol{\mathbf{D}}$ can diagonal or full, give $\boldsymbol{\mathbf{Q}}$
and $\boldsymbol{\mathbf{D}}$ different time-varying functions, and learn
$\nabla\boldsymbol{\mathbf{H}}$.

##  7 Discussion

We present an algorithm for training multivariate diffusions with linear time-
varying inference processes with a specified stationary distribution and any
number of auxiliary variables. This includes automating transition kernel
computation and providing a parameterization of diffusions that have a
specified stationary distribution, which facilitate working with new diffusion
processes, including learning the diffusion. The experiments show that
learning matches or surpasses the best fixed diffusion performance, bypassing
the need to choose a diffusion. mdms achieve bpds similar to univariate
diffusions, with as many as three times more score parameters. The proposed
mdm elbo reports a tighter bound for the existing cld relative to existing
probability flow-based evaluations. This work enables future directions
including interactions across data coordinates and using new stationary
distributions.

##  8 Acknowledgements

This work was generously funded by NIH/NHLBI Award R01HL148248, NSF Award
1922658 NRT-HDR: FUTURE Foundations, Translation, and Responsibility for Data
Science, and NSF CAREER Award 2145542. The authors would additionally like to
thank Chin-Wei Huang for helpful discussing regarding Huang et al. (2021).

## References

  * Barbour (1988) Andrew D Barbour.  Stein’s method and poisson process convergence.  _Journal of Applied Probability_ , 25(A):175–184, 1988.
  * Campbell et al. (2022) Andrew Campbell, Joe Benton, Valentin De Bortoli, Tom Rainforth, George Deligiannidis, and Arnaud Doucet.  A continuous time framework for discrete denoising models.  _arXiv preprint arXiv:2205.14987_ , 2022.
  * Chen et al. (2020) Nanxin Chen, Yu Zhang, Heiga Zen, Ron J Weiss, Mohammad Norouzi, and William Chan.  Wavegrad: Estimating gradients for waveform generation.  _arXiv preprint arXiv:2009.00713_ , 2020.
  * Child (2020) Rewon Child.  Very deep vaes generalize autoregressive models and can outperform them on images.  _arXiv preprint arXiv:2011.10650_ , 2020.
  * Dhariwal & Nichol (2021) Prafulla Dhariwal and Alexander Nichol.  Diffusion models beat gans on image synthesis.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Dockhorn et al. (2021) Tim Dockhorn, Arash Vahdat, and Karsten Kreis.  Score-based generative modeling with critically-damped langevin diffusion.  _arXiv preprint arXiv:2112.07068_ , 2021.
  * Dupont et al. (2019) Emilien Dupont, Arnaud Doucet, and Yee Whye Teh.  Augmented neural odes.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Efron (2011) Bradley Efron.  Tweedie’s formula and selection bias.  _Journal of the American Statistical Association_ , 106(496):1602–1614, 2011.
  * Ho et al. (2020) Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _arXiv preprint arXiv:2006.11239_ , 2020.
  * Ho et al. (2022) Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi, and Tim Salimans.  Cascaded diffusion models for high fidelity image generation.  _J. Mach. Learn. Res._ , 23:47–1, 2022.
  * Huang et al. (2020) Chin-Wei Huang, Laurent Dinh, and Aaron Courville.  Augmented normalizing flows: Bridging the gap between generative flows and latent variable models.  _arXiv preprint arXiv:2002.07101_ , 2020.
  * Huang et al. (2021) Chin-Wei Huang, Jae Hyun Lim, and Aaron C Courville.  A variational perspective on diffusion-based generative models and score matching.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Kingma & Welling (2013) Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Kingma et al. (2021) Diederik P Kingma, Tim Salimans, Ben Poole, and Jonathan Ho.  Variational diffusion models.  _arXiv preprint arXiv:2107.00630_ , 2021.
  * Kong et al. (2020) Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro.  Diffwave: A versatile diffusion model for audio synthesis.  _arXiv preprint arXiv:2009.09761_ , 2020.
  * Ma et al. (2015) Yi-An Ma, Tianqi Chen, and Emily Fox.  A complete recipe for stochastic gradient mcmc.  _Advances in neural information processing systems_ , 28, 2015.
  * Maaløe et al. (2019) Lars Maaløe, Marco Fraccaro, Valentin Liévin, and Ole Winther.  Biva: A very deep hierarchy of latent variables for generative modeling.  _Advances in neural information processing systems_ , 32, 2019.
  * Mittal et al. (2021) Gautam Mittal, Jesse Engel, Curtis Hawthorne, and Ian Simon.  Symbolic music generation with diffusion models.  _arXiv preprint arXiv:2103.16091_ , 2021.
  * Mou et al. (2019) Wenlong Mou, Yi-An Ma, Martin J Wainwright, Peter L Bartlett, and Michael I Jordan.  High-order langevin diffusion yields an accelerated mcmc algorithm.  _arXiv preprint arXiv:1908.10859_ , 2019.
  * Nichol & Dhariwal (2021) Alex Nichol and Prafulla Dhariwal.  Improved denoising diffusion probabilistic models.  _arXiv preprint arXiv:2102.09672_ , 2021.
  * Nichol et al. (2021) Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen.  Glide: Towards photorealistic image generation and editing with text-guided diffusion models.  _arXiv preprint arXiv:2112.10741_ , 2021.
  * Ramesh et al. (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.
  * Ranganath et al. (2014) Rajesh Ranganath, Sean Gerrish, and David Blei.  Black box variational inference.  In _Artificial intelligence and statistics_ , pp. 814–822. PMLR, 2014.
  * Ranganath et al. (2016) Rajesh Ranganath, Dustin Tran, and David Blei.  Hierarchical variational models.  In _International conference on machine learning_ , pp. 324–333. PMLR, 2016.
  * Rezende & Mohamed (2015) Danilo Rezende and Shakir Mohamed.  Variational inference with normalizing flows.  In _International Conference on Machine Learning_ , pp. 1530–1538. PMLR, 2015.
  * Saharia et al. (2022) Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al.  Photorealistic text-to-image diffusion models with deep language understanding.  _arXiv preprint arXiv:2205.11487_ , 2022.
  * Särkkä & Solin (2019) Simo Särkkä and Arno Solin.  _Applied stochastic differential equations_ , volume 10.  Cambridge University Press, 2019.
  * Sasaki et al. (2021) Hiroshi Sasaki, Chris G Willcocks, and Toby P Breckon.  Unit-ddpm: Unpaired image translation with denoising diffusion probabilistic models.  _arXiv preprint arXiv:2104.05358_ , 2021.
  * Shi et al. (2012) Jianghong Shi, Tianqi Chen, Ruoshi Yuan, Bo Yuan, and Ping Ao.  Relation of a new interpretation of stochastic differential equations to ito process.  _Journal of Statistical physics_ , 148:579–590, 2012.
  * Sohl-Dickstein et al. (2015) Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.  Deep unsupervised learning using nonequilibrium thermodynamics.  In _International Conference on Machine Learning_ , pp. 2256–2265. PMLR, 2015.
  * Sønderby et al. (2016) Casper Kaae Sønderby, Tapani Raiko, Lars Maaløe, Søren Kaae Sønderby, and Ole Winther.  Ladder variational autoencoders.  _Advances in neural information processing systems_ , 29, 2016.
  * Song et al. (2020a) Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon.  Sliced score matching: A scalable approach to density and score estimation.  In _Uncertainty in Artificial Intelligence_ , pp. 574–584. PMLR, 2020a.
  * Song et al. (2020b) Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-based generative modeling through stochastic differential equations.  _arXiv preprint arXiv:2011.13456_ , 2020b.
  * Song et al. (2021) Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon.  Maximum likelihood training of score-based diffusion models.  _Advances in Neural Information Processing Systems_ , 34:1415–1428, 2021.
  * Titsias & Lázaro-Gredilla (2014) Michalis Titsias and Miguel Lázaro-Gredilla.  Doubly stochastic variational bayes for non-conjugate inference.  In _International conference on machine learning_ , pp. 1971–1979. PMLR, 2014.
  * Vahdat & Kautz (2020) Arash Vahdat and Jan Kautz.  Nvae: A deep hierarchical variational autoencoder.  _Advances in Neural Information Processing Systems_ , 33:19667–19679, 2020.
  * Vincent (2011) Pascal Vincent.  A connection between score matching and denoising autoencoders.  _Neural computation_ , 23(7):1661–1674, 2011\.
  * Yin & Ao (2006) L Yin and P Ao.  Existence and construction of dynamical potential in nonequilibrium processes without detailed balance.  _Journal of Physics A: Mathematical and General_ , 39(27):8593, 2006.
  * Zhang & Chen (2013) Zhenzhong Zhang and Dayue Chen.  A new criterion on existence and uniqueness of stationary distribution for diffusion processes.  _Advances in Difference Equations_ , 2013(1):1–6, 2013.

##  Appendix A Does my model use auxiliary variables?

In section 3 we gave the example choice of
$q(\boldsymbol{\mathbf{y}}_{0}^{v}|x)=\mathcal{N}(0,\boldsymbol{\mathbf{I}})$
coordinate-wise. It is also a common choice to set
$\pi_{\theta}=\mathcal{N}(0,\boldsymbol{\mathbf{I}})$. Because the optimum in
diffusion models is $p_{\theta}=q$ for all $t$, we see a peculiar phenomenon
under this choice: the model has main and auxiliary dimensions independent at
both endpoints $0$ and $T$. Does this mean that the model does not use
auxiliary variables? We show that even when
$q_{\phi}(\boldsymbol{\mathbf{y}}_{0})$ and $\pi_{\theta}$ have main and
auxiliary variables independent, the model can use the auxiliary variables. A
sufficient condition is $\boldsymbol{\mathbf{Q}}+\boldsymbol{\mathbf{D}}$ is
non-diagonal.

To make this precise, we recall that we model with
$p_{\theta}(\boldsymbol{\mathbf{u}}_{T}^{z}=x)$. To show the model is using
auxiliary variables, we just need to show that
$\boldsymbol{\mathbf{u}}_{T}^{z}$ (main coordinate at $T$) depends on
$\boldsymbol{\mathbf{u}}_{t}^{v}$ (aux. coordinate at $t$) for $T>t$. At
optimum,
$p_{\theta}(\boldsymbol{\mathbf{u}}_{T}^{z},\boldsymbol{\mathbf{u}}_{t}^{v})=q_{\phi}(\boldsymbol{\mathbf{y}}_{0}^{z},\boldsymbol{\mathbf{y}}_{T-t}^{v})$.
Therefore it is sufficient to show that for some time $s$,
$q_{\phi}(\boldsymbol{\mathbf{y}}_{s}^{v}|\boldsymbol{\mathbf{y}}_{0}^{z})\neq
q_{\phi}(\boldsymbol{\mathbf{y}}_{s}^{v})$. Because
$\boldsymbol{\mathbf{y}}_{0}^{z}$, is determined by $x$ we need to show that
$q_{\phi}(\boldsymbol{\mathbf{y}}_{s}^{v}|x)\neq
q_{\phi}(\boldsymbol{\mathbf{y}}_{s}^{v})$. To do that, we first derive
$q(\boldsymbol{\mathbf{y}}_{s}|x)$ and then marginalize to get
$q(\boldsymbol{\mathbf{y}}_{s}^{v}|x)$ from it. Since the former is 2D
Gaussian, the latter is available in terms of the former’s mean and
covariance. Suppose $\mathbb{E}[\boldsymbol{\mathbf{y}}_{0}^{v}]=0$,
$\boldsymbol{\mathbf{Q}}=[[0,-1],[1,0]]$ and
$\boldsymbol{\mathbf{D}}=[[1,0],[0,1]]$ and we have $s=.1$ We have:

$\displaystyle\mathbb{E}[\boldsymbol{\mathbf{y}}_{s}|x]$
$\displaystyle=\exp\Big{[}-s(\boldsymbol{\mathbf{Q}}+\boldsymbol{\mathbf{D}})\Big{]}\begin{pmatrix}x\\\
0\end{pmatrix}=\exp\Big{[}\begin{bmatrix}-.1&.1\\\
-.1&-.1\end{bmatrix}\Big{]}\begin{pmatrix}x\\\
0\end{pmatrix}=\begin{pmatrix}0.9003x\\\ -0.090x\end{pmatrix}$ (18)

Regardless of the covariance any 1D of this 2D gaussian will have mean that is
a function of $x$, meaning that $q(\boldsymbol{\mathbf{y}}_{s}^{v}|x)$ does
not equal $q(\boldsymbol{\mathbf{y}}_{s}^{v})$ (which is also a Gaussian but
with mean depending on $\boldsymbol{\mathbf{x}}^{\prime}s$ mean rather than
$x$ itself. Therefore, even under the setup with independent endpoints, the
optimal model makes use of the intermediate auxiliary variables in its final
modeling distribution $p_{\theta}(\boldsymbol{\mathbf{u}}_{T}^{z}=x)$.

Are there choices of $\boldsymbol{\mathbf{Q}}$ and $\boldsymbol{\mathbf{D}}$
that lead to learning models that don’t make use of the extra dimensions? As
mentioned, in the inference process, $\boldsymbol{\mathbf{Q}}$ is responsible
for mixing information among the coordinates, and is the only source of this
when $\boldsymbol{\mathbf{D}}$ is diagonal. Then, if
$\boldsymbol{\mathbf{Q}}=\boldsymbol{0}$ and $\boldsymbol{\mathbf{D}}$ is
diagonal, none of the coordinates for a given feature
$\boldsymbol{\mathbf{x}}_{j}$ (including
$\boldsymbol{\mathbf{u}}_{tj}^{z},\boldsymbol{\mathbf{u}}_{tj}^{v_{1}},\ldots,\boldsymbol{\mathbf{u}}_{tj}^{v_{K-1}}$)
interact for any $t$. Then, since $p_{\theta}=q$ at optimum, independence of
the coodinates at all $t$ in $q$ imply the same in $p_{\theta}$ and the model
will not make use of any auxiliary variables when modeling the marginal $\log
p_{\theta}(\boldsymbol{\mathbf{u}}_{T}^{z}=x)$.

##  Appendix B Stationary parameterization

The non-linear time-homogeneous Itô process family is:

$\displaystyle
d\boldsymbol{\mathbf{y}}=f(\boldsymbol{\mathbf{y}})dt+g(\boldsymbol{\mathbf{y}})\boldsymbol{\mathbf{B}}_{t}.$
(19)

This family can be restricted to those with stationary distributions. Ma et
al. (2015) show a complete recipe to span the subset of this family with a
desired stationary distribution. Let $\boldsymbol{\mathbf{Q}}$ be skew-
symmetric ($-\boldsymbol{\mathbf{Q}}=\boldsymbol{\mathbf{Q}}^{\top}$) and
$\boldsymbol{\mathbf{D}}$ is positive semi-definite. Suppose the desired
stationary distribution is $q_{\infty}(\boldsymbol{\mathbf{y}})$. For a matrix
$\boldsymbol{\mathbf{A}}$, let $\sqrt{\boldsymbol{\mathbf{A}}}$ refer to the
matrix square root defined by
$\boldsymbol{\mathbf{a}}=\sqrt{\boldsymbol{\mathbf{A}}}\iff\boldsymbol{\mathbf{A}}=\boldsymbol{\mathbf{a}}\boldsymbol{\mathbf{a}}^{\top}$.
Then, Ma et al. (2015) show that, setting
$\boldsymbol{\mathbf{H}}(\boldsymbol{\mathbf{y}})=-\log
q_{\infty}(\boldsymbol{\mathbf{y}})$,
$g(\boldsymbol{\mathbf{y}})=\sqrt{2\boldsymbol{\mathbf{D}}(\boldsymbol{\mathbf{y}})}$,
and

$\displaystyle
f(\boldsymbol{\mathbf{y}})=-[\boldsymbol{\mathbf{D}}(\boldsymbol{\mathbf{y}})+\boldsymbol{\mathbf{Q}}(\boldsymbol{\mathbf{y}})]\nabla\boldsymbol{\mathbf{H}}(\boldsymbol{\mathbf{y}})+\boldsymbol{\Gamma}(\boldsymbol{\mathbf{y}}),\quad\quad\boldsymbol{\Gamma}_{i}(\boldsymbol{\mathbf{y}})=\sum_{j=1}^{d}\frac{\partial}{\partial\boldsymbol{\mathbf{z}}_{j}}(\boldsymbol{\mathbf{D}}_{ij}(\boldsymbol{\mathbf{y}})+\boldsymbol{\mathbf{Q}}_{ij}(\boldsymbol{\mathbf{y}})),$
(20)

yields a process $\boldsymbol{\mathbf{y}}_{t}$ with stationary distribution
$q_{\infty}$. We extend it to time-varying (time in-homogeneous) processes.

######  Theorem 2.

$q_{\infty}(\boldsymbol{\mathbf{y}})\propto\exp[-H(\boldsymbol{\mathbf{y}})]$
is a stationary distribution of

$\displaystyle
d\boldsymbol{\mathbf{y}}=\Bigg{(}-[\boldsymbol{\mathbf{D}}(\boldsymbol{\mathbf{y}},t)+\boldsymbol{\mathbf{Q}}(\boldsymbol{\mathbf{y}},t)]\nabla\boldsymbol{\mathbf{H}}(\boldsymbol{\mathbf{y}})+\boldsymbol{\Gamma}(\boldsymbol{\mathbf{y}},t)\Bigg{)}dt+\sqrt{2\boldsymbol{\mathbf{D}}(\boldsymbol{\mathbf{y}},t)}\boldsymbol{\mathbf{B}}_{t},$
(21)

for

$\displaystyle\boldsymbol{\Gamma}_{i}(\boldsymbol{\mathbf{y}},t)=\sum_{j=1}^{d}\frac{\partial}{\partial\boldsymbol{\mathbf{y}}_{j}}(\boldsymbol{\mathbf{D}}_{ij}(\boldsymbol{\mathbf{y}},t)+\boldsymbol{\mathbf{Q}}_{ij}(\boldsymbol{\mathbf{y}},t)).$
(22)

###### Proof.

The Fokker Planck equation is:

$\displaystyle\partial_{t}q(\boldsymbol{\mathbf{y}},t)=-\sum_{i}\frac{\partial}{\partial\boldsymbol{\mathbf{y}}_{i}}\Big{[}f_{i}(\boldsymbol{\mathbf{y}},t)q(\boldsymbol{\mathbf{y}},t)\Big{]}+\sum_{i,j}\frac{\partial^{2}}{\partial\boldsymbol{\mathbf{y}}_{i}\partial\boldsymbol{\mathbf{y}}_{j}}\Big{[}\boldsymbol{\mathbf{D}}_{ij}(\boldsymbol{\mathbf{y}},t)q(\boldsymbol{\mathbf{y}},t)\Big{]}$
(23)

A stationary distribution is one where the Fokker-Planck right hand side is
equal to $0$. To show that the stationary characterization also holds of time-
inhomogenous processes with
$\boldsymbol{\mathbf{D}}(\boldsymbol{\mathbf{y}},t)$ and
$\boldsymbol{\mathbf{Q}}(\boldsymbol{\mathbf{y}},t)$, we take two steps,
closely following Yin & Ao (2006); Shi et al. (2012); Ma et al. (2015), but
noting that there is no requirement for
$\boldsymbol{\mathbf{Q}},\boldsymbol{\mathbf{D}}$ to be free of $t$. First, we
show that the Fokker-Plack equation can be re-written as:

$\displaystyle\partial_{t}q(\boldsymbol{\mathbf{y}},t)=\nabla\cdot\Bigg{(}\Big{[}\boldsymbol{\mathbf{D}}(\boldsymbol{\mathbf{y}},t)+\boldsymbol{\mathbf{Q}}(\boldsymbol{\mathbf{y}},t)\Big{]}\Big{[}q(\boldsymbol{\mathbf{y}},t)\nabla
H(\boldsymbol{\mathbf{y}})+\nabla q(\boldsymbol{\mathbf{y}},t)\Big{]}\Bigg{)}$
(24)

Second, because the whole expression is set to $0$ when the inside expression
equals $0$

$\displaystyle q(\boldsymbol{\mathbf{y}},t)\nabla
H(\boldsymbol{\mathbf{y}})+\nabla q(\boldsymbol{\mathbf{y}},t)=0,$ (25)

we just need to show that this holds when
$q(\boldsymbol{\mathbf{y}},t)=\exp[-H(\boldsymbol{\mathbf{y}})]/{\mathbf{Z}}$.
The second step is concluded because

$\displaystyle\Big{[}q(\boldsymbol{\mathbf{y}},t)\nabla
H(\boldsymbol{\mathbf{y}})+\nabla q(\boldsymbol{\mathbf{y}},t)\Big{]}$
$\displaystyle=\frac{1}{{\mathbf{Z}}}\Big{[}\exp[-H(\boldsymbol{\mathbf{y}})]\nabla
H(\boldsymbol{\mathbf{y}})+\nabla\exp[-H(\boldsymbol{\mathbf{y}})]\Big{]}=0,$

where ${\mathbf{Z}}$ is the normalization constant of $\exp(-H(y))$.

It only remains to show that Fokker-Plack can be re-written in divergence form
with time-dependent $\boldsymbol{\mathbf{Q}},\boldsymbol{\mathbf{D}}$. In the
following let $Q_{ijt}$ denote
$\boldsymbol{\mathbf{Q}}_{ij}(\boldsymbol{\mathbf{y}},t)$ and likewise for
$D_{ijt}$. Let $\partial_{i}$ denote
$\frac{\partial}{\partial\boldsymbol{\mathbf{y}}_{i}}$ and let it denote
$\frac{d}{d\boldsymbol{\mathbf{y}}_{i}}$ for scalar functions. We will use
$[Ax]_{i}=\sum_{j}A_{ij}x_{j}$.

$\displaystyle\partial_{t}q_{t}$
$\displaystyle=\nabla\cdot\Big{(}[\boldsymbol{\mathbf{D}}(\boldsymbol{\mathbf{y}},t)+\boldsymbol{\mathbf{Q}}(\boldsymbol{\mathbf{y}},t)][q\nabla
H+\nabla q]\Big{)}$
$\displaystyle=\sum_{i}\partial_{i}\Big{(}\Big{[}[\boldsymbol{\mathbf{D}}(\boldsymbol{\mathbf{y}},t)+\boldsymbol{\mathbf{Q}}(\boldsymbol{\mathbf{y}},t)][q\nabla
H+\nabla q]\Big{]}_{i}\Big{)}$
$\displaystyle=\sum_{i}\partial_{i}\sum_{j}[D_{ijt}+Q_{ijt}][q\nabla H+\nabla
q]_{j}$
$\displaystyle=\sum_{i}\partial_{i}\sum_{j}[D_{ijt}+Q_{ijt}][q\partial_{j}H+\partial_{j}q]$
$\displaystyle=\sum_{i}\partial_{i}\sum_{j}[D_{ijt}+Q_{ijt}][q\partial_{j}H]+\sum_{i}\partial_{i}\sum_{j}[D_{ijt}+Q_{ijt}][\partial_{j}q]$
$\displaystyle=\sum_{i}\partial_{i}\sum_{j}[D_{ijt}+Q_{ijt}][q\partial_{j}H]+\sum_{i}\partial_{i}\sum_{j}D_{ijt}[\partial_{j}q]+\sum_{i}\partial_{i}\sum_{j}Q_{ijt}[\partial_{j}q]$

We re-write the 2nd and 3rd term. Holding $i$ fixed and noting $q$ is scalar,
we get the product rule
$\sum_{j}D_{ijt}(\partial_{j}q)=\sum_{j}\partial_{j}[D_{ijt}q]-q\sum_{j}\partial_{j}D_{ijt}$
for each $i$, and likewise for $q$:

$\displaystyle\sum_{i}\partial_{i}\sum_{j}[D_{ijt}+Q_{ijt}][q\partial_{j}H]+\sum_{i}\partial_{i}\sum_{j}D_{ijt}[\partial_{j}q]+\sum_{i}\partial_{i}\sum_{j}Q_{ijt}[\partial_{j}q]$
$\displaystyle=\sum_{i}\partial_{i}\sum_{j}[D_{ijt}+Q_{ijt}][q\partial_{j}H]+\sum_{i}\partial_{i}\sum_{j}\partial_{j}[D_{ijt}q]-q\sum_{j}\partial_{j}D_{ijt}$
$\displaystyle\quad\quad+\sum_{i}\partial_{i}\sum_{j}\partial_{j}[Q_{ijt}q]-q\sum_{j}\partial_{j}Q_{ijt}$

Because $\boldsymbol{\mathbf{Q}}(\boldsymbol{\mathbf{y}},t)$ is skew-
symmetric, we have that
$\sum_{i}\partial_{i}\sum_{j}\partial_{j}[Q_{ijt}q]=0$, leaving

$\displaystyle\partial_{t}q_{t}$
$\displaystyle=\sum_{i}\partial_{i}\Bigg{[}\sum_{j}[D_{ijt}+Q_{ijt}][q\partial_{j}H]\Bigg{]}+\sum_{i}\partial_{i}\Bigg{[}\sum_{j}\partial_{j}[D_{ijt}q]-q\sum_{j}\partial_{j}D_{ijt}-q\sum_{j}\partial_{j}Q_{ijt}\Bigg{]}$
$\displaystyle=\sum_{i}\partial_{i}\Bigg{[}\sum_{j}[D_{ijt}+Q_{ijt}][\partial_{j}H]q\Bigg{]}+\sum_{i}\partial_{i}\Bigg{[}\sum_{j}\partial_{j}[D_{ijt}q]-q\sum_{j}\partial_{j}(D_{ijt}+Q_{ijt})\Bigg{]}$
$\displaystyle=\sum_{i}\partial_{i}\Bigg{[}\Big{(}\sum_{j}[D_{ijt}+Q_{ijt}][\partial_{j}H]-\sum_{j}\partial_{j}(D_{ijt}+Q_{ijt})\Big{)}q\Bigg{]}+\sum_{i}\sum_{j}\frac{\partial^{2}}{\boldsymbol{\mathbf{y}}_{i}\boldsymbol{\mathbf{y}}_{j}}(D_{ijt}q)$

Recalling that $f_{i}(\boldsymbol{\mathbf{y}},t)=\Big{(}-[D+Q]\nabla
H+\Gamma\Big{)}_{i}$ and again that $[Ax]_{i}=\sum_{j}A_{ij}x_{j}$, we have
equality with the original Fokker-Planck

$\displaystyle=\sum_{i}\partial_{i}\Bigg{[}\Big{(}\sum_{j}[D_{ijt}+Q_{ijt}][\partial_{j}H]-\sum_{j}\partial_{j}(D_{ijt}+Q_{ijt})\Big{)}q\Bigg{]}+\sum_{ij}\frac{\partial^{2}}{\boldsymbol{\mathbf{y}}_{i}\boldsymbol{\mathbf{y}}_{j}}(D_{ijt}q)$
$\displaystyle=-\sum_{i}\frac{\partial}{\partial\boldsymbol{\mathbf{y}}_{i}}\Big{[}f_{i}(\boldsymbol{\mathbf{y}},t)q(\boldsymbol{\mathbf{y}},t)\Big{]}+\sum_{ij}\frac{\partial^{2}}{\boldsymbol{\mathbf{y}}_{i}\boldsymbol{\mathbf{y}}_{j}}\Big{[}\boldsymbol{\mathbf{D}}_{ij}(\boldsymbol{\mathbf{y}},t)q(\boldsymbol{\mathbf{y}},t)\Big{]}$
$\displaystyle=\partial_{t}q(\boldsymbol{\mathbf{y}},t)$

∎

We have shown $\exp[-H(\boldsymbol{\mathbf{y}})]/\boldsymbol{\mathbf{Z}}$ is a
stationary distribution of the time-varying non-linear Itô process:

$\displaystyle
d\boldsymbol{\mathbf{y}}=\Bigg{(}-[\boldsymbol{\mathbf{D}}(\boldsymbol{\mathbf{y}},t)+\boldsymbol{\mathbf{Q}}(\boldsymbol{\mathbf{y}},t)]\nabla
H(\boldsymbol{\mathbf{y}})+\boldsymbol{\Gamma}(\boldsymbol{\mathbf{y}},t)\Bigg{)}dt+\sqrt{2\boldsymbol{\mathbf{D}}(\boldsymbol{\mathbf{y}},t)}\boldsymbol{\mathbf{B}}_{t}.$
(26)

However, for some choices of
$\boldsymbol{\mathbf{Q}},\boldsymbol{\mathbf{D}}$,
$\exp[-H(\boldsymbol{\mathbf{y}})]/\boldsymbol{\mathbf{Z}}$ is not necessarily
the unique stationary distribution. One problematic case can occur as follows.
Suppose that row $i$ of $(\boldsymbol{\mathbf{Q}}+\boldsymbol{\mathbf{D}})$ is
all-zero; in this case, $d\boldsymbol{\mathbf{y}}_{i}=0$ which implies that
$({\mathbf{y}}_{i})_{t}=({\mathbf{y}}_{i})_{0}$ for all $t>0$. Then, the
initial distribution is also a stationary distribution. To rule out such
pathological diffusions, we make the assumption that
$\boldsymbol{\mathbf{Q}}+\boldsymbol{\mathbf{D}}$ is full rank. Then, for
uniqueness, recall that stationary distributions are the zeros of

$\displaystyle\partial_{t}q(\boldsymbol{\mathbf{y}},t)=\nabla\cdot\Bigg{(}\Big{[}\boldsymbol{\mathbf{D}}(\boldsymbol{\mathbf{y}},t)+\boldsymbol{\mathbf{Q}}(\boldsymbol{\mathbf{y}},t)\Big{]}\Big{[}q(\boldsymbol{\mathbf{y}},t)\nabla
H(\boldsymbol{\mathbf{y}})+\nabla q(\boldsymbol{\mathbf{y}},t)\Big{]}\Bigg{)}$

where the expression is of the form
$\boldsymbol{\mathbf{A}}\boldsymbol{\mathbf{v}}$ for
$\boldsymbol{\mathbf{A}}=\boldsymbol{\mathbf{D}}(\boldsymbol{\mathbf{y}},t)+\boldsymbol{\mathbf{Q}}(\boldsymbol{\mathbf{y}},t)$
and

$\displaystyle\boldsymbol{\mathbf{v}}=\Big{[}q(\boldsymbol{\mathbf{y}},t)\nabla
H(\boldsymbol{\mathbf{y}})+\nabla q(\boldsymbol{\mathbf{y}},t)\Big{]}.$

Under the assumption that $\boldsymbol{\mathbf{Q}}+\boldsymbol{\mathbf{D}}$ is
full rank, the expression can only be zero when $\boldsymbol{\mathbf{v}}$ is
zero. To show uniqueness under the full rank assumption, one must then show
that

$\displaystyle\nabla
q(\boldsymbol{\mathbf{y}},t)=-q(\boldsymbol{\mathbf{y}},t)\nabla
H(\boldsymbol{\mathbf{y}}).$

holds only if
$q(\boldsymbol{\mathbf{y}},t)=\exp[-H(\boldsymbol{\mathbf{y}})]/\boldsymbol{\mathbf{Z}}$.
Even if $\exp[-H(\boldsymbol{\mathbf{y}})]/\boldsymbol{\mathbf{Z}}$ is the
unique stationary distribution, convergence to that distribution is a
question. See Zhang & Chen (2013) for more details.

Learning $\boldsymbol{\mathbf{Q}}_{\phi},\boldsymbol{\mathbf{D}}_{\phi}$ in
the mdm elbo helps push $\boldsymbol{\mathbf{y}}_{T}$ to the model prior
$\pi_{\theta}$ and avoid issues like those discussed.

###  B.1 Linear Processes

Next, we specialize this general family to linear Itô processes to maintain
tractable transition distributions. A linear process is one where the drift
$f(\boldsymbol{\mathbf{y}},t)$ and diffusion $g(\boldsymbol{\mathbf{y}},t)$
are linear functions of $\boldsymbol{\mathbf{y}}$. We express the drift
function of a non-linear time-varying Itô process with stationary distribution
proportional to $\exp[-H(\boldsymbol{\mathbf{y}})]$ as

$\displaystyle-(\boldsymbol{\mathbf{Q}}(\boldsymbol{\mathbf{y}},t)+\boldsymbol{\mathbf{D}}(\boldsymbol{\mathbf{y}},t))\nabla
H(\boldsymbol{\mathbf{y}})+\Gamma(\boldsymbol{\mathbf{y}},t).$

Next, linear Itô processes have Gaussian stationary distributions (Särkkä &
Solin, 2019) so $H(\boldsymbol{\mathbf{y}})$ must be quadratic and $\nabla
H(\boldsymbol{\mathbf{y}})$ is linear, and neither are constant in
$\boldsymbol{\mathbf{y}}$. Because $\nabla H(\boldsymbol{\mathbf{y}})$ is
linear, it can be expressed as
$\boldsymbol{\mathbf{S}}\boldsymbol{\mathbf{y}}$ for some matrix
$\boldsymbol{\mathbf{S}}$ where $\boldsymbol{\mathbf{S}}$ is the inverse of
the covariance matrix. Because $\nabla H$ is multiplied by
$\boldsymbol{\mathbf{Q}},\boldsymbol{\mathbf{D}}$, this means that
$\boldsymbol{\mathbf{Q}},\boldsymbol{\mathbf{D}}$ must be free of
$\boldsymbol{\mathbf{y}}$. Recalling that $\Gamma$ is expressed as a sum of
derivatives w.r.t $\boldsymbol{\mathbf{y}}$ of
$\boldsymbol{\mathbf{Q}}+\boldsymbol{\mathbf{D}}$, this means that $\Gamma$
must satisfy $\Gamma=0$. Next, because of the stationary requirement that
$g(t)=\sqrt{2\boldsymbol{\mathbf{D}}(\boldsymbol{\mathbf{y}},t)}$, we can also
conclude by the restriction on $\boldsymbol{\mathbf{D}}$ that the diffusion
coefficient function must be independent of the state
$\boldsymbol{\mathbf{y}}$. Our final form for linear time-varying processes
with stationary distributions $\mathcal{N}(0,\boldsymbol{\mathbf{S}}^{-1})$
is:

$\displaystyle
d\boldsymbol{\mathbf{y}}=\underbrace{-\Big{[}\boldsymbol{\mathbf{Q}}(t)+\boldsymbol{\mathbf{D}}(t)\Big{]}\boldsymbol{\mathbf{S}}\boldsymbol{\mathbf{y}}}_{f(\boldsymbol{\mathbf{y}},t)}dt+\underbrace{\sqrt{2\boldsymbol{\mathbf{D}}(t)}}_{g(t)}d\boldsymbol{\mathbf{B}}_{t}$
(27)

###  B.2 Parameterizing $\boldsymbol{\mathbf{Q}}_{\phi}$

Suppose $b_{q}(s)$ is a positive scalar function defined on the time domain
with known integral. Suppose $\tilde{\boldsymbol{\mathbf{Q}}}_{\phi}$ is any
matrix. Then
$\tilde{\boldsymbol{\mathbf{Q}}}_{\phi}-\tilde{\boldsymbol{\mathbf{Q}}}_{\phi}^{\top}$
is skew-symmetric with
$\tilde{\boldsymbol{\mathbf{Q}}}_{\phi,ij}=-\tilde{\boldsymbol{\mathbf{Q}}}_{\phi,ji}$.
We can set $\boldsymbol{\mathbf{Q}}_{\phi}$ to

$\displaystyle\boldsymbol{\mathbf{Q}}_{\phi}(s)=b_{q}(s)\cdot\Big{[}\tilde{\boldsymbol{\mathbf{Q}}}_{\phi}-\tilde{\boldsymbol{\mathbf{Q}}}_{\phi}^{\top}\Big{]}$
(28)

This is a general parameterization of time-independent skew-symmetric
matrices, which have number of degrees of freedom equal to the number of
entries in one of the triangles of the matrix, excluding the diagonal.

###  B.3 Parameterizing $\boldsymbol{\mathbf{D}}_{\phi}$

Suppose $b_{d}(s)$ is a positive scalar function defined on the time domain
with known integral. Suppose $\tilde{\boldsymbol{\mathbf{D}}}_{\phi}$ is any
matrix. Then
$\tilde{\boldsymbol{\mathbf{D}}}_{\phi}\tilde{\boldsymbol{\mathbf{D}}}_{\phi}^{\top}$
is positive semi-definite and spans all time-independent positive semi-
definite matrices. We can set $\boldsymbol{\mathbf{D}}_{\phi}$ to

$\displaystyle\boldsymbol{\mathbf{D}}_{\phi}(s)=b_{d}(s)\cdot\Big{[}\tilde{\boldsymbol{\mathbf{D}}}_{\phi}\tilde{\boldsymbol{\mathbf{D}}}_{\phi}^{\top}\Big{]}$
(29)

To show
$\tilde{\boldsymbol{\mathbf{D}}}\tilde{\boldsymbol{\mathbf{D}}}^{\top}$ spans
all positive semi-definite matrices: suppose $\boldsymbol{\mathbf{M}}$ is
positive semi-definite. Then it is square. Then it can be eigen-decomposed
into
$\boldsymbol{\mathbf{M}}=\boldsymbol{\mathbf{V}}\boldsymbol{\mathbf{\Sigma}}\boldsymbol{\mathbf{V}}^{\top}$
The degrees of freedom in
$\boldsymbol{\mathbf{V}}\boldsymbol{\mathbf{\Sigma}}\boldsymbol{\mathbf{V}}^{\top}$
are just
$\boldsymbol{\mathbf{R}}=\boldsymbol{\mathbf{V}}\sqrt{\boldsymbol{\mathbf{\Sigma}}}$
since
$\boldsymbol{\mathbf{V}}\boldsymbol{\mathbf{\Sigma}}\boldsymbol{\mathbf{V}}^{\top}=\boldsymbol{\mathbf{R}}\boldsymbol{\mathbf{R}}^{\top}$
and the square root is taken element-wise because
$\boldsymbol{\mathbf{\Sigma}}$ is diagonal and is real because each
$\boldsymbol{\mathbf{\Sigma}}_{ij}\geq 0$, which is true because
$\boldsymbol{\mathbf{M}}$ is positive semi-definite. Take
$\boldsymbol{\mathbf{D}}=\boldsymbol{\mathbf{R}}$.

In our experiments we parameterize ${\mathbf{D}}$ as a diagonal-only matrix.

###  B.4 Integrals

The known integral requirement comes from the integrals required in the
transition kernel, and can be relaxed two possible ways:

  * •

numerical integration of function with unknown integral. This is expected to
have low error given that the function is scalar-in scalar-out.

  * •

Directly parameterize the integral and use auto-grad when needing the
functions not-integrated.

We stick with the known integrals. In conclusion, the underlying parameters
are positive scalar functions $b_{q}(s),b_{d}(s)$ defined on the time domain
and with known integral, and general matrices
$\tilde{\boldsymbol{\mathbf{Q}}}_{\phi},\tilde{\boldsymbol{\mathbf{D}}}_{\phi}$.

###  B.5 Instances

#####  vpsde.

vpsde has $K=1$. Consequently,
$\boldsymbol{\mathbf{Q}},\boldsymbol{\mathbf{D}}$ are $K\times K$. The only
$1\times 1$ skew-symmetric matrix is $0$, so $\boldsymbol{\mathbf{Q}}=0$.
Setting $\boldsymbol{\mathbf{D}}(t)=\frac{1}{2}\beta(t)$ recovers vpsde:

$\displaystyle
d\boldsymbol{\mathbf{y}}=-\frac{\beta(t)}{2}\boldsymbol{\mathbf{y}}dt+\sqrt{\beta(t)}d\boldsymbol{\mathbf{B}}_{t}$
(30)

$\nabla H(\boldsymbol{\mathbf{y}})=\boldsymbol{\mathbf{y}}$ so
$\boldsymbol{\mathbf{H}}(\boldsymbol{\mathbf{y}})=\frac{1}{2}\|\boldsymbol{\mathbf{y}}\|_{2}^{2}$.
The stationary distribution is $\mathcal{N}(0,\boldsymbol{\mathbf{I}})$.

#####  cld.

The cld process (eq $5$ in Dockhorn et al. (2021)) is defined as

$\displaystyle\begin{pmatrix}d{\mathbf{z}}_{t}\\\
d{\mathbf{v}}_{r}\end{pmatrix}=d{\mathbf{y}}_{t}=\begin{pmatrix}0&\frac{\beta}{M}\\\
-\beta&-\frac{\Gamma\beta}{M}\\\
\end{pmatrix}{\mathbf{y}}_{t}+\begin{pmatrix}0&0\\\
0&\sqrt{2\Gamma\beta}\end{pmatrix}d\boldsymbol{\mathbf{B}}_{t}.$

In $\boldsymbol{\mathbf{Q}}/\boldsymbol{\mathbf{D}}$ parameterization, we have

$\displaystyle H({\mathbf{y}})$
$\displaystyle=\frac{1}{2}\left\lVert{\mathbf{z}}\right\rVert_{2}^{2}+\frac{1}{2M}\left\lVert{\mathbf{v}}\right\rVert_{2}^{2},\qquad\nabla_{{\mathbf{u}}}H({\mathbf{y}})=\begin{pmatrix}{\mathbf{z}}\\\
\frac{1}{M}{\mathbf{v}}\end{pmatrix}$ $\displaystyle\boldsymbol{\mathbf{Q}}$
$\displaystyle=\begin{pmatrix}0&-{\beta}\\\
{\beta}&0\end{pmatrix},\qquad\boldsymbol{\mathbf{D}}=\begin{pmatrix}0&0\\\
0&{\Gamma\beta}\end{pmatrix}$

The stationary distribution of this process is:

$\displaystyle
q_{\phi,\infty}\propto\exp(-H({\mathbf{y}}))={\mathcal{N}}({\mathbf{z}};0,I_{d}){\mathcal{N}}({\mathbf{v}};0,MI_{d})$
(31)

#####  alda.

Mou et al. (2019) define a third-order diffusion process for the purpose of
gradient-based mcmc sampling. The alda diffusion process can be specified as

$\displaystyle\boldsymbol{\mathbf{Q}}$
$\displaystyle=\begin{pmatrix}0&-\frac{1}{L}I&0\\\ \frac{1}{L}I&0&-\gamma I\\\
0&\gamma I&0\end{pmatrix},\quad\boldsymbol{\mathbf{D}}=\begin{pmatrix}0&0&0\\\
0&0&0\\\ 0&0&{\frac{\xi}{L}}I\end{pmatrix}.$ (32)

Note that $\boldsymbol{\mathbf{Q}}$ is skew-symmetric and
$\boldsymbol{\mathbf{D}}$ is positive semi-definite, therefore we have that
$q_{t}({\mathbf{u}})\rightarrow q_{\phi,\infty}$. In this case,

$\displaystyle q_{\phi,\infty}$
$\displaystyle={\mathcal{N}}(\boldsymbol{\mathbf{z}};0,\boldsymbol{\mathbf{I}}_{d}){\mathcal{N}}({\mathbf{v}}_{1};0,\frac{1}{L}\boldsymbol{\mathbf{I}}_{d}){\mathcal{N}}({\mathbf{v}}_{2};0,\frac{1}{L}\boldsymbol{\mathbf{I}}_{d})$

#####  malda.

Similar to alda, we specify a diffusion process we term malda which we specify
as

$\displaystyle\boldsymbol{\mathbf{Q}}$
$\displaystyle=\begin{pmatrix}0&-\frac{1}{L}I&-\frac{1}{L}\\\
\frac{1}{L}I&0&-\gamma I\\\ \frac{1}{L}&\gamma
I&0\end{pmatrix},\quad\boldsymbol{\mathbf{D}}=\begin{pmatrix}0&0&0\\\
0&{\frac{1}{L}}I&0\\\ 0&0&{\frac{1}{L}}I\end{pmatrix}.$ (33)

Note that $\boldsymbol{\mathbf{Q}}$ is skew-symmetric and
$\boldsymbol{\mathbf{D}}$ is positive semi-definite. In this case this is

$\displaystyle q_{\phi,\infty}$
$\displaystyle={\mathcal{N}}(\boldsymbol{\mathbf{z}};0,\boldsymbol{\mathbf{I}}_{d}){\mathcal{N}}(\boldsymbol{\mathbf{v}}_{1};0,\frac{1}{L}I_{d}){\mathcal{N}}(\boldsymbol{\mathbf{v}}_{2};0,\frac{1}{L}I_{d})$

##  Appendix C Transitions for linear processes

For time variable $s$ and Brownian motion
$\widehat{\boldsymbol{\mathbf{B}}}_{s}$ driving diffusions of the form

$\displaystyle
d\boldsymbol{\mathbf{y}}=f(\boldsymbol{\mathbf{y}},s)ds+g(s)d\widehat{\boldsymbol{\mathbf{B}}}_{s},$
(34)

when $f_{\phi}(\boldsymbol{\mathbf{y}}_{s},s),g_{\phi}(s)$ are linear, the
transition kernel
$q_{\phi}(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})$ is always
normal (Särkkä & Solin, 2019). Therefore, we just find the mean
$\boldsymbol{\mathbf{m}}_{s|0}$ and covariance
$\boldsymbol{\mathbf{\Sigma}}_{s|0}$ of
$q(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})$. Let
$f(\boldsymbol{\mathbf{y}},s)=\boldsymbol{\mathbf{A}}(s)\boldsymbol{\mathbf{y}}$.
The un-conditional time $s$ mean and covariance are solutions to

$\displaystyle\begin{split}d\boldsymbol{\mathbf{m}}_{s}/ds&=\boldsymbol{\mathbf{A}}(s)\boldsymbol{\mathbf{m}}_{s}\\\
d\boldsymbol{\Sigma}_{s}/ds&=\boldsymbol{\mathbf{A}}(s)\boldsymbol{\Sigma}_{s}+\boldsymbol{\Sigma}_{s}\boldsymbol{\mathbf{A}}^{\top}(s)+g^{2}(s)\end{split}$
(35)

By (6.6) in Särkkä & Solin (2019), for computing conditionals
$q(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})$, we can take the
marginal distribution ODEs and compute conditionals by simply setting the time
$0$ mean and covariance initial conditions to the conditioning value and to
$\boldsymbol{\mathbf{0}}$ respectively. We take (6.36-6.39) and set
$\boldsymbol{\mathbf{m}}_{0}=\boldsymbol{\mathbf{u}}_{0}$ and
$\boldsymbol{\mathbf{\Sigma}}_{0}=0$ to condition. Let
$[\boldsymbol{\mathbf{A}}]_{s}=\int_{0}^{s}\boldsymbol{\mathbf{A}}(\nu)d\nu$.
The mean is

$\displaystyle\boldsymbol{\mathbf{m}}_{s|0}=\exp\left[\int_{0}^{s}\boldsymbol{\mathbf{A}}(\nu)d\nu\right]\boldsymbol{\mathbf{y}}_{0}=\exp\Big{(}\Big{[}A\Big{]}_{s}\Big{)}\underbrace{=\exp(s\boldsymbol{\mathbf{A}})\boldsymbol{\mathbf{y}}_{0}}_{\text{no
integration if $\boldsymbol{\mathbf{A}}(\nu)=\boldsymbol{\mathbf{A}}$}},$ (36)

where $\exp$ denotes matrix exponential. (6.36-6.39) state the covariance
$q(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})$ as a matrix
factorization, for which a derivation is provided below
$\boldsymbol{\mathbf{\Sigma}}_{s}=\boldsymbol{\mathbf{C}}_{s}({\boldsymbol{\mathbf{H}}_{s}})^{-1}$
for $\boldsymbol{\mathbf{C}}_{s},\boldsymbol{\mathbf{H}}_{s}$ being the
solutions of:

$\displaystyle\begin{pmatrix}\frac{d}{ds}\boldsymbol{\mathbf{C}}_{s}\\\
\frac{d}{ds}\boldsymbol{\mathbf{H}}_{s}\end{pmatrix}=\begin{pmatrix}\boldsymbol{\mathbf{A}}(s)&g^{2}(s)\\\
\boldsymbol{\mathbf{0}}&-\boldsymbol{\mathbf{A}}^{\top}(s)\end{pmatrix}\begin{pmatrix}\boldsymbol{\mathbf{C}}_{s}\\\
\boldsymbol{\mathbf{H}}_{s}\end{pmatrix}$ (37)

To condition and get $\boldsymbol{\mathbf{\Sigma}}_{s|0}$ from
$\boldsymbol{\mathbf{\Sigma}}_{s}$, we set
$\boldsymbol{\mathbf{\Sigma}}_{0}=\boldsymbol{\mathbf{0}}$, and initialize
${\mathbf{C}}_{s},{\mathbf{H}}_{s}$ by
$\boldsymbol{\mathbf{C}}_{0}=\boldsymbol{\mathbf{0}}$ and
$\boldsymbol{\mathbf{H}}_{0}=\boldsymbol{\mathbf{I}}$.

$\displaystyle\begin{pmatrix}\boldsymbol{\mathbf{C}}_{s}\\\
\boldsymbol{\mathbf{H}}_{s}\end{pmatrix}=\exp\Bigg{[}\begin{pmatrix}[\boldsymbol{\mathbf{A}}]_{s}&[g^{2}]_{s}\\\
\boldsymbol{\mathbf{0}}&-[\boldsymbol{\mathbf{A}}^{\top}]_{s}\end{pmatrix}\Bigg{]}\begin{pmatrix}\boldsymbol{\mathbf{0}}\\\
\boldsymbol{\mathbf{I}}\end{pmatrix}\underbrace{=\exp\Bigg{[}s\begin{pmatrix}\boldsymbol{\mathbf{A}}&g^{2}\\\
\boldsymbol{\mathbf{0}}&-\boldsymbol{\mathbf{A}}^{\top}\end{pmatrix}\Bigg{]}\begin{pmatrix}\boldsymbol{\mathbf{0}}\\\
\boldsymbol{\mathbf{I}}\end{pmatrix}}_{\text{no integration if
$\boldsymbol{\mathbf{A}}(\nu)=\boldsymbol{\mathbf{A}},g(\nu)=g$}}.$ (38)

Finally,
$\boldsymbol{\mathbf{\Sigma}}_{s|0}=\boldsymbol{\mathbf{C}}_{s}({\boldsymbol{\mathbf{H}}_{s}})^{-1}$.

###  C.1 Derivation of the Covariance matrix solution

Equation 35 gives an expression for $d\boldsymbol{\Sigma}_{s}/ds$. To derive
the matrix factorization technique used in eq. 37, we use eq. 35 and the
desired condition
$\boldsymbol{\Sigma}_{s}=\boldsymbol{\mathbf{C}}_{s}\boldsymbol{\mathbf{H}}_{s}^{-1}$
to derive expressions for $d\boldsymbol{\mathbf{C}}_{s}/ds$ and
$d\boldsymbol{\mathbf{H}}_{s}/ds$ and suitable intial conditions so that the
factorization also starts at the desired $\boldsymbol{\Sigma}_{0}$. Let
$\boldsymbol{\mathbf{\Sigma}}_{s}={\mathbf{C}}_{s}{\mathbf{H}}_{s}^{-1}$, then
note that ${\mathbf{C}}_{s},{\mathbf{H}}_{s}$ satisfies

$\displaystyle\frac{d}{ds}\boldsymbol{\mathbf{\Sigma}}_{s}$
$\displaystyle=\frac{d}{ds}{\mathbf{C}}_{s}{\mathbf{H}}_{s}^{-1}$
$\displaystyle={\mathbf{C}}_{s}\frac{d}{ds}{\mathbf{H}}_{s}^{-1}+\left(\frac{d}{ds}{\mathbf{C}}_{s}\right){\mathbf{H}}_{s}^{-1}$

And using the fact that

$\displaystyle\frac{d}{ds}{\mathbf{H}}_{s}{\mathbf{H}}_{s}^{-1}$
$\displaystyle=0$
$\displaystyle{\mathbf{H}}_{s}\frac{d}{ds}{\mathbf{H}}_{s}^{-1}+\frac{d}{ds}{\mathbf{H}}_{s}\left({\mathbf{H}}_{s}^{-1}\right)$
$\displaystyle=0$ $\displaystyle\frac{d}{ds}{\mathbf{H}}_{s}^{-1}$
$\displaystyle=-{\mathbf{H}}_{s}^{-1}\frac{d}{ds}{\mathbf{H}}_{s}\left({\mathbf{H}}_{s}^{-1}\right)$

we get that

$\displaystyle{\mathbf{C}}_{s}\frac{d}{ds}{\mathbf{H}}_{s}^{-1}+\left(\frac{d}{ds}{\mathbf{C}}_{s}\right){\mathbf{H}}_{s}^{-1}$
$\displaystyle=-{\mathbf{C}}_{s}{\mathbf{H}}_{s}^{-1}\frac{d}{ds}{\mathbf{H}}_{s}\left({\mathbf{H}}_{s}^{-1}\right)+\left(\frac{d}{ds}{\mathbf{C}}_{s}\right){\mathbf{H}}_{s}^{-1}$
$\displaystyle-{\mathbf{C}}_{s}{\mathbf{H}}_{s}^{-1}\frac{d}{ds}{\mathbf{H}}_{s}\left({\mathbf{H}}_{s}^{-1}\right)+\left(\frac{d}{ds}{\mathbf{C}}_{s}\right){\mathbf{H}}_{s}^{-1}$
$\displaystyle={\mathbf{A}}(s){\mathbf{C}}_{s}{\mathbf{H}}_{s}^{-1}+{\mathbf{C}}_{s}{\mathbf{H}}_{s}^{-1}{\mathbf{A}}^{\top}(s)+g^{2}(s)$
$\displaystyle={\mathbf{A}}(s){\mathbf{C}}_{s}{\mathbf{H}}_{s}^{-1}+{\mathbf{C}}_{s}{\mathbf{H}}_{s}^{-1}{\mathbf{A}}^{\top}(s){\mathbf{H}}_{s}{\mathbf{H}}_{s}^{-1}+g^{2}(s){\mathbf{H}}_{s}{\mathbf{H}}_{s}^{-1}$
$\displaystyle\left(-{\mathbf{C}}_{s}{\mathbf{H}}_{s}^{-1}\frac{d}{ds}{\mathbf{H}}_{s}+\frac{d}{ds}{\mathbf{C}}_{s}\right){\mathbf{H}}_{s}^{-1}$
$\displaystyle=\left({\mathbf{A}}(s){\mathbf{C}}_{s}+{\mathbf{C}}_{s}{\mathbf{H}}_{s}^{-1}{\mathbf{A}}^{\top}(s){\mathbf{H}}_{s}+g^{2}(s){\mathbf{H}}_{s}\right){\mathbf{H}}_{s}^{-1}$
$\displaystyle-{\mathbf{C}}_{s}{\mathbf{H}}_{s}^{-1}\frac{d}{ds}{\mathbf{H}}_{s}+\frac{d}{ds}{\mathbf{C}}_{s}$
$\displaystyle={\mathbf{A}}(s){\mathbf{C}}_{s}+{\mathbf{C}}_{s}{\mathbf{H}}_{s}^{-1}{\mathbf{A}}^{\top}(s){\mathbf{H}}_{s}+g^{2}(s){\mathbf{H}}_{s}$
$\displaystyle\left[{\mathbf{C}}_{s}{\mathbf{H}}_{s}^{-1}\quad{\mathbf{I}}_{d}\right]^{\top}\frac{d}{ds}\begin{pmatrix}{\mathbf{H}}_{s}\\\
{\mathbf{C}}_{s}\end{pmatrix}$
$\displaystyle=\left[{\mathbf{C}}_{s}{\mathbf{H}}_{s}^{-1}\quad{\mathbf{I}}_{d}\right]^{\top}\begin{pmatrix}-{\mathbf{A}}^{\top}(s){\mathbf{H}}_{s}\\\
{\mathbf{A}}(s){\mathbf{C}}_{s}+g^{2}(s){\mathbf{H}}_{s}\end{pmatrix}$

Now, we note ${\mathbf{C}}_{s},{\mathbf{H}}_{s}$ satisfy the following

$\displaystyle\frac{d}{ds}{\mathbf{H}}_{s}$
$\displaystyle=-{\mathbf{A}}^{\top}(s){\mathbf{H}}_{s}$
$\displaystyle\frac{d}{ds}{\mathbf{C}}_{s}$
$\displaystyle={\mathbf{A}}(s){\mathbf{C}}_{s}+g^{2}(s){\mathbf{H}}_{s}$

which implies that

$\displaystyle\frac{d}{ds}\begin{pmatrix}\boldsymbol{\mathbf{C}}_{s}\\\
\boldsymbol{\mathbf{H}}_{s}\end{pmatrix}=\begin{pmatrix}\boldsymbol{\mathbf{A}}(s)&g^{2}(s)\\\
\boldsymbol{\mathbf{0}}&-\boldsymbol{\mathbf{A}}^{\top}(s)\end{pmatrix}\begin{pmatrix}\boldsymbol{\mathbf{C}}_{s}\\\
\boldsymbol{\mathbf{H}}_{s}\end{pmatrix}$ (39)

with ${\mathbf{C}}_{0}=\boldsymbol{\mathbf{\Sigma}}_{0}$ and
${\mathbf{H}}_{0}={\mathbf{I}}_{d}$, as
${\mathbf{C}}_{0}{\mathbf{H}}_{0}^{-1}=\boldsymbol{\mathbf{\Sigma}}_{0}$.

###  C.2 Hybrid Score Matching

Instead of computing
$q(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})$, we can apply the
hybrid score matching principle (Dockhorn et al., 2021) to reduce variance by
compute objectives using $q(\boldsymbol{\mathbf{y}}_{s}|x)$ instead of
$q(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})$, which amounts to
integrating out $\boldsymbol{\mathbf{v}}_{0}$. To accomplish this, following
Särkkä & Solin (2019), we simply replace $\boldsymbol{\mathbf{y}}_{0}$ with
$[x,\mathbb{E}[\boldsymbol{\mathbf{v}}_{0}]]$ in the expression for
$\boldsymbol{\mathbf{m}}_{s|0}$, i.e. replace the conditioning value of
$\boldsymbol{\mathbf{v}}_{0}$ with the mean of its chosen initial
distribution:

$\displaystyle\mathbb{E}[\boldsymbol{\mathbf{y}}_{s}|x]=\exp\Bigg{[}\int_{0}^{s}A(\nu)d\nu\Bigg{]}\begin{pmatrix}x\\\
\mathbb{E}[\boldsymbol{\mathbf{v}}_{0}]\end{pmatrix}$ (40)

For the convariance, instead of using
$\boldsymbol{\mathbf{C}}_{0}=\boldsymbol{\mathbf{\Sigma}}_{0}=\boldsymbol{\mathbf{0}}$,
we use a block matrix to condition on $x$ but not
$\boldsymbol{\mathbf{v}}_{0}$. We decompose $\boldsymbol{\mathbf{\Sigma}}_{0}$
into its blocks $\boldsymbol{\mathbf{\Sigma}}_{0,xx}$,
$\boldsymbol{\mathbf{\Sigma}}_{0,vv}$ ,$\boldsymbol{\mathbf{\Sigma}}_{0,xv}$.
As before, to condition on $x$ we set
$\boldsymbol{\mathbf{\Sigma}}_{0,xx}=\boldsymbol{\mathbf{0}}$. Because
$q(\boldsymbol{\mathbf{v}}_{0})$ is set to be independent of $x$,
$\boldsymbol{\mathbf{\Sigma}}_{0,xv}$ is also set to
$\boldsymbol{\mathbf{0}}$. Finally, instead of $\boldsymbol{\mathbf{0}}$, to
marginalize out $\boldsymbol{\mathbf{v}}_{0}$,
$\boldsymbol{\mathbf{\Sigma}}_{0,vv}$ is set to the covariance of the chosen
initial time zero distribution for $\boldsymbol{\mathbf{v}}_{0}$. E.g. if
$\boldsymbol{\mathbf{v}}_{0,j}\sim N(0,\gamma)$ for each dimension, then
$\boldsymbol{\mathbf{\Sigma}}_{0,vv}=N(0,\gamma I)$.

We operationalize this in a simple piece of code, which makes the elbo
tractable and easy, i.e. skips both analytic derivations and numerical forward
integration during training.

###  C.3 Transitions in Stationary Parameterization

In terms of $\boldsymbol{\mathbf{Q}},\boldsymbol{\mathbf{D}}$, the transitions
$q(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})$ for time $s$ are
normal with mean $\boldsymbol{\mathbf{m}}_{s|0}$ and
$\boldsymbol{\mathbf{\Sigma}}_{s|0}$ equal to:

$\displaystyle\boldsymbol{\mathbf{m}}_{s|0}=\exp\Big{(}-\Big{[}\boldsymbol{\mathbf{Q}}+\boldsymbol{\mathbf{D}}\Big{]}_{s}\Big{)}\boldsymbol{\mathbf{y}}_{0},\quad\quad\begin{pmatrix}\boldsymbol{\mathbf{C}}_{s}\\\
\boldsymbol{\mathbf{H}}_{s}\end{pmatrix}=\exp\Bigg{[}\begin{pmatrix}-[\boldsymbol{\mathbf{Q}}+\boldsymbol{\mathbf{D}}]_{s}&[2\boldsymbol{\mathbf{D}}]_{s}\\\
\boldsymbol{\mathbf{0}}&[(\boldsymbol{\mathbf{Q}}+\boldsymbol{\mathbf{D}})^{\top}]_{s}\end{pmatrix}\Bigg{]}\begin{pmatrix}\boldsymbol{\mathbf{0}}\\\
\boldsymbol{\mathbf{I}}\end{pmatrix}$ (41)

where
$\boldsymbol{\mathbf{\Sigma}}_{s|0}=\boldsymbol{\mathbf{C}}_{s}({\boldsymbol{\mathbf{H}}_{s}})^{-1}$.
For the time invariant case, this simplifies to

$\displaystyle\boldsymbol{\mathbf{m}}_{s|0}=\exp[-s(\boldsymbol{\mathbf{Q}}+\boldsymbol{\mathbf{D}})]\boldsymbol{\mathbf{y}}_{0},\quad\quad\begin{pmatrix}\boldsymbol{\mathbf{C}}_{s}\\\
\boldsymbol{\mathbf{H}}_{s}\end{pmatrix}=\exp\Bigg{[}s\begin{pmatrix}-(\boldsymbol{\mathbf{Q}}+\boldsymbol{\mathbf{D}})&2\boldsymbol{\mathbf{D}}\\\
\boldsymbol{\mathbf{0}}&(\boldsymbol{\mathbf{Q}}+\boldsymbol{\mathbf{D}})^{\top}\end{pmatrix}\Bigg{]}\begin{pmatrix}\boldsymbol{\mathbf{0}}\\\
\boldsymbol{\mathbf{I}}\end{pmatrix}$ (42)

##  Appendix D Generic change of measure and Jensen’s for approximate
marginalization

Suppose
$\boldsymbol{\mathbf{u}}=[\boldsymbol{\mathbf{z}},\boldsymbol{\mathbf{v}}]$
and we have an expression for
$p(\boldsymbol{\mathbf{u}}=[z,v])=p(\boldsymbol{\mathbf{z}}=z,\boldsymbol{\mathbf{v}}=v)$.
By marginalization, we can get $p(\boldsymbol{\mathbf{z}}=z)$, and we can
introduce another distribution $q$ to pick a sampling distribution of our
choice:

$\displaystyle\begin{split}p(\boldsymbol{\mathbf{z}}=z)&=\int_{v}p(\boldsymbol{\mathbf{z}}=z,\boldsymbol{\mathbf{v}}=v)dv\\\
&=\int_{v}p(\boldsymbol{\mathbf{z}}=z|\boldsymbol{\mathbf{v}}=v)p(\boldsymbol{\mathbf{v}}=v)dv\\\
&=\int_{v}\frac{q(\boldsymbol{\mathbf{v}}=v|\boldsymbol{\mathbf{z}}=z)}{q(\boldsymbol{\mathbf{v}}=v|\boldsymbol{\mathbf{z}}=z)}p(\boldsymbol{\mathbf{z}}=z|\boldsymbol{\mathbf{v}}=v)p(\boldsymbol{\mathbf{v}}=v)dv\\\
&=\mathbb{E}_{q(\boldsymbol{\mathbf{v}}=v|\boldsymbol{\mathbf{z}}=z)}\Big{[}\frac{p(\boldsymbol{\mathbf{z}}=z,\boldsymbol{\mathbf{v}}=v)}{q(\boldsymbol{\mathbf{v}}=v|\boldsymbol{\mathbf{z}}=z)}\Big{]}\end{split}$
(43)

We often work with these expressions in log space, and need to pull the
expectation outside to use Monte Carlo. Jensen’s bound allows this:

$\displaystyle\log p(\boldsymbol{\mathbf{z}}=z)$
$\displaystyle=\log\mathbb{E}_{q(\boldsymbol{\mathbf{v}}=v|\boldsymbol{\mathbf{z}}=z)}\Big{[}\frac{p(\boldsymbol{\mathbf{z}}=z,\boldsymbol{\mathbf{v}}=v)}{q(\boldsymbol{\mathbf{v}}=v|\boldsymbol{\mathbf{z}}=z)}\Big{]}$
$\displaystyle\geq\mathbb{E}_{q(\boldsymbol{\mathbf{v}}=v|\boldsymbol{\mathbf{z}}=z)}\Big{[}\log\frac{p(\boldsymbol{\mathbf{z}}=z,\boldsymbol{\mathbf{v}}=v)}{q(\boldsymbol{\mathbf{v}}=v|\boldsymbol{\mathbf{z}}=z)}\Big{]}$

The following shows that the bound is tight when
$q(\boldsymbol{\mathbf{v}}=v|\boldsymbol{\mathbf{z}}=z)=p(\boldsymbol{\mathbf{v}}=v|\boldsymbol{\mathbf{z}}=z)$:

$\displaystyle\begin{split}\mathbb{E}_{q(\boldsymbol{\mathbf{v}}=v|\boldsymbol{\mathbf{z}}=z)}\Big{[}\log\frac{p(\boldsymbol{\mathbf{z}}=z,\boldsymbol{\mathbf{v}}=v)}{q(\boldsymbol{\mathbf{v}}=v|\boldsymbol{\mathbf{z}}=z)}\Big{]}&=_{\text{assume}}\mathbb{E}_{p(\boldsymbol{\mathbf{v}}=v|\boldsymbol{\mathbf{z}}=z)}\Big{[}\log\frac{p(\boldsymbol{\mathbf{z}}=z,\boldsymbol{\mathbf{v}}=v)}{p(\boldsymbol{\mathbf{v}}=v|\boldsymbol{\mathbf{z}}=z)}\Big{]}\\\
&=\mathbb{E}_{p(\boldsymbol{\mathbf{v}}=v|\boldsymbol{\mathbf{z}}=z)}\Big{[}\log\Big{(}\frac{p(\boldsymbol{\mathbf{z}}=z,\boldsymbol{\mathbf{v}}=v)}{p(\boldsymbol{\mathbf{v}}=v,\boldsymbol{\mathbf{z}}=z)}\cdot
p(\boldsymbol{\mathbf{z}}=z)\Big{)}\Big{]}\\\
&=\mathbb{E}_{p(\boldsymbol{\mathbf{v}}=v|\boldsymbol{\mathbf{z}}=z)}\Big{[}\log
p(\boldsymbol{\mathbf{z}}=z)\Big{]}\\\ &=\log
p(\boldsymbol{\mathbf{z}}=z)\end{split}$ (44)

##  Appendix E elbo for mdms

$\displaystyle\log p_{\theta}(x)$
$\displaystyle=\log\int_{v_{0}}p_{\theta}(x_{0},v_{0})dv_{0}$ (45)
$\displaystyle=\log\int_{v_{0}}p_{\theta}(u_{0}=[x,v_{0}])$ (46)
$\displaystyle=\log\int_{v_{0}}\frac{q(v_{0}|x)}{q(v_{0}|x)}p_{\theta}(u_{0}=[x,v_{0}])$
(47)
$\displaystyle=\log\mathbb{E}_{q(v_{0}|x)}\Bigg{[}\frac{p_{\theta}(u_{0}=[x,v_{0}])}{q(v_{0}|x)}\Bigg{]}$
(48) $\displaystyle\geq\mathbb{E}_{q(v_{0}|x)}\Bigg{[}\log
p_{\theta}(u_{0}=[x,v_{0}])-\log q(v_{0}|x)\Bigg{]}$ (49)
$\displaystyle\geq\mathbb{E}_{q(y|x)}\Bigg{[}\log\pi_{\theta}(y_{T})+\int_{0}^{T}-\|s_{\theta}\|_{g^{2}}^{2}-\nabla\cdot(g^{2}s_{\theta}-f)ds-\log
q(y^{v}_{0}|x)\Bigg{]}$ (50)

The first inequality holds due to Jensen’s inequality and the second due to an
application of Theorem 1 from Huang et al. (2021) or Theorem 3 from Song et
al. (2021) applied to the joint variable $\boldsymbol{\mathbf{u}}_{0}$.

###  E.1 ISM to DSM

####  E.1.1 Lemma: expectation by parts

We will need a form of multivariate integration by parts which gives us for
some $f$ and some $q(x)$, $E_{q(x)}[\nabla_{x}\cdot
f(x)]=-E_{q(x)}[f(x)^{\top}\nabla_{x}\log q(x)]$

$\displaystyle E_{q(x)}[\nabla_{x}\cdot f_{i}(x)]$ $\displaystyle=\int
q(x)\sum_{i=1}^{d}[\nabla_{x_{i}}f_{i}(x)]dx$
$\displaystyle=\int\sum_{i=1}^{d}q(x)\nabla_{x_{i}}f_{i}(x)dx$
$\displaystyle=\sum_{i=1}^{d}\int_{x_{-i}}\int_{x_{i}}q(x)\nabla_{x_{i}}f_{i}(x)dx_{i}dx_{-i}$
$\displaystyle=\sum_{i=1}^{d}\int\Bigg{[}\Big{[}q(x)\int\nabla_{x_{i}}f_{i}(x)dx_{i}\Big{]}_{-\infty}^{\infty}-\int\nabla_{x_{i}}q(x)\int\nabla_{x_{i}}f_{i}(x)dx_{i}]\Bigg{]}dx_{-i}$
$\displaystyle=\sum_{i=1}^{d}\int\Bigg{[}-\int\nabla_{x_{i}}q(x)f_{i}(x)dx_{i}\Bigg{]}dx_{-i}$
$\displaystyle=\sum_{i=1}^{d}\int\Bigg{[}-\int q(x)\nabla_{x_{i}}\log
q(x)f_{i}(x)dx_{i}\Bigg{]}dx_{-i}$ $\displaystyle=\sum_{i=1}^{d}-\int\int
q(x)\nabla_{x_{i}}\log q(x)f_{i}(x)dx_{i}dx_{-i}$
$\displaystyle=\sum_{i=1}^{d}-E_{q(x)}\Big{[}\nabla_{x_{i}}\log
q(x)f_{i}(x)\Big{]}$ $\displaystyle=-E_{q(x)}[f(x)^{\top}\nabla_{x}\log q(x)]$

This equality also follows directly from the Stein operator using the
generator method to the Langevin diffusion (Barbour, 1988).

####  E.1.2 DSM Elbo

Using the “expectation by parts”, we have:

$\displaystyle\mathbb{E}_{q(u_{t}|x)}[\nabla_{u_{t}}\cdot
g^{2}(t)s_{\theta}(u_{t},t)]=-\mathbb{E}_{q(u_{t}|x)}[(g^{2}(t)s_{\theta}(u_{t},t))^{\top}\nabla_{u_{t}}\log
q(u_{t}|x)]$

Also we have, for $s_{\theta}$ evaluated at $(u_{t},t)$, by completing the
square,

$\displaystyle-\frac{1}{2}||s_{\theta}||_{g^{2}(t)}+s_{\theta}^{\top}g^{2}(t)\nabla\log
q(u_{t}|x)=-\frac{1}{2}||s_{\theta}-\nabla\log
q(u_{t}|x)||_{g^{2}(t)}^{2}+.5||\nabla\log q(u_{t}|x)||_{g^{2}(t)}^{2}$

The two together give us:

$\displaystyle\begin{split}\log
p(x)&\geq\mathbb{E}_{q(u_{T}|x)}\Bigg{[}\log\pi\Bigg{]}+\int_{0}^{T}\Bigg{[}\mathbb{E}_{q(u_{t}|x)}\Big{[}-\nabla\cdot
g^{2}s_{\theta}-.5||s_{\theta}||_{g^{2}(t)}^{2}+\nabla\cdot
f\Big{]}dt\Bigg{]}\\\
&=\mathbb{E}_{q(u_{T}|x)}\Bigg{[}\log\pi\Bigg{]}+\int_{0}^{T}\Bigg{[}\mathbb{E}_{q(u_{t}|x)}\Big{[}(g^{2}s_{\theta})^{\top}\nabla_{u_{t}}\log
q(u_{t}|x)-.5||s_{\theta}||_{g^{2}(t)}^{2}+\nabla\cdot f\Big{]}dt\Bigg{]}\\\
&=\mathbb{E}_{q(u_{T}|x)}\Bigg{[}\log\pi\Bigg{]}+\int_{0}^{T}\Bigg{[}\mathbb{E}_{q(u_{t}|x)}\Big{[}-\frac{1}{2}||s_{\theta}-\nabla\log
q(u_{t}|x)||_{g^{2}(t)}^{2}\\\ &\hskip 180.00027pt+.5||\nabla\log
q(u_{t}|x)||_{g^{2}(t)}^{2}+\nabla_{u_{t}}\cdot f\Big{]}\Bigg{]}dt\end{split}$
(51)

###  E.2 Noise prediction

We have that for normal
$\mathcal{N}(\boldsymbol{\mathbf{y}}_{s};\boldsymbol{\mathbf{m}}_{s|0},\boldsymbol{\Sigma}_{s|0})$,
we can sample $\boldsymbol{\mathbf{y}}_{s}$ with normal noise
$\epsilon\sim\mathcal{N}(0,I)$ and
$\boldsymbol{\mathbf{y}}_{s}=\boldsymbol{\mathbf{m}}_{s|0}+\boldsymbol{\mathbf{L}}\epsilon$
where $\boldsymbol{\mathbf{L}}$ is the cholesky decomposition of
$\boldsymbol{\Sigma}_{s|0}$ Then, the score is

$\displaystyle\nabla_{\boldsymbol{\mathbf{y}}_{s}}$ $\displaystyle\log
q(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})\Bigg{|}_{\boldsymbol{\mathbf{y}}_{s}=\boldsymbol{\mathbf{m}}_{s|0}+\boldsymbol{\mathbf{L}}\epsilon}$
$\displaystyle=-\boldsymbol{\Sigma}_{s|0}^{-1}\Big{(}\boldsymbol{\mathbf{y}}_{s}-\boldsymbol{\mathbf{m}}_{s|0}\Big{)}$
$\displaystyle=-\boldsymbol{\Sigma}_{s|0}^{-1}\Big{(}\Big{[}\boldsymbol{\mathbf{m}}_{s|0}+\boldsymbol{\mathbf{L}}\epsilon\Big{]}-\boldsymbol{\mathbf{m}}_{s|0}\Big{)}$
$\displaystyle=-\boldsymbol{\Sigma}_{s|0}^{-1}\Big{(}\boldsymbol{\mathbf{L}}\epsilon\Big{)}$
$\displaystyle=-\Big{(}\boldsymbol{\mathbf{L}}\boldsymbol{\mathbf{L}}^{\top}\Big{)}^{-1}\Big{(}\boldsymbol{\mathbf{L}}\epsilon\Big{)}$
$\displaystyle=-\Big{(}\boldsymbol{\mathbf{L}}^{\top}\Big{)}^{-1}\boldsymbol{\mathbf{L}}^{-1}\boldsymbol{\mathbf{L}}\epsilon$
$\displaystyle=-\Big{(}\boldsymbol{\mathbf{L}}^{\top}\Big{)}^{-1}\epsilon=-\Big{(}\boldsymbol{\mathbf{L}}^{-1}\Big{)}^{\top}\epsilon=-\boldsymbol{\mathbf{L}}^{\top,-1}\epsilon$

Parameterize $s_{\theta}(\boldsymbol{\mathbf{y}}_{s},s)$ as
$s_{\theta}(\boldsymbol{\mathbf{y}}_{s},s)=-\boldsymbol{\mathbf{L}}^{\top,-1}\epsilon_{\theta}(\boldsymbol{\mathbf{y}},s)$.
This gives

$\displaystyle\frac{1}{2}\|-\boldsymbol{\mathbf{L}}^{\top,-1}\epsilon_{\theta}(\boldsymbol{\mathbf{y}},s)\quad-\quad-\boldsymbol{\mathbf{L}}^{\top,-1}\epsilon\|_{g_{\phi}^{2}(s)}^{2}$
$\displaystyle=\frac{1}{2}\|\boldsymbol{\mathbf{L}}^{\top,-1}\epsilon\quad-\quad\boldsymbol{\mathbf{L}}^{\top,-1}\epsilon_{\theta}(\boldsymbol{\mathbf{y}},s)\|_{g_{\phi}^{2}(s)}^{2}$
$\displaystyle=\frac{1}{2}\Bigg{(}\boldsymbol{\mathbf{L}}^{\top,-1}\epsilon\quad-\quad\boldsymbol{\mathbf{L}}^{\top,-1}\epsilon_{\theta}(\boldsymbol{\mathbf{y}},s)\Bigg{)}^{\top}g_{\phi}^{2}(s)\Bigg{(}\boldsymbol{\mathbf{L}}^{\top,-1}\epsilon\quad-\quad\boldsymbol{\mathbf{L}}^{\top,-1}\epsilon_{\theta}(\boldsymbol{\mathbf{y}},s)\Bigg{)}$
$\displaystyle=\frac{1}{2}\Bigg{(}\boldsymbol{\mathbf{L}}^{\top,-1}\Big{[}\epsilon-\epsilon_{\theta}(\boldsymbol{\mathbf{y}},s)\Big{]}\Bigg{)}^{\top}g_{\phi}^{2}(s)\Bigg{(}\boldsymbol{\mathbf{L}}^{\top,-1}\Big{[}\epsilon-\epsilon_{\theta}(\boldsymbol{\mathbf{y}},s)\Big{]}\Bigg{)}$

We can also use this insight to analytically compute the quadratic score term
(following is computed per data-dimension, so must be multiplied by $D$ when
computing the elbo):

$\displaystyle\mathbb{E}_{\boldsymbol{\mathbf{y}}_{0}}\mathbb{E}_{\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0}}\Bigg{[}\frac{1}{2}\|\nabla_{\boldsymbol{\mathbf{y}}_{s}}\log
q_{\phi}(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})\|^{2}_{g_{\phi}^{2}(s)}\Bigg{]}$
$\displaystyle=\mathbb{E}_{\boldsymbol{\mathbf{y}}_{0}}\mathbb{E}_{\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0}}\Bigg{[}\Big{(}\nabla_{\boldsymbol{\mathbf{y}}_{s}}\log
q_{\phi}(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})\Big{)}^{\top}g_{\phi}^{2}(s)\Big{(}\nabla_{\boldsymbol{\mathbf{y}}_{s}}\log
q_{\phi}(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})\Big{)}\Bigg{]}$
$\displaystyle=\mathbb{E}_{\boldsymbol{\mathbf{y}}_{0}}\mathbb{E}_{\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0}}\Bigg{[}\Big{(}-\boldsymbol{\mathbf{L}}^{\top,-1}\epsilon\Big{)}^{\top}g_{\phi}^{2}(s)\Big{(}-\boldsymbol{\mathbf{L}}^{\top,-1}\epsilon\Big{)}\Bigg{]}$
$\displaystyle=\mathbb{E}_{\boldsymbol{\mathbf{y}}_{0}}\mathbb{E}_{\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0}}\Bigg{[}\epsilon^{\top}(-\boldsymbol{\mathbf{L}}^{-1})g_{\phi}^{2}(s)(-\boldsymbol{\mathbf{L}}^{\top,-1})\epsilon\Bigg{]}$
$\displaystyle=\mathbb{E}_{\boldsymbol{\mathbf{y}}_{0}}\mathbb{E}_{\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0}}\Bigg{[}\epsilon^{\top}\Big{(}\boldsymbol{\mathbf{L}}^{-1}g_{\phi}^{2}(s)\boldsymbol{\mathbf{L}}^{\top,-1}\Big{)}\epsilon\Bigg{]}$
$\displaystyle=\mathbb{E}_{\boldsymbol{\mathbf{y}}_{0}}\mathbb{E}_{\epsilon}\Bigg{[}\epsilon^{\top}\Big{(}\boldsymbol{\mathbf{L}}^{-1}g_{\phi}^{2}(s)\boldsymbol{\mathbf{L}}^{\top,-1}\Big{)}\epsilon\Bigg{]}$
$\displaystyle=\mathbb{E}_{\epsilon}\Bigg{[}\epsilon^{\top}\Big{(}\boldsymbol{\mathbf{L}}^{-1}g_{\phi}^{2}(s)\boldsymbol{\mathbf{L}}^{\top,-1}\Big{)}\epsilon\Bigg{]}$
$\displaystyle=\text{Trace}\Bigg{(}\boldsymbol{\mathbf{L}}^{-1}g_{\phi}^{2}(s)\boldsymbol{\mathbf{L}}^{\top,-1}\Bigg{)}$

##  Appendix F elbos in stationary parameterization

We use the stationary parmeterization described in appendix B. We now
specialize the elbo to the linear stationary parameterization.

Recall
$f_{\phi}(\boldsymbol{\mathbf{y}},s)=-[\boldsymbol{\mathbf{Q}}_{\phi}(s)+\boldsymbol{\mathbf{D}}_{\phi}(s)]\boldsymbol{\mathbf{y}}$.
Recall $g_{\phi}(s)=\sqrt{2\boldsymbol{\mathbf{D}}_{\phi}(s)}$ We have
$g_{\phi}^{2}(s)=2\boldsymbol{\mathbf{D}}_{\phi}(s)$. We can write the mdm ism
elbo as

$\displaystyle\mathcal{L}^{\text{mism}}=\mathbb{E}_{v\sim
q_{\gamma}}\Bigg{[}\mathbb{E}_{s\sim\text{Unif}(0,T)}\Big{[}\ell_{s}^{(ism)}\Big{]}+\ell_{T}+\ell_{q}\Bigg{]}$
(52)

where

$\displaystyle\begin{split}\ell_{s_{\theta}}&=-\frac{1}{2}\|s_{\theta}(\boldsymbol{\mathbf{y}}_{s},s)\|^{2}_{\underbrace{2\boldsymbol{\mathbf{D}}_{\phi}(s)}_{g^{2}_{\phi}}}\\\
\ell_{\text{div-
fgs}}&=\nabla_{\boldsymbol{\mathbf{y}}_{s}}\cdot\Big{[}\underbrace{-[\boldsymbol{\mathbf{Q}}_{\phi}(s)+\boldsymbol{\mathbf{D}}_{\phi}(s)]\boldsymbol{\mathbf{y}}_{s}}_{f_{\phi}}-\underbrace{2\boldsymbol{\mathbf{D}}_{\phi}(s)}_{g_{\phi}^{2}}s_{\theta}(\boldsymbol{\mathbf{y}}_{s},s)\Big{]}\\\
\ell_{s}^{\text{ism}}&=\mathbb{E}_{\underbrace{q_{\phi,s,(x,v)}}_{\text{depends
on
}\boldsymbol{\mathbf{Q}},\boldsymbol{\mathbf{D}}}}\Big{[}\ell_{s_{\theta}}+\ell_{\text{div-
fgs}}\Big{]}\\\
\ell_{T}&=\mathbb{E}_{\underbrace{q_{\phi,T},(x,v)}_{{\text{depends on
}\boldsymbol{\mathbf{Q}},\boldsymbol{\mathbf{D}}}}}\Big{[}\log\pi_{\theta}(\boldsymbol{\mathbf{y}}_{T})\Big{]}\\\
\ell_{q}&=-\log q_{\gamma}(v|x)\end{split}$ (53)

For the dsm form,

$\displaystyle\mathcal{L}^{\text{mdsm}}=\mathbb{E}_{v\sim
q_{\gamma}}\Bigg{[}\mathbb{E}_{s\sim\text{Unif}(0,T)}\Big{[}\ell_{s}^{(dsm)}\Big{]}+\ell_{T}+\ell_{q}\Bigg{]}$
(54)

where

$\displaystyle\ell_{\text{div-f}}$
$\displaystyle=\nabla_{\boldsymbol{\mathbf{y}}_{s}}\cdot\underbrace{-[\boldsymbol{\mathbf{Q}}_{\phi}(s)+\boldsymbol{\mathbf{D}}_{\phi}(s)]\boldsymbol{\mathbf{y}}_{s}}_{f_{\phi}}$
$\displaystyle\ell_{\text{fwd-score}}$
$\displaystyle=\frac{1}{2}\Big{|}\Big{|}\underbrace{\nabla_{\boldsymbol{\mathbf{y}}_{s}}\log
q_{\phi}(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})}_{\text{depends
on
}\boldsymbol{\mathbf{Q}},\boldsymbol{\mathbf{D}}}\|^{2}_{\underbrace{2\boldsymbol{\mathbf{D}}_{\phi}(s)}_{g^{2}_{\phi}}}$
$\displaystyle\ell_{\text{neg-scorediff}}$
$\displaystyle=-\frac{1}{2}\|s_{\theta}(\boldsymbol{\mathbf{y}}_{s},s)-\underbrace{\nabla_{\boldsymbol{\mathbf{y}}_{s}}\log
q_{\phi}(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})}_{\text{depends
on
}\boldsymbol{\mathbf{Q}},\boldsymbol{\mathbf{D}}}\|^{2}_{\underbrace{2\boldsymbol{\mathbf{D}}_{\phi}(s)}_{g^{2}_{\phi}}}$
$\displaystyle\ell_{s}^{(dsm)}$
$\displaystyle=\mathbb{E}_{\underbrace{q_{\phi,s,(x,v)}}_{\text{depends on
}\boldsymbol{\mathbf{Q}},\boldsymbol{\mathbf{D}}}}\Bigg{[}\ell_{\text{neg-
scorediff}}+\ell_{\text{fwd-score}}+\ell_{\text{div-f}}\Bigg{]}$

##  Appendix G Algorithms

###  G.1 Generic Transition Kernel

Input: data $x$. time $s$. $\boldsymbol{\mathbf{A}},g$.

compute: $\boldsymbol{\mathbf{A}}(s)$ and $g(s)$

compute:
$\boldsymbol{\mathbf{M}}_{s}=\int_{0}^{s}\boldsymbol{\mathbf{A}}(t)dt$
(integrated drift)

compute: $\boldsymbol{\mathbf{N}}_{s}=\int_{0}^{s}g^{2}(t)dt$ (integrated
diffusions squared)

compute: $\gamma_{s|0}=\exp\Big{(}\boldsymbol{\mathbf{M}}_{s}\Big{)}$ (mean
coefficient)

set: $\boldsymbol{\mathbf{y}}_{0}=[x,0_{1},\ldots,0_{K-1}]$ ,
$\boldsymbol{\Sigma}_{0,zz}=\boldsymbol{0}$, and
$\boldsymbol{\Sigma}_{0,zv},\boldsymbol{\Sigma}_{0,vv}$ to chosen initial
distribution

compute:
$\boldsymbol{\mathbf{m}}_{s|0}=\gamma_{s|0}\boldsymbol{\mathbf{y}}_{0}$ (mean)

compute: $\displaystyle\begin{pmatrix}\boldsymbol{\mathbf{C}}_{s}\\\
\boldsymbol{\mathbf{H}}_{s}\end{pmatrix}=\exp\Bigg{[}\begin{pmatrix}\boldsymbol{\mathbf{M}}_{s}&\boldsymbol{\mathbf{N}}_{s}\\\
\boldsymbol{\mathbf{0}}&-\boldsymbol{\mathbf{M}}_{s}^{\top}\end{pmatrix}\Bigg{]}\begin{pmatrix}\boldsymbol{\mathbf{\Sigma}}_{0}\\\
\boldsymbol{\mathbf{I}}\end{pmatrix}\quad\text{(ingredients for cov.)}$ (55)

compute:
$\boldsymbol{\mathbf{\Sigma}}_{s|0}=\boldsymbol{\mathbf{C}}_{s}({\boldsymbol{\mathbf{H}}_{s}})^{-1}\quad(\textbf{cov.})$

Output: $\mathcal{N}(\boldsymbol{\mathbf{m}}_{s|0},\boldsymbol{\Sigma}_{s|0})$

Algorithm 2 Get transition distribution $\boldsymbol{\mathbf{y}}_{s}|x$

###  G.2 Transitions with $Q,D$

Current param matrices
$\tilde{\boldsymbol{\mathbf{Q}}}_{\phi},\tilde{\boldsymbol{\mathbf{D}}}_{\phi}$
and along with fixed time-in scalar-out functions $b_{q}(s),b_{d}(s)$ and
their known integrals $B_{q}(s),B_{d}(s)$. $q_{\gamma}(v_{0}|z_{0}=x)$ taken
to be parameterless so that $v_{0}\sim\mathcal{N}(0,I)$. Model params are
$s_{\theta}$ fixed $\pi_{\theta}$.

Input: time $s$ and current params $\phi$

compute: $[b_{q}]_{s}=\int_{0}^{s}b_{q}(\nu)d\nu$ using known integral
$B_{q}(s)-B_{q}(0)$

compute: $[b_{d}]_{s}=\int_{0}^{s}b_{d}(\nu)d\nu$ using known integral
$B_{d}(s)-B_{d}(0)$.

compute:
$[\boldsymbol{\mathbf{Q}}_{\phi}]_{s}=[b_{q}]_{s}\cdot\Big{[}\tilde{\boldsymbol{\mathbf{Q}}}_{\phi}-\tilde{\boldsymbol{\mathbf{Q}}}_{\phi}^{\top}\Big{]}$
for current params $\tilde{\boldsymbol{\mathbf{Q}}}_{\phi}$.

compute:
$[\boldsymbol{\mathbf{D}}_{\phi}]_{s}=[b_{d}]_{s}\cdot\Big{[}\tilde{\boldsymbol{\mathbf{D}}}_{\phi}\tilde{\boldsymbol{\mathbf{D}}}_{\phi}^{\top}\Big{]}$
for current params $\tilde{\boldsymbol{\mathbf{D}}}_{\phi}$.

compute:
$\boldsymbol{\mathbf{M}}_{s}=-([\boldsymbol{\mathbf{Q}}_{\phi}]_{s}+[\boldsymbol{\mathbf{D}}_{\phi}]_{s})$
($\boldsymbol{\mathbf{M}}$ just a variable name)

compute:
$\boldsymbol{\mathbf{N}}_{s}=[2\boldsymbol{\mathbf{D}}_{\phi}]_{s}=2\cdot[\boldsymbol{\mathbf{D}}_{\phi}]_{s}$
($\boldsymbol{\mathbf{N}}$ just a variable name)

compute:
$\boldsymbol{\mathbf{Q}}_{s}=b_{q}(s)\cdot\Big{[}\tilde{\boldsymbol{\mathbf{Q}}}_{\phi}-\tilde{\boldsymbol{\mathbf{Q}}}_{\phi}^{\top}\Big{]}$
(not integrated)

compute:
$\boldsymbol{\mathbf{D}}_{s}=b_{d}(s)\cdot\Big{[}\tilde{\boldsymbol{\mathbf{D}}}_{\phi}\tilde{\boldsymbol{\mathbf{D}}}_{\phi}^{\top}\Big{]}$
(not integrated)

compute: $A_{s}=-[\boldsymbol{\mathbf{Q}}_{s}+\boldsymbol{\mathbf{D}}_{s}]$
(drift coef.)

compute: $g_{s}^{2}=2\boldsymbol{\mathbf{D}}_{s}$ (diffusion coef. squared)

Output:
$\boldsymbol{\mathbf{A}}_{s},g^{2}_{s},\boldsymbol{\mathbf{M}}_{s},\boldsymbol{\mathbf{N}}_{s}$

Algorithm 3 Get $\boldsymbol{\mathbf{Q}},\boldsymbol{\mathbf{D}}$ and their
integrated terms $\boldsymbol{\mathbf{M}},\boldsymbol{\mathbf{N}}$

Input: Sample $\boldsymbol{\mathbf{y}}_{0}=(x,v)$ and time $s$. Current params
$\phi$

set:
$\boldsymbol{\mathbf{A}}_{s},g^{2}_{s},\boldsymbol{\mathbf{M}}_{s},\boldsymbol{\mathbf{N}}_{s}\leftarrow$
algorithm 3

compute:
$\boldsymbol{\mathbf{m}}_{s|0}=\exp\Big{(}\boldsymbol{\mathbf{M}}_{s}\Big{)}\boldsymbol{\mathbf{y}}_{0}$
(transition mean)

compute: ingredients for transition cov. matrix:
$\displaystyle\begin{pmatrix}\boldsymbol{\mathbf{C}}_{s}\\\
\boldsymbol{\mathbf{H}}_{s}\end{pmatrix}=\exp\Bigg{[}\begin{pmatrix}\boldsymbol{\mathbf{M}}_{s}&\boldsymbol{\mathbf{N}}_{s}\\\
\boldsymbol{\mathbf{0}}&-\boldsymbol{\mathbf{M}}_{s}^{\top}\end{pmatrix}\Bigg{]}\begin{pmatrix}\boldsymbol{\mathbf{0}}\\\
\boldsymbol{\mathbf{I}}\end{pmatrix}$ (56)

compute:
$\boldsymbol{\mathbf{\Sigma}}_{s|0}=\boldsymbol{\mathbf{C}}_{s}({\boldsymbol{\mathbf{H}}_{s}})^{-1}$
(transition cov).

instantiate:
$q_{\phi,s,(x,v)}=q_{\phi}(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})=\mathcal{N}(\boldsymbol{\mathbf{m}}_{s|0},\boldsymbol{\Sigma}_{s|0})$.

Output: $q_{\phi,s,(x,v)},A_{s},g_{s}^{2}$

Algorithm 4 Get transition distributions

###  G.3 elbo algorithms

input: Data point $x$ and current params $\theta,\phi,\gamma$

draw: an aux. sample $v\sim q_{\gamma}(v|x)$

draw: a sample $s\sim\text{Unif}(0,T)$

set: $\boldsymbol{\mathbf{y}}_{0}=(x,v)$

set: $q_{\phi,s,\boldsymbol{\mathbf{y}}_{0}},A_{s},g_{s}^{2}\leftarrow$
algorithm 4 called on $\boldsymbol{\mathbf{y}}_{0},s,\phi$

draw: $\boldsymbol{\mathbf{y}}_{s}\sim q_{\phi,s,\boldsymbol{\mathbf{y}}_{0}}$

compute: $\ell_{s}$ with $\text{dsm}(s)$ (algorithm 6) or $\text{ism}(s)$
(algorithm 7) on
$\boldsymbol{\mathbf{y}}_{s},\theta,A_{s},g_{s}^{2},q_{\phi,s,\boldsymbol{\mathbf{y}}_{0}}$

set: $q_{\phi,T,\boldsymbol{\mathbf{y}}_{0}},\\_\\_,\\_\\_\leftarrow$
algorithm 4 called on $\boldsymbol{\mathbf{y}}_{0},T,\phi$

draw: $\boldsymbol{\mathbf{y}}_{T}\sim q_{\phi,T,\boldsymbol{\mathbf{y}}_{0}}$

output: $\ell_{s}+\log\pi_{\theta}(\boldsymbol{\mathbf{y}}_{T})-\log
q_{\gamma}(v)$

Algorithm 5 Compute elbo with ism or dsm

input: $\boldsymbol{\mathbf{y}}_{s}$, $\theta$, $A_{s}$, $g_{s}^{2}$,
$q_{\phi,s,\boldsymbol{\mathbf{y}}_{0}}$.

compute: $\text{fwd-score}=\nabla_{\boldsymbol{\mathbf{y}}_{s}}\log
q_{\phi}(\boldsymbol{\mathbf{y}}_{s}|\boldsymbol{\mathbf{y}}_{0})$

compute: $\text{model-score}=s_{\theta}(\boldsymbol{\mathbf{y}}_{s},s)$

compute: $\text{fwd-score-term}=\frac{1}{2}(\text{fwd-
score})^{\top}g_{s}^{2}(\text{fwd-score})$

compute: $\text{score-diff}=\text{model-score}-\text{fwd-score}$

compute: $\text{diff-term}=-\frac{1}{2}\text{score-
diff}^{\top}g_{s}^{2}\text{score-diff}$

compute: $\text{div-f}=\nabla_{\boldsymbol{\mathbf{y}}_{s}}\cdot
A_{s}\boldsymbol{\mathbf{y}}_{s}$

output: $\text{dsm}(s)=\text{fwd-score-term}+\text{diff-term}+\text{div-f}$

Algorithm 6 Compute $\text{dsm}(s)$

input: $\boldsymbol{\mathbf{y}}_{s}$, $\theta$, $A_{s}$, $g_{s}^{2}$,
$q_{\phi,s,\boldsymbol{\mathbf{y}}_{0}}$.

compute: $\text{model-score}=s_{\theta}(\boldsymbol{\mathbf{y}}_{s},s)$

compute: $\text{score-term}=-\frac{1}{2}\text{model-
score}^{\top}g_{s}^{2}\text{model-score}$

compute: $\text{div-gs}=\nabla_{\boldsymbol{\mathbf{y}}_{s}}\cdot
g_{s}^{2}s_{\theta}(\boldsymbol{\mathbf{y}}_{s},s)$

compute: $\text{div-f}=\nabla_{\boldsymbol{\mathbf{y}}_{s}}\cdot
A_{s}\boldsymbol{\mathbf{y}}_{s}$

compute: $\text{div-term}=-\text{div-gs}+\text{div-f}$

output: $\text{ism}(s)=\text{score-term}+\text{div-term}$

Algorithm 7 Compute $\text{ism}(s)$

##  Appendix H Valid elbo with truncation

The integrand in the elbo and its gradients is not bounded at time $0$.
Therefore, following Sohl-Dickstein et al. (2015) and Song et al. (2021) the
integrand in eq. 7 is integrated from $[\epsilon,T]$, rather than $[0,T]$.
However, that integral is not a valid lower bound on $\log p_{\theta}(x)$.
Instead, it can be viewed as a proper lower bound on the prior for a latent
variable $\boldsymbol{\mathbf{y}}_{\epsilon}$. Therefore, to provide a bound
for the data, one can introduce a likelihood and substitute the prior lower
bound into a standard variational bound that integrates out the latent.

To provide a valid lower bound for multivariate diffusions, we extend theorem
6 in Song et al. (2021) from univariate to multivariate diffusions.

######  Theorem 3.

For transition kernel $q_{\phi}({\mathbf{y}}_{s}\mid{\mathbf{y}}_{0})$, we can
compute upper bound the model likelihood at time $0$ as follows, for any
$\epsilon>0$

$\displaystyle\log
p_{\theta}(x)\geq\mathbb{E}_{q_{\phi}({\mathbf{y}}^{v}_{0}\mid
x)}\mathbb{E}_{q_{\phi}({\mathbf{y}}_{\epsilon}|{\mathbf{y}}_{0})}\left[\log\frac{p_{\theta}({\mathbf{y}}_{0}\mid{\mathbf{y}}_{\epsilon})}{q_{\phi}({\mathbf{y}}_{\epsilon}\mid{\mathbf{y}}_{0})}+\mathcal{L}_{\text{mdm}}({\mathbf{y}}_{\epsilon},\epsilon)-\log
q_{\phi}({\mathbf{y}}_{0}^{v}\mid x)\right],$ (57)

where $\mathcal{L}_{\text{mdm}}({\mathbf{y}}_{\epsilon},\epsilon)$ is defined
as

$\displaystyle\mathcal{L}_{\text{mdm}}({\mathbf{y}}_{\epsilon},\epsilon)=\mathbb{E}_{q_{\phi}({\mathbf{y}}_{>\epsilon}\mid{\mathbf{y}}_{\epsilon})}\left[\log\pi_{\theta}({\mathbf{y}}_{T})-\int_{\epsilon}^{T}\frac{1}{2}\left\lVert
s_{q}\right\rVert_{g_{\phi}}^{2}-\frac{1}{2}\left\lVert
s_{\theta}-s_{q}\right\rVert^{2}_{g_{\phi}}+\nabla\cdot f_{\phi}\right].$

###### Proof.

For transition kernel $q_{\phi}({\mathbf{y}}_{s}\mid{\mathbf{y}}_{0})$, we can
compute upper bound the model likelihood at time $0$ following an application
of the variational bound

$\displaystyle\log p_{\theta}(x)$
$\displaystyle=\log\int_{v_{0}}p_{\theta}({\mathbf{y}}_{0}=[x,v_{0}])dv_{0}$
$\displaystyle=\log\int_{v_{0},{\mathbf{y}}_{\epsilon}}p_{\theta}({\mathbf{y}}_{0},{\mathbf{y}}_{\epsilon})dv_{0}d{\mathbf{y}}_{\epsilon}$
$\displaystyle=\log\int_{v_{0},{\mathbf{y}}_{\epsilon}}q_{\phi}({\mathbf{y}}_{\epsilon}\mid{\mathbf{y}}_{0})\frac{q(v_{0}\mid
x)}{q(v_{0}\mid
x)}\frac{p_{\theta}({\mathbf{y}}_{0},{\mathbf{y}}_{\epsilon})}{q_{\phi}({\mathbf{y}}_{\epsilon}\mid{\mathbf{y}}_{0})}dv_{0}d{\mathbf{y}}_{\epsilon}$
$\displaystyle=\log\int_{v_{0},{\mathbf{y}}_{\epsilon}}q_{\phi}({\mathbf{y}}_{\epsilon}\mid{\mathbf{y}}_{0})\frac{q(v_{0}\mid
x)}{q(v_{0}\mid
x)}\frac{p_{\theta}({\mathbf{y}}_{0}\mid{\mathbf{y}}_{\epsilon})p_{\theta}({\mathbf{y}}_{\epsilon})}{q_{\phi}({\mathbf{y}}_{\epsilon}\mid{\mathbf{y}}_{0})}dv_{0}d{\mathbf{y}}_{\epsilon}$
$\displaystyle\geq\mathbb{E}_{q(v_{0}\mid
x)q_{\phi}({\mathbf{y}}_{\epsilon}\mid{\mathbf{y}}_{0})}\left[\log\frac{p_{\theta}({\mathbf{y}}_{0}\mid{\mathbf{y}}_{\epsilon})}{q_{\phi}({\mathbf{y}}_{\epsilon}\mid{\mathbf{y}}_{0})}-\log
q_{\phi}({\mathbf{y}}_{0}^{v}\mid x)+\log
p_{\theta}({\mathbf{y}}_{\epsilon})\right]$

A lower bound for $\log p_{\theta}({\mathbf{y}}_{\epsilon})$ can be derived in
a similar manner to eq. 7, such that

$\displaystyle\log
p_{\theta}({\mathbf{y}}_{\epsilon})\geq\mathcal{L}_{\text{mdm}}({\mathbf{y}}_{\epsilon},\epsilon)=\mathbb{E}_{q_{\phi}({\mathbf{y}}_{>\epsilon}\mid{\mathbf{y}}_{\epsilon})}\left[\log\pi_{\theta}({\mathbf{y}}_{T})-\int_{\epsilon}^{T}\frac{1}{2}\left\lVert
s_{q}\right\rVert_{g_{\phi}}^{2}-\frac{1}{2}\left\lVert
s_{\theta}-s_{q}\right\rVert^{2}_{g_{\phi}}+\nabla\cdot f_{\phi}\right].$

The choice of $p_{\theta}({\mathbf{y}}_{0}\mid{\mathbf{y}}_{\epsilon})$ is
arbitrary, however following Sohl-Dickstein et al. (2015); Song et al. (2021)
we let $p_{\theta}({\mathbf{y}}_{0}\mid{\mathbf{y}}_{\epsilon})$ be Gaussian
with mean $\mu_{p_{\theta},\epsilon}$ and covariance
$\Sigma_{p_{\theta},\epsilon}$. Suppose
$q_{\phi}({\mathbf{y}}_{\epsilon}\mid{\mathbf{y}}_{0})={\mathcal{N}}({\mathbf{y}}_{\epsilon}\mid{\mathbf{A}}{\mathbf{y}}_{0},\Sigma)$,
then we select the following mean $\mu_{p_{\theta},\epsilon}$ and covariance
$\Sigma_{p_{\theta},\epsilon}$ for
$p_{\theta}({\mathbf{y}}_{0}\mid{\mathbf{y}}_{\epsilon})$

$\displaystyle\mu_{p_{\theta},\epsilon}$
$\displaystyle={\mathbf{A}}^{-1}\Sigma
s_{\theta}({\mathbf{y}}_{\epsilon},\epsilon)+{\mathbf{A}}^{-1}{\mathbf{y}}_{\epsilon}$
$\displaystyle\Sigma_{p_{\theta},\epsilon}$
$\displaystyle={\mathbf{A}}^{-1}\Sigma{\mathbf{A}}^{-\top}$

where $\mu_{p_{\theta},\epsilon},\Sigma_{p_{\theta},\epsilon}$ are derived
using Tweedie’s formula (Efron, 2011) by setting
$\mu_{\epsilon}=\mathbb{E}[{\mathbf{y}}_{0}\mid{\mathbf{y}}_{\epsilon}]$ and
$\Sigma_{\epsilon}=\text{Var}\left({\mathbf{y}}_{0}\mid{\mathbf{y}}_{\epsilon}\right)$.
∎

We next derive this choice as an approximation of the optimal Gaussian
likelihood.

###  H.1 Likelihood derivation

Suppose ${\mathbf{y}}_{0}\sim q_{0}({\mathbf{y}}_{0})$ and
${\mathbf{y}}_{\epsilon}\sim{\mathcal{N}}({\mathbf{y}}_{\epsilon}\mid
A{\mathbf{y}}_{0},\Sigma)$. Here, $A,\Sigma$ are the mean coefficient and
covariance derived from the transition kernel at time $\epsilon$. We use
Tweedie’s formula to get the mean and covariance of
$\boldsymbol{\mathbf{y}}_{0}$ given $\boldsymbol{\mathbf{y}}_{\epsilon}$ under
$q$. This mean and covariance feature the true score
$\nabla_{\boldsymbol{\mathbf{y}}_{\epsilon}}\log
q(\boldsymbol{\mathbf{y}}_{\epsilon})$. We replace the score with the score
model $s_{\theta}$ and then set
$p_{\theta}(\boldsymbol{\mathbf{y}}_{0}|\boldsymbol{\mathbf{y}}_{\epsilon})$
to have the resulting approximate mean and covariance. We make this choice
because the optimal
$p_{\theta}(\boldsymbol{\mathbf{y}}_{0}|\boldsymbol{\mathbf{y}}_{\epsilon})$
equals the true
$q(\boldsymbol{\mathbf{y}}_{0}|\boldsymbol{\mathbf{y}}_{\epsilon})$ as
discussed throughout the work.

Here ${\mathbf{y}}_{0}=[{\mathbf{x}}_{0},{\mathbf{v}}_{0}]$ where
${\mathbf{x}}_{0}\sim q_{\text{data}}$.

Let $\eta$ be the natural parameter for the multivariate Gaussian likelihood
${\mathcal{N}}({\mathbf{y}}_{\epsilon}\mid A{\mathbf{y}}_{0},\Sigma)$. Then,
Tweedie’s formula (Efron, 2011) states that:

$\displaystyle\mathbb{E}[\eta\mid{\mathbf{u}}_{\epsilon}]=\nabla_{{\mathbf{y}}_{\epsilon}}l({\mathbf{y}}_{\epsilon})-\nabla_{{\mathbf{y}}_{\epsilon}}l_{0}({\mathbf{y}}_{\epsilon})$

  * •

$l({\mathbf{y}}_{\epsilon})=\log q({\mathbf{y}}_{\epsilon})$

  * •

$s_{\theta}(\boldsymbol{\mathbf{y}}_{\epsilon},\epsilon)$ is taken to be the
true score $\nabla_{\boldsymbol{\mathbf{y}}_{\epsilon}}\log
q(\boldsymbol{\mathbf{y}}_{\epsilon})$ so that
$\nabla_{{\mathbf{y}}_{\epsilon}}l({\mathbf{y}}_{\epsilon})=s_{\theta}({\mathbf{y}}_{\epsilon},\epsilon)$

  * •

$l_{0}$ is the log of the base distribution defined in the exponential family
parameterization.

The base distribution is a multivariate Gaussian with mean $0$ and covariance
$\Sigma$, therefore
$\nabla_{{\mathbf{y}}_{\epsilon}}l_{0}({\mathbf{y}}_{\epsilon})=-\Sigma^{-1}{\mathbf{y}}_{\epsilon}$,

$\displaystyle\mathbb{E}[\eta\mid{\mathbf{y}}_{\epsilon}]=s_{\theta}({\mathbf{y}}_{\epsilon},\epsilon)+\Sigma^{-1}{\mathbf{y}}_{\epsilon}.$

However, Tweedie’s formula is not directly applicable since our
$\boldsymbol{\mathbf{y}}_{\epsilon}$ is not directly normal with mean
$\boldsymbol{\mathbf{y}}_{0}$. Instead, to derive the conditional mean of
${\mathbf{y}}_{0}$ given ${\mathbf{y}}_{\epsilon}$, we use the relation
$\eta=\Sigma^{-1}{\mathbf{A}}{\mathbf{y}}_{0}$ and the linearity of
conditional expectation to get

$\displaystyle\mathbb{E}[{\mathbf{y}}_{0}\mid{\mathbf{y}}_{\epsilon}]$
$\displaystyle=\mathbb{E}[A^{-1}\Sigma\eta|{\mathbf{y}}_{\epsilon}]$
$\displaystyle=A^{-1}\Sigma\mathbb{E}[\eta\mid{\mathbf{y}}_{\epsilon}]$
$\displaystyle=A^{-1}\Sigma\left(s_{\theta}({\mathbf{y}}_{\epsilon},\epsilon)+\Sigma^{-1}{\mathbf{y}}_{\epsilon}\right)$
$\displaystyle=A^{-1}\Sigma
s_{\theta}({\mathbf{y}}_{\epsilon},\epsilon)+A^{-1}{\mathbf{y}}_{\epsilon}.$

For the variance, we use the following relation
${\mathbf{y}}_{\epsilon}=A{\mathbf{y}}_{0}+\sqrt{\Sigma}\epsilon$, which
implies that

$\displaystyle{\mathbf{y}}_{0}$
$\displaystyle=A^{-1}{\mathbf{y}}_{\epsilon}-A^{-1}\sqrt{\Sigma}\epsilon$
$\displaystyle\text{Var}\left({\mathbf{y}}_{0}\mid{\mathbf{y}}_{\epsilon}\right)$
$\displaystyle=A^{-1}\Sigma A^{-T}.$

Therefore, for the model posterior distribution
$p_{\theta}({\mathbf{y}}_{0}\mid{\mathbf{y}}_{\epsilon})$ we choose a Normal
with mean and covariance

$\displaystyle\mu_{p_{\theta},\epsilon}$ $\displaystyle=A^{-1}\Sigma
s_{\theta}({\mathbf{y}}_{\epsilon},\epsilon)+A^{-1}{\mathbf{y}}_{\epsilon}$
$\displaystyle\Sigma_{p_{\theta},\epsilon}$ $\displaystyle=A^{-1}\Sigma
A^{-T}$