AIGC实战——扩散模型(Diffusion Model)

AIGC实战——扩散模型

0. 前言

与生成对抗网络 (Generative Adversarial Network, GAN)一样，扩散模型是过去十年中最有影响力的生成模型技术之一。在许多基准测试中，当前的扩散模型已经超过了以往最先进的 GAN 模型，并迅速成为生成模型的首选。扩散模型的核心理念与其他生成模型，例如去噪自编码器、能量模型等有诸多相似之处。事实上，扩散来源于热力学扩散。同时，基于评分的生成模型领域(即能量模型)也取得了重要的进展，其直接估计对数分布的梯度(也称为评分函数)，以训练模型。噪声条件得分网络 (Noise Conditional Score Network, NCSN) 使用多尺度噪声扰动应用于原始数据，以确保模型在低数据密度区域具有良好的性能表现。
扩散模型 (Diffusion Model) 在先前模型的基础上，揭示了扩散模型与基于评分的生成模型之间的联系，并训练了一个能够在多个数据集上与 GAN 相媲美的扩散模型，称为去噪扩散概率模型 (Denoising Diffusion Probabilistic ModelD, DDPM)。本节将介绍去噪扩散概率模型的工作原理。然后，学习如何使用 Keras 构建去噪扩散概率模型。

1. 去噪扩散概率模型

去噪扩散概率模型 (Denoising Diffusion Probabilistic Model, DDPM) 的核心思想是通过一系列小步骤训练一个深度学习模型去除图像中的噪声。如果我们从完全随机的噪声开始，理论上我们能够不断应用该模型，直到获得一幅看上去就像是从训练集中采样出来的图像。
首先，我们准备用于训练去噪扩散概率模型的数据集，然后分别介绍正向(加噪)和逆向(去噪)扩散过程。

1.1 Flowers 数据集

为了训练去噪扩散概率模型，使用 Kaggle 中的花卉数据集 Oxford 102 Flower，其中包含 8000 多张各种花卉的彩色图像。下载数据集后，将花卉图像解压并保存到 ./data 文件夹中。
使用 Keras 的 image_dataset_from_directory 函数加载图像，将图像尺寸调整为 64×64，并将像素值缩放到 [0,1] 范围内。将数据集重复五次，以增加训练时长，并将数据分成组进行批处理，其中每组包含 64 张图像。

# 使用 Keras 的 image_dataset_from_directory 函数加载数据集
train_data = utils.image_dataset_from_directory(
    "./data/dataset/train",
    labels=None,
    image_size=(IMAGE_SIZE, IMAGE_SIZE),
    batch_size=None,
    shuffle=True,
    seed=42,
    interpolation="bilinear",
)
# 将像素值缩放到[0, 1]范围内
def preprocess(img):
    img = tf.cast(img, "float32") / 255.0
    return img


train = train_data.map(lambda x: preprocess(x))
# 将数据集重复五次
train = train.repeat(DATASET_REPETITIONS)
# 将数据集分成组进行批处理，其中每组 64 张图像
train = train.batch(BATCH_SIZE, drop_remainder=True)

数据集中的示例图像如下图所示。

获取数据集后，我们继续介绍如何使用正向扩散过程向图像添加噪声。

1.2 正向扩散过程

假设有一幅图像 x 0 x_0 x0​，我们希望通过多个步骤(比如 T = 1 , 000 T=1,000 T=1,000 )逐渐破坏此图像，使得最终图像无法与标准的高斯噪声区分开(即 x T x_T xT​ 的均值为 0，方差为 1)。
可以定义一个函数 q q q，将方差为 β t β_t βt​ 的少量高斯噪声添加到图像 x t − 1 x_{t-1} xt−1​ 中，生成新图像 x t x_t xt​。如果我们反复应用这个函数，就会生成一系列噪声逐渐增加的图像 ( x 0 , . . . , x T ) (x_0, ..., x_T) (x0​,...,xT​)，如下图所示。

使用数学公式表示这个更新过程(其中， ϵ t − 1 ϵ_{t-1} ϵt−1​ 是均值为 0，方差为 1 的标准高斯分布)：
x t = 1 − β t x t − 1 + β t ϵ t − 1 x_t = \sqrt{1 - β_t} x_{t-1} + \sqrt{β_t} ϵ_{t-1} xt​=1−βt​ ​xt−1​+βt​ ​ϵt−1​
需要注意的是，我们还对输入图像 x t − 1 x_{t-1} xt−1​ 进行了缩放，以确保输出图像 x t x_t xt​ 的方差随时间保持不变。这样，如果我们将原始图像 x 0 x_0 x0​ 归一化为均值为 0 、方差为 1，那么当 T T T 足够大时，通过归纳法， x T x_T xT​ 将逼近标准高斯分布。
假设 x t − 1 x_{t-1} xt−1​ 均值为 0、方差为 1，那么 1 − β t x t − 1 \sqrt {1 - β_t} x_{t-1} 1−βt​ ​xt−1​ 的方差将为 1 − β t 1 - β_t 1−βt​，而 β t ϵ t − 1 \sqrt {β_t} ϵ_{t-1} βt​ ​ϵt−1​ 的方差将为 β t β_t βt​，根据方差的规则 V a r ( a X ) = a 2 V a r ( X ) Var(aX) = a^2 Var(X) Var(aX)=a2Var(X)。将它们相加，得到一个新的分布 x t x_t xt​，它的均值为 0，方差为 1 − β t + β t = 1 1 - β_t + β_t = 1 1−βt​+βt​=1，根据方差的规则 V a r ( X + Y ) = V a r ( X ) + V a r ( Y ) Var(X+Y) = Var(X) + Var(Y) Var(X+Y)=Var(X)+Var(Y) (其中 X X X 和 Y Y Y 是独立的)。因此，如果将 x 0 x_0 x0​ 归一化为均值为 0、方差为1，那么对于所有的 x t x_t xt​，包括最后的图像 x T x_T xT​，我们可以保证它们也满足这个条件，即逼近标准高斯分布。这样，我们就能够轻松地对 x T x_T xT​ 进行采样，然后通过训练好的神经网络模型应用逆扩散过程。换句话说，我们的正向添加噪声过程 q q q 也可以改写为：
q ( x t ∣ x t − 1 ) = N ( x t ; 1 − β t x t − 1 , β t I ) q(xt|x_{t-1}) = \mathcal N(x_t; \sqrt{1 - β_t} xt-1, β_t\mathbf I) q(xt∣xt−1​)=N(xt​;1−βt​ ​xt−1,βt​I)
其中 N \mathcal N N 表示高斯分布， I \mathbf I I 表示单位矩阵。

1.3 重参数化技巧

重新参数化技巧 (Reparameterization Trick) 是一种在不需要经过 t t t 次应用 q q q 的情况下，直接从图像 x 0 x_0 x0​ 跳转到任意噪声版本图像 x t x_t xt​ 的方法。
如果我们定义 α t = 1 − β t α_t=1-β_t αt​=1−βt​ 和 α ‾ t = ∏ i = 1 t α i \overline α_t=∏_{i=1}^tα_i αt​=∏i=1t​αi​，那么得到以下形式：
x t = α t x t − 1 + 1 − α t ϵ t − 1 = α t α t − 1 x t − 2 + 1 − α t α t − 1 ϵ t − 2 = . . . = α ‾ t x 0 + 1 − α ‾ t ϵ \begin{equation*} \begin{aligned} x_t &= \sqrt{α_t}x_{t-1} + \sqrt {1-α_t}ϵ_{t-1} \\ &= \sqrt{α_tα_{t-1}}x{t-2} + \sqrt{1-α_tα_{t-1}}ϵ_{t-2} \\ &=...\\ &= \sqrt{\overline \alpha_t}x_0+\sqrt{1-\overline\alpha_t}ϵ \\ \end{aligned} \end{equation*} xt​​=αt​ ​xt−1​+1−αt​ ​ϵt−1​=αt​αt−1​ ​xt−2+1−αt​αt−1​ ​ϵt−2​=...=αt​ ​x0​+1−αt​ ​ϵ​​
需要注意的是，根据定理：两个高斯分布相加得到新的高斯分布。因此，我们可以从原始图像 x 0 x_0 x0​ 跳转到前向扩散过程的任何步骤 x t x_t xt​。此外，我们可以使用 α ‾ t \overline \alpha_t αt​ 的值来定义扩散规划 (diffusion schedule)，而不是使用原始的 β t β_t βt​ 值，其中 α ‾ t \overline \alpha_t αt​ 可以表示为与信号(原始图像 x 0 x_0 x0​ )相关的方差，而 1 − α ‾ t 1-\overline α_t 1−αt​ 则是与噪声 ( ϵ ϵ ϵ) 相关的方差。
因此，正向扩散过程 q q q 也可以表达为：
q ( x t ∣ x 0 ) = N ( x t ; α ‾ t x 0 , ( 1 − α ‾ t ) I ) q(x_t|x_0) = \mathcal N(x_t; \sqrt{\overline α_t}x_0, (1-\overline α_t)\mathbf I) q(xt​∣x0​)=N(xt​;αt​ ​x0​,(1−αt​)I)

1.4 扩散规划

需要注意的是，我们可以自由地在每个时间步选择不同的 β t β_t βt​，它们不必全部相等。关于 β t β_t βt​ (或 α ‾ t \overline α_t αt​ )值如何随时间的变化称为扩散规划 (Di€usion Schedule)。
可以采用线性扩散规划 (linear diffusion schedule) 来定义 β t β_t βt​，即 β t β_t βt​ 随着 t t t 线性增加，从 β 1 = 0.0001 β_1=0.0001 β1​=0.0001 增加到 β T = 0.02 β_T=0.02 βT​=0.02。这可以确保在噪声处理的早期阶段，我们采取较小的噪声步长，而在图像已经包含大量噪声的后期阶段，采取较大的噪声步长。使用 Keras 实现线性扩散规划。

def linear_diffusion_schedule(diffusion_times):
    min_rate = 0.0001
    max_rate = 0.02
    betas = min_rate + diffusion_times * (max_rate - min_rate)
    alphas = 1 - betas
    alpha_bars = tf.math.cumprod(alphas)
    signal_rates = tf.sqrt(alpha_bars)
    noise_rates = tf.sqrt(1 - alpha_bars)
    return noise_rates, signal_rates

T = 1000
# 扩散时间是在 0 和 1 之间相等间隔的步长
diffusion_times = tf.convert_to_tensor([x / T for x in range(T)])
#  将线性扩散规划应用于扩散时间，以获得噪声和信号速率
linear_noise_rates, linear_signal_rates = linear_diffusion_schedule(
    diffusion_times
)

实践证明，余弦扩散规划优于线性规划，余弦规划根据以下公式计算 α ‾ t \overline α_t αt​ 值：
α ‾ t = c o s 2 ( t T ⋅ π 2 ) \overline α_t = cos^2(\frac t T·\frac π 2) αt​=cos2(Tt​⋅2π​)
因此，更新后的方程如下(使用三角恒等式 c o s 2 ( x ) + s i n 2 ( x ) = 1 cos^2(x) + sin^2(x) = 1 cos2(x)+sin2(x)=1)：
x t = c o s ( t T ⋅ π 2 ) x 0 + s i n ( t T ⋅ π 2 ) ϵ x_t = cos(\frac t T·\frac π 2)x_0 + sin(\frac t T·\frac π 2)ϵ xt​=cos(Tt​⋅2π​)x0​+sin(Tt​⋅2π​)ϵ
以上方程是原始余弦扩散规划的简化版本，可以在其中添加偏移项和比例因子，以防止扩散过程开始时的噪声步长过小。使用 Keras 实现余弦和偏移余弦扩散规划。

# 余弦扩散规划(不带偏移和比例因子)
def cosine_diffusion_schedule(diffusion_times):
    signal_rates = tf.cos(diffusion_times * math.pi / 2)
    noise_rates = tf.sin(diffusion_times * math.pi / 2)
    return noise_rates, signal_rates
#  偏移余弦扩散规划调整规划以确保在噪声处理开始时噪声步长不会太小
def offset_cosine_diffusion_schedule(diffusion_times):
    min_signal_rate = 0.02
    max_signal_rate = 0.95
    start_angle = tf.acos(max_signal_rate)
    end_angle = tf.acos(min_signal_rate)
    diffusion_angles = start_angle + diffusion_times * (end_angle - start_angle)
    signal_rates = tf.cos(diffusion_angles)
    noise_rates = tf.sin(diffusion_angles)
    return noise_rates, signal_rates

cosine_noise_rates, cosine_signal_rates = cosine_diffusion_schedule(diffusion_times)
(offset_cosine_noise_rates, offset_cosine_signal_rates,) = offset_cosine_diffusion_schedule(diffusion_times)

可以计算每个 t t t 的 α ‾ t \overline α_t αt​ 值，以获取在线性、余弦和偏移余弦扩散规划的每个时间步中通过的信号 ( α ‾ t \overline α_t αt​) 和噪声 ( 1 − α ‾ t 1-\overline α_t 1−αt​）的量，如下图所示。

需要注意的是，在余弦扩散规划中，噪声水平的增加速度比线性扩散规划更慢。余弦扩散规划比线性扩散规划更平滑地向图像添加噪声，这提高了训练效率和生成质量。使用线性和余弦扩散规划在途中添加噪声的效果如下图所示：

1.5 逆向扩散过程

接下来，我们继续介绍逆向扩散过程，构建神经网络 p θ ( x t − 1 ∣ x t ) p_θ(x_{t-1}|x_t) pθ​(xt−1​∣xt​) 可以移除噪声处理，即近似逆向分布 q ( x t − 1 ∣ x t ) q(x_{t-1}|x_t) q(xt−1​∣xt​)。如果我们能够构建此模型，就可以从 N ( 0 , I ) \mathcal N(0,I) N(0,I) 中采样随机噪声，然后多次应用逆向扩散过程，以生成一幅新的图像。

逆向扩散过程和变分自编码器 (Variational Autoencoder，VAE) 的解码器之间有许多相似之处，这两者都旨在使用神经网络将随机噪声转化为有意义的输出。扩散模型和 VAE 之间的区别在于，在 VAE 中，正向过程(将图像转换为噪声)是模型的一部分(即它是可学习的)，而在扩散模型中并没有参数化。
因此，可以在扩散模型中应用与变分自编码器相同的损失函数，原始的 DDM 论文推导了此损失函数的确切形式，并且表明，通过训练一个网络 ϵ θ ϵ_θ ϵθ​ 来预测已添加到给定图像 x 0 x_0 x0​ 的噪声 ϵ ϵ ϵ，即可以优化此损失函数。
换句话说，我们对图像 x 0 x_0 x0​ 进行采样，并通过 t t t 个噪声步骤将其转换为图像 x t = α ‾ t x 0 + ( 1 − α ‾ t ) ϵ x_t=\sqrt {\overline α_t}x_0 + \sqrt {(1-\overline α_t)}ϵ xt​=αt​ ​x0​+(1−αt​) ​ϵ。我们将这个新图像和噪声率 α ‾ t \overline α_t αt​ 提供给神经网络，并要求它预测 ϵ ϵ ϵ，计算预测 ϵ θ ( x t ) ϵθ_(x_t) ϵθ(​xt​) 与真实 ϵ ϵ ϵ 之间的平方误差的梯度，并使用梯度下降法进行优化。
需要注意的是，扩散模型实际上维护了网络的两个副本：一个是使用梯度下降主动训练的，另一个是指数移动平均 (Exponential Moving Average, EMA) 的网络，它是在先前的训练步骤中对主动训练网络权重进行的平均。EMA 网络不易受到训练过程中的短期波动和峰值的影响，因此在生成方面比被主动训练的网络更加稳健。因此，每当我们想要从网络中生成输出时，都会使用 EMA 网络。模型的训练过程如下所示。

使用 Keras 实现上述训练步骤。

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 = offset_cosine_diffusion_schedule

    def compile(self, **kwargs):
        super().compile(**kwargs)
        self.noise_loss_tracker = metrics.Mean(name="n_loss")

    @property
    def metrics(self):
        return [self.noise_loss_tracker]

    def denormalize(self, images):
        images = self.normalizer.mean + images * self.normalizer.variance**0.5
        return tf.clip_by_value(images, 0.0, 1.0)

    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 reverse_diffusion(self, initial_noise, diffusion_steps):
        num_images = initial_noise.shape[0]
        step_size = 1.0 / diffusion_steps
        current_images = initial_noise
        for step in range(diffusion_steps):
            diffusion_times = tf.ones((num_images, 1, 1, 1)) - step * step_size
            noise_rates, signal_rates = self.diffusion_schedule(diffusion_times)
            pred_noises, pred_images = self.denoise(current_images, noise_rates, signal_rates, training=False)
            next_diffusion_times = diffusion_times - step_size
            next_noise_rates, next_signal_rates = self.diffusion_schedule(next_diffusion_times)
            current_images = (next_signal_rates * pred_images + next_noise_rates * pred_noises)
        return pred_images

    def generate(self, num_images, diffusion_steps, initial_noise=None):
        if initial_noise is None:
            initial_noise = tf.random.normal(shape=(num_images, IMAGE_SIZE, IMAGE_SIZE, 3))
        generated_images = self.reverse_diffusion(initial_noise, diffusion_steps)
        generated_images = self.denormalize(generated_images)
        return generated_images

    def train_step(self, images):
        # 首先对图像进行归一化处理，使其均值为 0，方差为 1
        images = self.normalizer(images, training=True)
        # 采样噪声以匹配输入图像的形状
        noises = tf.random.normal(shape=(BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3))
        # 采样随机扩散时间
        diffusion_times = tf.random.uniform(shape=(BATCH_SIZE, 1, 1, 1), minval=0.0, maxval=1.0)
        # 根据余弦扩散规划生成噪声和信号率
        noise_rates, signal_rates = self.diffusion_schedule(diffusion_times)
        # 将信号和噪声权重应用于输入图像，以生成带有噪声的图像
        noisy_images = signal_rates * images + noise_rates * noises
        with tf.GradientTape() as tape:
            # 通过要求网络预测噪声并撤消加噪操作(使用提供的噪声率和信号率)，对带噪声的图像进行去噪
            pred_noises, pred_images = self.denoise(noisy_images, noise_rates, signal_rates, training=True)
            # 计算预测噪声与真实噪声之间的损失(平均绝对误差)
            noise_loss = self.loss(noises, pred_noises)  # used for training
        gradients = tape.gradient(noise_loss, self.network.trainable_weights)
        # 针对损失函数使用梯度下降优化网络权重
        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(EMA * ema_weight + (1 - EMA) * weight)
        return {m.name: m.result() for m in self.metrics}

    def test_step(self, images):
        images = self.normalizer(images, training=False)
        noises = tf.random.normal(shape=(BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3))
        diffusion_times = tf.random.uniform(shape=(BATCH_SIZE, 1, 1, 1), minval=0.0, maxval=1.0)
        noise_rates, signal_rates = self.diffusion_schedule(diffusion_times)
        noisy_images = signal_rates * images + noise_rates * noises
        pred_noises, pred_images = self.denoise(noisy_images, noise_rates, signal_rates, training=False)
        noise_loss = self.loss(noises, pred_noises)
        self.noise_loss_tracker.update_state(noise_loss)

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

2. U-Net 去噪模型

接下来，我们使用 U-Net 作为去噪模型架构，以预测添加到给定图像中的噪声。

2.1 U-Net 架构

DDPM 使用了 U-Net 体系结构，下图展示了该网络的结构，并给出了通过网络时张量的形状。

与变分自编码器 (Variational Autoencoder, VAE) 类似，U-Net 由两部分组成：1) 下采样部分，输入图像在空间上逐渐缩小，但通道逐渐增加；2) 上采样部分，潜表示在空间上逐渐扩大，而通道数逐渐减少。然而，与 VAE 不同的是，网络的上采样和下采样部分之间还存在跳跃连接。VAE 是顺序的，数据从输入流经网络传递到输出，一层接一层；而 U-Net 不同，因为跳跃连接允许信息绕过网络的某些部分直接流向后面的网络层。
在 U-Net 架构中，输出与输入具有相同的形状，在扩散模型中，添加到图像中的噪声与图像本身的形状完全相同，因此 U-Net 成为去噪扩散概率模型网络架构的自然选择。
使用 Keras 中构建 U-Net 架构。

# U-Net 的第一个输入是我们希望去噪的图像
noisy_images = layers.Input(shape=(IMAGE_SIZE, IMAGE_SIZE, 3))
# 图像经过一个 Conv2D 层，增加通道的数量
x = layers.Conv2D(32, kernel_size=1)(noisy_images)
# U-Net 的第二个输入是噪声方差(标量值)
noise_variances = layers.Input(shape=(1, 1, 1))
# 使用正弦嵌入对其进行编码
noise_embedding = layers.Lambda(sinusoidal_embedding)(noise_variances)
# 将该嵌入复制到空间维度上，以匹配输入图像的大小
noise_embedding = layers.UpSampling2D(size=IMAGE_SIZE, interpolation="nearest")(noise_embedding)
# 两个输入流在通道维度上进行连接
x = layers.Concatenate()([x, noise_embedding])
# skips 列表用于保存连接到下游 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 模型，以噪声图像和噪声方差作为输入，并输出预测的噪声图像
unet = models.Model([noisy_images, noise_variances], x, name="unet")

为了更深入地理解 U-Net，我们还需要了解四个概念：噪声方差的正弦嵌入，ResidualBlock，DownBlock 和 UpBlock。

2.2 正弦嵌入

正弦嵌入 (Sinusoidal embedding) 最早由 Vaswani 等人提出，其核心思想是将标量值(噪声方差)转换为一个独特的高维向量，能够在网络下游提供更复杂的表示。原始论文根据这一思想将句子中的离散位置编码为向量；NeRF 将其扩展到连续值。
具体来说，一个标量值 x x x 编码如下：
γ ( x ) = ( s i n ( 2 π e 0 f x ) , . . . , s i n ( 2 π e ( L − 1 ) f x ) , c o s ( 2 π e 0 f x ) , . . . , c o s ( 2 π e ( L − 1 ) f x ) ) γ(x) = (sin(2πe^{0f}x), ..., sin(2πe^{(L-1)f}x), cos(2πe^{0f}x), ..., cos(2πe^{(L-1)f}x)) γ(x)=(sin(2πe0fx),...,sin(2πe(L−1)fx),cos(2πe0fx),...,cos(2πe(L−1)fx))
其中使用 L = 16 L=16 L=16 作为所需噪声嵌入长度的一半， f = l n ( 1000 ) ( L − 1 ) f=\frac {ln(1000)}{(L-1)} f=(L−1)ln(1000)​ 作为频率的最大缩放因子，在这种情况下，嵌入模式如下图所示：

编写这个嵌入函数，将一个噪声方差标量值转换为长度为 32 的向量。

def sinusoidal_embedding(x):
    frequencies = tf.exp(
        tf.linspace(
            tf.math.log(1.0),
            tf.math.log(1000.0),
            NOISE_EMBEDDING_SIZE // 2,
        )
    )
    angular_speeds = 2.0 * math.pi * frequencies
    embeddings = tf.concat(
        [tf.sin(angular_speeds * x), tf.cos(angular_speeds * x)], axis=3
    )
    return embeddings

2.3 ResidualBlock

DownBlock 和 UpBlock 都使用了 ResidualBlock 层，我们已经学习了残差块的构建方法，在本节进行简单回顾。
残差块 (ResidualBlock) 是包含跳跃连接的一组神经网络层，将残差块输入添加到输出中。残差块可以用于构建更深的网络，学习更复杂的模式，而不会受到梯度消失和退化问题的严重影响。梯度消失问题是指随着网络深度增加，传播到较深层的梯度变得非常小，因此学习非常缓慢。退化问题是指随着神经网络变得更深，它们不一定比浅层网络更准确，准确率可能会在某个深度达到饱和，然后迅速退化。
He 等人在 ResNet 论文中引入了残差块，通过在网络层周围添加一个跳跃连接，模块有选择地绕过复杂的权重更新，并简单地通过恒等映射，使得网络可以在不牺牲梯度大小或网络准确性的情况下进行深度训练。
ResidualBlock 如下图所示，在某些残差块中，还需要在跳跃连接上使用核大小为 1 的 Conv2D 层，以使通道数与块的其余部分保持一致。

使用 Keras 实现 ResidualBlock块。

def ResidualBlock(width):
    def apply(x):
        input_width = x.shape[3]
        # 检查输入的通道数是否与该块预期输出的通道数匹配，如果不匹配，则在跳跃连接上添加一个额外的 Conv2D 层，以使通道数与 ResidualBlock 块的其余部分保持一致
        if input_width == width:
            residual = x
        else:
            residual = layers.Conv2D(width, kernel_size=1)(x)
        # 应用 BatchNormalization 层
        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)
        # 将 ResidualBlock 输入添加到输出中，以获得 ResidualBlock 块的最终输出
        x = layers.Add()([x, residual])
        return x

    return apply

2.4 DownBlocks 和 UpBlocks

每个连续的 DownBlock 通过 block_depth (本节所用模型中为 2 )个 ResidualBlocks 增加通道数，同时还在最后应用了一个 AveragePooling2D 层，以将图像的尺寸减半。每个 ResidualBlock 都被添加到一个列表中，用于连接到 U-Net 的 UpBlock 层作为跳跃连接。
UpBlock 首先应用一个 UpSampling2D 层，通过双线性插值将图像的尺寸扩大一倍。每个连续的 UpBlock 通过 block_depth 个 ResidualBlocks 减少通道数，同时还通过 U-Net 中的跳跃连接与 DownBlocks 的输出进行串联。该过程如下图所示。

使用 Keras 实现 DownBlock 和 UpBlock：

def DownBlock(width, block_depth):
    def apply(x):
        x, skips = x
        for _ in range(block_depth):
            # DownBlock 使用 ResidualBlock 增加图像的通道数
            x = ResidualBlock(width)(x)
            # 将每个 ResidualBlock 保存到一个列表 (skips) 中，供 UpBlock 使用
            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
        # UpBlock 以一个 UpSampling2D 层开始，将图像的尺寸扩大一倍
        x = layers.UpSampling2D(size=2, interpolation="bilinear")(x)
        for _ in range(block_depth):
            # 将当前输出与对应的 DownBlock 层的输出通过 Concatenate 层连接在一起
            x = layers.Concatenate()([x, skips.pop()])
            # 通过 UpBlock 时，使用 ResidualBlock 减少图像的通道数
            x = ResidualBlock(width)(x)
        return x

    return apply

3. 训练扩散模型

创建、编译并拟合去噪扩散概率模型。

# 实例化模型
ddm = DiffusionModel()
# 使用训练集计算归一化统计信息
ddm.normalizer.adapt(train)
# 使用 AdamW 优化器(类似于 Adam，但具有权重衰减，有助于稳定训练过程)和均方绝对误差损失函数编译模型
ddm.compile(
    optimizer=optimizers.experimental.AdamW(
        learning_rate=LEARNING_RATE, weight_decay=WEIGHT_DECAY
    ),
    loss=losses.mean_absolute_error,
)
# 在 50 个 epochs 上拟合模型
ddm.fit(train, epochs=EPOCHS)

4. 去噪扩散概率模型的采样

为了从训练好的模型中采样图像，我们需要应用逆扩散 (reverse diffusion) 过程，也就是说，我们需要从随机噪声开始，并使用模型逐步消除噪声，直到得到一张清晰的花朵图片。
扩散模型被训练用于预测给定噪声图像中添加的总噪声量，而不仅仅是在添加噪声过程的最后一个时间步中添加的噪声。但是，我们不希望一次性完全消除噪声，通过一次性从完全随机的噪声中预测图像显然并不可行。我们希望能够模仿正向过程，逐步地在多个步骤中消除预测的噪声，以允许模型根据预测结果进行调整。
为了生成逼真图像，我们可以使用两个步骤从 x t x_t xt​ 跳转到 x t − 1 x_{t-1} xt−1​。首先，使用模型的噪声预测来计算原始图像 x 0 x_0 x0​ 的估计值，然后将预测的噪声重新应用于该图像，得到 x t − 1 x_{t-1} xt−1​ 作为下一步迭代的图像，如下图所示。

多次重复此过程，最终得到对 x 0 x_0 x0​ 的估计，得到一个高质量样本。实践中，可以自由选择迭代的步数，且不必与训练噪声过程中的时间步数(本节构建的模型中为 1000 )相同，反向迭代的步数可以较小，本节中，我们使用 20 个迭代步数。可以使用以下数学方程描述此过程：
x t − 1 = α ‾ t − 1 ( x t − 1 − α ‾ t ϵ θ ( t ) ( x t ) α ‾ t ) ⏟ p r e d i c t e d   x 0 + 1 − α ‾ t − 1 − σ t 2 ϵ θ ( t ) ( x t ) ⏟ d i r e c t i o n   p o i n t i n   t o   x t + σ t ϵ t ⏟ r a n d o m   n o i s e x_{t-1} = \overline α_{t-1} \underbrace {(\frac {x_t-\sqrt{1 - \overlineα_t} ϵ_θ^{(t)}(x_t)} {\sqrt {\overline \alpha_t}})}_{predicted\ x_0} +\underbrace{ \sqrt{1 - \overline α_{t-1} - σ_t^2} ϵ_θ(t)(x_t)}_{direction\ pointin\ to\ x_t} + \underbrace {σ_t ϵ_t}_{random\ noise} xt−1​=αt−1​predicted x0​ (αt​ ​xt​−1−αt​ ​ϵθ(t)​(xt​)​)​​+direction pointin to xt​ 1−αt−1​−σt2​ ​ϵθ​(t)(xt​)​​+random noise σt​ϵt​​​
方程右侧括号内的第一项是使用网络预测的噪声计算得到的估计图像 x 0 x_0 x0​。然后，我们将其乘以 t − 1 t-1 t−1 时的信号率 α ‾ t − 1 \sqrt {\overline α_{t-1}} αt−1​ ​，并再一次使用预测噪声(通过乘以 t − 1 t-1 t−1 时的噪声率 1 − α ‾ t − 1 − σ t 2 \sqrt{1 - \overline α_{t-1} - σ_t^2} 1−αt−1​−σt2​ ​ 进行缩放)；并添加额外的高斯随机噪声 σ t ϵ t {σ_t ϵ_t} σt​ϵt​，其中 σ t σ_t σt​ 用于决定生成过程的随机性有多高。
当对于所有 t t t 有 σ t = 0 σ_t = 0 σt​=0 时，模型称为去噪扩散隐式模型 (Denoising Diffusion Implicit Model, DDIM)，使用 DDIM，生成过程完全是确定性的，也就是说，相同的随机噪声输入将始终产生相同的输出。这样我们就可以得到从潜空间到像素空间的明确定义的映射关系。
接下来，我们将构建 DDIM 以使生成过程具有确定性。实现 DDIM 采样过程(逆扩散)：

    def reverse_diffusion(self, initial_noise, diffusion_steps):
        num_images = initial_noise.shape[0]
        step_size = 1.0 / diffusion_steps
        current_images = initial_noise
        # 在固定的步数(例如 20 步)内生成观测样本
        for step in range(diffusion_steps):
            # 所有扩散时间都设为 1，即逆扩散过程开始时
            diffusion_times = tf.ones((num_images, 1, 1, 1)) - step * step_size
            # 噪声率和信号率根据扩散规划进行计算
            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
            # 计算新的噪声率和信号率
            next_noise_rates, next_signal_rates = self.diffusion_schedule(next_diffusion_times)
            # 根据 t-1 扩散规划重新应用预测的噪声到预测的图像，计算 t-1 时的图像
            current_images = (next_signal_rates * pred_images + next_noise_rates * pred_noises)
        # 完成 20 步后，返回最终的 x 预测图像 x0
        return pred_images

5. 扩散模型分析

接下来，我们学习如何使用训练好的模型进行三种不同的操作：生成新图像，测试逆扩散步数对生成图像质量的影响，以及在潜空间中两个图像之间进行插值。

5.1 生成图像

为了使用训练后的模型生成图像样本，我们只需运行逆扩散过程，并确保在最后对输出进行反归一化处理(即将像素值恢复到 [0, 1] 范围内)：

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

    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
        )
    def generate(self, num_images, diffusion_steps, initial_noise=None):
        # 生成初始噪声图像
        if initial_noise is None:
            initial_noise = tf.random.normal(shape=(num_images, IMAGE_SIZE, IMAGE_SIZE, 3))
        # 应用逆扩散过程
        generated_images = self.reverse_diffusion(initial_noise, diffusion_steps)
        # 网络输出的图像均值为 0，方差为 1，因此需要重新应用从训练数据计算得出的均值和方差来执行反归一化处理
        generated_images = self.denormalize(generated_images)
        return generated_images

在下图中，可以看到训练过程中扩散模型的生成的图像样本。

5.2 调整逆扩散步数

我们还可以测试调整逆扩散步数对图像质量的影响。直觉上，扩散过程中采取的步数越多，图像生成的质量越高。

在上图中可以看到，随着逆扩散步数的增加，生成的质量确实得到了改善。最初的噪声样本中，模型只能预测出一个模糊的颜色块，随着步数的增加，模型能够改善并锐化其生成图像。然而，生成图像所需的时间与逆扩散步数成正比，因此需要在生成质量和生成速度之间进行权衡。

5.3 在图像之间进行插值

在变分自编码器中，我们可以在高斯潜空间的两点之间进行插值，以在像素空间中平滑地过渡图像。在扩散模型，我们使用球面插值方法，确保方差保持恒定，同时将两个高斯噪声图像混合在一起。具体而言，每个步骤的初始噪声图像由 a sin ⁡ ( π 2 t ) + b cos ⁡ ( π 2 t ) a \sin(\frac π2t)+b \cos(\frac π2t) asin(2π​t)+bcos(2π​t) 确定，其中 t t t 在 0 到 1 之间平滑变化， a a a 和 b b b 是我们希望进行插值的两个随机采样的高斯噪声张量，插值结果如下图所示。

小结

本节中，我们介绍了最近最先进的生成模型之一，扩散模型。介绍了去噪扩散概率模型 (Denoising Diffusion Probabilistic Model, DDPM)，并利用去噪扩散隐式模型 (Denoising Diffusion Implicit Model, DDIM) 的思想，使生成过程具备完全的确定性。扩散模型由前向扩散过程和逆扩散过程组成，前向扩散过程通过一系列小步骤向训练数据添加噪声，而逆扩散过程中模型的目标是预测添加的噪声。

系列链接

AIGC实战——生成模型简介
AIGC实战——深度学习 (Deep Learning, DL)
AIGC实战——卷积神经网络(Convolutional Neural Network, CNN)
AIGC实战——自编码器(Autoencoder)
AIGC实战——变分自编码器(Variational Autoencoder, VAE)
AIGC实战——使用变分自编码器生成面部图像
AIGC实战——生成对抗网络(Generative Adversarial Network, GAN)
AIGC实战——WGAN(Wasserstein GAN)
AIGC实战——条件生成对抗网络(Conditional Generative Adversarial Net, CGAN)
AIGC实战——自回归模型(Autoregressive Model)
AIGC实战——改进循环神经网络
AIGC实战——像素卷积神经网络(PixelCNN)
AIGC实战——归一化流模型(Normalizing Flow Model)
AIGC实战——能量模型(Energy-Based Model)