ICML 2023 | PFGM++: 一键升级你的扩散模型

最近，我们见证了由物理启发的深度生成模型的大放异彩，如基于热力学的扩散模型（diffusion model）和基于静电学的泊松流模型（PFGM）。比如，扩散模型被广泛应用到文生图（text-to-image), 文生3D（text-to-3D)或者语音、分子生成等等。然而，PFGM在图片生成上有着比扩散模型更优、更鲁棒的表现。

在今年的ICML上，来自MIT的研究者通过扩展PFGM中静电学的框架，提出了新的一族生成模型，统一了扩散模型与PFGM。在他们新提出的扩展泊松流模型（PFGM++）这一框架中，扩散模型和PFGM分别有着各自的优缺点。该论文展示了如何在PFGM++框架中获得同时拥有扩散模型和PFGM两者优点的新模型。此外，PFGM++有着比PFGM更高效的训练方法，并且支持一键升级已有的扩散模型，使其更加鲁棒。在标准的图片生成数据集CIFAR-10，FFHQ 64x64与LSUN Churches 256x256上, PFGM++都优于之前表现最好的扩散模型。

论文链接：https://arxiv.org/abs/2302.04265
代码链接：https://github.com/Newbeeer/pfgmpp

一、PFGM背景

如同扩散模型，PFGM也属于归一流模型 (normalizing flow)。归一流模型通过某个简单分布（比如高斯分布）和复杂数据分布间的双射来生成数据：它首先从简单分布中采样，接着通过双射将初始样本点映射到数据分布中的一点。

图1: PFGM的基本原理：电场线在数据分布和大半球上的均匀分布之间定义了一个双射关系

二、PFGM++: 扩展泊松流模型

如图3所示，PFGM++不但统一了之前的PFGM与扩散模型，还允许人们找到兼顾两者优点的新模型。这离不开PFGM++提出的两个新技术：(1)更高维的增广空间；（2）更高效的训练方式。下文将对其一一介绍。

图3: PFGM++将PFGM和扩散模型统一起来，同时也具备结合它们各自优点的潜力

1.在N+D维空间中的电场

2.更高效的训练方式

如图6所示，在PFGM中，人们需要用很多训练数据来逼近空间中每点的电场方向；但PFGM++借鉴了扩散模型中denoising score-matching [2] 的思想, 通过一个加在干净样本上的扰动函数来学空间中的电场方向。这个新的训练方法不再需要逼近电场方向，并且与常用的条件生成框架相兼容，比如文生图。