變分自動編碼機(VAE)理解和實現(Tensorflow)

你需要知道的：

1. 自動編碼機Auto-Encoder (AE)由兩部分encoder和decoder組成，encoder輸入x數據，輸出潛在變量z，decoder輸入z然後輸出一個x’，目的是讓x’與x的分佈儘量一致，當兩者完全一樣時，中間的潛在變量z可以看作是x的一種壓縮狀態，包含了x的全部feature特徵，此時監督信號就是原數據x本身。
2. 變分自動編碼機VAE是自動編碼機的一種擴展，它假設輸出的潛在變量z服從一種先驗分佈，如高斯分佈。這樣，在訓練完模型後，我們可以通過採樣這種先驗分佈得到z’，這個z’可能是訓練過程中沒有出現過的，但是我們依然能在解碼器中通過這個z’獲得x’，從而得到一些符合原數據x分佈的新樣本，具有“生成“新樣本的能力。
3. VAE是一種生成模型，它的目標是要得到p(z|x)$p(z|x)$ 分佈，即給定輸入數據x的分佈，得到潛在變量x的分佈， 與其他的生成模型一樣，它計算的是x和z的聯合概率分佈p(x,z)$p(x,z)$ （如樸素貝葉斯模型通過計算p(x,z)/p(x)$p(x,z)/p(x)$ 得到p(z|x)$p(z|x)$ ），當然它不是直接計算這個聯合概率分佈，而是藉助一些公式變換求解。

從簡單的例子理解VAE/AE的意義：

前面講過，變分自動編碼機的目的是想知道觀測數據x背後的潛在變量z分佈，即p(z|x)$p(z|x)$ ，舉個簡單的例子，比如天氣是我們的觀測數據x，但我們想知道影響天氣變化背後的一些無法觀測的因素z，這個z就像自然法則一樣能夠左右最後觀測到的天氣，這樣我們以後描述某個天氣，就可以完全量化爲對應的潛在變量z。對於這個例子，VAE/AE都能完成這個事情，但如果現在我們想生成一些新的天氣樣本來作爲研究，這個時候只有VAE可以很容易做這個事情：擬合現有樣本分佈的一個潛在變量的先驗分佈，通過採樣這個先驗分佈來獲得新的樣本；而對於AE這個事情就比較難了：由於每個樣本x被固定編碼爲對應的z，我們無法知道潛在樣本的分佈（若此時我們知道了z的分佈，就等於知道了真實數據x的分佈，這顯然是不可能的，相比VAE的解決方案是把真實數據x對應的潛在分佈映射到一個先驗分佈上），若AE硬要獲得新樣本怎麼做呢，此時只能隨機採樣z了，很顯然我們無法驗證：根據這個z是否能正確地還原出一個符合真實樣本x的新樣本。

除了單純“生成“新的樣本用途，生成模型還可以用來去噪聲，比如現在的圖片裏有霧霾，我們想把圖片裏的霧霾去掉，還原沒有霧霾的樣子，就可以用VAE/AE做：把有霧霾的圖片當作輸入x，對應的無霧霾的圖片（假設我們能夠在天氣好的時候獲得）作爲最後要還原的x’訓練VAE模型，如果訓練的足夠好的話，以後再任意拿一張有霧霾的圖片，VAE能夠還原出這個圖片沒有霧霾的樣子，這就是生成模型的優勢。當然，判別模型也能做這個事情：在給定原圖像的情況下，儘量擬合原圖像的變換圖像，但是若測試時出現了之前訓練過程中沒有出現的圖像，效果會不好，因爲判別模型是基於條件概率p(x′|x)$p(x^{\prime }|x)$ ，若新的條件x模型都沒見過，效果肯定不好啊，所以判別模型更注重泛化能力。而生成模型會去擬合x和x’聯合概率分佈p(x,x′)$p(x,x^{\prime })$ ，因此p(x′|x)$p(x^{\prime }|x)$ 的計算只需要除以邊緣概率分佈p(x)$p(x)$ 即可，而對於VAE來說，它擬合的其實是x和潛在變量z的聯合概率分佈p(x,z)$p(x,z)$ ，獲得p(z|x)$p(z|x)$ 從而間接生成x’

VAE推導

爲了求解真實的後驗p(z|x)$p(z|x)$ 概率分佈，VAE引入一個識別模型q(z|x)$q(z|x)$ 去近似p(z|x)$p(z|x)$ ，那麼衡量這兩個分佈之間的差異自然就是相對墒了，也就是KL散度，VAE的目的就是要讓這個相對墒越小，因此推導從相對墒開始：

KL(q(z|x)||p(z|x))=∫q(z|x)logq(z|x)p(z|x)dz=∫q(z|x)(logq(z|x)−logp(z,x)p(x))dz=∫q(z|x)(logq(z|x)−logp(z,x)+logp(x))dz=∫q(z|x)(logq(z|x)−logp(z,x))dz+logp(x)=Ez∼q(z|x)(logq(z|x)−logp(z,x))+logp(x)(1)(2)(3)(4)(5)$$\begin{array}{}\text{(1)}& KL(q(z|x)||p(z|x))& =\int q(z|x)\log \frac{q(z|x)}{p(z|x)}dz\text{(2)}& & =\int q(z|x)\left(\log q(z|x)-\log \frac{p(z,x)}{p(x)}\right)dz\text{(3)}& & =\int q(z|x)\left(\log q(z|x)-\log p(z,x)+\log p(x)\right)dz\text{(4)}& & =\int q(z|x)\left(\log q(z|x)-\log p(z,x)\right)dz+\log p(x)\text{(5)}& & =E_{z\sim q(z|x)}\left(\log q(z|x)-logp(z,x)\right)+\log p(x)\end{array}$$

,

我們把兩個分佈的KL散度展開後得到了兩項，第一項是一個期望，第二個是真實樣本概率的對數logp(x)$\log p(x)$ ，雖然我們不知道它的值是多少，但是我們知道它的值是一個定值。我們將上述結果稍微調換位置得到如下：

L(x)=Ez∼q(z|x)(−logq(z|x)+logp(z,x))=logp(x)−KL(q(z|x)||p(z|x))$$L(x)=E_{z\sim q(z|x)}(-\log q(z|x)+\log p(z,x))=\log p(x)-KL(q(z|x)||p(z|x))$$

令L(x)$L(x)$ 爲上述期望, 它等於一個固定值減去KL散度，由於KL散度值是恆大於0的（當兩個分佈完全一致時，KL散度爲0），因此有L(x)⩽logp(x)$L(x)⩽\log p(x)$ ，此時L(x)$L(x)$ 可以看作是真實概率log值的一個下界，原文叫做變分下界(variational lower bound)。我們目的當然是最優化這個下界，當下界越靠近logp(x)$\log p(x)$ 時，KL散度越小，此時我們q(z|x)$q(z|x)$ 就能夠越準確地估計p(z|x)$p(z|x)$ 。

現在我們繼續研究這個下界L，發現裏面有個聯合概率分佈p(z,x)$p(z,x)$ ，這個東西可不好求，因此繼續把它用貝葉斯公式展開，然後合併成如下樣子：

L(x)=Ez∼q(z|x)(−logq(z|x)+logp(z,x))=Ez∼q(z|x)(−logq(z|x)+logp(x|z)+logp(z))=∫q(z|x)(−logq(z|x)+logp(x|z)+logp(z))dz=−∫q(z|x)(logq(z|x)p(z))dz+∫q(z|x)logp(x|z))dz=−KL(q(z|x)||p(z))+Ez∼q(z|x)(logp(x|z))(6)(7)(8)(9)(10)$$\begin{array}{}\text{(6)}& L(x)& =E_{z\sim q(z|x)}(-\log q(z|x)+\log p(z,x))\text{(7)}& & =E_{z\sim q(z|x)}(-\log q(z|x)+\log p(x|z)+\log p(z))\text{(8)}& & =\int q(z|x)(-\log q(z|x)+\log p(x|z)+\log p(z))dz\text{(9)}& & =-\int q(z|x)\left(\log \frac{q(z|x)}{p(z)}\right)dz+\int q(z|x)\log p(x|z))dz\text{(10)}& & =-KL(q(z|x)||p(z))+E_{z\sim q(z|x)}(\log p(x|z))\end{array}$$

經過變換，我們把這個變分下界L(x)$L(x)$ 用一個期望和KL散度的差表示，我們先看這個期望怎麼求，這個期望表示的是在已獲得的z變量的情況下輸出x的log似然期望，這也可以看作是編碼器的損失函數，因爲我們希望編碼器能通過z儘量的還原出x，也就是儘量使這個對數似然在z服從q(z|x)$q(z|x)$ 分佈情況下最大，那麼這個期望怎麼求呢？最簡單的就是蒙特卡洛採樣了：對於樣本x，用q(z|x)$q(z|x)$ 分佈採樣出L個z，對於每個z算出p(x|z)$p(x|z)$ 概率的log值，然後取平均即爲所求期望，而且當採樣次數L越大，這個均值越接近於真實的期望值：

Ez∼q(z|x)(logp(x|z))≈1L∑l=1Llogp(x|zl),zl∼q(zl|x)$$E_{z\sim q(z|x)}(\log p(x|z))\approx \frac{1}{L}\sum _{l=1}^L\log p(x|z_l),z_l\sim q(z_l|x)$$

但是這種簡單的蒙特卡洛採樣的缺點是估計出來的值方差太大(high variance)，也就是說採樣出的z與z之間相差比較大，導致最後估計值波動性太大，而且這種直接採樣的方法通常是不可求導的，所以不實用。因此，VAE把對z的採樣分成兩部分來求：一部分是固定的值比如標準差σ$\sigma$ 和均值μ$\mu$ ，另一部分是一個隨機的高斯噪聲ϵ$ϵ$ ；具體來說，用一個函數g(x,ϵ)$g(x,ϵ)$ 表示最後採樣出的z值，這個函數由兩部分的和組成：g(x,ϵ)=μx+σx⊙ϵ$g(x,ϵ)={\mu }_x+{\sigma }_x\odot ϵ$ ，其中ϵ∼N(0,1)$ϵ\sim N(0,1)$ ，μx${\mu }_x$ 和σx${\sigma }_x$ 是兩個關於x$x$ 的向量，一般可以理解爲網絡在輸入x樣本後輸出的兩個向量，⊙$\odot$ 表示點乘；這樣，z的採樣由於被固定的μx${\mu }_x$ 和σx${\sigma }_x$ 值決定着其均值和方差，而隨機的部分只由高斯分佈決定，因此減小了方差，而且這種情況下，我們還能計算μx${\mu }_x$ 和σx${\sigma }_x$ 的梯度用於更新，這種trick叫做重參數化(reparameterization trick)，當然上述只是g(x,ϵ)$g(x,ϵ)$ 的一種形式，論文給出了構造g(x,ϵ)$g(x,ϵ)$ 的一般約束。

其實上述的期望換一種角度理解，本質上描述瞭解碼器的性能，z相當於是從編碼器獲得的潛在變量，而解碼器要做的就是儘量讓z能還原出原來的x，也就是儘可能讓logp(x|z)$\log p(x|z)$ 最大化，因此它的損失函數就是p(x|z)$p(x|z)$ 與真實x分佈的交叉熵。

那麼我們回到變分下界L(x)$L(x)$ ，我們已經知道了如何最大化式子中第二項的期望，那麼如何最小化第一項呢？我們知道KL散度是恆大於0的，因此我們只需要最小化KL散度即可，此時變分下界最大。由於KL散度描述着兩個分佈之間的差距，VAE因此讓p(z)$p(z)$ 服從一個先驗的高斯分佈N(0,1)$N(0,1)$ ，便直接可以展開式子計算q(z|x)$q(z|x)$ 與p(z)$p(z)$ 的KL散度，這是因爲q(z|x)$q(z|x)$ 其實就是一種高斯均值爲μx${\mu }_x$ ，方差爲σx${\sigma }_x$ 的高斯分佈（由上述g(x,ϵ)$g(x,ϵ)$ 的求法可得），衡量兩個高斯分佈的差異可以通過它們的密度函數展開推導出來，有興趣的可以嘗試推一下：

−KL(q(z|x)||p(z))=12∑j=1J(1+log((σj)2)−(μj)2−(σj)2)$$-KL(q(z|x)||p(z))=\frac{1}{2}\sum _{j=1}^J(1+\log (({\sigma }_j{)}^2)-({\mu }_j{)}^2-({\sigma }_j{)}^2)$$

這裏σj${\sigma }_j$ 和μj${\mu }_j$ 分別表示向量σx${\sigma }_x$ 和μx${\mu }_x$ 的第j$j$ 個值，這個KL散度本質上描述了編碼器的損失：VAE強制讓輸出的z變量服從先驗的高斯分佈N(0,1)$N(0,1)$ ，因此損失函數即爲當前輸出的z分佈與標準高斯分佈之間的距離，也就是這個KL散度。

最後L(x)$L(x)$ 被寫成：

L(x)=−KL(q(z|x)||p(z))+Ez∼q(z|x)(logp(x|z))=12∑j=1J(1+log((σj)2)−(μj)2−(σj)2)+1L∑l=1Llogp(x|zl)(11)(12)$$\begin{array}{}\text{(11)}& L(x)& =-KL(q(z|x)||p(z))+E_{z\sim q(z|x)}(\log p(x|z))\text{(12)}& & =\frac{1}{2}\sum _{j=1}^J(1+\log (({\sigma }_j{)}^2)-({\mu }_j{)}^2-({\sigma }_j{)}^2)+\frac{1}{L}\sum _{l=1}^L\log p(x|z_l)\end{array}$$

總結：最優化L(x)$L(x)$ 變分下界意味着讓編碼器輸出的z值符合先驗的高斯分佈的情況下，同時也讓解碼器能夠最大可能的用z還原出原來的x，這就是VAE的整個流程，有非常漂亮的理論依據。

VAE的實現(Tensorflow)

這裏主要寫一下實現中比較重要的部分，源碼請參考這個github，使用的mnist手寫體識別的數據集，輸入的是一張張手寫圖片，輸出的是經過潛在變量z還原後的圖片。

編碼器：

def gaussian_MLP_encoder(...):
    # 1st hidden layer
    ...

    # 2nd hidden layer
    ...

    # output layer
    wo = tf.get_variable('wo', [h1.get_shape()[1], n_output * 2], initializer=w_init)
    bo = tf.get_variable('bo', [n_output * 2], initializer=b_init)
    gaussian_params = tf.matmul(h1, wo) + bo

    # The mean parameter is unconstrained
    mean = gaussian_params[:, :n_output]
    # The standard deviation must be positive. Parametrize with a softplus and
    # add a small epsilon for numerical stability
    stddev = 1e-6 + tf.nn.softplus(gaussian_params[:, n_output:])

編碼器的輸出分兩部分，一部分表示mean，一部分表示標準差std，其中由於標準差是恆大於0，因此用了softplus激活函數：

解碼器：

def bernoulli_MLP_decoder(...):
    # 1st hidden layer
    ...

    # 2nd hidden layer
    w1 = tf.get_variable('w1', [h0.get_shape()[1], n_hidden], initializer=w_init)
    b1 = tf.get_variable('b1', [n_hidden], initializer=b_init)
    h1 = tf.matmul(h0, w1) + b1
    h1 = tf.nn.elu(h1)
    h1 = tf.nn.dropout(h1, keep_prob)

    # output layer-mean
    wo = tf.get_variable('wo', [h1.get_shape()[1], n_output], initializer=w_init)
    bo = tf.get_variable('bo', [n_output], initializer=b_init)
    y = tf.sigmoid(tf.matmul(h1, wo) + bo)

輸出的大小與輸入一致，其中每個元素代表着此位置的像素值爲0的概率(或者255，根據輸入來定)，所以用sigmoid激活函數

損失函數

# 編碼器得到標準差和均值向量
mu, sigma = gaussian_MLP_encoder(x_hat, n_hidden, dim_z, keep_prob)

# reparameterization 重參數採樣得到z
z = mu + sigma * tf.random_normal(tf.shape(mu), 0, 1, dtype=tf.float32)

# 解碼器傳入z，輸出y
y = bernoulli_MLP_decoder(z, n_hidden, dim_img, keep_prob)
y = tf.clip_by_value(y, 1e-8, 1 - 1e-8)

# marginal_likelihood loss爲y與輸入數據x之間交叉墒，即解碼器的損失
marginal_likelihood = tf.reduce_sum(x * tf.log(y) + (1 - x) * tf.log(1 - y), 1)
marginal_likelihood = tf.reduce_mean(marginal_likelihood)

# KL_divergence爲z與標準高斯分佈之間的差距，即編碼器的損失
KL_divergence = 0.5 * tf.reduce_sum(tf.square(mu) + tf.square(sigma) - tf.log(1e-8 + tf.square(sigma)) - 1, 1)
KL_divergence = tf.reduce_mean(KL_divergence)

# 變分下界L(x)，目標最大化
ELBO = marginal_likelihood - KL_divergence

# 令損失函數爲-L(x)，目標梯度下降最小化
loss = -ELBO

訓練過程

# 定義更新器，最小化loss
train_op = tf.train.AdamOptimizer(learn_rate).minimize(loss)

with tf.Session() as sess:

    sess.run(tf.global_variables_initializer(), feed_dict={keep_prob : 0.9})

    for epoch in range(n_epochs):

        # Random shuffling
        np.random.shuffle(train_total_data)
        train_data_ = train_total_data[:, :-mnist_data.NUM_LABELS]

        # Loop over all batches
        for i in range(total_batch):
            # Compute the offset of the current minibatch in the data.
            offset = (i * batch_size) % (n_samples)
            batch_xs_input = train_data_[offset:(offset + batch_size), :]

            # 輸出label等於輸入值
            batch_xs_target = batch_xs_input

            # 可以在輸入中加入噪音，讓VAE從帶有噪音的x還原真實的x
            if ADD_NOISE:
                batch_xs_input = batch_xs_input * np.random.randint(2, size=batch_xs_input.shape)
                batch_xs_input += np.random.randint(2, size=batch_xs_input.shape)

            # forward + backword 過程，記錄總的loss，編碼器和解碼器loss
            _, tot_loss, loss_likelihood, loss_divergence = sess.run(
                (train_op, loss, neg_marginal_likelihood, KL_divergence),
                feed_dict={x_hat: batch_xs_input, x: batch_xs_target, keep_prob : 0.9})

結果

輸入數據：

輸出（第0個epoch）：

輸出（第59個epoch）：

更多結果參考github