AIGC: PNDM (Pseudo Numerical Method for Diffusion Models) 筆記

2023-08-04 09:58 作者:剎那-Ksana- 0人讀過 | 我要投稿

又是一種新的采樣方法，又是一個新的公式地獄？

Runge-Kutta 方法

Runge-Kutta 方法是一種在歐拉方法的基礎上，比歐拉方法精度更高的解常微分方程（Ordinary Differential Equation） $%5Cfrac%7Bdx%7D%7Bdt%7D%3Df(x%2Ct)$ 的數(shù)值解法（Numerical Method）。

其一般的形式為

$x_%7Bt%2B1%7D%3Dx_t%2BhAk_1%2BhBk_2%2BhCk_3%2B...$

$k_1%2Ck_2%2C...$ ?有其固定的形式，A, B, C,... 則是待定系數(shù)。這里有多少個? $k$ ?就代表了多少階?Runge-Kutta.?有關 Runge-Kutte 的公式推導，鏈接將放在文章末尾。

論文里面使用的是4階?Runge-Kutta 方法：

$%5Cbegin%7Balign*%7D%0A%26%20x_%7Bt%2B%5Cdelta%7D%3Dx_t%2B%5Cfrac%7B%5Cdelta%7D%7B6%7D(k_1%2B2k_2%2B2k_3%2Bk_4)%5C%5C%0A%26%20k_1%20%3D%20f(x_t%2Ct)%20%5C%5C%0A%26%20k_2%20%3D%20f(x_t%2B%5Cfrac%7B%5Cdelta%7D%7B2%7Dk_1%2Ct%2B%5Cfrac%7B%5Cdelta%7D%7B2%7D)%20%5C%5C%0A%26%20k_3%20%3D%20f(x_t%2B%5Cfrac%7B%5Cdelta%7D%7B2%7Dk_2%2C%20t%2B%5Cfrac%7B%5Cdelta%7D%7B2%7D)%20%5C%5C%0A%26%20k_4%20%3D%20f(x_t%2B%5Cdelta%20k_3%2C%20t%2B%5Cdelta)%0A%5Cend%7Balign*%7D$

Linear Multistep 方法

Linear Multistep 和?Runge-Kutta 差不多，都是 ODE 的數(shù)值解法，其一般形式為：

$%5Cbegin%7Balign*%7D%0A%26%20%5Calpha_s%20y_%7Bt%2Bs%7D%20%2B%20%5Calpha_%7Bs-1%7D%20y_%7Bt%2Bs-1%7D%20%2B%20...%20%2B%20%5Calpha_0%20y_0%3D%20%5C%5C%0A%26%20h%5Cbeta_s%20f_%7Bt%2Bs%7D%2Bh%20%5Cbeta_%7Bs-1%7D%20f_%7Bt%2Bs-1%7D%20%2B%20...%20%2B%20h%20%5Cbeta_0%20f_t%0A%5Cend%7Balign*%7D$

論文里面采用的形式為：

$x_%7Bt%2B%5Cdelta%7D%3Dx_t%2B%5Cfrac%7B%5Cdelta%7D%7B24%7D(55f_t-59f_%7Bt-%5Cdelta%7D%2B37f_%7Bt-2%5Cdelta%7D-9f_%7Bt-3%5Cdelta%7D)%2C%5C%20f_%7Bt%7D%3Df(x_t%2C%20t)$

ODE 與數(shù)值解

之前在 DDIM?的筆記里面，最后一節(jié)里面也給出了當? $%5Csigma%3D0$ ?時，DDIM 的?ODE 形式，在這里，論文里面對于公式做了一些變形：

$%5Cfrac%7Bdx%7D%7Bdt%7D%20%3D%20-%5Cbar%7B%5Calpha%7D'(t)%5Cleft(%5Cfrac%7Bx(t)%7D%7B2%5Cbar%7B%5Calpha%7D(t)%7D-%5Cfrac%7B%5Cepsilon_%5Ctheta(x(t)%2Ct)%7D%7B2%5Cbar%7B%5Calpha%7D(t)%5Csqrt%7B1-%5Cbar%7B%5Calpha%7D(t)%7D%7D%5Cright)$

至于公式是怎么來的，不是特別重要，故省略步驟了。

有了上述的 ODE 公式以后，可以套用一些數(shù)值解法進行求解。但是論文發(fā)現(xiàn)，數(shù)值解法下如果大幅減少步數(shù)（提速）會引入大量的噪音。

針對這種情況，研究團隊做了調查，發(fā)現(xiàn)問題來源有兩個。

首先，上述 ODE 公式和? $%5Cepsilon_%5Ctheta$ ?只在一小片弧形的區(qū)域里面明確定義（well-defined），在這塊弧形區(qū)域之外，將沒有足夠的樣本讓? $%5Cepsilon_%5Ctheta$ ?去貼合噪音，而傳統(tǒng)的數(shù)值法都是沿著直線計算的，所以會引入誤差。

只有在紅色的弧線區(qū)域里面，才明確定義。區(qū)域之外，都會有很大的誤差。

其次，線性計劃下的? $%5Cbeta_t$ ，當 t 趨近于 0 時，ODE 的值將會趨于無窮，于是違反了數(shù)值解法的先決條件。

解決方法

針對以上問題，論文里面提出了相應的解決方案。

針對第一點，論文提出，應該在一個流形上去求解（把弧線拉直，或者說把直線拉彎）。這里，我們將傳統(tǒng)的數(shù)值解法（即 Runge-Kutta, Linear Multistep 等等）分為兩部分——梯度部分（Gradient Part） $f'$ 和線性的轉移部分（Transfer Part） $%5Ccolor%7Bpurple%7D%7Bx_%7Bt%2B%5Cdelta%7D%7D%3D%5Ccolor%7Bgreen%7D%7Bx_%7Bt%7D%7D%2B%5Cdelta%20%5Ccolor%7Bbrown%7D%7Bf'%7D$ . 如果轉移部分不是線性的，那么這種數(shù)值解法論文里面稱作偽數(shù)值方法（Pseudo Numerical Method）。

很明顯，如果我們的目標分布是弧形的話，我們不能使用線性的轉移方法。并且，如果我們的梯度部分如果是精確的（即，有足夠的數(shù)據(jù)去衡量噪音 $%5Cepsilon_%5Ctheta$ ），那么我們的轉移部分也應該盡量地貼合流形。

所以，在這里如果假設? $x_t%3D%20%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7Dx_0%20%2B%20%5Csqrt%7B1-%5Cbar%7B%5Calpha%7D_t%7D%5Cepsilon$ , 且我們的噪音 $%5Cepsilon_%5Ctheta$ 是精確的，即? $%5Cepsilon_%5Ctheta%3D%5Cepsilon$ , 那么根據(jù) DDIM 的逆向過程：

? $x_%7Bt-%5Cdelta%7D%20%3D%20%5Csqrt%7B%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D%7D%5Cleft(%5Cfrac%7Bx_t-%5Csqrt%7B1-%5Cbar%7B%5Calpha%7D_t%7D%5Cepsilon_%5Ctheta(x_t%2C%20t)%7D%7B%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7D%7D%5Cright)%20%2B%20%5Csqrt%7B1-%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D-%5Csigma%5E2_t%7D%5Cepsilon_%5Ctheta(x_t%2C%20t)%20%2B%20%5Csigma_t%20%5Cepsilon_t$

（注意這里? $%5Csigma_t%3D0$ ）

可以推導得出，對于任意的 $t'%5Cleq%20t$ , 有：

$%20%20%20%5Cbegin%7Bsplit%7D%0A%20%20%20%20%20%20x_%7Bt'%7D%20%26%3D%20%5Csqrt%7B%5Cbar%7B%5Calpha%7D_%7Bt'%7D%7D%5Cleft(%5Cfrac%7Bx_t-%5Csqrt%7B1-%5Cbar%7B%5Calpha%7D_t%7D%5Cepsilon_%5Ctheta%7D%7B%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7D%7D%5Cright)%20%2B%20%5Csqrt%7B1-%5Cbar%7B%5Calpha%7D_%7Bt'%7D%7D%5Cepsilon_%5Ctheta%20%5C%5C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%26%3D%20%5Csqrt%7B%5Cbar%7B%5Calpha%7D_%7Bt'%7D%7D%5Cleft(%5Cfrac%7B%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7Dx_0%20%2B%20%5Csqrt%7B1-%5Cbar%7B%5Calpha%7D_t%7D%5Cepsilon-%5Csqrt%7B1-%5Cbar%7B%5Calpha%7D_t%7D%5Cepsilon%7D%7B%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7D%7D%5Cright)%20%2B%20%5Csqrt%7B1-%5Cbar%7B%5Calpha%7D_%7Bt'%7D%7D%5Cepsilon_%5Ctheta(x_t%2C%20t)%20%5C%5C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%26%3D%20%5Csqrt%7B%5Cbar%7B%5Calpha%7D_%7Bt'%7D%7Dx_0%20%2B%20%5Csqrt%7B1-%5Cbar%7B%5Calpha%7D_%7Bt'%7D%7D%5Cepsilon.%0A%20%20%20%5Cend%7Bsplit%7D$

所以，得出的結論就是，如果? $%5Cepsilon$ ?是精確的，那么? $x_%7Bt'%7D$ ?也是精確的。

如果將上面逆向過程的式子，兩邊同時減去? $x_t$ , 那么就得到了如下的形式：

$x_%7Bt-%5Cdelta%7D%20-%20%7Bx_t%7D%3D%20%5Cfrac%7B(%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D-%5Cbar%7B%5Calpha%7D_t)%20x_t%7D%7B%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7D(%5Csqrt%7B%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D%7D%2B%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7D)%7D%20-%20%0A%20%20%20%5Cfrac%7B(%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D-%5Cbar%7B%5Calpha%7D_t)%20%5Cepsilon_%5Ctheta(x_t%2C%20t)%7D%7B%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7D(%5Csqrt%7B(1-%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D)%5Cbar%7B%5Calpha%7D_%7Bt%7D%7D%20%2B%20%5Csqrt%7B(1-%5Cbar%7B%5Calpha%7D_%7Bt%7D)%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D%7D)%7D$

然后我們對公式做一下變形，將? $(%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D-%5Cbar%7B%5Calpha%7D_t)$ ?拆分成? $(%5Csqrt%7B%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D%7D-%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7D)(%5Csqrt%7B%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D%7D%2B%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7D)$ , 把等式左邊的? $x_t$ ?移到等式右邊，所以：

$x_t%2B%5Cfrac%7B(%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D-%5Cbar%7B%5Calpha%7D_t)%20x_t%7D%7B%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7D(%5Csqrt%7B%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D%7D%2B%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7D)%7D%20%3D%5Cfrac%7B%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7Dx_t%7D%7B%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7D%7D%2B%5Cfrac%7B(%5Csqrt%7B%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D%7D-%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7D)%20x_t%7D%7B%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7D%7D$

進一步化簡一下就得到了論文里面的形式：

$x_%7Bt-%5Cdelta%7D%3D%20%5Cfrac%7B%5Csqrt%7B%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D%7D%7D%7B%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7D%7Dx_t%20-%20%0A%20%20%20%5Cfrac%7B(%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D-%5Cbar%7B%5Calpha%7D_t)%7D%7B%5Csqrt%7B%5Cbar%7B%5Calpha%7D_t%7D(%5Csqrt%7B(1-%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D)%5Cbar%7B%5Calpha%7D_%7Bt%7D%7D%20%2B%20%5Csqrt%7B(1-%5Cbar%7B%5Calpha%7D_%7Bt%7D)%5Cbar%7B%5Calpha%7D_%7Bt-%5Cdelta%7D%7D)%7D%5Cepsilon_t$

所以，這就是我們轉移部分的新公式，并且我們把這個公式命名為? $%5Cphi(x_t%2C%20%5Cepsilon_t%2C%20t%2C%20t-%5Cdelta)$ 。而? $%5Cepsilon_t$ 則是梯度部分。

由于以上的公式是由 DDPM/DDIM 轉換過來的，所以具備的一個特性是，當 t 趨向于 0 時， $%5Cepsilon_%7B%5Ctheta%7D$ 將會越來越精確。于是順帶著解決了上面的第二個問題。

梯度部分

如果只是上面的部分的話，其實沒有什么特別的地方。因為整個公式其實是 DDIM 逆向過程的變形（說了那么多，原來都是白說了?。?/span>。在這里，論文認為，梯度部分的計算可以借用傳統(tǒng)的數(shù)值方法。

例如，在 linear multistep 方法下，我們有：

$%5Cbegin%7Balign*%7D%0A%26e_t%20%3D%20%5Cepsilon_%5Ctheta(x_t%2C%20t)%5C%5C%0A%26e_t'%20%3D%20%5Cfrac%7B1%7D%7B24%7D(55e_t-59e_%7Bt-%5Cdelta%7D%2B37e_%7Bt-2%5Cdelta%7D-9e_%7Bt-3%5Cdelta%7D)%20%5C%5C%0A%26x_%7Bt%2B%5Cdelta%7D%20%3D%20%5Cphi(x_t%2C%20e_t'%2C%20t%2C%20t%2B%5Cdelta)%0A%5Cend%7Balign*%7D$

論文里面把這種方法稱作 PLMS（Pseudo?Linear Multistep）。

由于 LMS 依賴于前三步，所以最初的三步需要用 Runge-Kutta 方法來算出來。Runge-Kutta 方法也有，稱作 PRK，但是太長了不搬過來了，詳見原論文。

基于 PLMS 的 F-PNDM 算法，F(xiàn)代表每一次計算要用到4步的數(shù)據(jù)

Second-order PNDM，每一步的計算利用的是2步的數(shù)據(jù)，但是原理差不多。只不過用的是2階 Runge-Kutta 方法以及 2步形式的 Linear Multistep. 因為大同小異，就不搬過來了，內容見論文附錄 A.3.?

這里是從官方的 Github上面 F-PNDM 摘錄的上述算法的代碼，以及一些個人加入的注釋

一些推薦參考的資料

Runge-Kutte 的資料（英文），Numerical Approximations of Solutions to First-Order Equations -- Other Numerical Methods -- The Runge-Kutta Method 小節(jié)，里面還有例子

：https://www.sciencedirect.com/science/article/abs/pii/B9780128047767000024

Diffusers, 目前來說注釋最完整，最容易理解的代碼庫：https://github.com/huggingface/diffusers/tree/main

官方的 PNDM Github, 里面還有 FON, 歐拉方法等其他采樣方法：https://github.com/luping-liu/PNDM/tree/master

完。

標簽：