当前位置: 首页 > 科技观察

卷!MIT泊松流生成模型打败扩散模型,兼顾质量和速度

时间:2023-03-21 21:43:14 科技观察

简介扩散模型起源于物理学中的热力学,但最近在人工智能领域大放异彩。还有哪些物理理论可以推动生成模型研究的发展?最近,麻省理工学院的研究人员在高维电磁理论的启发下,提出了一种称为泊松流的生成模型。理论上,这种模型形象直观,理论严谨;实验上,它在生成质量、生成速度和鲁棒性方面往往优于扩散模型。该论文已被NeurIPS2022录用。论文地址:https://arxiv.org/abs/2209.11178代码地址:https://github.com/Newbeeer/Poisson_flow受静电力学启发,研究人员提出了一种新的生成模型,称为泊松流模型(PoissonFlowGenerativeModels,或PFGM)。直观上,本研究可以将N维数据点看作是N+1维空间中新维z=0平面上的一组正电荷,它们在高维空间中产生电场。从z=0平面开始,沿着它们产生的电场线,该研究能够将样品传送到半球表面(如图1所示)。这些电场线的方向对应于高维空间中泊松方程解的梯度。研究人员证明,当半球的半径足够大时,电场线可以将z=0平面上的电荷分布(即数据分布)转变为半球上的均匀分布(图2)。PFGM利用电场??线的可逆性,在z=0平面上产生数据分布:首先,研究人员在一个大的半球上均匀采样,然后让样本沿着电场线从球体移动到z=0平面生成数据。由于沿电场线的运动可以用常微分方程(ODE)来描述,因此在实际采样中,研究人员只需求解一个由电场线方向确定的ODE即可。通过电场,PFGM将球体上的简单分布转换为复杂的数据分布。从这个角度来看,PFGM可以被认为是一个连续的归一化流。在图像生成实验中,PFGM是目前标准数据集CIFAR-10上最好的标准化流模型,达到了2.35的FID分数(衡量图像质量的指标)。研究人员还展示了PFGM的一些其他用途,例如它能够计算图像似然(likelihood)、图像编辑以及缩放到高分辨率图像数据集。此外,研究人员还发现PFGM相较于近期的热扩散模型(DiffusionModels)具有三大优势:(1)在相同的网络结构上,PFGM的ODE生成的样本质量远优于扩散模型的ODE;(2)当SDE(随机微分方程)生成质量与扩散模型相似时,PFGM的ODE速度快10-20倍;(3)PFGM在网络结构上比扩散模型鲁棒性更强,表达能力更弱。图1:采样点沿电场线移动。顶部:数据分布呈心形;底部:数据分布呈PFGM图2:左:泊松场在3D中的轨迹;右图:在图像上使用PFGM的前向ODE和反向ODE方法的概述请注意,上述过程将N维数据嵌入到N+1维(附加z维)的空间中。为了方便区分,研究人员用x和表示N维数据和N+1维。为了得到上述高维电场线,需要求解如下泊松方程:其中是在z=0平面上要生成的数据分布;是势函数,这是研究人员要解决的目标。由于只需要知道电场线的方向,研究人员推导出电场线梯度(势函数的梯度)的解析形式:电场线的轨迹(见图2)可以由以下ODE描述:在下面的定理中,研究人员证明了上述ODE定义了高维半球上的均匀分布和z=0平面上数据分布的双射。这个结论与图1和图2的直觉相同:可以通过电场线恢复数据分布。PFGM训练给定一个数据分布采样的数据集,研究者利用数据集对应的电场线梯度来近似数据分布对应的电场线梯度:电场线梯度就是学习目标。本研究使用扰动函数在空间中选择点,平方损失函数让神经网络学习空间中的归一化电场线梯度。具体算法如下:PFGM的采样学习归一化后,去学习归一化空间电场线梯度归一化后,数据分布可以通过以下ODE进行采样:ODE使样本逐渐从通过减小z,沿电场线到z=0平面的大球体。此外,该研究提出将大球体上的均匀分布投影到某个z平面上以方便ODE模拟,并通过变量替换进一步加速采样。具体步骤请参考文章3.3节。实验结果如表1所示。该研究使用标准数据集CIFAR-10来评估不同的模型。在此数据集上,PFGM是性能最好的可逆归一化流动模型,实现了2.35的FID分数。使用相同的网络结构(DDPM++/DDPM++deep),PFGM优于扩散模型。研究人员还观察到,PFGM在与扩散模型的SDE(随机微分方程)生成相同质量的情况下实现了10到20倍的加速,更好地平衡了生成质量和速度。此外,研究人员发现PFGM在表现力较低的网络结构上比扩散模型更稳健,并且在相同条件下在高维数据集上仍优于扩散模型。具体请参考文章实验部分。在图3中,该研究可视化了PFGM生成图像的过程。表1:CIFAR-10数据的样本质量(FID、Inception)和采样步骤数(NFE)图3:CIFAR-10、CelebA64x64、LSUN卧室256x256松散方程生成模型PFGM的PFGM采样过程结论。该模型预测N+1维扩展空间中的归一化电场线梯度,由与电场线对应的ODE采样。在实验中,本研究的模型是目前最好的标准化流模型,在相同的网络结构上取得了比扩散模型更好的生成效果和更快的采样速度。PFGM的采样过程对噪声具有更强的鲁棒性,也可以扩展到更高维的数据集。研究人员期望PFGM在其他应用领域也能表现出色,例如分子生成和3D数据生成。