Elucidating the Design Space of Diffusion-Based Generative Models 阅读笔记

文章用一种新的设计框架统一diffusion-based model，并使用模块化（modular）的思想，分别从采样、训练、score network设计三个方面分析和改进diffusion-based model。

之前的工作¹已经把diffusion-based model统一到SDE或者ODE框架下了，这篇文章的作者同样也从SDE和ODE的角度出发，不过换了一种SDE和ODE的表示形式。

假设有方差是${\sigma }_{data}$的数据分布$p_{data}(\mathbf{x})$。考虑一族柔化后的分布（mollified distribution）$p(\mathbf{x};\sigma (t))$，其通过对数据添加方差为$\sigma$的高斯噪声产生，$p(\mathbf{x};\sigma )=p_{data}\ast \mathcal{N}(0,{\sigma }^2\mathbf{I})=\int p_{data}({\mathbf{x}}_0)\mathcal{N}(\mathbf{x}-{\mathbf{x}}_0;0,{\sigma }^2\mathbf{I})d{\mathbf{x}}_0=\int p_{data}({\mathbf{x}}_0)\mathcal{N}(\mathbf{x};{\mathbf{x}}_0,{\sigma }^2\mathbf{I})d{\mathbf{x}}_0$，这里$\ast$表示卷积运算。理想情况下，我们希望$\sigma (t)$的选择使得$p_{{\sigma }_{min}}\approx p_{data}$，$p_{{\sigma }_{max}}\approx \mathcal{N}(0,{\sigma }_{max}^2\mathbf{I})$。可以注意到上面添加噪声的方式并不能得到均值为0的噪声分布，但在实际中如果${\sigma }_{max}>>{\sigma }_{data}$，我们可以认为$p(\mathbf{x};{\sigma }_{max})$接近方差为${\sigma }_{max}$的纯高斯噪声。

上面添加噪声的过程没有考虑$\mathbf{x}$的缩放，但在一些已有的工作中（比如DDPM）会给$\mathbf{x}$添加缩放。考虑对未缩放变量$\hat{\mathbf{x}}$添加缩放$\mathbf{x}=s(t)\hat{\mathbf{x}}$。
则有下面的ODE描述数据分布随着时间的变化：
$$\mathrm{d}\mathbf{x}=\left[\frac{\dot{s}(t)}{s(t)}\mathbf{x}-s(t{)}^2\dot{\sigma }(t)\sigma (t){\nabla }_{\mathbf{x}}\log p(\frac{\mathbf{x}}{s(t)};\sigma (t))\right]dt\text{\hspace{1em}(4)}$$注意$p(\mathbf{x};\sigma )$是针对未缩放样本的分布，对于缩放样本的分布作者表示为$p_t(\mathbf{x})=s(t{)}^{-d}p(\mathbf{x}/s(t);\sigma )$。
对于上面的ODE，perturbation kernel的形式是：
$$p_{0t}(\mathbf{x}(t)|\mathbf{x}(0))=\mathcal{N}(\mathbf{x}(t);s(t)\mathbf{x}(0),s(t{)}^2\sigma (t{)}^2\mathbf{I})\text{\hspace{1em}(11)}$$信噪比$SNR(t)=\frac{s(t{)}^2}{s(t{)}^2{\sigma }^2(t)}=\frac{1}{\sigma (t)}$。

diffusion-based model都是逐步添加高斯噪声的过程，主要就是两个不同，一个是缩放$s(t)$，一个是噪声的大小$\sigma (t)$。引用论文中的原话：

Our formulation (Eq. 4) highlights the fact that every realization of the probability flow ODE is simply a reparameterization of the same canonical ODE; changing σ(t) corresponds to reparameterizing t, whereas changing s(t) corresponds to reparameterizing x.

如果没有缩放，即$s(t)=1$，那么公式(4)的ODE简化为：
$$\mathrm{d}\mathbf{x}=-\dot{\sigma }(t)\sigma (t){\nabla }_{\mathbf{x}}\log p(\mathbf{x};\sigma (t))dt\text{\hspace{1em}(1)}$$在之前的工作¹中SDE的形式是：
$$\mathrm{d}\mathbf{x}=f(t)\mathbf{x}+g(t)d{w}_t\text{\hspace{1em}(10)}$$其中$s(t)=\exp ({\int }_o^{t}f(\xi )d\xi )$，$\sigma (t)=\sqrt{{\int }_o^t\frac{g(\xi {)}^2}{s(\xi {)}^2}d\xi }$。

不同于之前的论文，这篇文章考虑的是一个直接估计去噪样本的去噪函数$D(\mathbf{x};\sigma )$：
$${\mathbb{E}}_{y\sim p_{data}}{\mathbb{E}}_{\mathbf{n}\sim \mathcal{N}(0,{\sigma }^2\mathbf{I})}\Vert D(\mathbf{y}+\mathbf{n};\sigma )-\mathbf{y}{\Vert }_2^2,{\nabla }_{\mathbf{x}}\log p(\mathbf{x};\sigma )=(D(\mathbf{x};\sigma )-\mathbf{x})/{\sigma }^2\text{\hspace{1em}(2,3)}$$其中$\mathbf{y}$是训练样本，$\mathbf{n}$是添加的噪声。在这种设置下，score function的计算变成了用$D(\mathbf{x};\sigma )$估计添加的噪声。用网络$D_{\theta }(\mathbf{x};\sigma )$按照公式（2）可以估计$D(\mathbf{x};\sigma )$。需要注意的是，$D_{\theta }(\mathbf{x};\sigma )$可能包括额外的预处理步骤和后处理步骤。

ODE解轨迹的形状由$\sigma (t)$和$s(t)$决定。因为在求解微分方程的时候截断误差（truncation error）和$dx/dt$的曲率有关，作者认为最好的选择是$\sigma (t)=t$和$s(t)=1$，带入公式(1)得到$dx/dt=(\mathbf{x}-D(\mathbf{x};t))/t$。这种情况下$\sigma$和$t$是相同的，两个符号可以串着用。好处是在任何$x,t$位置，一个到$t=0$的Euler步就是对去噪图像的估计$D_{\theta }(\mathbf{x};t)$，解估计的切线总是指向去噪图像。如下图所示（c）也就是$\sigma (t)=t$和$s(t)=1$的情况，这和DDIM相同。

作者将SDE表示成下面的形式，这个公式是没有缩放的版本，即$s(t)=1$：

这揭示了为什么随机性在实践中有帮助：隐式朗之万扩散驱动样本在给定时间朝向所需的边际分布，主动纠正早期采样步骤中产生的任何错误。

直接用网络$D_{\theta }$预测$D(\mathbf{x};\sigma )$在实际中效果并不好。对于有监督训练神经网络而言，每个样本的梯度幅度的变化不应该太大，而直接学习$D_{\theta }$变化会很大。作者考虑用网络$F_{\theta }$添加预处理步骤和后处理步骤来代表$D_{\theta }$：
$$D_{\theta }(\mathbf{x};\sigma )=c_{skip}(\sigma )\mathbf{x}+c_{out}(\sigma )F_{\theta }(c_{in}(\sigma )\mathbf{x};c_{noise}(\sigma ))$$

Table1是常见的diffusion-based model在作者提出的框架下的总结。

---

1. Score-Based Generative Modeling through Stochastic Differential Equations ↩︎ ↩︎