人人都能看懂的DDPM加噪公式推导

本文是作者对DDPM的一些个人理解，可能有理解错误的地方，欢迎评论指正。在网上也搜索了很多人写的文章，都是相互抄袭，或者是通篇公式。符号一大堆，不理解作者为什么要这样设置参数，不理解内在动机。

本文尝试以笔者自己的观点去理解推导公式，隐含着笔者对作者动机的猜测，可能不够准确，但是对于理解来说非常有用。首先看下面这张图，相信大家对加噪和去噪过程容易理解。我简单描述一下：

加噪过程：分为T步，每步在上一次加完噪的基础上，从正态分布中随机采样噪声加入。
降噪过程：与加噪过程相反，不再赘述。

最近看到一个巨牛的人工智能教程，分享一下给大家。教程不仅是零基础，通俗易懂，而且非常风趣幽默，像看小说一样！觉得太牛了，所以分享给大家。平时碎片时间可以当小说看，【点这里可以去膜拜一下大神的“小说”】。

先给出DDPM最终公式：

1. ${\alpha }_t+{\beta }_t=1$，其中 ${\beta }_t$ 是正态分布方差，即第 $t$ 步产生的噪声从 $N(0,{\beta }_t)$ 采样。
2. $X_t=\sqrt{{\overline{\alpha }}_t}{X}_0+\sqrt{1-{\overline{\alpha }}_t}{z}_t$，其中$X_t$表示第t步加噪后的图像，$X_0$表示初始图像。${\overline{\alpha }}_t={\alpha }_t{\alpha }_{t-1}{\alpha }_{t-2}...{\alpha }_1$，$z_t$ ~$N(0,1)$

因此只需确定N步加噪中的每一步的噪点采样方差，就能一步实现叠加N随机噪声。

0 数学知识

对于DDPM加噪过程，其实只需一个公式：

如果$x$ ~ $N(0,1)$，那么，$ax$ ~ $N(0,a^2)$

1 前向公式推导

1.1 噪点采样分布

噪点从正态分布中采样，那么，如何确定正态分布的均值和方差？

1. 均值： $u$。理论上可以取任意值，但是考虑到让添加到图像上，希望噪点值有正有负，因此取$u=0$是最合适的，因为可以保证正负概率相同。
2. 方差：${\sigma }^2$。理论上可以取任意值，这里我们不具体设置值。

第 $t$ 次往图像上添加正态分布噪点$Z_t$，则$Z_t$ ~ $N(0,{\sigma }_t^2)$。令$Z_t={\sigma }_t{z}_t$，通过正态分布性质可知：$z_t$~$N(0,1)$

因此，只需从标准正态分布采样 $z_t$，然后乘以 ${\sigma }_t$ 即可得到 $Z_t$

1.2如何添加噪声

第 $t$ 步加噪后的图像$X_t$是在第 $t-1$ 步得到的图像$X_{t-1}$基础上加噪点$Z_t$。即：
$$X_t=a_t{X}_{t-1}+b_t{Z}_t$$
其中， $a_t$ 和 $b_t$ 为第 $t$ 步的两个常数系数。$Z_t$替换成${\sigma }_t{z}_t$，得：
$$X_t=a_t{X}_{t-1}+b_t{\sigma }_t{z}_t$$
其中$b_t{\sigma }_t{z}_t$ ~ $N(0,b_t^2{\sigma }_t^2)$，即，方差为$b_t^2{\sigma }_t^2$的正态分布。不妨令${\beta }_t=b_t^2{\sigma }_t^2$，那么：
$$X_t=a_t{X}_{t-1}+\sqrt{{\beta }_t}{z}_t$$
可以如下递推：
$$\begin{array}{rl}{X}_t& =a_t{X}_{t-1}+\sqrt{{\beta }_t}{z}_t\\ & =a_t(a_{t-1}{X}_{t-2}+\sqrt{{\beta }_{t-1}}{z}_{t-1})+\sqrt{{\beta }_t}{z}_t\\ & =a_t{a}_{t-1}{X}_{t-2}+a_t\sqrt{{\beta }_{t-1}}{z}_{t-1}+\sqrt{{\beta }_t}{z}_t\\ & =a_t{a}_{t-1}(a_{t-2}{X}_{t-3}+\sqrt{{\beta }_{t-2}}{z}_{t-2})+a_t\sqrt{{\beta }_{t-1}}{z}_{t-1}+\sqrt{{\beta }_t}{z}_t\\ & =a_t{a}_{t-1}{a}_{t-2}{X}_{t-3}+a_t{a}_{t-1}\sqrt{{\beta }_{t-2}}{z}_{t-2}+a_t\sqrt{{\beta }_{t-1}}{z}_{t-1}+\sqrt{{\beta }_t}{z}_t\\ & =(a_t{a}_{t-1}{a}_{t-2}...a_1)X_0+(a_t{a}_{t-1}...a_2)\sqrt{{\beta }_1}{z}_1+...+a_t{a}_{t-1}\sqrt{{\beta }_{t-2}}{z}_{t-2}+a_t\sqrt{{\beta }_{t-1}}{z}_{t-1}+\sqrt{{\beta }_t}{z}_t\end{array}$$
上式绿色部分，每一项都是独立的正态分布，且均值为0，方差为各自系数的平方。根据正态分布叠加公式，它们的和也服从正态分布，且均值为0，方差为各项方差之和。

因此上式绿色部分的方差为：

$$(a_t{a}_{t-1}...a_2{)}^2{\beta }_1+(a_t{a}_{t-1}...a_3{)}^2{\beta }_2+...+a_t^2{a}_{t-1}^2{\beta }_{t-2}+a_t^2{\beta }_{t-1}+{\beta }_t$$

前面有提到： $a_t$ 和 $b_t$ 为第 $t$ 步的两个常数系数，${\beta }_t$与 $b_t$是一个概念，${\beta }_t$也是个常数系数。

为了便于计算，令$a_t=\sqrt{1-{\beta }_t}$ ，即，$a_t^2+{\beta }_t=1$，上式为：

$$\begin{array}{rl}& (a_t{a}_{t-1}...a_2{)}^2{\beta }_1+(a_t{a}_{t-1}...a_3{)}^2{\beta }_2+...+a_t^2{a}_{t-1}^2{\beta }_{t-2}+a_t^2{\beta }_{t-1}+{\beta }_t\\ & \\ & =(a_t{a}_{t-1}...a_2{)}^2(1-a_1^2)+(a_t{a}_{t-1}...a_3{)}^2(1-a_2^2)+...+a_t^2{a}_{t-1}^2(1-a_{t-2}^2)+a_t^2(1-a_{t-1}^2)+(1-a_t^2)\\ & \\ & =(a_t{a}_{t-1}...a_2{)}^2-(a_t{a}_{t-1}...a_1{)}^2+(a_t{a}_{t-1}...a_3{)}^2-(a_t{a}_{t-1}...a_2{)}^2+...+(a_t{a}_{t-1}{)}^2-(a_t{a}_{t-1}{a}_{t-2}{)}^2\\ & +a_t^2-a_t^2{a}_{t-1}^2+1-a_t^2\\ & =1-(a_t{a}_{t-1}...a_1{)}^2\end{array}$$

即：
$$X_t=(a_t{a}_{t-1}{a}_{t-2}...a_1)X_0+\sqrt{1-(a_t{a}_{t-1}...a_1{)}^2}{z}_t$$
其中， $z_t$ ~$N(0,1)$，令$\overline{a_t}=a_t{a}_{t-1}{a}_{t-2}...a_1$，那么：
$$X_t=\overline{a_t}{X}_0+\sqrt{1-{\overline{a_t}}^2}{z}_t$$

1.3 对齐论文

注意，在DDPM 论文中，给出的由$X_0$生成$X_t$公式如下：
$$q(x_t|x_0)=N(x_t;\sqrt{\overline{{\alpha }_t}}{x}_0,(1-\overline{{\alpha }_t})I)$$
翻译一下就是$X_t$服从均值为$\sqrt{\overline{{\alpha }_t}}{x}_0$，方差为$(1-\overline{{\alpha }_t})$的正态分布。进一步标准化即：
$$x_t=\sqrt{\overline{{\alpha }_t}}{x}_0+\sqrt{1-\overline{{\alpha }_t}}{z}_t$$
其中， $z_t$ ~$N(0,1)$

注意到，本文推导出来的公式与论文稍微有些区别，但是本质是一样的。前面1.2小节我们得到两个公式：

1. $a_t^2+{\beta }_t=1$，其中 ${\beta }_t$ 是正态分布方差，即第 $t$ 步噪声从 $N(0,{\beta }_t)$ 采样。
2. $X_t=(a_t{a}_{t-1}{a}_{t-2}...a_1)X_0+\sqrt{1-(a_t{a}_{t-1}...a_1{)}^2}{z}_t$

为了对齐论文，我们把符号稍微调整。令${\alpha }_t=a_t^2$，则有：

1. ${\alpha }_t+{\beta }_t=1$，其中 ${\beta }_t$ 是第 $t$ 步噪声从 $N(0,{\beta }_t)$ 采样。
2. $X_t=\sqrt{{\alpha }_t{\alpha }_{t-1}{\alpha }_{t-2}...{\alpha }_1}{X}_0+\sqrt{1-{\alpha }_t{\alpha }_{t-1}...{\alpha }_1}{z}_t$

令${\overline{\alpha }}_t={\alpha }_t{\alpha }_{t-1}{\alpha }_{t-2}...{\alpha }_1$，则有
$$X_t=\sqrt{{\overline{\alpha }}_t}{X}_0+\sqrt{1-{\overline{\alpha }}_t}{z}_t$$

1.3 小结

目前为止，得到如下两个公式：

1. ${\alpha }_t+{\beta }_t=1$，其中 ${\beta }_t$ 是正态分布方差，即第 $t$ 步产生的噪声从 $N(0,{\beta }_t)$ 采样。
2. $X_t=\sqrt{{\overline{\alpha }}_t}{X}_0+\sqrt{1-{\overline{\alpha }}_t}{z}_t$，其中${\overline{\alpha }}_t={\alpha }_t{\alpha }_{t-1}{\alpha }_{t-2}...{\alpha }_1$，$z_t$ ~$N(0,1)$

2 代码实现

代码实现这里我们首先分析diffusers库里实现的scheduling_ddpm源码，下一篇文章手撕代码实现，敬请关注。从上面公式推理分析过程大概可以总结实现步骤如下：

1. 确定总的步数num_train_timesteps
2. 按照某种策略，确定第t步的噪音采样的方差${\beta }_t$
3. 计算 $\overline{{\alpha }_t}={\alpha }_t{\alpha }_{t-1}{\alpha }_{t-2}...{\alpha }_1$，即 $\overline{{\alpha }_t}=(1-{\beta }_t)(1-{\beta }_{t-1}(1-{\beta }_{t-2})...(1-{\beta }_1)$
4. 计算第 $t$ 步加噪后的图像$X_t=\sqrt{\overline{{\alpha }_t}}{X}_0+\sqrt{1-\overline{{\alpha }_t}}{z}_t$

分析diffusers库实现的scheduling_ddpm源码，源码地址https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_ddpm.py

对应上述的4步，我们从源码中找到对应的实现。首先在__init__函数中实现如下几步：

1. 确定总步数num_train_timesteps，作为参数传入到函数__init__中（下面代码第3行）：

 def __init__(
        self,
        num_train_timesteps: int = 1000,   # 总步数
        beta_start: float = 0.0001,        # 方差beta的最小值
        beta_end: float = 0.02,            # 方差beta的最大值
        beta_schedule: str = "linear",
       ...其他参数略
    ):

2. 按照某种策略，确定第t步的噪音采样的方差${\beta }_t$ ，这一步依然在函数__init__中实现：

if trained_betas is not None:     # 如果直接传递方差beta，那就直接用
     self.betas = torch.tensor(trained_betas, dtype=torch.float32)
 elif beta_schedule == "linear":  # 线性生成策略，从最小到最大值，等间距采样作为每一步的方差。
     self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
 elif beta_schedule == "scaled_linear":
     # this schedule is very specific to the latent diffusion model.
     self.betas = (
         torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
     )
 elif beta_schedule == "squaredcos_cap_v2":
     # Glide cosine schedule
     self.betas = betas_for_alpha_bar(num_train_timesteps)
 elif beta_schedule == "sigmoid":
     # GeoDiff sigmoid schedule
     betas = torch.linspace(-6, 6, num_train_timesteps)
     self.betas = torch.sigmoid(betas) * (beta_end - beta_start) + beta_start
 else:
     raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")

此时可以得到每一步的采样方差，存放在self.betas属性中。

3. 计算$\overline{{\alpha }_t}=(1-{\beta }_t)(1-{\beta }_{t-1}(1-{\beta }_{t-2})...(1-{\beta }_1)$，这一步依然在函数__init__中实现：

 self.alphas =1.0 - self.betas
 self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
 self.one = torch.tensor(1.0)

注意，上面代码属性self.alphas相当于本文中的${\alpha }_t^2$，因为他这里计算的是$1-{\beta }_t$，还没开根号。另外，这里是一次性计算了所有步的${\alpha }_t^2$。

这里调用了torch.cumprod函数，此函数是计算当前位置前面所有元素相乘。说的抽象，以一个具体例子说明：

>>> a = torch.randn(10)
>>> a
tensor([ 0.6001,  0.2069, -0.1919,  0.9792,  0.6727,  1.0062,  0.4126,
        -0.2129, -0.4206,  0.1968])
>>> torch.cumprod(a, dim=0)
tensor([ 0.6001,  0.1241, -0.0238, -0.0233, -0.0157, -0.0158, -0.0065,
         0.0014, -0.0006, -0.0001])

>>> a[5] = 0.0
>>> torch.cumprod(a, dim=0)
tensor([ 0.6001,  0.1241, -0.0238, -0.0233, -0.0157, -0.0000, -0.0000,
         0.0000, -0.0000, -0.0000])

因此，这里self.alphas_cumprod属性相当于是本文中的 $\overline{{\alpha }_t}$

4. 计算第 $t$ 步加噪后的图像$X_t$

公式：$X_t=\sqrt{\overline{{\alpha }_t}}{X}_0+\sqrt{1-\overline{{\alpha }_t}}{z}_t$
这里封装到函数add_noise中：

 def add_noise(
        self,
        original_samples: torch.FloatTensor,
        noise: torch.FloatTensor,
        timesteps: torch.IntTensor,
    ) -> torch.FloatTensor:
        ...略
        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
		...略
        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
        ...略
        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
        return noisy_samples

add_noise函数中，original_samples参数表示初始图片 $X_0$。 noise参数表示从标准正态分布采样得到的噪声即 $z_t$ ；timesteps表示采样加噪step次数，注意与num_train_timesteps区分，为了提升泛化能力，在训练时，会动态设置加噪次数。接下来我们只看函数内部实现的关键代码，其他代码忽略。

重点代码解析：

sqrt_alpha_prod：相当于$\sqrt{\overline{{\alpha }_t}}$

sqrt_one_minus_alpha_prod：即$\sqrt{1-\overline{{\alpha }_t}}$

noisy_samples = $\sqrt{\overline{{\alpha }_t}}\ast$ original_samples +$\sqrt{1-\overline{{\alpha }_t}}\ast$ noise

与第1节推导公式完全一致。