引言

擴散模型最早來源于物理中的熱力學，最近卻在人工智能領域大放異彩。還有什么物理理論可以推動生成模型研究的發(fā)展呢？最近，來自 MIT 的研究者受到高維電磁理論的啟發(fā)，提出了一種稱作泊松流（Poisson Flow）的生成模型。理論上，這種模型具有直觀的圖像和嚴謹?shù)睦碚摚粚嶒炆?，它在生成質量、生成速度和魯棒性上往往比擴散模型更好。本文已被NeurIPS 2022接收。

• 論文地址：https://arxiv.org/abs/2209.11178  
• 代碼地址：https://github.com/Newbeeer/Poisson_flow

受到靜電力學的啟發(fā)，研究人員提出了一種新的生成模型，名為泊松流模型 (Poisson Flow Generative Models, or PFGM)。直觀上，該研究可以把 N 維的數(shù)據(jù)點看成在 N+1 維空間中新增維度 z=0 平面上的一群正電荷，它們產(chǎn)生了高維空間中的電場。從 z=0 平面開始沿著它們產(chǎn)生的電場線往外走，該研究能夠把樣本送到一個半球面上（如圖一所示）。這些電場線的方向對應于高維空間中泊松方程 (Poisson Equation)的解的梯度。研究人員證明了當半球的半徑足夠大的時候，電場線能夠把在 z=0 平面上的電荷分布（也就是數(shù)據(jù)分布）轉換為一個在半球面上的均勻分布（圖二）。

PFGM 利用了電場線的可逆性來生成 z=0 平面上的數(shù)據(jù)分布：首先研究人員在大的半球面上均勻采樣，接著讓樣本沿著電場線從球面往 z=0 平面運動，從而生成數(shù)據(jù)。由于沿著電場線的運動可以由一個常微分方程（ODE）描述，因此在實際的采樣中研究人員只需要解一個由電場線方向決定的 ODE。通過電場，PFGM 將一個球面上的簡單分布轉換為一個復雜的數(shù)據(jù)分布。從這個角度來看，PFGM 可以被認為是一種連續(xù)的標準化流(Normalizing Flow)。

在圖像生成實驗中，PFGM 是當前在標準數(shù)據(jù)集 CIFAR-10 上表現(xiàn)最好的標準化流模型，取得了 2.35 的 FID score （圖片質量的度量）。研究人員也展示了 PFGM 的其他一些用途，比如它能夠計算圖片的似然 (likelihood)、進行圖片編輯和擴展到高分辨率的圖片數(shù)據(jù)集上。此外，研究人員發(fā)現(xiàn) PFGM 比近期大熱的擴散模型 (Diffusion Models)有著三個優(yōu)點：

（1）在相同的網(wǎng)絡結構上，PFGM 的 ODE 生成的樣本質量遠好于擴散模型的 ODE；（2）在與擴散模型的 SDE （隨機微分方程）生成質量差不多的情況下，PFGM 的 ODE 達到了 10 倍 - 20 倍的加速；

（3）PFGM 在表達能力更弱的網(wǎng)絡結構上比擴散模型魯棒。

圖一：樣本點沿著電場線運動 。上圖：數(shù)據(jù)分布呈愛心狀；下圖：數(shù)據(jù)分布呈 PFGM 狀

圖二：左圖：泊松場在三維中的軌跡；右圖：在圖像上使用 PFGM 的前向 ODE 和反向 ODE

方法概覽

注意到上述的過程將 N 維數(shù)據(jù)嵌入到了在 N+1 維（多了 z 維度）的空間中。為了方便區(qū)分，研究人員把 N 維數(shù)據(jù)和 N+1 維用 x 和表示。為了得到上述的高維電場線，需要解如下的泊松方程：

其中是位于 z=0 平面上想要生成的數(shù)據(jù)分布；是勢函數(shù)，也就是研究人員求解的目標。由于只需要知道電場線的方向，研究人員推導出了電場線的梯度（勢函數(shù)的梯度）的解析形式：

電場線的軌跡（見圖二）能夠被下面的 ODE 所描述：

在下面的定理中，研究人員證明了上述 ODE 定義了一個高維半球面上的均勻分布和 z=0 平面上的數(shù)據(jù)分布的雙射。這個結論與圖一、圖二的直觀相同：可以通過電場線來還原數(shù)據(jù)分布。

PFGM 的訓練?

給定一個從數(shù)據(jù)分布中采樣得到數(shù)據(jù)集 ，研究人員用該數(shù)據(jù)集所對應的電場線梯度，來近似數(shù)據(jù)分布所對應的電場線梯度：

該電場線梯度是學習目標。該研究通過 perturb 函數(shù)在空間中進行選點，并且平方損失函數(shù)讓神經(jīng)網(wǎng)絡去學習空間中歸一化的電場線梯度, 具體算法如下：

PFGM 的采樣?

當學習完歸一化去學習空間中歸一化的電場線梯度后，可以通過如下的 ODE 對數(shù)據(jù)分布進行采樣：

該 ODE 通過減小 z，使得樣本從大球面沿著電場線逐漸運動到 z=0 平面。此外，該研究提出了將大球面上的均勻分布投影到某個 z 平面以方便 ODE 模擬，并進一步通過變量替換來進一步加速采樣。具體步驟請參見文章的 3.3 節(jié)。

實驗結果

在表一中，該研究使用標準數(shù)據(jù)集 CIFAR-10 來評估不同模型。在該數(shù)據(jù)集上，PFGM 是表現(xiàn)最好的可逆標準化流模型，取得了 2.35 的 FID score。在使用相同的網(wǎng)絡結構 (DDPM++/DDPM++ deep) 的條件下，PFGM 的表現(xiàn)優(yōu)于擴散模型。研究人員同時觀測到，在與擴散模型的 SDE （隨機微分方程）生成質量差不多的情況下，PFGM 達到了 10 倍 - 20 倍的加速，更好地兼顧了生成質量與速度。此外，研究人員發(fā)現(xiàn) PFGM 在表達能力更弱的網(wǎng)絡結構上比擴散模型魯棒，并且在更高維的數(shù)據(jù)集上依然優(yōu)于同等條件下的擴散模型。具體請見文章的實驗章節(jié)。在圖三中，該研究可視化了 PFGM 生成圖片的過程。

表一：CIFAR-10 數(shù)據(jù)上的樣本質量（FID, Inception）與采樣步數(shù) (NFE)

圖三：PFGM 在 CIFAR-10, CelebA 64x64, LSUN bedroom 256x256 上的采樣過程

結論

該研究提出了一個基于泊松方程的生成模型 PFGM。這個模型預測 N+1 維的擴展空間中的歸一化電場線梯度，并通過電場線對應的 ODE 來采樣。實驗中，該研究的模型是當前最好的標準化流模型，并在相同的網(wǎng)絡結構上取得了比擴散模型更好的生成效果與更快的采樣速度。PFGM 的采樣過程對噪聲更魯棒，也能擴展到更高維的數(shù)據(jù)集中。研究人員期望 PFGM 能夠在其他應用領域中也能取得亮眼表現(xiàn)，比如分子生成和 3D 數(shù)據(jù)生成。