Auto-Encoding Variational Bayes， PGM(概率圖模型)

1. PGM 概率圖模型

AVEB（Auto-Encoding Variational Bayes）的思想可以從概率圖模型來理解。它是建立在概率圖模型Inference中的Variational Inference基礎上的。

1.1 Variational Inference

概率圖的inference，是一種從概率圖模型中獲取信息的方式（通俗的來說，就是向模型提問）。inference一般由兩種常見的形式：

• Marginal inference : 用來得到整個系統狀態的某一部分（從其他狀態變量中marginalize之後）的概率。（以圖像處理爲例，如果扣掉圖像的一塊，如何使用其他的圖像信息補全這一塊缺失的內容）。
• MAP inference ： 最大後驗概率，從模型中得到結果。（以圖像爲例，輸入一張圖像，通過我們的概率模型，得到它是貓的概率和它是狗的概率）

Variational Inference是把inference問題，考慮成一個優化問題 ：
我們有一組候選的概率分佈Q，真實的概率分佈記爲p（p是未知的，但是我們有服從它的觀測結果）。我們的目標則是在Q中找到一個元素q，使得q最接近於p。最後，我們使用p（完全已知的一個分佈）來回答上面的inference問題。

這樣的問題就涉及到了兩個方面：

• 一個是如何判斷p與q的相似程度。
• 如果找到這個最相似的q。

比較常用的方法是使用KL divergence衡量相似程度，並且使用gradient descent的方法找到最佳的q。

1.2 KL divergence

定義：
$KL(q||p) = \sum_{x} q(x)\log \frac{q(x)}{p(x)}$

重要性質：
$KL(q||p) \ge 0$

$KL(q||p) = 0 ,\ if \ p =q$

$KL(q||p) \ne KL(p||q) ,\ if \ p \ne q$

在這裏，p和q都是歸一化的概率分佈。但是實際上我們還是不能衡量KL，因爲p對我們來說是未知的（intractable）。我們考慮一個更加一般的p分佈， 它可能不是歸一化的，所以p的表達式爲：

$p(x_{1}, x_{2},...,x_{n};\theta) = \frac{\bar p(x_{1}, x_{2},...,x_{n};\theta) }{Z(\theta)} = \frac{1}{Z(\theta)} \prod_{k} \phi_{k}(x_{k};\theta)$

$\phi_{k}$是一個factor（因子）。看這個表達式我們也可以更好的理解爲什麼p是intractable。這是由於$Z(\theta)$是intractable，爲了求這個歸一化常數，我們需要對所有可能的狀態x進行計算，但是鑑於p其實是未知的分佈，我們不可能分析它對於所有x狀態的可能結果（當然我們可以用monte carto的方法sampling近似）。這裏也其實表達了對於一般的分佈，我們無法用直接KL衡量它與另一個分佈q的相似程度。

1.3 The Variational Lower bound

我們先分析下面的等式（直接把新的$\bar p$帶入KL的表達式）：
$J(q) = \sum_{x} q(x)\log \frac{q(x)}{\bar p(x)}$

$J(q) = \sum_{x} q(x)\log \frac{q(x)}{p(x)} - \log Z(\theta) = KL(q||p) - \log Z(\theta)$

我們可以得到下面的等式不等式：

$\log Z(\theta) = KL(q||p) - J(q) \ge - J(q)$

這表示了，$- J(q)$ 其實是$\log Z(\theta)$的一個下界（lower bound）。而$KL(q||p)$恰好是下界與真實值之間的gap。如果我們嘗試最大化這個下界$- J(q)$ ，那麼會促使$- J(q)$和$\log Z(\theta)$變得越來越接近，也就是說他們之間的gap越來越小，也就是達到了最小化$KL(q||p)$的目的。

$- J(q)$被稱作 the variational lower bound 或者 the evidence lower bound (ELBO)。

$-J(q) = -\sum_{x} q(x)\log q(x) + \sum_{x} q(x)\log \bar p(x)$

$-J(q) = E_{q(x)}[\log \bar p(x) - \log q(x)]$

通常上面的結論也被寫爲：

$\log Z(\theta) \ge \mathbb{E}_{q(x)}[\log \bar p(x) - \log q(x)] = \mathcal{L}(p,q)$
對這個表達式求相對於q分佈的期望值。

2. Black-box variational inference

2.1 gradient descent

這是一個用來最大化ELBO的算法，它的核心思想很簡單，它不在乎p和q的分佈具體是什麼，而單純的只考慮他們的梯度，然後使用梯度下降（gradient descent）來優化目標函數。

• q分佈的參數爲$\phi$，p分佈的參數爲$\theta$。所以我們需要分別對ELBO對$\phi$和$\theta$進行求導。
• 也就引出的它的唯一一個前提條件（或者說假設）：q是相對於$\phi$可微的（differentiable）。

我們需要計算的梯度爲：
$\nabla_{\theta, \phi} \mathbb{E}_{q_{\phi}(z)}[\log p_{\theta}(z) - \log q_{\phi}(z) ]$
注意這裏的標記和前面的不一樣了，在這裏z是我們在這個inference中想要求的量，x是它的輸入值，p可以理解爲我們的觀測數據，q是我們預測模型。

2.2 Score function gradient estimator

我們把上面的梯度公式分開$\phi$和$\theta$兩部分來求：

$\theta$部分：
$\nabla_{\theta} \mathbb{E}_{q_{\phi}(z)}[\log p_{\theta}(z) - \log q_{\phi}(z) ] = \nabla_{\theta} \mathbb{E}_{q_{\phi}(z)}[\log p_{\theta}(z)] = \mathbb{E}_{q_{\phi}(z)}[\nabla_{\theta} \log p_{\theta}(z)]$

這樣我們可以通過Monte Carlo來對$\Delta_{\theta} \log p_{\theta}(z)$進行sample，得到它的結果。

$\phi$部分：
$\nabla_{\phi} \mathbb{E}_{q_{\phi}(z)}[\log p_{\theta}(z) - \log q_{\phi}(z) ]$

這一個部分則更爲複雜，因爲它涉及到的期望的分佈也是和$\phi$相關的。我們可以通過score function estimator求（可以在這篇文章的 附錄B中得到下面公式的證明 paper）：

$\nabla_{\phi} \mathbb{E}_{q_{\phi}(z)}[\log p_{\theta}(z) - \log q_{\phi}(z) ] = \mathbb{E}_{q_{\phi}(z)} [(\log p_{\theta}(z) - \log q_{\phi}(z)) \nabla_{\phi}\log q_{\phi}(z)]$

然後我們就可以通過Monte Carlo來sample期望的結果了。
但是這個算法存在一個很大的問題 ：high variance。

3. Our problem - Latent variable model

在繼續講解算法之前，我們先把我們的問題整理一下，並且把前面的公式，使用符合問題的記號重新在這裏表達：

• 我們假設我們在這裏的問題是生成隨機的人臉圖片。
• 準備階段時候，我們有一系列的人臉圖片數據集{x}。
• 我們的目的是，使用算法在這些圖片中找到一定的規律，encode這些圖片爲一系列的參數（特徵）z：
  $P(z|x)$
• 我們隨機給出一組參數（特徵）z，可以通過算法生成一張新的人臉的圖像x：
  $P(x|z)$

我們在這裏的p分佈爲：

$p_{\theta}(z | x) = \frac{ p_{\theta}(x, z) }{p_{\theta}(x)}$

$J(q) = \sum_{(z|x)} q_{\phi}(z | x)\log \frac{q_{\phi}(z | x)}{p_{\theta}(z, x)} = \sum_{(z|x)} q_{\phi}(z | x)\log \frac{q_{\phi}(z | x)}{p_{\theta}(z | x)} - \log p_{\theta}(x) = -\mathcal{L}(p,q)$

我們可以得到下面的等式不等式：

$\log p_{\theta}(x) = KL(q_{\phi}(z | x)||p_{\theta}(z | x)) + \mathcal{L}(p,q)$

4. Auto-encoding reformulation

這裏對ELBO的公式進行了簡單的reform:

$\log p_{\theta}(x) = KL(q_{\phi}(z | x)||p_{\theta}(z | x)) - \sum_{(z|x)} q_{\phi}(z | x)\log \frac{q_{\phi}(z | x)}{p_{\theta}(z, x)}$

$= \sum_{(z|x)} q_{\phi}(z | x)\log [p_{\theta}(x |z)\frac{q_{\phi}(z | x)}{p_{\theta}(z | x) p_{\theta}(x |z)}] - \sum_{(z|x)} q_{\phi}(z | x)\log \frac{q_{\phi}(z | x)}{p_{\theta}(z)} \frac{p_{\theta}(z)}{p_{\theta}(z, x)}$

$= \mathbb{E}_{q_{\phi}(z | x)}[\log [p_{\theta}(x |z)] - KL(q_{\phi}(z | x) || p_{\theta}(z)) + \sum_{(z|x)} q_{\phi}(z | x) \log \frac{q_{\phi}(z | x)}{p_{\theta}(z | x) p_{\theta}(x |z)}\frac{p_{\theta}(z, x)}{p_{\theta}(z)}$

$= \mathbb{E}_{q_{\phi}(z | x)}[\log [p_{\theta}(x |z)] - KL(q_{\phi}(z | x) || p_{\theta}(z)) + KL(q_{\phi}(z | x) || p_{\theta}(z | x))$

（因爲上面最右邊一個很長的表達式也是一個KL，而且我們知道KLdivergence是非負的）

$\log p_{\theta}(x) \ge \mathbb{E}_{q_{\phi}(z | x)}[\log [p_{\theta}(x |z)] - KL(q_{\phi}(z | x) || p_{\theta}(z))$

• 上面的右邊是在z~$q(z|x)上sample的期望和divergence。相當於前面我們提到的，提供一個x（圖像），把它encode成一個特徵表達z，所以在這裏q就是相當於是一個encoder。
• $\mathbb{E}_{q_{\phi}(z | x)}[\log [p_{\theta}(x |z)]$這一項則是一個maximal log-likelihood函數，它表達瞭如果給到我們的提取的z特徵，能不能最大程度恢復原本的圖像x。p就相當於一個decoder。
• $KL(q_{\phi}(z | x) || p_{\theta}(z))$這一項是一個的divergence。用來限制$q_{\phi}(z | x)$，以保證它不會每次都生成一樣的特徵$q_{\phi}(z)$。

5. The Reparameterization Trick

這裏考慮對分佈重新進行參數化，以簡化求梯度的難度。
原本的分佈爲：

$q_{\theta}(z | x)$

假設誤差是正態分佈：

$\epsilon \sim p(\epsilon)$

假設z的分佈是依賴於x的函數，同時包含一定的誤差$\epsilon$:

$z = g_{\phi}(\epsilon, x)$

$z = g_{\phi, \sigma}(\epsilon) = \mu + \epsilon \cdot \sigma$

所以對任意一個函數的期望的梯度可以表示爲：

$\nabla_{\phi} \mathbb{E} _{z \sim q_{\phi}(z|x)} [f(x,z)] = \nabla_{\phi} \mathbb{E} _{\epsilon \sim p(\epsilon )} [f(x,g_{\phi}(\epsilon, x))] = \mathbb{E} _{\epsilon \sim p(\epsilon )} [ \nabla_{\phi} f(x,g_{\phi}(\epsilon, x))]$

用這樣的重新參數化，可以更加輕鬆的解決前面的梯度問題。而且在這種方法下沒有那麼多的variance （這個在文章中有更加細節的討論）

6. Neural Network pick q and p

接下來就是使用神經網絡來選擇適當的decoder（p）和encoder（q）的問題了。

7. Loss function in Neural Network

7.1 reformulation of ELBO

從上面的推導，我們得到：

$\log p_{\theta}(x) = \mathbb{E}_{q_{\phi}(z | x)}[\log [p_{\theta}(x |z)] - KL(q_{\phi}(z | x) || p_{\theta}(z)) + KL(q_{\phi}(z | x) || p_{\theta}(z | x))$

可以寫爲：

$\log p_{\theta}(x) - \mathbb{E}_{q_{\phi}(z | x)}[\log [p_{\theta}(x |z)] + KL(q_{\phi}(z | x) || p_{\theta}(z)) = KL(q_{\phi}(z | x) || p_{\theta}(z | x))$

所以其實我們的目標函數其實可以看作爲最小化：

$\log p_{\theta}(x) - \mathbb{E}_{q_{\phi}(z | x)}[\log [p_{\theta}(x |z)] + KL(q_{\phi}(z | x) || p_{\theta}(z))$

• 第一項是最小化從x通過decode p恢復的圖像和原本的輸入圖像的差異。
• 第二項則是前面解釋過的regulatization項。

7.2 Neural Network distribution

使用前面提到的reparameterization之後，我們的分佈則變爲：

$p(x|z) = \mathcal{N} (x; \vec\mu(z), diag(\vec\sigma(z))^{2})$

$p(z) = \mathcal{N}(z;0,I)$

$q(z|x) = \mathcal{N}(z;\vec\mu(x), diag(\vec\sigma(x))^{2})$

7.3 Loss function

帶入到我們的reformulation ELBO表達式中：

$\log p_{\theta}(x) - \mathbb{E}_{q_{\phi}(z | x)}[\log [p_{\theta}(x |z)] + KL(\mathcal{N}(z;\vec\mu(x), diag(\vec\sigma(x))^{2}) || \mathcal{N}(z;0,I))$

下圖來自Medium