【深度學習】VAE(Variational Auto-Encoder)原理

下面的內容是從李宏毅2017年機器學習課程中關於VAE一節[1]中整理的。課程講解的非常細緻，再整理一遍方便理解查閱。

---

一、AE與VAE

AE(Auto-Encoder)是一個應用很廣泛的機器學習方法。
主要內容即是：將輸入(Input)經過編碼器(encoder)壓縮爲一個編碼(code)，再通過解碼器(decoder)將編碼(code)解碼爲輸出(Output)。
學習的目標即：要使得輸出(Output)與輸入(Input)越接近越好。
以輸入爲圖像爲例，結構圖如下：

AE中間階段生成的編碼向量，並不是隨機、沒有意義的。編碼中攜帶着與輸入有關的信息，編碼中的某些維度代表着輸入數據的某些特徵。例如生成人臉圖像時，編碼可以表示人臉表情、頭髮樣子、是否有鬍子等等。
VAE變分自動編碼器作爲AE的變體，它主要的變動是對編碼(code)的生成上。編碼(code)不再像AE中是唯一映射的，而是具有某種分佈，使得編碼(code)在某範圍內波動時都可產生對應輸出。藉助下面這個例子進行理解：

如上圖AE示意圖，左側是對滿月圖像編解碼，右側是對弦月圖像編解碼，而像中間的編碼對解碼器來說並不知道要生成何種圖像。在VAE示意圖中，左右兩側對圖像編解碼過程中，編碼有不同程度的擾動(即圖中noise)，解碼器利用擾動範圍內的編碼同樣可以生成相應的圖像，而對交界處的編碼，編碼器既想生成滿月圖像，又想生成弦月圖像，爲此做出折中，生成位於兩者之間的圖像。
這就是VAE一個較爲直觀的想法。

---

二、VAE原理

VAE是一個深度生成模型，其最終目的是生成出概率分佈$P(x)$，$x$即輸入數據。
在VAE中，我們通過高斯混合模型(Gaussian Mixture Model)來生成$P(x)$，也就是說$P(x)$是由一系列高斯分佈疊加而成的，每一個高斯分佈都有它自己的參數$\mu$和$\sigma$。

那我們藉助一個變量$z\sim N(0,I)$(注意$z$是一個向量，生成自一個高斯分佈)，找一個映射關係，將向量$z$映射成這一系列高斯分佈的參數向量$\mu (z)$和$\sigma (z)$。有了這一系列高斯分佈的參數我們就可以得到疊加後的$P(x)$的形式，即$x|z \sim N \big(\mu(z), \sigma(z)\big)$。(這裏的“形式”僅是對某一個向量$z$所得到的)。
那麼要找的這個映射關係怎麼獲得呢？就拿神經網絡來做唄，只要神經元足夠想要啥樣的函數得不到呢。如下圖形式：

輸入向量$z$，得到參數向量$\mu (z)$和$\sigma (z)$。這個映射關係是要在訓練過程中更新NN權重得到的。這部分作用相當於最終的解碼器(decoder)。

對於某一個向量$z$我們知道了如何找到$P(x)$。那麼對連續變量$z$依據全概率公式有：
$P(x)=\int _{z} P(z)P(x|z)dz$但是很難直接計算積分部分。我們用極大似然估計來估計$P(x)$，有似然函數$L$：$L=\sum_{x}\log P(x)$這裏我們額外引入一個分佈$q(z|x)$，$z|x \sim N\big(\mu^\prime(x), \sigma^\prime(x)\big)$。這個分佈表示形式如下：
這個分佈同樣是用一個神經網絡來完成，向量$z$根據NN輸出的參數向量$\mu '(x)$和$\sigma '(x)$運算得到，注意這三個向量具有相同的維度。這部分作用相當於最終的編碼器(encoder)。
之後就開始推導了。
$\begin{aligned} \log P(x)&=\int_{z}q(z|x)\log P(x)dz \qquad \because \int_{z}q(z|x)dz=1 \\ &=\int_{z} q(z|x)\log \frac{P(z, x)}{P(z|x)}dz \\ &=\int_z q(z|x)\log \big(\frac{P(z,x)}{q(z|x)} \cdot \frac{q(z|x)}{P(z|x)}\big)dz \\ &=\int_z q(z|x)\log \frac{q(z|x)}{P(z|x)}dz + \int_z q(z|x)\log \frac{P(z,x)}{q(z|x)}dz \\ &=D_{KL}\big(q(z|x)||P(z|x)\big) + \int_z q(z|x)\log \frac{P(z,x)}{q(z|x)}dz \\ &\eqslantgtr \int_z q(z|x)\log \frac{P(z,x)}{q(z|x)}dz \qquad \because D_{KL}(q||P) \eqslantgtr 0 \end{aligned}$
我們將$\int_z q(z|x)\log \frac{P(z,x)}{q(z|x)}dz$稱爲$\log P(x)$的$\textit{\textbf{(variational) lower bound}}$ (下界)，簡稱爲$L_b$。最大化$L_b$就等價於最大化似然函數$L$。那麼接下來具體看$L_b$，
$\begin{aligned} L_b&=\int_z q(z|x)\log \frac{P(z,x)}{q(z|x)}dz \\ &=\int_z q(z|x)\log \big(\frac{P(z)}{q(z|x)} \cdot P(x|z) \big)dz \\ &=\int_z q(z|x)\log \frac{P(z)}{q(z|x)}dz + \int_z q(z|x)\log P(x|z)dz \\ &=-D_{KL}\big( q(z|x)||P(z)\big) + \int_z q(z|x)\log P(x|z)dz \\ &=-D_{KL}\big( q(z|x)||P(z)\big) + E_{q(z|x)}[\log P(x|z)] \end{aligned}$
最大化$L_b$包括下面兩部分：

• $\textit{minimizing}$ $D_{KL}\big( q(z|x)||P(z)\big)$，使後驗分佈近似值$q(z|x)$接近先驗分佈$P(z)$。也就是說通過$q(z|x)$生成的編碼$z$不能太離譜，要與某個分佈相當才行，這裏是對中間編碼生成起了限制作用。
  當$q(z|x)$和$P(z)$都是高斯分佈時，推導式有([2]中Appendix B)：$D_{KL}\big( q(z|x)||P(z)\big)=-\frac{1}{2}\sum_{j}^{J}\big( 1+\log (\sigma_{j})^2 - (\mu_j)^2-(\sigma_j)^2\big)$其中$J$表示向量$z$的總維度數，$\sigma_j$和$\mu_j$表示$q(z|x)$輸出的參數向量$\sigma$和$\mu$的第$j$個元素。(這裏的$\sigma$和$\mu$等於前文中$\mu '(x)$和$\sigma '(x)$)
• $\textit{maximizing}$ $E_{q(z|x)}[\log P(x|z)]$，即在給定編碼器輸出$q(z|x)$下解碼器輸出$P(x|z)$越大越好。這部分也就相當於常規的Reconstruction Loss(重建損失)。

由此我們可以得出VAE的原理圖：

通常忽略掉decoder輸出的$\sigma(x)$一項，僅要求$\mu(x)$與$x$越接近越好。
VAE的損失函數即：
$\min Loss_{VAE} = D_{KL}\big( q(z|x)||P(z)\big) -E_{q(z|x)}[\log P(x|z)]$
附：
極大似然估計$P(x)$的時候還有一種寫法，即通過$P(x)=\int_z P(x,z)dz$來推導。如圖[3]：

裏邊有提到術語ELBO，Evidence Lower Bound(證據下界)，有興趣的可以自行查閱瞭解（偷個懶）。

---

三、reparameterization trick

由上文中VAE原理圖可以看出，編碼$z$是由分佈$q(z|x)$採樣產生，而採樣操作是不可微分的，因此反向傳播做不了。[2]中提到了reparameterization trick來解決，藉助[4]中的示意圖理解下：

將上圖左圖原來的採樣操作通過reparameterization trick變換爲右圖的形式。引入一個外部向量$\epsilon \sim N(0, 1)$，通過$z=\mu + \sigma \odot \epsilon$計算編碼$z$（$\odot$表示element-wise乘法，$\epsilon$的每一維都服從標準高斯分佈），由此loss的梯度可以通過$\mu$和$\sigma$分支傳遞到encoder model處（$\epsilon$並不需要梯度信息來更新）。
最終的VAE實際形式如下圖所示：

---

四、不足

VAE在產生新數據的時候是基於已有數據來做的，或者說是對已有數據進行某種組合而得到新數據的，它並不能生成或創造出新數據。另一方面是VAE產生的圖像比較模糊。
而大名鼎鼎的GAN利用對抗學習的方式，既能生成新數據，也能產生較清晰的圖像。後續的更是出現了很多種變形。

---

五、參考文獻

[1] Unsupervised Learning: Deep Generative Model (2017/04/27)
[2] VAE原著Auto-encoding variational bayes
[3] VAE的三種不同推導方法
[4] https://www.jeremyjordan.me/variational-autoencoders/