AIGC实战——能量模型(Energy-Based Model)

AIGC实战——能量模型

• • 0. 前言
  • 1. 能量模型
  • • 1.1 模型原理
    • 1.2 MNIST 数据集
    • 1.3 能量函数
  • 2. 使用 Langevin 动力学进行采样
  • • 2.1 随机梯度 Langevin 动力学
    • 2.2 实现 Langevin 采样函数
  • 3. 利用对比散度训练
  • 小结
  • 系列链接

0. 前言

能量模型 (Energy-based Model, EBM) 是一类常见的生成模型，其借鉴了物理系统建模的一个关键思想，即事件的概率可以用玻尔兹曼分布来表示。玻尔兹曼分布是一种将实值能量函数归一化到 0 和 1 之间的函数，该分布最早由 Ludwig Boltzmann 于 1868 年提出，用于描述处于热平衡状态的气体系统。在本节中，我们将利用这一思想来训练一个生成模型，用于生成 MNIST 手写数字的图像。

1. 能量模型

1.1 模型原理

能量模型是使用玻尔兹曼分布对真实的数据生成分布进行建模的方法，其中 $E(x)$ 被称为观测值 $x$ 的能量函数( energy function，或得分)。玻尔兹曼分布如下：
$$p(x)=\frac{e^{-E(x)}}{{\int }_{\hat{x}\in X}{e}^{-E(\hat{x})}}$$
在实际应用中，需要训练一个神经网络 $E(x)$，使其对可能出现的观测值输出较低的分数(使 $p(x)$ 接近 1)，对不太可能出现的观测值输出较高的分数(使 $p(x)$ 接近 0)。
使用这种方式对数据进行建模存在两个挑战。首先，我们不清楚如何使用模型来生成新的观测值——我们可以使用它为给定观测值生成一个分数，但尚不清楚如何生成一个具有较低分数(即合理的观测值)的观测值。
其次，上式中的规范化分母包含一个积分，这通常难以计算。如果我们无法计算这个积分，那么我们就不能使用极大似然估计来训练模型，因为模型训练要求 $p(x)$ 是一个有效的概率分布。
能量模型的关键思想在于，使用近似技术确保永远不需要计算这个无法处理的分母。这与归一化流方法不同，在归一化流方法中，我们需要确保对标准高斯分布进行的转换不会改变输出仍然是有效的概率分布。
我们通过使用对比散度(用于能量模型 (Energy-based Model, EBM) 是一类常见的生成模型，其借鉴了物理系统建模的一个关键思想，即事件的概率可以用玻尔兹曼分布来表示。玻尔兹曼分布是一种将实值能量函数归一化到 0 和 1 之间的函数，该分布最早由 Ludwig Boltzmann 于 1868 年提出，用于描述处于热平衡状态的气体系统。在本节中，我们将利用这一思想来训练一个生成模型，用于生成 MNIST 手写数字的图像。训练)和 Langevin 动力学(用于采样)技术来绕过无法处理的分母问题，我们会在本节后面详细探讨这些技术，并构建 EBM。
首先，准备一个数据集并设计一个简单的神经网络来表示实值能量函数 $E(x)$。

1.2 MNIST 数据集

本节，我们将使用 MNIST 数据集，其中包含手写数字的灰度图像，数据集中的示例图像如下所示。

该数据集已预置于 TensorFlow 中，可以通过以下代码进行下载：

# Load the data
(x_train, _), (x_test, _) = datasets.mnist.load_data()

将像素值缩放到 [-1, 1] 的范围内，将图像的尺寸填充为 32×32 像素，并将其转换为 TensorFlow Dataset：

# Preprocess the data
def preprocess(imgs):
    """
    Normalize and reshape the images
    """
    imgs = (imgs.astype("float32") - 127.5) / 127.5
    imgs = np.pad(imgs, ((0, 0), (2, 2), (2, 2)), constant_values=-1.0)
    imgs = np.expand_dims(imgs, -1)
    return imgs

x_train = preprocess(x_train)
x_test = preprocess(x_test)

x_train = tf.data.Dataset.from_tensor_slices(x_train).batch(BATCH_SIZE)
x_test = tf.data.Dataset.from_tensor_slices(x_test).batch(BATCH_SIZE)

1.3 能量函数

构建了数据集后，就可以构建表示能量函数 $E(x)$ 的神经网络。能量函数 $E_{\theta }(x)$ 是一个具有参数 $\theta$ 的神经网络，它可以将输入图像 $x$ 转换为一个标量值。在整个网络中，我们使用 swish 激活函数。

Swish 激活函数
Swish 是 Google 在 2017 年引入的替代 ReLU 的激活函数，其定义如下：
$$Swish(x)=x\cdot sigmoid(x)=\frac{x}{e^{-x}+1}$$
Swish 和 ReLU 函数曲线非常相似，但关键区别在于 Swish 是平滑的，这有助于缓解梯度消失问题，这对于基于能量的模型特别重要，Swish 函数曲线如下图所示。

表示能量函数的网络由一组堆叠的 Conv2D 层组成，逐渐减小图像的尺寸并增加通道数。最后一层是一个具有线性激活函数的单个全连接单元，因此网络可以输出范围在 $(-\infty ,\infty )$ 的值。

ebm_input = layers.Input(shape=(IMAGE_SIZE, IMAGE_SIZE, CHANNELS))
# 能量函数是一组堆叠的 Conv2D 层，使用 swish 激活函数
x = layers.Conv2D(
    16, kernel_size=5, strides=2, padding="same", activation=activations.swish
)(ebm_input)
x = layers.Conv2D(
    32, kernel_size=3, strides=2, padding="same", activation=activations.swish
)(x)
x = layers.Conv2D(
    64, kernel_size=3, strides=2, padding="same", activation=activations.swish
)(x)
x = layers.Conv2D(
    64, kernel_size=3, strides=2, padding="same", activation=activations.swish
)(x)
x = layers.Flatten()(x)
x = layers.Dense(64, activation=activations.swish)(x)
# 最后一层是一个具有线性激活函数的单个全连接单元
ebm_output = layers.Dense(1)(x)
model = models.Model(ebm_input, ebm_output)
# 构建 Keras 模型将输入图像转换为标量能量值
print(model.summary())

2. 使用 Langevin 动力学进行采样

能量函数仅针对给定的输入输出一个得分，问题在于我们如何利用这个函数生成具有较低能量得分的新样本。
我们将使用 Langevin 动力学 (Langevin dynamics) 技术，利用了计算能量函数相对于其输入的梯度。如果我们从样本空间中的一个随机点出发，并朝着计算得到的梯度相反的方向移动一小段步长，将逐渐降低能量函数。如果神经网络能够正确训练，那么随机噪声应该会转化成一个类似于训练集观测值的图像。

2.1 随机梯度 Langevin 动力学

在穿越样本空间时，我们还必须向输入添加一小部分随机噪声；否则，有可能陷入局部最小值。因此，这种技术称为随机梯度 Langevin 动力学 (Stochastic Gradient Langevin Dynamics)，如下图所示，使用三维空间表示二维空间，其中第三个维度是能量函数的值。路径是一个带有噪声的下坡，沿着输入 $x$ 相对于能量函数 $E(x)$ 的负梯度方向进行。在 MNIST 图像数据集中，有 1024 个像素，因此需要我们在一个 1024 维空间中计算，但是所用的原理是相通的。

需要注意的是，这种梯度下降的方式与我们通常用于训练神经网络的梯度下降方式之间存在差异。
在训练神经网络时，我们使用反向传播计算损失函数相对于网络参数(即权重)的梯度。然后，缓慢调整参数朝着负梯度的方向更新，以便在多次迭代中逐渐最小化损失。
使用 Langevin 动力学时，我们保持神经网络的权重不变，并计算输出相对于输入的梯度。然后，我们稍微调整输入朝着负梯度的方向更新，以便在多次迭代中逐渐最小化输出(能量得分)。
这两个过程都利用了相同的思想(梯度下降)，但应用于不同的函数和实体。
形式上，Langevin 动力学可以用以下方程描述：
$$x^k=x^{k-1}-\eta {\nabla }_x{E}_{\theta }(x^{k-1})+\omega$$
其中 $\omega \sim \mathcal{N}(0,\sigma )，x_0\sim \mu (–1,1)$。$\eta$ 是需要调整的步长超参数，如果太大，步长会越过最小值，如果太小，算法将收敛过慢。$x_0\sim \mu (–1,1)$ 表示在 [-1, 1] 范围内的均匀分布。

2.2 实现 Langevin 采样函数

接下来，使用 Keras 实现 Langevin 采样函数。

# Function to generate samples using Langevin Dynamics
def generate_samples(
    model, inp_imgs, steps, step_size, noise, return_img_per_step=False
):
    imgs_per_step = []
    # 循环给定的步数
    for _ in range(steps):
        # 为图像添加少量噪音
        inp_imgs += tf.random.normal(inp_imgs.shape, mean=0, stddev=noise)
        inp_imgs = tf.clip_by_value(inp_imgs, -1.0, 1.0)
        with tf.GradientTape() as tape:
            tape.watch(inp_imgs)
            # 将图像通过模型，得到能量得分
            out_score = model(inp_imgs)
        # 计算输出对输入的梯度
        grads = tape.gradient(out_score, inp_imgs)
        grads = tf.clip_by_value(grads, -GRADIENT_CLIP, GRADIENT_CLIP)
        # 将一小部分梯度添加到输入图像中
        inp_imgs += step_size * grads
        inp_imgs = tf.clip_by_value(inp_imgs, -1.0, 1.0)
        if return_img_per_step:
            imgs_per_step.append(inp_imgs)
    if return_img_per_step:
        return tf.stack(imgs_per_step, axis=0)
    else:
        return inp_imgs

3. 利用对比散度训练

我们已经知道如何从样本空间中采样出一个新的低评分能量点，接下来介绍如何训练模型。
能量模型无法应用最大似然估计，因为能量函数不输出概率，而是输出一个在样本空间上的分数。我们将应用 Geoffrey Hinton 在 2002 年提出的对比散度 (contrastive divergence) 技术，用于训练非规范化评分模型。我们需要最小化数据的负对数似然：
$$\mathcal{L}=-{\mathbb{E}}_{x\sim data}[log{p}_{\theta }(x)]$$
当 $p_{\theta }(x)$ 采用玻尔兹曼分布形式，具有能量函数 $E\theta (x)$ 时，可以通过以下方式表示该值的梯度：
$${\nabla }_{\theta }\mathcal{L}={\mathbb{E}}_{x\sim data}[{\nabla }_{\theta }{E}_{\theta }(x)]-{\mathbb{E}}_{x\sim model}[{\nabla }_{\theta }{E}_{\theta }(x)]$$
直观上讲，我们希望训练模型对真实观测值输出较大的负能量得分，对生成的伪造观测值输出较大的正能量得分，以便使这两类观测值之间的差距尽可能大。
换句话说，我们可以计算真实样本和伪造样本的能量得分之间的差，并将其作为损失函数。
要计算伪造样本的能量得分，我们需要能够从分布 $p_{\theta }(x)$ 中精确采样，但由于难以求解玻尔兹曼分布的规范化分母，因此难以实现；因此，我们可以使用 Langevin 采样过程生成一组能量得分较低的观测值。这个过程需要无限多的步骤才能生成一个完美的样本(显然并不现