PyTorch生成式人工智能（21）——归一化流模型(Normalizing Flow Model)

0. 前言

我们已经学习了三类生成模型：变分自动编码器 (Variational Autoencoder, VAE)、生成对抗网络 (Generative Adversarial Network, GAN) 和自回归模型 (Autoregressive Model)。每种模型都使用不同的方式对分布 $p(x)$ 进行建模，其中 VAE 和 GAN 引入了易于使用的潜变量，自回归模型则将分布建模为先前元素值的函数。在本节中，我们将介绍一类新的生成模型，归一化流模型 (Normalizing Flow Model)。

1. 归一化流模型

1.1 归一化流模型基本原理

归一化流模型的目标与变分自编码器类似，变分自编码器学习一个编码器映射函数，用于将复杂分布映射到可以从中采样的简单分布，并学习了一个从简单分布到复杂分布的解码器映射函数，以便通过从简单分布中采样一个点z并应用学到的转换来生成新的数据点。从概率角度来看，解码器用于建模 $p(x|z)$，但编码器只是真实分布 $p(z|x)$ 的近似，编码器和解码器是两个完全不同的神经网络。
在归一化流模型中，解码函数被设计为编码函数的精确反函数，且计算速度快，使得归一化流具有可计算性。而神经网络默认情况下并不是可逆函数，因此需要利用深度学习的灵活性创建一个可逆过程，将一个复杂分布转化为一个简单分布。
为此，我们首先需要了解变量转换 (change of variables) 的技术。在本节中，我们使用一个只有两个维度的简单例子，以便了解归一化流的工作原理，将其进行扩展即可适应复杂问题。

1.2 变量变换

假设有一个在二维空间 $x=(x_1,x_2)$ 中定义的矩形 $X$ 上的概率分布 $p_X(x)$，如下图所示。

此函数在分布的定义域(即 $x_1$ 在 [1, 4] 范围内，$x_2$ 在 [0, 2] 范围内)上积分为 1，因此它表示一个明确定义的概率分布，可以写成如下形式：
$${\int }_0^2{\int }_1^4{p}_X(x)d{x}_{1}d{x}_2=1$$
假设我们想要将这个分布平移和缩放，使其在一个单位正方形 $Z$ 上定义。我们可以通过定义一个新变量 $z=(z_1,z_2)$ 和一个函数 $f$，将 $X$ 中的每个点映射到 $Z$ 中的点，如下所示：
$$\begin{gathered} z=f(x) \\ {z}_1=\frac{x_1-1}{3} \\ {z}_2=\frac{x_2}{2} \end{gathered}$$
此函数是可逆的，也就是说，存在一个函数 $g$，能够将每个 $z$ 映射回其对应的点 $x$，即若一函数有反函数，此函数便称为可逆的 (invertible)。这对于变量变换来说是至关重要的，否则我们无法在两个空间之间进行一致的反向和正向映射，我们可以通过重新排列定义f的方程来找到 $g$，如下图所示。

接下来，我们继续了解从 $X$ 到 $Z$ 的变量变换如何影响概率分布 $p_X(x)$。我们可以通过将定义 $g$ 的方程代入 $p_X(x)$ 来进行转换，得到一个以 $z$ 为自变量的函数 $p_Z(z)$：
$$p_Z(z)=\frac{((3z_1+1)-1)(2z_2)}{9}=\frac{2z_1{z}_2}{3}$$
然而，如果在单位正方形上对 $p_Z(z)$进行积分，会出现以下问题：
$${\int }_0^1{\int }_0^1\frac{2z_1{z}_2}{3}d{z}_{1}d{z}_2=\frac{1}{6}$$
即转换后的函数 $p_Z(z)$ 不再是一个有效的概率分布，因为它的积分等于 $\frac{1}{6}$。如果我们想要将复杂的数据概率分布转换为一个简单的分布以便进行采样，就必须确保其积分为 $1$。
缺失的因子 $6$ 源于转换后的概率分布的定义域比原始定义域小六倍——原始矩形 $X$ 的面积为 6，变换后被压缩为面积为 1 的单位正方形 $Z$。因此，我们需要将新的概率分布乘以一个归一化因子，该因子等于面积(或在更高维度中的体积)的相对变化，可以使用变换的雅可比行列式 (Jacobian determinant) 的绝对值计算给定变换的体积变化。

1.3 雅可比行列式

函数 $z=f(x)$ 的雅可比矩阵是其一阶偏导数的矩阵，如下所示：
$$J=\frac{\partial z}{\partial x}=\left[\begin{array}{ccc}\frac{\partial z_1}{\partial x_1}& \cdots & \frac{\partial z_1}{\partial x_n}\\ ⋮& \ddots & ⋮\\ \frac{\partial z_m}{\partial x_1}& \cdots & \frac{\partial z_m}{\partial x_n}\end{array}\right]$$
在以上示例中，如果计算 $z_1$ 关于 $x_1$ 的偏导，结果为 $\frac{1}{3}$，如果计算 $z_1$ 关于 $x_2$ 的偏导，结果为 0；同样地，如果计算 $z_2$ 关于 $x_1$ 的偏导，结果为 0；最后，如果计算 $z_2$ 关于 $x_2$ 的偏导，结果为 $\frac{1}{2}$。因此，函数 $f(x)$ 的雅可比矩阵如下：
$$J=\left(\begin{array}{rr}\frac{1}{3}& 0\\ 0& \frac{1}{2}\end{array}\right)$$
行列式 (determinant) 仅对方阵有定义，并且它等于通过将矩阵表示的变换应用到单位(超)立方体上所创建的平行六面体的有向体积 (signed volume)。在二维情况下，这等于将矩阵表示的变换应用到单位正方形上所创建的平行四边形的有向面积 (signed area)。
用于计算 $n$ 维矩阵的行列式的时间复杂度为 $O(n^3)$。对于以上示例，只需要两维形式的公式，即：
$$det\left(\begin{array}{rr}a& b\\ c& d\end{array}\right)=ad-bc$$
因此，对于以上示例，雅可比矩阵的行列式为：
$$\frac{1}{3}\times \frac{1}{2}-0\times 0=\frac{1}{6}$$
此缩放因子 ($\frac{1}{6}$) 用于确保变换后的概率分布积分仍然为 1。根据定义，行列式带有符号，也就是说它可以是负数。因此，我们需要取雅可比矩阵行列式的绝对值才能得到体积的相对变化。

1.4 变量变换方程

我们可以使用变量变换方程描述 $X$ 和 $Z$ 之间的变量变换过程：
$$p_X(x)=p_Z(z)|det(\frac{\partial z}{\partial x})|$$
基于以上理论构建生成模型，关键在于理解，如果 $p_Z(z)$ 是一个简单的分布，我们可以很容易地从中进行采样(例如高斯分布)，那么理论上，我们只需要找到一个适当的可逆函数 $f(x)$，可以将数据 $X$ 映射到 $Z$，以及相应的反函数 $g(z)$ 用于将采样的 $z$ 映射回原始域中的点 $x$。我们可以利用上述包含雅可比行列式的方程来找到数据分布 $p(x)$ 的精确的、可计算的公式。
然而，在实际应用中存在两个主要问题。首先，计算高维矩阵的行列式的计算代价较高，其时间复杂度为 $O(n^3)$。在实践中，这并不可行，因为即使是小尺寸的 32x32 像素灰度图像也有 1,024 个维度。其次，通常并无显式计算可逆函数 $f(x)$ 的方法，我们可以使用神经网络来找到某个函数 $f(x)$，但我们不能保证能够逆转此网络，因为神经网络只能单向工作。
为了解决这两个问题，我们需要使用一种特殊的神经网络架构，确保变量变换函数 $f$ 具有易于计算的行列式，并且是可逆的。接下来，我们将介绍如何使用实值非体积保持 (Real-valued Non-volume Preserving, RealNVP) 变换的技术解决上述问题。

2. RealNVP

实值非体积保持 (Real-valued Non-volume Preserving, RealNVP) 由 Dinh 等人于 2017 年首提出，RealNVP 可以将复杂的数据分布转换为简单的高斯分布，同时还具有可逆性和易于计算的雅可比行列式。

2.1 Two Moons 数据集

在本节中，我们使用由 sklearn 的 make_moons 函数创建的数据集，make_moons 函数能够创建包含噪声的二维点数据集，类似于两个新月形状。

from sklearn.datasets import make_moons
import torch
from torch.utils.data import Dataset, DataLoaderclass TwoMoonsDataset(Dataset):def __init__(self, n_samples=10000, noise=0.05):self.data = torch.tensor(make_moons(n_samples=n_samples, noise=noise)[0], dtype=torch.float32)def __len__(self):return len(self.data)def __getitem__(self, idx):return self.data[idx]dataset = TwoMoonsDataset()
dataloader = DataLoader(dataset, batch_size=256, shuffle=True)device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

接下来，我们将构建一个 RealNVP 模型，可以生成与两个新月形数据集类似分布的二维点。使用此简单示例能够帮助我们详细了解归一化化流模型在实践中的工作原理。首先，我们介绍一种称为耦合层的新层。

2.2 耦合层

耦合层 (coupling layer) 为其输入的每个元素输出一个缩放因子和一个平移因子。换句话说，它会输出两个与输入大小完全相同的张量，一个用于缩放因子 s，一个用于平移因子 t。

为了构建自定义耦合层处理 Two Moons 数据集，我们可以通过堆叠一组 Dense 层来创建缩放输出，并使用另一组 Dense 层来创建平移输出；而对于图像，耦合层可以使用 Conv2D 层替代 Dense 层。

import torch.nn as nnclass CouplingLayer(nn.Module):def __init__(self, dim, hidden_dim, mask):super().__init__()self.dim = dimself.register_buffer('mask', mask)# 缩放和平移网络self.s = nn.Sequential(nn.Linear(dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, dim),nn.Tanh())self.t = nn.Sequential(nn.Linear(dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, dim))def forward(self, x, invert=False):x_masked = x * self.masks = self.s(x_masked) * (1 - self.mask)t = self.t(x_masked) * (1 - self.mask)if not invert:y = x_masked + (1 - self.mask) * (x * torch.exp(s) + t)log_det = s.sum(dim=1)return y, log_detelse:y = x_masked + (1 - self.mask) * ((x - t) * torch.exp(-s))return y

在将输出维度还原回与输入维度相同之前，通过增加维度学习数据更复杂的特征。还可以对每一层使用正则化器以惩罚较大的权重。

2.3 通过耦合层传递数据

耦合层的架构并不复杂，它的独特之处在于输入数据在被送入耦合层时如何进行掩码和变换。

只有数据的前 $d$ 个维度被输入到第一个耦合层中，剩下的 D-d 个维度被掩码(即被设为零)。在本节构建的模型中，我们假设 $D=2$ 且 $d=1$，这意味着耦合层不会看到全部值 $(x_1,x_2)$，而只能看到 $(x_1,0)$。
耦合层的输出是缩放因子和平移因子，输出也会被掩码，但使用的是之前的相反掩码，以便只有后半部分被传递出去，在以上模型中，输出结果为 $(0,s_2)$ 和 $(0,t_2)$。然后，这些输出被逐元素地应用于输入的后半部分 $x_2$，而输入的前半部分 $x_1$ 则直接传递，没有进行任何更新。综上，对于一个维度为 $D$ 且 $d<D$ 的向量，更新方程如下：
$$\begin{gathered} z_{1:d}=x_{1:d} \\ {z}_{d+1:D}=x_{d+1:D}\odot exp(s(x_{1:d}))+t(x_{1:d}) \end{gathered}$$
想要知道为什么需要构建一个掩码了这么多信息的层，我们首先观察此函数的雅可比矩阵的结构：
$$\frac{\partial z}{\partial x}=\left[\begin{array}{cc}\mathbf{I}& 0\\ \frac{\partial z_d+1:D}{\partial x1:d}& diag(exp[s(x_{1:d})])\end{array}\right]$$
左上角的 $d\times d$ 子矩阵是单位矩阵 $\mathbf{I}$，因为 $z_{1:d}=x_{1:d}$，这些元素直接传递而不进行更新。因此，左上角的子矩阵是 0，因为 $z_{d+1:D}$ 不依赖于 $x_{d+1:D}$。
右下角的子矩阵只是一个对角矩阵，使用元素 $exp(s(x1:d))$ 进行填充，因为 $z_{d+1:D}$ 线性依赖于 $x_{d+1:D}$，梯度仅依赖于缩放因子(不依赖于平移因子)。下图展示了该矩阵形式，只有非零元素被填充上颜色。

需要注意的是，对角线上方没有非零元素，因此这种矩阵称为下三角矩阵。这种结构化矩阵(下三角矩阵)的行列式等于对角线元素的乘积。换句话说，行列式不依赖于左下角的复杂导数。因此，可以将这个矩阵的行列式写为：
$$det(J)=exp(\sum _{j}s(x_{1:d}{)}_j)$$
上式十分易于计算的，解决了雅可比行列式的计算问题后，我们继续考虑另一问题，即目标是函数必须易于求其反函数。我们可以通过重新排列正向方程得到可逆函数，如下所示：
$$x_{1:d}=z_{1:d}{x}_{d+1:D}=(z_{d+1:D}-t(x_{1:d}))\odot exp(-s(x_{1:d}))$$
如下图所示。

解决了 RealNVP 模型的构建问题后，我们继续考虑如何更新输入的前 $d$ 个元素。

2.4 堆叠耦合层

为了更新输入的前 $d$ 个元素，我们可以将耦合层堆叠在一起，但交替使用掩码模式，在某一个层保持不变的元素将在下一层中进行更新。这种架构能够学习数据更复杂的表示，因为它是一个更深层的神经网络。
耦合层组合的雅可比矩阵仍然很容易计算，根据线性代数的原理，矩阵乘积的行列式是行列式的乘积。同样地，两个函数组合的反函数就是各自反函数的组合，如下方程所示：
$$\begin{gathered} det(A\cdot B)=det(A)det(B) \\ (f_b\circ f_a{)}^{-1}=f_a^{-1}\circ f_b^{-1} \end{gathered}$$
因此，如果我们堆叠耦合层，并且每次改变掩码模式，就可以构建一个能够转换整个输入张量的神经网络，同时保持具有简单雅可比矩阵行列式和可逆性的关键特性。总体结构如下图所示。

2.5 训练 RealNVP 模型

构建了 RealNVP 模型后，便可以对其进行训练，以学习 Two Moons 数据集的复杂分布。我们希望模型最小化数据的负对数似然 $-log{p}_X(x)$：
$$-log{p}_X(x)=-log{p}_Z(z)-log|det(\frac{\partial z}{\partial x})|$$
选择正向过程f的目标输出分布 $p_Z(z)$ 为标准高斯分布，因为我们可以很容易地从该分布中进行采样。然后，通过应用逆过程 $g$，我们可以将从高斯分布中采样得到的点转换回原始图像域，如下图所示。

接下来，使用 PyTorch 构建 RealNVP 网络。

class RealNVP(nn.Module):def __init__(self, dim, hidden_dim, num_layers):super().__init__()self.layers = nn.ModuleList()mask = torch.zeros(dim, device=device)mask[::2] = 1  # 交替mask模式for _ in range(num_layers):self.layers.append(CouplingLayer(dim, hidden_dim, mask))mask = 1 - mask  # 交替maskself.prior = torch.distributions.MultivariateNormal(torch.zeros(dim).to(device), torch.eye(dim).to(device))def forward(self, x):log_det = 0for layer in self.layers:x, ld = layer(x)log_det += ldreturn x, log_detdef inverse(self, z):for layer in reversed(self.layers):z = layer(z, invert=True)return zmodel = RealNVP(dim=2, hidden_dim=64, num_layers=4).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)for epoch in range(100):total_loss = 0for x in dataloader:x = x.to(device)z, log_det = model(x)# 计算负对数似然log_prob_z = model.prior.log_prob(z)loss = -(log_prob_z + log_det).mean()optimizer.zero_grad()loss.backward()optimizer.step()total_loss += loss.item()print(f"Epoch {epoch+1}, Loss: {total_loss/len(dataloader):.4f}")

3. RealNVP 模型分析

模型训练完成后，使用模型生成样本：

import matplotlib.pyplot as pltwith torch.no_grad():z_samples = model.prior.sample((1000,)).to(device)x_samples = model.inverse(z_samples).cpu()plt.scatter(dataset.data[:,0], dataset.data[:,1], alpha=0.5, label='Real')
plt.scatter(x_samples[:,0], x_samples[:,1], alpha=0.5, label='Generated')
plt.legend()
plt.show()

潜在空间可视化：

z_real = model(dataset.data.to(device))[0].detach().cpu()plt.figure(figsize=(10,4))
plt.subplot(121)
plt.scatter(z_real[:,0], z_real[:,1])
plt.title("Latent Space")
plt.subplot(122)
plt.scatter(dataset.data[:,0], dataset.data[:,1])
plt.title("Original Data")
plt.show()

小结

归一化流模型是由神经网络定义的可逆函数，通过变量变换，直接对数据密度函数进行建模。在一般情况下，变量变换方程需要计算高度复杂的雅可比行列式，但这并不实际。为了解决这一问题，RealNVP 模型限制了神经网络的形式，使其满足两个基本条件：可逆性和易于计算的雅可比行列式。

系列链接

PyTorch生成式人工智能实战：从零打造创意引擎
PyTorch生成式人工智能（1）——神经网络与模型训练过程详解
PyTorch生成式人工智能（2）——PyTorch基础
PyTorch生成式人工智能（3）——使用PyTorch构建神经网络
PyTorch生成式人工智能（4）——卷积神经网络详解
PyTorch生成式人工智能（5）——分类任务详解
PyTorch生成式人工智能（6）——生成模型（Generative Model）详解
PyTorch生成式人工智能（7）——生成对抗网络实践详解
PyTorch生成式人工智能（8）——深度卷积生成对抗网络
PyTorch生成式人工智能（9）——Pix2Pix详解与实现
PyTorch生成式人工智能（10）——CyclelGAN详解与实现
PyTorch生成式人工智能（11）——神经风格迁移
PyTorch生成式人工智能（12）——StyleGAN详解与实现
PyTorch生成式人工智能（13）——WGAN详解与实现
PyTorch生成式人工智能（14）——条件生成对抗网络（conditional GAN，cGAN）
PyTorch生成式人工智能（15）——自注意力生成对抗网络（Self-Attention GAN, SAGAN）
PyTorch生成式人工智能（16）——自编码器（AutoEncoder）详解
PyTorch生成式人工智能（17）——变分自编码器详解与实现
PyTorch生成式人工智能（18）——循环神经网络详解与实现
PyTorch生成式人工智能（19）——自回归模型详解与实现
PyTorch生成式人工智能（20）——像素卷积神经网络(PixelCNN)