2
$\begingroup$

I have been studying diffusion models from this tutorial: https://arxiv.org/abs/2403.18103 and trying to derive all results as I read it. Although this tutorial is very comprehensive, it skips many of the derivation steps. I am currently stuck in the derivation of the Evidence Lower Bound (ELBO) terms for variational diffusion model (Page 17). The ELBO consists of 3 terms: reconstruction loss term, prior matching and consistency term. I am struggling with the derivation for the consistency term as I am trying to expand the expectation as an integral over the specified probability distribution $q_{\phi}(x_{1:T}|x_0)$. To simplify the expression, I am using the Chapman-Kolmogorov property for the Markov process, as explained here

$$\int dx_t q_{\phi}(x_{t+1}|x_t)q_{\phi}(x_t|x_{t-1}) = q_{\phi}(x_{t+1}|x_{t-1})$$

Solving the consistency term,

$$E_{q_{\phi}(x_{1:T}|x_0)}\Big[\text{log} \prod_{t=1}^{T-1} \frac{p(x_t|x_{t+1})}{q_{\phi}(x_{t}|x_{t-1})}\Big] = \sum_{t=1}^{T-1} E_{q_{\phi}(x_{1:T}|x_0)}\Big[\text{log} \frac{p(x_t|x_{t+1})}{q_{\phi}(x_{t}|x_{t-1})} \Big] = \sum_{t=1}^{T-1} \int dx_1 dx_2 \ldots dx_{t-2}dx_{t-1}dx_{t}dx_{t+1}dx_{t+2} \ldots dx_{T} q_{\phi}(x_{T}|x_{T-1}) \ldots q_{\phi}(x_{t+2}|x_{t+1}) q_{\phi}(x_{t+1}|x_{t}) q_{\phi}(x_{t}|x_{t-1}) q_{\phi}(x_{t-1}|x_{t-2}) q_{\phi}(x_{t-2}|x_{t-3}) \ldots q_{\phi}(x_{1}|x_{0}) \text{log} \frac{p(x_t|x_{t+1})}{q_{\phi}(x_t | x_{t-1})}$$

Using the Chapman-Kolmogorov property, I can perform the integral for the indices $t \in$ {$1, \ldots, (t-2)$} $\cup$ { $(t+2),\ldots, (T-1)$}. After these integrations, the expression becomes,

$$ \sum_{t=1}^{T-1} \int dx_{t-1} dx_{t} dx_{t+1} dx_{T} q_{\phi}(x_T | x_{t+1}) q_{\phi}(x_{t+1}|x_{t}) q_{\phi}(x_{t}|x_{t-1}) q_{\phi}(x_{t-1}|x_{0}) \text{log} \frac{p(x_t|x_{t+1})}{q_{\phi}(x_t | x_{t-1})} $$

Question 1: In the above expression, can I perform the integral over the index $T$, as in $\int dx_T q_{\phi}(x_T|x_{t+1}) = 1$? Is this correct?

Let's say, I did perform the integral over index $T$, and proceed further as,

$$ \sum_{t=1}^{T-1} \int dx_{t-1} dx_{t} dx_{t+1} q_{\phi}(x_{t+1}|x_{t}) q_{\phi}(x_{t}|x_{t-1}) q_{\phi}(x_{t-1}|x_{0}) \text{log} \frac{p(x_t|x_{t+1})}{q_{\phi}(x_t | x_{t-1})} $$ $$ = -\sum_{t=1}^{T-1} \int dx_{t-1} dx_{t+1} q_{\phi}(x_{t+1}|x_{t}) q_{\phi}(x_{t-1}|x_{0}) \int dx_{t} q_{\phi}(x_{t}|x_{t-1}) \text{log} \frac{q_{\phi}(x_t | x_{t-1})}{p(x_t|x_{t+1})} $$ $$ = -\sum_{t=1}^{T-1} \int dx_{t-1} dx_{t+1} q_{\phi}(x_{t+1}|x_{t}) q_{\phi}(x_{t-1}|x_{0}) D_{KL}(q_{\phi}(x_t | x_{t-1}) \Vert p(x_t | x_{t+1})) $$

Question 2: Is the above step correct?

The final expression is supposed to look like the following, where the expectation is over $q_{\phi}(x_{t+1, t-1}|x_0)$ but I am not sure how to reduce my above expression to this:

$$ -\sum_{t=1}^{T-1} E_{q_{\phi}(x_{t+1, t-1}|x_0)}\Big [ D_{KL}(q_{\phi}(x_t | x_{t-1}) \Vert p(x_t | x_{t+1})) \Big] $$

Any feedback is appreciated. Thanks in advance!

$\endgroup$
1
  • $\begingroup$ Please, ask only one question per post. You can always create multiple posts and reference other posts. $\endgroup$
    – nbro
    Commented Jul 23 at 9:24

1 Answer 1

0
$\begingroup$

ELBO derivation for DDPM doesn't need the convoluted Chapman-Kolmogorov property for the Markov process, your tutorial also doesn't mention it at all. They key difficulty is to reverse the direction of $q_{\phi}(x_t|x_{t-1})$ in the original consistency term of the ELBO via Bayes rule and one more conditional rv $x_0$. Then the consistency term just takes expectation wrt to the tractable distribution $q_{\phi}(x_t|x_0)$ instead of the original odd $q_{\phi}(x_{t-1},x_{t+1}|x_0)$. Further on you can use the nice property of Gaussians (every type of conditional distribution here is Gaussian) of the forward and reverse consistent Markov chains to get a simple closed form training objective.

$\endgroup$

Not the answer you're looking for? Browse other questions tagged .