Diffusion, a class of generative models, boils down to an MSE between added noise and network predicted noise.
Why would learning to predict the noise help with image generation (which is what it is most used for)? How did we arrive at MSE? This post dives deep into the math to answer these questions.

Background

One way to interpret Diffusion is as a continuous VAE (Variational Auto Encoder).
A VAE computes a lower bound on the likelihood of generating real data samples ( $log p_{θ} (x)$ ) by approximating the unknown posterior $p_{θ} (z | x)$ with a learnable one $q_{ϕ} (z | x)$ [Fig. 1]:

$\begin{matrix} - log p_{θ} (x) & = - log p_{θ} (x) \leq - log p_{θ} (x) + D_{K L} (q_{ϕ} (z | x) ∥ p_{θ} (z | x)) & [KL is always positive] \leq - log p_{θ} (x) + \int q_{ϕ} (z | x) log (\frac{q_{ϕ} (z | x)}{p_{θ} (z | x)}) d z & [Definition of KL] \leq - log p_{θ} (x) + \int q_{ϕ} (z | x) log (\frac{q_{ϕ} (z | x) p_{θ} (x)}{p_{θ} (z, x)}) d z & [conditional to joint] \leq - log p_{θ} (x) + \int q_{ϕ} (z | x) (log p_{θ} (x) + log (\frac{q_{ϕ} (z | x)}{p_{θ} (z, x)})) d z \leq - log p_{θ} (x) + log p_{θ} (x) + \int q_{ϕ} (z | x) (log (\frac{q_{ϕ} (z | x)}{p_{θ} (x | z) p_{θ} (z)})) d z & [p_{θ} independent of z; joint to conditional] = - E_{z_{\sim} q_{ϕ} (z | x)} [log (\frac{q_{ϕ} (z | x)}{p_{θ} (z)}) - log p_{θ} (x | z)] & [Definition of E for continuous variable z] = - E_{z_{\sim} q_{ϕ} (z | x)} log p_{θ} (x | z) + D_{K L} (q_{ϕ} (z | x) ∥ p_{θ} (z)) & [Tractable] \end{matrix}$

The process of encoding ( $q_{ϕ} (z | x)$ ) and decoding ( $p_{θ} (x | z)$ ) leads to regularized compression of images ( $x$ ) in the $z$ space (hence the name auto encoder; variational comes from approximating distributions). During inference, random images can be generated from the decoder, based on sampling of $z$ .

Derivation

Figure 2: Markov chain of forward [ $q$ ] (reverse [ $p_{θ}$ ]) diffusion process of generating a sample by slowly adding (removing) noise. (Image source: Ho et al. 2020)

Similar to VAE, diffusion models compute a lower bound on the likelihood of generating real data samples ( $log p_{θ} (x_{0})$ ) by approximating the unknown posterior $p_{θ} (x_{t - 1} | x_{t})$ with a computable one $q_{ϕ} (x_{t - 1} | x_{t}, x_{0})$ , for all $t$ [Fig. 2] (It follows the same math as VAE for all $t$ , therefore could be thought of as its continuous version):

$\begin{matrix} - log p_{θ} (x_{0}) & \leq - log p_{θ} (x_{0}) \leq - log p_{θ} (x_{0}) + D_{K L} (q (x_{1 : T} | x_{0}) ∥ p_{θ} (x_{1 : T} | x_{0})) & [KL is always positive] = - log p_{θ} (x_{0}) + E_{x_{1 : T} \sim q (x_{1 : T} | x_{0})} [log (\frac{q (x_{1 : T} | x_{0})}{p_{θ} (x_{1 : T}) | x_{0}})] = - log p_{θ} (x_{0}) + E_{x_{1 : T} \sim q (x_{1 : T} | x_{0})} ⎡ ⎢ ⎢ ⎣ log ⎛ ⎜ ⎜ ⎝ \frac{q (x_{1 : T} | x_{0})}{(\frac{p_{θ} (x_{0} | x_{1 : T}) p_{θ} (x_{1 : T})}{p_{θ} (x_{0})})} ⎞ ⎟ ⎟ ⎠ ⎤ ⎥ ⎥ ⎦ & [Bayes' rule] = - log p_{θ} (x_{0}) + E_{x_{1 : T} \sim q (x_{1 : T} | x_{0})} ⎡ ⎢ ⎢ ⎣ log ⎛ ⎜ ⎜ ⎝ \frac{q (x_{1 : T} | x_{0})}{(\frac{p_{θ} (x_{0}, x_{1 : T})}{p_{θ} (x_{0})})} ⎞ ⎟ ⎟ ⎠ ⎤ ⎥ ⎥ ⎦ & [conditional to joint] = - log p_{θ} (x_{0}) + E_{x_{1 : T} \sim q (x_{1 : T} | x_{0})} ⎡ ⎢ ⎢ ⎣ log ⎛ ⎜ ⎜ ⎝ \frac{q (x_{1 : T} | x_{0})}{(\frac{p_{θ} (x_{0 : T})}{p_{θ} (x_{0})})} ⎞ ⎟ ⎟ ⎠ ⎤ ⎥ \end{matrix}$ ...

Posts

Wikitag Contributions

Comments