论文阅读笔记：Denoising Diffusion Probabilistic Models (3)

论文阅读笔记：Denoising Diffusion Probabilistic Models (1)
论文阅读笔记：Denoising Diffusion Probabilistic Models (2)
论文阅读笔记：Denoising Diffusion Probabilistic Models (3)

4、损失函数逐项分析

可以看出$L$总共分为了3项，首先考虑第一项$L_1$。
$$\begin{array}{c}\begin{array}{rl}{L}_1& =E_{x_{1:T}\sim q(x_{1:T}|x_0)}\left(log\right[\frac{q(x_T|x_0)}{p(x_T)}\left]\right)\\ & =\int d{x}_{1:T}\cdot q(x_{1:T}|x_0)\cdot log[\frac{q(x_T|x_0)}{p(x_T)}]\\ & =\int d{x}_{1:T}\cdot \frac{q(x_{1:T}|x_0)}{q(x_T|x_0)}\cdot q(x_T|x_0)\cdot log[\frac{q(x_T|x_0)}{p(x_T)}]\\ & =\int d{x}_{1:T}\cdot \underset{q(x_{1:T}|x_0)=q(x_T|x_0)\cdot q(x_{1;T-1}|x_0,x_T)}{\underbrace{q(x_{1:T-1}|x_0,x_T)}}\cdot q(x_T|x_0)\cdot log[\frac{q(x_T|x_0)}{p(x_T)}]\\ & =\int \left(\underset{二重积分化为两个定积分相乘，并且=1}{\underbrace{\int q(x_{1:T-1}|x_0,x_T)\cdot \prod _{k=1}^{T-1}d{x}_k}}\right)\cdot q(x_T|x_0)\cdot log[\frac{q(x_T|x_0)}{p(x_T)}]\cdot d{x}_T\\ & =\int q(x_T|x_0)\cdot log[\frac{q(x_T|x_0)}{p(x_T)}]\cdot d{x}_T\\ & =E_{x^T\sim q(x_T|x_0)}log\left[\frac{q(x_T|x_0)}{p(x_T)}\right]\\ & =KL(q(x_T|x_0)||p(x_T))\end{array}\end{array}$$

可以看出，$L_1$是$q(x_T|x_0)$和$p(x_T)$的散度。$q(x_T|x_0)$是前向加噪过程的终点，是无限趋向于标准正态分布。而$p(x_T)$是高斯分布，这在论文《Denoising Diffusion Probabilistic Models》中的2 Background的第四行中有说明。由 两个高斯分布KL散度推导可以计算出$L_1$，也就是说$L_1$是一个定值。因此，在损失函数中$L_1$可以被忽略掉。

接着考虑第二项$L_2$。

$$\begin{array}{c}\begin{array}{rl}{L}_2& =E_{x_{1:T}\sim q(x_{1:T}|x_0)}\left(\sum _{t=2}^{T}log\right[\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}\left]\right)\\ & =\sum _{t=2}^T{E}_{x_{1:T}\sim q(x_{1:T}|x_0)}\left(log\right[\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}\left]\right)\\ & =\sum _{t=2}^T(\int d{x}_{1:T}\cdot q(x_{1:T}|x_0)\cdot log\left[\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}\right])\\ & =\sum _{t=2}^T(\int d{x}_{1:T}\cdot \frac{q(x_{1:T}|x_0)}{q(x_{t-1}|x_t,x_0)}\cdot q(x_{t-1}|x_t,x_0)\cdot log\left[\frac{q(x_{t-1}|x_t,x_0)}{p(x_{t-1}|x_t)}\right])\\ & =\sum _{t=2}^T(\int d{x}_{1:T}\cdot \underset{q(x_{0:T})=q(x_0)\cdot q(x_{1:T}|x_0)}{\underbrace{\frac{q(x_{0:T})}{q(x_0)}}}\cdot \underset{q(x_t,x_{t-1},x_0)=q(x_t,x_0)\cdot q(x_{t-1}|x_t,x_0)}{\underbrace{\frac{q(x_t,x_0)}{q(x_t,x_{t-1},x_0)}}}\cdot q(x_{t-1}|x_t,x_0)\cdot log\left[\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}\right])\\ & =\sum _{t=2}^T(\int d{x}_{1:T}\cdot \frac{q(x_{0:T})}{q(x_0)}\cdot \frac{q(x_t,x_0)}{q(x_{t-1},x_0)\cdot q(x_t|x_{t-1},x_0)}\cdot q(x_{t-1}|x_t,x_0)\cdot log\left[\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}\right])\\ & =\sum _{t=2}^T(\int [\int \frac{q(x_{0:T})}{q(x_0)}\cdot \frac{q(x_t,x_0)}{q(x_{t-1},x_0)\cdot q(x_t|x_{t-1},x_0)}\prod _{k\ge 1,k\ne t-1}d{x}_k]\cdot q(x_{t-1}|x_t,x_0)\cdot log\left[\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}d{x}_{t-1}\right])\\ & =\sum _{t=2}^T(\int [\int \frac{q(x_{0:T})}{q(x_{t-1},x_0)}\cdot \frac{q(x_t,x_0)}{q(x_0)\cdot q(x_t|x_{t-1},x_0)}\prod _{k\ge 1,k\ne t-1}d{x}_k]\cdot q(x_{t-1}|x_t,x_0)\cdot log\left[\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}d{x}_{t-1}\right])\\ & =\sum _{t=2}^T(\int [\underset{q(x_{0;T})=q(x_{t-1},x_0)\cdot q(x_{k:k\ge 1,k\ne t-1}|x_{t-1},x_0)}{\underbrace{\int q(x_{k:k\ge 1,k\ne t-1}|x_{t-1},x_0)}}\cdot \underset{q(x_t,x_0)=q(x_0)\cdot q(x_t|x_0)}{\underbrace{\frac{q(x_t|x_0)}{q(x_t|x_{t-1},x_0)}}}\prod _{k\ge 1,k\ne t-1}d{x}_k]\cdot q(x_{t-1}|x_t,x_0)\cdot log\left[\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}d{x}_{t-1}\right])\\ & =\sum _{t=2}^T(\int [\int q(x_{k:k\ge 1,k\ne t-1}|x_{t-1},x_0)\cdot \underset{=1}{\underbrace{\frac{q(x_t|x_0)}{q(x_t|x_{t-1},x_0)}}}\prod _{k\ge 1,k\ne t-1}d{x}_k]\cdot q(x_{t-1}|x_t,x_0)\cdot log[\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}d{x}_{t-1}\left]\right)\\ & =\sum _{t=2}^T(\int [\int q(x_{k:k\ge 1,k\ne t-1}|x_{t-1},x_0)\cdot \prod _{k\ge 1,k\ne t-1}d{x}^k]\cdot q(x_{t-1}|x_t,x_0)\cdot log[\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}d{x}_{t-1}\left]\right)\\ & =\sum _{t=2}^T(\int [\underset{=1}{\underbrace{\int q(x_{k:k\ge 1,k\ne t-1}|x_{t-1},x^0)\cdot \prod _{k\ge 1,k\ne t-1}d{x}_k}}]\cdot q(x_{t-1}|x_t,x_0)\cdot log\left[\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}d{x}_{t-1}\right])\\ & =\sum _{t=2}^T(\int q(x_{t-1}|x_t,x_0)\cdot log\left[\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}d{x}_{t-1}\right])\\ & =\sum _{t=2}^T\left(E_{x_{t-1}\sim q(x_{t-1}|x_t,x_0)}log\right[\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}\left]\right)\\ & =\sum _{t=2}^{T}KL(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))\end{array}\end{array}$$
最后考虑$L_3$，事实上，在论文《Deep Unsupervised Learning using Nonequilibrium Thermodynamics》中提到为了防止边界效应，强制另$p(x^0|x^1)=q(x^1|x^0)$，因此这一项也是个常数。

由以上分析可知道，损失函数可以写为公式（3）。
$$\begin{array}{c}\begin{array}{rl}L& :=L_1+L_2+L_3\\ & =KL(q(x_T|x_0)||p(x_T))+\sum _{t=2}^{T}KL(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))-log[p_{\theta }(x_0|x_1)]\end{array}\end{array}$$

忽略掉$L_1$和$L_3$，损失函数可以写为公式（4）。
$$\begin{array}{c}\begin{array}{r}L:=\sum _{t=2}^{T}KL(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))\end{array}\end{array}$$

可以看出 损失函数$L$是两个高斯分布$q(x_{t-1}|x_t,x_0)$和$p_{\theta }(x_{t-1}|x_t)$的KL散度。$q(x_{t-1}|x_t,x_0)$的均值和方差由论文阅读笔记：Denoising Diffusion Probabilistic Models (1)可知，分别为

$$\begin{array}{c}\begin{array}{rl}{\sigma }_1& =\sqrt{\frac{{\beta }_t\cdot (1-\overline{{\alpha }_{t-1}})}{(1-\overline{{\alpha }_t})}}\\ {\mu }_1& =\frac{1}{\sqrt{{\alpha }_t}}\cdot (x_t-\frac{{\beta }_t}{\sqrt{1-\overline{{\alpha }_t}}}\cdot z_t)\\ 或者{\mu }_1& =\frac{\sqrt{{\alpha }_t}\cdot (1-\overline{{\alpha }_{t-1}})}{1-\overline{{\alpha }_t}}\cdot x_t+\frac{{\beta }_t\cdot \sqrt{\overline{{\alpha }_{t-1}}}}{1-\overline{{\alpha }_t}}\cdot x_0\end{array}\end{array}$$

而$p_{\theta }(x_{t-1}|x_t)$则由模型（深度学习模型或者其他模型）估算出其均值和方差，分别记作${\mu }_2,{\sigma }_2$。
因此损失函数$L$可以进一步写为公式12。
$$\begin{array}{c}\begin{array}{r}L:=log\left[\frac{{\sigma }_2}{{\sigma }_1}\right]+\frac{{\sigma }_1^2+({\mu }_1-{\mu }_2{)}^2}{2{\sigma }_2^2}-\frac{1}{2}\end{array}\end{array}$$

5、代码解析

最后结合原文中的代码diffusion-https://github.com/hojonathanho/diffusion来理解一下训练过程和推理过程。
首先是训练过程

class GaussianDiffusion2:"""Contains utilities for the diffusion model.Arguments:- what the network predicts (x_{t-1}, x_0, or epsilon)- which loss function (kl or unweighted MSE)- what is the variance of p(x_{t-1}|x_t) (learned, fixed to beta, or fixed to weighted beta)- what type of decoder, and how to weight its loss? is its variance learned too?"""# 模型中的一些定义def __init__(self, *, betas, model_mean_type, model_var_type, loss_type):self.model_mean_type = model_mean_type  # xprev, xstart, epsself.model_var_type = model_var_type  # learned, fixedsmall, fixedlargeself.loss_type = loss_type  # kl, mseassert isinstance(betas, np.ndarray)self.betas = betas = betas.astype(np.float64)  # computations here in float64 for accuracyassert (betas > 0).all() and (betas <= 1).all()timesteps, = betas.shapeself.num_timesteps = int(timesteps)alphas = 1. - betasself.alphas_cumprod = np.cumprod(alphas, axis=0)self.alphas_cumprod_prev = np.append(1., self.alphas_cumprod[:-1])assert self.alphas_cumprod_prev.shape == (timesteps,)# calculations for diffusion q(x_t | x_{t-1}) and othersself.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)self.sqrt_one_minus_alphas_cumprod = np.sqrt(1. - self.alphas_cumprod)self.log_one_minus_alphas_cumprod = np.log(1. - self.alphas_cumprod)self.sqrt_recip_alphas_cumprod = np.sqrt(1. / self.alphas_cumprod)self.sqrt_recipm1_alphas_cumprod = np.sqrt(1. / self.alphas_cumprod - 1)# calculations for posterior q(x_{t-1} | x_t, x_0)self.posterior_variance = betas * (1. - self.alphas_cumprod_prev) / (1. - self.alphas_cumprod)# below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chainself.posterior_log_variance_clipped = np.log(np.append(self.posterior_variance[1], self.posterior_variance[1:]))self.posterior_mean_coef1 = betas * np.sqrt(self.alphas_cumprod_prev) / (1. - self.alphas_cumprod)self.posterior_mean_coef2 = (1. - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1. - self.alphas_cumprod)# 在模型Model类当中的方法def train_fn(self, x, y):B, H, W, C = x.shapeif self.randflip:x = tf.image.random_flip_left_right(x)assert x.shape == [B, H, W, C]# 随机生成第t步t = tf.random_uniform([B], 0, self.diffusion.num_timesteps, dtype=tf.int32)# 计算第t步时对应的损失函数losses = self.diffusion.training_losses(denoise_fn=functools.partial(self._denoise, y=y, dropout=self.dropout), x_start=x, t=t)assert losses.shape == t.shape == [B]return {'loss': tf.reduce_mean(losses)}# 根据x_start采样到第t步的带噪图像def q_sample(self, x_start, t, noise=None):"""Diffuse the data (t == 0 means diffused for 1 step)"""if noise is None:noise = tf.random_normal(shape=x_start.shape)assert noise.shape == x_start.shapereturn (self._extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +self._extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)# 计算q(x^{t-1}|x^t,x^0)分布的均值和方差def q_posterior_mean_variance(self, x_start, x_t, t):"""Compute the mean and variance of the diffusion posterior q(x_{t-1} | x_t, x_0)"""assert x_start.shape == x_t.shapeposterior_mean = (self._extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +self._extract(self.posterior_mean_coef2, t, x_t.shape) * x_t)posterior_variance = self._extract(self.posterior_variance, t, x_t.shape)posterior_log_variance_clipped = self._extract(self.posterior_log_variance_clipped, t, x_t.shape)assert (posterior_mean.shape[0] == posterior_variance.shape[0] == posterior_log_variance_clipped.shape[0] ==x_start.shape[0])return posterior_mean, posterior_variance, posterior_log_variance_clipped# 由深度学习模型UNet估算出p(x^{t-1}|x^t)分布的方差和均值def p_mean_variance(self, denoise_fn, *, x, t, clip_denoised: bool, return_pred_xstart: bool):B, H, W, C = x.shapeassert t.shape == [B]model_output = denoise_fn(x, t)# Learned or fixed variance?if self.model_var_type == 'learned':assert model_output.shape == [B, H, W, C * 2]model_output, model_log_variance = tf.split(model_output, 2, axis=-1)model_variance = tf.exp(model_log_variance)elif self.model_var_type in ['fixedsmall', 'fixedlarge']:# below: only log_variance is used in the KL computationsmodel_variance, model_log_variance = {# for fixedlarge, we set the initial (log-)variance like so to get a better decoder log likelihood'fixedlarge': (self.betas, np.log(np.append(self.posterior_variance[1], self.betas[1:]))),'fixedsmall': (self.posterior_variance, self.posterior_log_variance_clipped),}[self.model_var_type]model_variance = self._extract(model_variance, t, x.shape) * tf.ones(x.shape.as_list())model_log_variance = self._extract(model_log_variance, t, x.shape) * tf.ones(x.shape.as_list())else:raise NotImplementedError(self.model_var_type)# Mean parameterization_maybe_clip = lambda x_: (tf.clip_by_value(x_, -1., 1.) if clip_denoised else x_)if self.model_mean_type == 'xprev':  # the model predicts x_{t-1}pred_xstart = _maybe_clip(self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output))model_mean = model_outputelif self.model_mean_type == 'xstart':  # the model predicts x_0pred_xstart = _maybe_clip(model_output)model_mean, _, _ = self.q_posterior_mean_variance(x_start=pred_xstart, x_t=x, t=t)elif self.model_mean_type == 'eps':  # the model predicts epsilonpred_xstart = _maybe_clip(self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output))model_mean, _, _ = self.q_posterior_mean_variance(x_start=pred_xstart, x_t=x, t=t)else:raise NotImplementedError(self.model_mean_type)assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shapeif return_pred_xstart:return model_mean, model_variance, model_log_variance, pred_xstartelse:return model_mean, model_variance, model_log_variance# 损失函数的计算过程def training_losses(self, denoise_fn, x_start, t, noise=None):assert t.shape == [x_start.shape[0]]# 随机生成一个噪音if noise is None:noise = tf.random_normal(shape=x_start.shape, dtype=x_start.dtype)assert noise.shape == x_start.shape and noise.dtype == x_start.dtype# 将随机生成的噪音加到x_start上得到第t步的带噪图像x_t = self.q_sample(x_start=x_start, t=t, noise=noise)# 有两种损失函数的方法，'kl'和'mse'，并且这两种方法差别并不明显。if self.loss_type == 'kl':  # the variational boundlosses = self._vb_terms_bpd(denoise_fn=denoise_fn, x_start=x_start, x_t=x_t, t=t, clip_denoised=False, return_pred_xstart=False)elif self.loss_type == 'mse':  # unweighted MSEassert self.model_var_type != 'learned'target = {'xprev': self.q_posterior_mean_variance(x_start=x_start, x_t=x_t, t=t)[0],'xstart': x_start,'eps': noise}[self.model_mean_type]model_output = denoise_fn(x_t, t)assert model_output.shape == target.shape == x_start.shapelosses = nn.meanflat(tf.squared_difference(target, model_output))else:raise NotImplementedError(self.loss_type)assert losses.shape == t.shapereturn losses# 计算两个高斯分布的KL散度，代码中的logvar1，logvar2为方差的对数def normal_kl(mean1, logvar1, mean2, logvar2):"""KL divergence between normal distributions parameterized by mean and log-variance."""return 0.5 * (-1.0 + logvar2 - logvar1 + tf.exp(logvar1 - logvar2)+ tf.squared_difference(mean1, mean2) * tf.exp(-logvar2))# 使用'kl'方法计算损失函数def _vb_terms_bpd(self, denoise_fn, x_start, x_t, t, *, clip_denoised: bool, return_pred_xstart: bool):true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(x_start=x_start, x_t=x_t, t=t)model_mean, _, model_log_variance, pred_xstart = self.p_mean_variance(denoise_fn, x=x_t, t=t, clip_denoised=clip_denoised, return_pred_xstart=True)kl = normal_kl(true_mean, true_log_variance_clipped, model_mean, model_log_variance)kl = nn.meanflat(kl) / np.log(2.)decoder_nll = -utils.discretized_gaussian_log_likelihood(x_start, means=model_mean, log_scales=0.5 * model_log_variance)assert decoder_nll.shape == x_start.shapedecoder_nll = nn.meanflat(decoder_nll) / np.log(2.)# At the first timestep return the decoder NLL, otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))assert kl.shape == decoder_nll.shape == t.shape == [x_start.shape[0]]output = tf.where(tf.equal(t, 0), decoder_nll, kl)return (output, pred_xstart) if return_pred_xstart else output

接下来是推理过程。

def p_sample(self, denoise_fn, *, x, t, noise_fn, clip_denoised=True, return_pred_xstart: bool):"""Sample from the model"""# 使用深度学习模型，根据x^t和t估算出x^{t-1}的均值和分布model_mean, _, model_log_variance, pred_xstart = self.p_mean_variance(denoise_fn, x=x, t=t, clip_denoised=clip_denoised, return_pred_xstart=True)noise = noise_fn(shape=x.shape, dtype=x.dtype)assert noise.shape == x.shape# no noise when t == 0nonzero_mask = tf.reshape(1 - tf.cast(tf.equal(t, 0), tf.float32), [x.shape[0]] + [1] * (len(x.shape) - 1))# 当t>0时，模型估算出的结果还要加上一个高斯噪音，因为要继续循环。当t=0时，循环停止，因此不需要再添加噪音了，输出最后的结果。sample = model_mean + nonzero_mask * tf.exp(0.5 * model_log_variance) * noiseassert sample.shape == pred_xstart.shapereturn (sample, pred_xstart) if return_pred_xstart else sampledef p_sample_loop(self, denoise_fn, *, shape, noise_fn=tf.random_normal):"""Generate samples"""assert isinstance(shape, (tuple, list))# 生成总的布数Ti_0 = tf.constant(self.num_timesteps - 1, dtype=tf.int32)# 随机生成一个噪音作为p(x^T)img_0 = noise_fn(shape=shape, dtype=tf.float32)# 循环T次，得到最终的图像_, img_final = tf.while_loop(cond=lambda i_, _: tf.greater_equal(i_, 0),body=lambda i_, img_: [i_ - 1,self.p_sample(denoise_fn=denoise_fn, x=img_, t=tf.fill([shape[0]], i_), noise_fn=noise_fn, return_pred_xstart=False)],loop_vars=[i_0, img_0],shape_invariants=[i_0.shape, img_0.shape],back_prop=False)assert img_final.shape == shapereturn img_final