引言

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

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

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

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

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

（1）在相同的網(wǎng)絡(luò)結(jié)構(gòu)上，PFGM 的 ODE 生成的樣本質(zhì)量遠(yuǎn)好于擴(kuò)散模型的 ODE；（2）在與擴(kuò)散模型的 SDE （隨機(jī)微分方程）生成質(zhì)量差不多的情況下，PFGM 的 ODE 達(dá)到了 10 倍 - 20 倍的加速；

（3）PFGM 在表達(dá)能力更弱的網(wǎng)絡(luò)結(jié)構(gòu)上比擴(kuò)散模型魯棒。

圖一：樣本點(diǎn)沿著電場線運(yùn)動 。上圖：數(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ù)，也就是研究人員求解的目標(biāo)。由于只需要知道電場線的方向，研究人員推導(dǎo)出了電場線的梯度（勢函數(shù)的梯度）的解析形式：

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

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

PFGM 的訓(xùn)練?

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

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

PFGM 的采樣?

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

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

實(shí)驗(yàn)結(jié)果

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

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

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

結(jié)論

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