论文阅读笔记：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$可以进一步写为公式（6）。
$$\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}$$
论文《Denoising Diffusion Probabilistic Models》中是直接给方差${\sigma }_2$（论文中为${\sigma }_t$）设置了固定的参数，分别为${\beta }_t$或者$\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}\cdot {\beta }_t$。后文中的公式（8）（9）（10）都是这样的。并且论文指出：从实验来看，这两者之间的结果是相似。然而，与论文提到的，代码中使用了两种损失函数，第一种损失函数就是公式（6），如下方代码中的函数normal_kl中所示。第二种损失函数是忽略掉方差 。直接对均值计算平方差，如下方代码代码中函数train_losses中的mse部分所示。

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, eps
	    self.model_var_type = model_var_type  # learned, fixedsmall, fixedlarge
	    self.loss_type = loss_type  # kl, mse
	
	    assert isinstance(betas, np.ndarray)
	    self.betas = betas = betas.astype(np.float64)  # computations here in float64 for accuracy
	    assert (betas > 0).all() and (betas <= 1).all()
	    timesteps, = betas.shape
	    self.num_timesteps = int(timesteps)
	
	    alphas = 1. - betas
	    self.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 others
	    self.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 chain
	    self.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.shape
	    if 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.shape
	    return (
	        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.shape
	    posterior_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.shape
	    assert 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 computations
	      model_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_output
	    elif self.model_mean_type == 'xstart':  # the model predicts x_0
	      pred_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 epsilon
	      pred_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.shape
	    if return_pred_xstart:
	      return model_mean, model_variance, model_log_variance, pred_xstart
	    else:
	      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 bound
	      losses = 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 MSE
	      assert 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.shape
	      losses = nn.meanflat(tf.squared_difference(target, model_output))
	    else:
	      raise NotImplementedError(self.loss_type)
	
	    assert losses.shape == t.shape
	    return losses
	    
	# 计算两个高斯分布的KL散度，代码中的logvar1，logvar2为方差的对数. 上文中的公式（6）
	def normal_kl(mean1, logvar1, mean2, logvar2):
	  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.shape
	    decoder_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 == 0
    nonzero_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) * noise
    assert sample.shape == pred_xstart.shape
    return (sample, pred_xstart) if return_pred_xstart else sample

def p_sample_loop(self, denoise_fn, *, shape, noise_fn=tf.random_normal):
    """
    Generate samples
    """
    assert isinstance(shape, (tuple, list))
	# 生成总的布数T
    i_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 == shape
    return img_final