DDPM Math

Main Expression

\(\log p_\theta(\mathbf{x}) = \log \int_{\mathbf{z}_{1,…,T}} p_\theta(\mathbf{x},\mathbf{z}_{1,…,T} )d\mathbf{z}_{1,…,T}\)

Focusing on the Numerator in the log

Now,

We have parametrized the decoder such that \(p_\theta( \mathbf{x} \mid \mathbf{z}_1)\) is conditionally independent from \(\mathbf{z}_2,….,\mathbf{z}_T\). We call this the Markov property of \(p\). This is:

Expanding by using the Chain Rule:

Once again, using the Markov property of \(p\):

Continuing, we get

Thus, the probability of taking a particular path from \(\mathbf{z}_T\) to \(\mathbf{x}\) is given by the above expression. Which is typical, this is just the chain rule applied to Markov chain.

One way to put this:

The integrals iterate over all possible values of \(\mathbf{z}_1, \dots, \mathbf{z}_T\). They are then plugged in for the expression for the joint distribution \(p_\theta(\mathbf{x}, \mathbf{z}_{1,\dots,T})\) to get a probability value, which accumulates across the loop. \(p_\theta(\mathbf{x}, \mathbf{z}_{1,\dots,T})\) maps a tuple of values \((\mathbf{x}, \mathbf{z}_{1,\dots,T})\) to a density value based on the joint distribution.

However, given this same tuple, \((\mathbf{x}, \mathbf{z}_{1,\dots,T})\), we can evaluate the density by evaluating the probability density of the ‘path’ that goes from \(\mathbf{z}_T\) to \(\mathbf{x}\), which is this expression:

Which is the probability that the decoder takes that ‘path’ from \(\mathbf{z}_T\) to \(\mathbf{x}\).

Focusing on the Denominator in the log

Now consider:

Since the forward process is a Markov Chain:

Continuing:

Thus, evaluating \(q( \mathbf{z}_{1,…,T} \mid \mathbf{x})\) involves starting with \(\mathbf{x}\), and evaluating the probability fo the chain of events leading from \(\mathbf{x}\) to \(\mathbf{z}_T\).

Using Bayes’ rule:

The first step seems like a hack – since \(\mathbf{z}_t\) conditioned on \(\mathbf{z}_{t-1}\) is independent from \(\mathbf{x}\), we can add in the extra condition on \(\mathbf{x}\) without worrying.

Thus,

Things cancel in the fraction:

Focusing on the log part

Putting this together, we have

Assume the last term goes to zero, since the distribution after all the forward diffusion steps should be very similar to \(p_\theta\left(\mathbf{z}_T\right)=\mathcal{N}\left(0,\mathbf{I}\right)\).

Why is the distribution for \(p_\theta\left(\mathbf{z}_T\right)=N\left(0,\mathbf{I}\right)?\) Well, we can choose it to be that way, by making the ‘decoder’ evaluate it as such. Or, we can think of p as attempting to fit the ‘true’ distribution of the data, which is done approximately in this case by being \(\mathcal{N}\left(0,\mathbf{I}\right)\).

Let us put this in the integral:

Back to Main Expression

We can marginalize out a lot of stuff. Here is an example:

Single Term

The stuff in the brackets goes to 1, since all conditional distributions are still distributions, they integrate to 1 over their input variable(s).

\(=\int_{\mathbf{z}_t}{\left\{\int_{\mathbf{z}_{t-1}}{\log{\left[\frac{p_\theta\left(\mathbf{z}_{t-1}| \mathbf{z}_t\right)}{q\left(\mathbf{z}_{t-1}| \mathbf{z}_t\right)}\right]}q\left(\mathbf{z}_{t-1}| \mathbf{z}_t,\ \mathbf{x}\right)d\mathbf{z}_{t-1}}\right\}q\left(\mathbf{z}_t| \mathbf{x}\right)}{dz}_t\) The term in the middle is a KL expression: \(=-\int_{\mathbf{z}_t}{D_{KL}\left[q\left(\mathbf{z}_{t-1}| \mathbf{z}_t\right)\ || \ p_\theta\left(\mathbf{z}_{t-1}| \mathbf{z}_t\right)\right]q\left(\mathbf{z}_t| \mathbf{x}\right)}{dz}_t\)

Back to Main Expression

Marginalizing the first term (won’t show the whole thing this time):

Turning the first term into an expectation:

Focusing on the Later KL Terms

According to our parametrization, we have

And we have from other computations that:

Thus, since we are computing the KL between two Gaussian distributions, we can actually compute the KL divergence here in closed form:

Which is proportional to the squared difference between the means.

First Term:

Back to Main Expression

Maximizing this means minimizing this:

We can imagine minimizing this term regarding a specific \(\mathbf{x}\) by sampling \(\mathbf{z}_1,\ldots,\mathbf{z}_{T}\) from \(\mathbf{x}\), computing the expression, and changing the parameters of \(\mathbf{f}\). We can see this more clearly by adding stuff to the expectations.

Now using a Monte Carlo estimate:

The loss function minimizes the difference between the estimated mean \(\mathbf{f}(\mathbf{z}_t)\) of \(\mathbf{z}_{t-1}\), and the most likely value (mean) it took, given \(\mathbf{z}_t\) and \(\mathbf{x}\).