Original: W HERE TO D IFFUSE , H OW TO D IFFUSE , AND H OW TO G ET B ACK : A UTOMATED L EARNING FOR M ULTIVARIATE D IFFUSIONS
Revised: W HERE TO D IFFUSE , H OW TO D IFFUSE , AND H OW TO G ET B ACK : A UTOMATED L EARNING FOR M ULTIVARIATE D IFFUSIONS

Original: A BSTRACT
Revised: A BSTRACT

Original: Diffusion-based generative models ( DBGM s) perturb data to a target noise distribution and reverse thisinference diffusion process to generate samples.
Revised: Diffusion-based generative models ( DBGM s) perturb data to a target noise distribution and reverse this process to generate samples.

Original: The choice of  inference diffusion  affects both likelihoods and sample quality.
Revised: The choice of noising process, or inference diffusion process , affects both likelihoods and sample quality.

Original: For example, extending the inference process with auxiliary variables leads to improved sample quality.
Revised: For example, extending the inference process with auxiliary variables leads to improved sample quality.

Original: While there are many such multivariate diffusions to explore, each new one requires significant model-specific analysis, hindering rapid prototyping and evaluation.
Revised: While there are many such multivariate diffusions to explore, each new one requires significant model-specific analysis, hindering rapid prototyping and evaluation.

Original: In this work, we study Multivariate Diffusion Models ( MDM s).
Revised: In this work, we study Multivariate Diffusion Models ( MDM s).

Original: For any number of auxiliary variables, we provide a recipe for maximizing a lowerbound on the MDM s likelihood without requiring any model-specific analysis.
Revised: For any number of auxiliary variables, we provide a recipe for maximizing a lower-bound on the MDM s likelihood without requiring any model-specific analysis.

Original: We then demonstrate how to parameterize the diffusion for a specified target noise distribution; these two points together enable optimizing the inference diffusion process.
Revised: We then demonstrate how to parameterize the diffusion for a specified target noise distribution; these two points together enable optimizing the inference diffusion process.

Original: Optimizing the diffusion expands easy experimentation from just a few well-known processes to an automatic search over all linear diffusions.
Revised: Optimizing the diffusion expands easy experimentation from just a few well-known processes to an automatic search over all linear diffusions.

Original: To demonstrate these ideas, we introduce two new specific diffusions as well as learn a diffusion process on the MNIST , CIFAR 10, and IMAGENET 32 datasets.
Revised: To demonstrate these ideas, we introduce two new specific diffusions as well as learn a diffusion process on the MNIST , CIFAR 10, and IMAGENETdatasets.

Original: We show learned MDM s match or surpass bits-per-dims ( BPD s) relative to fixed choices of diffusions for a given dataset and model architecture.
Revised: We show learned MDM s match or surpass bits-per-dims ( BPD s) relative to fixed choices of diffusions for a given dataset and model architecture.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Diffusion-based generative models ( DBGM s) perturb data to a target noise distribution and reverse this process to generate samples.
Revised: Diffusion-based generative models ( DBGM s) perturb data to a target noise distribution and reverse this process to generate samples.

Original: They have achieved impressive performance in image generation, editing, translation (Dhariwal & Nichol, 2021; Nichol & Dhariwal, 2021; Sasaki et al., 2021;
Revised: They have achieved impressive performance in image generation, editing, translation (Dhariwal & Nichol, 2021; Nichol & Dhariwal, 2021; Sasaki et al., 2021;

Original: 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).
Revised: 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).

Original: 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).
Revised: 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).

Original: The choice of this inference process affects both likelihoods and sample quality.
Revised: The choice of this inference process affects both likelihoods and sample quality.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: (2021) instead introduce the critically-damped langevin diffusion ( CLD ), showing that significant improvements in sample generation can be gained by carefully designing new processes.
Revised: (2021) instead introduce the critically-damped langevin diffusion ( CLD ), showing that significant improvements in sample generation can be gained by carefully designing new processes.

Original: CLD pairs each data dimension with an auxiliary “velocity” variable and diffuses them jointly using second-order Langevin dynamics.
Revised: CLD pairs each data dimension with an auxiliary “velocity” variable and diffuses them jointly using second-order Langevin dynamics.

Original: 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 DBGM s?
Revised: 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 DBGM s?

Original: For instance, are there other auxiliary variable diffusions that would lead to improvements like CLD ?
Revised: For instance, are there other auxiliary variable diffusions that would lead to improvements like CLD ?

Original: This avenue seems promising as auxiliary variables have improved other generative models and inferences, such as normalizing flows
Revised: This avenue seems promising as auxiliary variables have improved other generative models and inferences, such as normalizing flows

Original: (Huang et al., 2020), neural ordinary differential equations ( ODE s) (Dupont et al., 2019), hierarchical variational models (Ranganath et al., 2016), ladder variational autoencoder (Sønderby et al., 2016), among others.
Revised: (Huang et al., 2020), neural ordinary differential equations ( ODE s) (Dupont et al., 2019), hierarchical variational models (Ranganath et al., 2016), ladder variational autoencoder (Sønderby et al., 2016), among others.

Original: Despite its success, CLD also provides evidence that each new process requires significant modelspecific analysis.
Revised: Despite its success, CLD also provides evidence that each new process requires significant modelspecific analysis.

Original: 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).
Revised: 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).

Original: Auxiliary variables seemingly complicate this process further; computing conditionals of the inference process necessitates solving matrix Lyupanov equations (section 3.3).
Revised: Auxiliary variables seemingly complicate this process further; computing conditionals of the inference process necessitates solving matrix Lyupanov equations (section 3.3).

Original: Deriving the inference stationary distribution—which helps the model and inference match—can be intractable.
Revised: Deriving the inference stationary distribution—which helps the model and inference match—can be intractable.

Original: These challenges limit rapid prototyping and evaluation of new inference processes.
Revised: These challenges limit rapid prototyping and evaluation of new inference processes.

Original: Concretely, training a diffusion model requires:
Revised: Concretely, training a diffusion model requires:

Original: ( R1):
Revised: ( R1):

Original: Selecting an inference and model process pair such
Revised: Selecting an inference and model process pair such

Original: that the inference
Revised: that the inference

Original: process converges to the model prior ( R2):
Revised: process converges to the model prior ( R2):

Original: Deriving the ELBO for this pair ( R3): Estimating the ELBO and its gradients by deriving and computing the inference process’
Revised: Deriving the ELBO for this pair ( R3): Estimating the ELBO and its gradients by deriving and computing the inference process’

Original: transition
Revised: transition

Original: kernel
Revised: kernel

Original: In this work, we introduce Multivariate Diffusion Models ( MDM s) and a method for training and evaluating them.
Revised: In this work, we introduce Multivariate Diffusion Models ( MDM s) and a method for training and evaluating them.

Original: MDM s are diffusion-based generative models trained with auxiliary variables.
Revised: MDM s are diffusion-based generative models trained with auxiliary variables.

Original: We provide a recipe for training MDM s beyond specific instantiations–like VPSDE and CLD —to all linear inference processes that have a stationary distribution, with any number of auxiliary variables.
Revised: We provide a recipe for training MDM s beyond specific instantiations–like VPSDE and CLD —to all linear inference processes that have a stationary distribution, with any number of auxiliary variables.

Original: First, we bring results from gradient-based MCMC (Ma et al., 2015) to diffusion modeling to construct MDM s that converge to a chosen model prior ( R1 ); this tightens the ELBO .
Revised: First, we bring results from gradient-based MCMC (Ma et al., 2015) to diffusion modeling to construct MDM s that converge to a chosen model prior ( R1 ); this tightens the ELBO .

Original: Secondly, for any number of auxiliary variables, we derive the MDM ELBO ( R2 ).
Revised: Secondly, for any number of auxiliary variables, we derive the MDM ELBO ( R2 ).

Original: Finally, we show that the transition kernel of linear MDM s, necessary for the ELBO , can be computed automatically and generically, for higher-dimensional auxiliary systems ( R3 ).
Revised: Finally, we show that the transition kernel of linear MDM s, necessary for the ELBO , can be computed automatically and generically, for higher-dimensional auxiliary systems ( R3 ).

Original: With these tools, we explore a variety of new inference processes for diffusion-based generative models.
Revised: With these tools, we explore a variety of new inference processes for diffusion-based generative models.

Original: We then note that the automatic transitions and fixed stationary distributions facilitate directly learning the inference to maximize the MDM ELBO .
Revised: We then note that the automatic transitions and fixed stationary distributions facilitate directly learning the inference to maximize the MDM ELBO .

Original: Learning turns diffusion model training into a search not only over score models but also inference processes, at no extra derivational cost.
Revised: Learning turns diffusion model training into a search not only over score models but also inference processes, at no extra derivational cost.

Original: Methodological Contributions.
Revised: Methodological Contributions.

Original: In summary, our methodological contributions are:
Revised: In summary, our methodological contributions are:

Original: Deriving ELBO s for training and evaluating multivariate diffusion models ( MDM s) with auxiliary variables.
Revised: Deriving ELBO s for training and evaluating multivariate diffusion models ( MDM s) with auxiliary variables.

Original: Showing that the diffusion transition covariance does not need to be manually derived for each new diffusion.
Revised: Showing that the diffusion transition covariance does not need to be manually derived for each new diffusion.

Original: We instead demonstrate that a matrix factorization technique, previously unused in diffusion models, can automatically compute the covariance analytically for any linear MDM .
Revised: We instead demonstrate that a matrix factorization technique, previously unused in diffusion models, can automatically compute the covariance analytically for any linear MDM .

Original: Using results from gradient-based Markov chain
Revised: Using results from gradient-based Markov chain

Original: Monte
Revised: Monte

Original: Carlo ( MCMC ) to construct MDM s with a complete parameterization of inference processes whose stationary distribution matches the model prior.
Revised: Carlo ( MCMC ) to construct MDM s with a complete parameterization of inference processes whose stationary distribution matches the model prior.

Original: Combining the above into an algorithm called Automatic Multivariate Diffusion Training ( AMDT ) that enables training without diffusion-specific derivations.
Revised: Combining the above into an algorithm called Automatic Multivariate Diffusion Training ( AMDT ) that enables training without diffusion-specific derivations.

Original: AMDT enables training score models for any linear diffusion, including optimizing the diffusion and score jointly.
Revised: AMDT enables training score models for any linear diffusion, including optimizing the diffusion and score jointly.

Original: To demonstrate these ideas, we develop MDM s with two specific diffusions as well as learned multivariate diffusions.
Revised: To demonstrate these ideas, we develop MDM s with two specific diffusions as well as learned multivariate diffusions.

Original: The specific diffusions are accelerated Langevin diffusion ( ALDA ) (introduced in Mou et al.
Revised: The specific diffusions are accelerated Langevin diffusion ( ALDA ) (introduced in Mou et al.

Original: as a higher-order scheme for gradient-based MCMC ) and an alteration, modified accelerated Langevin diffusion ( MALDA ).
Revised: as a higher-order scheme for gradient-based MCMC ) and an alteration, modified accelerated Langevin diffusion ( MALDA ).

Original: Previously, using these diffusions for generative modeling would require significant model-specific analysis.
Revised: Previously, using these diffusions for generative modeling would require significant model-specific analysis.

Original: Instead, AMDT for these diffusions is derivation-free.
Revised: Instead, AMDT for these diffusions is derivation-free.

Original: Empirical contributions.
Revised: Empirical contributions.

Original: We train MDM s on the MNIST , IMAGENET 32 and CIFAR -10 datasets.
Revised: We train MDM s on the MNIST , IMAGENET 32 and CIFAR -10 datasets.

Original: In the experiments, we show that:
Revised: In the experiments, we show that:

Original: Training new and existing fixed diffusions, such as ALDA and MALDA , is easy with the proposed algorithm AMDT .
Revised: Training new and existing fixed diffusions, such as ALDA and MALDA , is easy with the proposed algorithm AMDT .

Original: 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.
Revised: 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.

Original: There are new and existing MDM s, 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.
Revised: There are new and existing MDM s, 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.

Original: 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.
Revised: 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.

Original: We additionally show the utility of the MDM ELBO by showing on a dataset that CLD achieves better bits-per-dims ( BPD s) than previously reported with the probability flow ODE (Dockhorn et al., 2021).
Revised: We additionally show the utility of the MDM ELBO by showing on a dataset that CLD achieves better bits-per-dims ( BPD s) than previously reported with the probability flow ODE (Dockhorn et al., 2021).

Original:  S ETUP
Revised:  S ETUP

Original: 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;
Revised: 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;

Original: Kingma et al., 2021).
Revised: Kingma et al., 2021).

Original: Diffusions sample from a model prior z 0 ∼ π θ and then evolve a continuous-time stochastic process z t ∈ R d :
Revised: Diffusions sample from a model prior z 0 ∼ π θ and then evolve a continuous-time stochastic process z t ∈ R d :

Original: [Equation]
Revised: [Equation]

Original: where B t is a d -dimensionsal Brownian motion.
Revised: where B t is a d -dimensionsal Brownian motion.

Original: The model is trained so that z T approximates the data x ∼ q data .
Revised: The model is trained so that z T approximates the data x ∼ q data .

Original: 1 Maximum likelihood training of diffusion models is intractable (Huang et al., 2021;
Revised: 1 Maximum likelihood training of diffusion models is intractable (Huang et al., 2021;

Original: Song et al., 2021; Kingma et al., 2021).
Revised: Song et al., 2021; Kingma et al., 2021).

Original: Instead, they are trained using a variational lower bound on log p θ ( z T = x ) .
Revised: Instead, they are trained using a variational lower bound on log p θ ( z T = x ) .

Original: The bound requires an inference process q ϕ ( y s | x = x ) :
Revised: The bound requires an inference process q ϕ ( y s | x = x ) :

Original: [Equation]
Revised: [Equation]

Original: where (cid:98) B s is another Brownian motion independent of B t .
Revised: where (cid:98) B s is another Brownian motion independent of B t .

Original: The inference process is usually taken to be specified rather than learned, and chosen to be i.i.d.
Revised: The inference process is usually taken to be specified rather than learned, and chosen to be i.i.d.

Original: for each y tj conditional on each x j .
Revised: for each y tj conditional on each x j .

Original: This leads to the interpretation of the y tj as noisy versions of features x j (Ho et al., 2020).
Revised: This leads to the interpretation of the y tj as noisy versions of features x j (Ho et al., 2020).

Original: While the diffusion ELBO is challenging to derive in general, Huang et al.
Revised: While the diffusion ELBO is challenging to derive in general, Huang et al.

Original: ; Song et al.
Revised: ; Song et al.

Original: (2021) show that when the model process takes the form:
Revised: (2021) show that when the model process takes the form:

Original: [Equation]
Revised: [Equation]

Original: the ELBO is:
Revised: the ELBO is:

Original: [Equation]
Revised: [Equation]

Original: where f ϕ , g ϕ , s θ are evaluated at ( y s , s ) , ∥ x ∥ 2 A = x ⊤ Ax and g 2 = gg ⊤ .
Revised: where f ϕ , g ϕ , s θ are evaluated at ( y s , s ) , ∥ x ∥ 2 A = x ⊤ Ax and g 2 = gg ⊤ .

Original: Equation (4) features the Implicit Score Matching ( ISM ) loss (Song et al., 2020a), and can be re-written as an ELBO L dsm featuring Denoising Score Matching ( DSM ) (Vincent, 2011; Song et al., 2020b), see appendix E.1.
Revised: Equation (4) features the Implicit Score Matching ( ISM ) loss (Song et al., 2020a), and can be re-written as an ELBO L dsm featuring Denoising Score Matching ( DSM ) (Vincent, 2011; Song et al., 2020b), see appendix F.1.

Original:  A RECIPE FOR M ULTIVARIATE D IFFUSION M ODELS
Revised:  A RECIPE FOR M ULTIVARIATE D IFFUSION M ODELS

Original: As has been shown in prior work (Song et al., 2021; Dockhorn et al., 2021), the choice of diffusion matters.
Revised: As has been shown in prior work (Song et al., 2021; Dockhorn et al., 2021), the choice of diffusion matters.

Original: 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.
Revised: 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.

Original: At first glance, training such diffusions can seem challenging.
Revised: At first glance, training such diffusions can seem challenging.

Original: First, one needs an ELBO that includes auxiliary variables.
Revised: First, one needs an ELBO that includes auxiliary variables.

Original: This ELBO will require sampling from the transition kernel, and setting the model prior to the specified inference stationary distribution.
Revised: This ELBO will require sampling from the transition kernel, and setting the model prior to the specified inference stationary distribution.

Original: But doing such diffusion-specific analysis manually is challenging and hinders rapid prototyping.
Revised: But doing such diffusion-specific analysis manually is challenging and hinders rapid prototyping.

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

Revised: In appendix A we discuss how the presented methods can also be used to automate and improve simplified score matching and noise prediction objectives used to train diffusion models.

Original: 3. M ULTIVARIATE M ODEL AND I NFERENCE
Revised: 3. M ULTIVARIATE M ODEL AND I NFERENCE

Original: For the j th data coordinate at each time t , MDM s pair z tj ∈ R with a vector of auxiliary variables v tj ∈ R K − 1 into a joint vector u t and diffuse in the extended space:
Revised: For the j th data coordinate at each time t , MDM s pair z tj ∈ R with a vector of auxiliary variables v tj ∈ R K − 1 into a joint vector u t and diffuse in the extended space:

Original: [Equation]
Revised: [Equation]

Original: MDM s model the data x with z T , a coordinate in u T ∼ p θ .
Revised: MDM s model the data x with z T , a coordinate in u T ∼ p θ .

Original: For the j th feature x j , each u tj ∈ R K consists of a “data” dimension u ztj and auxiliary variable u vtj .
Revised: For the j th feature x j , each u tj ∈ R K consists of a “data” dimension u ztj and auxiliary variable u vtj .

Original: Therefore u ∈ R dK .
Revised: Therefore u ∈ R dK .

Original: We extend the drift coefficient h θ from a function in R d × R + → R d to the extended space R
Revised: We extend the drift coefficient h θ from a function in R d × R + → R d to the extended space R

Original: dK × R + → R dK .
Revised: dK × R + → R dK .

Original: We likewise extend the diffusion coefficient to a matrix β θ acting on Brownian motion B t ∈ R dK .
Revised: We likewise extend the diffusion coefficient to a matrix β θ acting on Brownian motion B t ∈ R dK .

Original: Because the MDM model is over the extended space, the inference distribution y must be too.
Revised: Because the MDM model is over the extended space, the inference distribution y must be too.

Original: We then set q ( y v 0 | y z 

Revised: We then set q ( y v 0 | y z

Original: = x ) to any chosen initial distribution, e.g. N ( 0 , I ) and discuss this choice in section 4.
Revised: = x ) to any chosen initial distribution, e.g. N ( 0 , I ) and discuss this choice in section 4.

Original: Then y s evolves according to the auxiliary variable inference process:
Revised: Then y s evolves according to the auxiliary variable inference process:

Original: [Equation]
Revised: [Equation]

Original: where the inference drift and diffusion coefficients f ϕ , g ϕ are now over the extended space y
Revised: where the inference drift and diffusion coefficients f ϕ , g ϕ are now over the extended space y

Original: =
Revised: =

Original: [ y z , y v ] .
Revised: [ y z , y v ] .

Original: The function f ϕ lets the z and v coordinates of y tj interact in the inference process.
Revised: The function f ϕ lets the z and v coordinates of y tj interact in the inference process.

Original: A SSUMPTIONS
Revised: A SSUMPTIONS

Original: This work demonstrates how to parameterize time-varying It ˆ o processes, used for diffusion modeling, to have a stationary distribution that matches the given model prior.
Revised: This work demonstrates how to parameterize time-varying Itˆo processes, used for diffusion modeling, to have a stationary distribution that matches the given model prior.

Original: To take advantage of the automatic transition kernels also presented, the inferences considered for modeling are linear time-varying processes and take the form:
Revised: To take advantage of the automatic transition kernels also presented, the inferences considered for modeling are linear time-varying processes and take the form:

Original: [Equation]
Revised: [Equation]

Original: where A ϕ ( s ) : R + →
Revised: where A ϕ ( s ) : R + →

Original: dK × dK and g ϕ ( s ) : R + → dK × dK are matrix-valued functions.
Revised: dK × dK and g ϕ ( s ) : R + → dK × dK are matrix-valued functions.

Original: 3. ELBO FOR MDM s
Revised: 3. ELBO FOR MDM s

Original: We now show how to train MDM s to optimize a lower bound on the log likelihood of the data.
Revised: We now show how to train MDM s to optimize a lower bound on the log likelihood of the data.

Original: Like in the univariate case, we use the parameterization in eq.
Revised: Like in the univariate case, we use the parameterization in eq.

Original: (3) to obtain a tractable ELBO .
Revised: (3) to obtain a tractable ELBO .

Original: Theorem 1.
Revised: Theorem 1.

Original: The MDM log marginal likelihood of the data is lower-bounded by:log p
Revised: The MDM log marginal likelihood of the data is lower-bounded by:

Original: [Equation]
Revised: [Equation]

Original: where divergences and gradients are taken with respect to y sand s q = ∇y s log q ( y s | x ) .
Revised: where divergences and gradients are taken with respect to y s ∇

Original: Proof.
Revised: Proof.

Original: The proof for the  L mismbound is in appendix E.
Revised: The proof for the MDM ISM ELBO L mism is in appendix F.

Original: In short, we introduce auxiliary variables, apply Theorem 1 of Huang et al.
Revised: In short, we introduce auxiliary variables, apply Theorem 1 of Huang et al.

Original: (2021) (equivalently, Theorem 3 of Song et al.
Revised: (2021) (equivalently, Theorem 3 of Song et al.

Original: (2021) or appendix

Original:  E of Kingma et al.
Revised: or appendix E of Kingma et al.

Original: (2021)) to the joint space, and then apply an additional variational bound to v 0 .
Revised: (2021)) to the joint space, and then apply an additional variational bound to v 0 .

Original: The  L mdsm bound follows by analogy of the ISM - DSM connection from Huang et al.
Revised: The MDM DSM ELBO L mdsm is likewise derived in appendix F, similarly to Huang et al.

Original: ; Song et al. (2021) for univariate diffusion models.
Revised: (2021); Song et al. but extended to multivariate diffusions.

Original: We call the first bound the MDM ISM bound L mism and the second the MDM DSM bound L mdsm .

Original: In practice, estimates of L mism can be computationally prohibitive. We train MDM ’s by estimating the gradients of L mdsm .
Revised: We train MDM ’s by estimating the gradients of L mdsm , as estimates of L mism can be computationally prohibitive.

Original: For numerical stability, the integral in eq.
Revised: For numerical stability, the integral in eq.

Original: (7) is computed on [ ϵ, T ] rather than [0 , T ] .
Revised: (7) is computed on [ ϵ, T ] rather than [0 , T ] .

Original: One can regard this as a bound for a variable u ϵ .
Revised: One can regard this as a bound for a variable u ϵ .

Original: To maintain a proper likelihood bound for the data, one can choose a likelihood u 0 | u ϵ and compose bounds as we demonstrate in appendix H.
Revised: To maintain a proper likelihood bound for the data, one can choose a likelihood u 0 | u ϵ and compose bounds as we demonstrate in appendix I.

Original: 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.
Revised: 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.

Original: 3. I NGREDIENT 1: C OMPUTING THE TRANSITION q ϕ ( y s | x )
Revised: 3. I NGREDIENT 1: C OMPUTING THE TRANSITION q ϕ ( y s | x )

Original: To estimate eq.
Revised: To estimate eq.

Original: (7) and its gradients, we need samples from q ( y s | x ) and to compute ∇ log q ( y s | x ) .
Revised: (7) and its gradients, we need samples from q ( y s | x ) and to compute ∇ log q ( y s | x ) .

Original: While an intractable problem for MDM s in general, we provide two ingredients for tightening and optimizing these bounds in a generic fashion for linear inference MDM s.
Revised: While an intractable problem for MDM s in general, we provide two ingredients for tightening and optimizing these bounds in a generic fashion for linear inference MDM s.

Original: We first show how to compute the mean and covariance of q ( y s | y 0 ) and thenhow to adapt this to the mean and covariance of q ( y s | x ) .
Revised: We first show how to automate computation of q ( y s | y 0 ) and then q ( y s | x ) .

Original: For linear MDM s of the form:
Revised: For linear MDM s of the form:

Original: [Equation]
Revised: [Equation]

Original: the transition kernel q ( y s | y 0 ) is Gaussian (S¨arkk¨a & Solin, 2019).
Revised: the transition kernel q ( y s | y 0 ) is Gaussian (S ¨ arkk ¨ a & Solin, 2019).

Original: Let f ( y , s )
Revised: Let f ( y , s )

Original: =
Revised: =

Original: A ( s ) y .
Revised: A ( s ) y .

Original: Then, the mean and covariance are solutions to the following ODE s:
Revised: Then, the mean and covariance are solutions to the following ODE s:

Original: [Equation]
Revised: [Equation]

Original: The mean can be solved analytically:
Revised: The mean can be solved analytically:

Original: [Equation]
Revised: [Equation]

Original: The covariance equation does not have as simple a solution because eq.
Revised: The covariance equation does not have as simple a solution because eq.

Original: (9) as the unknown matrix Σ s | 0 is being multiplied both from the left and the right.
Revised: (9) as the unknown matrix Σ s | 0 is being multiplied both from the left and the right.

Original: Instead of solving eq.
Revised: Instead of solving eq.

Original: (8) for a specific diffusion manually , as done in previous work (e.g. pages 5054 of Dockhorn et al.
Revised: (8) for a specific diffusion manually , as done in previous work (e.g. pages 5054 of Dockhorn et al.

Original: , we show that a matrix factorization technique (S ¨ arkk ¨ a & Solin (2019), sec. 6 . 3 ) previously unused in diffusion-based generative models can automatically compute Σ s |generically for any linear MDM .
Revised: , we show that a matrix factorization technique (S ¨ arkk ¨ a & Solin (2019), sec. 6 . 3 ) previously unused in diffusion-based generative models can automatically compute Σ s |generically for any linear MDM .

Original: Define C s , H s that evolve according to:
Revised: Define C s , H s that evolve according to:

Original: [Equation]
Revised: [Equation]

Original: then Σ s | 

Revised: then Σ s |

Original: =
Revised: =

Original: C s H s − 1 for C 0 = Σ 0 and H 

Revised: C s H s − 1 for C 0 = Σ 0 and H

Original: = I (Appendix C).
Revised: = I (Appendix D).

Original: These equations can be solved in closed-form,
Revised: These equations can be solved in closed-form,

Original: [Equation]
Revised: [Equation]

Original: no integration if A (

Original: [Equation] A [Equation] ) where [ A ] s = (cid:82) s 0 A ( ν ) dν .
Revised: where [ A ] s = (cid:82) s 0 A ( ν ) dν .

Original: To condition on y 

Revised: To condition on y

Original: = ( x, v ) , we set Σ 0 = 0 .
Revised: = ( x, v ) , we set Σ 0 = 0 .

Original: Computing q ϕ ( y s | x ) .
Revised: Computing q ϕ ( y s | x ) .

Original: For the covariance Σ s | 0 , to condition on x instead of y 0 , we set Σ 0 toΣ 0 = (cid:18) [Equation] (cid:19)
Revised: For the covariance Σ s | 0 , to condition on x instead of y 0 , we set Σ 0 to

Original: [Equation]
Revised: [Equation]

Original: To compute the mean, it is the same expression as for q ( y s | y 0 ) , but with a different initial condition:
Revised: To compute the mean, it is the same expression as for q ( y s | y 0 ) , but with a different initial condition:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: See appendix C for more details.
Revised: See appendix D for more details.

Original: Input: Data { x i } , inference process matrices Q ϕ , D ϕ , model prior π θ , initial distribution q ϕ
Revised: Input: Data { x i } , inference process matrices Q ϕ , D ϕ , model prior π θ , initial distribution q ϕ

Original: ( y v 0 | x ) , and score model architecture s θ
Revised: ( y v 0 | x ) , and score model architecture s θ

Original: Returns: Trained score model s θ while s θ not converged do
Revised: Returns: Trained score model s θ while s θ not converged do

Original: Sample x ∼ (cid:80) Ni =1 1 N δ x i , v 0 ∼ q ϕ
Revised: Sample x ∼ (cid:80) Ni =1 1 N δ x i , v 0 ∼ q ϕ

Original: ( y v 0 | x )
Revised: ( y v 0 | x )

Original: Sample s ∼ U [0 , T ] and y s , y T ∼ q ϕ ( y s | x ) using algorithm
Revised: Sample s ∼ U [0 , T ] and y s , y T ∼ q ϕ ( y s | x ) using algorithm

Original: Estimate the stochastic gradient of the MDM ELBO , ∇ θ L ( θ, ϕ ) , using eq.
Revised: Estimate the stochastic gradient of the MDM ELBO , ∇ θ L ( θ, ϕ ) , using eq.

Original: (7) θ ← θ + α ∇ θ L ( θ, ϕ ) .
Revised: (7) θ ← θ + α ∇ θ L ( θ, ϕ ) .

Original: if learning inference then ϕ ← ϕ + α ∇ ϕ L ( θ, ϕ ) end if end while
Revised: if learning inference then ϕ ← ϕ + α ∇ ϕ L ( θ, ϕ ) end if end while

Original: Output s θ
Revised: Output s θ

Original: [Table]
Revised: [Table]

Original: A fast and simple algorithm.
Revised: A fast and simple algorithm.

Original: We show in algorithm(appendix G) that computing the transition kernel only requires knowing f, g and requires no diffusion-specific analysis.
Revised: We show in algorithm(appendix H) that computing the transition kernel only requires knowing f, g and requires no diffusion-specific analysis.

Original: For K − 1 auxiliary variables, A , g are K × K .
Revised: For K − 1 auxiliary variables, A , g are K × K .

Original: Like for scalar diffusions, these parameters are shared across data coordinates.
Revised: Like for scalar diffusions, these parameters are shared across data coordinates.

Original: This means matrix exponentials and inverses are done on K × K matrices, where K is only 2 or 3 in our experiments.
Revised: This means matrix exponentials and inverses are done on K × K matrices, where K is only 2 or 3 in our experiments.

Original: In table 1, we compare the time to sample a batch of size 256 from the transition kernel for CIFAR 10 and MNIST .
Revised: In table 1, we compare the time to sample a batch of size 256 from the transition kernel for CIFAR 10 and MNIST .

Original: The table shows the extra computational cost of the automated algorithm is negligible.
Revised: The table shows the extra computational cost of the automated algorithm is negligible.

Revised: This automation likewise applies to simplified score matching and noise prediction objectives, since all rely on q ϕ ( y s | x ) (appendix A).

Original: 3. I NGREDIENT 2: MDM P ARAMETERIZATION
Revised: 3. I NGREDIENT 2: MDM P ARAMETERIZATION

Original: The MDM ELBO (eq. (7)) is tighter when the inference y T tends toward the model’s prior π θ .
Revised: The MDM ELBO (eq. (7)) is tighter when the inference y T tends toward the model’s prior π θ .

Original: Here we construct inference processes with the model prior π θ as a specified stationary distribution q ∞ .
Revised: Here we construct inference processes with the model prior π θ as a specified stationary distribution q ∞ .

Original: Ma et al.
Revised: Ma et al.

Original: (2015) provide a complete recipe for constructing gradient-based MCMC samplers; the recipe constructs non-linear time-homogeneous It ˆ o processes with a given stationary distribution, and show that the parameterization spans all such It ˆ o processes with that stationary distribution.
Revised: (2015) provide a complete recipe for constructing gradient-based MCMC samplers; the recipe constructs non-linear time-homogeneous Itˆo processes with a given stationary distribution, and show that the parameterization spans all such Itˆo processes with that stationary distribution.

Original: Diffusion models usually have time-varying drift and diffusion coefficients (e.g. use of the β ( t ) function).
Revised: Diffusion models usually have time-varying drift and diffusion coefficients (e.g. use of the β ( t ) function).

Original: To build diffusion models that match the model prior, we first extend Theorem 1 from
Revised: To build diffusion models that match the model prior, we first extend Theorem 1 from

Original: Ma et al.
Revised: Ma et al.

Original: (2015) to construct non-linear It ˆ
Revised: (2015) to construct non-linear It ˆ

Original: o processes with time-varying drift and diffusion coefficients with a given stationary distribution (Appendix B).
Revised: o processes with time-varying drift and diffusion coefficients with a given stationary distribution (Appendix C).

Original: Then, to keep transitions tractable (per
Revised: Then, to keep transitions tractable (per

Original: Section 3.3), we specialize this result to linear Itˆo diffusions.
Revised: Section 3.3), we specialize this result to linear Itˆo diffusions.

Original: We directly state the result for linear time-varying diffusions with stationary distributions.
Revised: We directly state the result for linear time-varying diffusions with stationary distributions.

Original: The parameterization requires a skew-symmetric matrix − Q ( s )
Revised: The parameterization requires a skew-symmetric matrix − Q ( s )

Original: =
Revised: =

Original: Q ( s ) ⊤ , a positive semi-definite matrix
Revised: Q ( s ) ⊤ , a positive semi-definite matrix

Original: D ( s ) , and a function ∇ H ( y ) such that the desired stationary distribution q ∞ is proportional to exp[ − H ( y )] .
Revised: D ( s ) , and a function ∇ H ( y ) such that the desired stationary distribution q ∞ is proportional to exp[ − H ( y )] .

Original: Linear Itˆo diffusions have Gaussian stationary distributions (S¨arkk¨a & Solin, 2019) meaning that ∇ H is linear and can be expressed as Sy for some matrix S .
Revised: Linear Itˆo diffusions have Gaussian stationary distributions (S¨arkk¨a & Solin, 2019) meaning that ∇ H is linear and can be expressed as Sy for some matrix S .

Original: For a matrix A , let √ A refer to the matrix square root defined by a = √
Revised: For a matrix A , let √ A refer to the matrix square root defined by a = √

Original: A ⇐⇒
Revised: A ⇐⇒

Original: A = aa ⊤ .
Revised: A = aa ⊤ .

Original: Then, the Itˆo diffusion:
Revised: Then, the Itˆo diffusion:

Original: [Equation]
Revised: [Equation]

Original: has Gaussian stationary distribution N ( 0 , S − 1 ) where Q ( s ) , D ( s ) and S are parameters.
Revised: has Gaussian stationary distribution N ( 0 , S − 1 ) where Q ( s ) , D ( s ) and S are parameters.

Original: For a discussion ofuniqueness 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.
Revised: For a discussion of convergence to the stationary distribution, as well as skew-symmetric and positive semi-definite parameterizations, see appendix C, where we also show that existing diffusion processes such as VPSDE and CLD are included in Q / D parameterization.

Original: We display the ELBO in terms of Q / D in appendix F and an algorithm in appendix G.
Revised: We display the ELBO in terms of Q / D in appendix G and an algorithm in appendix H.

Revised: For score matching and noise prediction losses and a given q ϕ , achieving a minimizing value with respect to s θ does not imply that the generative model score will match the inference score.

Revised: Modeling the data also requires the marginal distribution of q ϕ,T to approximate π .

Revised: When q ϕ is constant, it is important to confirm the stationary distribution is appropriately set, and the tools used here for the

Revised: ELBO can be used to satisfy this requirement for score matching and noise prediction (appendix A).

Original: 3. L EARNING THE INFERENCE PROCESS
Revised: 3. L EARNING THE INFERENCE PROCESS

Original: The choice of diffusion matters, and the ELBO s in eq.
Revised: The choice of diffusion matters, and the ELBO s in eq.

Original: (7) have no requirement for fixed q ϕ .
Revised: (7) have no requirement for fixed q ϕ .

Original: We therefore learn the inference process jointly with s θ .
Revised: We therefore learn the inference process jointly with s θ .

Original: Under linear transitions (ingredient 1 ), no algorithmic details change as the diffusion changes during training.
Revised: Under linear transitions (ingredient 1 ), no algorithmic details change as the diffusion changes during training.

Original: Under stationary parameterization (ingredient 2 ), we can learn without the stationary distribution going awry.
Revised: Under stationary parameterization (ingredient 2 ), we can learn without the stationary distribution going awry.

Original: In the experiments, learning matches or surpasses BPD s of fixed diffusions for a given dataset and score architecture.
Revised: In the experiments, learning matches or surpasses BPD s of fixed diffusions for a given dataset and score architecture.

Original: In L mdsm or L mism , q ϕ, ∞ may be set to equal π θ , but it is y T ∼ q ϕ,T for the chosen T that is featured in the ELBO .
Revised: In L mdsm or L mism , q ϕ, ∞ may be set to equal π θ , but it is y T ∼ q ϕ,T for the chosen T that is featured in the ELBO .

Original: Learning q ϕ can choose y T to reduce the cross-entropy:
Revised: Learning q ϕ can choose y T to reduce the cross-entropy:

Original: [Equation]
Revised: [Equation]

Original: Minimizing eq.
Revised: Minimizing eq.

Original: (14) will tighten the ELBO for any s θ .
Revised: (14) will tighten the ELBO for any s θ .

Original: Next, q ϕ is featured in the remaining terms; optimizing for q ϕ will tighten and improve the ELBO alongside s θ .
Revised: Next, q ϕ is featured in the remaining terms that feature s θ ; optimizing for q ϕ will tighten and improve the ELBO alongside s θ .

Original: Finally, q ϕ is featured in the expectations and the − log q ϕ term:
Revised: Finally, q ϕ is featured in the expectations and the − log q ϕ term:

Original: [Equation]
Revised: [Equation]

Original: The q ϕ ( y v 0 | x ) terms impose an optimality condition that p θ ( u v T | u z T ) =
Revised: The q ϕ ( y v 0 | x ) terms impose an optimality condition that p θ ( u vT | u zT ) =

Original: q ϕ ( y v 0 | y z 0 ) (appendix D),
Revised: q ϕ ( y v 0 | y z 0 ) (appendix E),

Original: When it is satisfied, no looseness in the ELBO is due to the initial time zero auxiliary variables.
Revised: When it is satisfied, no looseness in the ELBO is due to the initial time zero auxiliary variables.

Original: To learn, Q , D need to be specified with parameters
Revised: To learn, Q , D need to be specified with parameters

Original: ϕ that enable gradients.
Revised: ϕ that enable gradients.

Original: We keep S fixed at inverse covariance of π θ .
Revised: We keep S fixed at inverse covariance of π θ .

Original: The choice of Q , D generates the distribution parameters of the Gaussian distributions q ϕ ( y s | x )  through the meanm s | 0 and covariance Σ s | 0 .
Revised: The transition kernel q ϕ ( y s | x ) depends on Q , D through its mean and covariance.

Original: Gaussian distributions permit gradient estimation with reparameterization or score-function gradients (Kingma & Welling, 2013; Ranganath et al., 2014; Rezende & Mohamed, 2015; Titsias & L ´ azaro-Gredilla, 2014).
Revised: Gaussian distributions permit gradient estimation with reparameterization or scorefunction gradients (Kingma & Welling, 2013; Ranganath et al., 2014; Rezende & Mohamed, 2015; Titsias & L´azaro-Gredilla, 2014).

Original: Reparameterization is accomplished via:
Revised: Reparameterization is accomplished via:

Original: [Equation]
Revised: [Equation]

Original: 16) where ϵ ∼ N (0 , I dK ) and L s | [-
-] satisfies L s | 

Revised: ) where ϵ ∼ N (0 , I dK ) and L s | [-
-]0 satisfies L s |

Original: L ⊤ s | 

Revised: L ⊤ s |

Original: = Σ s | 0 , derived using coordinate-wise Cholesky decomposition.
Revised: = Σ s | 0 , derived using coordinate-wise Cholesky decomposition.

Original: Gradients flow through eq.
Revised: Gradients flow through eq.

Original: (16) from y s to m s | 0 and Σ s | 0 to Q , D to parameters ϕ .
Revised: (16) from y s to m s | 0 and Σ s | 0 to Q , D to parameters ϕ .

Original: Algorithm 1 displays Automatic Multivariate Diffusion Training ( AMDT ).
Revised: Algorithm 1 displays Automatic Multivariate Diffusion Training ( AMDT ).

Original: A MDT provides a training method for diffusion-based generative models for either fixed Q , D matrices or for learning the
Revised: A MDT provides a training method for diffusion-based generative models for either fixed Q , D matrices or for learning the

Original: Q ϕ , D ϕ matrices, without requiring any diffusion-specific analysis.
Revised: Q ϕ , D ϕ matrices, without requiring any diffusion-specific analysis.

Revised: Learning in other diffusion objectives.

Revised: Like in the ELBO , learning in score matching or noise prediction objectives can improve the match between the inference process and implied generative model (appendix A).

Original:  I NSIGHTS INTO M ULTIVARIATE D IFFUSIONS
Revised:  I NSIGHTS INTO M ULTIVARIATE D IFFUSIONS

Original: Scalar versus Multivariate Processes.
Revised: Scalar versus Multivariate Processes.

Original: Equation (13) clarifies what can change while preserving q ∞ .
Revised: Equation (13) clarifies what can change while preserving q ∞ .

Original: Recall that Q and D are K × K for K − 1 auxiliary variables.
Revised: Recall that Q and D are K × K for K − 1 auxiliary variables.

Original: Because 0 is the only 1 ×skew-symmetric matrix, scalar processes must set Q = 0 .
Revised: Because 0 is the only 1 ×skew-symmetric matrix, scalar processes must set Q = 0 .

Original: With q ϕ, ∞ =
Revised: With q ϕ, ∞ =

Original: N (0 , I ) , the process is:
Revised: N (0 , I ) , the process is:

Original: [Equation]
Revised: [Equation]

Original: What is left is the VPSDE process used widely in diffusion models where D ( s ) = 12 β ( s ) is 1 × 1 (Song et al., 2020b).
Revised: What is left is the VPSDE process used widely in diffusion models where D ( s ) = 12 β ( s ) is 1 × 1 (Song et al., 2020b).

Original: This reveals that the VPSDE process is the only scalar diffusion with a stationary distribution.
Revised: This reveals that the VPSDE process is the only scalar diffusion with a stationary distribution.

Original: This also clarifies the role of Q : it accounts for mixing between dimensions in multivariate processes, as do non-diagonal entries in D for K > 1 .
Revised: This also clarifies the role of Q : it accounts for mixing between dimensions in multivariate processes, as do non-diagonal entries in D for K > 1 .

Original: CLD optimizes a  lower bound on the data log-likelihood.
Revised: CLD optimizes a log-likelihood lower bound.

Original: Taking the derivative of L 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 θ ( x )  without appealing to the probability flow ODE .
Revised: Differentiating L mdsm (eq. (7)) with respect to the score model parameters, we show that the objective for CLD (Dockhorn et al., 2021) maximizes a lower bound on log p θ ( x ) , not just log p θ ( u 0 ) , without appealing to the probability flow ODE .

Original: Does my model use auxiliary variables?
Revised: Does my model use auxiliary variables?

Original: An example initial distribution is q ( y v 0 | x ) =
Revised: An example initial distribution is q ( y v 0 | x ) =

Original: N (0 , I ) .
Revised: N (0 , I ) .

Original: It is also common to set π θ = N (0 , I ) .
Revised: It is also common to set π θ = N (0 , I ) .

Original: Because the optimum for diffusions is p θ = q , the optimal model has main and auxiliary dimensions independent at endpoints 0 and T .
Revised: Because the optimum for diffusions is p θ = q , the optimal model has main and auxiliary dimensions independent at endpoints 0 and T .

Original: Does this mean that the model does not use auxiliary variables?
Revised: Does this mean that the model does not use auxiliary variables?

Original: In appendix A, we show that in this case the model can still use auxiliary variables at intermediate times.
Revised: In appendix B, we show that in this case the model can still use auxiliary variables at intermediate times.

Original: A sufficient condition is non-diagonal Q + D .
Revised: A sufficient condition is non-diagonal Q + D .

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We test the MDM framework with handcrafted and learned diffusions.
Revised: We test the MDM framework with handcrafted and learned diffusions.

Original: The handcrafted diffusions are (a) A LDA , used in (Mou et al., 2019) for accelerated gradient-based MCMC sampling (eq.
Revised: The handcrafted diffusions are (a) A LDA , used in (Mou et al., 2019) for accelerated gradient-based MCMC sampling (eq.

Original: (32)) and (b) M ALDA : a modified version of ALDA (eq.
Revised: (32)) and (b) M ALDA : a modified version of ALDA (eq.

Original: Both have two auxiliary variables.
Revised: Both have two auxiliary variables.

Original: We also learn diffusions with 1 and 2 auxiliary variables.
Revised: We also learn diffusions with 1 and 2 auxiliary variables.

Original: We compare with VPSDE and ELBO -trained CLD .
Revised: We compare with VPSDE and ELBO -trained CLD .

Original: [Table] [Table]
Revised: [Table] [Table]

Original: Following prior work, we train DBGM s for image generation.
Revised: Following prior work [Table] we train DBGM s for image generation.

Original: We use the U-Net from Ho et al.
Revised: We use the U-Net from Ho et al.

Original: 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 IMAGENET 32 and only 865 more for MNIST ).
Revised: 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 IMAGENET 32 and only 865 more for MNIST ).

Original: We use simple uniform dequantization.
Revised: We use simple uniform dequantization.

Original: We report estimates of L mdsm (which reduces to the standard L dsm for K = 1 ).
Revised: We report estimates of L mdsm (which reduces to the standard L dsm for K = 1 ).

Original: We sample times using the importance sampling distribution from Song et al.
Revised: We sample times using the importance sampling distribution from Song et al.

Original: (2021) with truncation set to ϵ = 10 − 3 .
Revised: (2021) with truncation set to ϵ = 10 − 3 .

Original: To ensure the truncated bound is proper, we use a likelihood described in appendix H.
Revised: To ensure the truncated bound is proper, we use a likelihood described in appendix I.

Original: Results.
Revised: Results.

Original: Table 2 shows that the inference process matters and displays.
Revised: Table 2 shows that the inference process matters and displays.

Original: It displays DBGM s that we train and evaluate on CIFAR -10, IMAGENET 32 and MNIST .
Revised: It displays DBGM s that we train and evaluate on CIFAR -10, IMAGENET 32 and MNIST .

Original: This includes the existing VPSDE and
Revised: This includes the existing VPSDE and

Original: CLD , the new MALDA and ALDA , and the new learned inference processes.
Revised: CLD , the new MALDA and ALDA , and the new learned inference processes.

Original: All are trained with the 35.7M parameter architecture.
Revised: All are trained with the 35.7M parameter architecture.

Original: For CIFAR -10, learning outperforms CLD , and both outperform the standard choice of VPSDE .
Revised: For CIFAR -10, learning outperforms CLD , and both outperform the standard choice of VPSDE .

Original: For MNIST , learned diffusions match VPSDE while the three fixed auxiliary diffusions are worse.
Revised: For MNIST , learned diffusions match VPSDE while the three fixed auxiliary diffusions are worse.

Original: On IMAGENET 32, all perform similarly.
Revised: On IMAGENET 32, all perform similarly.

Original: 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.
Revised: 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.

Original: In Figure 1 we plot the generated samples from CIFAR 10.
Revised: In Figure 1 we plot the generated samples from CIFAR 10.

Original: Table 3’s first two rows display diffusion models from previous work: VPSDE (Song et al., 2021) and
Revised: Table 3’s first two rows display diffusion models from previous work: VPSDE (Song et al., 2021) and

Original: CLD (Dockhorn et al., 2021) both with the 108 million score model from Song et al. (2021) (labeled “large”).
Revised: CLD (Dockhorn et al., 2021) both with the 108 million score model from Song et al. (2021) (labeled “large”).

Original: The rest are DBGM s that we train using the U-Net with 35.7 million parameters for CIFAR10 and IMAGENET 32 and 1.1 million for MNIST .
Revised: The rest are DBGM s that we train using the U-Net with 35.7 million parameters for CIFAR10 and IMAGENET 32 and 1.1 million for MNIST .

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Table 3 also demonstrates a tighter bound for CLD trained and evaluated with the MDM
Revised: Table 3 also demonstrates a tighter bound for CLD trained and evaluated with the MDM

Original: ELBO ( ≤ 3 . 11 ) relative to existing probability flow-based evaluations ( 3 . 31 ).
Revised: ELBO ( ≤ 3 . 11 ) relative to existing probability flow-based evaluations ( 3 . 31 ).

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Evidence Lower Bounds.
Revised: Evidence Lower Bounds.

Original: Song et al.
Revised: Song et al.

Original: ; Huang et al.
Revised: ; Huang et al.

Original: (2021) derive the ISM and DSM lowerbounds on the model log likelihood.
Revised: (2021) derive the ISM and DSM lowerbounds on the model log likelihood.

Original: 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.
Revised: 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.

Original: Auxiliary variables.
Revised: Auxiliary variables.

Original: Dupont et al.
Revised: Dupont et al.

Original: (2019) shows that augmented neural ODE s model a richer set of functions and Huang et al.
Revised: (2019) shows that augmented neural ODE s model a richer set of functions and Huang et al.

Original: (2020) uses this principle for normalizing flows.
Revised: (2020) uses this principle for normalizing flows.

Original: Hierarchical variational models and auto-encoders marginalize auxiliary variables to build expressive distributions (Ranganath et al., 2016;
Revised: Hierarchical variational models and auto-encoders marginalize auxiliary variables to build expressive distributions (Ranganath et al., 2016;

Original: Sønderby et al., 2016; Maaløe et al., 2019; Vahdat & Kautz, 2020; Child, 2020).
Revised: Sønderby et al., 2016; Maaløe et al., 2019; Vahdat & Kautz, 2020; Child, 2020).

Original: We apply this principle to DBGM s, including and extending CLD (Dockhorn et al., 2021).
Revised: We apply this principle to DBGM s, including and extending CLD (Dockhorn et al., 2021).

Original: Learning inference.
Revised: Learning inference.

Original: Learning q ϕ with p θ is motivated in previous work (Kingma & Welling, 2013; Sohl-Dickstein et al., 2015; Kingma et al., 2021).
Revised: Learning q ϕ with p θ is motivated in previous work (Kingma & Welling, 2013; Sohl-Dickstein et al., 2015; Kingma et al., 2021).

Original: Kingma et al.
Revised: Kingma et al.

Original: (2021) learn the noise schedule for VPSDE .
Revised: (2021) learn the noise schedule for VPSDE .

Original: For MDM s, there are parameters to learn beyond the noise schedule; Q can be non-zero, D can diagonal or full, give Q and D different time-varying functions, and learn ∇ H .
Revised: For MDM s, there are parameters to learn beyond the noise schedule; Q can be non-zero, D can diagonal or full, give Q and D different time-varying functions, and learn ∇ H .

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: The experiments show that learning matches or surpasses the best fixed diffusion performance, bypassing the need to choose a diffusion.
Revised: The experiments show that learning matches or surpasses the best fixed diffusion performance, bypassing the need to choose a diffusion.

Original: MDM s achieve BPD s similar to univariate diffusions, with as many as three times more score parameters.
Revised: MDM s achieve BPD s similar to univariate diffusions, with as many as three times more score parameters.

Original: The proposed MDM ELBO reports a tighter bound for the existing CLD relative to existing probability flow-based evaluations.
Revised: The proposed MDM ELBO reports a tighter bound for the existing CLD relative to existing probability flow-based evaluations.

Original: This work enables future directions including interactions across data coordinates and using new stationary distributions.
Revised: This work enables future directions including interactions across data coordinates and using new stationary distributions.

Original:  A CKNOWLEDGEMENTS
Revised:  A CKNOWLEDGEMENTS

Original: This work was generously funded by NIH/NHLBI Award R01HL148248, NSF Award
Revised: This work was generously funded by NIH/NHLBI Award R01HL148248, NSF Award

Original: NRT-HDR: FUTURE Foundations, Translation, and Responsibility for Data Science, and NSF CAREER Award 2145542.
Revised: NRT-HDR: FUTURE Foundations, Translation, and Responsibility for Data Science and NSF CAREER Award 2145542.

Original: The authors would additionally like to thank Chin-Wei Huang for helpful discussing regarding Huang et al.
Revised: The authors would additionally like to thank Chin-Wei Huang for helpful discussing regarding Huang et al.

Original: A D OES MY MODEL USE AUXILIARY VARIABLES ?
Revised: A A UTOMATED S CORE M ATCHING WITH L EARNED I NFERENCE

Revised: Like for the MDM ELBO , the methods in this work apply to training with the score matching loss: [Equation] where λ : [0 , T ] → R + is a weighing function.

Revised: The score-matching loss is often optimized in its simplified noise prediction form:

Revised: [Equation]

Revised: where s θ = − L −⊤ t ϵ θ and y t = µ t + L t ϵ and ϵ is the noise used in sampling y t .

Revised: We describe here how the improvements to the ELBO studied in this work carry over to L SM and L NP .

Revised: In the following let q 0 be the data distribution, let p ( θ,ϕ ) , 0 be the model’s distribution of the data, and recall that the model is defined by ( s θ , f ϕ , g ϕ ) and prior π via a continuous-time stochastic process with drift coefficient g 2 ϕ s θ − f ϕ and and diffusion coefficient g ϕ .

Revised: First, minimizing L SM or L NP so that ∇ y t log q ϕ ( y t ) =

Revised: s θ ( y t , t ) does not alone imply that p ( θ,ϕ ) ,will equal q 0 ; it must also be that q ϕ,T ≈

Revised: π .

Revised: Foregoing this requirement means π will produce samples that the generative model may not be able to push onto the path the model was trained on (formally, the score of the generative model would not equal the time-reversal of the forward score even if s θ equals the forward score).

Revised: This condition can be satisfied if q ϕ can be chosen with stationary distribution π .

Revised: Section 3.4 describes how to accomplish this.

Revised: Next, for any fixed q ϕ , automatic transitions from section 3.3 streamline the computation of the score matching loss, allowing for simple score computation for a wide class of diffusions beyond VP.

Revised: Finally, for a fixed q ϕ with q ϕ,T ≈ π and a score architecture s θ , minimizing L SM or L NP w.r.t θ may be suboptimal.

Revised: Optimization, like for the elbo, carries over to score matching and can close this gap; learning w.r.t.

Revised: both θ, ϕ increases the ability to successfully minimize the loss at each t (section 3.5).

Revised: In other words, since the generative model is defined by ( s θ , f ϕ , g ϕ ) , learning q ϕ means the loss trains all three components of the generative model rather than just one.

Revised: In summary, score matching is automatic and can learn over the space of linear diffusions that tend to the model prior.

Revised: B D OES MY MODEL USE AUXILIARY VARIABLES ?

Original: In section 3 we gave the example choice of q ( y v 0 | x ) =
Revised: In section 3 we gave the example choice of q ( y v 0 | x ) =

Original: N (0 , I ) coordinate-wise.
Revised: N (0 , I ) coordinate-wise.

Original: It is also a common choice to set π θ = N (0 , I ) .
Revised: It is also a common choice to set π θ = N (0 , I ) .

Original: Because the optimum in diffusion models is p θ = 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 .
Revised: Because the optimum in diffusion models is p θ = 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 .

Original: Does this mean that the model does not use auxiliary variables?
Revised: Does this mean that the model does not use auxiliary variables?

Original: We show that even when q ϕ ( y 0 ) and π θ have main and auxiliary variables independent, the model can use the auxiliary variables.
Revised: We show that even when q ϕ ( y 0 ) and π θ have main and auxiliary variables independent, the model can use the auxiliary variables.

Original: A sufficient condition is Q + D is non-diagonal.
Revised: A sufficient condition is Q + D is non-diagonal.

Original: To make this precise, we recall that we model with p θ ( u zT =
Revised: To make this precise, we recall that we model with p θ ( u zT =

Original: x
Revised: x

Original: To show the model is using auxiliary variables, we just need to show that u zT (main coordinate at T ) depends on u vt (aux. coordinate at t ) for T > t .
Revised: To show the model is using auxiliary variables, we just need to show that u zT (main coordinate at T ) depends on u vt (aux. coordinate at t ) for T > t .

Original: At optimum, p θ ( u zT , u vt ) =
Revised: At optimum, p θ ( u zT , u vt ) =

Original: q ϕ ( y z 0 , y vT − t ) .
Revised: q ϕ ( y z 0 , y vT − t ) .

Original: Therefore it is sufficient to show that for some time s , q ϕ ( y vs | y z 0 ) ̸
Revised: Therefore it is sufficient to show that for some time s , q ϕ ( y vs | y z 0 ) ̸

Original: = q ϕ
Revised: = q ϕ

Original: ( y vs ) .
Revised: ( y vs ) .

Original: Because y z 0 , is determined by x we need to show that q ϕ ( y vs | x ) ̸ = q ϕ
Revised: Because y z 0 , is determined by x we need to show that q ϕ ( y vs | x ) ̸ = q ϕ

Original: ( y vs ) .
Revised: ( y vs ) .

Original: To do that, we first derive q ( y s | x ) and then marginalize to get q ( y vs | x ) from it.
Revised: To do that, we first derive q ( y s | x ) and then marginalize to get q ( y vs | x ) from it.

Original: Since the former is 2D Gaussian, the latter is available in terms of the former’s mean and covariance.
Revised: Since the former is 2D Gaussian, the latter is available in terms of the former’s mean and covariance.

Original: Suppose E
Revised: Suppose E

Original: [ y v 0 ] = 0 , Q =
Revised: [ y v 0 ] = 0 , Q =

Original: [[0 , − 1] , [1 , 0]] and D =
Revised: [[0 , − 1] , [1 , 0]] and D =

Original: and we have s = .
Revised: and we have s = .

Original: 1 We have:
Revised: 1 We have:

Original: [Equation]
Revised: [Equation]

Original: (cid:105) (cid:18) x

Original: (cid:19)

Original: = (cid:18)

Original: [Equation]

Original: 9003 x

Original: [Equation] [Equation] 090 x (cid:19) (18) Regardless of the covariance any 1D of this 2D gaussian will have mean that is a function of x , meaning that q ( y vs | x ) does not equal q ( y vs ) (which is also a Gaussian but with mean depending on x ′ s mean rather than x itself.
Revised: Regardless of the covariance any 1D of this 2D gaussian will have mean that is a function of x , meaning that q ( y v s | x ) does not equal q ( y vs ) (which is also a Gaussian but with mean depending on x ′ s mean rather than x itself.

Original: Therefore, even under the setup with independent endpoints, the optimal model makes use of the intermediate auxiliary variables in its final modeling distribution p θ ( u zT = x ) .
Revised: Therefore, even under the setup with independent endpoints, the optimal model makes use of the intermediate auxiliary variables in its final modeling distribution p θ ( u zT = x ) .

Original: Are there choices of Q and D that lead to learning models that don’t make use of the extra dimensions?
Revised: Are there choices of Q and D that lead to learning models that don’t make use of the extra dimensions?

Original: As mentioned, in the inference process, Q is responsible for mixing information among the coordinates, and is the only source of this when D is diagonal.
Revised: As mentioned, in the inference process, Q is responsible for mixing information among the coordinates, and is the only source of this when D is diagonal.

Original: Then, if Q = 

Revised: Then, if Q =

Original: and D is diagonal, none of the coordinates for a given feature x j (including u ztj , u v 1 tj , . . .
Revised: and D is diagonal, none of the coordinates for a given feature x j (including u ztj , u v 1 tj , . . .

Original: , u v K − 1 tj ) interact for any t .
Revised: , u v K − 1 tj ) interact for any t .

Original: Then, since p θ = q at optimum, independence of the coodinates at all t in q imply the same in p θ and the model will not make use of any auxiliary variables when modeling the marginal log p θ ( u zT = x ) .
Revised: Then, since p θ = q at optimum, independence of the coodinates at all t in q imply the same in p θ and the model will not make use of any auxiliary variables when modeling the marginal log p θ ( u zT = x ) .

Original: B S TATIONARY PARAMETERIZATION
Revised: C S TATIONARY PARAMETERIZATION

Original: The non-linear time-homogeneous It ˆ o process family is:
Revised: The non-linear time-homogeneous Itˆo process family is:

Original: [Equation]
Revised: [Equation]

Original: This family can be restricted to those with stationary distributions.
Revised: This family can be restricted to those with stationary distributions.

Original: Ma et al.
Revised: Ma et al.

Original: (2015) show a complete recipe to span the subset of this family with a desired stationary distribution.
Revised: (2015) show a complete recipe to span the subset of this family with a desired stationary distribution.

Original: Let Q be skewsymmetric ( − Q = Q ⊤ ) and D is positive semi-definite.
Revised: Let Q be skewsymmetric ( − Q = Q ⊤ ) and D is positive semi-definite.

Original: Suppose the desired stationary distribution is q ∞ ( y ) .
Revised: Suppose the desired stationary distribution is q ∞ ( y ) .

Original: For a matrix A , let √ A refer to the matrix square root defined by a = √
Revised: For a matrix A , let √ A refer to the matrix square root defined by a = √

Original: A ⇐⇒
Revised: A ⇐⇒

Original: A = aa ⊤ .
Revised: A = aa ⊤ .

Original: Then, Ma et al.
Revised: Then, Ma et al.

Original: (2015) show that, setting H ( y )
Revised: (2015) show that, setting H ( y )

Original: = − log q ∞ ( y ) , g ( y ) =
Revised: = − log q ∞ ( y ) , g ( y ) =

Original: (cid:112) 2 D ( y ) , and 
Revised: (cid:112) 2 D ( y ) , and f ( y

Revised: [Equation]

Revised: = − [ D

Revised: [Equation]

Revised: y

Original: [Equation]
Revised: [Equation]

Revised: + Q

Revised: [Equation]

Revised: y

Revised: [Equation]

Revised: [Equation]

Revised: H ( y

Revised: [Equation]

Revised: + Γ ( y

Revised: [Equation]

Revised: [Equation]

Original: yields a process y t with stationary distribution q ∞ .
Revised: yields a process y t with stationary distribution q ∞ .

Original: We extend it to time-varying (time inhomogeneous) processes.
Revised: We extend it to time-varying (time inhomogeneous) processes.

Original: Theorem 2.
Revised: Theorem 2.

Original: q ∞ ( y )
Revised: q ∞ ( y )

Original: ∝ exp[ − H
Revised: ∝ exp[ − H

Original: ( y )] is a stationary distribution of
Revised: ( y )] is a stationary distribution of

Original: [Equation]
Revised: [Equation]

Original: for
Revised: for

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: The Fokker Planck equation is:
Revised: The Fokker Planck equation is:

Original: [Equation]
Revised: [Equation]

Original: A stationary distribution is one where the Fokker-Planck right hand side is equal to 0 .
Revised: A stationary distribution is one where the Fokker-Planck right hand side is equal to 0 .

Original: To show that the stationary characterization also holds of time-inhomogenous processes with D ( y , t ) and
Revised: To show that the stationary characterization also holds of time-inhomogenous processes with D ( y , t ) and

Original: Q ( y , t ) , we take two steps, closely following Yin & Ao (2006); Shi et al.
Revised: Q ( y , t ) , we take two steps, closely following Yin & Ao (2006); Shi et al.

Original: ; Ma et al.
Revised: ; Ma et al.

Original: but noting that there is no requirement for Q , D to be free of t .
Revised: but noting that there is no requirement for Q , D to be free of t .

Original: First, we show that the Fokker-Plack equation can be re-written as: 
Revised: First, we show that the Fokker-Plack equation can be re-written as: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Second, because the whole expression is set to 0 when the inside expression equals
Revised: Second, because the whole expression is set to 0 when the inside expression equals

Original: [Equation]
Revised: [Equation]

Original:  we just need to show that this holds when q ( y , t ) =
Revised: (25) we just need to show that this holds when q ( y , t ) =

Original: exp[ − H ( y )] / Z .
Revised: exp[ − H ( y )] / Z .

Original: The second step is concluded because
Revised: The second step is concluded because

Original: [Equation]
Revised: [Equation]

Original: where Z is the normalization constant of exp( − H ( y )) .
Revised: where Z is the normalization constant of exp( − H ( y )) .

Original: It only remains to show that Fokker-Plack can be re-written in divergence form with time-dependent
Revised: It only remains to show that Fokker-Plack can be re-written in divergence form with time-dependent

Original: Q , D .
Revised: Q , D .

Original: In the following let Q ijt denote Q ij ( y , t ) and likewise for D ijt .
Revised: In the following let Q ijt denote Q ij ( y , t ) and likewise for D ijt .

Original: Let ∂ i denote ∂∂ y i and let it denote dd y i for scalar functions.
Revised: Let ∂ i denote ∂∂ y i and let it denote dd y i for scalar functions.

Original: We will use [ Ax ] i = (cid:80)
Revised: We will use [ Ax ] i = (cid:80)

Original: j A ij x j .
Revised: j A ij x j .

Original: [Equation]
Revised: [Equation]

Original: We re-write the 2nd and 3rd term.
Revised: We re-write the 2nd and 3rd term.

Original: Holding i fixed and noting q is scalar, we get the product rule (cid:80)j D ijt ( ∂ j q )
Revised: Holding i fixed and noting q is scalar, we get the product rule (cid:80)

Original: =
Revised: [Equation]

Original: (cid:80)j
Revised: (cid:80)

Original: ∂ j [ D ijt q ] − q (cid:80) j

Original:  ∂ j D ijt for each i , and likewise for q : 
Revised: j ∂ j D ijt for each i , and likewise for q : (cid:88)

Original: [Equation]
Revised: [Equation]

Original: i ∂ i (cid:88) j [ D ijt + Q ijt ][ q∂ j H ] + (cid:88) i ∂ i (cid:88) j D ijt [ ∂ j q ] + (cid:88) i ∂i (cid:88)
Revised: Because Q ( y , t ) is skew-symmetric, we have that (cid:80) i ∂

Original:  jQ ijt [ ∂ jq ] = (cid:88) i ∂ i (cid:88) j [D ijt + Q ijt ][ q∂ j H ] + (cid:88) i ∂ i(cid:88) j ∂ j
Revised: i (cid:80) j ∂ j [ Q ijt q ] = 0 , leaving ∂ t q t = (cid:88) i ∂ i

Original: [ D ijt q ] − q (cid:88) j ∂ j D ijt + (cid:88) i ∂ i (cid:88)

Original: j ∂ j [ Q ijt q ] − q (cid:88)

Original: j ∂ j Q ijt

Original: Because Q ( y , t ) is skew-symmetric, we have that (cid:80) i ∂

Original: i (cid:80) j ∂ j [ Q ijt q ] = 0 , leaving ∂ t q t = (cid:88)

Original: [Equation]
Revised: [Equation]

Original: = (cid:88) 
Revised: = (cid:88) i ∂

Original: [Equation]
Revised: [Equation]

Original: = [Equation] D ijt q
Revised: i

Original: [Equation]
Revised: [Equation]

Original: Recalling that f i ( y , t ) =
Revised: Recalling that f i ( y , t ) =

Original: (cid:16) − [ D + Q ] ∇ H +Γ (cid:17) i and again that [ Ax ] i = (cid:80)
Revised: (cid:16) − [ D + Q ] ∇ H +Γ (cid:17) i and again that [ Ax ] i = (cid:80)

Original: j A ij x j , we have equality with the original Fokker-Planck= (cid:88)
Revised: j A ij x j , we have equality with the original Fokker-Planck

Original: [Equation]
Revised: [Equation]

Original: We have shown exp[ − H ( y )]
Revised: We have shown exp[ − H ( y )]

Original: / Z is a stationary distribution of the time-varying non-linear Itˆo process:
Revised: / Z is a stationary distribution of the time-varying non-linear Itˆo process:

Original: [Equation]
Revised: [Equation]

Original: However, for some choices of Q [Equation] D , exp[ − H ( y )] / Z is not necessarily the unique stationary distribution.
Revised: However, for some choices of Q , D , exp[ − H ( y )] / Z is not necessarily the unique stationary distribution.

Original: One problematic case can occur as follows.
Revised: One problematic case can occur as follows.

Original: Suppose that row i of ( Q + D ) is all-zero; in this case, d y i = 

Revised: Suppose that row i of ( Q + D ) is all-zero; in this case, d y i =

Original: which implies that ( y i )
Revised: which implies that ( y i )

Original: t
Revised: t

Original: =
Revised: =

Original: ( y i ) 

Revised: ( y i )

Original: for all t > 0 .
Revised: for all t > 0 .

Original: Then, the initial distribution is also a stationary distribution.
Revised: Then, the initial distribution is also a stationary distribution.

Original: To rule out such pathological diffusions [Equation] we make the assumption that
Revised: To rule out such pathological diffusions, we make the assumption that

Original: Q + D is full rank.
Revised: Q + D is full rank.

Original: Then, for uniqueness, recall that stationary distributions are the zeros of
Revised: Then, for uniqueness, recall that stationary distributions are the zeros of

Original: [Equation]
Revised: [Equation]

Original: where the expression is of the form Av for A
Revised: where the expression is of the form Av for A

Original: = D ( y , t ) + Q ( y , t ) and
Revised: = D ( y , t ) + Q ( y , t ) and

Original: [Equation]
Revised: [Equation]

Original: Under the assumption that Q + D is full rank, the expression can only be zero when v is zero.
Revised: Under the assumption that Q + D is full rank, the expression can only be zero when v is zero.

Original: To show uniqueness under the full rank assumption, one must then show that
Revised: To show uniqueness under the full rank assumption, one must then show that

Original: [Equation]
Revised: [Equation]

Original: holds only if q ( y , t ) =
Revised: holds only if q ( y , t ) =

Original: exp[ − H ( y )] / Z .
Revised: exp[ − H ( y )] / Z .

Original: Even if exp[ − H ( y )] / Z is the unique stationary distribution, convergence to that distribution is a question.
Revised: Even if exp[ − H ( y )] / Z is the unique stationary distribution, convergence to that distribution is a question.

Original: See Zhang & Chen (2013) for more details.
Revised: See Zhang & Chen (2013) for more details.

Original: Learning Q ϕ , D ϕ in the MDM ELBO helps push y T to the model prior π θ and avoid issues like those discussed.
Revised: Learning Q ϕ , D ϕ in the MDM ELBO helps push y T to the model prior π θ and avoid issues like those discussed.

Original: B. L INEAR P ROCESSES
Revised: C. L INEAR P ROCESSES

Original: Next, we specialize this general family to linear Itˆo processes to maintain tractable transition distributions.
Revised: Next, we specialize this general family to linear Itˆo processes to maintain tractable transition distributions.

Original: A linear process is one where the drift f
Revised: A linear process is one where the drift f

Original: ( y , t ) and diffusion g ( y , t ) are linear functions of y .
Revised: ( y , t ) and diffusion g ( y , t ) are linear functions of y .

Original: We express the drift function of a non-linear time-varying Itˆo process with stationary distribution proportional to exp[ − H ( y )] as
Revised: We express the drift function of a non-linear time-varying It ˆ o process with stationary distribution proportional to exp[ − H ( y )] as

Original: [Equation]
Revised: [Equation]

Original: Next, linear Itˆo processes have Gaussian stationary distributions (S¨arkk¨a & Solin, 2019)
Revised: Next, linear It ˆ o processes have Gaussian stationary distributions (S ¨ arkk ¨ a & Solin, 2019)

Original: so H ( y ) must be quadratic and ∇ H ( y ) is linear, and neither are constant in y .
Revised: so H ( y ) must be quadratic and ∇ H ( y ) is linear, and neither are constant in y .

Original: Because ∇ H ( y ) is linear, it can be expressed as Sy for some matrix S where S is the inverse of the covariance matrix.
Revised: Because ∇ H ( y ) is linear, it can be expressed as Sy for some matrix S where S is the inverse of the covariance matrix.

Original: Because ∇ H is multiplied by Q , D , this means that Q , D must be free of y .
Revised: Because ∇ H is multiplied by Q , D , this means that Q , D must be free of y .

Original: Recalling that Γ is expressed as a sum of derivatives w.r.t y of Q + D , this means that Γ must satisfy Γ = 0 .
Revised: Recalling that Γ is expressed as a sum of derivatives w.r.t y of Q + D , this means that Γ must satisfy Γ = 0 .

Original: Next, because of the stationary requirement that g ( t ) = (cid:112) 2 D ( y , t ) , we can also conclude by the restriction on D that the diffusion coefficient function must be independent of the state y .
Revised: Next, because of the stationary requirement that g ( t ) = (cid:112) 2 D ( y , t ) , we can also conclude by the restriction on D that the diffusion coefficient function must be independent of the state y .

Original: Our final form for linear time-varying processes with stationary distributions N (0 , S − 1 ) is:
Revised: Our final form for linear time-varying processes with stationary distributions N (0 , S − 1 ) is:

Original: [Equation]
Revised: [Equation]

Original: B. P ARAMETERIZING Q ϕ
Revised: C. P ARAMETERIZING Q ϕ

Original: Suppose b q ( s ) is a positive scalar function defined on the time domain with known integral.
Revised: Suppose b q ( s ) is a positive scalar function defined on the time domain with known integral.

Original: Suppose ˜ Q ϕ is any matrix.
Revised: Suppose ˜ Q ϕ is any matrix.

Original: Then ˜ Q ϕ − ˜ Q ⊤ ϕ is skew-symmetric with ˜ Q ϕ,ij = − ˜ Q
Revised: Then ˜ Q ϕ − ˜ Q ⊤ ϕ is skew-symmetric with ˜ Q ϕ,ij = − ˜ Q

Original: ϕ,ji .
Revised: ϕ,ji .

Original: We can set Q ϕ to
Revised: We can set Q ϕ to

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: 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.

Original: B. P ARAMETERIZING D ϕ
Revised: C.

Revised: P ARAMETERIZING D ϕ

Original: Suppose b d ( s ) is a positive scalar function defined on the time domain with known integral.
Revised: Suppose b d ( s ) is a positive scalar function defined on the time domain with known integral.

Original: Suppose ˜ D ϕ is any matrix.
Revised: Suppose ˜ D ϕ is any matrix.

Original: Then ˜ D
Revised: Then ˜ D

Original: ϕ ˜
Revised: ϕ ˜

Original: D ⊤ ϕ is positive semi-definite and spans all time-independent positive semi-definite matrices.
Revised: D ⊤ ϕ is positive semi-definite and spans all time-independent positive semi-definite matrices.

Original: We can set D ϕ to
Revised: We can set D ϕ to

Original: [Equation]
Revised: [Equation]

Original: To show ˜ D
Revised: To show ˜ D

Original: ˜ D
Revised: ˜ D

Original: ⊤ spans all positive semi-definite matrices: suppose M is positive semi-definite.
Revised: ⊤ spans all positive semi-definite matrices: suppose M is positive semi-definite.

Original: Then it is square.
Revised: Then it is square.

Original: Then it can be eigen-decomposed into M =
Revised: Then it can be eigen-decomposed into M =

Original: VΣV ⊤
Revised: VΣV ⊤

Original: The degrees of freedom in VΣV ⊤ are just R = V √ Σ since VΣV ⊤ = RR ⊤ and the square root is taken element-wise because Σ is diagonal and is real because each Σ ij ≥ 0 , which is true because M is positive semi-definite.
Revised: The degrees of freedom in VΣV ⊤ are just R = V √ Σ since VΣV ⊤ = RR ⊤ and the square root is taken element-wise because Σ is diagonal and is real because each Σ ij ≥ 0 , which is true because M is positive semi-definite.

Original: Take
Revised: Take

Original: D = R .
Revised: D = R .

Original: In our experiments we parameterize D as a diagonal-only matrix.
Revised: In our experiments we parameterize D as a diagonal-only matrix.

Original: B. I NTEGRALS
Revised: C. I NTEGRALS

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

Original: • numerical integration of function with unknown integral.
Revised: • numerical integration of function with unknown integral.

Original: This is expected to have low error given that the function is scalar-in scalar-out.
Revised: This is expected to have low error given that the function is scalar-in scalar-out.

Original: • Directly parameterize the integral and use auto-grad when needing the functions notintegrated.
Revised: • Directly parameterize the integral and use auto-grad when needing the functions notintegrated.

Original: We stick with the known integrals.
Revised: We stick with the known integrals.

Original: 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 ˜ Q ϕ , ˜ D ϕ .
Revised: 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 ˜ Q ϕ , ˜ D ϕ .

Original: B. I NSTANCESVPSDE .
Revised: C. I NSTANCES

Revised: VPSDE .

Original: VPSDE has K = 1 .
Revised: VPSDE has K = 1 .

Original: Consequently, Q , D are K × K .
Revised: Consequently, Q , D are K × K .

Original: The only 1 × 1 skew-symmetric matrix is 0 , so Q = 0 .
Revised: The only 1 × 1 skew-symmetric matrix is 0 , so Q = 0 .

Original: Setting D ( t ) = 12 β ( t ) recovers VPSDE :
Revised: Setting D ( t ) = 12 β ( t ) recovers VPSDE :

Original: [Equation]
Revised: [Equation]

Original: [Equation] ∇ H ( y ) =
Revised: (30) ∇ H ( y ) =

Original: y so H ( y ) =
Revised: y so H ( y ) =

Original: 12 ∥ y ∥ 22 .
Revised: 12 ∥ y ∥ 22 .

Original: The stationary distribution is N (0 , I ) .
Revised: The stationary distribution is N (0 , I ) .

Original: CLD .
Revised: CLD .

Original: The CLD process (eq 5 in Dockhorn et al.
Revised: The CLD process (eq 5 in Dockhorn et al.

Original: (2021)) is defined as
Revised: (2021)) is defined as

Original: [Equation]
Revised: [Equation]

Original: In Q / D parameterization, we have
Revised: In Q / D parameterization, we have

Original: [Equation]
Revised: [Equation]

Original: The stationary distribution of this process is:
Revised: The stationary distribution of this process is:

Original: [Equation]
Revised: [Equation]

Original: ALDA .
Revised: ALDA .

Original: Mou et al.
Revised: Mou et al.

Original: (2019) define a third-order diffusion process for the purpose of gradient-based
Revised: (2019) define a third-order diffusion process for the purpose of gradient-based

Original: MCMC sampling.
Revised: MCMC sampling.

Original: The ALDA diffusion process can be specified as
Revised: The ALDA diffusion process can be specified as

Original: [Equation]
Revised: [Equation]

Original: Note that Q is skew-symmetric and D is positive semi-definite, therefore we have that q t ( u ) → q ϕ, ∞ .
Revised: Note that Q is skew-symmetric and D is positive semi-definite, therefore we have that q t ( u ) → q ϕ, ∞ .

Original: In this case,
Revised: In this case,

Original: [Equation]
Revised: [Equation]

Original: MALDA .
Revised: MALDA .

Original: Similar to ALDA , we specify a diffusion process we term MALDA which we specify as 
Revised: Similar to ALDA , we specify a diffusion process we term MALDA which we specify as [Equation]

Original: [Equation]
Revised: [Equation]

Original: Note that Q is skew-symmetric and D is positive semi-definite.
Revised: Note that Q is skew-symmetric and D is positive semi-definite.

Original: In this case this is
Revised: In this case this is

Original: [Equation]
Revised: [Equation]

Original: C T RANSITIONS FOR LINEAR PROCESSES
Revised: D T RANSITIONS FOR LINEAR PROCESSES

Original: For time variable s and Brownian motion (cid:98) B s driving diffusions of the form
Revised: For time variable s and Brownian motion (cid:98) B s driving diffusions of the form

Original: [Equation]
Revised: [Equation]

Original: (34) when f ϕ ( y s , s ) , g ϕ ( s ) are linear, the transition kernel q ϕ ( y s | y 0 ) is always normal (S ¨ arkk ¨ a & Solin, 2019).
Revised: 34) when f ϕ ( y s , s ) , g ϕ ( s ) are linear, the transition kernel q ϕ ( y s | y 0 ) is always normal (S¨arkk¨a & Solin, 2019).

Original: Therefore, we just find the mean m s | 0 and covariance Σ s | 0 of q ( y s | y 0 ) .
Revised: Therefore, we just find the mean m s | 0 and covariance Σ s | 0 of q ( y s | y 0 ) .

Original: Let f ( y , s ) =
Revised: Let f ( y , s ) =

Original: A ( s ) y .
Revised: A ( s ) y .

Original: The un-conditional time s mean and covariance are solutions to
Revised: The un-conditional time s mean and covariance are solutions to

Original: [Equation]
Revised: [Equation]

Original: By (6.6) in S¨arkk¨a & Solin (2019), for computing conditionals q ( y s | 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 0 respectively.
Revised: By (6.6) in S¨arkk¨a & Solin (2019), for computing conditionals q ( y s | 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 0 respectively.

Original: We take (6.36-6.39) and set m 0 = uand Σ 

Revised: We take (6.36-6.39) and set m 0 = uand Σ

Original: = 0 to condition.
Revised: = 0 to condition.

Original: Let [ A ] s = (cid:82) s 

Revised: Let [ A ] s = (cid:82) s

Original: A ( ν ) dν .
Revised: A ( ν ) dν .

Original: The mean is
Revised: The mean is

Original: [Equation]
Revised: [Equation]

Original:  where exp denotes matrix exponential.
Revised: 36) where exp denotes matrix exponential.

Original: (6.36-6.39) state the covariance q
Revised: (6.36-6.39) state the covariance q

Original: ( y s | y 0 ) as a matrix factorization, for which a derivation is provided below Σ s =
Revised: ( y s | y 0 ) as a matrix factorization, for which a derivation is provided below Σ s =

Original: C s ( H s ) − 1 for C s , H s being the solutions of:
Revised: C s ( H s ) − 1 for C s , H s being the solutions of:

Original: [Equation]
Revised: [Equation]

Original: To condition and get Σ s | 0 from Σ s , we set Σ 

Revised: To condition and get Σ s | 0 from Σ s , we set Σ

Original: = 0 , and initialize C s , H s by C 0 = 0 and H 

Revised: = 0 , and initialize C s , H s by C 0 = 0 and H

Original: = I .
Revised: = I .

Original: [Equation] (cid:18) 
Revised: (cid:18) C s

Original: A g − A ⊤ (cid:19) (cid:35) (cid:18)
Revised: H s

Original: I (cid:19) (cid:124) (cid:123)(cid:122) (cid:125) no integration if A ( ν ) = 
Revised: [Equation] no integration if A ( ν ) = A [Equation]

Original: A ,g ( ν ) = g.
Revised: g ( ν [Equation] = g

Original: Finally
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original:  Σ s | 

Revised: Finally, Σ s |

Original: =
Revised: =

Original: C s ( H s ) − 1 .
Revised: C s ( H s ) − 1 .

Original: C. D ERIVATION OF THE C OVARIANCE MATRIX SOLUTION
Revised: D. D ERIVATION OF THE C OVARIANCE MATRIX SOLUTION

Original: Equation (35) gives an expression for d Σ s /ds .
Revised: Equation (35) gives an expression for d Σ s /ds .

Original: To derive the matrix factorization technique used in eq.
Revised: To derive the matrix factorization technique used in eq.

Original: (37), we use eq.
Revised: (37), we use eq.

Original: (35) and the desired condition Σ s =
Revised: (35) and the desired condition Σ s =

Original: C s H − 1 s to derive expressions for d C s /ds and d H s /ds and suitable intial conditions so that the factorization also starts at the desired Σ 0 .
Revised: C s H − 1 s to derive expressions for d C s /ds and d H s /ds and suitable intial conditions so that the factorization also starts at the desired Σ 0 .

Original: Let Σ s =
Revised: Let Σ s =

Original: C s H − 1 s , then note that C s , H s satisfies
Revised: C s H − 1 s , then note that C s , H s satisfies

Original: [Equation]
Revised: [Equation]

Original: And using the fact that
Revised: And using the fact that

Original: [Equation]
Revised: [Equation]

Original: we get that
Revised: we get that

Original: [Equation]
Revised: [Equation]

Original: Now, we note C s , H s satisfy the following
Revised: Now, we note C s , H s satisfy the following

Original: [Equation]
Revised: [Equation]

Original: which implies that
Revised: which implies that

Original: [Equation]
Revised: [Equation]

Original: with C 0 = Σ 0 and H 

Revised: with C 0 = Σ 0 and H

Original: = I d , as C 0 H − 1 0 = Σ 0 .
Revised: = I d , as C 0 H − 1 0 = Σ 0 .

Original: C. H YBRID S CORE M ATCHING
Revised: D. H YBRID S CORE M ATCHING

Original: Instead of computing q ( y s | y 0 ) , we can apply the hybrid score matching principle (Dockhorn et al., 2021) to reduce variance by compute objectives using q ( y s | x ) instead of q ( y s | y 0 ) , which amounts to integrating out v 0 .
Revised: Instead of computing q ( y s | y 0 ) , we can apply the hybrid score matching principle (Dockhorn et al., 2021) to reduce variance by compute objectives using q ( y s | x ) instead of q ( y s | y 0 ) , which amounts to integrating out v 0 .

Original: To accomplish this, following S¨arkk¨a & Solin (2019), we simply replace ywith [ x, E [ v 0 ]] in the expression for m s | 0 , i.e. replace the conditioning value of v 0 with the mean of its chosen initial distribution:
Revised: To accomplish this, following S¨arkk¨a & Solin (2019), we simply replace ywith [ x, E [ v 0 ]] in the expression for m s | 0 , i.e. replace the conditioning value of v 0 with the mean of its chosen initial distribution:

Original: [Equation]
Revised: [Equation]

Original: For the convariance, instead of using C 0 = Σ 0 = 0 , we use a block matrix to condition on x but not v 0 .
Revised: For the convariance, instead of using C 0 = Σ 0 = 0 , we use a block matrix to condition on x but not v 0 .

Original: We decompose Σ 0 into its blocks Σ 0 ,xx , Σ 0 ,vv , Σ 0 ,xv .
Revised: We decompose Σ 0 into its blocks Σ 0 ,xx , Σ 0 ,vv , Σ 0 ,xv .

Original: As before, to condition on x we set Σ 0 ,xx = 0 .
Revised: As before, to condition on x we set Σ 0 ,xx = 0 .

Original: Because q ( v 0 ) is set to be independent of x , Σ 0 ,xv is also set to 0 .
Revised: Because q ( v 0 ) is set to be independent of x , Σ 0 ,xv is also set to 0 .

Original: Finally, instead of 0 , to marginalize out v 0 , Σ 0 ,vv is set to the covariance of the chosen initial time zero distribution for v 0 .
Revised: Finally, instead of 0 , to marginalize out v 0 , Σ 0 ,vv is set to the covariance of the chosen initial time zero distribution for v 0 .

Original: E.g. if v 0 ,j ∼ N (0 , γ ) for each dimension, then Σ 0 ,vv = N (0 , γI ) .
Revised: E.g. if v 0 ,j ∼ N (0 , γ ) for each dimension, then Σ 0 ,vv = N (0 , γI ) .

Original: 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.
Revised: 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.

Original: C. T RANSITIONS IN S TATIONARY P ARAMETERIZATION
Revised: D. T RANSITIONS IN S TATIONARY P ARAMETERIZATION

Original: In terms of Q , D , the transitions q ( y s | y 0 ) for time s are normal with mean m s | 0 and Σ s | 0 equal to:
Revised: In terms of Q , D , the transitions q ( y s | y 0 ) for time s are normal with mean m s | 0 and Σ s | 0 equal to:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where Σ s | 

Revised: where Σ s |

Original: =
Revised: =

Original: C s ( H s ) − 1 .
Revised: C s ( H s ) − 1 .

Original: For the time invariant case, this simplifies tom s |
Revised: For the time invariant case, this simplifies to

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: D G ENERIC CHANGE OF MEASURE AND J ENSEN ’ S FOR APPROXIMATE MARGINALIZATION
Revised: E G ENERIC CHANGE OF MEASURE AND J ENSEN ’ S FOR APPROXIMATE MARGINALIZATION

Original: Suppose u =
Revised: Suppose u =

Original: [ z , v ] and we have an expression for p ( u =
Revised: [ z , v ] and we have an expression for p ( u =

Original: [ z, v ])
Revised: [ z, v ])

Original: =
Revised: =

Original: p ( z = z, v = v ) .
Revised: p ( z = z, v = v ) .

Original: By marginalization, we can get p ( z = z ) , and we can introduce another distribution q to pick a sampling distri- bution of our choice:
Revised: By marginalization, we can get p ( z = z ) , and we can introduce another distribution q to pick a sampling distribution of our choice:

Original: [Equation]
Revised: [Equation]

Original: We often work with these expressions in log space
Revised: We often work with these expressions in log space

Original: [Equation]
Revised: [Equation]

Original: and need to pull the expectation outside to use
Revised: and need to pull the expectation outside to use

Original: Monte Carlo.
Revised: Monte Carlo.

Original: Jensen’s bound allows this:
Revised: Jensen’s bound allows this:

Original: [Equation]
Revised: [Equation]

Original: The following shows that the bound is tight when q ( v = v | z = z ) =
Revised: The following shows that the bound is tight when q ( v = v | z = z ) =

Original: p ( v = v | z = z ) :
Revised: p ( v = v | z = z ) :

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: E ELBO FOR MDM s
Revised: F ELBO FOR MDM s

Original: [Equation]
Revised: [Equation]

Original: The first inequality holds due to Jensen’s inequality and the second due to an application of Theorem 1 from Huang et al.
Revised: The first inequality holds due to Jensen’s inequality and the second due to an application of Theorem 1 from Huang et al.

Original: (2021) or Theorem 3 from Song et al.
Revised: (2021) or Theorem 3 from Song et al.

Original: (2021) applied to the joint variable u 0 .
Revised: (2021) applied to the joint variable u 0 .

Original: E. ISM TO DSM E.1. L EMMA : EXPECTATION BY PARTS
Revised: F. ISM TO DSM F.1. L EMMA : EXPECTATION BY PARTS

Original: We will need a form of multivariate integration by parts which gives us for some f and some q ( x ) ,
Revised: We will need a form of multivariate integration by parts which gives us for some f and some q ( x ) ,

Original: E q ( x ) [ ∇ x · f ( x )] = − E q ( x ) [ f ( x ) ⊤ ∇ x log q ( x )]

Original: [Equation]
Revised: [Equation]

Original: This equality also follows directly from the Stein operator using the generator method to the
Revised: This equality also follows directly from the Stein operator using the generator method to the

Original: Langevin diffusion (Barbour, 1988).
Revised: Langevin diffusion (Barbour, 1988)

Revised: [Equation]

Original: E.1.
Revised: F.1. DSM E LBO

Original: DSM E LBO

Original: Using the “expectation by parts”
Revised: Using the “expectation by parts”

Original: [Equation]
Revised: [Equation]

Original: we have:
Revised: we have:

Original: [Equation]
Revised: [Equation]

Original: Also we have, for s θ evaluated at ( u t , t ) , by completing the square, −
Revised: Also we have, for s θ evaluated at ( u t , t ) [Equation] by completing the square, −

Original: [Equation]
Revised: [Equation]

Original: The two together give us:
Revised: The two together give us:

Original: [Equation]
Revised: [Equation]

Original: dt (51)

Original: E. N OISE PREDICTION
Revised: F. N OISE PREDICTION

Original: We have that for normal N ( y s ; m s | 0 , Σ s | 0 ) , we can sample y s with normal noise ϵ ∼ N (0 , I ) and y s = m s | 

Revised: We have that for normal N ( y s ; m s | 0 , Σ s | 0 ) , we can sample y s with normal noise ϵ ∼ N (0 , I ) and y s = m s |

Original: + L ϵ where L is the cholesky decomposition of Σ s | 

Revised: + L ϵ where L is the cholesky decomposition of Σ s |

Original: Then, the score is
Revised: Then, the score is

Original: [Equation]
Revised: [Equation]

Original: Parameterize s θ ( y s , s ) as s θ ( y s , s )
Revised: Parameterize s θ ( y s , s ) as s θ ( y s , s )

Original: =
Revised: =

Original: − L ⊤ , − 1 ϵ θ ( y , s ) .
Revised: − L ⊤ , − 1 ϵ θ ( y , s ) .

Original: This gives
Revised: This gives

Original: [Equation]
Revised: [Equation]

Original: 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 ):
Revised: 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 ):

Original: [Equation]
Revised: [Equation]

Original: F ELBOS IN STATIONARY PARAMETERIZATION
Revised: G ELBOS IN STATIONARY PARAMETERIZATION

Original: We use the stationary parmeterization described in appendix B. We now specialize the ELBO to the linear stationary parameterization
Revised: We use the stationary parmeterization described in appendix C. We now specialize the ELBO to the linear stationary parameterization. Recall

Original: [Equation]
Revised: [Equation]

Original: Recall f ϕ ( y , s ) = − [ Q ϕ ( s ) + D ϕ ( s )]
Revised: D ϕ ( s )]

Original: y .
Revised: y .

Original: Recall g ϕ ( s ) =
Revised: Recall g ϕ ( s ) =

Original: (cid:112) 2 D ϕ
Revised: (cid:112) 2 D ϕ

Original: ( s )
Revised: ( s )

Original: We have g 2 ϕ
Revised: We have g 2 ϕ

Original: ( s ) =
Revised: ( s ) =

Original: 2 D ϕ ( s ) .
Revised: 2 D ϕ ( s ) .

Original: We can write the MDM ISM ELBO as
Revised: We can write the MDM ISM ELBO as

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: For the DSM form,
Revised: For the DSM form

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: G A LGORITHMS G. G ENERIC T RANSITION K ERNEL
Revised: H A LGORITHMS H. G ENERIC T RANSITION K ERNEL

Original: Input: data x . time s .
Revised: Input: data x . time s .

Original: A , g . compute:
Revised: A , g . compute:

Original: A ( s ) and g ( s ) compute:
Revised: A ( s ) and g ( s ) compute:

Original: M s = (cid:82) s 0 A ( t ) dt (integrated drift) compute:
Revised: M s = (cid:82) s 0 A ( t ) dt (integrated drift) compute:

Original: N s = (cid:82) s 0 g 2 ( t )
Revised: N s = (cid:82) s 0 g 2 ( t )

Original: dt (integrated diffusions squared) compute: γ s | 0 = exp (cid:16) M s (cid:17) (mean coefficient) set: y 

Revised: dt (integrated diffusions squared) compute: γ s | 0 = exp (cid:16) M s (cid:17) (mean coefficient) set: y

Original: =
Revised: =

Original: [ x, 0 1 , . .
Revised: [ x, 0 1 , . .

Original: , 0 K − 1 ] , Σ 0 ,zz = 0 , and Σ 0 ,zv , Σ 0 ,vv to chosen initial distribution compute: m s | 

Revised: , 0 K − 1 ] , Σ 0 ,zz = 0 , and Σ 0 ,zv , Σ 0 ,vv to chosen initial distribution compute: m s |

Original: = γ s | 

Revised: = γ s |

Original: y 

Revised: y

Original: ( mean ) compute: (cid:18) C s
Revised: ( mean ) compute: (cid:18) C s

Original: H s (cid:19) = exp (cid:34) (cid:18) M s
Revised: H s (cid:19) = exp (cid:34) (cid:18) M s

Original: N s − M ⊤ s (cid:19) (cid:35) (cid:18) Σ
Revised: N s − M ⊤ s (cid:19) (cid:35) (cid:18) Σ

Original: I (cid:19) (ingredients for cov.)
Revised: I (cid:19) (ingredients for cov.)

Original: (55) compute: Σ s | 

Revised: (55) compute: Σ s |

Original: =
Revised: =

Original: C s ( H s ) −( cov. )
Revised: C s ( H s ) −( cov. )

Original: Output: N ( m s | 0 , Σ s | 0 )
Revised: Output: N ( m s | 0 , Σ s | 0 )

Original: G.
Revised: H.

Original: T RANSITIONS WITH Q, D
Revised: T RANSITIONS WITH Q, D

Original: Current param matrices ˜ Q ϕ , ˜ D
Revised: Current param matrices ˜ Q ϕ , ˜ D

Original: ϕ and along with fixed time-in scalar-out functions
Revised: ϕ and along with fixed time-in scalar-out functions

Original: b q ( s ) , b d ( s ) and their known integrals B q ( s ) , B d ( s ) .
Revised: b q ( s ) , b d ( s ) and their known integrals B q ( s ) , B d ( s ) .

Original: q γ ( v 0 |
Revised: q γ ( v 0 |

Original: z
Revised: z

Original: =
Revised: =

Original: x ) taken to be parameterless so that v 0 ∼ N (0 , I ) .
Revised: x ) taken to be parameterless so that v 0 ∼ N (0 , I ) .

Original: Model params are s θ fixed π θ .
Revised: Model params are s θ fixed π θ .

Original: Algorithm 3 Get Q , D and their integrated terms M , N
Revised: Algorithm 3 Get Q , D and their integrated terms M , N

Original: Input: time s and current params ϕ compute: [ b q ] s = (cid:82) s 0 b q ( ν ) dν using known integral B q ( s ) − B q (0) compute:
Revised: Input: time s and current params ϕ compute: [ b q ] s = (cid:82) s 0 b q ( ν ) dν using known integral B q ( s ) − B q (0) compute:

Original: [ b d ] s = (cid:82) s 0 b d ( ν )
Revised: [ b d ] s = (cid:82) s 0 b d ( ν )

Original: dν using known integral B d ( s ) − B d (0) .
Revised: dν using known integral B d ( s ) − B d (0) .

Original: compute: [ Q ϕ ] s =
Revised: compute: [ Q ϕ ] s =

Original: [ b q ] s · (cid:104)
Revised: [ b q ] s · (cid:104)

Original: ˜ Q ϕ − ˜ Q ⊤ ϕ (cid:105) for current params ˜ Q ϕ .
Revised: ˜ Q ϕ − ˜ Q ⊤ ϕ (cid:105) for current params ˜ Q ϕ .

Original: compute:
Revised: compute:

Original: [ D ϕ ] s =
Revised: [ D ϕ ] s =

Original: [ b d ] s · (cid:104)
Revised: [ b d ] s · (cid:104)

Original: ˜ D
Revised: ˜ D

Original: ϕ ˜ D ⊤ ϕ (cid:105) for current params ˜ D ϕ .
Revised: ϕ ˜ D ⊤ ϕ (cid:105) for current params ˜ D ϕ .

Original: compute: M s = − ([ Q ϕ ] s + [ D ϕ ] s )
Revised: compute: M s = − ([ Q ϕ ] s + [ D ϕ ] s )

Original: ( M just a variable name) compute:
Revised: ( M just a variable name) compute:

Original: N s =
Revised: N s =

Original: [2 D ϕ ] s = 2 · [ D ϕ ] s ( N just a variable name) compute:
Revised: [2 D ϕ ] s = 2 · [ D ϕ ] s ( N just a variable name) compute:

Original: Q s = b q ( s ) · (cid:104) ˜
Revised: Q s = b q ( s ) · (cid:104) ˜

Original: Q ϕ − ˜ Q ⊤ ϕ (cid:105) (not integrated) compute:
Revised: Q ϕ − ˜ Q ⊤ ϕ (cid:105) (not integrated) compute:

Original: D s = b d ( s ) · (cid:104) ˜ D ϕ ˜ D ⊤ ϕ (cid:105) (not integrated) compute:
Revised: D s = b d ( s ) · (cid:104) ˜ D ϕ ˜ D ⊤ ϕ (cid:105) (not integrated) compute:

Original: A s = − [ Q s + D s ] (drift coef.)
Revised: A s = − [ Q s + D s ] (drift coef.)

Original: compute: g 2 s = 2 D s (diffusion coef.
Revised: compute: g 2 s = 2 D s (diffusion coef.

Original: squared)
Revised: squared)

Original: Output: A s , g 2 s , M s , N s
Revised: Output: A s , g 2 s , M s , N s

Original: Algorithm
Revised: Algorithm

Original: Input: Sample y 

Revised: Input: Sample y

Original: = ( x, v ) and time s .
Revised: = ( x, v ) and time s .

Original: Current params ϕ set: A s , g 2 s , M s , N s ← algorithmcompute: m s | 0 =
Revised: Current params ϕ set: A s , g 2 s , M s , N s ← algorithmcompute: m s | 0 =

Original: exp (cid:16) M s (cid:17) y 0 (transition mean) compute: ingredients for transition cov.
Revised: exp (cid:16) M s (cid:17) y 0 (transition mean) compute: ingredients for transition cov.

Original: matrix:
Revised: matrix:

Original: [Equation]
Revised: [Equation]

Original: (cid:34) (cid:18)
Revised: (cid:34) (cid:18)

Original: [Equation]
Revised: [Equation]

Original: (cid:19) (cid:35) (cid:18)
Revised: (cid:19) (cid:35) (cid:18)

Original: 0I (cid:19) (56) compute: Σ s | 

Revised: 0I (cid:19) (56) compute: Σ s |

Original: =
Revised: =

Original: C s ( H s ) − 1 (transition cov).
Revised: C s ( H s ) − 1 (transition cov).

Original: instantiate: q ϕ,s, ( x,v ) =
Revised: instantiate: q ϕ,s, ( x,v ) =

Original: q ϕ ( y s | y 0 ) =
Revised: q ϕ ( y s | y 0 ) =

Original: N ( m s | 0 , Σ s | 0 ) .
Revised: N ( m s | 0 , Σ s | 0 ) .

Original: Output: q ϕ,s, ( x,v ) , A s , g 2 s
Revised: Output: q ϕ,s, ( x,v ) , A s , g 2 s

Original: Algorithm input: Data point x and current params θ, ϕ, γ draw: an aux. sample v ∼ q γ ( v | x ) draw: a sample s ∼ Unif (0 , T ) set: y 

Revised: Algorithm input: Data point x and current params θ, ϕ, γ draw: an aux. sample v ∼ q γ ( v | x ) draw: a sample s ∼ Unif (0 , T ) set: y

Original: = ( x, v ) set: q ϕ,s, y 0 , A s , g 2 s ← algorithm 4 called on y 0 , s, ϕ draw: y s ∼ q ϕ,s, ycompute: ℓ s with dsm ( s ) (algorithm 6) or ism ( s ) (algorithm 7) on y s , θ, A s , g 2 s , q ϕ,s, yset: q ϕ,T, y 0 , , ← algorithm 4 called on y 0 , T, ϕ draw: y T ∼ q ϕ,T, y 0 output: ℓ s + log π θ ( y T ) − log q γ ( v )
Revised: = ( x, v ) set: q ϕ,s, y 0 , A s , g 2 s ← algorithm 4 called on y 0 , s, ϕ draw: y s ∼ q ϕ,s, ycompute: ℓ s with dsm ( s ) (algorithm 6) or ism ( s ) (algorithm 7) on y s , θ, A s , g 2 s , q ϕ,s, yset: q ϕ,T, y 0 , , ← algorithm 4 called on y 0 , T, ϕ draw: y T ∼ q ϕ,T, y 0 output: ℓ s + log π θ ( y T ) − log q γ ( v )

Original: Algorithm input: y s , θ , A s , g 2 s , q ϕ,s, y 0 .
Revised: Algorithm input: y s , θ , A s , g 2 s , q ϕ,s, y 0 .

Original: compute: fwd-score = ∇ y s log q ϕ ( y s | y 0 ) compute: model-score =
Revised: compute: fwd-score = ∇ y s log q ϕ ( y s | y 0 ) compute: model-score =

Original: s θ
Revised: s θ

Original: ( y s , s ) compute: fwd-score-term = 12 ( fwd-score )
Revised: ( y s , s ) compute: fwd-score-term = 12 ( fwd-score )

Original: ⊤ g 2 s ( fwd-score ) compute: score-diff = model-score − fwd-score compute: diff-term = − 12 score-diff ⊤ g 2 s score-diff compute: div-f = ∇ y s · A s y s output: dsm ( s ) =
Revised: ⊤ g 2 s ( fwd-score ) compute: score-diff = model-score − fwd-score compute: diff-term = − 12 score-diff ⊤ g 2 s score-diff compute: div-f = ∇ y s · A s y s output: dsm ( s ) =

Original: fwd-score-term
Revised: fwd-score-term

Original: + diff-term + div-f
Revised: + diff-term + div-f

Original: Algorithm input: y s , θ , A s , g 2 s , q ϕ,s, y 0 .
Revised: Algorithm input: y s , θ , A s , g 2 s , q ϕ,s, y 0 .

Original: compute: model-score =
Revised: compute: model-score =

Original: s θ
Revised: s θ

Original: ( y s , s ) compute: score-term = − 12 model-score ⊤ g 2 s model-score compute: div-gs = ∇ y s · g 2 s s θ ( y s , s ) compute: div-f = ∇ y s · A s y s compute: div-term = − div-gs + div-f output: ism ( s ) =
Revised: ( y s , s ) compute: score-term = − 12 model-score ⊤ g 2 s model-score compute: div-gs = ∇ y s · g 2 s s θ ( y s , s ) compute: div-f = ∇ y s · A s y s compute: div-term = − div-gs + div-f output: ism ( s ) =

Original: score-term + div-term
Revised: score-term + div-term

Original: H V ALID ELBO WITH TRUNCATION
Revised: I V ALID ELBO WITH TRUNCATION

Original: The integrand in the ELBO and its gradients is not bounded at time 0 .
Revised: The integrand in the ELBO and its gradients is not bounded at time 0 .

Original: Therefore, following SohlDickstein et al.
Revised: Therefore, following SohlDickstein et al.

Original: and Song et al.
Revised: and Song et al.

Original: the integrand in eq.
Revised: the integrand in eq.

Original: (7) is integrated from [ ϵ, T ] , rather than [0 , T ] .
Revised: (7) is integrated from [ ϵ, T ] , rather than [0 , T ] .

Original: However, that integral is not a valid lower bound on log p θ ( x ) .
Revised: However, that integral is not a valid lower bound on log p θ ( x ) .

Original: Instead, it can be viewed as a proper lower bound on the prior for a latent variable y ϵ .
Revised: Instead, it can be viewed as a proper lower bound on the prior for a latent variable y ϵ .

Original: 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.
Revised: 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.

Original: To provide a valid lower bound for multivariate diffusions, we extend theorem 6 in Song et al.
Revised: To provide a valid lower bound for multivariate diffusions, we extend theorem 6 in Song et al.

Original: from univariate to multivariate diffusions.
Revised: from univariate to multivariate diffusions.

Original: Theorem 3.
Revised: Theorem 3.

Original: For transition kernel q ϕ ( y s | y 0 ) , we can compute upper bound the model likelihood at time 0 as follows, for any ϵ >
Revised: For transition kernel q ϕ ( y s | y 0 ) , we can compute upper bound the model likelihood at time 0 as follows, for any ϵ >

Original: [Equation]
Revised: [Equation]

Original: where L mdm ( y ϵ , ϵ ) is defined as
Revised: where L mdm ( y ϵ , ϵ ) is defined as

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: For transition kernel q ϕ ( y s | y 0 ) , we can compute upper bound the model likelihood at time 0 following an application of the variational bound
Revised: For transition kernel q ϕ ( y s | y 0 ) , we can compute upper bound the model likelihood at time 0 following an application of the variational bound

Original: [Equation]
Revised: [Equation]

Original: A lower bound for log p
Revised: A lower bound for log p

Original: θ ( y ϵ ) can be derived in a similar manner to eq.
Revised: θ ( y ϵ ) can be derived in a similar manner to eq.

Original: (7), such that
Revised: (7), such that

Original: [Equation]
Revised: [Equation]

Original: The choice of p θ ( y 0 | y ϵ ) is arbitrary, however following Sohl-Dickstein et al.
Revised: The choice of p θ ( y 0 | y ϵ ) is arbitrary, however following Sohl-Dickstein et al.

Original: ; Song et al.
Revised: ; Song et al.

Original: (2021) we let p θ ( y 0 | y ϵ ) be Gaussian with mean µ p θ ,ϵ and covariance Σ
Revised: (2021) we let p θ ( y 0 | y ϵ ) be Gaussian with mean µ p θ ,ϵ and covariance Σ

Original: p θ ,ϵ .
Revised: p θ ,ϵ .

Original: Suppose q ϕ
Revised: Suppose q ϕ

Original: ( y ϵ | y 0 )
Revised: ( y ϵ | y 0 )

Original: =
Revised: =

Original: N ( y ϵ |
Revised: N ( y ϵ |

Original: Ay 0 , Σ) , then we select the following mean µ p θ ,ϵ and covariance Σ
Revised: Ay 0 , Σ) , then we select the following mean µ p θ ,ϵ and covariance Σ

Original: p θ ,ϵ for p θ
Revised: p θ ,ϵ for p θ

Original: ( y 0 | y ϵ )
Revised: ( y 0 | y ϵ )

Original: [Equation]
Revised: [Equation]

Original: where µ p θ ,ϵ , Σ p θ ,ϵ are derived using Tweedie’s formula (Efron, 2011) by setting µ ϵ =
Revised: where µ p θ ,ϵ , Σ p θ ,ϵ are derived using Tweedie’s formula (Efron, 2011) by setting µ ϵ =

Original: E
Revised: E

Original: [ y 0 | y ϵ ] and Σ ϵ =
Revised: [ y 0 | y ϵ ] and Σ ϵ =

Original: Var ( y 0 | y ϵ ) .
Revised: Var ( y 0 | y ϵ ) .

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

Original: H.
Revised: I.

Original: L IKELIHOOD DERIVATION
Revised: L IKELIHOOD DERIVATION

Original: Suppose y 0 ∼ q 0 ( y 0 ) and y ϵ ∼ N
Revised: Suppose y 0 ∼ q 0 ( y 0 ) and y ϵ ∼ N

Original: ( y ϵ |
Revised: ( y ϵ |

Original: A y 0 , Σ) .
Revised: A y 0 , Σ) .

Original: Here, A, Σ are the mean coefficient and covariance derived from the transition kernel at time ϵ .
Revised: Here, A, Σ are the mean coefficient and covariance derived from the transition kernel at time ϵ .

Original: We use Tweedie’s formula to get the mean and covariance of y 0 given y ϵ under q .
Revised: We use Tweedie’s formula to get the mean and covariance of y 0 given y ϵ under q .

Original: This mean and covariance feature the true score ∇ y ϵ log q ( y ϵ ) .
Revised: This mean and covariance feature the true score ∇ y ϵ log q ( y ϵ ) .

Original: We replace the score with the score model s θ and then set p θ ( y 0 | y ϵ ) to have the resulting approximate mean and covariance.
Revised: We replace the score with the score model s θ and then set p θ ( y 0 | y ϵ ) to have the resulting approximate mean and covariance.

Original: We make this choice because the optimal p θ ( y 0 | y ϵ ) equals the true q ( y 0 | y ϵ ) as discussed throughout the work.
Revised: We make this choice because the optimal p θ ( y 0 | y ϵ ) equals the true q ( y 0 | y ϵ ) as discussed throughout the work.

Original: Here y 

Revised: Here y

Original: =
Revised: =

Original: [ x 0 , v 0 ] where x 0 ∼ q data .
Revised: [ x 0 , v 0 ] where x 0 ∼ q data .

Original: Let η be the natural parameter for the multivariate Gaussian likelihood N ( y ϵ | A y 0 , Σ) .
Revised: Let η be the natural parameter for the multivariate Gaussian likelihood N ( y ϵ | A y 0 , Σ) .

Original: Then,
Revised: Then,

Original: Tweedie’s formula (Efron, 2011) states that:
Revised: Tweedie’s formula (Efron, 2011) states that:

Original: [Equation]
Revised: [Equation]

Original: • l [Equation] y ϵ ) 
Revised: • l ( y ϵ ) =

Original: = log q
Revised: log q

Original: [Equation]
Revised: [Equation]

Original: y ϵ ) • s θ ( y ϵ , ϵ ) is taken to be the true score ∇ y ϵ log q
Revised: y ϵ ) • s θ ( y ϵ , ϵ ) is taken to be the true score ∇ y ϵ log q

Original: ( y ϵ )
Revised: ( y ϵ )

Original: so that ∇ y ϵ l ( y ϵ )
Revised: so that ∇ y ϵ l ( y ϵ )

Original: =
Revised: =

Original: s θ
Revised: s θ

Original: ( y ϵ , ϵ ) • l 0 is the log of the base distribution defined in the exponential family parameterization.
Revised: ( y ϵ , ϵ ) • l 0 is the log of the base distribution defined in the exponential family parameterization.

Original: The base distribution is a multivariate Gaussian with mean 0 and covariance Σ , therefore ∇ y ϵ l 0 ( y ϵ )
Revised: The base distribution is a multivariate Gaussian with mean 0 and covariance Σ , therefore ∇ y ϵ l 0 ( y ϵ )

Original: = − Σ − 1 y ϵ 
Revised: = − Σ − 1 y ϵ ,

Original: [Equation]
Revised: [Equation]

Original: However, Tweedie’s formula is not directly applicable since our y ϵ is not directly normal with mean y 0 .
Revised: However, Tweedie’s formula is not directly applicable since our y ϵ is not directly normal with mean y 0 .

Original: Instead, to derive the conditional mean of y 0 given y ϵ , we use the relation η = Σ − 1 Ay 0 and the linearity of conditional expectation to get
Revised: Instead, to derive the conditional mean of y 0 given y ϵ , we use the relation η = Σ − 1 Ay 0 and the linearity of conditional expectation to get

Original: [Equation]
Revised: [Equation]

Original: For the variance, we use the following relation y ϵ
Revised: For the variance, we use the following relation y ϵ

Original: =
Revised: =

Original: A y 0 + √ Σ ϵ , which implies that
Revised: A y 0 + √ Σ ϵ , which implies that

Original: [Equation]
Revised: [Equation]

Original: Therefore, for the model posterior distribution p θ ( y 0 | y ϵ )
Revised: Therefore, for the model posterior distribution p θ ( y 0 | y ϵ )

Original: we choose a Normal with mean and covariance
Revised: we choose a Normal with mean and covariance

Original: [Equation]
Revised: [Equation]