扩散模型DDPM根据代码尝试进行形象的理解
扩散模型可以用于生成图形。SBM（Score-Based Model）和DDPM（Denoising Diffusion Probabilistic Model）是两种常见的扩散模型。本文依据Github上的极简代码，尝试理解DDPM。
DDPM 的基本思想是在训练阶段将数据逐渐加上噪声（扩散过程），然后在预测阶段再一步步去除噪声（反向扩散过程），得到真实数据。
训练阶段
假设一张图片通过加噪的过程完全变成噪声需要T步，每一步加一点点噪声，直到最后完全变成噪声。
训练阶段，给模型的输入值是加噪的步数t、以及加噪后的带噪声图像x_t，预测输出值是所加的噪声。这样，在预测阶段，根据每一步预测的噪声，就可以一步步去噪、还原图像。
初始化
以下是训练器的类的初始化定义：
class GaussianDiffusionTrainer(nn.Module):
def __init__(self, model, beta_1, beta_T, T):
super().__init__()
self.model = model
self.T = T
self.register_buffer(
'betas', torch.linspace(beta_1, beta_T, T).double())
alphas = 1. - self.betas
alphas_bar = torch.cumprod(alphas, dim=0)
# calculations for diffusion q(x_t | x_{t-1}) and others
self.register_buffer(
'sqrt_alphas_bar', torch.sqrt(alphas_bar))
self.register_buffer(
'sqrt_one_minus_alphas_bar', torch.sqrt(1. - alphas_bar))
其中：
- self.model：传入的
model是U-Net。U-Net是一种全卷积网络，先通过编码器提取深层特征，然后通过解码器逐步恢复图像的空间分辨率，并通过跳跃连接来融合编码器和解码器的特征，以生成高质量的输出图像。 - self.T：扩散的步数，此处默认是1000。
- self.betas：从1e-4到2e-2的长度为
T的等差数列。 - alphas_bar：
alphas的累乘，是0.9999到约4e-5的递减数列。 - self.sqrt_alphas_bar：
alphas_bar的开根号，递减数列。 - self.sqrt_one_minus_alphas_bar：
(1 - alphas_bar)的开根号，递增数列。
前向传播
def forward(self, x_0):
"""
Algorithm 1.
"""
t = torch.randint(self.T, size=(x_0.shape[0], ), device=x_0.device)
noise = torch.randn_like(x_0)
x_t = (
extract(self.sqrt_alphas_bar, t, x_0.shape) * x_0 +
extract(self.sqrt_one_minus_alphas_bar, t, x_0.shape) * noise)
loss = F.mse_loss(self.model(x_t, t), noise, reduction='none')
return loss
其中：
- x_0：传入的图像矩阵，形为
B C H W，如[80, 3, 32, 32]。 - t：为每张图像随机选取一个
T以内的扩散步数，训练这一步的模型预测能力。 - noise：随机噪声。
- x_t：第t步的带噪声的图像。计算方式是\(x_{t} = \sqrt{1-\beta_t} x_{0} + \sqrt{\beta_t} \mathbf{\epsilon}_t\)，其中 \(\beta_t\) 是一个时间相关的递增参数，\(\mathbf{\epsilon}_t\) 是一个标准高斯噪声。随着步数t的增加，真实图像\(x_0\)的权重越来越小，噪声的权重越来越大，象征着一步步加噪的过程。
通过这部分的训练，可以使self.model具有预测噪声的能力。
预测阶段
class GaussianDiffusionSampler(nn.Module):
def __init__(self, model, beta_1, beta_T, T):
super().__init__()
self.model = model
self.T = T
self.register_buffer('betas', torch.linspace(beta_1, beta_T, T).double())
alphas = 1. - self.betas
alphas_bar = torch.cumprod(alphas, dim=0)
alphas_bar_prev = F.pad(alphas_bar, [1, 0], value=1)[:T]
self.register_buffer('coeff1', torch.sqrt(1. / alphas))
self.register_buffer('coeff2', self.coeff1 * (1. - alphas) / torch.sqrt(1. - alphas_bar))
self.register_buffer('posterior_var', self.betas * (1. - alphas_bar_prev) / (1. - alphas_bar))
def predict_xt_prev_mean_from_eps(self, x_t, t, eps):
assert x_t.shape == eps.shape
return (
extract(self.coeff1, t, x_t.shape) * x_t -
extract(self.coeff2, t, x_t.shape) * eps
)
def p_mean_variance(self, x_t, t):
# below: only log_variance is used in the KL computations
var = torch.cat([self.posterior_var[1:2], self.betas[1:]])
var = extract(var, t, x_t.shape)
eps = self.model(x_t, t)
xt_prev_mean = self.predict_xt_prev_mean_from_eps(x_t, t, eps=eps)
return xt_prev_mean, var
def forward(self, x_T):
"""
Algorithm 2.
"""
x_t = x_T
for time_step in reversed(range(self.T)):
print(time_step)
t = x_t.new_ones([x_T.shape[0], ], dtype=torch.long) * time_step
mean, var= self.p_mean_variance(x_t=x_t, t=t)
# no noise when t == 0
if time_step > 0:
noise = torch.randn_like(x_t)
else:
noise = 0
x_t = mean + torch.sqrt(var) * noise
assert torch.isnan(x_t).int().sum() == 0, "nan in tensor."
x_0 = x_t
return torch.clip(x_0, -1, 1)
关注Algorithm 2部分。首先，输入值x_T是完全的噪声。去噪过程通过逆向时间步循环来实现，从最高时间步 T-1 开始，逐步回到第0时间步。
在每个时间步，计算均值和方差。然后，最新的x_t一部分由均值mean组成，另一部分由一个新的噪声组成，噪声权重是\(\sqrt{var}\)。可以预见，随着时间步的推移，噪声权重会越来越低，直至0。最后一步的均值就是去噪完成的图像。
关于均值和方差，参见函数 p_mean_variance()。
方差是从一开始就计算好的方差序列中取值，该序列 self.posterior_var == [0, 5.4532e-5, ..., 2e-2] 是一个递增数列，那么逆向时间步循环下就是个递减数列，代表方差越来越小。
eps 是模型预测出的图像附带的噪声。
均值的计算利用函数 predict_xt_prev_mean_from_eps() 进行。均值 mean = coeff1 * x_t - coeff2 * eps，表示图像去除一部分噪声。随着时间步的推移，coeff1 从1.0102逐渐递减到1.0001，coeff2 从0.0202缓慢递减到0.0063再快速增加到0.01。
通俗来讲，去噪过程就是去一点噪、加一点噪、去一点噪、加一点噪，加的噪声越来越小，最终实现完全去噪的过程。
推荐后续阅读
Stable Diffusion 解读（三）：原版实现及Diffusers实现源码解读 | 周弈帆的博客 (zhouyifan.net)
