Diffusion Models：生成扩散模型

https://www.zhihu.com/video/1538310598659530752

用 KL 散度来表示变分上界

根据 Diffusion Models 最早提出的一篇文章[1]，进一步对 L_{vlb} 推导，可以得到变分上限为熵与多个KL散度的累加，根据 KL 散度重写变分上限有：

L_{vlb}=L_{0}+L_{1}+\ldots+L_{T-1}+L_{T} \tag{7}

其中有：

L_{0}=-\log p_{\theta}(x_{0} \mid x_{1}) \tag{8}

L_{t-1}=D_{KL}(q(x_{t-1} \mid x_{t}, x_{0}) \| p_{\theta}(x_{t-1} \mid x_{t})) \tag{9}

L_{T}=D_{KL}(q(x_{T} \mid x_{0}) \| p(x_{T})) \tag{10}

x_0 都会出现在扩散过程中的 L_{t-1} ，现在所有 KL 散度都是在高斯概率分布之间进行比较。 这意味着可以使用闭包表达式，而不是采样的蒙特卡洛估计方式来精确计算变分上界。

到这里看不懂没关系，想表达的是最小化 L_{vlb} 即可最小化 Diffusion Models 的目标损失，而求解 L_{vlb} 则可以通过计算 KL 散度来代替。

有了目标函数的数学基础后，现在需要就如何实现扩散模型训练过程有几个细节：

1. 对于正向扩散过程，唯一需要的选择是概率相关的向量（均值和方差），其值在扩散过程中在隐变量 x_t 中直接添加高斯参数 \beta_{t} 。
   
2. 对于逆扩散过程，需要选择能够表达高斯分布的模型结构，神经网络模型的拟合能力很强，于是就可以引入神经网络模型啦。
   
3. 最后就是对于神经网络模型有一个简单的要求，模型的输入、输出、中间隐变量必须要有相同的维度 dims。
   

损失函数和 L_T

既然有神经网络模型，那自然离不开损失函数，有了损失函数就有了优化的方向和目标，下面来展开损失函数的定义。

在扩散过程公式（1）中，其中 \beta 1, \ldots, \beta T 是高斯分布方差的超参数。且实际中 \beta_t 随着 t 增大是递增的，即 \beta_1 < \ldots < \beta_T 。在实际代码中，我们设置 \beta_t 是从 0.0001 到0.02 线性插值。

在逆扩散过程公式（10）中，时间步 T 中前，由于前向扩散过程中 q 没有可学习参数，只是单纯添加高斯噪声，而最后时间步 T 得到的 x_T 则是纯高斯噪声，因此在训练的过程中 L_T 可以当做常量忽略。即下面公式不需要：

L_{T}=D_{KL}(q(x_{T} \mid x_{0}) \| p(x_{T})):= 0

回顾公式（3），把反向马尔科夫过程转变为高斯分布来表示：

p_{\theta}(x_{t-1} \mid x_{t}):=\mathcal{N}(x_{t-1} ; \boldsymbol{\mu}_{\theta}(x_{t}, t), \mathbf{\Sigma}_{\theta}(x_{t}, t)) \tag{3}

在求解（训练）的过程需要知道得到高斯分布的均值 \boldsymbol{\mu}{\theta}(x{t}, t) 和方差 \mathbf{\Sigma}{\theta}(x{t}, t) 。对于方差 \Sigma 不好求解，DDPM 论文中直接使用 \beta_{t} 来代替：

\boldsymbol{\Sigma}_{\theta}(x_{t}, t)=\sigma_{t}^{2} \mathbb{I} \sigma_{t}^{2}=\beta_{t} \tag{11}

假设多元高斯是具有相同方差的独立高斯的乘积，方差值可以随时间变化 。基于这个假设可以将逆扩散过程的方差，设置为与正向扩散过程方差相同，得到公式（11）。不过有意思的是在实际代码中因为引入了神经网络，所以可以通过训练的方式来得到方差 \Sigma 。

数学上，假设使用 \beta_{t} 来代替方差 {\Sigma}_{\theta} ，因此有：

p_{\theta}(x_{t-1} \mid x_{t}):=\mathcal{N}(x_{t-1} ; \boldsymbol{\mu}_{\theta}(x_{t}, t), \boldsymbol{\Sigma}_{\theta}(x_{t}, t) \\ :=\mathcal{N}(x_{t-1} ; \boldsymbol{\mu}_{\theta}(x_{t}, t), \sigma_{t}^{2} \mathbf{I}) \tag{12}

那现在就更清楚啦，可以将KL散度：

L_{t-1}=D_{KL}(q(x_{t-1} \mid x_{t}, x_{0}) \| p_{\theta}(x_{t-1} \mid x_{t})) \tag{13}

转换为：

L_{t-1} \propto\|\tilde{\mu}_{t}(x_{t}, x_{0})-\mu_{\theta}(x_{t}, t)\|^{2} \tag{14}

由正向扩散过程我们知道，任意时刻的 x_t 可以由 x_0 和 \beta 表示。参数 \mu_{\theta} 最直接的表示是预测扩散模型的后验概率的均值。不过在实际测试当中，DDPM 作者发现通过训练以在任何给定时间步 t 预测噪声分量 \mu_{\theta} 会产生更好的结果。设：

\boldsymbol{\mu}_{\theta}(x_{t}, t)=\frac{1}{\sqrt{\alpha_{t}}}(x_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}} \boldsymbol{\epsilon}_{\theta}(x_{t}, t)) \tag{15}

其中：

\alpha_{t}:=1-\beta_{t} \text { and } \bar{\alpha}_{t}:=\prod_{s=1}^{t} \alpha_{s} \tag{16}

因此可以得到损失函数 L，使得训练更加稳定：

L_{\text {simple }}(\theta):=\mathbb{E}_{t, x_{0}, \epsilon}[\|\boldsymbol{\epsilon}-\boldsymbol{\epsilon}_{\theta}(\sqrt{\bar{\alpha}_{t}} x_{0}+\sqrt{1-\bar{\alpha}_{t}} \boldsymbol{\epsilon}, t)\|^{2}] \tag{17}

高斯分布 {\epsilon} 为神经网络模型所预测的噪声（用于去噪），可看做为 {\epsilon}{\theta}(x_t, t) 。 Diffusion Models 训练的核心就是取学习高斯噪声 {\epsilon}, {\epsilon}{\theta} 之间均方误差 MSE 。

实际训练过程和 L_0

上面已经说过啦，在逆扩散过程中，马尔科夫过程表示为由连续条件高斯分布下的累积变换组成。有了总体的优化策略，还要看每个像素的计算方式，在逆扩散过程结束时，我们希望得到一张生成好的图像，因此需要设计一种方法，使得图像上每个像素值都满足离散的对数似然。

为了达到这个目的，将逆扩散过程中的最后从 x_1 到 x_0 的转换设置为独立的离散计算方式。 即在最后一个转换过程在给定 x_1$ x_1 下得到图像 x_0 满足对数似然，假设像素与像素之间是相互独立的：

p_{\theta}(x_{0} \mid x_{1})=\prod_{i=1}^{D} p_{\theta}(x_{0}^{i} \mid x_{1}^{i}) \tag{18}

式 (18) 中 D 是输入数据的维数，上标 i 表示图像中的一个坐标位置。现在的目标是确定给定像素的值可能性有多大，也就是想要知道对应时间步 t=1 下噪声图像 x 中相应像素值的分布：

\mathcal{N}(x ; \mu_{\theta}^{i}(x_{1}, 1), \sigma_{1}^{2}) \tag{19}

其中 t = 1 的像素分布来自多元高斯分布，其对角协方差矩阵允许我们将分布拆分为单变量高斯的乘积：

\mathcal{N}(x ; \mu_{\theta}(x_{1}, 1), \sigma_{1}^{2} \mathbb{I})=\prod_{i=1}^{D} \mathcal{N}(x ; \mu_{\theta}^{i}(x_{1}, 1), \sigma_{1}^{2}) \tag{20}

现在假设图像已经从0-255的数值之间，经过归一化在[-1,1]的范围内。在 t=0 时给定每个像素的像素值，最后一个时间步 t=1 的转换概率分布 p_{\theta}(x_{0} \mid x_{1}) 的值就是每个像素值的乘积。简而言之，这个过程由等式简洁 (18) 地表示：

p_{\theta}(x_{0} \mid x_{1})=\prod_{i=1}^{D} p_{\theta}(x_{0}^{i} \mid x_{1}^{i}) \\ =\prod_{i=1}^{D} \int_{\delta_{-}(x_{0}^{i})}^{\delta_{+}(x_{i}^{i})} \mathcal{N}(x ; \mu_{\theta}^{i}(x_{1}, 1), \sigma_{1}^{2}) dx \tag{21}

其中约束有：

\delta_{-}(x)= \begin{cases}-\infty & x=-1 \\ x-\frac{1}{255} & x>-1\end{cases}

和：

\delta_{+}(x)= \begin{cases}\infty & x=1 \\ x+\frac{1}{255} & x<1\end{cases}

那现在我们就可以计算公式 (8) 中，最后的转换概率分布 p_{\theta}(x_{0} \mid x_{1}) 的值。

训练过程如图左边 Algorithm 1 Training部分：

1. 从标准高斯分布采样一个噪声 \epsilon \sim \mathcal{N}(0, \mathbf{I}) ；
   
2. 通过梯度下降最小化损失 \nabla_{\theta}|\boldsymbol{\epsilon}-\mathbf{z}{\theta}(\sqrt{\bar{\alpha}{t}} x_{0}+\sqrt{1-\bar{\alpha}_{t}} \boldsymbol{\epsilon}, t)|^{2} ；
   
3. 训练到收敛为止（训练时间比较长，T代码中设置为1000）
   

测试（采样）如图右边 Algorithm 2 Sampling部分：

1. 从标准高斯分布采样一个噪声 x_T \sim \mathcal{N}(0, \mathbf{I}) ；
   
2. 从时间步 T 开始正向扩散迭代到时间步 1；
   
3. 如果时间步不为1，则从标准高斯分布采样一个噪声 z \sim \mathcal{N}(0, \mathbf{I}) ，否则 z=0；
   
4. 根据高斯分布计算每个时间步 t 的噪声图；
   

网络模型

虽然通过上面一连串很复杂或者看不懂的公式，其实就是为了得到损失函数，或者要知道 Diffusion Model 的优化目标。简化损失函数后的目标变得更加简单了，主要是训练模型 {\epsilon}_{\theta} 。但是看了很多文章，其实还停留在复杂的数学公式推导，仍然没有定义 Diffusion Model 的网络模型结构。

在定义网络模型之前呢，再次声明，Diffusion Model 中对模型的唯一要求是输入和输出的数据维度 dims 需要相同。有点意思，看到这句话，估计就会想到类似于 U-Net的模型架构，而不是 VAE 那种 Encoder 和 Decoder 对 Latency sapce 压缩的结构啦。

Pytroch 实现

扩散过程

在扩散过程公式（1）中，其中 \beta 1, \ldots, \beta T 是高斯分布方差的超参数。且实际中 \beta_t 随着 t 增大是递增的，即 \beta_1 < \ldots < \beta_T 。在下面实际代码中，我们使用numpy设置 \beta_t 是从 0.0001 到0.02 线性插值linspace，然后再转换称为Pytroch的tensor。对应时间步T，即num_steps设置为1000。

回顾扩散过程的公式(1)：

q(x_{1:T} \mid x_{0}) := \prod_{t=1}^{T} q(x_{t} \mid x_{t-1}) := \prod_{t=1}^{T} \mathcal{N}(x_{t} ; \sqrt{1-\beta_{t}} x_{t-1}, \beta_{t} \mathbf{I}) \tag{1}

为了方便表达，假设 \alpha_{t}=1-\beta_{t} ，和 \bar{\alpha}{t}=\prod_{s=1}^{t} \alpha_{s} ，那么有：

q(x_{t} \mid x_{0})=\mathcal{N}\left(\mathbf{x}_{t} ; \sqrt{\bar{\alpha}_{t}} \mathbf{x}_{t-1},\left(1-\bar{\alpha}_{t}\right) \mathbf{I}\right) \tag{22}

所以代码中 alphas 即对应 \alpha_{t} ，alphas_prod对应 \bar{\alpha}_{t} 。然后先计算得到公式 (22) 中的变量。

num_steps = 1000
beta = torch.tensor(np.linspace(1e-5, 0.2e-2, num_steps))
alphas = 1 - betas
alphas_prod = torch.cumprod(alphas, 0)
alphas_prod_p = torch.cat([torch.tensor([1]).float(), alphas_prod[:-1]], 0)
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)

既然有了公式(22) 中的变量， q(x_{1:T} \mid x_{0}) 的函数表示就变得比较简单啦，直接使用下面函数 q_x 来表示。

def q_x(x_0, t, noise=None):
    if not noise: noise = torch.randn_like(x_0)
    alphas_t = extract(alphas_bar_sqrt, t, x_0)
    alphas_1_m_t = extract(one_minus_alphas_bar_sqrt, t, x_0)
    return (alphas_t * x_0 + alphas_1_m_t * noise)

正向过程正式的计算比较简单直接，正如上面理论部分提到的，通过时间步 T 在每次马尔科夫链的转换过程对样本数据 dataset 添加噪声：

for i in range(num_steps):
    q_i = q_x(dataset, torch.tensor([i]))

对应公式求解 q(x_{t} \mid x_{0}) 则为：

posterior_mean_coef_1 = (betas * torch.sqrt(alphas_prod_p) / (1 - alphas_prod))
posterior_mean_coef_2 = ((1 - alphas_prod_p) * torch.sqrt(alphas) / (1 - alphas_prod))
posterior_variance = betas * (1 - alphas_prod_p) / (1 - alphas_prod)
posterior_log_variance_clipped = torch.log(torch.cat((posterior_variance[1].view(1, 1), posterior_variance[1:].view(-1, 1)), 0)).view(-1)
def q_posterior_mean_variance(x_0, x_t, t):
    coef_1 = extract(posterior_mean_coef_1, t, x_0)
    coef_2 = extract(posterior_mean_coef_2, t, x_0)
    mean = coef_1 * x_0 + coef_2 * x_t
    var = extract(posterior_log_variance_clipped, t, x_0)
    return mean, var

训练过程

与正向扩散过程不同，逆扩散过程需要训练神经网络模型，这里通过定义损失函数和训练参数，然后进行训练。

损失函数刚才原理部分已经讲过啦，通过最小化数据的负对数似然的变分上限进行求解，也就是对应公式(17)。而网络模型直接使用Unet，这里面就不再展开Unet的编写啦，网上一大堆。

from model import Unet
from ema import EMA
import torch.optim as optim
model = Unet()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# Create EMA model
ema = EMA(0.9)
ema.register(model)
# Batch size
batch_size = 128
for t in range(num_steps):
    # X is a torch Variable
    permutation = torch.randperm(dataset.size()[0])
    for i in range(0, dataset.size()[0], batch_size):
        # Retrieve current batch
        indices = permutation[i:i+batch_size]
        batch_x = dataset[indices]
        # Compute the loss.
        loss = noise_estimation_loss(model, batch_x, alphas_bar_sqrt, one_minus_alphas_bar_sqrt, num_steps)
        # Before the backward pass, zero all of the network gradients
        optimizer.zero_grad()
        # Backward pass: compute gradient of the loss with respect to parameters
        loss.backward()
        # Perform gradient clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
        # Calling the step function to update the parameters
        optimizer.step()
        # Update the exponential moving average
        ema.update(model)
    # Print loss
    if (t % 100 == 0):
        print(loss)

参考公式(17)的实现损失函数，在上面训练的过程中其实没有展开，下面给出简单的代码示例，有空再详细解读，今天有点肚子痛，赶着收工。

def noise_estimation_loss(model, x_0, alphas_bar_sqrt, one_minus_alphas_bar_sqrt, n_steps):
    batch_size = x_0.shape[0]
    # Select a random step for each example
    t = torch.randint(0, n_steps, size=(batch_size // 2 + 1,))
    t = torch.cat([t, n_steps - t - 1], dim=0)[:batch_size].long()
    # x0 multiplier
    a = extract(alphas_bar_sqrt, t, x_0)
    # eps multiplier