Derivations

This is some notes I took while learning diffusion models. Need to clean it up & turn it into a proper blog post.

$q(x_0)$ is our data-distribution. It is a datapoint, $x_0$, before any forward process diffusion. $p_\theta$ is what we're trying to learn

We can't optimize the 1st part of the final expression, so all that's left is to optimize $\mathbb{E}_{q(x_0)}[-\log p_\theta]$. $\mathbb{E}_{q(x_0)}$ is expectation under the data distribution. This is equivalent to sampling from the data distribution.

Expectation definition:

We currently have: $p_\theta(x_0)$. But the reverse diffusion process depends on latents $x_1$ through $x_T$. We need to pull these into our expression. To do that, we use marginalization.

What's marginalization?

• Let $A$ be a random variable that represents the probability of it raining tomorrow.
• Let $B$ be a random variable that represents some other event. Maybe: Tom catches Jerry today. (Note how the 2 events don't necessarily have to be related). The probability of it raining tomorrow can also be expressed as "probability of it raining tomorrow & probability that Tom catches Jerry today" + "probability of it raining tomorrow & probability that Tom does not catch Jerry today "

In this case, both $A$ and $B$ can take on 2 possible values.

If this was continuous, we would have:

Side note, in this case, $A$ and $B$ are independent, so we can do:

We now reach out 1st excerpt from the paper!

We needed to "pull in" all the latent variables that $\textbf{x}_0$ depends on. To do this, we marginalized $\textbf{x}_0$ with respect to all the latent variables in the reverse diffusion process.

Note:

Aside: Forward Diffusion

For the forward diffusion process:

We scale the mean by $\sqrt{1-\beta_t}$ in order to prevent the values/variance from exploding. TODO: Be more precise here. Write some tests & verify what is happening.

I originally thought it was b/c we want to scale the mean down to 0. Indeed, if we put in data composed of all 1s into the forward diffusion process, the end result would have a mean of 0 and a stdev of 1.

But, that isn't the goal of scaling the mean. In diffusion models, we assume our input has a mean of 0. Because if our forward diffusion process shifts the mean to 0, then our neural net must learn to shift the mean back. That's probably bad since our neural net uses the same weights for each reverse diffusion step.

Aside: Training-Inference Compute Asymetry

Diffusion models allow you to train with much less compute, while still utilizing a ton of compute during sampling time.

Aside: Functional Form

Both processes have the same functional form when $\beta_t$ are small

This is saying "we need a lot of timesteps in the forward/reverse diffusion processes" in order for both forward/reverse processes to be gaussian.

Example: Imagine your dataset consists of cats & dogs. If you do a single reverse step to go from noise to "full cat/dog" then your reverse step can't produce a gaussian distribution of outcomes. How do you fit a gaussian to cats & dogs? You would need a multimodal distribution instead.

ELBO

1. Marginalization to pull-in the latent variables
2. Bring in the forward diffusion process
3. Integral => Expectation
4. $\mathbb{E}[\log X] \leq \log \mathbb{E}[X]$ using Jensen's Inequality

Jensen's Inequality

How do we prove $\mathbb{E}[\log X] < \log \mathbb{E}[X]$?

Proof By Example

average(log(1), log(2), log(3)) $\simeq 0.59$ log(average(1, 2, 3)) $\simeq 0.69$

Picture Proof

Desmos ![[Pasted image 20220812143558.png]]

---

1. Use the fact that $p_\theta$ and $q$ are defined as products + log rules.

A note about $p(x_T)$:

Technically, it should be $p_\theta(x_T)$, but we assume that after enough forward diffusion steps, $p(x_T)$ is identical to the normal distribution (hence the equation earlier in the paper: $p(x_T) = \mathcal{N}(x_T;\textbf{0},\textbf{I})$).

Concretely, to calculate $\log p(x_T)$ we could:

1. Take an image from our data set, $x_0$
2. Run forward diffusion for T steps to get $T$.
3. Calculate the probability of $x_T$ appearing from the normal distribution, $\mathcal{N}$.
   • TODO: I'm guessing you would do this by calculating the probability of the normal distribution discretizing to the given image?

Rewriting ELBO

For simplicity, we'll write:

as

Now we'll simplify it:

1. Extract the 1st term out of the summation. This is necessary because when we apply Bayes' Theorem to $q$, we will get a non-sensical result if we also apply it to the 1st term in the summation.
2. Apply Bayes' rule to $q(x_t|x_{t-1})$
3. We need to condition the reverse conditional probability, $q$, on $x_0$. Why? $q(x_{t-1}|x_t)$ needs to give the probability distribution of $x_{t-1}$s given $x_t$, but this might be extremely difficult if, e.g, $x_t$ has a lot of noise (which it will near the end of the diffusion process). If we know the original image, $x_0$, this process becomes easy. This also makes the reverse conditional probability tractable. I.e., we can compute it.
4. Log rules.
5. Expand the 2nd summation, apply log rules to get the log of a cumulative product, cancel terms.
6. Log rules
7. Definition of KL-divergence

Reparameterization

We can do the forward process in 1-step. Need help simplifying an iterative diffusion process to a 1-step process

1. We go from $\beta_t$ to $\sqrt{\beta_t}$ because $\beta_tI$ is a covariance matrix. Specifically, it's the diagonal covariance matrix, so it only has covariance values of the form $\text{COV}(X, X) = \text{VAR}(X) = \sigma^2$
2. Substitute $\beta_t$ with $a_t$
3. Substitute using $x_{t-1} = \sqrt{\alpha_{t-1}}x_{t-2} + \sqrt{1 - \alpha_{t-1}}\epsilon$
4. Algebra
5. $\epsilon$ is sampled from the normal distribution $\epsilon = \mathcal{N}(0, I)$ Thus, the last 2 terms are gaussians with mean 0 and standard deviation $\sqrt{a_t}\sqrt{1-a_{t-1}}$ and $\sqrt{1-a_t}$. $\text{VAR}(X + Y) = \text{VAR}(X) + \text{VAR}(Y)$ if $X$ and $Y$ are independent. Thus, the sum of those two terms is a gaussian with mean 0 and variance $\alpha_t(1 - \alpha_{t-1}) + (1 - \alpha_t) = 1 - \alpha_t\alpha_{t-1}$

Deriving The Posterior

A lot of the math here comes from the excellent blog post What are Diffusion Models?

The reverse conditional probability is:

We need to calculate $\tilde\mu$ and $\tilde\beta$.

1. Bayes' Rule
2. This is the "proportional to" symbol. We're using these facts:
   1. The PDF of a gaussian is $\frac{1}{\sqrt{2\sigma^2\pi}} e^{-\frac{(x-\mu)^2}{2\sigma^2}}$
   2. $e^a \cdot \frac{e^b}{e^c} = e^{a+b-c}$
      • $q(x_t|x_{t-1},x_0) = \mathcal{N}(x_t; \sqrt{1-\beta_t}x_{t-1}, \beta_tI)$ (definition of $q$)
        • $\mu = \sqrt{1-\beta_t}x_{t-1} = \sqrt{a_t}x_{t-1}$
        • $\sigma = \sqrt{\beta_t}$
      • $q(x_{t}|x_0) = \mathcal{N}(x_{t}; \sqrt{\bar{a}_t}x_0, (1-\bar{a}_t)I)$ (reparameterization trick)
        • $\mu = \sqrt{\bar{a}_t}x_0$
        • $\sigma = \sqrt{1-\bar{a}_t}$
      • Same idea for $q(x_{t-1}|x_0)$

Let's examine part of the PDF of the gaussian more closely.

From this, we know:

One last note, $x$ is actually $x_{t-1}$ since that is the output of $q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \color{blue}{\tilde{\boldsymbol{\mu}}}(\mathbf{x}_t, \mathbf{x}_0), \color{red}{\tilde{\beta}_t} \mathbf{I})$. Now we can derive the mean and variance of $q(x_{t-1}|x_t,x_0)$:

Simplifying The Posterior Mean

From reparameterization we know:

Aside: Implementation

We can do a bit more algebra to derive the terms used in the OpenAI codebase:

---

We can use our definition of $x_0$ and substitute:

I could not figure out the algebra for this step. But, I did verify the results numerically. See these tests.

Closed Form KL-Divergence

We want to calculate the KL-divergence efficiently. Normally, we would need to Monte Carlo estimation. But, since $L_{t-1}$ is comparing 2 gaussians, we can use the closed form of the KL Divergence:

The authors ignore the 1st & last terms since they don't learn the variance for the reverse diffusion process.

• Question: Why bother including $\sigma_t$ in the denominator given that we're not learning? (My guess: resembles denoised score matching)

Extensions

DDIM

DDIM allows deterministic sampling without modifying the training process. Let's analyze the OpenAI code for DDIM sampling:

    def ddim_sample(
        self,
        model,
        x,
        t,
        clip_denoised=True,
        denoised_fn=None,
        model_kwargs=None,
        eta=0.0,
    ):
        """
        Sample x_{t-1} from the model using DDIM.
        Same usage as p_sample().
        """
        out = self.p_mean_variance(
            model,
            x,
            t,
            clip_denoised=clip_denoised,
            denoised_fn=denoised_fn,
            model_kwargs=model_kwargs,
        )
        # Usually our model outputs epsilon, but we re-derive it
        # in case we used x_start or x_prev prediction.
        eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])
        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
		# ===MARKER 1===
        sigma = (
            eta
            * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
            * th.sqrt(1 - alpha_bar / alpha_bar_prev)
        )
        # Equation 12.
        noise = th.randn_like(x)
		# ===MARKER 2===
        mean_pred = (
            out["pred_xstart"] * th.sqrt(alpha_bar_prev)
            + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps
        )
        nonzero_mask = (
            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
        )  # no noise when t == 0
        sample = mean_pred + nonzero_mask * sigma * noise
        return {"sample": sample, "pred_xstart": out["pred_xstart"]}

Marker 1

1. Utilize the facts that A. $\bar{a}_t$ is a product B. $a_t \coloneqq 1 - \beta_t$

TODO: Why isn't the OpenAI code just using square root of $\beta_t$ directly? My guess, float64/float32 numeric stuff.

Marker 2

mean_pred = (
    out["pred_xstart"] * th.sqrt(alpha_bar_prev)
    + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps
)
# Skip a few lines
sample = mean_pred + nonzero_mask * sigma * noise

Where $\epsilon_\theta$ is the random noise predicted by the network and $\epsilon$ is random noise we generate. Note, when $\text{eta} = 0$ , then this simplifies to:

which is a deterministic sampling process! In essence, instead of sampling from the predicted posterior mean/variance, we just use the reparameterization trick and stop 1-step earlier.