DDPM Diffusion Model Formula Reasoning----Loss Function

DDPM Diffusion Model Formula Reasoning----Loss Function

2.3 Loss function derivation

We construct the loss function with the idea of ​​maximum likelihood estimation:
$$L=-\log p_{}\left(x_0\right)$$ is the parameter θ \theta
of the inverse diffusion network$\theta$ makes the data x 0 x_0just started sampling$x_0$most likely to occur.

Next, we need to transform the above formula, using some ELBO and VAE content.

2.3.1 ELBA

Known maximum likelihood estimate $p_{}\left(x_0\right)$ and observation$x_0$, and x 1 obtained by diffusion of observations $x_{1:T}$, by the marginal probability distribution formula:
$$p_{}({\mathbf{x}}_0)=\int p_{}(x_{0:T})d{x}_{1:T}$$
therefore
$$\begin{array}{rl}\log p_{}({\mathbf{x}}_0)& =\log \int p_{}(x_{0:\mathbf{T}})d{\mathbf{x}}_{1:\mathbf{T}}\\ & =\log \int \frac{p_{}(x_{0:T})q_{}(x_{1:T}|x_0)}{q_{}(x_{1:T}\mid {\mathbf{x}}_0)}d{\mathbf{x}}_{1:T}\\ & =\log {\mathbb{E}}_{q_{}(x_{1:T}\mid {\mathbf{x}}_0)}\left[\frac{p_{}\left(x_{0:T}\right)}{q_{}(x_{1:T}|x_0)}\right]\\ & \ge E_{q_{}(x_{1:T}|x_0)}\left[\log \frac{p_{}\left(x_{0:T}\right)}{q_{}(x_{1:T}|x_0)}\right]\end{array}$$
The last step is determined by the piano sound inequality $(Jense{n}^{\prime }sInequality)$.
This does not seem very intuitive, and there is another way of derivation that is simpler:
$$\begin{array}{rl}\log p(\mathbf{x})& =\log p\left(x\right)\int q_{}(\mathbf{z}\mid x)dz\\ & =\int q_{}(z|\mathbf{x})(\log p(x))dz\\ & ={\mathbb{E}}_{q_{}(\mathbf{z}\mid \mathbf{x})}[\log p(\mathbf{x})]\\ & =E_{q_{}(z|x)}\left[\log \frac{p(\mathbf{x},\mathbf{z})}{p(\mathbf{z}|\mathbf{x})}\right]\\ & =E_{q_{}(z|x)}\left[\log \frac{p(x,\mathbf{z})q_{}(z|\mathbf{x})}{p(z|\mathbf{x})q_{}(z|x)}\right]\\ & =E_{q_{}(z|x)}\left[\log \frac{p(x,z)}{q_{}(z|x)}\right]+E_{q_{}(z|x)}\left[\log \frac{q_{}(z|x)}{p(z|x)}\right]\\ & =E_{q_{}(z|x)}\left[\log \frac{p(x,z)}{q_{}(z|x)}\right]+D_{\mathrm{KL}}\left(q_{}(z|\mathbf{x})\Vert p(\mathbf{z}|\mathbf{x})\right)\\ & \ge {\mathbb{E}}_{q_{}(\mathbf{z}\mid \mathbf{x})}\left[\log \frac{p(\mathbf{x},\mathbf{z})}{q_{}(\mathbf{z}\mid \mathbf{x})}\right]\\ & =E_{q_{}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\left[\log \frac{p_{}\left(x_{0:\mathbf{T}}\right)}{q_{}(x_{1:T}|{\mathbf{x}}_0)}\right]\end{array}$$
zz here$z$ stands for$x_{1:T}$，$x$ stands for$x_0$. Bayesian formula is used in the middle.

The conclusions derived by the two methods here are:
$$\log p_{}({\mathbf{x}}_0)\ge {\mathbb{E}}_{q_{}(x_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\left[\log \frac{p_{}\left(x_{0:\mathbf{T}}\right)}{q_{}(x_{1:T}\mid {\mathbf{x}}_0)}\right]$$

其中 ${\mathbb{E}}_{q_{}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\left[\log \frac{p_{}\left(x_{0:\mathbf{T}}\right)}{q_{}(x_{1:T}|x_0)}\right]$就是ELBO$(EvidenceLowerBound)\; ,thatis$ , the variational lower bound$.$

We want to make the loss function $-\log p_{}\left(x_0\right)$ 最小，就是另ELBO ${\mathbb{E}}_{q_{}(x_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\left[\log \frac{p_{}\left(x_{0:\mathbf{T}}\right)}{q_{}({\mathbf{x}}_{1:\mathbf{T}}|x_0)}\right]$ Max.

我们令 $L_{VL}=-{\mathbb{E}}_{q_{}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\left[\log \frac{p_{}\left(x_{0:\mathbf{T}}\right)}{q_{}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\right]={\mathbb{E}}_{q_{}(x_{1:T}|x_0)}\left[\log \frac{q_{}(x_{1:T}|x_0)}{p_{}\left(x_{0:T}\right)}\right]$ , and then convert to solving the loss function. Break it down further:
$$\begin{array}{rl}{L}_{}& ={\mathbb{E}}_{q\left({\mathbf{x}}_{0:T}\right)}\left[\log \frac{q\left(x_{1:T}\mid {\mathbf{x}}_0\right)}{p_{}\left({\mathbf{x}}_{0:T}\right)}\right]\\ & =E_q\left[\log \frac{{\prod }_{t=1}^{T}q\left(x_t|x_{t-1}\right)}{p_{}\left(x_T\right){\prod }_{t=1}^T{p}_{}\left(x_{t-1}|x_t\right)}\right]\\ & =E_q\left[-\log p_{}\left(x_T\right)+\sum _{t=1}^T\log \frac{q\left(x_t|x_{t-1}\right)}{p_{}\left(x_{t-1}|x_t\right)}\right]\\ & =E_q\left[-\log p_{}\left(x_T\right)+\sum _{t=2}^T\log \frac{q\left(x_t|x_{t-1}\right)}{p_{}\left(x_{t-1}|x_t\right)}+\log \frac{q\left(x_1|x_0\right)}{p_{}\left(x_0|x_1\right)}\right]\\ & =E_q\left[-\log p_{}\left(x_T\right)+\sum _{t=2}^T\log \left(\frac{q\left(x_{t-1}|x_t,x_0\right)}{p_{}\left(x_{t-1}|x_t\right)}\cdot \frac{q\left(x_t|x_0\right)}{q\left(x_{t-1}|x_0\right)}\right)+\log \frac{q\left(x_1|x_0\right)}{p_{}\left(x_0|x_1\right)}\right]\\ & =E_q\left[-\log p_{}\left(x_T\right)+\sum _{t=2}^T\log \frac{q\left(x_{t-1}|x_t,x_0\right)}{p_{}\left(x_{t-1}|x_t\right)}+\sum _{t=2}^T\log \frac{q\left(x_t|x_0\right)}{q\left(x_{t-1}|x_0\right)}+\log \frac{q\left(x_1|x_0\right)}{p_{}\left(x_0|x_1\right)}\right]\\ & =E_q\left[-\log p_{}\left(x_T\right)+\sum _{t=2}^T\log \frac{q\left(x_{t-1}|x_t,x_0\right)}{p_{}\left(x_{t-1}|x_t\right)}+\log \frac{q\left(x_T|x_0\right)}{q\left(x_1|x_0\right)}+\log \frac{q\left(x_1|x_0\right)}{p_{}\left(x_0|x_1\right)}\right]\\ & =E_q\left[\log \frac{q\left(x_T|x_0\right)}{p_{}\left(x_T\right)}+\sum _{t=2}^T\log \frac{q\left(x_{t-1}|x_t,x_0\right)}{p_{}\left(x_{t-1}|x_t\right)}-\log p_{}\left(x_0|x_1\right)\right]\\ & =E_q[\underset{L_T}{\underbrace{D_{\mathrm{KL}}\left(q\left(x_T|x_0\right)\Vert p_{}\left(x_T\right)\right)}}+\sum _{t=2}^T\underset{L_{t-1}}{\underbrace{D_{KL}\left(q\left(x_{t-1}|x_t,x_0\right)\Vert p_{}\left(x_{t-1}|x_t\right)\right)}}-\underset{L_0}{\underbrace{\log p_{}\left(x_0|x_1\right)}}]\end{array}$$

Here the initial subscript starts from $q_{}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)$ is changed to$q\left(x_{0:T}\right)$ , I think$x_0$is known, so the two expressions are equivalent.

There is another confusing point in the middle, $q\left(x_{1:T}|x_0\right)$ represents the distribution of forward propagation,$q\left(x_{t-1}|x_t\right)$ represents the true distribution of the inverse diffusion process, where$q$ has no forward or reverse meaning, it just represents probability and distribution, which can be understood as the probability PPin probability theory$P$。$p_{}\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)$ represents the inverse diffusion distribution we want to solve for.

Next we have $L_T,L_{t-1},L_0$These three situations are classified and discussed:

2.3.2 $L_T$

$q\left(x_{1:T}|x_0\right)$ represents the forward diffusion process, there is no learnable parameter;$p_{}(x_T$x T x_Tin$)$$x_T$is the noise that obeys the standard Gaussian distribution, $p_{}$is the inverse diffusion process, for the inverse diffusion process, $x_T$is known, so this term $L_T$Can be used as a constant.

2.3.3 $L_{t-1}$

And $L_{t-1}$It can be seen that the real inverse diffusion distribution $q\left({\mathbf{x}}_{t-1}\mid x_t,x_0\right)$ and the inverse diffusion distribution we require$p_{}\left(x_{t-1}|x_t\right)$ KL divergence.

1. The function $q\left(x_{t-1}|x_t,x_0\right)$ mean and variance we have obtained:
   $${\tilde{m}}_t=\frac{1}{\sqrt{a_t}}\left(x_t-\frac{1-a_t}{\sqrt{1-{\bar{a}}_t}}{e}_t\right),{\tilde{b}}_t=\frac{1-{\bar{a}}_{t-1}}{1-{\bar{a}}_t}\cdot b_t$$
2. The second distribution $p_{}\left(x_{t-1}\mid {\mathbf{x}}_t\right)$ is the target distribution we want to fit, it is also a Gaussian distribution, the mean is estimated by the network, and the variance is set to be$b_t$Form:
   $$p_{}\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)=\mathcal{N}\left(x_{t-1};m_{}\left(x_t,t\right),S_{}\left(x_t,t\right)\right)$$

Therefore, in order to make these two distributions close, we can ignore the variance, and only need to minimize the distance between the means of the two distributions. We use the second norm to express:
$$\begin{array}{rl}{L}_t& ={\mathbb{E}}_q\left[{\Vert {\tilde{m}}_t\left({\mathbf{x}}_t,{\mathbf{x}}_0\right)-m_{}\left(x_t,t\right)\Vert }^2\right]\\ & ={\mathbb{E}}_{x_0,}\left[{\Vert \frac{1}{\sqrt{a_t}}(x_t(x_0,)\_-\frac{b_t}{\sqrt{1-{\bar{a}}_t}})\_-m_{}\left(x_t(x_0,)\_,t\right)\Vert }^2\right]ϵ\sim \mathcal{N}(0,1)\end{array}$$
It can be observed in this formula that we need to use $m_{}\left(x_t(x_0,)\_,t\right)$ within1$\frac{1}{\sqrt{a_t}}\left(x_t(x_0,)\_-\frac{b_t}{\sqrt{1-{\bar{a}}_t}}ϵ\right)$ ,define:
$$m_{}\left({\mathbf{x}}_t,t\right)=\frac{1}{\sqrt{a_t}}\left(x_t-\frac{b_t}{\sqrt{1-{\bar{a}}_t}}{ϵ}_{}\left(x_t,t\right)\right)$$
is to directly use the neural network${ϵ}_{}\left(x_t,t\right)$ to predict the noise$ϵ$ . Then bring the predicted noise into the defined expression to calculate the predicted mean.

So the loss function becomes:
$$\begin{array}{rl}{L}_t& ={\mathbb{E}}_{{\mathbf{x}}_0,}\left[{\Vert \frac{1}{\sqrt{a_t}}({\mathbf{x}}_t-\frac{b_t}{\sqrt{1-{\bar{a}}_t}})\_-\frac{1}{\sqrt{a_t}}\left(x_t-\frac{b_t}{\sqrt{1-{\bar{a}}_t}}{ϵ}_{}\left(x_t,t\right)\right)\Vert }^2\right]ϵ\sim \mathcal{N}(0,1)\\ & =E_{x_0,}\left[{\Vert ϵ-{ϵ}_{}\left(x_t,t\right)\Vert }^2\right]ϵ\sim N(0,1)\text{ Throw away the coefficients of the constant items , the author said that it is better to train }\\ & =E_{x_0,}\left[{\Vert ϵ-{ϵ}_{}\left(\sqrt{{\bar{a}}_t}{x}_0+\sqrt{1-{\bar{a}}_t}ϵ,t\right)\Vert }^2\right],ϵ\sim N(0,1)\end{array}$$
The input to the network is a picture xt x_t linearly combined with noise$x_t$, the true value of the noise combined with it is $ϵ$ ,the free function${ϵ}_{}\left(\sqrt{{\bar{a}}_t}{x}_0+\sqrt{1-{\bar{a}}_t}ϵ,t\right)$ to fit this noise.

2.3.4 $L_0$

The final $L_0=-\log p_{}\left(x_0\mid x_1\right)$ is the noised image of the last step$x_1$Generate denoising image $x_0$The maximum likelihood estimation of , in order to generate a better image, we need to use the maximum likelihood estimation for each pixel, so that each pixel value on the image satisfies the discrete log likelihood.

To achieve this, the last part of the inverse diffusion process is changed from $x_1$to $x_0$The transformations are set to independent discrete calculations. That is, in the last conversion process at a given $x_1$get image $x_0$Satisfy the log likelihood, assuming that the pixels are independent of each other:
$$p_{}\left(x_0\mid x_1\right)=\prod _{i=1}^D{p}_{}\left(x_0^i|x_1^i\right)$$
$D$ is for$dimension\; of\; x$ , superscript$i$ represents a coordinate position in the image. The goal now is to determine how likely the value of a given pixel is, that is, to know the corresponding time step$t=1$ lower noise imagexxDistribution of corresponding pixel values ​​in$x$
$$\mathcal{N}\left(x;m_i^i\left(x_1,1\right),p_1^2\right)$$
where$t=The\; pixel\; distribution\; of\; 1$ comes from a multivariate Gaussian distribution whose diagonal covariance matrix allows us to split the distribution into a product of univariate Gaussians: N (
$$\mathcal{N}\left(x;m_{}\left(x_1,1\right),p_1^2\mathbb{I}\right)=\prod _{i=1}^{D}N\left(x;m_i^i\left(x_1,1\right),p_1^2\right)$$
Now assume that the image has been normalized in the range [-1,1] from the value of 0-255. Given the pixel value of each pixel at t=0, the transition probability distribution$p_{}\left(x_0\mid x_1\right)$ 的值就是每个像素值的乘积。所以：
$$\begin{array}{rl}{p}_{}\left({\mathbf{x}}_0\mid {\mathbf{x}}_1\right)& =\prod _{i=1}^D{\int }_{d_{-}\left(x_0^i\right)}^{d_{+}\left(x_0^i\right)}\mathcal{N}\left(x;m_i^i\left(x_1,1\right),p_1^2\right)dx\\ d_{+}(x)& =\{\begin{array}{ll}\infty & \text{ if }x=1\\ x+\frac{1}{255}& \text{ if }x<1\end{array}{d}_{-}(x)=\{\begin{array}{ll}-\infty & \text{ if }x=-1\\ x-\frac{1}{255}& \text{ if }x>-1\end{array}\end{array}$$
This formula comes from the original paper, here is an analysis of its meaning. That is, we want to add the noise image of the last step $x_1$Fit denoised image $x_0$, set each pixel of the image to a Gaussian distribution, a total of $D$ pixels. while$x_0$The original value range of each pixel is { 0 , 1 , … , 255 } \{0,1, \ldots, 255\}{ 0,1,…,255 } , mapped to [ − 1 , 1 ] [-1,1]after normalization$[-1,1]$ range.

Now we take out a single $x_1$Pixels on $x_1^i$,given the function $\mathcal{N}\left(x;m_i^i\left(x_1,1\right),p_1^2\right)$ , the target to be fitted is$x_0$The corresponding position pixel point $x_0^i$, and $x_0^i$The value range of the original discrete space { 0 , 1 , … , 255 } \{0,1, \ldots, 255\}{ 0,1,…,255 } is mapped to the continuous space$[-1,1]$ , so each original discrete value corresponds to an interval in the continuous space, and the formula for interval mapping is:
$$\begin{array}{rl}{d}_{+}(x)& =\{\begin{array}{ll}\infty & \text{ if }x=1\\ x+\frac{1}{255}& \text{ if }x<1\end{array}{d}_{-}(x)=\{\begin{array}{ll}-\infty & \text{ if }x=-1\\ x-\frac{1}{255}& \text{ if }x>-1\end{array}\end{array}$$

The above is the analysis and derivation of all the formulas involved in the diffusion and inverse diffusion process of DDPM, including the loss function construction part.

References and blog

Understanding Diffusion Models: A Unified Perspective
Denoising Diffusion Probabilistic Models
https://yinglinzheng.netlify.app/diffusion-model-tutorial
https://zhuanlan.zhihu.com/p/549623622