探究扩散模型的前向扩散过程权重系数设计

前言

目前常见的前向扩散过程的公式表达如下：
$$x_t=\sqrt{1-{\alpha }_t}\cdot x_{t-1}+\sqrt{{\alpha }_t}\cdot {ϵ}_t$$前段时间投稿一篇有关扩散模型的文章，被审稿人问道为什么要采用这种 $\sqrt{1-{\alpha }_t}+\sqrt{\alpha }$ 的加权设计？遂前去搜寻资料和博客，大多数回答都讲是为了最终能够采样到高斯分布，但都不具体说原因。总不能这样给审稿人回复吧，怕不是分分钟被拒。因此，没有办法，只能自己研究一下原因。

先声明：奈何本人能力有限，这种方法只能算是一种探究，欢迎大家指正和交流。

分析

这里我不去关注使用这样加权的原因，而是关注使用不同加权方式分别会产生什么效果。首先我们假设存在 $a_t$ 和 $b_t$ 满足：
$$x_t=a_t{x}_{t-1}+b_t{ϵ}_t\text{\hspace{1em}(1)}$$
对上式按照递推的方式逐步展开得：
$$\begin{array}{c}\begin{array}{rl}{x}_t& =a_t{x}_{t-1}+b_t{ϵ}_t\\ & =a_t\left(a_{t-1}{x}_{t-2}+b_{t-1}{ϵ}_{t-1}\right)+b_t{ϵ}_t\\ & =a_t\left[a_{t-1}\left(a_{t-2}{x}_{t-3}+b_{t-2}{ϵ}_{t-2}\right)+b_{t-1{ϵ}_{t-1}}\right]+b_t{ϵ}_t\\ & =\cdots \cdots \\ & =\left(a_t{a}_{t-1}\cdots a_1\right)x_0+\left(a_t{a}_{t-1}\cdots a_2\right)b_1{ϵ}_1+\left(a_t{a}_{t-1}\cdots a_3\right)b_2{ϵ}_2+\cdots +a_t{b}_{t-1}{ϵ}_{t-1}+b_t{ϵ}_t\end{array}\end{array}$$由于高斯噪声的叠加性质，当$p\sim \mathcal{N}(0,{\sigma }_1^2)$，$q\sim \mathcal{N}(0,{\sigma }_2^2)$时，$p+q\sim \mathcal{N}(0,{\sigma }_1^2+{\sigma }_2^2)$，而上式中的${ϵ}_i$均是标准高斯分布，即${ϵ}_i\sim \mathcal{N}(0,\text{I})$。因此上式可以进一步简化为
$$\begin{array}{c}\begin{array}{rl}{x}_t& =\left(a_t{a}_{t-1}\cdots a_1\right)x_0+\left(a_t{a}_{t-1}\cdots a_2\right)b_1{ϵ}_1+\left(a_t{a}_{t-1}\cdots a_3\right)b_2{ϵ}_2+\cdots +a_t{b}_{t-1}{ϵ}_{t-1}+b_t{ϵ}_t\\ & =\left(a_t{a}_{t-1}\cdots a_1\right)x_0+\sqrt{{\left(a_t{a}_{t-1}\cdots a_2\right)}^2{b}_1^2+\cdots +a_t^2{b}_{t-1}^2+b_t^2}ϵ\end{array}\end{array}$$其中，$ϵ\sim \mathcal{N}(0,\text{I})$，根据扩散模型的参数设置可知 $0<a_t,b_t<1$，因此$x_0$的系数会很快收敛到0，即
$${\overline{\alpha }}_t=a_t{a}_{t-1}\cdots a_1\approx 0$$，因此前向扩散过程能否收敛到一个标准高斯分布的判断就是噪声的系数能否收敛到1。对此，我这里分别判断三种加权方式下的收敛情况：

1. $a_t+b_t=1$；
2. $a_t^2+b_t^2=1$；
3. $a_t=b_t$。

这里取$b_t\in (0,1)$，实际上是与原扩散模型保持一致，$b_t\in (0.0001,0.02)$。对于第一种情况，可以发现其在不同扩散步长时的收敛情况如图：


显然不论扩散步长如何设置，最终的噪声系数还是无法收敛到1，即前向扩散过程无法收敛为一个标准的高斯噪声分布。对于第二种加权方式，其在不同扩散步长时的收敛情况如图：


对于第二种加权方式，噪声分布的系数最终会收敛到1，而且过程是平稳的，时间无关的。同时，在选择第二种加权方式的时候，我们可以对噪声的系数表达式进行简化：
$$\begin{array}{c}\begin{array}{rl}& \sqrt{{\left(a_t{a}_{t-1}\cdots a_2\right)}^2{b}_1^2+\cdots +a_t^2{b}_{t-1}^2+b_t^2}\\ =& \sqrt{a_t^2\{a_{t-1}^2\left[a_{t-2}^2\cdots <a_2^2{b}_1^2+b_2^2>\cdots +b_{t-2}^2\right]+b_{t-1}^2\}+b_t^2}\\ =& \sqrt{a_t^2\{a_{t-1}^2\left[a_{t-2}^2\cdots <a_2^2\left(1-a_1^2\right)+b_2^2>\cdots +b_{t-2}^2\right]+b_{t-1}^2\}+b_t^2}\\ =& \sqrt{a_t^2\{a_{t-1}^2\left[a_{t-2}^2\cdots <a_2^2-a_1^2{a}_2^2+b_2^2>\cdots +b_{t-2}^2\right]+b_{t-1}^2\}+b_t^2}\\ =& \sqrt{a_t^2\{a_{t-1}^2\left[a_{t-2}^2\cdots <1-a_1^2{a}_2^2>\cdots +b_{t-2}^2\right]+b_{t-1}^2\}+b_t^2}\\ & \cdots \cdots \\ =& \sqrt{1-a_t^2{a}_{t-1}^2{a}_{t-2}^2\cdots a_1^2}\end{array}\end{array}$$显然这种加权方式所得的结果也更加优雅，从表达式可以直接看出其最终会收敛到1。对于第三种加权方式，不同扩散步长时的收敛情况如图：


很明显，这种加权方式是无法让噪声系数收敛到1的，也就无法令前向扩散过程收敛为一个标注的高斯分布。以下是我们做该部分仿真是用的 Python 代码：

import numpy as np
import matplotlib.pyplot as plt
from pylab import mpl
mpl.rcParams['font.sans-serif'] = ['SimSun'] # 指定默认字体
mpl.rcParams['axes.unicode_minus'] = False # 解决保存图像是负号'-'显示为方块的问题

num = 100000

b = np.linspace(0.0001, 0.02, num, dtype=np.float64)
a = np.sqrt(1 - np.square(b))
# a = b
# a = 1 - b
c = []
result = 0
for i in range(num):
    if i == 0:
        result = np.square(b[0])
        c.append(np.sqrt(result))
    else:
        result = np.square(a[i]) * result + np.square(b[i])
        c.append(np.sqrt(result))

plt.figure()
plt.plot(range(num), c)
plt.xlabel('扩散步数')
plt.ylabel('噪声系数值')
plt.show()

虽然我们上边的方法从结果导向地解释了使用这种加权方式的合理性，但其具体的原因没有分析。而我作为一个初学者，也存在几个问题仍然存在疑惑，希望能与大家交流。

• 为什么一定要前向扩散到一个标准高斯分布？
• 为什么一定要加标准高斯噪声？