超级会员免费看
基于能量的模型：原理、训练与应用
1. 引言
在深度学习领域，我们已经探讨了多种用于对可观测变量（如图像）的边缘分布 $p(x)$ 进行建模的深度生成模型，包括自回归模型（ARMs）、基于流的模型（简称为流）、变分自编码器（VAEs）以及诸如分层 VAEs 和基于扩散的深度生成模型（DDGMs）等分层模型。
然而，我们一直倡导在寻找可观测变量和决策变量的联合分布 $p(x, y) = p(y|x)p(x)$ 的背景下使用深度生成建模。对联合分布取对数后，我们得到两个相加的部分：$\ln p(x, y) = \ln p(y|x) + \ln p(x)$。在混合建模设置中，我们概述了如何构建和训练这样的联合模型，但混合建模的缺点是需要对两个分布进行加权，即 $\ell(x, y_{\lambda}) = \ln p(y|x) + \lambda \ln p(x)$，并且当 $\lambda \neq 1$ 时，这个目标函数并不对应于联合分布的对数似然。那么，是否有可能构建一个模型使得 $\lambda = 1$ 呢？这里我们将讨论使用概率能量基模型（EBMs）来解决这个问题。
EBMs 的历史可以追溯到上世纪 80 年代，当时提出了被称为玻尔兹曼机的模型。有趣的是，玻尔兹曼机的思想来源于统计物理学，由认知科学家提出。简单来说，我们可以定义一个能量函数 $E(x)$，它为给定状态分配一个值（能量），而不是像高斯分布或伯努利分布那样指定一个特定的分布。能量函数没有限制，因此可以用神经网络进行参数化。然后，通过将能量转换为未归一化的概率 $e^{-E(x)}$，并通过 $Z = \sum_{x} e^{-E(x)}$（即配分函数）进行归一化，我们可以得到玻尔兹曼（也称为吉布斯）分布：
[p(x) = \frac{e^{-E(x)}}{Z}]
如果考虑连续随机变量，求和符号应替换为积分。例如，高斯分布也可以表示为具有解析可处理的配分函数的玻尔兹曼分布，其能量函数形式为：
[E(x; \mu, \sigma^2) = \frac{1}{2\sigma^2} (x - \mu)^2]
从而得到：
[p(x) = \frac{e^{-E(x)}}{\int e^{-E(x)}dx} = \frac{e^{-\frac{1}{2\sigma^2} (x - \mu)^2}}{\sqrt{2\pi\sigma^2}}]
在实践中，大多数能量函数不会产生易于计算的配分函数，并且配分函数通常是学习基于能量的模型时的关键问题。此外，一般来说，从这些模型中采样也很困难，因为没有像 ARMs、流或 VAEs 那样的生成过程。
尽管存在这些问题，使用 EBMs 仍有至少三个好处：
- 原则上，能量函数是无约束的，可以是任何函数，比如神经网络。
- 能量函数可以是多模态的，而无需像混合分布那样进行定义。
- 无论定义在离散还是连续变量上，能量函数都没有区别。
2. 模型构建
我们使用一些参数 $\theta$ 构建一个关于可观测随机变量和决策随机变量的能量函数 $E(x, y; \theta)$，它为一对 $(x, y)$ 分配一个值（能量），其中 $x \in \mathbb{R}^D$ 且 $y \in {0, 1, \ldots, K - 1}$。假设 $E(x, y; \theta)$ 由一个神经网络 $NN_{\theta}(x)$ 参数化，该网络返回 $K$ 个值：$NN_{\theta}: \mathbb{R}^D \to \mathbb{R}^K$。那么，我们可以将能量定义为：
[E(x, y; \theta) = -NN_{\theta}(x)[y]]
其中 $[y]$ 表示神经网络 $NN_{\theta}(x)$ 的特定输出。
联合概率分布定义为玻尔兹曼分布：
[p_{\theta}(x, y) = \frac{\exp{NN_{\theta}(x)[y]}}{\sum_{x,y} \exp{NN_{\theta}(x)[y]}} = \frac{\exp{NN_{\theta}(x)[y]}}{Z_{\theta}}]
其中配分函数定义为 $Z_{\theta} = \sum_{x,y} \exp{NN_{\theta}(x)[y]}$。
我们可以通过联合分布计算边缘分布和条件分布：
-
边缘分布 $p_{\theta}(x)$
：
[p_{\theta}(x) = \sum_{y} p_{\theta}(x, y) = \frac{\sum_{y} \exp{NN_{\theta}(x)[y]}}{Z_{\theta}}]
我们还可以将分子改写为：
[\sum_{y} \exp{NN_{\theta}(x)[y]} = \exp\left{\log\left(\sum_{y} \exp{NN_{\theta}(x)[y]}\right)\right} = \exp\left{\text{LogSumExp}
y{NN
{\theta}(x)[y]}\right}]
其中 $\text{LogSumExp}
y{f(y)} = \ln \sum
{y} \exp{f(y)}$。因此，边缘分布可以定义为：
[p_{\theta}(x) = \frac{\exp\left{\text{LogSumExp}
y{NN
{\theta}(x)[y]}\right}}{Z_{\theta}}]
-
条件分布 $p_{\theta}(y|x)$
：
[p_{\theta}(y|x) = \frac{p_{\theta}(x, y)}{p_{\theta}(x)} = \frac{\exp{NN_{\theta}(x)[y]}}{\sum_{y} \exp{NN_{\theta}(x)[y]}}]
这实际上就是 softmax 函数。这表明基于能量的模型既可以用作分类器，也可以用作边缘分布，并且只需要定义一个神经网络即可。
联合分布的对数为：
[\ln p_{\theta}(x, y) = \ln p_{\theta}(y|x) + \ln p_{\theta}(x) = \ln \text{softmax}{NN_{\theta}(x)[y]} + \left(\text{LogSumExp}
y{NN
{\theta}(x)[y]} - \ln Z_{\theta}\right)]
我们可以看到，该模型需要一个共享的神经网络来计算两个分布，通过选择最终的激活函数可以得到特定的分布。模型的示意图如下：
graph LR
classDef startend fill:#F5EBFF,stroke:#BE8FED,stroke-width:2px
classDef process fill:#E5F6FF,stroke:#73A6FF,stroke-width:2px
classDef decision fill:#FFF6CC,stroke:#FFBC52,stroke-width:2px
A([输入 x]):::startend --> B(神经网络 NNθ):::process
B --> C(LogSumExpy):::process
C --> D(˜pθ(x)):::process
B --> E(softmax):::process
E --> F(pθ(y|x)):::process
D --> G(归一化):::process
F --> G
G --> H([输出分布]):::startend
3. 训练
我们有一个单一的神经网络需要训练，训练目标是联合分布的对数。由于训练目标是条件分布 $p_{\theta}(y|x)$ 的对数和边缘分布 $p_{\theta}(x)$ 的对数之和，因此计算关于参数 $\theta$ 的梯度需要分别对每个组件求梯度。
我们知道学习分类器没有问题，所以让我们仔细看看第二个组件：
[\nabla_{\theta} \ln p_{\theta}(x) = \nabla_{\theta} \text{LogSumExp}
y{NN
{\theta}(x)[y]} - \nabla_{\theta} \ln Z_{\theta}]
[= \nabla_{\theta} \text{LogSumExp}
y{NN
{\theta}(x)[y]} - \nabla_{\theta} \ln \sum_{x} \exp\left{\text{LogSumExp}
y{NN
{\theta}(x)[y]}\right}]
[= \nabla_{\theta} \text{LogSumExp}
y{NN
{\theta}(x)[y]} - \sum_{x’} \frac{\exp\left{\text{LogSumExp}
y{NN
{\theta}(x’)[y]}\right}}{\sum_{x’‘,y’‘} \exp{NN_{\theta}(x’‘)[y’‘]}} \nabla_{\theta} \text{LogSumExp}
y{NN
{\theta}(x’)[y]}]
[= \nabla_{\theta} \text{LogSumExp}
y{NN
{\theta}(x)[y]} - \mathbb{E}
{x’ \sim p
{\theta}(x)} \left[\nabla_{\theta} \text{LogSumExp}
y{NN
{\theta}(x’)[y]}\right]]
这里，第一部分 $\nabla_{\theta} \text{LogSumExp}
y{NN
{\theta}(x)[y]}$ 的梯度可以针对给定的数据点 $x$ 进行计算，因为 log-sum-exp 函数是可微的，所以我们可以使用自动求导工具。然而，第二部分 $\mathbb{E}
{x’ \sim p
{\theta}(x)} \left[\nabla_{\theta} \text{LogSumExp}
y{NN
{\theta}(x’)[y]}\right]$ 存在两个问题：
- 配分函数对数的梯度变成了根据模型分布的 $x$ 的期望值，而这个期望值无法进行解析计算，并且从边缘分布 $p_{\theta}(x)$ 中采样也并非易事。
- 我们需要计算 $NN_{\theta}(x)$ 的 log-sum-exp 的期望值，不过好消息是我们可以使用自动求导工具来完成。
通常，这个期望值通过蒙特卡罗样本进行近似，但不清楚如何有效地从 EBM 中采样。Grathwohl 等人建议使用朗之万动力学（Langevin dynamics），这是一种马尔可夫链蒙特卡罗（MCMC）方法。在我们的情况下，朗之万动力学从随机初始化的 $x_0$ 开始，然后利用能量函数的景观信息（即梯度）来寻找新的 $x$，即：
[x_{t + 1} = x_t + \alpha \nabla_{x_t} \text{LogSumExp}
y {NN
{\theta}(x)[y]} + \sigma \cdot \epsilon]
其中 $\alpha > 0$，$\sigma > 0$，且 $\epsilon \sim \mathcal{N}(0, I)$。朗之万动力学可以看作是在可观测空间中进行的随机梯度下降，并且在每一步添加了一个小的高斯噪声。
应用这个过程 $\eta$ 步后，我们可以近似梯度如下：
[\nabla_{\theta} \ln p_{\theta}(x) \approx \nabla_{\theta} \text{LogSumExp}
y{NN
{\theta}(x)[y]} - \nabla_{\theta} \text{LogSumExp}
y{NN
{\theta}(x_{\eta})[y]}]
其中 $x_{\eta}$ 表示朗之万动力学过程的最后一步。
最终的训练过程如下：
1. 从数据集中采样 $x_n$ 和 $y_n$。
2. 计算 $NN_{\theta}(x_n)[y]$。
3. 使用例如均匀分布初始化 $x_0$。
4. 运行朗之万动力学 $\eta$ 步：
[x_{t + 1} = x_t + \alpha \nabla_{x_t} \text{LogSumExp}
y {NN
{\theta}(x)[y]} + \sigma \cdot \epsilon]
5. 计算目标函数：
[L_{clf}(\theta) = \sum_{y} \mathbb{1}[y = y_n] \theta_y \ln {NN_{\theta}(x_n)[y]}]
[L_{gen}(\theta) = \text{LogSumExp}
y{NN
{\theta}(x)[y]} - \text{LogSumExp}
y{NN
{\theta}(x_{\eta})[y]}]
[L(\theta) = L_{clf}(\theta) + L_{gen}(\theta)]
6. 应用自动求导工具计算 $\nabla_{\theta} L(\theta)$ 并更新神经网络。
需要注意的是，$L_{clf}(\theta)$ 实际上就是交叉熵损失，而 $L_{gen}(\theta)$ 是对 $x$ 的对数边缘分布的（粗略）近似。
4. 代码实现
以下是实现 EBM 的代码：
import torch
import torch.nn as nn
class EBM(nn.Module):
def __init__(self, energy_net, alpha, sigma, ld_steps, D):
super(EBM, self).__init__()
print('EBM by JT.')
# 用于 EBM 的神经网络
self.energy_net = energy_net
# 分类损失
self.nll = nn.NLLLoss(reduction='none')
# 超参数
self.D = D
self.sigma = sigma
self.alpha = torch.FloatTensor([alpha])
self.ld_steps = ld_steps
def classify(self, x):
f_xy = self.energy_net(x)
y_pred = torch.softmax(f_xy, 1)
return torch.argmax(y_pred, dim=1)
def class_loss(self, f_xy, y):
# 计算 logits（用于分类）
y_pred = torch.softmax(f_xy, 1)
return self.nll(torch.log(y_pred), y)
def gen_loss(self, x, f_xy):
# 使用朗之万动力学采样
x_sample = self.sample(x=None, batch_size=x.shape[0])
# 计算 f(x_sample)[y]
f_x_sample_y = self.energy_net(x_sample)
return -(torch.logsumexp(f_xy, 1) - torch.logsumexp(f_x_sample_y, 1))
def forward(self, x, y, reduction='avg'):
# 前向传播通过网络
# 计算 f(x)[y]
f_xy = self.energy_net(x)
# 判别部分
# 计算判别损失：交叉熵
L_clf = self.class_loss(f_xy, y)
# 生成部分
# 计算生成损失：E(x) - E(x_sample)
L_gen = self.gen_loss(x, f_xy)
# 最终目标
if reduction == 'sum':
loss = (L_clf + L_gen).sum()
else:
loss = (L_clf + L_gen).mean()
return loss
def energy_gradient(self, x):
self.energy_net.eval()
# 复制不需要梯度的原始数据
x_i = torch.FloatTensor(x.data)
x_i.requires_grad = True
# 计算梯度
x_i_grad = torch.autograd.grad(torch.logsumexp(self.energy_net(x_i), 1).sum(), [x_i], retain_graph=True)[0]
self.energy_net.train()
return x_i_grad
def langevine_dynamics_step(self, x_old, alpha):
# 计算关于 x_old 的梯度
grad_energy = self.energy_gradient(x_old)
# 采样 eta ~ Normal(0, alpha)
epsilon = torch.randn_like(grad_energy) * self.sigma
# 新样本
x_new = x_old + alpha * grad_energy + epsilon
return x_new
def sample(self, batch_size=64, x=None):
# 1) 从均匀分布采样
x_sample = 2. * torch.rand([batch_size, self.D]) - 1.
# 2) 运行朗之万动力学
for i in range(self.ld_steps):
x_sample = self.langevine_dynamics_step(x_sample, alpha=self.alpha)
return x_sample
5. 受限玻尔兹曼机
通过能量函数定义模型的思想是广泛的玻尔兹曼机（BMs）家族的基础。玻尔兹曼机定义的能量函数为：
[E(x; \theta) = -\left(x^{\top} W x + b^{\top} x\right)]
其中 $\theta = {W, b}$，$W$ 是权重矩阵，$b$ 是偏置向量。BM 的问题在于由于配分函数的存在，它们很难训练。然而，我们可以通过引入潜在变量并限制可观测变量之间的连接来缓解这个问题。
5.1 受限玻尔兹曼机的定义
考虑一个由二进制可观测变量 $x \in {0, 1}^D$ 和二进制潜在（隐藏）变量 $z \in {0, 1}^M$ 组成的 BM。变量之间的关系通过以下能量函数指定：
[E(x, z; \theta) = -x^{\top} W z - b^{\top} x - c^{\top} z]
其中 $\theta = {W, b, c}$ 是一组参数，$W \in \mathbb{R}^{D \times M}$，$b \in \mathbb{R}^D$，$c \in \mathbb{R}^M$ 分别是权重、可观测偏置和隐藏偏置。对于上述能量函数，受限玻尔兹曼机（RBM）由吉布斯分布定义：
[p(x, z|\theta) = \frac{1}{Z_{\theta}} \exp\left{-E(x, z; \theta)\right}]
其中配分函数为：
[Z_{\theta} = \sum_{x} \sum_{z} \exp\left{-E(x, z; \theta)\right}]
可观测变量的边缘概率（观测的似然）为：
[p(x|\theta) = \frac{1}{Z_{\theta}} \exp\left{-F(x; \theta)\right}]
其中 $F(\cdot)$ 是自由能：
[F(x; \theta) = -b^{\top} x - \sum_{j} \log\left(1 + \exp\left{c_j + (W_{\cdot j})^{\top} x\right}\right)]
RBM 具有一个非常有用的性质，即在给定可观测变量的情况下，隐藏变量的条件分布可以分解，反之亦然，这导致：
[p(z_m = 1|x, \theta) = \text{sigm}\left(c_m + (W_{\cdot m})^{\top} x\right)]
[p(x_d = 1|z, \theta) = \text{sigm}(b_d + W_{d \cdot} z)]
5.2 训练 RBMs
对于给定的数据 $D = {x_n}
{n = 1}^N$，我们可以使用最大似然方法训练 RBM，该方法寻求对数似然函数的最大值：
[\ell(\theta) = \frac{1}{N} \sum
{x_n \in D} \log p(x_n|\theta)]
学习目标 $\ell(\theta)$ 关于 $\theta$ 的梯度形式如下：
[\nabla_{\theta} \ell(\theta) = -\frac{1}{N} \sum_{n = 1}^N \left(\nabla_{\theta} F(x_n; \theta) - \sum_{\hat{x}} p(\hat{x}|\theta) \nabla_{\theta} F(\hat{x}; \theta)\right)]
一般来说，上述梯度无法进行解析计算，因为第二项需要对所有可观测变量的配置进行求和。一种解决方法是使用标准的随机近似，用根据 $p(x|\theta)$ 抽取的 $S$ 个样本 ${\hat{x}
1, \ldots, \hat{x}_S}$ 的和来代替 $p(x|\theta)$ 下的期望：
[\nabla
{\theta} \ell(\theta) \approx -\left(\frac{1}{N} \sum_{n = 1}^N \nabla_{\theta} F(x_n; \theta) - \frac{1}{S} \sum_{s = 1}^S \nabla_{\theta} F(\hat{x}
s; \theta)\right)]
另一种方法是对比散度（contrastive divergence），它通过对从应用 $K$ 步块吉布斯采样过程获得的分布中抽取的样本 $\bar{x}_n$ 求和来近似式中的期望：
[\nabla
{\theta} \ell(\theta) \approx -\frac{1}{N} \sum_{n = 1}^N \left(\nabla_{\theta} F(x_n; \theta) - \nabla_{\theta} F(\bar{x}_n; \theta)\right)]
原始的 CD 使用 $K$ 步吉布斯链，从每个数据点 $x_n$ 开始获得一个样本 $\bar{x}_n$，并在每次参数更新后重新启动。另一种方法是持久对比散度（Persistent Contrastive Divergence，PCD），它在每次更新后不重新启动链，这通常会导致收敛速度较慢，但最终性能更好。
5.3 通过能量函数定义高阶关系
能量函数是一个有趣的概念，因为它允许对变量之间的高阶依赖关系进行建模。例如，二进制 RBM 可以通过引入两种隐藏变量（即子空间单元和门单元）扩展到三阶乘法交互。子空间单元是反映特征变化的隐藏变量，因此它们对不变性更具鲁棒性。门单元负责激活子空间单元，可以看作是由子空间特征组成的池化特征。
考虑以下随机变量：$x \in {0, 1}^D$，$h \in {0, 1}^M$，$S \in {0, 1}^{M \times K}$。我们关注一个可观测变量 $x_i$ 和两种类型的隐藏二进制单元（一个门单元 $h_j$ 和一个子空间单元 $s_{jk}$）相连的情况。每个门单元与一组子空间隐藏单元相关联。联合配置的能量函数定义如下：
[E(x, h, S; \theta) = -\sum_{i = 1}^D \sum_{j = 1}^M \sum_{k = 1}^K W_{ijk} x_i h_j s_{jk} - \sum_{i = 1}^D b_i x_i - \sum_{j = 1}^M c_j h_j - \sum_{j = 1}^M h_j \sum_{k = 1}^K D_{jk} s_{jk}]
综上所述，基于能量的模型为深度生成建模提供了一种独特的方法，尽管存在一些挑战，但通过合理的设计和训练方法，它们在分类和生成任务中都具有潜在的应用价值。受限玻尔兹曼机作为基于能量模型的一个重要实例，展示了如何通过能量函数来定义复杂的变量关系和进行有效的训练。在实际应用中，我们可以根据具体问题选择合适的模型和训练方法，以获得更好的性能。
基于能量的模型：原理、训练与应用
6. 模型应用与效果展示
在完成模型的训练后，我们可以通过运行代码（可参考：https://github.com/jmtomczak/intro_dgm）来观察基于能量的模型（EBM）的实际效果。以下是训练后可能出现的一些结果示例：
| 效果类型 | 描述 |
| ---- | ---- |
| 随机真实图像 | 随机从数据集中选取的真实图像，用于对比模型生成的图像。 |
| 无条件生成图像 | 应用 $\eta = 20$ 步的朗之万动力学后，从 EBM 中进行的无条件生成图像。这些图像展示了模型学习到的数据分布特征。 |
| 目标验证曲线（$L_{clf} + L_{gen}$） | 反映了训练过程中分类损失和生成损失之和的变化情况。通过观察该曲线，我们可以判断模型是否收敛以及训练的稳定性。 |
| 生成目标验证曲线（$L_{gen}$） | 专门展示生成损失的变化曲线，有助于我们了解模型在生成数据方面的性能提升过程。 |
graph LR
classDef startend fill:#F5EBFF,stroke:#BE8FED,stroke-width:2px
classDef process fill:#E5F6FF,stroke:#73A6FF,stroke-width:2px
classDef decision fill:#FFF6CC,stroke:#FFBC52,stroke-width:2px
A([训练 EBM 模型]):::startend --> B(生成图像):::process
B --> C(计算损失):::process
C --> D(绘制验证曲线):::process
D --> E([观察效果]):::startend
7. 总结与展望
基于能量的模型（EBM）为深度生成建模提供了一种独特的视角和方法。与之前讨论的自回归模型（ARMs）、基于流的模型、变分自编码器（VAEs）等相比，EBM 具有一些显著的优势，同时也面临着一些挑战。
7.1 优势总结
- 函数灵活性 ：能量函数原则上是无约束的，可以是任何函数，这使得我们可以使用神经网络等强大的工具来进行参数化，从而能够捕捉复杂的数据分布。
- 多模态建模 ：能量函数可以自然地实现多模态，而无需像混合分布那样进行复杂的定义，这对于处理具有多个模式的数据非常有用。
- 变量类型通用性 ：无论是离散变量还是连续变量，EBM 都可以进行有效的建模，具有广泛的适用性。
- 模型多功能性 ：通过一个共享的神经网络，EBM 既可以用作分类器，也可以用作边缘分布，实现了分类和生成任务的统一。
7.2 挑战分析
- 配分函数计算 ：大多数能量函数不会产生易于计算的配分函数，这使得在学习过程中计算归一化概率变得困难，成为训练 EBM 的关键难题。
- 采样困难 ：由于缺乏像 ARMs、流或 VAEs 那样的明确生成过程，从 EBM 中进行有效的采样是一个挑战，通常需要使用如朗之万动力学等复杂的方法。
7.3 未来展望
尽管存在这些挑战，但随着研究的不断深入和技术的不断发展，EBM 仍然具有很大的潜力。未来的研究方向可能包括：
-
配分函数近似方法的改进
：寻找更高效、更准确的配分函数近似方法，以降低计算复杂度，提高训练效率。
-
采样算法的优化
：开发更有效的采样算法，能够快速、准确地从 EBM 中生成样本，从而更好地应用于实际任务。
-
模型架构的创新
：探索新的模型架构和能量函数形式，进一步提高 EBM 的性能和表达能力。
-
多领域应用拓展
：将 EBM 应用到更多的领域，如计算机视觉、自然语言处理、生物信息学等，解决实际问题，验证其有效性和实用性。
总之，基于能量的模型为我们提供了一个富有挑战性和创新性的研究方向。通过不断地探索和改进，我们有望克服现有的困难，充分发挥 EBM 的优势，为深度生成建模带来新的突破和发展。
在实际应用中，我们可以根据具体问题的特点和需求，选择合适的模型和训练方法。对于一些对数据分布建模要求较高、需要处理复杂多模态数据的任务，EBM 可能是一个不错的选择。同时，结合其他模型的优势，如将 EBM 与 VAEs 或流模型相结合，可能会取得更好的效果。希望本文能够为读者提供一个全面的了解和参考，激发更多的研究和实践。
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
