基于泊松方程的生成模型 PFGM

引言

扩散模型最早来源于物理中的热力学，最近却在人工智能领域大放异彩。还有什么物理理论可以推动生成模型研究的发展呢？最近，来自 MIT 的研究者受到高维电磁理论的启发，提出了一种称作泊松流（Poisson Flow）的生成模型。理论上，这种模型具有直观的图像和严谨的理论；实验上，它在生成质量、生成速度和鲁棒性上往往比扩散模型更好。

受到静电力学的启发，研究人员提出了一种新的生成模型，名为泊松流模型 (Poisson Flow Generative Models, or PFGM)。直观上，该研究可以把 N 维的数据点看成在 N+1 维空间中新增维度 z=0 平面上的一群正电荷，它们产生了高维空间中的电场。从 z=0 平面开始沿着它们产生的电场线往外走，该研究能够把样本送到一个半球面上（如图一所示）。这些电场线的方向对应于高维空间中泊松方程 (Poisson Equation)的解的梯度。研究人员证明了当半球的半径足够大的时候，电场线能够把在 z=0 平面上的电荷分布（也就是数据分布）转换为一个在半球面上的均匀分布（图二）。

PFGM 利用了电场线的可逆性来生成 z=0 平面上的数据分布：首先研究人员在大的半球面上均匀采样，接着让样本沿着电场线从球面往 z=0 平面运动，从而生成数据。由于沿着电场线的运动可以由一个常微分方程（ODE）描述，因此在实际的采样中研究人员只需要解一个由电场线方向决定的 ODE。通过电场，PFGM 将一个球面上的简单分布转换为一个复杂的数据分布。从这个角度来看，PFGM 可以被认为是一种连续的标准化流(Normalizing Flow)。

在图像生成实验中，PFGM 是当前在标准数据集 CIFAR-10 上表现最好的标准化流模型，取得了 2.35 的 FID score （图片质量的度量）。研究人员也展示了 PFGM 的其他一些用途，比如它能够计算图片的似然 (likelihood)、进行图片编辑和扩展到高分辨率的图片数据集上。此外，研究人员发现 PFGM 比近期大热的扩散模型 (Diffusion Models)有着三个优点：

（1）在相同的网络结构上，PFGM 的 ODE 生成的样本质量远好于扩散模型的 ODE；（2）在与扩散模型的 SDE （随机微分方程）生成质量差不多的情况下，PFGM 的 ODE 达到了 10 倍 - 20 倍的加速；

（3）PFGM 在表达能力更弱的网络结构上比扩散模型鲁棒。

图一：样本点沿着电场线运动 。上图：数据分布呈爱心状；下图：数据分布呈 PFGM 状

图二：左图：泊松场在三维中的轨迹；右图：在图像上使用 PFGM 的前向 ODE 和反向 ODE

方法概览

注意到上述的过程将 N 维数据嵌入到了在 N+1 维（多了 z 维度）的空间中。为了方便区分，研究人员把 N 维数据和 N+1 维用 x 和表示。为了得到上述的高维电场线，需要解如下的泊松方程：

其中是位于 z=0 平面上想要生成的数据分布；是势函数，也就是研究人员求解的目标。由于只需要知道电场线的方向，研究人员推导出了电场线的梯度（势函数的梯度）的解析形式：

电场线的轨迹（见图二）能够被下面的 ODE 所描述：

在下面的定理中，研究人员证明了上述 ODE 定义了一个高维半球面上的均匀分布和 z=0 平面上的数据分布的双射。这个结论与图一、图二的直观相同：可以通过电场线来还原数据分布。

PFGM 的训练

给定一个从数据分布中采样得到数据集，研究人员用该数据集所对应的电场线梯度，来近似数据分布所对应的电场线梯度：

该电场线梯度是学习目标。该研究通过 perturb 函数在空间中进行选点，并且平方损失函数让神经网络去学习空间中归一化的电场线梯度, 具体算法如下：

PFGM 的采样

当学习完归一化去学习空间中归一化的电场线梯度后，可以通过如下的 ODE 对数据分布进行采样：

该 ODE 通过减小 z，使得样本从大球面沿着电场线逐渐运动到 z=0 平面。此外，该研究提出了将大球面上的均匀分布投影到某个 z 平面以方便 ODE 模拟，并进一步通过变量替换来进一步加速采样。具体步骤请参见文章的 3.3 节。

实验结果

在表一中，该研究使用标准数据集 CIFAR-10 来评估不同模型。在该数据集上，PFGM 是表现最好的可逆标准化流模型，取得了 2.35 的 FID score。在使用相同的网络结构 (DDPM++/DDPM++ deep) 的条件下，PFGM 的表现优于扩散模型。研究人员同时观测到，在与扩散模型的 SDE （随机微分方程）生成质量差不多的情况下，PFGM 达到了 10 倍 - 20 倍的加速，更好地兼顾了生成质量与速度。此外，研究人员发现 PFGM 在表达能力更弱的网络结构上比扩散模型鲁棒，并且在更高维的数据集上依然优于同等条件下的扩散模型。具体请见文章的实验章节。在图三中，该研究可视化了 PFGM 生成图片的过程。

表一：CIFAR-10 数据上的样本质量（FID, Inception）与采样步数 (NFE)

‍

图三：PFGM 在 CIFAR-10, CelebA 64x64, LSUN bedroom 256x256 上的采样过程

结论

该研究提出了一个基于泊松方程的生成模型 PFGM。这个模型预测 N+1 维的扩展空间中的归一化电场线梯度，并通过电场线对应的 ODE 来采样。实验中，该研究的模型是当前最好的标准化流模型，并在相同的网络结构上取得了比扩散模型更好的生成效果与更快的采样速度。PFGM 的采样过程对噪声更鲁棒，也能扩展到更高维的数据集中。研究人员期望 PFGM 能够在其他应用领域中也能取得亮眼表现，比如分子生成和 3D 数据生成。

审核编辑：彭静