[人工智能]Diffusion model—扩散模型

摘要

The diffusion model is a generative model of the Encoder-Decoder architecture, which is divided into a diffusion stage and an inverse diffusion stage. In the diffusion stage, by continuously adding noise to the original data, the data is changed from the original distribution to the distribution we expect, for example, the original data distribution is changed to a normal distribution by continuously adding Gaussian noise. During the inverse diffusion stage, a neural network is used to restore the data from a normal distribution to the original data distribution. Its advantage is that each point on the normal distribution is a mapping of the real data, and the model has better interpretability. The disadvantage is that iterative sampling is slow, resulting in low model training and prediction efficiency.

扩散模型是Encoder-Decoder架构的生成模型，分为扩散阶段和逆扩散阶段。 在扩散阶段，通过不断对原始数据添加噪声，使数据从原始分布变为我们期望的分布，例如通过不断添加高斯噪声将原始数据分布变为正态分布。 在逆扩散阶段，使用神经网络将数据从正态分布恢复到原始数据分布。 它的优点是正态分布上的每个点都是真实数据的映射，模型具有更好的可解释性。 缺点是迭代采样速度慢，导致模型训练和预测效率低。

参考资料

一、简介

Diffusion model模型分为扩散过程和逆扩散过程，扩散过程通过对原始数据不断加入高斯噪音，使原始数据变为高斯分布的数据，即从$X_0$ $?>$ $X_T$。逆扩散过程通过高斯噪声还原出图片，即从$X_T$ $?>$ $X_0$。

二、扩散过程

2.1 定义扩散过程

在设定扩散过程是一个马尔可夫链的条件下，向原始信息中不断添加高斯噪声，每一步添加高斯噪声的过程是从$X_{t?1}?>X_t$，于是定义公式：
$$q(x_t|x_{t?1})=N(x_t;\sqrt{1?{\beta }_t}{x}_{t?1},{\beta }_{t}I)$$

该公式表示从$x_{t?1}?>x_t$是一个以$\sqrt{1?{\beta }_t}{x}_{t?1}$为均值${\beta }_t$为方差的高斯分布变换。

2.2 重参数技巧得到迭代公式

利用重参数技巧得到每一次添加高斯噪声的公式如下：
$$X_t=\sqrt{1?{\beta }_t}{X}_{t?1}+{\sqrt{\beta }}_t{Z}_t$$

• $X_t$表示 t 时刻的数据分布
• $Z_t$表示 t 时刻添加的高斯噪音，一般固定是均值为0方差为1的高斯分布
• $\sqrt{1?{\beta }_t}{X}_{t?1}$ 表示当前时刻分布的均值
• ${\sqrt{\beta }}_t$表示当前时刻分布的标准差(标准差=$\sqrt{方差}$)

注意：其中${\beta }_t$是预先设定0~1之间的常量，故扩散过程不含参。

2.3 得到全局扩散公式

在2.2的迭代公式中可知，扩散过程中只有一个参数$\beta$，而$\beta$是预先设置的常量，故扩散过程中无未知的需要学习的参数，所以只需要知道初始数据分布$X_0$和${\beta }_t$就可以得到任意时刻的分布$X_t$，具体公式如下：

• $X_0$为原始数据的分布
• ${\alpha }_t=1?{\beta }_t$
• $\overline{{\alpha }_t}={\prod }_{i=1}^t{\alpha }_i$
• Z为均值为0方差为1的高斯分布

2.4 扩散过程实现代码

2.4.1 总结扩散公式

由2.3可知扩散过程公式为：
$$X_t=\sqrt{\overline{{\alpha }_t}}{X}_0+\sqrt{1?\bar{\alpha }}Z$$其中：

• $X_0$为原始数据的分布
• ${\alpha }_t=1?{\beta }_t$
• $\overline{{\alpha }_t}={\prod }_{i=1}^t{\alpha }_i$
• Z为均值为0方差为1的高斯分布

2.4.2 代码

1. 用make_s_curve生成数据为例得到$X_0$

   # 得到数据X0
   s_curve, _ = make_s_curve(10**4, noise=0.1)
   x_0 = s_curve[:, [0, 2]]/10.0
   # 查看形状
   print(np.shape(x_0))
   # 绘图
   data = x_0.T
   fig, ax = plt.subplots()
   ax.scatter(*data, color='red', edgecolor='white')
   ax.axis('off')
   dataset = torch.Tensor(data)
   

   

2. 假定有100个时刻设置， 所有时刻的$\beta$

   num_steps = 100
   betas = torch.linspace(-6, 6, num_steps)
   betas = torch.sigmoid(betas)*(0.5e-2 - 1e-5)+1e-5
   

   $\beta$为0-1之前很小的数，最大值为0.5e-2，最小值为1e-5

3. 得到$\alpha$（$\alpha =1?\beta$）

   alphas = 1 - betas
   

4. 得到各个时刻的$\overline{{\alpha }_t}$（$\overline{{\alpha }_t}={\prod }_{i=1}^t{\alpha }_i$）

   alphas_prod = torch.cumprod(alphas, 0)
   

5. 得到$\sqrt{{\alpha }_t}$

   alphas_bar_sqrt = torch.sqrt(alphas_bar)
   

6. 得到$\sqrt{1?\overline{{\alpha }_t}}$

   one_minus_alphas_bar_sqrt = torch.sqrt(1-alphas_bar)
   

7. 输入$X_0$与时刻t，得到$X_t$，即$$X_t=\sqrt{\overline{{\alpha }_t}}{X}_0+\sqrt{1?\overline{{\alpha }_t}}Z$$

   def x_t(x_0, t):
       noise = torch.randn_like(x_0)
       return (alphas_bar_sqrt[t]*x_0 + one_minus_alphas_bar_sqrt[t]*noise)
   

8. 扩散过程演示

   num_shows = 20
   fig, axs = plt.subplots(2, 10, figsize=(28, 3))
   plt.rc('text', color='blue')
   
   for i in range(num_shows):
       j = i//10
       k = i%10
       num_x_t = x_t(dataset, torch.tensor([i*num_steps//num_shows]))
       axs[j, k].scatter(*num_x_t, color='red', edgecolor='white')
       axs[j, k].set_axis_off()
       axs[j, k].set_title('$q(\mathbf{x}_{'+str(i*num_steps//num_shows)+'})$')
   

   

三、逆扩散过程

3.1 目标公式

扩散过程是将原始数据不断加噪得到高斯噪声，逆扩散过程是从高斯噪声中恢复原始数据，我们假定逆扩散过程仍然是一个马尔可夫链的过程，要做的是$X_T?>X_0$，用公式表达如下：
$$p_{\theta }(x_{t?1}|x_t)=N(x_{t?1};u_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t))$$

3.2 后验条件概率

推导得到后验条件概率$q(x_{t?1}|x_t,x_0)$

其方差$\overline{{\beta }_t}$为：
$$\overline{{\beta }_t}=\frac{1?\overline{{\alpha }_{t?1}}}{1?\overline{{\alpha }_t}}{\beta }_t$$
均值$\bar{u}(x_{t?1},x_0)$为：
$$\bar{u}(x_{t?1},x_0)=\frac{\sqrt{{\alpha }_t}(1?{\bar{\alpha }}_{t?1})}{1?\overline{{\alpha }_t}}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t?1}}{\beta }_t}{1?\overline{{\alpha }_t}}{x}_0$$
逆扩散过程模型不应当事先知道$x_0$，故需将$x_0$用$x_t$代替，根据2.4得到：

代入均值公式中，化简后得到后验条件均值：
$${\bar{u}}_t=\frac{1}{\sqrt{{\alpha }_t}}(x_t?\frac{{\beta }_t}{\sqrt{1?\overline{{\alpha }_t}}}{z}_t)$$

四、优化目标

4.1 损失函数公式推导

得到损失函数如下：

4.2 损失函数代码实现

def diffusion_loss_fn(model, x_0, alphas_bar_sqrt, one_minus_alphas_bar_sqrt, n_steps):
    batch_size = x_0.shape[0]
    # 生成时间随机值，大小是(batch_size//2)
    t = torch.randint(0, n_steps, size=(batch_size//2,))
    
    t = torch.cat([t, num_steps-1-t], dim=0)
    t = t.unsqueeze(-1) # t.shape为（batch_size, 1）
    
    a = alphas_bar_sqrt[t].to(device)
    
    aml = one_minus_alphas_bar_sqrt[t].to(device)
    
    e = torch.randn_like(x_0).to(device)
    
    x = x_0 * a + e * aml
    
    output = model(x, t.squeeze(-1).to(device))
    
    return (e - output).square().mean()

五、算法流程

5.1 模型训练代码

print('训练模型...')

batch_size = 128
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
num_epoch = 4000
plt.rc('text', color='blue')

model = MLPDiffusion(num_steps)
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for t in range(num_epoch):
    for idx, batch_x in enumerate(dataloader):
        batch_x = batch_x.to(device)
        loss = diffusion_loss_fn(model,batch_x,alphas_bar_sqrt,one_minus_alphas_bar_sqrt,num_steps)
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
        optimizer.step()
        
    if(t%100==0):
        print(loss)
torch.save(model, "model.h5")

5.2 模型采样代码

def p_sample_loop(model, shape, n_steps, betas, one_minus_alphas_bar_sqrt):
    cur_x = torch.randn(shape).to(device)
    x_seq = [cur_x]
    for i in reversed(range(n_steps)):
        cur_x = p_sample(model, cur_x, i, betas.to(device), one_minus_alphas_bar_sqrt.to(device))
        x_seq.append(cur_x)
    return x_seq
        
def p_sample(model, x, t, betas, one_minus_alphas_bar_sqrt):
    t = torch.tensor([t]).to(device)
    coeff = betas[t]/one_minus_alphas_bar_sqrt[t]
    eps_theta = model(x, t)
    # 计算均值
    mean = (1 / (1-betas[t]).sqrt())*(x - (coeff*eps_theta))
    z = torch.randn_like(x).to(device)
    # 计算标准差
    sigma_t = betas[t].sqrt().to(device)
    sample = mean + sigma_t * z
    return (sample)


model = torch.load("model.h5")
x_seq = p_sample_loop(model, dataset.shape, num_steps, betas, one_minus_alphas_bar_sqrt)   
fig, axs = plt.subplots(1, 10, figsize=(28, 3))
for i in range(1, 11):
    cur_x = x_seq[i*10].detach()
    axs[i-1].scatter(cur_x[:, 0].cpu(), cur_x[:, 1].cpu(), color='red', edgecolor='white');
    axs[i-1].set_axis_off();
    axs[i-1].set_title('$q(\mathbf{x}_{'+str(i*10)+'})$')

5.3 训练好的模型效果