卷！MIT泊松流天生模型击败聚集模型，兼顾品质与速度

引言

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

论文地址：https://arxiv.org/abs/2209.11178  

代码地址：https://github.com/Newbeeer/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 64×64, LSUN bedroom 256×256 上的采样过程

结论

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