10.3 Diffusion Models

One desideratum that has emerged from our investigation of directed generative models is for the distribution of latent variables to be essentially structureless. For one thing, this makes it easy to sample the latent variables, and therefore to generate data (since directed models can be sampled with a single “ancestral pass” from root to leaves). For another, it accords with our high-level conceptual goal of explaining observations in terms of a simpler set of independent causes.

Seen from the perspective of the recognition model, this desideratum presents a paradoxical character: it seems that the corresponding goal for the recognition model is to destroy structure. Clearly, it cannot do so in an irreversible way (e.g., by multiplying the observations by zero and adding noise), or there will be no mutual information between observed and latent variables. However, there exists a class of models, known in statistical physics as diffusion processes, which gradually delete structure from distributions but which are nevertheless reversible [69].

A diffusion process can be described by a (very long) Markov chain that (very mildly) corrupts the data at every step. Now, for sufficiently mild data corruption, the process is reversible [69]. Still, we cannot simply apply Bayes’ rule, since it requires the prior distribution over the original, uncorrupted data—i.e., the data distribution, precisely what we want to sample from. On the other hand, it turns out that each reverse-diffusion step must take the same distributional form as the forward-diffusion step [69]. We still don’t know how to convert noisy observations into the parameters of this distribution, but perhaps this mapping can be learned.

In particular, suppose we pair a recognition model describing such a diffusion process with a generative model that is a Markov chain of the same length, and with the same form for its conditional distributions, but pointing in the other direction. Then training the generative model to assign high probability to the observation—or more precisely, lower the joint relative entropy—while making inferences under the recognition model will effectively oblige the generative model to learn to denoise the data at every step. That is, it will become a model of the reverse-diffusion process.

Notice that for this process to be truly reversible, the dimensionality of the data must stay constant across all steps of the Markov chain (including the observations themselves). Also note that, as lately described, the recognition distribution is fixed and requires no learnable parameters. We revisit this idea below.

The generative and recognition models.

The model can be expressed most elegantly if we use ${\bm{\check{X}}}_{0}$ for the observed data ${\bm{Y}}{}$ , and likewise ${\bm{\hat{X}}}_{0}$ for their counterparts in the generative model, ${\bm{\hat{Y}}}$ —so in this section we do. Then the diffusion generative model can be written simply as

equation (10.28) (10.28)

{\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{1}},\ldots,{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{L}},{% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }_{0}};\bm{\theta}}\right)}={\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]% {.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{L}}}% \right)}\prod_{l=1}^{L}{\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l-1}}% \middle|{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{% .75,0,.25}\bm{x}_{l}};\bm{\theta}}\right)};

i.e., a Markov chain. Notice that the generative prior is not parameterized. This accords with the intuition lately discussed that the model should convert structureless noise into the highly structured distribution of interest. The recognition model likewise simplifies according to the independence statements for a Markov chain:

equation (10.29) (10.29)

{\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{1}},\ldots,{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{L}}% \middle|{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{% .75,0,.25}\bm{x}_{0}}}\right)}=\prod_{l=1}^{L}{\check{p}\mathopen{}\mathclose{% {}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{% .75,0,.25}\bm{x}_{l}}\middle|{\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l-1}}}\right)}.

We have omitted the usual $\bm{\phi}$ from this model because by assumption it has no learnable parameters.

It remains, of course, to specify distributional forms for the factors in Eqs. 10.28 and 10.29. We explore one choice below, but first derive the loss for the more general case.

The joint relative entropy.

The joint relative entropy, we recall once again, is the difference between the joint cross entropy (between hybrid distribution and generative model) and the entropy of the hybrid distribution, ${\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}}\middle|{\color[rgb]{.75,0,.25}% \definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}}{}}\right)}{p% \mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}}{}}\right)}$ . However, in the diffusion model, the recognition model is parameterless, so the entire entropy term is a constant as far as our optimization goes. That is, our optimization is essentially “all M step” (improving the generative model), and therefore can be carried out on the joint cross entropy. This would be a mistake if the recognition model were (as usual) to be intepreted as merely an approximation to the posterior under the generative model. In the diffusion model, in contrast, the recognition model is interpreted as ground truth, for which the generative model provides the approximation. The overall loss can be reduced by improving the generative fit either to the data or to the recognition model, but in this case both are desirable per se.

Given the conditional independencies of a Markov chain, then, the joint relative entropy reduces to a sum of conditional cross entropies (plus constant terms):

equation (10.30) (10.30)

\begin{split}\mathcal{L}_{\text{JRE}}(\bm{\theta})&{}\stackrel{{\scriptstyle% \text{+}}}{{=}}\mathbb{E}_{{\bm{\check{X}}}_{1},\ldots,{\bm{\check{X}}}_{L}{},% {\bm{\check{X}}}_{0}{}}{\mathopen{}\mathclose{{}\left[\log{\check{p}\mathopen{% }\mathclose{{}\left({\bm{\check{X}}}_{1},\ldots,{\bm{\check{X}}}_{L}\middle|{% \bm{\check{X}}}_{0}}\right)}-\log{\hat{p}\mathopen{}\mathclose{{}\left({\bm{% \check{X}}}_{1},\ldots,{\bm{\check{X}}}_{L},{\bm{\check{X}}}_{0};\bm{\theta}}% \right)}}\right]}\\ &{}\stackrel{{\scriptstyle\text{+}}}{{=}}\mathbb{E}_{{\bm{\check{X}}}_{1},% \ldots,{\bm{\check{X}}}_{L}{},{\bm{\check{X}}}_{0}{}}{\mathopen{}\mathclose{{}% \left[-\log{\hat{p}\mathopen{}\mathclose{{}\left({\bm{\check{X}}}_{L}}\right)}% -\sum_{l=1}^{L}\log{\hat{p}\mathopen{}\mathclose{{}\left({\bm{\check{X}}}_{l-1% }\middle|{\bm{\check{X}}}_{l};\bm{\theta}}\right)}}\right]}\\ &{}\stackrel{{\scriptstyle\text{+}}}{{=}}-\sum_{l=1}^{L}\int_{\bm{\check{x}}_{% 0}{}}\int_{\bm{\check{x}}_{l-1}{}}\int_{\bm{\check{x}}_{l}{}}\check{p}% \mathopen{}\mathclose{{}\left(\bm{\check{x}}_{l},\bm{\check{x}}_{l-1}\middle|% \bm{\check{x}}_{0}}\right)\log{\hat{p}\mathopen{}\mathclose{{}\left(\bm{\check% {x}}_{l-1}\middle|\bm{\check{x}}_{l};\bm{\theta}}\right)}\mathop{}\!\mathrm{d}% {\bm{\check{x}}_{l}{}}\mathop{}\!\mathrm{d}{\bm{\check{x}}_{l-1}{}}{p\mathopen% {}\mathclose{{}\left(\bm{\check{x}}_{0}}\right)}\mathop{}\!\mathrm{d}{\bm{% \check{x}}_{0}{}}.\end{split}

(The omitted additive constants are not all.) Now we need to choose generative and recognition models that make the integrals (or sums) in these cross entropies tractable.

10.3.1 Gaussian diffusion models

Probably the most intuitive diffusion process is based on Gaussian noise; for example,

{\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l-1}}}% \right)}\mathrel{\vbox{\hbox{.}\hbox{.} }}=\mathcal{N}\mathopen{}\mathclose{{}\left(\beta_{l}{\color[rgb]{.75,0,.25}% \definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l-1}},\>\gamma_{l}^% {2}\mathbf{I}}\right)

for some parameters $\beta_{l}$ and $\gamma_{l}$ . Essentially the process scales (down) the data and adds isotropic noise. However, we defer specifying these parameters for the moment, and turn directly to the generative model. Suffice to say, if $\beta_{l}$ is sufficiently close to 1 and $\gamma_{l}$ is sufficiently small, then the generative transitions are likewise Gaussian (see discussion above). Furthermore, after many diffusion steps, the distribution of the state will be Gaussian and isotropic. With the appropriate selection of the recognition parameters, we can force this distribution to have zero mean and unit variance. Therefore we define the generative model to be

equation (10.31) (10.31)

\displaystyle{\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}% \definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{L}}}\right)}% \mathrel{\vbox{\hbox{.}\hbox{.} }}=\mathcal{N}\mathopen{}\mathclose{{}\left(\bm{0},\>\mathbf{I}}\right),

\displaystyle{\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}% \definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l-1}}\middle|{% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }_{l}};\bm{\theta}}\right)}\mathrel{\vbox{\hbox{.}\hbox{.} }}=\mathcal{N}\mathopen{}\mathclose{{}\left(\bm{\mu}({\color[rgb]{.75,0,.25}% \definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}},l,\bm{\theta}),% \>\sigma^{2}({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{% .75,0,.25}\bm{x}_{l}},l,\bm{\theta})\mathbf{I}}\right).

The mean and variance of this denoising distribution can depend on the corrupted sample ( ${\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{% x}_{l}}$ ) in a complicated way, so in general we can let $\bm{\mu}$ and $\sigma^{2}$ be neural networks. Nevertheless, for simplicity in the derivation, and because learning variances is significantly more challenging than learning means, let us further replace the variance function $\sigma^{2}({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{% .75,0,.25}\bm{x}_{l}},l,\bm{\theta})$ with a set of $L$ fixed (i.e., not learned), data-independent parameters, $\sigma^{2}_{l}$ . (To save space, we also write the mean function, $\bm{\mu}({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{% .75,0,.25}\bm{x}_{l}},l,\bm{\theta})=\bm{\mu}_{l}$ , but it certainly does depend on the data and generative-model parameters.) Then the joint relative entropy (Eq. 10.30) can be expressed as

equation (10.32) (10.32)

\begin{split}\mathcal{L}_{\text{JRE}}(\bm{\theta})&{}\stackrel{{\scriptstyle% \text{+}}}{{=}}\sum_{l=1}^{L}\frac{1}{2\sigma^{2}_{l}}\int_{\bm{\check{x}}_{0}% {}}\int_{\bm{\check{x}}_{l-1}{}}\int_{\bm{\check{x}}_{l}{}}\check{p}\mathopen{% }\mathclose{{}\left(\bm{\check{x}}_{l},\bm{\check{x}}_{l-1}\middle|\bm{\check{% x}}_{0}}\right){p\mathopen{}\mathclose{{}\left(\bm{\check{x}}_{0}}\right)}% \mathopen{}\mathclose{{}\left\lVert\bm{\check{x}}_{l-1}-\bm{\mu}_{l}}\right% \rVert^{2}\mathop{}\!\mathrm{d}{\bm{\check{x}}_{l}{}}\mathop{}\!\mathrm{d}{\bm% {\check{x}}_{l-1}{}}\mathop{}\!\mathrm{d}{\bm{\check{x}}_{0}{}}.\end{split}

The interpretation is now clear. Minimizing the joint relative entropy obliges the generative model to learn how to “undo” one step of corruption with Gaussian noise, for all steps $l\in[1,L]$ . At the model’s disposal is the arbitrarily powerful function (neural network) $\bm{\mu}({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{% .75,0,.25}\bm{x}_{l}},l,\bm{\theta})$ . Since the amount of corruption can vary with $l$ , optimizing Eq. 10.32 obliges $\bm{\mu}$ to be able to remove noise of possibly different sizes.

As for the implementation, we see that the integrals in Eq. 10.32 can be approximated with samples: first a draw ( $\bm{\check{x}}_{0}$ ) from the data distribution, followed by draws ( $\bm{\check{x}}_{1},\ldots,\bm{\check{x}}_{{L}}$ ) down the length of the recognition model. Notice, however, that under this scheme, each summand would be estimated with samples from three random variables, ${\bm{\check{X}}}_{0},{\bm{\check{X}}}_{l-1},{\bm{\check{X}}}_{l}$ . We can reduce this by one, and thereby reduce the variance of our Monte Carlo estimator, by exploiting some properties of Gaussian noise. In particular, we will reverse the order of expansion in applying the chain rule of probability to the recognition model:

equation (10.33) (10.33)

{\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}},{\color[rgb]{.75,0,.25}% \definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l-1}}\middle|{% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }_{0}}{}}\right)}={\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l-1}}% \middle|{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{% .75,0,.25}\bm{x}_{l}},{\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}}}\right)}{\check{p}\mathopen{}% \mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{% rgb}{.75,0,.25}\bm{x}_{l}}\middle|{\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}}}\right)}.

Then we will carry out the expectation under ${\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l-1}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}},{% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }_{0}}}\right)}$ in closed form. In preparation, we now turn to ${\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l-1}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}},{% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }_{0}}}\right)}$ and ${\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}}}% \right)}$ .

The recognition marginals.

A very useful upshot of defining the recognition model to consist only of scaling and the addition of Gaussian noise is that the distribution of any random variable under this model is Gaussian (conditioned, that is, on $\bm{\check{x}}_{0}$ ). That includes ${\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}}}% \right)}$ , which we might call the “recognition marginal.” For reasons that will soon become clear, we make these marginals the starting point of our definition of the recognition model [38], and then work backwards to the transition probabilities:

equation (10.34) (10.34)

{\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}}}% \right)}\mathrel{\vbox{\hbox{.}\hbox{.} }}=\mathcal{N}\mathopen{}\mathclose{{}\left(\upsilon_{l}\rho_{l}{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}},\>% \upsilon^{2}_{l}\mathbf{I}}\right).

Under this definition, $\rho_{l}^{2}$ is a kind of signal-to-noise ratio. We require it to decrease monotonically with $l$ .

Consistency with these marginals requires linear-Gaussian transitions:

\displaystyle{\bm{\check{X}}}_{m}=\beta_{ml}{\bm{\check{X}}}_{l}+\gamma_{ml}{% \bm{\check{Z}}}_{m},

\displaystyle{\bm{\check{Z}}}_{m}\sim\mathcal{N}\mathopen{}\mathclose{{}\left(% \bm{0},\>\mathbf{I}}\right).

Note that $m$ and $l$ need not even be consecutive steps in the Markov chain, although we require $m>l$ . Furthermore, by the law of total expectation,

\upsilon_{m}\rho_{m}{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}% {rgb}{.75,0,.25}\bm{x}_{0}}=\mathbb{E}_{{\bm{\check{X}}}_{m}|{\bm{\check{X}}}_% {0}}{\mathopen{}\mathclose{{}\left[{\bm{\check{X}}}_{m}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}}}% \right]}=\beta_{ml}\mathbb{E}_{{\bm{\check{X}}}_{l}|{\bm{\check{X}}}_{0}}{% \mathopen{}\mathclose{{}\left[{\bm{\check{X}}}_{l}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}}}% \right]}=\upsilon_{l}\rho_{l}\beta_{ml}{\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}}\implies\beta_{ml}=\frac{% \upsilon_{m}}{\upsilon_{l}}\frac{\rho_{m}}{\rho_{l}}.

Likewise, by the law of total covariance,

\begin{split}\upsilon^{2}_{m}\mathbf{I}&{}=\text{Cov}_{{\bm{\check{X}}}_{m}|{% \bm{\check{X}}}_{0}}{\mathopen{}\mathclose{{}\left[{\bm{\check{X}}}_{m}\middle% |{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm% {x}_{0}}}\right]}=\beta_{ml}^{2}\text{Cov}_{{\bm{\check{X}}}_{l}|{\bm{\check{X% }}}_{0}}{\mathopen{}\mathclose{{}\left[{\bm{\check{X}}}_{l}\middle|{\color[rgb% ]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}}}% \right]}+\gamma_{ml}^{2}\mathbf{I}=\mathopen{}\mathclose{{}\left(\upsilon^{2}_% {l}\beta_{ml}^{2}+\gamma_{ml}^{2}}\right)\mathbf{I}\\ \implies\gamma_{ml}^{2}&{}=\upsilon^{2}_{m}-\upsilon^{2}_{l}\mathopen{}% \mathclose{{}\left(\frac{\upsilon_{m}}{\upsilon_{l}}\frac{\rho_{m}}{\rho_{l}}}% \right)^{2}=\upsilon^{2}_{m}\mathopen{}\mathclose{{}\left(1-\frac{\rho_{m}^{2}% }{\rho_{l}^{2}}}\right).\end{split}

Notice (what our notation implied) that $\beta_{ml}$ and $\gamma_{ml}$ are necessarily scalars. In fine, the conditional recognition probabilities are given by

equation (10.35) (10.35)

{\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{m}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}}}% \right)}=\mathcal{N}\mathopen{}\mathclose{{}\left(\frac{\upsilon_{m}}{\upsilon% _{l}}\frac{\rho_{m}}{\rho_{l}}{\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}},\>\upsilon^{2}_{m}\mathopen{}% \mathclose{{}\left(1-\frac{\rho_{m}^{2}}{\rho_{l}^{2}}}\right)\mathbf{I}}% \right).

The recognition “posterior transitions.”

The other recognition distribution we require in order to use Eq. 10.33 is the “reverse-transition” ${\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l-1}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}},{% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }_{0}}}\right)}$ . Here we again solve for the more generic case of ${\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{m}},{% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }_{0}}}\right)}$ in which the state $l$ precedes $m$ but not necessarily directly. This distribution is again normal (all recognition distributions are), although this time the calculation of the cumulants is slightly more complicated, since it requires Bayes’ rule. Here the “prior” is ${\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}}}% \right)}$ (and given by the definition of the recognition marginals, Eq. 10.34); and the “likelihood” (or emission) is ${\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{m}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}},{% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }_{0}}}\right)}={\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25% }\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{m}}\middle|{\color% [rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}}% }\right)}$ (the resulting conditional recognition probabilities, given by Eq. 10.35). We have worked out the general case of Bayes rule for jointly Gaussian random variables in Section 3.1.2. From Eq. 3.14, the posterior precision is the sum of the (unnormalized) prior and likelihood precisions (in the space of ${\bm{\check{X}}}_{l}$ ):

\begin{split}\text{Cov}_{{\bm{\check{X}}}_{l}|{\bm{\check{X}}}_{m},{\bm{\check% {X}}}_{0}}{\mathopen{}\mathclose{{}\left[{\bm{\check{X}}}_{l}\middle|{\color[% rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{m}},% {\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{% x}_{0}}}\right]}&{}=\mathopen{}\mathclose{{}\left(\frac{1}{\upsilon^{2}_{l}}+% \frac{\rho_{l}^{2}}{\upsilon^{2}_{m}(\rho_{l}^{2}-\rho_{m}^{2})}\mathopen{}% \mathclose{{}\left(\frac{\upsilon_{m}}{\upsilon_{l}}\frac{\rho_{m}}{\rho_{l}}}% \right)^{2}}\right)^{-1}\mathbf{I}\\ &{}=\mathopen{}\mathclose{{}\left(\frac{1}{\upsilon^{2}_{l}}+\frac{\rho_{m}^{2% }}{\upsilon^{2}_{l}(\rho_{l}^{2}-\rho_{m}^{2})}}\right)^{-1}\mathbf{I}=\frac{% \upsilon^{2}_{l}(\rho_{l}^{2}-\rho_{m}^{2})}{\rho_{l}^{2}}\mathbf{I}.\end{split}

From Eq. 3.13, the posterior mean is a convex combination of the information from the prior and likelihood:

\mathbb{E}_{{\bm{\check{X}}}_{l}|{\bm{\check{X}}}_{m},{\bm{\check{X}}}_{0}}{% \mathopen{}\mathclose{{}\left[{\bm{\check{X}}}_{l}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{m}},{% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }_{0}}}\right]}=\frac{\rho_{l}^{2}-\rho_{m}^{2}}{\rho_{l}^{2}}\upsilon_{l}\rho% _{l}{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}% \bm{x}_{0}}+\frac{\upsilon^{2}_{l}}{\upsilon^{2}_{m}}\frac{\upsilon_{m}}{% \upsilon_{l}}\frac{\rho_{m}}{\rho_{l}}{\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{m}}=\frac{\upsilon_{l}}{\rho_{l}% }\mathopen{}\mathclose{{}\left((\rho_{l}^{2}-\rho_{m}^{2}){\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}}+\frac% {\rho_{m}}{\upsilon_{m}}{\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{m}}}\right).

Assembling the cumulants, we have

equation (10.36) (10.36)

{\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{m}},{% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }_{0}}}\right)}=\mathcal{N}\mathopen{}\mathclose{{}\left(\frac{\upsilon_{l}}{% \rho_{l}}\mathopen{}\mathclose{{}\left((\rho_{l}^{2}-\rho_{m}^{2}){\color[rgb]% {.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}}+% \frac{\rho_{m}}{\upsilon_{m}}{\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{m}}}\right),\>\frac{\upsilon^{2}_{l}(% \rho_{l}^{2}-\rho_{m}^{2})}{\rho_{l}^{2}}\mathbf{I}}\right).

The reverse-transition cross entropies, revisited.

We noted above that one way to evaluate the joint relative entropy for the diffusion model is to form Monte Carlo estimates of each of the summands in Eq. 10.32. Naïvely, we could evaluate each summand with samples from three random variables ( ${\bm{\check{X}}}_{0},{\bm{\check{X}}}_{l-1},{\bm{\check{X}}}_{l}$ ), but as we also noted, one of the expectations can actually be taken in closed-form. In particular, if we expand ${\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}},{\color[rgb]{.75,0,.25}% \definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l-1}}\middle|{% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }_{0}}{}}\right)}$ according the chain rule of probability given by Eq. 10.33—namely, in terms of the two distributions just derived, Eqs. 10.34 and 10.36 (letting $l$ be $l-1$ and $m$ be $l$ )—then we can write

equation (10.37) (10.37)

\begin{split}\mathcal{L}_{\text{JRE}}(\bm{\theta})&{}\stackrel{{\scriptstyle% \text{+}}}{{=}}\sum_{l=1}^{L}\frac{1}{2\sigma^{2}_{l}}\int_{\bm{\check{x}}_{0}% {}}\int_{\bm{\check{x}}_{l-1}{}}\int_{\bm{\check{x}}_{l}{}}{\check{p}\mathopen% {}\mathclose{{}\left(\bm{\check{x}}_{l-1}\middle|\bm{\check{x}}_{l},\bm{\check% {x}}_{0}}\right)}{\check{p}\mathopen{}\mathclose{{}\left(\bm{\check{x}}_{l}% \middle|\bm{\check{x}}_{0}}\right)}{p\mathopen{}\mathclose{{}\left(\bm{\check{% x}}_{0}}\right)}\mathopen{}\mathclose{{}\left\lVert\bm{\check{x}}_{l-1}-\bm{% \mu}_{l}}\right\rVert^{2}\mathop{}\!\mathrm{d}{\bm{\check{x}}_{l}{}}\mathop{}% \!\mathrm{d}{\bm{\check{x}}_{l-1}{}}\mathop{}\!\mathrm{d}{\bm{\check{x}}_{0}{}% }\\ &{}\stackrel{{\scriptstyle\text{+}}}{{=}}\sum_{l=1}^{L}\frac{1}{2\sigma^{2}_{l% }}\int_{\bm{\check{x}}_{0}{}}\int_{\bm{\check{x}}_{l}{}}{\check{p}\mathopen{}% \mathclose{{}\left(\bm{\check{x}}_{l}\middle|\bm{\check{x}}_{0}}\right)}{p% \mathopen{}\mathclose{{}\left(\bm{\check{x}}_{0}}\right)}\mathopen{}\mathclose% {{}\left\lVert\frac{\upsilon_{l-1}}{\rho_{l-1}}\mathopen{}\mathclose{{}\left((% \rho_{l-1}^{2}-\rho_{l}^{2})\bm{\check{x}}_{0}+\frac{\rho_{l}}{\upsilon_{l}}% \bm{\check{x}}_{l}}\right)-\bm{\mu}_{l}}\right\rVert^{2}\mathop{}\!\mathrm{d}{% \bm{\check{x}}_{l}{}}\mathop{}\!\mathrm{d}{\bm{\check{x}}_{0}{}}.\end{split}

In moving to the second line, the expectation of the quadratic form was taken under ${\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l-1}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}},{% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }_{0}}}\right)}$ using the identity B.16 from the appendix, except that the trace term, again a function of the fixed variance, was absorbed into the (now different) constant $c$ . The remaining expectations can be approximated with Monte Carlo estimates, since we have in hand samples from the data distribution, ${p\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}}{}}\right)}$ , and it is straightforward to generate samples of ${\bm{\check{X}}}_{l}$ from Eq. 10.34.

How are we to interpret Eq. 10.37? It would be convenient if this could also be expressed as the mean squared error between an uncorrupted sample and a predictor—call it $\bm{\hat{m}}({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{% .75,0,.25}\bm{x}_{l}},l,\bm{\theta})$ —that has access only to corrupted samples ( $\bm{\check{x}}_{l}$ ), as in Eq. 10.32. Of course it can, if we simply reparameterize the generative mean function on analogy with the mean of ${\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l-1}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}},{% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }_{0}}}\right)}$ (Eq. 10.36):

\bm{\mu}({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{% .75,0,.25}\bm{x}_{l}},l,\bm{\theta})\mathrel{\vbox{\hbox{.}\hbox{.} }}=\frac{\upsilon_{l-1}}{\rho_{l-1}}\mathopen{}\mathclose{{}\left((\rho_{l-1}^% {2}-\rho_{l}^{2})\bm{\hat{m}}({\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}},l,\bm{\theta})+\frac{\rho_{l}}{% \upsilon_{l}}{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{% .75,0,.25}\bm{x}_{l}}}\right).

Note that this reparameterization loses no generality. In terms of this predictor, the joint relative entropy then becomes

equation (10.38) (10.38)

\begin{split}\mathcal{L}_{\text{JRE}}(\bm{\theta})&{}\stackrel{{\scriptstyle% \text{+}}}{{=}}\sum_{l=1}^{L}\frac{\upsilon_{l-1}^{2}}{2\sigma^{2}_{l}\rho_{l-% 1}^{2}}(\rho_{l-1}^{2}-\rho_{l}^{2})^{2}\int_{\bm{\check{x}}_{0}{}}\int_{\bm{% \check{x}}_{l}{}}{\check{p}\mathopen{}\mathclose{{}\left(\bm{\check{x}}_{l}% \middle|\bm{\check{x}}_{0}}\right)}{p\mathopen{}\mathclose{{}\left(\bm{\check{% x}}_{0}}\right)}\mathopen{}\mathclose{{}\left\lVert\bm{\check{x}}_{0}-\bm{\hat% {m}}_{l}}\right\rVert^{2}\mathop{}\!\mathrm{d}{\bm{\check{x}}_{l}{}}\mathop{}% \!\mathrm{d}{\bm{\check{x}}_{0}{}}.\end{split}

So as in Eq. 10.32, minimizing the loss amounts to optimizing a denoising function. But in this case, it is the completely uncorrupted data samples, $\bm{\check{x}}_{0}$ , that are to be recovered, and accordingly a different (but related) denoising function/neural network ( $\bm{\hat{m}}_{l}$ ) that is to be used. The integrals in Eq. 10.38 are again to be estimated with samples, but from only two rather than three random variables.

Up till now we have refrained from specifying $\sigma^{2}_{l}$ . However, we now note that the expectations carried out in Eq. 10.37 amount to computing the cross entropy between ${\check{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l-1}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}},{% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }_{0}}}\right)}$ and ${\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l-1}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}};\bm{% \theta}}\right)}$ . Since cross entropy is minimized when the distributions are equal, it seems sensible simply to equate their variances.⁹⁹ 9 Nevertheless, this is not quite optimal. We will not in general be able to set the means of these distributions precisely equal, so the variance $\sigma^{2}_{l}$ really ought to soak up the difference. Comparing Eqs. 10.36 and 10.31, we have

\sigma^{2}_{l}\stackrel{{\scriptstyle\text{set}}}{{=}}\upsilon^{2}_{l-1}(\rho_% {l-1}^{2}-\rho_{l}^{2})/\rho_{l-1}^{2},

in which case Eq. 10.38 simplifies to the even more elegant

equation (10.39) (10.39)

\begin{split}\mathcal{L}_{\text{JRE}}(\bm{\theta})&{}\stackrel{{\scriptstyle% \text{+}}}{{=}}\sum_{l=1}^{L}\frac{\rho_{l-1}^{2}-\rho_{l}^{2}}{2}\int_{\bm{% \check{x}}_{0}{}}\int_{\bm{\check{x}}_{l}{}}{\check{p}\mathopen{}\mathclose{{}% \left(\bm{\check{x}}_{l}\middle|\bm{\check{x}}_{0}}\right)}{p\mathopen{}% \mathclose{{}\left(\bm{\check{x}}_{0}}\right)}\mathopen{}\mathclose{{}\left% \lVert\bm{\check{x}}_{0}-\bm{\hat{m}}_{l}}\right\rVert^{2}\mathop{}\!\mathrm{d% }{\bm{\check{x}}_{l}{}}\mathop{}\!\mathrm{d}{\bm{\check{x}}_{0}{}}\\ &{}\stackrel{{\scriptstyle\text{+}}}{{\approx}}\sum_{l=1}^{L}\frac{\rho_{l-1}^% {2}-\rho_{l}^{2}}{2}{\mathopen{}\mathclose{{}\left\langle{\mathopen{}% \mathclose{{}\left\lVert{\bm{\check{X}}}_{0}-\bm{\hat{m}}({\bm{\check{X}}}_{l}% ,l,\bm{\theta})}\right\rVert^{2}}}\right\rangle_{{\bm{\check{X}}}_{l}{},{\bm{% \check{X}}}_{0}{}}}.\end{split}

In words, each summand in Eq. 10.39 computes the mean squared error between the uncorrupted data ${\bm{\check{X}}}_{0}$ and a denoised version of the corrupted data, ${\bm{\check{X}}}_{l}$ . But before summing, the MSE at step $l$ is weighted by the amount of SNR lost in transitioning from step $l-1$ to $l$ . In fact, Eq. 10.39 tells us that fitting a Gaussian reverse-diffusion model is equivalent to fitting a (conceptually) different generative model:

equation (10.40) (10.40)

{\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{0}}\middle|{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}};\bm{% \theta}}\right)}=\mathcal{N}\mathopen{}\mathclose{{}\left(\bm{\hat{m}}({\color% [rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}}% ,l,\bm{\theta}),\>\frac{1}{\rho_{l-1}^{2}-\rho_{l}^{2}}\mathbf{I}}\right),% \quad l\in 1,\ldots,L.

Implementation

The pathwise gradient.

In actually carrying out the sample average for Eq. 10.39, we would typically reparameterize ${\bm{\check{X}}}_{l}$ along the lines of Eq. 10.20, i.e. as a scaled and shifted standard normal variate ( ${\bm{\check{Z}}}$ ), using Eq. 10.34, and then apply the LotUS:

equation (10.41) (10.41)

\mathcal{L}_{\text{JRE}}(\bm{\theta})\stackrel{{\scriptstyle\text{+}}}{{=}}% \sum_{l=1}^{L}\frac{\rho_{l-1}^{2}-\rho_{l}^{2}}{2}{\mathopen{}\mathclose{{}% \left\langle{\mathopen{}\mathclose{{}\left\lVert{\bm{\check{X}}}_{0}-\bm{\hat{% m}}\mathopen{}\mathclose{{}\left(\upsilon_{l}(\rho_{l}{\bm{\check{X}}}_{0}+{% \bm{\check{Z}}}),l,\bm{\theta}}\right)}\right\rVert^{2}}}\right\rangle_{{\bm{% \check{Z}}}{},{\bm{\check{X}}}_{0}{}}}.

The continuous-time limit.

Eq. 10.34 tells us that the data can be corrupted to an arbitrary position in the Markov chain with a single computation. Consequently, it is not actually necessary to run the chain sequentially from 1 to $L$ during training, which is critical for parallelized implementations. Indeed, we need not even limit ourselves to an integer number of steps. Suppose we allow the SNR to be a monotonically decreasing function $h$ of a continuous variable $u$ that ranges from 0 to 1, such that $h(l/L)=\rho_{l}^{2}$ . For consistency, we will define another function on $[0,1]$ for the marginal variance, such that $g(l/L)=\upsilon^{2}_{l}$ , although $g$ need not be monotonic. Then if we scale the joint relative entropy in Eq. 10.39 by the “step size” $1/L$ and take the limit as $L\to\infty$ , the loss becomes

\begin{split}\lim_{L\to\infty}\frac{1}{L}\mathcal{L}_{\text{JRE}}(\bm{\theta})% &{}\stackrel{{\scriptstyle\text{+}}}{{=}}\lim_{L\to\infty}\sum_{l=1}^{L}\frac{% 1}{2L}\mathopen{}\mathclose{{}\left(h\mathopen{}\mathclose{{}\left(\frac{l-1}{% L}}\right)-h\mathopen{}\mathclose{{}\left(\frac{l}{L}}\right)}\right)\\ &{}\qquad\qquad{\mathopen{}\mathclose{{}\left\langle{\mathopen{}\mathclose{{}% \left\lVert{\bm{\check{X}}}_{0}-\bm{\hat{m}}\mathopen{}\mathclose{{}\left(% \sqrt{g(l/L)h(l/L)}{\bm{\check{X}}}_{0}+\sqrt{g(l/L)}{\bm{\check{Z}}},l/L,\bm{% \theta}}\right)}\right\rVert^{2}}}\right\rangle_{{\bm{\check{Z}}}{},{\bm{% \check{X}}}_{0}{}}}\\ &{}=\frac{1}{2}\int_{0}^{1}\frac{\mathrm{d}{h}}{\mathrm{d}{u}}(u){\mathopen{}% \mathclose{{}\left\langle{\mathopen{}\mathclose{{}\left\lVert{\bm{\check{X}}}_% {0}-\bm{\hat{m}}\mathopen{}\mathclose{{}\left(\sqrt{g(u)h(u)}{\bm{\check{X}}}_% {0}+\sqrt{g(u)}{\bm{\check{Z}}},u,\bm{\theta}}\right)}\right\rVert^{2}}}\right% \rangle_{{\bm{\check{Z}}}{},{\bm{\check{X}}}_{0}{}}}\mathop{}\!\mathrm{d}{u}\\ &{}=\frac{1}{2}{\mathopen{}\mathclose{{}\left\langle{\frac{\mathrm{d}{h}}{% \mathrm{d}{u}}(U)\mathopen{}\mathclose{{}\left\lVert{\bm{\check{X}}}_{0}-\bm{% \hat{m}}\mathopen{}\mathclose{{}\left(\sqrt{g(U)h(U)}{\bm{\check{X}}}_{0}+% \sqrt{g(U)}{\bm{\check{Z}}},U,\bm{\theta}}\right)}\right\rVert^{2}}}\right% \rangle_{{\bm{\check{Z}}}{},{\bm{\check{X}}}_{0}{},U{}}}\end{split}

with $U\sim\mathcal{U}\mathopen{}\mathclose{{}\left(0,1}\right)$ a uniformly distributed random variable. This is the preferred implementation of diffusion models [36]. But the second line also suggests the change of variables $\lambda=h(u)$ , under which the integral becomes

\lim_{L\to\infty}\frac{1}{L}\mathcal{L}_{\text{JRE}}(\bm{\theta})=\frac{1}{2}% \int_{\lambda{}}{\mathopen{}\mathclose{{}\left\langle{\mathopen{}\mathclose{{}% \left\lVert{\bm{\check{X}}}_{0}-\bm{\hat{m}}\mathopen{}\mathclose{{}\left(% \sqrt{g(h^{-1}(\lambda))\lambda}{\bm{\check{X}}}_{0}+\sqrt{g(h^{-1}(\lambda))}% {\bm{\check{Z}}},h^{-1}(\lambda),\bm{\theta}}\right)}\right\rVert^{2}}}\right% \rangle_{{\bm{\check{Z}}}{},{\bm{\check{X}}}_{0}{}}}\mathop{}\!\mathrm{d}{% \lambda{}}.

Notice that $h$ can be safely removed from this equation, since we have not yet committed to any particular $g$ , and $\bm{\hat{m}}$ is assumed to be arbitrarily flexible. This shows that in the continuous-time limit, any choice of SNR function yields the same joint relative entropy in expectation, as long as (1) it is monotonically decreasing and (2) it has well-chosen endpoints, $\lambda_{\text{min}}$ and $\lambda_{\text{max}}$ . However, this choice does affect the variance of this sample average, and various SNR “schedules” have been experimented with in practice [36].

A connection to denoising score matching

There is in fact another illuminating reparameterization, this time in terms of the negative energy gradient^†^†margin: negative energy gradient , $-\partial{U}/\partial{{\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}}}$ . Because this quantity can be written (generically) as $\partial{\mathopen{}\mathclose{{}\left[\log{p\mathopen{}\mathclose{{}\left({% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }}}\right)}}\right]}/\partial{{\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}{\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}}}}$ , it is sometimes called the score function for its resemblance to $\partial{\mathopen{}\mathclose{{}\left[\log p({\color[rgb]{.75,0,.25}% \definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}};\bm{\theta})}\right% ]}/\partial{\bm{\theta}}$ . We will call it the force^†^†margin: force . Intuitively, the force points toward the modal $\bm{x}$ . Furthermore, if our goal in fitting a generative model is merely to synthesize new data, then the force suffices, because the iteration

equation (10.42) (10.42)

{\bm{X}}_{i+1}={\bm{X}}_{i}-\epsilon\frac{\partial{U}}{\partial{{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}}}}}({\bm{X% }}_{i})+\sqrt{2\epsilon}{\bm{Z}}_{i},

(with ${\bm{Z}}_{l}\sim\mathcal{N}\mathopen{}\mathclose{{}\left(0,\>\mathbf{I}}\right)$ , and step size $\epsilon$ ) can be shown to generate samples approximately from the distribution ${p\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}}}\right)}\propto\exp\mathopen{}% \mathclose{{}\left\{-U({\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}})}\right\}$ . This iteration is known as Langevin dynamics^†^†margin: Langevin dynamics , and we return to it in Section 12.2.2. For now we simply ask how we might get or estimate the force.

Consider a random variable ${\bm{Y}}$ that was created by corrupting the random variable of interest, ${\bm{X}}$ , with some kind of additive noise. The resulting marginal distribution of ${\bm{Y}}$ ,

p(\bm{y})=\int_{\bm{x}{}}{p\mathopen{}\mathclose{{}\left(\bm{x}}\right)}p_{% \text{noise}}(\bm{y}-\bm{x})\mathop{}\!\mathrm{d}{\bm{x}{}}\approx\sum_{n}^{{N% }}p_{\text{noise}}(\bm{y}-\bm{x}_{n})

can be thought of as a kernel-density estimate of the distribution of interest, ${p\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}}}\right)}$ , with ${N}$ data samples and kernel $p_{\text{noise}}$ . So perhaps we can use the former in the place of the latter in our Langevin dynamics, Eq. 10.42. But then how are we to get the force of $p(\bm{y})$ ? Expanding it, we find that

equation (10.43) (10.43)

\begin{split}\frac{\partial{}}{\partial{{\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}}^{\text{T}}}\log p(\bm{y})=\frac% {1}{p(\bm{y})}\frac{\partial{}}{\partial{{\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}}^{\text{T}}}\int_{\bm{x}{}}p(\bm% {y}|\bm{x})p(\bm{x})\mathop{}\!\mathrm{d}{\bm{x}{}}&{}=\frac{1}{p(\bm{y})}\int% _{\bm{x}{}}p(\bm{y}|\bm{x})\frac{\partial{}}{\partial{{\color[rgb]{.75,0,.25}% \definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}}^{\text{T}}}\log% \mathopen{}\mathclose{{}\left(p(\bm{y}|\bm{x})}\right)p(\bm{x})\mathop{}\!% \mathrm{d}{\bm{x}{}}\\ &{}=\mathbb{E}_{{\bm{X}}{}|{\bm{Y}}}{\mathopen{}\mathclose{{}\left[\frac{% \partial{}}{\partial{{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor% }{rgb}{.75,0,.25}\bm{y}}}^{\text{T}}}\log\mathopen{}\mathclose{{}\left(p(\bm{y% }|{\bm{X}})}\right)\middle|\bm{y}{}}\right]}.\end{split}

This equation says that the force of the marginal $p(\bm{y})$ equals the expected (under $p({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}% \bm{x}}|\bm{y})$ ) force of the conditional, $p(\bm{y}|{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{% .75,0,.25}\bm{x}})$ .¹⁰¹⁰ 10 Cf. Eq. 8.3. This is the exact same manoeuver used to introduce latent variables in the derivation of EM, although there the derivative is with respect to parameters. The latter can often be computed. For example, if the data are corrupted by scaling and then adding zero-mean Gaussian noise, then the conditional energy and its expected negative gradient (force) are

U(\bm{y}|{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{% .75,0,.25}\bm{x}})=\frac{(\bm{y}-\alpha{\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}})^{\text{T}}\Sigma^{-1}(\bm{y}-% \alpha{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{% .75,0,.25}\bm{x}})}{2}\implies\mathbb{E}_{{\bm{X}}{}|{\bm{Y}}}{\mathopen{}% \mathclose{{}\left[-\frac{\partial{}}{\partial{{\color[rgb]{.75,0,.25}% \definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}}^{\text{T}}}U(\bm{y% }|{\bm{X}})\middle|\bm{y}{}}\right]}=\Sigma^{-1}\mathopen{}\mathclose{{}\left(% \alpha\mathbb{E}{\mathopen{}\mathclose{{}\left[{\bm{X}}|\bm{y}}\right]}-\bm{y}% }\right).

Putting this together with Eq. 10.43, we see that for additive Gaussian noise,

equation (10.44) (10.44)

\frac{\partial{}}{\partial{{\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}}^{\text{T}}}\log p(\bm{y})=\Sigma^{-1}% \mathopen{}\mathclose{{}\left(\alpha\mathbb{E}{\mathopen{}\mathclose{{}\left[{% \bm{X}}|\bm{y}}\right]}-\bm{y}}\right)

This is sometimes called Tweedie’s formula^†^†margin: Tweedie’s formula . This looks helpful—if we had in hand the posterior mean!

Now it is a fact that, of all estimators for ${\bm{X}}$ , the posterior mean has minimum mean squared error [62]. Putting all these pieces together [31, 76, Kadkhodaie2021] yields the following procedure to generate samples from the distribution of interest, ${p\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}}}\right)}$ : (1) Find an estimator for ${\bm{X}}$ that minimizes mean square error; (2) use this in place of the posterior mean in Eq. 10.44 to compute the expected conditional force and, consequently, the marginal force; (3) use the marginal force, $\partial{\mathopen{}\mathclose{{}\left[\log p(\bm{y})}\right]}/\partial{{% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y% }}}$ , as a proxy for the data force, $\partial{\mathopen{}\mathclose{{}\left[\log{p\mathopen{}\mathclose{{}\left({% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }}}\right)}}\right]}/\partial{{\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}{\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}}}}$ , in Langevin dynamics (Eq. 10.42). This method of density estimation is known as denoising score matching^†^†margin: denoising score matching .

Now we examine the diffusion model in light of this procedure. Fitting the generative model to the data has turned out to be equivalent to minimizing the mean squared error between ${\bm{\check{X}}}_{0}$ and $\bm{\hat{m}}({\bm{\check{X}}}_{l},l,\bm{\theta})$ (Eq. 10.39). Therefore we can interpret $\bm{\hat{m}}_{l}$ as (an estimator for) the posterior mean, $\mathbb{E}{\mathopen{}\mathclose{{}\left[{\bm{\check{X}}}_{0}|{\bm{\check{X}}}% _{l}}\right]}$ . The samples ${\bm{\check{X}}}_{l}$ are generated by a recognition model that corrupts the data samples ${\bm{\check{X}}}_{0}$ with Gaussian noise (Eq. 10.34). Therefore we can use the posterior-mean estimator $\bm{\hat{m}}_{l}$ and Tweedie’s formula (Eq. 10.44) to construct a force estimator:

equation (10.45) (10.45)

\frac{\upsilon_{l}\rho_{l}\bm{\hat{m}}({\color[rgb]{.75,0,.25}\definecolor[% named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}},l,\bm{\theta})-{\color[rgb]{% .75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}}}{% \upsilon^{2}_{l}}=\mathrel{\vbox{\hbox{.}\hbox{.} }}\bm{\hat{f}}({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}% {.75,0,.25}\bm{x}_{l}},l,\bm{\theta}).

The force estimator $\bm{\hat{f}}_{l}$ also provides a good proxy for the data force, $\partial{\mathopen{}\mathclose{{}\left[\log{p\mathopen{}\mathclose{{}\left({% \color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x% }_{0}}{}}\right)}}\right]}/\partial{{\color[rgb]{.75,0,.25}\definecolor[named]% {pgfstrokecolor}{rgb}{.75,0,.25}\bm{\check{x}}_{0}}}$ , and consequently can be used to generate data with Langevin dynamics (Eq. 10.42).

Alternatively, the force can be fit directly, rather than indirectly via the posterior mean. That is easily done here as well, simply by rearranging Eq. 10.45 to reparameterize $\bm{\hat{m}}$ (again without loss of generality):

\bm{\hat{m}}({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{% .75,0,.25}\bm{x}_{l}},l,\bm{\theta})=\mathrel{\vbox{\hbox{.}\hbox{.} }}\frac{\upsilon_{l}}{\rho_{l}}\bm{\hat{f}}({\color[rgb]{.75,0,.25}% \definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}},l,\bm{\theta})+% \frac{1}{\upsilon_{l}\rho_{l}}{\color[rgb]{.75,0,.25}\definecolor[named]{% pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}_{l}}

for some arbitrary function (neural network) $\bm{\hat{f}}$ . Under this reparameterization, Eq. 10.39 becomes

equation (10.46) (10.46)

\begin{split}\mathcal{L}_{\text{JRE}}(\bm{\theta})&{}\stackrel{{\scriptstyle% \text{+}}}{{=}}\sum_{l=1}^{L}\frac{\rho_{l-1}^{2}-\rho_{l}^{2}}{2}{\mathopen{}% \mathclose{{}\left\langle{\mathopen{}\mathclose{{}\left\lVert{\bm{\check{X}}}_% {0}-\frac{\upsilon_{l}}{\rho_{l}}\bm{\hat{f}}\mathopen{}\mathclose{{}\left({% \bm{\check{X}}}_{l},l,\bm{\theta}}\right)-\frac{1}{\upsilon_{l}\rho_{l}}{\bm{% \check{X}}}_{l}}\right\rVert^{2}}}\right\rangle_{{\bm{\check{X}}}_{l}{},{\bm{% \check{X}}}_{0}{}}}\\ &{}\stackrel{{\scriptstyle\text{+}}}{{=}}\sum_{l=1}^{L}\frac{\rho_{l-1}^{2}-% \rho_{l}^{2}}{2}\frac{\upsilon^{2}_{l}}{\rho_{l}^{2}}{\mathopen{}\mathclose{{}% \left\langle{\mathopen{}\mathclose{{}\left\lVert\frac{\upsilon_{l}\rho_{l}{\bm% {\check{X}}}_{0}-{\bm{\check{X}}}_{l}}{\upsilon^{2}_{l}}-\bm{\hat{f}}\mathopen% {}\mathclose{{}\left({\bm{\check{X}}}_{l},l,\bm{\theta}}\right)}\right\rVert^{% 2}}}\right\rangle_{{\bm{\check{X}}}_{l}{},{\bm{\check{X}}}_{0}{}}}\\ &{}\stackrel{{\scriptstyle\text{+}}}{{=}}\sum_{l=1}^{L}\frac{\rho_{l-1}^{2}-% \rho_{l}^{2}}{2\rho_{l}^{2}}{\mathopen{}\mathclose{{}\left\langle{\mathopen{}% \mathclose{{}\left\lVert{\bm{\check{Z}}}_{l}+\upsilon_{l}\bm{\hat{f}}\mathopen% {}\mathclose{{}\left(\upsilon_{l}(\rho_{l}{\bm{\check{X}}}_{0}+{\bm{\check{Z}}% }_{l}),l,\bm{\theta}}\right)}\right\rVert^{2}}}\right\rangle_{{\bm{\check{Z}}}% _{l}{},{\bm{\check{X}}}_{0}{}}}.\end{split}

The last step follow from reparameterization. Intuitively, $\bm{\hat{f}}$ learns, like $\bm{\hat{m}}$ , how to uncorrupt data. But rather than transforming the corrupted sample ( $\bm{\check{x}}_{l}$ ) into an estimate of the (scaled) uncorrupted sample itself ( $\upsilon_{l}\rho_{l}\bm{\check{x}}_{0}$ ), a good $\bm{\hat{f}}_{l}$ produces (second line of Eq. 10.46) an estimate of the vector that points back to $\upsilon_{l}\rho_{l}\bm{\check{x}}_{0}$ from the corrupted sample $\bm{\check{x}}_{l}$ . This is consistent with our conclusion that any $\bm{\hat{f}}$ that satisfies Eq. 10.45 provides an estimator for the force of the data distribution. Alternatively, the final line of Eq. 10.46 tells us that the force estimator $\bm{\hat{f}}_{l}$ must try to recover each realizaton of noise ( $\bm{\check{z}}_{l}$ ) that corrupted each observed datum ( $\bm{\check{x}}_{0}$ ). But notice that the negative force, i.e. the positive energy gradient, must point in the direction of ${\bm{\check{Z}}}_{l}$ . This makes sense: we expect the noise to be “uphill.”

In either case, each summand in Eq. 10.46 corresponds to an objective for denoising score matching; or, put the other way around, fitting a Gaussian (reverse-)diffusion model (Eqs. 10.28, 10.29, 10.35, and 10.31) amounts to running denoising score matching for many different kernel widths. And indeed, such a learning procedure has been proposed and justified independently under this description [71].