生成式深度学习(第二版)-译文-第八章-扩散模型(II)

… … 接上篇博文

逆向扩散过程

现在，让我们再来看一下逆向扩散过程。回忆一下，我们想要的是构建一个神经网络 $p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$, 其能对加噪过程进行反向操作 — 也就是说，近似反向分布$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$。如果我们这么做我们可以从 $\mathcal{N}(0,\mathbf{I})$ 中采样随机噪声，并多次应用此逆向扩散过程，以生成一幅新的图像。整个过程可视化出来如图8-6所示。

逆向扩散过程与变分自编码器中的解码器有诸多相似之处。对两者而言，我们都是期待使用一个神经网络将随机噪声转变为有意义的输出。差别在于: 在VAE中，前向过程(将图像转变为噪声)是模型的一部分(也即，它是学到的)，而在扩散模型中，它是非参数化 (unparameterized) 的。

因此，在扩散模型中应用与VAE相同的损失函数就非常合理了。原始的 DDPM 论文导出了完全一致的损失函数，并表明该损失可通过 训练一个神经网络 ${ϵ}_{\theta }$ 来预测在时间步 $t$ 加入到给定图像${\mathbf{x}}_0$之噪声 $ϵ$ 而得到。

换句话说，我们采样一幅图像 ${\mathbf{x}}_0$，并通过 $t$ 个加噪步骤将之转换到一幅一幅图像 ${\mathbf{x}}_t=\sqrt{{\bar{a}}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{a}}_t}ϵ$。我们将这幅新的 (添加了噪声的) 图像以及噪声水平 ${\bar{a}}_t$ 到神经网络，并要求它针对预测${ϵ}_{\theta }({\mathbf{x}}_t)$和真实$ϵ$的平方误差采取梯度步骤来预测出 $ϵ$。

在下一小节，我们将看看该神经网络的结构。这里值得注意的是，扩散模型实际上维护了神经网络的两个副本: 一是使用梯度下降主动训练的，另一个(EMA网络) 则是一个主动训练网络权重在前序训练步上的指数移动平均(exponential moving average, EMA)。EMA 网络不太容易受到训练过程中的短期波动和峰值的影响，这使得它比主动训练的网络更具有鲁棒性。因此，当我们想要从网络生成输出时我们使用的是EMA网络。

模型的训练过程如下图 8-7 所示。

算法 1 训练
1: repeat
2: ${\mathbf{x}}_0\sim q({\mathbf{x}}_0)$
3: $t\sim Uniform(\{1,\cdots ,T\})$
4: $ϵ\sim \mathcal{N}(0,\mathbf{I})$
5: $Takegradientdescentstepon$
${\nabla }_{\theta }||ϵ-{ϵ}_{\theta }(\sqrt{{\bar{a}}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{a}}_t}ϵ,t)||$
6: until converged

在Keras中，我们可以如样例8-5对训练步进行编码。

# Keras 扩散模型中的 Train_step 函数
class DiffusionModel(models.Model):
    def __init__(self):
        super().__init__()
        self.normalizer = layers.Normalization()
        self.network = unet
        self.ema_network = models.clone_model(self.network)
        self.diffusion_schedule = cosine_diffusion_schedule

    ...

    def denoise(self, noisy_images, noise_rates, signal_rates, training):
        if training:
            network = self.network
        else:
            network = self.ema_network
        pred_noises = network(
            [noisy_images, noise_rates**2], training=training
        )
        pred_images = (noisy_images - noise_rates * pred_noises) / signal_rates

        return pred_noises, pred_images

    def train_step(self, images):
    	# 我们首先将输入批次图像转换为 0 均值 单位方程
        images = self.normalizer(images, training=True) 
        # 下一步，我们采样噪声来匹配输入图像的尺寸
        noises = tf.random.normal(shape=tf.shape(images)) 
        batch_size = tf.shape(images)[0]
        # 我们也对 diffusion times 进行采样
        diffusion_times = tf.random.uniform(
            shape=(batch_size, 1, 1, 1), minval=0.0, maxval=1.0
        ) 
        # 并用diffusion times根据余弦扩散计划来生成 噪声和信号 rates
        noise_rates, signal_rates = self.cosine_diffusion_schedule(
            diffusion_times
        ) 
        # 然后我们在输入图像上使用 信号 和 噪声 权重来生成噪声图像
        noisy_images = signal_rates * images + noise_rates * noises 
        
        with tf.GradientTape() as tape:
        	# 使用提供的 noise_rates和signal_rates 我们对噪声图像进行去噪: 要求网络预测噪声 并 去除噪声操作
            pred_noises, pred_images = self.denoise(
                noisy_images, noise_rates, signal_rates, training=True
            ) 
            # 我们计算预测噪声和真实噪声之间的损失 (均方误差)
            noise_loss = self.loss(noises, pred_noises)  
        gradients = tape.gradient(noise_loss, self.network.trainable_weights)
        # 并对损失函数 take a gradient step
        self.optimizer.apply_gradients(
            zip(gradients, self.network.trainable_weights)
        ) 
        self.noise_loss_tracker.update_state(noise_loss)

        for weight, ema_weight in zip(
            self.network.weights, self.ema_network.weights
        ):
        	# EMA 网络权重更新为 当前EMA权重和训练网络权重的一个加权平均
            ema_weight.assign(0.999 * ema_weight + (1 - 0.999) * weight) 

        return {m.name: m.result() for m in self.metrics}

U-Net 去噪模型

既然我们已经见过需要构建的那种神经网络( 或者 预测添加到给定图像的噪声)，我们可以仔细看看使之成为可能的网络架构。

DDPM 论文的作者使用一种名为*“U-Net”*的架构。网络的框图如下图8-8所示，显式的给出了张量在通过网络时的形状。

与变分自编码器类似，U-Net 包含两半: 一半是降采样，输入图像在空间上压缩，在通道上扩展；另一半是上采样，其中表示在空间上扩展，而通道数则减少。但是，与VAE不同的是，在网络的降采样和上采样部分中，对于同样空间形状的层，U-Net是存在跳跃连接的。VAE是序列的: 数据从网络的输入逐层流向输出。而 U-Net则不同，因为跳跃连接允许信息在网络的一部分上形成短路并直接流向后面的层。

当我们希望输入和输出具有相同形状时，U-Net 尤为有用。在我们的扩散模型样例中，我们希望预测加到图像上的噪声，该噪声与图像本身恰好有一样的形状，因此 U-Net 是网络结构的自然选择。

首先，我们看下 Keras 中构建 U-Net 的代码，如下样例 8-6 所示。

# 样例8-6 Keras 中的 U-Net 模型
# U-Net的第一个输入是我们希望去噪的图像
noisy_images = layers.Input(shape=(64, 64, 3)) 
# 图像通过 Conv2D层来增加通道数
x = layers.Conv2D(32, kernel_size=1)(noisy_images) 
# U-Net的第二个输入是噪声方差 (一个标量)
noise_variances = layers.Input(shape=(1, 1, 1)) 
# 这里用 sinusoidal embedding进行编码
noise_embedding = layers.Lambda(sinusoidal_embedding)(noise_variances) 
# 这个embedding 在空间维度上复制以匹配输入图像尺寸
noise_embedding = layers.UpSampling2D(size=64, interpolation="nearest")(
    noise_embedding
)
# 两个输入流在通道上连接起来
x = layers.Concatenate()([x, noise_embedding]) 
# skips list 会保持 那些我们想连接到下游 UpBlock层 的 DownBlock 层的输出
skips = [] 
# 张量传入一系列 DownBlock层，图像尺寸不断减小，通道数随着增加
x = DownBlock(32, block_depth = 2)([x, skips]) 
x = DownBlock(64, block_depth = 2)([x, skips])
x = DownBlock(96, block_depth = 2)([x, skips])
# 张量传入两个 ResidualBlock 层，保持图像尺寸和通道数不变
x = ResidualBlock(128)(x) 
x = ResidualBlock(128)(x)
# 张量传入一系列 UpBlock层，图像尺寸不断增加，通道数随着减少。跳跃连接接入了前面DownBlock层的输出
x = UpBlock(96, block_depth = 2)([x, skips]) 
x = UpBlock(64, block_depth = 2)([x, skips])
x = UpBlock(32, block_depth = 2)([x, skips])

# 最后的Conv2D层将通道数减少到3 (RGB)
x = layers.Conv2D(3, kernel_size=1, kernel_initializer="zeros")(x) 
# U-Net 是一个 Keras 模型，其接收 噪声图像 及 噪声方差作为输出，并输出预测的噪声map
unet = models.Model([noisy_images, noise_variances], x, name="unet")

为了进一步从细节上理解 U-Net，我们需要深入探索四个概念:

1. 噪声方差的正弦嵌入(the sinusoidal embedding of the noise variance)；
2. 残差模组 (the ResidualBlock)；
3. 下采样模组 (the DownBlock)；
4. 上采样模组 (the UpBlock)。

正弦嵌入 (sinusoidal embedding)

正弦嵌入 首先由 Vaswani 等在一篇论文中引入。我们将使用原始思想的一个改编，其应用于 Mildenhall 等一篇名为《NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis》的论文。

正弦嵌入的核心想法在于: 我们想将一个标量数值 (噪声方差) 转换到一个不同的高维向量，该向量可以为网络的下游提供一个更复杂的表示。原始论文使用这个想法来编码一个句子中单词的离散位置到向量；NeRF 论文将这个思想拓展到连续值。

特别的，一个标量 $x$ 可以使用下面的方程进行编码:
$\gamma (x)=(sin(2\pi e^{0f}x),\cdots ,sin(2\pi e^{(L-1)f}x),cos(2\pi e^{0f}x),\cdots ,cos(2\pi e^{(L-1)f}x))$

其中，我们选择 $L=16$ （我们期望噪声嵌入长度的一半），$f=\frac{ln(1000)}{L-1}$ 作为频率的最大放缩因子。

这产生了如下图 8-9 所示的嵌入模式。

如样例 8-7所示，我们可以对正弦嵌入函数编码实现。这个将单一噪声方差标量值转化为一个长为 32 的向量。

# 编码噪声方差的 sinusoidal_embedding 函数
def sinusoidal_embedding(x):
    frequencies = tf.exp(
        tf.linspace(
            tf.math.log(1.0),
            tf.math.log(1000.0),
            16,
        )
    )
    angular_speeds = 2.0 * math.pi * frequencies
    embeddings = tf.concat(
        [tf.sin(angular_speeds * x), tf.cos(angular_speeds * x)], axis=3
    )
    return embeddings

残差模组 (ResidualBlock)

降采样模组 和 上采样模组 中都包含有残差模组层，所以我们从这里开始。在第五章里，当我们构建PixelCNN时，我们已经探索了残差模组，但这里我们为了完整性再次回顾下。

一个残差模组是一组这样的层: 其包含从输入连向输出的跳跃连接。残差模块帮助我们构建 可以学习更复杂模式而无需遭受梯度消失/退化问题的更深网络。
梯度消失问题: 当网络加深时，在更深层传播的梯度非常小，从而让学习变得很慢。
梯度退化问题: 当网络加深时，它们反而不如对应的前层网络精准 — 精度似乎在一定的深度上出现饱和，然后快速退化。

退化(degradation)
退化问题某种意义上是反直觉的，但是却能在实际中观察到。因为更深的层理论上至少应该学会identity mapping，这是有意义的 — 尤其是考虑到更深网络面临的其它问题，如梯度消失问题。

该问题的解法非常简单，最早由何恺明等在2015年的 ResNet 论文中提出。通过为主要的权重层引入一个跳跃连接通路，模块可以选择跳过复杂的网络权重更新，而直接进行identity mapping。这让网络在不牺牲梯度尺寸或网络精度的同时可以训练的更深。

残差模块的框图如图8-10所示。注意在一些残差模块中，我们也在跳跃连接上 引入了额外的 kernel size 为 1的Conv2D层，从而让通道数目与模块的其余部分保持一致。

如下样例8-8所示，我们可以在Keras中编码实现一个 ResidualBlock。

# U-Net 中残差模组之代码实现
def ResidualBlock(width):
    def apply(x):
        input_width = x.shape[3]
        # 检查输入的通道数，看看它是否与我们想要模块输出的通道数匹配，如不是，则在跳跃连接上引入一个额外的 Conv2D层来将通道数与模块其余的部分对齐
        if input_width == width: 
            residual = x
        else:
        	# 应用 BN 层
            residual = layers.Conv2D(width, kernel_size=1)(x)
        x = layers.BatchNormalization(center=False, scale=False)(x) 
        # 应用两个 Conv2D 层
        x = layers.Conv2D(
            width, kernel_size=3, padding="same", activation=activations.swish
        )(x) 
        x = layers.Conv2D(width, kernel_size=3, padding="same")(x)
        # 将原始模块的输入添加到输出来提供最终的模块输出
        x = layers.Add()([x, residual]) 
        return x

    return apply

下采样模组 和 上采样模组

每个连续的下采样模组通过 block_depth (在我们的例子中 = 2) 个残差模组 来增加通道数，同时也最终接入一个 AveragePooling2D层来将图像尺寸减半。每个残差模组都被加入到一个列表里，以便通过U-Net跳跃连接直连的对应上采样层能够利用起来。

一个上采样模组首先应用 UpSampling2D 层来将尺寸通过双线性插值翻倍。每个连续的上采样模组通过 block_depth (在我们的例子中 = 2) 个残差模组 来减小通道数，同时也将通过U-Net跳跃连接直连的下采样模组的输出连接起来。这一过程的框图如下图8-11所示。

在Keras中，我们将下采样和上采样模块用样例8-9编码如下。

def DownBlock(width, block_depth):
    def apply(x):
        x, skips = x
        for _ in range(block_depth):
        	# 下采样模组使用 给定宽度的 残差模组来增加通道数
            x = ResidualBlock(width)(x) 
            # 保存到一个skips列表中，便于后续上采样模组的使用
            skips.append(x) 
        # AveragePooling2D层降低图像本身的维度到一半
        x = layers.AveragePooling2D(pool_size=2)(x) 
        return x

    return apply

def UpBlock(width, block_depth):
    def apply(x):
        x, skips = x
        # 上采样模组以 UpSampling层开始，对图像尺寸翻倍
        x = layers.UpSampling2D(size=2, interpolation="bilinear")(x) 
        for _ in range(block_depth):
         	# 下采样模组的输出通过 Concatenate层与当前输出相连
            x = layers.Concatenate()([x, skips.pop()]) 
            # 当传入上采样模组时，残差模组用来减低图像中的通道数
            x = ResidualBlock(width)(x) 
        return x

    return apply

扩散模型之训练

现在，我们集齐了训练去噪扩散模型的所有元素！样例8-10 创建、编译 并训练了一个扩散模型。

# 初始化模型
model = DiffusionModel() 
# 模型编译，使用AdamW优化器 (与Adam类似，但是包含weight decay 用以稳定训练过程)，另外，其损失函数为均方误差损失函数
model.compile(
    optimizer=optimizers.experimental.AdamW(learning_rate=1e-3, weight_decay=1e-4),
    loss=losses.mean_absolute_error,
) 
# 使用训练集计算归一化统计
model.normalizer.adapt(train) 
# 使用50 epochs来拟合模型
model.fit(
    train,
    epochs=50,
)

损失曲线 (噪声均方误差 [MAE] ) 如下图8-12所示。

从去噪扩散模型 (DDM)采样

为了从我们训练的模型采样图像，我们需要应用逆扩散过程 — 也即，我们需要从随机噪声出发，使用模型来逐渐去除噪声，直到我们得到一幅可识别的花朵图片。

我们脑中要有这样的概念: 我们的模型是训练来预测对训练集中一幅给定图像所施加的总噪声量的，而不仅仅是是我们在加噪过程的最后一个时间步所加的噪声。 然而，我们不希望一次性去除噪声 — 从纯噪声一步预测图像显然不能work！我们将模仿前向过程，逐渐的一小步一小步来去除预测噪声，从而允许模型依据自己的预测进行调整。

为了实现这一点，我们可以从 $x_t$ 用两步跳到 $x_{t-1}$ — 首先使用我们的模型噪声预测来计算原始图像 $x_0$的一个估计，然后通过再次在此图像上运用预测噪声，但是仅仅只是 $t-1$步，来产生 $x_{t-1}$。这一思想如图 8-13所示。

如果我们多步重复这一过程，我们最终将得到$x_0$的一个估计，该估计是通过多个小的步骤逐渐引导得到的。实际上，我们可以自由选择所需要采用的步数，关键的是，它并不需要像加噪训练过程的步数那么大 (比如，1000)。这里可以小的多 — 在本例中我们选择 20。

下面的方程 (Song 等于2020年提出) 从数学上描述了这一过程:

${\mathbf{x}}_{t-1}=\sqrt{a_{t-1}}(\frac{{\mathbf{x}}_t-\sqrt{1-a_t}{ϵ}_{\theta }^{(t)}({\mathbf{x}}_t)}{\sqrt{a_t}})+\sqrt{1-a_{t-1}-{\sigma }_t^2}\cdot {ϵ}_{\theta }^{(t)}({\mathbf{x}}_t)+{\sigma }_t{ϵ}_t$

让我们对这一公式进行拆解: 公式右边括号内的第一项是估计的图像$x_0$, 使用网络预测的噪声 ${ϵ}_{\theta }^t$ 计算得到。然后，我们对它按照 $t-1$ signal rate $\sqrt{a_{t-1}}$ 进行放缩，我们再次应用预测噪声，但是这次是用 $t-1$ noise rate $\sqrt{1-a_{t-1}-{\sigma }_t^2}$ 进行放缩。另外，还添加了额外的高斯随机噪声 ${\sigma }_t{ϵ}_t$, 因子${\sigma }_t$ 用来觉得我们希望生成过程有多随机。

特例 对所有的t，${\sigma }_t=0$ 对应于一种名为 去噪扩散隐模型 (Denoising Diffusion Implicit Model, DDIM)，最早由 Song 等在2020年提出。对于一个DDIM，生成过程是完全确定的 — 也即，相同的随机噪声输入将总是得到想通的输出。这种特性是理想的，因为我们有了一个从隐空间样本到像素空间生成输出的定义良好的映射。

在我们的例子中，我们实现了一个 DDIM，使得我们的生成过程完全确定。DDIM 采样过程 (逆扩散) 的代码如下样例 8-11所示。

# 样例8-11 从扩散模型进行采样
class DiffusionModel(models.Model):
...
    def reverse_diffusion(self, initial_noise, diffusion_steps):
        num_images = initial_noise.shape[0]
        step_size = 1.0 / diffusion_steps
        current_images = initial_noise
        # 固定数目的steps，例如20
        for step in range(diffusion_steps): 
        	# diffusion times 都设置为1 （也即，逆扩散过程的开始）
            diffusion_times = tf.ones((num_images, 1, 1, 1)) - step * step_size 
            # noise 和 signal rates 根据扩散计划进行计算
            noise_rates, signal_rates = self.diffusion_schedule(diffusion_times) 
            # 使用 U-Net预测噪声，这允许我们计算去噪图像估计
            pred_noises, pred_images = self.denoise(
                current_images, noise_rates, signal_rates, training=False
            ) 
            # 扩散次数每步递减
            next_diffusion_times = diffusion_times - step_size 
            # 计算新的 noise 和 signal rates
            next_noise_rates, next_signal_rates = self.diffusion_schedule(
                next_diffusion_times
            ) 
            # 通过再次应用预估噪声得到 t-1 图像，根据 t-1 扩散计划 rates
            current_images = (
                next_signal_rates * pred_images + next_noise_rates * pred_noises
            ) 
        # 20步以后，返回最终预测图像 $\mathbf{x}_0$
        return pred_images

扩散模型分析

现在，让我们看看使用所训练模型的三个不同方式: 生成新的图像，测试逆向扩散步数对生成质量的影响，以及在隐空间的两幅图像间进行插值。

生成图像

为了从我们训练的模型生成样本，我们可以简单运行逆向扩散过程，确保我们最终对输出进行了去归一化 (denormalize，也即，让像素值回到 [0,1] 范围)。我们可以通过使用样例8-12 中的代码实现这点。

class DiffusionModel(models.Model):
...
    def denormalize(self, images):
    	# 生成一些初始的噪声maps
        images = self.normalizer.mean + images * self.normalizer.variance**0.5 
        return tf.clip_by_value(images, 0.0, 1.0)

    def generate(self, num_images, diffusion_steps):
        initial_noise = tf.random.normal(shape=(num_images, 64, 64, 3)) 
        # 应用逆扩散过程
        generated_images = self.reverse_diffusion(initial_noise, diffusion_steps) 
        # 网络图像的输出均值为1，单位方差，因此我们需要通过利用训练数据的均值和方差对它进行去归一化
        generated_images = self.denormalize(generated_images) 
        return generated_images

如图8-14， 我们可以看到扩散模型在训练过程的不同epochs上生成的样本。

调整扩散步数

我们也可以测试在逆扩散过程中调整扩散步数是如何影响图像质量的。直觉上，采用越多的扩散步数，生成图像的质量越高。

如图8-15，我们可以看到生成质量确实随着扩散步数的增加而提升。从最开始的取样噪声出发，模型只能预测朦胧的色块。随着步数更多，模型可以逐步对生成对象进行细化和锐化。但是，生成图像所需的时间跟扩散步数是线性相关的，因此这里面有个折中。我们看到，对扩散步数 20 和 100 而言，提升是很小的，因此本例中我们选择 20 作为质量和速度的合理妥协。

图像插值

最后，如我们之前在VAEs中所见，我们可以对高斯隐空间进行点之间的插值，以使得像素空间中的图像变化平缓。这里，我们选择使用一种球形插值来确保当两种高斯噪声maps混合在一起时，方差依然为常量。特别的，在每个时间步，初始的噪声map 定义为 $asin(\frac{\pi }{2}t)+bcos(\frac{\pi }{2}t)$,其中$t$在0到1范围内平滑分布，$a$ 和 $b$ 则是我们希望进行内插的两个随机采样的噪声张量。

结果图像如图8-16所示。

本章小结

在本章中，我们探索了当下最令人兴奋、最有前景的生成式建模技术: 扩散模型。尤其的，我们使用了生成式扩散模型领域一篇关键论文(Ho 等于2020年发表)的核心思想，该论文发明了原始的去噪扩散概率模型(DDPM)。然后，我们利用去噪隐模型(DDIM)论文的思想对其进行拓展，使得生成过程完全确定。

我们已经看到扩散模型是如何由前向扩散 和 逆向扩散过程组成。前向扩散过程通过一系列的小步给训练数据添加噪声，而逆向扩散过程包含了一个努力预测所添加噪声的模型。

我们使用重参数化技巧来计算前向过程任何时间步的加噪图像，同时无需真的经历诸多的加噪步骤。我们看到用以加噪的参数任务选择在整个模型的成功中扮演重要的角色。

逆向扩散过程由 U-Net 组成，给定每步的加噪图像和noise rate，它在每个时间步尝试预测噪声。一个 U-Net包含下采样模组来在增加通道数的同时减小图片尺寸，以及一个上采样模组来在减小通道数的同时扩大图片尺寸。Noise rate则使用正弦嵌入进行编码。

从扩散模型采样也需要通过一系列步骤完成。对于一个给定的加噪图像，我们用U-Net预测所添加的噪声，然后该预测被用于计算原始图像的一个估计。我们以一个小点的noise rate应用预测噪声。这一过程在一系列步骤中重复 (步骤数目显著少于训练所需的步数)，这可以使得我们从一个标准高斯噪声分布的随机采样点出发获得最终的生成。

我们看到逆向过程中扩散步数的增大会提升生成图像质量，代价是推理速度的下降。我们也在隐空间进行了算术运算以进行两幅图像的内插。

译者注:
至此，关于生成式建模的主要技术细节已经介绍完成，读者可以通过以下链接进行相关技术的回溯:
[1] 生成式建模概述
[2] Transformer I，Transformer II
[3] 变分自编码器
[4] 生成对抗网络，高级生成对抗网络 I，高级生成对抗网络 II
[5] 自回归模型
[6] 归一化流模型
[7] 基于能量的模型
[8] 扩散模型 I, 扩散模型 II