扩散模型（Diffusion Model）原理与代码解析（二）

扩散模型（Diffusion Model）原理与代码解析（一）
扩散模型（Diffusion Model）原理与代码解析（二）

四、损失函数

我们已经明确了要训练$p_{\Theta }(X_{t-1}|X_t)$，那要怎么确定目标函数呢？有两个很直接的想法，一个是负对数的最大似然概率，即$-\log p_{\Theta }(X_0)$，另一个是真实分布与预测分布的交叉熵，即$-E_{q(X_0)}\log p_{\Theta }(X_0)$，然而，类似于VAE，由于我们很难对噪声空间进行积分，因此直接优化$-\log p_{\Theta }(X_0)$或$-E_{q(X_0)}\log p_{\Theta }(X_0)$是很困难的，因此我们不会直接优化它们，而是优化它们的变分上界（Variational Lower Bound）$L_{VLB}$，$L_{VLB}$的定义如下:$$L_{VLB}=E_{q(X_{0:T})}\left[\log \frac{q(X_{1:T}|X_0)}{p_{\Theta }(X_{0:T})}\right]$$下面证明$L_{VLB}$是$-\log p_{\Theta }(X_0)$和$-E_{q(X_0)}\log p_{\Theta }(X_0)$的上界，即证明$L_{VLB}\ge -\log p_{\Theta }(X_0)\&L_{VLB}\ge -E_{q(X_0)}\log p_{\Theta }(X_0)$:$$\begin{array}{rl}-\log p_{\Theta }(X_0)& \le -\log p_{\Theta }(X_0)+D_{KL}\left(q(X_{1:T}|X_0)||p_{\Theta }(X_{1:T}|X_0)\right)(KL散度大于等于零)\\ & =-\log p_{\Theta }(X_0)+E_{X_{1:T}\sim q(X_{1:T}|X_0)}\left(\log \frac{q(X_{1:T}|X_0)}{p_{\Theta }(X_{1:T}|X_0)}\right)\\ & =-\log p_{\Theta }(X_0)+E_{X_{1:T}\sim q(X_{1:T}|X_0)}\left(\log \frac{q(X_{1:T}|X_0)p_{\Theta }(X_0)}{p_{\Theta }(X_{0:T})}\right)\\ & =-\log p_{\Theta }(X_0)+E_{X_{1:T}\sim q(X_{1:T}|X_0)}\left(\log \frac{q(X_{1:T}|X_0)}{p_{\Theta }(X_{0:T})}+\log p_{\Theta }(X_0)\right)\\ & =E_{X_{0:T}\sim q(X_{0:T})}\left(\log \frac{q(X_{1:T}|X_0)}{p_{\Theta }(X_{0:T})}\right)\\ & =L_{VLB}\end{array}$$
$$\begin{array}{rl}{L}_{CE}& =-\int q(X_0)\log p_{\Theta }(X_0)d{X}_0\\ & =-E_{q(X_0)}\log p_{\Theta }(X_0)\\ & =-E_{q(X_0)}\log \left(\int p_{\Theta }(X_{1:T}|X_0)p_{\Theta }(X_0)d{X}_{1:T}\right)\\ & =-E_{q(X_0)}\log \left(\int p_{\Theta }(X_{0:T})d{X}_{1:T}\right)\\ & =-E_{q(X_0)}\log \left(\int q(X_{1:T}|X_0)\frac{p_{\Theta }(X_{0:T})}{q(X_{1:T}|X_0)}d{X}_{1:T}\right)\\ & =-E_{q(X_0)}\log \left(E_{q(X_{1:T}|X_0)}\frac{p_{\Theta }(X_{0:T})}{q(X_{1:T}|X_0)}\right)\\ & \le -E_{q(X_0)}\left(E_{q_{(}{X}_{1:T}|X_0)}\log \frac{p_{\Theta }(X_{0:T})}{q(X_{1:T}|X_0)}\right)\\ & =E_{q(X_{0:T})}\left[\log \frac{q(X_{1:T}|X_0)}{p_{\Theta }(X_{0:T})}\right]\\ & =L_{VLB}\end{array}$$至此，证明了$L_{VLB}$是$-\log p_{\Theta }(X_0)$和$-E_{q(X_0)}\log p_{\Theta }(X_0)$的上界。
下面，对$L_{VLB}$化简：$$\begin{array}{rl}{L}_{VLB}& =E_{q(X_{0:T})}\left[\log \frac{q(X_{1:T}|X_0)}{p_{\Theta }(X_{0:T})}\right]\\ & =E_{q(X_{0:T})}\left[\log \frac{{\prod }_{t=1}^{T}q(X_t|X_{t-1})}{p_{\Theta }(X_T){\prod }_{t=1}^T{p}_{\Theta }(X_{t-1}|X_t)}\right]\\ & =E_{q(X_{0:T})}\left[-\log p_{\Theta }(X_T)+\sum _{t=1}^T\log \frac{q(X_t|X_{t-1})}{p_{\Theta }(X_{t-1}|X_t)}\right]\\ & =E_{q(X_{0:T})}\left[-\log p_{\Theta }(X_T)+\sum _{t=2}^T\log \frac{q(X_t|X_{t-1})}{p_{\Theta }(X_{t-1}|X_t)}+\log \frac{q(X_1|X_0)}{p_{\Theta }(X_0|X_1)}\right]\#(第二项分子q(X_t|X_{t-1}=\frac{q(X_t,X_{t-1},X_0)}{q(X_{t-1},X_0)}=\frac{q(X_{t-1}|X_t,X_0)q(X_t|X_0)q(X_0)}{q(X_{t-1},X_0)}=\frac{q(X_{t-1}|X_t,X_0)q(X_t|X_0)}{q(X_{t-1}|X_0)})\\ & =E_{q(X_{0:T})}\left[-\log p_{\Theta }(X_T)+\sum _{t=2}^T\log \left(\frac{q(X_{t-1}|X_t{X}_0)}{p_{\Theta }(X_{t-1}|X_t)}\ast \frac{q(X_t|X_0)}{q(X_{t-1}|X_0)}\right)+\log \frac{q(X_1|X_0)}{p_{\Theta }(X_0|X_1)}\right]\\ & =E_{q(X_{0:T})}\left[-\log p_{\Theta }(X_T)+\sum _{t=2}^T\log \frac{q(X_{t-1}|X_t{X}_0)}{p_{\Theta }(X_{t-1}|X_t)}+\sum _{t=2}^T\log \frac{q(X_t|X_0)}{q(X_{t-1}|X_0)}+\log \frac{q(X_1|X_0)}{p_{\Theta }(X_0|X_1)}\right]\\ & =E_{q(X_{0:T})}\left[-\log p_{\Theta }(X_T)+\sum _{t=2}^T\log \frac{q(X_{t-1}|X_t{X}_0)}{p_{\Theta }(X_{t-1}|X_t)}+\log \frac{q(X_T|X_0)}{q(X_1|X_0)}+\log \frac{q(X_1|X_0)}{p_{\Theta }(X_0|X_1)}\right]\\ & =E_{q(X_{0:T})}\left[-\log p_{\Theta }(X_T)+\sum _{t=2}^T\log \frac{q(X_{t-1}|X_t{X}_0)}{p_{\Theta }(X_{t-1}|X_t)}+\log \frac{q(X_T|X_0)}{p_{\Theta }(X_0|X_1)}\right]\\ & =E_{q(X_{0:T})}\left[\log \frac{q(X_T|X_0)}{p_{\Theta }(X_T)}+\sum _{t=2}^T\log \frac{q(X_{t-1}|X_t{X}_0)}{p_{\Theta }(X_{t-1}|X_t)}-\log p_{\Theta }(X_0|X_1)\right]\\ & =D_{KL}(q(X_T|X_0)||p_{\Theta }(X_T))+\sum _{t=2}^T{D}_{KL}(q(X_{t-1}|X_t{X}_0)||p_{\Theta }(X_{t-1}|X_t))-\log p_{\Theta }(X_0|X_1)\\ & =L_T+L_{T-1}+...+L_0\\ & where:L_T=D_{KL}(q(X_T|X_0)||p_{\Theta }(X_T))\\ & L_t=D_{KL}(q(X_t|X_{t+1}{X}_0)||p_{\Theta }(X_t|X_{t+1})),1\le t\le T\\ & L_0=-\log p_{\Theta }(X_0|X_1)\end{array}$$从$L_t$即可看出，对$p_{\Theta }(X_t|X_{t+1})$的监督就是最小化$p_{\Theta }(X_t|X_{t+1})$和$q(X_t|X_{t+1}{X}_0)$的KL散度。

简单的说，我们的目的是希望学习出一个$p_{\Theta }(X_{t-1}|X_t)$，即能够从噪声图恢复出原图。为了达到这一个目的，我们使用$q(X_{t-1}|X_t{X}_0)$来监督$p_{\Theta }(X_{t-1}|X_t)$进行训练，$q(X_{t-1}|X_t{X}_0)$是可以用$q(X_t|X_0)$和$q(X_t|X_{t-1})$表示的，即$q(X_{t-1}|X_t{X}_0)$是已知的。

DDPM论文将$p_{\Theta }(X_{t-1}|X_t)=N(X_{t-1};{\mu }_{\theta }(X_t,t),{\Sigma }_{\theta }(X_t,t))$中的方差设置为${\beta }_t$，所以可学习的参数就只在均值中。对于两个单一变量的高斯分布$p$和$q$而言，$$KL(p,q)=\log \frac{{\sigma }_2}{{\sigma }_1}+\frac{{\sigma }_1^2+({\mu }_1-{\mu }_2{)}^2}{2{\sigma }_2^2}-\frac{1}{2}$$优化式子：$$\sum _{t=1}^T{D}_{KL}(q(X_{t-1}|X_t{X}_0)||p_{\Theta }(X_{t-1}|X_t))$$其中，$q$为已知有偏高斯分布，$p_{\Theta }$为所要拟合的分布，由于假设$p_{\Theta }$的方差为常数，则我们只需逼近$p_{\Theta }$和$q$的均值即可，等价于最小化式子：$$\begin{array}{rl}Loss& ={\mathbb{E}}_q\left(\frac{1}{2{\sigma }_t^2}||{\bar{\mu }}_t(X_t,X_0)-{\mu }_{\theta }(X_t,t)|{|}^2\right)+C\#(X_t是由X_0和噪声ϵ决定的变量)\\ & ={\mathbb{E}}_{X_0,ϵ}\left(\frac{1}{2{\sigma }_t^2}||{\bar{\mu }}_t(X_t(X_0,ϵ),X_0)-{\mu }_{\theta }(X_t,t)|{|}^2\right)\#(由于X_t=\sqrt{{\bar{\alpha }}_t}{X}_0+\sqrt{1-{\bar{\alpha }}_t}Z,Z\sim N(0,I)，即可求出X_0的表达式)\\ & ={\mathbb{E}}_{X_0,ϵ}\left(\frac{1}{2{\sigma }_t^2}||\frac{1}{\sqrt{{\alpha }_t}}(X_t(X_0,ϵ)-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}ϵ)-{\mu }_{\theta }(X_t,t)|{|}^2\right)\#(带入\bar{\mu }(X_t,X_0)的表达式，上述的X_0可不做替换)\\ & \#作者认为直接预测恢复的数据效果不好，转而预测噪声，这个想法有点像预测残差连接中的残差\\ & \#一个网络输入X_0,{\bar{\alpha }}_t,高斯噪声ϵ和t,然后预测高斯噪声ϵ\\ & ={\mathbb{E}}_{X_0,ϵ}\left(\frac{1}{2{\sigma }_t^2}||\frac{1}{\sqrt{{\alpha }_t}}(X_t(X_0,ϵ)-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}ϵ)-\frac{1}{\sqrt{{\alpha }_t}}(X_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(X_t,t))|{|}^2\right)\\ & ={\mathbb{E}}_{X_0,ϵ}\left(\frac{{\beta }_t^2}{2{\sigma }_t^2{\alpha }_t(1-{\bar{\alpha }}_t)}||ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{X}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t)|{|}^2\right)\end{array}$$作者在训练时发现，去掉Loss前的系数，可使训练稳定，所以简化后的Loss为$$Loss={\mathbb{E}}_{X_0,ϵ}\left(||ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{X}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t)|{|}^2\right)$$拟合出来的均值为${\mu }_{\theta }(x_t,t)=\frac{1}{\sqrt{{\alpha }_t}}\left(X_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(X_t,t)\right)$可用于采样。

五、训练和采样过程

训练过程：从数据集中采样$x_0$，从均匀分布采样 $t$，可使模型鲁棒，采样噪声，计算 Loss 更新模型。

采样过程：从标准正态分布采样$x_T$，迭代计算$x_{t-1}$，已知均值${\mu }_{\theta }(x_t,t)$和常数方差${\beta }_t$利用参数重整化可计算出$x_{t-1}$，直到$x_0$。

六、Pytorch示例代码

import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import make_moons
import torch
import torch.nn as nn
from PIL import Image

moons_curve, _ = make_moons(10 ** 4, noise=0.05)
print("shape of moons:", np.shape(moons_curve))

data = moons_curve.T

fig, ax = plt.subplots()
ax.scatter(*data, color='blue', edgecolor='white')
ax.axis('off')
dataset = torch.Tensor(moons_curve).float()

num_steps = 100  # 扩散100步
# 制定每一步的beta
betas = torch.linspace(-6, 6, num_steps)
betas = torch.sigmoid(betas) * (0.5e-2 - 1e-5) + 1e-5  # 先压缩到0~1，再乘以0.005

# 计算alpha、alpha_prod、alpha_prod_previous、alpha_bar_sqrt等变量的值
alphas = 1 - betas
alphas_prod = torch.cumprod(alphas, 0)
alphas_prod_p = torch.cat([torch.tensor([1]).float(), alphas_prod[:-1]], 0)  # 插入第一个数1，丢掉最后一个数，previous连乘
alphas_bar_sqrt = torch.sqrt(alphas_prod)
one_minus_alphas_bar_log = torch.log(1 - alphas_prod)
one_minus_alphas_bar_sqrt = torch.sqrt(1 - alphas_prod)

assert alphas.shape == alphas_prod.shape == alphas_prod_p.shape == alphas_bar_sqrt.shape == one_minus_alphas_bar_log.shape == one_minus_alphas_bar_sqrt.shape
print("all the same shape", betas.shape)  # 所有值都是同等维度，且都是常值


# 计算任意时刻的x采样值，基于X_0和重参数化 ， 扩散过程
def q_x(X_0, t):
    noise = torch.randn_like(X_0)
    alphas_t = alphas_bar_sqrt[t]
    alphas_1_m_t = one_minus_alphas_bar_sqrt[t]
    return (alphas_t * X_0 + alphas_1_m_t * noise)  # 在X_0的基础上添加噪声


num_shows = 20
fig, axs = plt.subplots(2, 10, figsize=(28, 3))
plt.rc('text', color='black')

# 共有10000个点，每个点包含两个坐标。生成100步中每隔5步加噪声后的图像，最终应该会成为一个各向同性的高斯分布
for i in range(num_shows):
    j = i // 10
    k = i % 10
    q_i = q_x(dataset, torch.tensor([i * num_steps // num_shows]))  # 生成t时刻的采样数据
    axs[j, k].scatter(q_i[:, 0], q_i[:, 1], color='red', edgecolor='white')
    axs[j, k].set_axis_off()
    axs[j, k].set_title('$q(\mathbf{X}_{' + str(i * num_steps // num_shows) + '})$')
fig.show()

class MLPDiffusion(nn.Module):  # 定义一个 MLP 模型
    def __init__(self, n_steps, num_units=128):
        super(MLPDiffusion, self).__init__()

        self.linears = nn.ModuleList(
            [
                nn.Linear(2, num_units),
                nn.ReLU(),
                nn.Linear(num_units, num_units),
                nn.ReLU(),
                nn.Linear(num_units, num_units),
                nn.ReLU(),
                nn.Linear(num_units, 2),
            ]
        )
        self.step_embeddings = nn.ModuleList(
            [
                nn.Embedding(n_steps, num_units),
                nn.Embedding(n_steps, num_units),
                nn.Embedding(n_steps, num_units),
            ]
        )

    def forward(self, x, t):
        # x = x_0
        for idx, embedding_layer in enumerate(self.step_embeddings):  # 三层
            t_embedding = embedding_layer(t)
            x = self.linears[2 * idx](x)
            x += t_embedding
            x = self.linears[2 * idx + 1](x)

        x = self.linears[-1](x)  # 输出维度与输入一致

        return x


def diffusion_loss_fn(model, x_0, alphas_bar_sqrt, one_minus_alphas_bar_sqrt, n_steps):
    """对任意时刻t进行采样计算loss"""
    batch_size = x_0.shape[0]

    # 对一个batchsize样本生成随机的时刻t
    t = torch.randint(0, n_steps, size=(batch_size // 2,))  # 为了 t 不重复，先采样一半
    t = torch.cat([t, n_steps - 1 - t], dim=0)
    t = t.unsqueeze(-1)

    # x0的系数
    a = alphas_bar_sqrt[t]

    # 随机噪声eps的系数
    aml = one_minus_alphas_bar_sqrt[t]

    # 生成随机噪音eps
    e = torch.randn_like(x_0)

    # 构造模型的输入
    x = x_0 * a + e * aml

    # 送入模型，得到t时刻的随机噪声预测值
    output = model(x, t.squeeze(-1))

    # 与真实噪声一起计算误差，求平均值
    return (e - output).square().mean()


def p_sample_loop(model, shape, n_steps, betas, one_minus_alphas_bar_sqrt):
    """从x[T]恢复x[T-1]、x[T-2]|...x[0]"""
    cur_x = torch.randn(shape)
    x_seq = [cur_x]
    for i in reversed(range(n_steps)):
        cur_x = p_sample(model, cur_x, i, betas, one_minus_alphas_bar_sqrt)
        x_seq.append(cur_x)
    return x_seq


def p_sample(model, x, t, betas, one_minus_alphas_bar_sqrt):
    """从 x_t 开始生成 t-1 时刻的重构值"""
    t = torch.tensor([t])

    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)
    sigma_t = betas[t].sqrt()

    sample = mean + sigma_t * z

    return (sample)


# 开始训练模型
seed = 1234
print('Training model...')
batch_size = 128
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
num_epoch = 4001
plt.rc('text', color='blue')

model = MLPDiffusion(num_steps)  # 输出维度是2，输入是x和step
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for t in range(num_epoch):
    for idx, batch_x in enumerate(dataloader):
        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)
        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], cur_x[:, 1], color='red', edgecolor='white');
            axs[i - 1].set_axis_off();
            axs[i - 1].set_title('$q(\mathbf{x}_{' + str(i * 10) + '})$')
        fig.show()

第0个epoch，100次扩散，每10次输出一次

第1000个epoch：

第2000个epoch：

第3000个epoch：

# 生成扩散和逆扩散的 GIF
imgs = []
for i in range(100):
    plt.clf()
    q_i = q_x(dataset,torch.tensor([i]))
    plt.scatter(q_i[:,0],q_i[:,1],color='red',edgecolor='white',s=5);
    plt.axis('off');
    
    img_buf = io.BytesIO()
    plt.savefig(img_buf,format='png')
    img = Image.open(img_buf)
    imgs.append(img)
reverse = []
for i in range(100):
    plt.clf()
    cur_x = x_seq[i].detach()
    plt.scatter(cur_x[:,0],cur_x[:,1],color='red',edgecolor='white',s=5);
    plt.axis('off')
    
    img_buf = io.BytesIO()
    plt.savefig(img_buf,format='png')
    img = Image.open(img_buf)
    reverse.append(img)
imgs = imgs + reverse
imgs[0].save("diffusion.gif", format='GIF', append_images=imgs, save_all=True, duration=100, loop=0)