論文閱讀筆記：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{equation} \begin{split} L_1&=E_{x_{1:T} \sim q(x_{1:T} | x_0)} \Bigg(log \Big[ \frac{q(x_{T}|x_0)}{ p(x_T)}\Big]\Bigg) \\ &=\int dx_{1:T} \cdot q(x_{1:T}| x_0) \cdot log \Big[ \frac{q(x_{T}|x_0)}{ p(x_T)}\Big] \\ &=\int dx_{1:T} \cdot \frac{q(x_{1:T}| x_0)}{q(x_T|x_0)} \cdot q(x_T|x_0) \cdot log \Big[ \frac{q(x_{T}|x_0)}{ p(x_T)}\Big] \\ &=\int dx_{1:T} \cdot \underbrace{ q(x_{1:T-1}| x_0, x_T) }_{q(x_{1:T}| x_0)=q(x_{T}|x_0) \cdot q(x_{1;T-1}| x_0, x_T)} \cdot q(x_T|x_0) \cdot log \Big[ \frac{q(x_{T}|x_0)}{ p(x_T)}\Big] \\ &=\int \Bigg( \underbrace{ \int q(x_{1:T-1}| x_0, x_T) \cdot \prod_{k=1}^{T-1} dx_k }_{二重積分化為兩個定積分相乘，并且=1} \Bigg) \cdot q(x_T|x_0) \cdot log \Big[ \frac{q(x_{T}|x_0)}{ p(x_T)} \Big] \cdot dx_{T} \\ &=\int q(x_T|x_0) \cdot log \Big[ \frac{q(x_{T}|x_0)}{ p(x_T)} \Big] \cdot dx_{T} \\ &=E_{x^T\sim q(x_T|x_0)} log \Big[ \frac{q(x_{T}|x_0)}{ p(x_T)} \Big]\\ &= KL\Big(q(x_T|x_0)||p(x_T)\Big) \end{split} \end{equation}$

可以看出， $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{equation} \begin{split} L_2&=E_{x_{1:T} \sim q(x_{1:T} | x_0)} \Bigg(\sum_{t=2}^{T} log \Big[\frac{q(x_{t-1}|x_t,x_0)}{ p_{\theta}(x_{t-1}|x_t)} \Big]\Bigg)\\ &=\sum_{t=2}^{T} E_{x_{1:T} \sim q(x_{1:T} | x_0)} \Bigg(log \Big[\frac{q(x_{t-1}|x_t,x_0)}{ p_{\theta}(x_{t-1}|x_t)} \Big]\Bigg)\\ &=\sum_{t=2}^{T} \Bigg( \int dx_{1:T} \cdot q(x_{1:T}| x_0) \cdot log \Big[\frac{q(x_{t-1}|x_t,x_0)}{ p_{\theta}(x_{t-1}|x_t)} \Big] \Bigg)\\ &=\sum_{t=2}^{T} \Bigg( \int dx_{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 \Big[\frac{q(x_{t-1}|x_t,x_0)}{ p(x_{t-1}|x_t)} \Big] \Bigg)\\ &=\sum_{t=2}^{T} \Bigg( \int dx_{1:T} \cdot \underbrace{ \frac{q(x_{0:T})}{q(x_0)}}_{q(x_{0:T})=q(x_0)\cdot q(x_{1:T}| x_0)} \cdot \underbrace{ \frac{q(x_t,x_0)}{q(x_t,x_{t-1},x_0)}}_{q(x_t,x_{t-1},x_0)=q(x_t,x_0)\cdot q(x_{t-1}|x_t,x_0)} \cdot q(x_{t-1}|x_t,x_0) \cdot log \Big[\frac{q(x_{t-1}|x_t,x_0)}{ p_{\theta}(x_{t-1}|x_t)} \Big] \Bigg)\\ &=\sum_{t=2}^{T} \Bigg( \int dx_{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 \Big[\frac{q(x_{t-1}|x_t,x_0)}{ p_{\theta}(x_{t-1}|x_t)} \Big] \Bigg)\\ &=\sum_{t=2}^{T} \Bigg( \int \bigg[ \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\geq1 ,k\neq t-1} dx_k \bigg] \cdot q(x_{t-1}|x_t,x_0) \cdot log \Big[\frac{q(x_{t-1}|x_t,x_0)}{ p_{\theta}(x_{t-1}|x_t)} dx_{t-1} \Big] \Bigg)\\ &=\sum_{t=2}^{T} \Bigg( \int \bigg[ \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\geq1 ,k\neq t-1} dx_k \bigg] \cdot q(x_{t-1}|x_t,x_0) \cdot log \Big[\frac{q(x_{t-1}|x_t,x_0)}{ p_{\theta}(x_{t-1}|x_t)} dx_{t-1} \Big] \Bigg)\\ &=\sum_{t=2}^{T} \Bigg( \int \bigg[ \underbrace{ \int q(x_{k:k\geq1,k\neq t-1}|x_{t-1},x_0)}_{q(x_{0;T})=q(x_{t-1},x_0)\cdot q(x_{k:k\geq1,k\neq t-1}|x_{t-1},x_0)} \cdot \underbrace {\frac{q(x_t|x_0)}{ q(x_t|x_{t-1},x_0)}}_{q(x_t,x_0)=q(x_0)\cdot q(x_t|x_0)} \prod_{k\geq1 ,k\neq t-1} dx_k \bigg] \cdot q(x_{t-1}|x_t,x_0) \cdot log \Big[\frac{q(x_{t-1}|x_t,x_0)}{ p_{\theta}(x_{t-1}|x_t)} dx_{t-1} \Big] \Bigg)\\ &=\sum_{t=2}^{T} \Bigg( \int \bigg[\int q(x_{k:k\geq1,k\neq t-1}|x_{t-1},x_0)\cdot \underbrace {\frac{q(x_t|x_0)}{ q(x_t|x_{t-1},x_0)}}_{=1} \prod_{k\geq1 ,k\neq t-1} dx_k \bigg] \cdot q(x_{t-1}|x_t,x_0) \cdot log \Big[\frac{q(x_{t-1}|x_t,x_0)}{ p_{\theta}(x_{t-1}|x_t)} dx_{t-1} \Big] \Bigg)\\ &=\sum_{t=2}^{T} \Bigg( \int \bigg[\int q(x_{k:k\geq1,k\neq t-1}|x_{t-1},x_0)\cdot \prod_{k\geq1 ,k\neq t-1} dx^k \bigg] \cdot q(x_{t-1}|x_t,x_0) \cdot log \Big[\frac{q(x_{t-1}|x_t,x_0)}{ p_{\theta}(x_{t-1}|x_t)} dx_{t-1} \Big] \Bigg)\\ &=\sum_{t=2}^{T} \Bigg( \int \bigg[\underbrace{ \int q(x_{k:k\geq1,k\neq t-1}|x_{t-1},x^0)\cdot \prod_{k\geq1 ,k\neq t-1} dx_k }_{=1}\bigg] \cdot q(x_{t-1}|x_t,x_0) \cdot log \Big[\frac{q(x_{t-1}|x_t,x_0)}{ p_{\theta}(x_{t-1}|x_t)} dx_{t-1} \Big] \Bigg)\\ &=\sum_{t=2}^{T} \Bigg( \int q(x_{t-1}|x_t,x_0) \cdot log \Big[\frac{q(x_{t-1}|x_t,x_0)}{ p_{\theta}(x_{t-1}|x_t)} dx_{t-1} \Big] \Bigg)\\ &=\sum_{t=2}^{T} \Bigg( E_{x_{t-1}\sim q(x_{t-1}|x_t,x_0)} log \Big[\frac{q(x_{t-1}|x_t,x_0)}{ p_{\theta}(x_{t-1}|x_t)} \Big] \Bigg)\\ &=\sum_{t=2}^{T}KL\bigg(q(x_{t-1}|x_t,x_0)||p_{\theta}(x_{t-1}|x_t) \bigg) \end{split} \end{equation}$
最后考慮 $L_3$ ，事實上，在論文《Deep Unsupervised Learning using Nonequilibrium Thermodynamics》中提到為了防止邊界效應，強制另 $p(x^0|x^1)=q(x^1|x^0)$ ，因此這一項也是個常數。

由以上分析可知道，損失函數可以寫為公式（3）。
$\begin{equation} \begin{split} L&:=L_1+L_2+L_3 \\ &=KL\Big(q(x_T|x_0)||p(x_T)\Big) + \sum_{t=2}^{T}KL\bigg(q(x_{t-1}|x_t,x_0)||p_{\theta}(x_{t-1}|x_t) \bigg)-log \Big[p_{\theta}(x_{0}|x_1)\Big] \end{split} \end{equation}$

忽略掉 $L_1$ 和 $L_3$ ，損失函數可以寫為公式（4）。
$\begin{equation} \begin{split} L:=\sum_{t=2}^{T}KL\bigg(q(x_{t-1}|x_t,x_0)||p_{\theta}(x_{t-1}|x_t) \bigg) \end{split} \end{equation}$

可以看出損失函數 $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{equation} \begin{split} \sigma_1&=\sqrt{\frac{\beta_t\cdot (1-\bar{\alpha_{t-1}})}{(1-\bar{\alpha_{t}})}}\\ \mu_1&=\frac{1}{\sqrt{\alpha_t}}\cdot (x_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha_t}}}\cdot z_t) \\ 或者 \mu_1&=\frac{\sqrt{\alpha_t}\cdot(1-\bar{\alpha_{t-1}})}{1-\bar{\alpha_t}}\cdot x_t+\frac{\beta_t\cdot \sqrt{\bar{\alpha_{t-1}}}}{1-\bar{\alpha_t}} \cdot x_0 \end{split} \end{equation}$

而 $p_{\theta}(x_{t-1}|x_t)$ 則由模型（深度學習模型或者其他模型）估算出其均值和方差，分別記作 $\mu_2,\sigma_2$ 。
因此損失函數 $L$ 可以進一步寫為公式12。
$\begin{equation} \begin{split} L:=log \Big[\frac{\sigma_2}{\sigma_1}\Big]+\frac{\sigma_1^2 +(\mu_1-\mu_2)^2}{2\sigma_2^2}-\frac{1}{2} \end{split} \end{equation}$

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