Diffusion Models

Links: https://theaisummer.com/diffusion-models/

Markovian Hierachical VAE

rvs:

• data: $x_0$,
• representation: $x_T$

$$(p(x_0,x_1,\cdots ,x_T),q(x_1,\cdots ,x_T|x_0))$$
where $x_1,\cdots ,x_T$ is unobservable, and

• generative model/backward trajectory:
  $$p(x_0,x_1,\cdots ,x_T)=p(x_T)\prod _{t}p(x_{t-1}|x_t)$$
• forward trajectory(Markov process):
  $$q(x_1,\cdots ,x_T|x_0))=\prod _{t}q(x_t|x_{t-1})$$

$$\begin{gathered} ELBO:=\int q(x_T|x_0)\log \frac{p(x_T)}{q(x_T|x_0)}\mathrm{d}{x}_T \\ +\sum _{t=2}^T\int q(x_{t-1},x_t|x_0)\log \frac{p(x_{t-1}|x_t)}{q(x_{t-1}|x_t,x_0)}\mathrm{d}{x}_{t-1}{x}_t \\ +\int q(x_1|x_0)\log p(x_1|x_0)\mathrm{d}{x}_1 \end{gathered}$$

Loss

$$\begin{gathered} Loss:=-ELBO=D_{KL}(q(x_T|x_0)\Vert p(x_T)) \\ +\sum _{t=2}^T\int q(x_t|x_0)\mathrm{d}{x}_t{D}_{KL}(q(x_{t-1}|x_t,x_0)\Vert p(x_{t-1}|x_t)) \\ -\int q(x_1|x_0)\log p(x_1|x_0)\mathrm{d}{x}_1 \end{gathered}$$

• prior matching term
• denoising matching term
• reconstruction term

Diffusion Models

basic assumption

• tractable distr: $p(x_T)$
• forward trajectory(Markov process): $q(x_t|x_{t-1})$ is fixed (has no unlearned parameter)

Definition(Diffusion Model)

• tractable distr: $p(x_T)\sim N(0,1)$
• generative model/backward trajectory: $p(x_{t-1}|x_t)\sim N(\mu (t),\Sigma (t))$
• forward trajectory(Gaussian diffusion): $q(x_t|x_{t-1})\sim N(x_{t-1}\sqrt{1-{\beta }_t},{\beta }_t)$,

Parameters:

• ${\beta }_t=1-{\alpha }_t$ or ${\bar{\alpha }}_t:={\prod }_t{\alpha }_t$: noise schedule, where ${\alpha }_t$ is small
• $\sqrt{{\bar{\alpha }}_t}$: signal rate

Fact.

• $q(x_t|x_0)\sim N(x_0\sqrt{{\bar{\alpha }}_t},1-{\bar{\alpha }}_t)$
• $q(x_{t-1}|x_t,x_0)\sim N({\mu }_q(x_t,x_0),{\sigma }^2(t))$ where
  $$\begin{gathered} {\mu }_q(x_t,x_0):=\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})x_t-\sqrt{{\bar{\alpha }}_{t-1}}(1-{\alpha }_t)x_0}{1-{\bar{\alpha }}_t} \\ =\frac{1}{\sqrt{{\alpha }_t}}{x}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}\sqrt{{\alpha }_t}}{ϵ}_0 \end{gathered}$$
  and ${\sigma }^2(t):=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t$.

Design I: $p(x_{t-1}|x_t)\sim N(\mu (t),\Sigma (t))$:
$$\begin{gathered} \mu (t)=\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})x_t-{\beta }_t\sqrt{{\bar{\alpha }}_{t-1}}\hat{x}(x_t,t)}{1-{\bar{\alpha }}_t} \\ \Sigma (t)={\sigma }^2(t) \end{gathered}$$

Design II: $p(x_{t-1}|x_t)\sim N(\mu (t),\Sigma (t))$:
$$\begin{gathered} \mu (t)=\frac{1}{\sqrt{{\alpha }_t}}{x}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}\sqrt{{\alpha }_t}}\widehat{ϵ}(x_t,t) \\ \Sigma (t)={\sigma }^2(t) \end{gathered}$$

Fact.
Under the design I:
$$\begin{gathered} D_{KL}(q(x_{t-1}|x_t,x_0)\Vert p_{\theta }(x_{t-1}|x_t))=\frac{1}{2{\sigma }_t^2}\frac{(1-{\bar{\alpha }}_{t-1}){\beta }_t^2}{(1-{\bar{\alpha }}_t{)}^2}\Vert \hat{x}(x_t,t)-x_0{\Vert }^2 \\ =\frac{1}{2}(\frac{1}{1-{\bar{\alpha }}_{t-1}}-\frac{1}{1-{\bar{\alpha }}_t})\Vert \hat{x}(x_t,t)-x_0{\Vert }^2 \end{gathered}$$

Under the design II:
$$D_{KL}(q(x_{t-1}|x_t,x_0)\Vert p_{\theta }(x_{t-1}|x_t))=\frac{1}{2{\sigma }_t^2}\frac{{\beta }_t^2}{(1-{\bar{\alpha }}_t){\alpha }_t^2}\Vert \widehat{ϵ}(x_t,t)-{ϵ}_0{\Vert }^2$$

Algorithm

Loss:
$$\begin{gathered} L=\sum _t{L}_t \\ {L}_t\approx \sum _{ϵ\sim N(0,1)}\Vert ϵ-\widehat{ϵ}(x_t,t){\Vert }^2,(0\le t<T) \end{gathered}$$
where $x_t:=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ$.

train NN $\widehat{ϵ}$ by data $\{(\widehat{ϵ}(x_t(x_{0,i},{ϵ}_{il}),t),{ϵ}_{il}),{ϵ}_{il}\sim N(0,1),l=1,\cdots ,L\}$ with size of $NL$ for each $t$。

---

Exercise

1. Given a latent variable model $p(x,z)$ with variational distr. $q(z|x)$. $q(x)$ represents data distr. and let $q(x,z)=q(z|x)q(x)$.
   $$\int q(x)L_x=\int q(x,z)\log \frac{p(x,z)}{q(z|x)}\sim D_{KL}(q(x,z)\Vert p(x,z))$$
   where $L_x$ is LEBO.

References

1. Jonathan Ho, Ajay Jain, Pieter Abbeel. Denoising Diffusion Probabilistic Models, 2020.
2. Calvin Luo, Understanding Diffusion Models: A Unified Perspective, 2022