引言

擴散模型最早來源于物理中的熱力學，最近卻在人工智能領域大放異彩。還有什么物理理論可以推動生成模型研究的發展呢？最近，來自 MIT 的研究者受到高維電磁理論的啟發，提出了一種稱作泊松流（Poisson Flow）的生成模型。理論上，這種模型具有直觀的圖像和嚴謹的理論；實驗上，它在生成質量、生成速度和魯棒性上往往比擴散模型更好。

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

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

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

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

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

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

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

方法概覽

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

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

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

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

PFGM 的訓練

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

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

PFGM 的采樣

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

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

實驗結果

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

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

?

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

結論

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