0%

论文阅读-diffusion model-Improved Denoising Diffusion Probabilistic Models

动机

  • DDPM可以产生高质量的生成样本,但没有文章显示DDPMlog-likelihood能力,引发一些问题,比如DDPM是否能够捕获分布的所有模式。

贡献

  • 通过参数化和混合目标函数,表明DDPM在保证高质量图像质量的同时可以提高log-likelihood,且可加速采样过程;
  • 通过precisionrecall指标评估,发现DDPM可以匹配GAN的样本质量,并获得更好的模式覆盖;
  • 发现模型容量和训练计算的增加,可以扩展模型的样本质量和log-likelihood能力。

本文的方法

log-likelihood被广泛应用于生成模型,且认为其可以促使生成模型获取数据分布中的所有模式,并提升特征学习能力和样本质量,但DDPM在该策略上表现较差,且模式覆盖很差。本文探索算法上的改进,使得DDPM在图像数据集上能获得更好的log-likelihood,并由此受益。

Learning \Sigma_{\theta}(x_t,t)

DDPM固定 \Sigma_{\theta}(x_t,t)=\sigma_t^2\mathbf{I} ,虽然 \tilde{\beta}_t \beta_t 作为两个极值,但发现固定 \sigma_t^2=\beta_t \sigma_t^2=\tilde{\beta}_t 的生成质量相似。

如上图所示,除了时间 t\rightarrow 0 时, \tilde{\beta}_t \beta_t 几乎等价,而此时模型正在处理难以察觉的细节,并且随着总步数 T 的增加, \tilde{\beta}_t \beta_t 接近的时间占比更长。这表明随着时间步的无限增加, \sigma_t 的影响很小, \mu_{\theta}(x_t, t) \Sigma_{\theta}(x_t,t) 更影响分布。这表明固定 \sigma_t 对于图像质量而言是合理的,但下图显示,开始的少量时间步对变分下界的影响最大,那么我们可以通过选择更好的 \Sigma_{\theta}(x_t,t) 来提升log-likelihood

前面的图显示 \Sigma_{\theta}(x_t,t) 的合理范围非常小,即使在对数域中,神经网络也很难直接预测。本文直接在对数域中通过插值 \tilde{\beta}_t \beta_t 来参数化方差,即通过模型输出向量 v ,然后参数化得到方差:

\Sigma_{\theta}(x_t,t) = \exp\left(v \log \beta_t + (1-v) \log \tilde{\beta}_t\right)

理想情况下模型可以预测整个插值范围,但实际情况下并没有,表明边界值足够 \Sigma_{\theta}(x_t,t) 去表征。由于 L_{\text{simple}} 不依赖于 \Sigma_{\theta}(x_t,t) ,则定义新损失函数:

L_{\text{hybrid}} = L_{\text{simple}} + \lambda L_{\text{vlb}}

为了避免 L_{\text{vlb}} 过度影响 L_{\text{simple}} ,设置 \lambda=0.001 ,且 L_{\text{vlb}} \mu_{\theta}(x_t, t) 进行stop-gradient

Improving the Noise Schedule

如上图所示,在前向加噪过程的后期,linear schedule的噪声很多导致对训练影响不大,而cosine schedule加噪更慢,且如下图所示,当忽略 20\% 的反向过程时,用linear schedule训练的模型并没有变得更差。

所以本文使用cosine schedule控制 \bar{\alpha}_t

\bar{\alpha}_t=\frac{f(t)}{f(0)}, \quad f(t)=\cos \left(\frac{t / T+s}{1+s} \cdot \frac{\pi}{2}\right)^2

其中, s 是一个很小的偏移量,防止在开始阶段 \beta_t 过小导致模型预测困难,在论文中设置 s=\sqrt{\beta_0}=0.008 。又根据定义:

\bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s = \alpha_t\bar{\alpha}_{t-1} \Longrightarrow \alpha_t = \frac{\bar{\alpha}_t}{\bar{\alpha}_{t-1}} \\ \beta_t = 1 - \alpha_t = 1 - \frac{\bar{\alpha}_t}{\bar{\alpha}_{t-1}}

本文将 \beta_t 截断为不大于 0.999 以避免扩散过程结束时(接近 t\rightarrow T )出现奇点。下图也显示linear schedule下降过快,过快地破坏了数据信息。

Reducing Gradient Noise

为了实现更好的log-likelihood,可以直接优化 L_{\text{vlb}} ,但作者发现 L_{\text{vlb}} 很难优化,如上图所示, L_{\text{vlb}} L_{\text{hybrid}} 都噪声较大,但 L_{\text{hybrid}} 能实现更好的log-likelihood,作者通过下图评估目标函数的gradient noise scales,发现 L_{\text{vlb}} 的梯度相比 L_{\text{hybrid}} 尺度更大,故需要寻找新方法减少 L_{\text{vlb}} 的方差,以便直接优化。

本文优化 L_{\text{vlb}} 时采用重要性采样,即:

L_{\mathrm{vlb}}=E_{t \sim p_t}\left[\frac{L_t}{p_t}\right], \text { where } p_t \propto \sqrt{E\left[L_t^2\right]} \text { and } \sum p_t=1

其中, E\left[L_t^2\right] 事先未知,且在训练时动态变化,故为每个损失项记录历史的前 10 个值并在训练时动态更新,训练初始阶段则从 t\in[0, T-1] 均匀采样 10 次。通过这个重要性采样,可以通过优化 L_{\text{vlb}} 实现最佳的log-likelihood,且重要性采样的目标函数噪声比原始均匀采样的目标函数噪声小很多,但在直接优化 L_{\text{hybrid}} 时,发现重要性采样技术无帮助。

Improving Sampling Speed

训练时本文将 T 提高到 4000 ,通常采样时使用相同的时间步( 1,2,\dots,T ),但为了提升采样速度,希望能够获取一个时间步子序列 S 。对于时间序列 S_t ,可通过训练时加噪策略获得 \bar{\alpha}_{S_t} ,并计算对应的采样方差:

\beta_{S_t} = 1 - \frac{\bar{\alpha}_{S_t}}{\bar{\alpha}_{S_{t-1}}}, \ \ \ \tilde{\beta}_{S_t} = \frac{1-\bar{\alpha}_{S_{t-1}}}{1-\bar{\alpha}_{S_{t}}}\beta_{S_t}

\Sigma_{\theta}(x_{S_t},S_t) 被参数化为 [\bar{\beta}_{S_t}, \beta_{S_t}] 之间,因此可以计算 p(x_{S_{t−1}} \mid x_{S_t}) \sim \mathcal{N}(\mu_\theta(x_{S_t}, S_t), \Sigma_\theta(x_{S_t}, S_t)) 。所以在 [1,T] 之间使用 K 个间隔均匀的实数,并四舍五入到最近的整数作为新的时间步子序列,下图显示了不同方法在不同的采样子序列下的结果,本文可使用仅 100 个采样时间步就可以实现接近最佳效果。

部分实验结果