卷!MIT提出泊松流生成模型击败扩散模型,兼顾质量与速度

卷!MIT提出泊松流生成模型击败扩散模型,兼顾质量与速度_第1张图片

©作者 | 机器之心编辑部

来源 | 机器之心

来自 MIT 的研究者提出了一种新的生成模型,名为泊松流模型。它在图像生成质量、生成速度和鲁棒性上比扩散模型更好。本文已被 NeurIPS 2022 接收。

32d83f823fdc5168d0b8f8b3d26bd1b8.png


引言

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

卷!MIT提出泊松流生成模型击败扩散模型,兼顾质量与速度_第2张图片

论文标题:

Poisson Flow Generative Models

论文链接:

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 状

卷!MIT提出泊松流生成模型击败扩散模型,兼顾质量与速度_第3张图片

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

d252453a5a883d51976ef4e28205c32c.png

方法概览

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

2b9bc044ea54810851bb04819c6c099e.png

其中28ab362bdc086ec41596a660f2a73ea9.png是位于 z=0 平面上想要生成的数据分布;d25e51f291a24c100401ddf6329835b0.png是势函数,也就是研究人员求解的目标。由于只需要知道电场线的方向,研究人员推导出了电场线的梯度(势函数的梯度)的解析形式:

83297a6283d8f6b70b872f5efd791df6.png

电场线的轨迹(见图二)能够被下面的 ODE 所描述:

a77803c701d8b240f28b8e396e8ccdb0.png

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

f6bb2df3de4851d36fe3415f0ef4affa.png

2.1 PFGM的训练

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

144418a1c680ff4e05802cea832e7cf6.png

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

卷!MIT提出泊松流生成模型击败扩散模型,兼顾质量与速度_第4张图片

2.2 PFGM的采样

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

958baec9d7fa479416e77891962ef8ed.png

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

6c47cff82fcfa48b692b6cff6a8b95e4.png

实验结果

在表一中,该研究使用标准数据集 CIFAR-10 来评估不同模型。在该数据集上,PFGM 是表现最好的可逆标准化流模型,取得了 2.35 的 FID score。在使用相同的网络结构 (DDPM++/DDPM++ deep) 的条件下,PFGM 的表现优于扩散模型。研究人员同时观测到,在与扩散模型的 SDE (随机微分方程)生成质量差不多的情况下,PFGM 达到了 10 倍 - 20 倍的加速,更好地兼顾了生成质量与速度。

此外,研究人员发现 PFGM 在表达能力更弱的网络结构上比扩散模型鲁棒,并且在更高维的数据集上依然优于同等条件下的扩散模型。具体请见文章的实验章节。在图三中,该研究可视化了 PFGM 生成图片的过程。

卷!MIT提出泊松流生成模型击败扩散模型,兼顾质量与速度_第5张图片

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

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

e2a2bdb64af26652b1e2f056a5efd5c3.png


结论

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

更多阅读

卷!MIT提出泊松流生成模型击败扩散模型,兼顾质量与速度_第6张图片

卷!MIT提出泊松流生成模型击败扩散模型,兼顾质量与速度_第7张图片

卷!MIT提出泊松流生成模型击败扩散模型,兼顾质量与速度_第8张图片

010efa8aa792b01c2121d1b430d3b3a3.gif

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢?答案就是:你不认识的人。

总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是最新论文解读,也可以是学术热点剖析科研心得竞赛经验讲解等。我们的目的只有一个,让知识真正流动起来。

 稿件基本要求:

• 文章确系个人原创作品,未曾在公开渠道发表,如为其他平台已发表或待发表的文章,请明确标注 

• 稿件建议以 markdown 格式撰写,文中配图以附件形式发送,要求图片清晰,无版权问题

• PaperWeekly 尊重原作者署名权,并将为每篇被采纳的原创首发稿件,提供业内具有竞争力稿酬,具体依据文章阅读量和文章质量阶梯制结算

 投稿通道:

• 投稿邮箱:[email protected] 

• 来稿请备注即时联系方式(微信),以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信(pwbot02)快速投稿,备注:姓名-投稿

卷!MIT提出泊松流生成模型击败扩散模型,兼顾质量与速度_第9张图片

△长按添加PaperWeekly小编

现在,在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·

6356072ae9f536d2e1b2b847bcdba6a3.jpeg

你可能感兴趣的:(人工智能,机器学习,深度学习,计算机视觉,神经网络)