變分自編碼器（二）：從貝葉斯觀點出發

源起 #

前幾天寫了博文《變分自編碼器（一）：原來是這麼一回事》，從一種比較通俗的觀點來理解變分自編碼器（VAE），在那篇文章的視角中，VAE跟普通的自編碼器差別不大，無非是多加了噪聲並對噪聲做了約束。然而，當初我想要弄懂VAE的初衷，是想看看究竟貝葉斯學派的機率圖模型究竟是如何與深度學習結合來發揮作用的，如果僅僅是得到一個通俗的理解，那顯然是不夠的。

所以我對VAE繼續思考了幾天，試圖用更一般的、機率化的語言來把VAE說清楚。事實上，這種思考也能回答通俗理解中無法解答的問題，比如重構損失用MSE好還是交叉熵好、重構損失和KL損失應該怎麼平衡，等等。

建議在閱讀《變分自編碼器（一）：原來是這麼一回事》後對本文進行閱讀，本文在內容上儘量不與前文重複。

準備 #

在進入對VAE的描述之前，我覺得有必要把一些概念性的內容講一下。

數值計算vs取樣計算 #

對於不是很熟悉機率統計的讀者，容易混淆的兩個概念應該是數值計算和取樣計算，也有讀者在《三味Capsule：矩陣Capsule與EM路由》出現過同樣的疑惑。比如已知機率密度函式$p(x)$，那麼$x$的期望也就定義為
$$\mathbb{E}[x] = \int x p(x)dx\tag{1}$$
如果要對它進行數值計算，也就是數值積分，那麼可以選若干個有代表性的點$x_0 < x_1 < x_2 < \dots < x_n$，然後得到
$$\mathbb{E}[x] \approx \sum_{i=1}^n x_i p(x_i) \left(x_i - x_{i-1}\right)\tag{2}$$
這裡不討論“有代表性”是什麼意思，也不討論提高數值計算精度的方法。這樣寫出來，是為了跟取樣計算對比。如果從$p(x)$中取樣若干個點$x_1,x_2,\dots,x_n$，那麼我們有
$$\mathbb{E}[x] \approx \frac{1}{n}\sum_{i=1}^n x_i,\quad x_i \sim p(x)\tag{3}$$
我們可以比較$(2)$跟$(3)$，它們的主要區別是$(2)$中包含了機率的計算而$(3)$中僅有$x$的計算，這是因為在$(3)$中$x_i$是從$p(x)$中依機率取樣出來的，機率大的$x_i$出現的次數也多，所以可以說取樣的結果已經包含了$p(x)$在裡邊，就不用再乘以$p(x_i)$了。

更一般地，我們可以寫出
$$\mathbb{E}_{x\sim p(x)}[f(x)] = \int f(x)p(x)dx \approx \frac{1}{n}\sum_{i=1}^n f(x_i),\quad x_i\sim p(x)\tag{4}$$
這就是蒙特卡洛模擬的基礎。

KL散度及變分 #

我們通常用KL散度來度量兩個機率分佈$p(x)$和$q(x)$之間的差異，定義為
$$KL\Big(p(x)\Big\Vert q(x)\Big) = \int p(x)\ln \frac{p(x)}{q(x)} dx=\mathbb{E}_{x\sim p(x)}\left[\ln \frac{p(x)}{q(x)}\right]\tag{5}$$
KL散度的主要性質是非負性，如果固定$p(x)$，那麼$KL\Big(p(x)\Big\Vert q(x)\Big)=0 \Leftrightarrow p(x)=q(x)$；如果固定$q(x)$，同樣有$KL\Big(p(x)\Big\Vert q(x)\Big)=0 \Leftrightarrow p(x)=q(x)$，也就是不管固定哪一個，最小化KL散度的結果都是兩者儘可能相等。這一點的嚴格證明要用到變分法，而事實上VAE中的V（變分）就是因為VAE的推導就是因為用到了KL散度（進而也包含了變分法）。

當然，KL散度有一個比較明顯的問題，就是當$q(x)$在某個區域等於0，而$p(x)$在該區域不等於0，那麼KL散度就出現無窮大。這是KL散度的固有問題，我們只能想辦法規避它，比如隱變數的先驗分佈我們用高斯分佈而不是均勻分佈，原因便在此，這一點我們在前文《變分自編碼器（一）：原來是這麼一回事》中也提到過了。

順便說點題外話，度量兩個機率分佈之間的差異只有KL散度嗎？當然不是，我們可以看維基百科的Statistical Distance一節，裡邊介紹了不少分佈距離，比如有一個很漂亮的度量，我們稱之為巴氏距離（Bhattacharyya distance），定義為
$$D_B\Big(p(x), q(x)\Big)=-\ln\int \sqrt{p(x)q(x)} dx\tag{6}$$
這個距離不僅對稱，還沒有KL散度的無窮大問題。然而我們還是選用KL散度，因為我們不僅要理論上的漂亮，還要實踐上的可行，KL散度可以寫成期望的形式，這允許我們對其進行取樣計算，相反，巴氏距離就沒那麼容易了，讀者要是想把下面計算過程中的KL散度替換成巴氏距離，就會發現寸步難行了。

本文的符號表 #

講解VAE免不了出現大量的公式和符號，這裡將部分式子的含義提前列舉如下：

$$\begin{array}{c|c}
\hline
x_k, z_k & \text{表示隨機變數}x,z\text{的第}k\text{個樣本}\\
\hline
x_{(k)}, z_{(k)} & \text{表示多元變數}x,z\text{的第}k\text{個分量}\\
\hline
\mathbb{E}_{x\sim p(x)}[f(x)] & \text{表示對}f(x)\text{算期望，其中}x\text{的分佈為}p(x)\\
\hline
KL\Big(p(x)\Big\Vert q(x)\Big)& \text{兩個分佈的}KL\text{散度}\\
\hline
\Vert x\Vert^2& \text{向量}x\text{的}l^2\text{範數，也就是我們通常說的模長的平方}\\
\hline
\mathcal{L}& \text{本文的損失函式的符號}\\
\hline
D,d & D\text{是輸入}x\text{的維度，}d\text{是隱變數}z\text{的維度}\\
\hline
\end{array}$$

框架 #

這裡透過直接對聯合分佈進行近似的方式，簡明快捷地給出了VAE的理論框架。

直面聯合分佈 #

出發點依然沒變，這裡再重述一下。首先我們有一批資料樣本$\{x_1,\dots,x_n\}$，其整體用$x$來描述，我們希望藉助隱變數$z$描述$x$的分佈$\tilde{p}(x)$：
$$q(x)=\int q(x|z)q(z)dz,\quad q(x,z) = q(x|z)q(z)\tag{7}$$
這裡$q(z)$是先驗分佈（標準正態分佈），目的是希望$q(x)$能逼近$\tilde{p}(x)$。這樣（理論上）我們既描述了$\tilde{p}(x)$，又得到了生成模型$q(x|z)$，一舉兩得。

接下來就是利用KL散度進行近似。但我一直搞不明白的是，為什麼從原作《Auto-Encoding Variational Bayes》開始，VAE的教程就聚焦於後驗分佈$p(z|x)$的描述？也許是受了EM演算法的影響，這個問題上不能應用EM演算法，就是因為後驗分佈$p(z|x)$難以計算，所以VAE的作者就聚焦於$p(z|x)$的推導。

但事實上，直接來對$p(x,z)$進行近似是最為乾脆的。具體來說，定義$p(x,z)=\tilde{p}(x)p(z|x)$，我們設想用一個聯合機率分佈$q(x,z)$來逼近$p(x,z)$，那麼我們用KL散度來看它們的距離：
$$KL\Big(p(x,z)\Big\Vert q(x,z)\Big) = \iint p(x,z)\ln \frac{p(x,z)}{q(x,z)} dzdx\tag{8}$$
KL散度是我們的終極目標，因為我們希望兩個分佈越接近越好，所以KL散度越小越好。當然，由於現在$p(x,z)$也有引數，所以不單單是$q(x,z)$來逼近$p(x,z)$，$p(x,z)$也會主動來逼近$q(x,z)$，兩者是相互接近。

於是我們有
$$\begin{aligned}KL\Big(p(x,z)\Big\Vert q(x,z)\Big) =& \int \tilde{p}(x) \left[\int p(z|x)\ln \frac{\tilde{p}(x)p(z|x)}{q(x,z)} dz\right]dx\\
=& \mathbb{E}_{x\sim \tilde{p}(x)} \left[\int p(z|x)\ln \frac{\tilde{p}(x)p(z|x)}{q(x,z)} dz\right]
\end{aligned}\tag{9}$$
這樣一來利用$(4)$式，把各個$x_i$代入就可以進行計算了，這個式子還可以進一步簡化，因為$\ln \frac{\tilde{p}(x)p(z|x)}{q(x,z)}=\ln \tilde{p}(x) + \ln \frac{p(z|x)}{q(x,z)}$，而
$$\begin{aligned}\mathbb{E}_{x\sim \tilde{p}(x)} \left[\int p(z|x)\ln \tilde{p}(x)dz\right] =& \mathbb{E}_{x\sim \tilde{p}(x)} \left[\ln \tilde{p}(x)\int p(z|x)dz\right]\\
=&\mathbb{E}_{x\sim \tilde{p}(x)} \big[\ln \tilde{p}(x)\big]
\end{aligned}\tag{10}$$
注意這裡的$\tilde{p}(x)$是根據樣本$x_1,x_2,\dots,x_n$確定的關於$x$的先驗分佈，儘管我們不一定能準確寫出它的形式，但它是確定的、存在的，因此這一項只是一個常數，所以可以寫出
$$\mathcal{L}=KL\Big(p(x,z)\Big\Vert q(x,z)\Big) - \text{常數}= \mathbb{E}_{x\sim \tilde{p}(x)} \left[\int p(z|x)\ln \frac{p(z|x)}{q(x,z)} dz\right]\tag{11}$$
目前最小化$KL\Big(p(x,z)\Big\Vert q(x,z)\Big)$也就等價於最小化$\mathcal{L}$。注意減去的常數為$\mathbb{E}_{x\sim \tilde{p}(x)} \big[\ln \tilde{p}(x)\big]$，所以$\mathcal{L}$擁有下界$-\mathbb{E}_{x\sim \tilde{p}(x)} \big[\ln \tilde{p}(x)\big]$～注意到$\tilde{p}(x)$不一定是機率，在連續情形時$\tilde{p}(x)$是機率密度，它可以大於1也可以小於1，所以$-\mathbb{E}_{x\sim \tilde{p}(x)} \big[\ln \tilde{p}(x)\big]$不一定是非負，即loss可能是負數。

你的VAE已經送達 #

到這裡，我們回顧初衷——為了得到生成模型，所以我們把$q(x,z)$寫成$q(x|z)q(z)$，於是就有
$$\begin{aligned}\mathcal{L} =& \mathbb{E}_{x\sim \tilde{p}(x)} \left[\int p(z|x)\ln \frac{p(z|x)}{q(x|z)q(z)} dz\right]\\
=&\mathbb{E}_{x\sim \tilde{p}(x)} \left[-\int p(z|x)\ln q(x|z)dz+\int p(z|x)\ln \frac{p(z|x)}{q(z)}dz\right]\end{aligned}\tag{12}$$
再簡明一點，那就是
$$\begin{aligned}\mathcal{L} = &\mathbb{E}_{x\sim \tilde{p}(x)} \left[\mathbb{E}_{z\sim p(z|x)}\big[-\ln q(x|z)\big]+\mathbb{E}_{z\sim p(z|x)}\Big[\ln \frac{p(z|x)}{q(z)}\Big]\right]\\
= &\mathbb{E}_{x\sim \tilde{p}(x)} \Bigg[\mathbb{E}_{z\sim p(z|x)}\big[-\ln q(x|z)\big]+KL\Big(p(z|x)\Big\Vert q(z)\Big)\Bigg]
\end{aligned}\tag{13}$$
看，括號內的不就是VAE的損失函式嘛？只不過我們換了個符號而已。我們就是要想辦法找到適當的$q(x|z)$和$q(z)$使得$\mathcal{L}$最小化。

再回顧一下整個過程，我們幾乎都沒做什麼“讓人難以想到”的形式變換，但VAE就出來了。所以，沒有必要去對後驗分佈進行分析，直面聯合分佈，我們能更快捷地到達終點。

不能搞分裂～ #

鑑於$(13)$式的特點，我們也許會將$\mathcal{L}$分開為兩部分看：$\mathbb{E}_{z\sim p(z|x)}\big[-\ln q(x|z)\big]$的期望和$KL\Big(p(z|x)\Big\Vert q(z)\Big)$的期望，並且認為問題變成了兩個loss的分別最小化。

然而這種看法是不妥的，因為$KL\Big(p(z|x)\Big\Vert q(z)\Big)=0$意味著$z$沒有任何辨識度，所以$-\ln q(x|z)$不可能小（預測不準），而如果$-\ln q(x|z)$小則$q(x|z)$大，預測準確，這時候$p(z|x)$不會太隨機，即$KL\Big(p(z|x)\Big\Vert q(z)\Big)$不會小，所以這兩部分的loss其實是相互拮抗的。所以，$\mathcal{L}$不能割裂來看，而是要整體來看，整個的$\mathcal{L}$越小模型就越接近收斂，而不能只單獨觀察某一部分的loss。

事實上，這正是GAN模型中夢寐以求的——有一個總指標能夠指示生成模型的訓練程序，在VAE模型中天然就具備了這種能力了，而GAN中要到WGAN才有這麼一個指標～

實驗 #

截止上面的內容，其實我們已經完成了VAE整體的理論構建。但為了要將它付諸於實驗，還需要做一些工作。事實上原論文《Auto-Encoding Variational Bayes》也在這部分做了比較充分的展開，但遺憾的是，網上很多VAE教程都只是推導到$(13)$式就沒有細說了。

後驗分佈近似 #

現在$q(z),q(x|z),p(z|x)$全都是未知的，連形式都還沒確定，而為了實驗，就得把$(13)$式的每一項都明確寫出來。

首先，為了便於取樣，我們假設$z\sim N(0,I)$，即標準的多元正態分佈，這就解決了$q(z)$。那$q(x|z),p(z|x)$呢？一股腦用神經網路擬合吧。

注：本來如果已知$q(x|z)$和$q(z)$，那麼$p(z|x)$最合理的估計應該是：
$$\hat{p}(z|x) = q(z|x) = \frac{q(x|z)q(z)}{q(x)} = \frac{q(x|z)q(z)}{\int q(x|z)q(z)dz}\tag{14}$$
這其實就是EM演算法中的後驗機率估計的步驟，具體可以參考《從最大似然到EM演算法：一致的理解方式》。但事實上，分母的積分幾乎不可能完成，因此這是行不通的。所以乾脆用一般的網路去近似它，這樣不一定能達到最優，但終究是一個可用的近似。

具體來說，我們假設$p(z|x)$也是（各分量獨立的）正態分佈，其均值和方差由$x$來決定，這個“決定”，就是一個神經網路：
$$p(z|x)=\frac{1}{\prod\limits_{k=1}^d \sqrt{2\pi \sigma_{(k)}^2(x)}}\exp\left(-\frac{1}{2}\left\Vert\frac{z-\mu(x)}{\sigma(x)}\right\Vert^2\right)\tag{15}$$
這裡的$\mu(x),\sigma^2(x)$是輸入為$x$、輸出分別為均值和方差的神經網路，其中$\mu(x)$就起到了類似encoder的作用。既然假定了高斯分佈，那麼$(13)$式中的KL散度這一項就可以先算出來：
$$KL\Big(p(z|x)\Big\Vert q(z)\Big)=\frac{1}{2} \sum_{k=1}^d \Big(\mu_{(k)}^2(x) + \sigma_{(k)}^2(x) - \ln \sigma_{(k)}^2(x) - 1\Big)\tag{16}$$
也就是我們所說的KL loss，這在上一篇文章已經給出。

生成模型近似 #

現在只剩生成模型部分$q(x|z)$了，該選什麼分佈呢？論文《Auto-Encoding Variational Bayes》給出了兩種候選方案：伯努利分佈或正態分佈。

什麼？又是正態分佈？是不是太過簡化了？然而並沒有辦法，因為我們要構造一個分佈，而不是任意一個函式，既然是分佈就得滿足歸一化的要求，而要滿足歸一化，又要容易算，我們還真沒多少選擇。

伯努利分佈模型 #

首先來看伯努利分佈，眾所周知它其實就是一個二元分佈：
$$p(\xi)=\left\{\begin{aligned}&\rho,\, \xi = 1;\\
&1-\rho,\, \xi = 0\end{aligned}\right.\tag{17}$$
所以伯努利分佈只適用於$x$是一個多元的二值向量的情況，比如$x$是二值影像時（mnist可以看成是這種情況）。這種情況下，我們用神經網路$\rho(z)$來算引數$\rho$，從而得到
$$q(x|z)=\prod_{k=1}^D \Big(\rho_{(k)}(z)\Big)^{x_{(k)}} \Big(1 - \rho_{(k)}(z)\Big)^{1 - x_{(k)}}\tag{18}$$
這時候可以算出
$$-\ln q(x|z) = \sum_{k=1}^D \Big[- x_{(k)} \ln \rho_{(k)}(z) - (1-x_{(k)}) \ln \Big(1 -\rho_{(k)}(z)\Big)\Big]\tag{19}$$
這表明$\rho(z)$要壓縮到0～1之間（比如用sigmoid啟用），然後用交叉熵作為損失函式，這裡$\rho(z)$就起到了類似decoder的作用。

正態分佈模型 #

然後是正態分佈，這跟$p(z|x)$是一樣的，只不過$x,z$交換了位置：
$$q(x|z)=\frac{1}{\prod\limits_{k=1}^D \sqrt{2\pi \tilde{\sigma}_{(k)}^2(z)}}\exp\left(-\frac{1}{2}\left\Vert\frac{x-\tilde{\mu}(z)}{\tilde{\sigma}(z)}\right\Vert^2\right)\tag{20}$$
這裡的$\tilde{\mu}(z),\tilde{\sigma}^2(z)$是輸入為$z$、輸出分別為均值和方差的神經網路，$\tilde{\mu}(z)$就起到了decoder的作用。於是
$$-\ln q(x|z) = \frac{1}{2}\left\Vert\frac{x-\tilde{\mu}(z)}{\tilde{\sigma}(z)}\right\Vert^2 + \frac{D}{2}\ln 2\pi + \frac{1}{2}\sum_{k=1}^D \ln \tilde{\sigma}_{(k)}^2(z)\tag{21}$$
很多時候我們會固定方差為一個常數$\tilde{\sigma}^2$，這時候
$$-\ln q(x|z) \sim \frac{1}{2\tilde{\sigma}^2}\Big\Vert x-\tilde{\mu}(z)\Big\Vert^2\tag{22}$$
這就出現了MSE損失函式。

所以現在就清楚了，對於二值資料，我們可以對decoder用sigmoid函式啟用，然後用交叉熵作為損失函式，這對應於$q(x|z)$為伯努利分佈；而對於一般資料，我們用MSE作為損失函式，這對應於$q(x|z)$為固定方差的正態分佈。

取樣計算技巧 #

前一節做了那麼多的事情，無非是希望能$(13)$式明確地寫下來。當我們假設$p(z|x)$和$q(z)$都是正態分佈時，$(13)$式的KL散度部分就已經算出來了，結果是$(16)$式；當我們假設$q(x|z)$是伯努利分佈或者高斯分佈時，$-\ln q(x|z)$也能算出來了。現在缺什麼呢？

取樣！

$p(z|x)$的作用分兩部分，一部分是用來算$KL\Big(p(z|x)\Big\Vert q(z)\Big)$，另一部分是用來算$\mathbb{E}_{z\sim p(z|x)}\big[-\ln q(x|z)\big]$的，而$\mathbb{E}_{z\sim p(z|x)}\big[-\ln q(x|z)\big]$就意味著
$$-\frac{1}{n}\sum_{i=1}^n \ln q(x|z_i),\quad z_i \sim p(z|x)\tag{23}$$
我們已經假定了$p(z|x)$是正態分佈，均值和方差由模型來算，這樣一來，藉助“重引數技巧”就可以完成取樣。

但是取樣多少個才適合呢？VAE非常直接了當：一個！所以這時候$(13)$式就變得非常簡單了：
$$\mathcal{L} = \mathbb{E}_{x\sim \tilde{p}(x)} \Bigg[-\ln q(x|z) + KL\Big(p(z|x)\Big\Vert q(z)\Big)\Bigg],\quad z\sim p(z|x)\tag{24}$$
該式中的每一項，可以在把$(16),(19),(21),(22)$式找到。注意對於一個batch中的每個$x$，都需要從$p(z|x)$取樣一個“專屬”於$x$的$z$出來才去算$-\ln q(x|z)$。而正因為VAE在$p(z|x)$這裡只取樣了一個樣本，所以它看起來就跟普通的AE差不多了。

那麼最後的問題就是取樣一個究竟夠了嗎？事實上我們會執行多個epoch，每次的隱變數都是隨機生成的，因此當epoch數足夠多時，事實上是可以保證取樣的充分性的。我也實驗過取樣多個的情形，感覺生成的樣本並沒有明顯變化。

致敬 #

這篇文章從貝葉斯理論的角度出發，對VAE的整體流程做了一個梳理。用這種角度考察的時候，我們心裡需要緊抓住兩個點：“分佈”和“取樣”——寫出分佈形式，並且透過取樣來簡化過程。

簡單來說，由於直接描述複雜分佈是難以做到的，所以我們透過引入隱變數來將它變成條件分佈的疊加。而這時候我們對隱變數的分佈和條件分佈都可以做適當的簡化（比如都假設為正態分佈），並且在條件分佈的引數可以跟深度學習模型結合起來（用深度學習來算隱變數的引數），至此，“深度機率圖模型”就可見一斑了。

讓我們一起致敬貝葉斯大神，以及眾多研究機率圖模型的大牛，他們都是真正的勇者。

轉載到請包括本文地址：https://spaces.ac.cn/archives/5343

更詳細的轉載事宜請參考：《科學空間FAQ》