基於流的生成模型

基於流的生成模型（flow-based generative model）是機器學習中的一類生成模型，利用歸一化流（normalizing flow）顯式建模機率分布。這是一種使用機率密度變量變換法將簡單分布轉換為複雜分布的統計方法。

直接建模似然函數具有很多優點。例如，可以直接計算得到負對數似然並將其作為損失函數最小化。此外，通過從初始分布中採樣並應用流變換可以生成新的樣本。

相比之下，變分自編碼器、生成對抗網絡等其他生成模型無法顯式地表示似然函數。

## 方法

<image>

考慮隨機變量$z_1$和$z_0$，其中$z_0 = f^{-1}_1(z_1)$。對於$i = 1, ..., K$，定義一系列由$z_0$變換得到的隨機變量$z_i = f_i(z_{i-1})$。其中函數$f_1, ..., f_K$滿足可逆性，即存在反函數$f^{-1}_i$。最終輸出變量$z_K$用於對目標分布進行建模。

$z_K$的對數似然為（參見下方推導過程）：

$$\log p_K(z_K) = \log p_0(z_0) - \sum_{i=1}^{K} \log \left|\det \frac{df_i(z_{i-1})}{dz_{i-1}}\right|$$

為了高效計算對數似然，函數$f_1, ..., f_K$應既易於求逆，也易於計算其雅可比矩陣的行列式。在實踐中，這些函數通常使用深度神經網路建模，並通過訓練以最小化目標分布數據樣本的負對數似然。這些架構一般設計為只需神經網路的正向傳播即可完成逆運算與雅可比行列式的計算，例如NICE、RealNVP、Glow等。

### 對數似然的推導

考慮$z_1$與$z_0$，兩者之間滿足$z_0 = f^{-1}_1(z_1)$。通過機率密度變量變換公式，$z_1$的分布為：

$$p_1(z_1) = p_0(z_0)\left|\det \frac{df_1^{-1}(z_1)}{dz_1}\right|$$

其中$\det \frac{df_1^{-1}(z_1)}{dz_1}$是$f^{-1}_1$的雅可比矩陣的行列式。

由反函數定理可以得到

$$p_1(z_1) = p_0(z_0)\left|\det \left(\frac{df_1(z_0)}{dz_0}\right)^{-1}\right|$$

使用行列式性質$\det(A^{-1}) = \det(A)^{-1}$（其中$A$是可逆矩陣），則有：

$$p_1(z_1) = p_0(z_0)\left|\det \frac{df_1(z_0)}{dz_0}\right|^{-1}$$

對上式取對數後得到對數似然：

$$\log p_1(z_1) = \log p_0(z_0) - \log \left|\det \frac{df_1(z_0)}{dz_0}\right|$$

對於任意$z_i$和$z_{i-1}$都能得到類似結論。最終由遞歸關係可以得到：

$$\log p_K(z_K) = \log p_0(z_0) - \sum_{i=1}^{K} \log \left|\det \frac{df_i(z_{i-1})}{dz_{i-1}}\right|$$

## 訓練方法

與訓練其他一些深度學習模型類似，歸一化流的目標是最小化模型的似然分布與目標分布之間的KL散度。將模型的似然分布記為$p_\theta$，要學習的目標分布記為$p^*$，則（正向）KL散度為：

$$D_{KL}[p^{*}(x)||p_{\theta}(x)] = -\mathbb{E}_{p^{*}(x)}[\log(p_{\theta}(x))] + \mathbb{E}_{p^{*}(x)}[\log(p^{*}(x))]$$

上式右邊第二項表示目標分布的熵，與模型參數$\theta$無關，因此在優化時可以忽略，留下需要最小化的項是目標分布下的負對數似然的期望。由於此項難以直接計算，可以通過重要性採樣的蒙特卡洛方法來近似。若已從目標分布$p^*(x)$中獨立採樣得到的數據集$\{x_{i}\}_{i=1:N}$從的樣本，則該項可近似估計為：

$$-\hat{\mathbb{E}}_{p^{*}(x)}[\log(p_{\theta}(x))] = -\frac{1}{N} \sum_{i=0}^{N} \log(p_{\theta}(x_{i}))$$

因此，可以將學習目標

$$\underset{\theta}{\operatorname{arg\,min}}\ D_{KL}[p^{*}(x)||p_{\theta}(x)]$$

替換為

$$\underset{\theta}{\operatorname{arg\,max}}\ \sum_{i=0}^{N} \log(p_{\theta}(x_{i}))$$

換句話說，最小化KL相度相當於最大化模型在觀測樣本下的似然。

訓練歸一化流的偽代碼如下：

- 輸入：數據集$x_{1:n}$，歸一化流模型$f_\theta(\cdot), p_0$
- 求解：通過梯度下降法最化$\max_\theta \sum_j \ln p_\theta(x_j)$
- 輸出：優化後的參數$\hat\theta$

## 變體

### 平面流

平面流（planar flow）是最早的歸一化流方法。給定某個激活函數$h$，以及具有適當維度的參數$\theta = (u, w, b)$，可以定義$$x = f_\theta(z) = z + u h(\langle w, z \rangle + b)$$一般而言，逆映射$f_\theta^{-1}$沒有解析解。

相應的雅可比行列式為$|\det (I + h'(\langle w, z \rangle + b) uw^T)| = |1 + h'(\langle w, z \rangle + b) \langle u, w\rangle|$。

為保證其處處可逆，行列式必須在整個定義域內非零。例如當$h = \tanh$、$\langle u, w \rangle > -1$時可以滿足可逆性要求。

### 非線性獨立成分估計（NICE）

非線性獨立成分估計（nonlinear independent components estimation，簡稱NICE）假設$x, z\in \R^{2n}$是偶數維變量，並將其從中間分成兩部分。此時歸一化流的定義為

$$x = \begin{bmatrix} <pre> x_1 \\ x_2 \end{bmatrix}= f_\theta(z) = \begin{bmatrix} z_1 \\z_2 \end{bmatrix} + \begin{bmatrix} 0 \\ m_\theta(z_1) \end{bmatrix}$$

其中$m_\theta$是任何帶有權重$\theta$的神經網路。

其逆映射$f_\theta^{-1}$為$z_1 = x_1, z_2 = x_2 - m_\theta(x_1)$，雅可比行列式為1，即該歸一化流動是體積保持（volume-preserving）的。

當$n=1$時，這一映身可以視為沿$x_2$方向的一種曲線剪切。

### 實值非體積保持（Real NVP）

實值非體積保持（real non-volume preserving，簡稱Real NVP）是NICE模型的一種推廣，定義為：

$$x = \begin{bmatrix} <pre> x_1 \\ x_2 \end{bmatrix}= f_\theta(z) = \begin{bmatrix} z_1 \\ e^{s_\theta(z_1)} \odot z_2 \end{bmatrix} + \begin{bmatrix} 0 \\ m_\theta(z_1) \end{bmatrix}$$

其逆映射是$z_1 = x_1, z_2 = e^{-s_\theta (x_1)}\odot (x_2 - m_\theta (x_1))$，相應的雅可比行列式為$\prod^n_{i=1} e^{s_\theta(z_{1, })}$。當$s_\theta = 0$時，退化為NICE模型。由於Real NVP映射將向量$x$的兩部分分開處理，通常需要在每一層後添加一個置換操作$(x_1, x_2) \mapsto (x_2, x_1)$。

### 生成流（Glow）

生成流（generative flow，簡稱Glow）模型中每層由三個部分組成：

- 通道方向的仿射變換$$y_{cij} = s_c(x_{cij} + b_c)$$相應的雅可比行列式為$\prod_c s_c^{HW}$。
- 可逆1x1卷積$$z_{cij} = \sum_{c'} K_{cc'} y_{cij}$$相應的雅可比行列式為$\det(K)^{HW}$，其中$K$是任意可逆矩陣。
- Real NVP部分，其雅可比行列式如前所述。

生成流通過引入可逆1x1卷積，改進了Real NVP中僅僅置換前後兩部分的方式，而是對所有層的通道進行總體上的置換。

### 掩碼自回歸流（MAF）

掩碼自回歸流（masked autoregresssive flow，簡稱MAF）基於自回歸模型，其定義了一個分布在$\R^n$上的隨機過程：

$$\begin{align} <pre> x_1 \sim& N(\mu_1, \sigma_1^2)\\ x_2 \sim& N(\mu_2(x_1), \sigma_2(x_1)^2)\\ &\cdots \\ x_n \sim& N(\mu_n(x_{1:n-1}), \sigma_n(x_{1:n-1})^2)\\ </pre> \end{align}$$

其中$\mu_i: \R^{i-1} \to \R$和$\sigma_i: \R^{i-1} \to (0, \infty)$是定義自回歸模型的固定函數。

使用重參數化技巧，該自回歸模型可以被推廣為歸一化流：

$$\begin{align} <pre> x_1 =& \mu_1 + \sigma_1 z_1\\ x_2 =& \mu_2(x_1) + \sigma_2(x_1) z_2\\ &\cdots \\ x_n =& \mu_n(x_{1:n-1}) + \sigma_n(x_{1:n-1}) z_n\\ </pre> \end{align}$$

令$z \sim N(0, I_{n})$ 可以重新得到自回歸模型。

正向映射由於是順序性的因而會很慢，但反向映射因為可以並列而會比較很快。

相應的雅可比矩陣是下對角矩陣，其行列式為$\sigma_1 \sigma_2(x_1)\cdots \sigma_n(x_{1:n-1})$。

通過反轉$f_\theta$和$f_\theta^{-1}$這兩個映射，可以得到逆自回歸流 (inverse autoregressive flow，簡稱IAF)。與MAF相反，IAF的正向映射較快、反向映射較慢。

### 連續歸一化流（CNF）

除了通過函數組合來構建流的方法，另一種方法是將流表示為連續時間動力學，得到連續歸一化流（continuous normalizing flow，簡稱CNF）。設$z_0$為具有分布$p(z_0)$的潛變量，使用以下流函數將此潛變量映射到數據空間：

$$x = F(z_0) = z_T = z_0 + \int_0^T f(z_t, t) dt$$

其中$f$是任意函數，可以使用神經網路等進行建模。其反函數為：

$$z_0 = F^{-1}(x) = z_T + \int_T^0 f(z_t, t) dt = z_T - \int_0^T f(z_t,t) dt$$

於是可以得到$x$的對數似然：

$$\log(p(x)) = \log(p(z_0)) - \int_0^T \text{Tr}\left[\frac{\partial f}{\partial z_t} \right] dt$$

由於上式中的跡僅取決於雅可比矩陣$\partial_{z_t} f$的對角線，這意味著對雅可比矩陣的形式沒有任何限制。這與先前的離散歸一化模型不同，後者將雅可比矩陣設計為上對角或下對角形式，以便更高效地計算其行列式。

其中的跡可以使用「哈欽森技巧」（Hutchinson's trick）來估計：給定任意矩陣$W\in \R^{n\times n}$以及滿足$E[uu^T] = I$的任意隨機向量$u\in \R^n$，可以得到$E[u^T W u] = tr(W)$。

通常，隨機向量$u$可以從正態分布$N(0, I)$ 或拉德馬赫分布$\{\pm n^{-1/2}\}^n$中進行採樣。

當$f$由神經網路實現時，需要使用神經常微分方程。實際上，CNF最早是與神經常微分方程在同一篇論文中提出的。

CNF主要存在兩個缺陷：一是連續流必須是同胚的，從而保持方向性和環境同痕（例如，不可能通過空間連續變形將左手翻轉為右手，也不可能將球面外翻或解開一個結），二是學習到的流$f$可能會由於退化而表現不佳（有無數個可能的$f$都能解決同一問題)。

通過增加額外的維度，CNF可以獲得足夠的自由度來反轉方向並突破環境同痕（類似於可以在三維空間中翻轉一個多邊形，或在四維空間中解開一個結），從而得到「增強神經常微分方程」。

通過結合流形的惠特尼嵌入定理和神經網路的通用近似定理，能夠證明$\R^n$的任何同胚可以由$\R^{2n+1}$上的神經常微分方程近似。

還可以為流$f$引入正則化損失，例如基於最優傳輸理論的正則化損失：

$$\lambda_{K} \int_{0}^{T}\left\|f(z_t, t)\right\|^{2} dt +\lambda_{J} \int_{0}^{T}\left\|\nabla_z f(z_t, t)\right\|_F^{2} dt$$

其中$\lambda_K, \lambda_J >0$是超參數。第一項懲罰模型隨時間變化流場的振盪，第二項則懲罰模型隨空間變化流場的振盪。這兩項共同引導模型生成在空間和時間上平滑的流。

## 缺點

儘管歸一化流在估計高維機率密度函數方面取得了成功，但其設計仍然存在一些缺陷。首先，歸一化流的潛空間並不是一個低維空間，因此基於流的模型默認情況下不支持數據壓縮，需要很大的計算量。不過，仍有辦法可以用它們進行圖像壓縮。

此外，基於流的模型在估計分布外樣本（即非訓練集分布中抽取的樣本）的似然值時通常表現不佳。學者提出了一些假設來解釋這一現象，其中包括典型集假設、模型訓練中的估計問題或由於數據分布熵引起的基礎性問題。

歸一化流最有趣的特性之一是其學習到的雙射映射的可逆性。這一特性通過模型設計中的約束得以保證，從而確保理論上的可逆性。逆映射對確保變量變換定理的適用性、雅可比行列式的計算以及模型採樣都至關重要。然而在實踐中，由於數值不精確性，這種可逆性可能被破壞，進而導致逆映射爆炸。

## 應用

基於流的生成模型已應用於多種場景，例如：

- 音訊生成
- 圖像生成
- 分子圖生成
- 點雲建模
- 影片生成
- 有損圖像壓縮
- 異常檢測