基于泊松方程的生成模型 PFGM

引言

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

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

PFGM 利用了電場(chǎng)線的可逆性來(lái)生成 z=0 平面上的數(shù)據(jù)分布：首先研究人員在大的半球面上均勻采樣，接著讓樣本沿著電場(chǎng)線從球面往 z=0 平面運(yùn)動(dòng)，從而生成數(shù)據(jù)。由于沿著電場(chǎng)線的運(yùn)動(dòng)可以由一個(gè)常微分方程（ODE）描述，因此在實(shí)際的采樣中研究人員只需要解一個(gè)由電場(chǎng)線方向決定的 ODE。通過(guò)電場(chǎng)，PFGM 將一個(gè)球面上的簡(jiǎn)單分布轉(zhuǎn)換為一個(gè)復(fù)雜的數(shù)據(jù)分布。從這個(gè)角度來(lái)看，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 的其他一些用途，比如它能夠計(jì)算圖片的似然 (likelihood)、進(jìn)行圖片編輯和擴(kuò)展到高分辨率的圖片數(shù)據(jù)集上。此外，研究人員發(fā)現(xiàn) PFGM 比近期大熱的擴(kuò)散模型 (Diffusion Models)有著三個(gè)優(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)沿著電場(chǎng)線運(yùn)動(dòng) 。上圖：數(shù)據(jù)分布呈愛(ài)心狀；下圖：數(shù)據(jù)分布呈 PFGM 狀

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

方法概覽

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

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

電場(chǎng)線的軌跡（見(jiàn)圖二）能夠被下面的 ODE 所描述：

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

PFGM 的訓(xùn)練

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

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

PFGM 的采樣

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

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

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

在表一中，該研究使用標(biāo)準(zhǔn)數(shù)據(jù)集 CIFAR-10 來(lái)評(píng)估不同模型。在該數(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ò)散模型。研究人員同時(shí)觀測(cè)到，在與擴(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ò)散模型。具體請(qǐng)見(jiàn)文章的實(shí)驗(yàn)章節(jié)。在圖三中，該研究可視化了 PFGM 生成圖片的過(guò)程。

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

?

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

結(jié)論

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