從Autoencoder to Beta-VAE

jasonzhangxianrong發表於2024-07-15

[2019-07-18更新:增加VQ-VAE和VQ-VAE-2的部分。
[更新於2019-07-26:增加有關TD-VAE的部分。

Autocoder 的發明是為了使用中間有一個狹窄瓶頸層的神經網路模型來重建高維資料(哎呀,對於變分自動編碼器來說可能不是這樣,我們將在後面的章節中詳細研究它)。一個很好的副產品是降維:瓶頸層捕獲壓縮的潛在編碼。這種低維表示可以用作各種應用程式(即搜尋)中的嵌入向量,幫助資料壓縮或揭示底層資料生成因素。

表示法

象徵意味 著
$\mathcal{D}$ 資料集 $\mathcal{D} = \{ \mathbf{x}^{(1)}, \mathbf{x}^{(2)}, \dots, \mathbf{x}^{(n)} \}$ 包含 $n$ 資料樣本;$\vert\mathcal{D}\vert =n $。
$\mathbf{x}^{(i)}$ 每個資料點都是 $d$ 維的向量,$\mathbf{x}^{(i)} = [x^{(i)}_1, x^{(i)}_2, \dots, x^{(i)}_d]$。
$\mathbf{x}$ 資料集中的一個資料樣本 $\mathbf{x} \in \mathcal{D}$。
$\mathbf{x}'$ $\mathbf{x}$ 的重構版本。
$\波浪號{\mathbf{x}}$ $\mathbf{x}$ 的損壞版本。
$\mathbf{z}$ 在瓶頸層中學習的壓縮程式碼。
$a_j^{(l)}$ 第 $l 個隱藏層中第 $j 個神經元的啟用函式。
$g_{\phi}(.)$ 由 $\phi$ 引數化的編碼函式。
$f_{\theta}(.)$ 由 $\theta$ 引數化的解碼函式。
$q_{\phi}(\mathbf{z}\vert\mathbf{x})$ 估計後驗機率函式,也稱為機率編碼器
$p_{\theta}(\mathbf{x}\vert\mathbf{z})$ 給定潛在程式碼(也稱為機率解碼器)生成真實資料樣本的可能性。

自動編碼器

Autoencoder 是一種神經網路,旨在以無監督的方式學習身份函式,以重建原始輸入,同時壓縮過程中的資料,從而發現更有效和壓縮的表示。這個想法起源於1980年代,後來由Hinton和Salakhutdinov在2006年發表的開創性論文推廣。

它由兩個網路組成:

  • 編碼器網路:將原有的高維輸入轉換為潛伏的低維碼。輸入大小大於輸出大小。
  • 解碼器網路:解碼器網路從程式碼中恢復資料,可能具有越來越大的輸出層。

從Autoencoder to Beta-VAE圖 1.自動編碼器模型體系結構的圖示。

編碼器網路基本上實現了降維,就像我們使用主成分分析(PCA)或矩陣分解(MF)一樣。此外,自動編碼器針對程式碼中的資料重建進行了顯式最佳化。一個好的中間表示不僅可以捕獲潛在變數,而且有利於完整的減壓過程。

該模型包含編碼器函式 $g(.)$ 引數化為 $\phi$ 和解碼器函式 $f(.)$ 引數化為 $\theta$。在瓶頸層中,輸入 $\mathbf{x}$ 學習的低維碼為 $\mathbf{z} = g_\phi(\mathbf{x})$,重構的輸入為 $\mathbf{x}' = f_\theta(g_\phi(\mathbf{x}))$。

將引數 $(\theta, \phi)$ 一起學習,以輸出與原始輸入 $\mathbf{x} \approx f_\theta(g_\phi(\mathbf{x}))$ 相同的重構資料樣本,或者換句話說,用於學習恆等函式。有各種指標可以量化兩個向量之間的差異,例如啟用函式為 S 形時的交叉熵,或者像 MSE 損失一樣簡單:

$$ L_\text{AE}(\theta, \phi) = \frac{1}{n}\sum_{i=1}^n (\mathbf{x}^{(i)} - f_\theta(g_\phi(\mathbf{x}^{(i)})))^2 $$

去噪自動編碼器

由於自編碼器學習恆等函式,因此當網路引數多於資料點數時,我們面臨著“過擬合”的風險。

為了避免過擬合和提高魯棒性,Denoising Autoencoder (Vincent et al. 2008) 提出了對基本自編碼器的修改。透過以隨機方式向輸入向量新增噪聲或遮蔽輸入向量的某些值,輸入部分損壞,$\tilde{\mathbf{x}} \sim \mathcal{M}_\mathcal{D}(\tilde{\mathbf{x}} \vert \mathbf{x})$。然後對模型進行訓練以恢復原始輸入(注意:不是損壞的輸入)。

$$ \begin{aligned} \tilde{\mathbf{x}}^{(i)} &\sim \mathcal{M}_\mathcal{D}(\tilde{\mathbf{x}}^{(i)} \vert \mathbf{x}^{(i)})\\ L_\text{DAE}(\theta, \phi) &= \frac{1}{n} \sum_{i=1}^n (\mathbf{x}^{(i)} - f_\theta(g_\phi(\tilde{\mathbf{x}}^{(i)})))^2 \end{對齊} $$

其中 $\mathcal{M}_\mathcal{D}$ 定義了從真實資料樣本到噪聲或損壞資料樣本的對映。

從Autoencoder to Beta-VAE圖 2.去噪自動編碼器模型架構的圖示。

這種設計的動機是,即使視野被部分遮擋或損壞,人類也可以輕鬆識別物體或場景。為了“修復”部分被破壞的輸入,去噪自動編碼器必須發現並捕獲輸入維度之間的關係,以便推斷缺失的部分。

對於具有高冗餘的高維輸入(如影像),模型可能依賴於從多個輸入維度的組合中收集的證據來恢復去噪版本,而不是過度擬合一個維度。這為學習魯棒潛在表示奠定了良好的基礎。

噪聲由隨機對映 $\mathcal{M}_\mathcal{D}(\tilde{\mathbf{x}} \vert \mathbf{x})$ 控制,並且它不特定於特定型別的損壞過程(即掩蔽噪聲、高斯噪聲、鹽和胡椒噪聲等)。當然,腐敗過程可以配備先驗知識

在原始 DAE 論文的實驗中,噪聲是這樣應用的:隨機選擇固定比例的輸入維度,並將其值強制為 0。聽起來很像輟學,對吧?好吧,去噪自動編碼器是在 2008 年提出的,比輟學論文早 4 年(Hinton 等人,2012) ;)

稀疏自動編碼器

稀疏自動編碼器對隱藏單元啟用應用“稀疏”約束,以避免過擬合併提高魯棒性。它強制模型同時啟用少量隱藏單元,或者換句話說,一個隱藏的神經元應該在大部分時間被啟用。

回想一下,常見的啟用函式包括 sigmoid、tanh、relu、leaky relu 等。當值接近 1 時,神經元被啟用,當值接近 0 時,神經元失活。

假設第 $l 個隱藏層中有 $s_l$ 個神經元,並且該層中第 $j 個神經元的啟用函式標記為 $a^{(l)}_j(.)$,$j=1,\dots,s_l$。這個神經元 $\hat{\rho}_j$ 的啟用分數預計是一個小數 $\rho$,稱為稀疏引數;通用配置是 $\rho = 0.05$。

$$ \hat{\rho}_j^{(l)} = \frac{1}{n} \sum_{i=1}^n [a_j^{(l)}(\mathbf{x}^{(i)})] \approx \rho $$

這種約束是透過在損失函式中新增一個懲罰項來實現的。KL 散度 $D_\text{KL}$ 測量兩個伯努利分佈之間的差異,一個具有均值 $\rho$,另一個具有均值 $\hat{\rho}_j^{(l)}$。超引數 $\beta$ 控制我們希望對稀疏性損失應用的懲罰強度。

$$ \begin{aligned} L_\text{SAE}(\theta) &= L(\theta) + \beta \sum_{l=1}^L \sum_{j=1}^{s_l} D_\text{KL}(\rho \| \hat{\rho}_j^{(l)}) \\ &= L(\theta) + \beta \sum_{l=1}^L \sum_{j=1}^{s_l} \rho\log\frac{\rho}{\hat{\rho}_j^{(l)}} + (1-\rho)\log\frac{1-\rho}{1-\hat{\rho}_j^{(l)}} \end{對齊} $$

從Autoencoder to Beta-VAE圖 4.均值為 \(\rho=0.25\) 的伯努利分佈與均值為 \(0 \leq \hat{\rho} \leq 1\) 的伯努利分佈之間的 KL 散度。

$k$-稀疏自動編碼器

在 $k$-Sparse Autoencoder (Makhzani and Frey, 2013) 中,稀疏性是透過僅將前 k 個最高啟用保留在具有線性啟用函式的瓶頸層中來強制執行的。 首先,我們透過編碼器網路執行前饋以獲得壓縮程式碼:$\mathbf{z} = g(\mathbf{x})$。 對程式碼向量 $\mathbf{z}$ 中的值進行排序。僅保留 k 個最大值,而其他神經元設定為 0。這也可以在具有可調節閾值的 ReLU 層中完成。現在我們有一個稀疏化程式碼:$\mathbf{z}' = \text{Sparsify}(\mathbf{z})$。 計算稀疏化程式碼的輸出和損失,$L = |\mathbf{x} - f(\mathbf{z}') |_2^2$。 而且,反向傳播只透過頂部 k 啟用的隱藏單位!

從Autoencoder to Beta-VAE圖 5.不同稀疏級別 k 的 k 稀疏自編碼器的濾波器,從 MNIST 學習,具有 1000 個隱藏單元。(圖片來源:Makhzani 和 Frey,2013)

收縮式自動編碼器

與稀疏自動編碼器類似,收縮自動編碼器Rifai等人,2011)鼓勵學習的表示保持在收縮空間中,以獲得更好的魯棒性。

它在損失函式中新增了一個項,以懲罰對輸入過於敏感的表示,從而提高對訓練資料點周圍小擾動的魯棒性。靈敏度是透過編碼器啟用的雅可比矩陣相對於輸入的 Frobenius 範數來測量的:

$$ \|J_f(\mathbf{x})\|_F^2 = \sum_{ij} \Big( \frac{\partial h_j(\mathbf{x})}{\partial x_i} \Big)^2 $$

其中 $h_j$ 是壓縮程式碼中的一個單位輸出 $\mathbf{z} = f(x)$。

此懲罰項是學習編碼的所有偏導數相對於輸入維度的平方和。作者聲稱,根據經驗,發現這種懲罰雕刻出對應於低維非線性流形的表示,同時對與流形正交的大多數方向保持更不變。

VAE: 變分自動編碼器

變分自動編碼器Kingma & Welling,2014)是VAE的縮寫,實際上與上述所有自動編碼器模型不太相似,但深深植根於變分貝葉斯和圖形模型的方法。

我們不想將輸入對映到固定向量中,而是將其對映到分佈中。讓我們將此分佈標記為 $p_\theta$,引數化為 $\theta$。資料輸入 $\mathbf{x}$ 和潛在編碼向量 $\mathbf{z}$ 之間的關係可以完全由下式定義:

  • 先驗 $p_\theta(\mathbf{z})$
  • 似然 $p_\theta(\mathbf{x}\vert\mathbf{z})$
  • 後驗 $p_\theta(\mathbf{z}\vert\mathbf{x})$

假設我們知道這個分佈的實引數 $\theta^{*}$。為了生成一個看起來像真實資料點 $\mathbf{x}^{(i)}$ 的樣本,我們按照以下步驟操作:

  1. 首先,從先驗分佈 $p_{\theta^*}(\mathbf{z})$ 中取樣 $\mathbf{z}^{(i)}$。
  2. 然後從條件分佈生成值 $\mathbf{x}^{(i)}$ $p_{\theta^*}(\mathbf{x} \vert \mathbf{z} = \mathbf{z}^{(i)})$。

最優引數 $\theta^{*}$ 是使生成真實資料樣本的機率最大化的引數:

$$ \theta^{*} = \arg\max_\theta \prod_{i=1}^n p_\theta(\mathbf{x}^{(i)}) $$

通常,我們使用對數機率將 RHS 上的乘積轉換為總和:

$$ \theta^{*} = \arg\max_\theta \sum_{i=1}^n \log p_\theta(\mathbf{x}^{(i)}) $$

現在讓我們更新等式以更好地演示資料生成過程,以便涉及編碼向量:

$$ p_\theta(\mathbf{x}^{(i)}) = \int p_\theta(\mathbf{x}^{(i)}\vert\mathbf{z}) p_\theta(\mathbf{z}) d\mathbf{z} $$

不幸的是,以這種方式計算 $p_\theta(\mathbf{x}^{(i)})$ 並不容易,因為檢查 $\mathbf{z}$ 的所有可能值並將它們相加非常昂貴。為了縮小值空間以方便更快的搜尋,我們想引入一個新的近似函式來輸出給定輸入 $\mathbf{x}$, $q_\phi(\mathbf{z}\vert\mathbf{x})$ 的可能程式碼,引數化為 $\phi$。

從Autoencoder to Beta-VAE圖 6.變分自動編碼器中涉及的圖形模型。實線表示生成分佈 \(p\_\theta(.)\) 和虛線表示分佈 \(q\_\phi (\mathbf{z}\vert\mathbf{x})\) 近似於難處理的後驗 \(p\_\theta (\mathbf{z}\vert\mathbf{x})\)

現在,該結構看起來很像自動編碼器:

  • 條件機率 $p_\theta(\mathbf{x} \vert \mathbf{z})$ 定義了一個生成模型,類似於上面介紹的解碼器 $f_\theta(\mathbf{x} \vert \mathbf{z})$。$p_\theta(\mathbf{x} \vert \mathbf{z})$ 也稱為機率解碼器
  • 近似函式 $q_\phi(\mathbf{z} \vert \mathbf{x})$ 是機率編碼器,其作用與上面的 $g_\phi(\mathbf{z} \vert \mathbf{x})$ 類似。

損失函式:ELBO

估計的後驗 $q_\phi(\mathbf{z}\vert\mathbf{x})$ 應該非常接近實數 $p_\theta(\mathbf{z}\vert\mathbf{x})$。我們可以使用 Kullback-Leibler 散度來量化這兩個分佈之間的距離。KL 散度 $D_\text{KL}(X|Y)$ 測量使用分佈 Y 表示 X 時丟失的資訊量。

在我們的例子中,我們希望最小化 $D_\text{KL}( q_\phi(\mathbf{z}\vert\mathbf{x}) | p_\theta(\mathbf{z}\vert\mathbf{x}) )$ 相對於 $\phi$。

但是為什麼要使用 $D_\text{KL}(q_\phi | p_\theta)$ (反轉 KL) 而不是 $D_\text{KL}(p_\theta | q_\phi)$ (正向 KL)?Eric Jang在他關於貝葉斯變分方法的帖子中有一個很好的解釋。快速回顧一下:

從Autoencoder to Beta-VAE圖 7.正向和反向 KL 散度對如何匹配兩個分佈有不同的要求。(圖片來源:blog.evjang.com/2016/08/variational-bayes.html)

  • 遠向KL背離:$D_\text{KL}(P|Q) = \mathbb{E}_{z\sim P(z)} \log\frac{P(z)}{Q(z)}$;我們必須確保 Q(z)>0 在 P(z)>0 中。最佳化的變分分佈 $q(z)$ 必須覆蓋整個 $p(z)$。
  • 反向 KL 散度:$D_\text{KL}(Q|P) = \mathbb{E}_{z\sim Q(z)} \log\frac{Q(z)}{P(z)}$;最小化反向 KL 背離會將 $Q(z)$ 壓縮到 $P(z)$ 以下。

現在讓我們擴充套件等式:

$$ \begin{aligned} & D_\text{KL}( q_\phi(\mathbf{z}\vert\mathbf{x}) \| p_\theta(\mathbf{z}\vert\mathbf{x}) ) & \\ &=\int q_\phi(\mathbf{z} \vert \mathbf{x}) \log\frac{q_\phi(\mathbf{z} \vert \mathbf{x})}{p_\theta(\mathbf{z} \vert \mathbf{x})} d\mathbf{z} & \\ &=\int q_\phi(\mathbf{z} \vert \mathbf{x}) \log\frac{q_\phi(\mathbf{z} \vert \mathbf{x})p_\theta(\mathbf{x})}{p_\theta(\mathbf{z}, \mathbf{x})} d\mathbf{z} & \scriptstyle{\text{;因為 }p(z \vert x) = p(z, x) / p(x)} \\ &=\int q_\phi(\mathbf{z} \vert \mathbf{x}) \big( \log p_\theta(\mathbf{x}) + \log\frac{q_\phi(\mathbf{z} \vert \mathbf{x})}{p_\theta(\mathbf{z}, \mathbf{x})} \big) d\mathbf{z} & \\ &=\log p_\theta(\mathbf{x}) + \int q_\phi(\mathbf{z} \vert \mathbf{x})\log\frac{q_\phi(\mathbf{z} \vert \mathbf{x})}{p_\theta(\mathbf{z}, \mathbf{x})} d\mathbf{z} & \scriptstyle{\text{;因為 }\int q(z \vert x) dz = 1}\\ &=\log p_\theta(\mathbf{x}) + \int q_\phi(\mathbf{z} \vert \mathbf{x})\log\frac{q_\phi(\mathbf{z} \vert \mathbf{x})}{p_\theta(\mathbf{x}\vert\mathbf{z})p_\theta(\mathbf{z})} d\mathbf{z} & \scriptstyle{\text{;因為 }p(z, x) = p(x \vert z) p(z)} \\ &=\log p_\theta(\mathbf{x}) + \mathbb{E}_{\mathbf{z}\sim q_\phi(\mathbf{z} \vert \mathbf{x})}[\log \frac{q_\phi(\mathbf{z} \vert \mathbf{x})}{p_\theta(\mathbf{z})} - \log p_\theta(\mathbf{x} \vert \mathbf{z})] &\\ &=\log p_\theta(\mathbf{x}) + D_\text{KL}(q_\phi(\mathbf{z}\vert\mathbf{x}) \| p_\theta(\mathbf{z})) - \mathbb{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x})}\log p_\theta(\mathbf{x}\vert\mathbf{z}) & \end{對齊} $$

所以我們有:

$$ D_\text{KL}( q_\phi(\mathbf{z}\vert\mathbf{x}) \| p_\theta(\mathbf{z}\vert\mathbf{x}) ) =\log p_\theta(\mathbf{x}) + D_\text{KL}(q_\phi(\mathbf{z}\vert\mathbf{x}) \| p_\theta(\mathbf{z})) - \mathbb{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x})}\log p_\theta(\mathbf{x}\vert\mathbf{z}) $$

重新排列等式的左右兩側後,

$$ \log p_\theta(\mathbf{x}) - D_\text{KL}( q_\phi(\mathbf{z}\vert\mathbf{x}) \| p_\theta(\mathbf{z}\vert\mathbf{x}) ) = \mathbb{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x})}\log p_\theta(\mathbf{x}\vert\mathbf{z}) - D_\text{KL}(q_\phi(\mathbf{z}\vert\mathbf{x}) \| p_\theta(\mathbf{z})) $$

方程的 LHS 正是我們在學習真實分佈時想要最大化的:我們希望最大化生成真實資料的 (log-) 可能性(即 $\log p_\theta(\mathbf{x})$),並最小化實數和估計後驗分佈之間的差異(術語 $D_\text{KL}$ 的工作原理類似於正則化器)。請注意,$p_\theta(\mathbf{x})$ 相對於 $q_\phi$ 是固定的。

上述的否定定義了我們的損失函式:

$$ \begin{aligned} L_\text{VAE}(\theta, \phi) &= -\log p_\theta(\mathbf{x}) + D_\text{KL}( q_\phi(\mathbf{z}\vert\mathbf{x}) \| p_\theta(\mathbf{z}\vert\mathbf{x}) )\\ &= - \mathbb{E}_{\mathbf{z} \sim q_\phi(\mathbf{z}\vert\mathbf{x})} \log p_\theta(\mathbf{x}\vert\mathbf{z}) + D_\text{KL}( q_\phi(\mathbf{z}\vert\mathbf{x}) \| p_\theta(\mathbf{z}) ) \\ \theta^{*}, \phi^{*} &= \arg\min_{\theta, \phi} L_\text{VAE} \end{對齊} $$

在變分貝葉斯方法中,這種損失函式稱為變分下界證據下界。名稱中的“下界”部分來自以下事實:KL 散度始終為非負數,因此 $-L_\text{VAE}$ 是 $\log p_\theta (\mathbf{x})$ 的下限。

$$ -L_\text{VAE} = \log p_\theta(\mathbf{x}) - D_\text{KL}( q_\phi(\mathbf{z}\vert\mathbf{x}) \| p_\theta(\mathbf{z}\vert\mathbf{x}) ) \leq \log p_\theta(\mathbf{x}) $$

因此,透過最小化損失,我們最大化了生成真實資料樣本的機率下限。

重新引數化技巧

損失函式中的期望項呼叫從 $\mathbf{z} \sim q_\phi(\mathbf{z}\vert\mathbf{x})$ 生成樣本。取樣是一個隨機過程,因此我們不能反向傳播梯度。為了使其可訓練,引入了重新引數化技巧:通常可以將隨機變數 $\mathbf{z}$ 表示為確定性變數 $\mathbf{z} = \mathcal{T}_\phi(\mathbf{x}, \boldsymbol{\epsilon})$,其中 $\boldsymbol{\epsilon}$ 是輔助獨立隨機變數,由 $\phi$ 引數化的變換函式 $\mathcal{T}_\phi$ 將 $\boldsymbol{\epsilon}$ 轉換為 $\mathbf{z}$。

例如,$q_\phi(\mathbf{z}\vert\mathbf{x})$ 形式的常見選擇是具有對角協方差結構的多元高斯:

$$ \begin{aligned} \mathbf{z} &\sim q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)}) = \mathcal{N}(\mathbf{z}; \boldsymbol{\mu}^{(i)}, \boldsymbol{\sigma}^{2(i)}\boldsymbol{I}) & \\ \mathbf{z} &= \boldsymbol{\mu} + \boldsymbol{\sigma} \odot \boldsymbol{\epsilon} \text{, where } \boldsymbol{\epsilon} \sim \mathcal{N}(0, \boldsymbol{I}) & \scriptstyle{\text{;重新引數化技巧。 \end{對齊} $$

其中 $\odot$ 是指按元素乘積。

從Autoencoder to Beta-VAE圖 8.說明重引數化技巧如何使 \(\mathbf{z}\) 取樣過程可訓練。(圖片來源:Kingma NIPS 2015研討會演講中的幻燈片12)

重引數化技巧也適用於其他型別的分佈,而不僅僅是高斯分佈。 在多元高斯情況下,我們透過學習分佈的均值和方差 $\mu$ 和 $\sigma$ 來使模型可訓練,顯式使用重引數化技巧,而隨機性保留在隨機變數 $\boldsymbol{\epsilon} \sim \mathcal{N}(0, \boldsymbol{I})$ 中。

從Autoencoder to Beta-VAE圖 9.具有多元高斯假設的變分自編碼器模型的圖示。

β-VAE型

如果推斷的潛在表示 $\mathbf{z}$ 中的每個變數只對一個生成因子敏感,而對其他因子相對不變,我們將說這個表示是解糾纏或因式分解的。解糾纏表示通常帶來的一個好處是良好的可解釋性和易於泛化到各種任務。

例如,在人臉照片上訓練的模型可能會在不同的維度中捕捉溫和、膚色、頭髮顏色、頭髮長度、情感、是否戴眼鏡和許多其他相對獨立的因素。這種解開的表示對面部影像的生成非常有益。

β-VAE(Higgins等人,2017)是變分自動編碼器的改進,特別強調發現解纏的潛在因素。遵循 VAE 中的相同激勵,我們希望最大限度地提高生成真實資料的機率,同時保持真實分佈和估計後驗分佈之間的距離較小(例如,在小常數 $\delta$ 下):

$$ \begin{aligned} &\max_{\phi, \theta} \mathbb{E}_{\mathbf{x}\sim\mathcal{D}}[\mathbb{E}_{\mathbf{z} \sim q_\phi(\mathbf{z}\vert\mathbf{x})} \log p_\theta(\mathbf{x}\vert\mathbf{z})]\\ &\text{subject to } D_\text{KL}(q_\phi(\mathbf{z}\vert\mathbf{x})\|p_\theta(\mathbf{z})) < \delta \end{對齊} $$

我們可以在 KKT 條件下將其改寫為具有拉格朗日乘數 $\beta$ 的拉格朗日量。上述只有一個不等式約束的最佳化問題等價於最大化以下方程 $\mathcal{F}(\theta, \phi, \beta)$:

$$ \begin{aligned} \mathcal{F}(\theta, \phi, \beta) &= \mathbb{E}_{\mathbf{z} \sim q_\phi(\mathbf{z}\vert\mathbf{x})} \log p_\theta(\mathbf{x}\vert\mathbf{z}) - \beta(D_\text{KL}(q_\phi(\mathbf{z}\vert\mathbf{x})\|p_\theta(\mathbf{z})) - \delta) & \\ & = \mathbb{E}_{\mathbf{z} \sim q_\phi(\mathbf{z}\vert\mathbf{x})} \log p_\theta(\mathbf{x}\vert\mathbf{z}) - \beta D_\text{KL}(q_\phi(\mathbf{z}\vert\mathbf{x})\|p_\theta(\mathbf{z})) + \beta \delta & \\ & \geq \mathbb{E}_{\mathbf{z} \sim q_\phi(\mathbf{z}\vert\mathbf{x})} \log p_\theta(\mathbf{x}\vert\mathbf{z}) - \beta D_\text{KL}(q_\phi(\mathbf{z}\vert\mathbf{x})\|p_\theta(\mathbf{z})) & \scriptstyle{\text{;因為 }\beta,\delta\geq 0} \end{對齊} $$

$\beta$-VAE 的損失函式定義為:

$$ L_\text{BETA}(\phi, \beta) = - \mathbb{E}_{\mathbf{z} \sim q_\phi(\mathbf{z}\vert\mathbf{x})} \log p_\theta(\mathbf{x}\vert\mathbf{z}) + \beta D_\text{KL}(q_\phi(\mathbf{z}\vert\mathbf{x})\|p_\theta(\mathbf{z})) $$

其中拉格朗日乘數 $\beta$ 被視為超引數。

由於 $L_\text{BETA}(\phi, \beta)$ 的否定是拉格朗日量 $\mathcal{F}(\theta, \phi, \beta)$ 的下界。最小化損耗等同於最大化拉格朗日量,因此適用於我們的初始最佳化問題。

當$\beta=1$時,它與VAE相同。當 $\beta > 1$ 時,它會對潛在瓶頸施加更強的約束,並限制 $\mathbf{z}$ 的表示能力。對於一些有條件獨立的生成因素,保持它們的解開是最有效的表示。因此,較高的 $\beta$ 會鼓勵更有效的潛在編碼,並進一步鼓勵解糾纏。同時,更高的 $\beta$ 可能會在重建質量和解纏的程度之間產生權衡。

Burgess, et al. (2017)資訊瓶頸理論的啟發下,深入討論了 $\beta$-VAE 中的膨脹,並進一步提出了對 $\beta$-VAE 的修改,以更好地控制編碼表示能力。

VQ-VAE 和 VQ-VAE-2

VQ-VAE(“向量化變分自動編碼器”;van den Oord等人,2017)模型透過編碼器學習離散潛在變數,因為離散表示可能更自然地適合語言、語音、推理等問題。

向量量化 (VQ) 是一種將 $K$ 維向量對映到一組有限的“程式碼”向量的方法。該過程與 KNN 演算法非常相似。樣本應對映到的最佳質心碼向量是具有最小歐幾里得距離的質心向量。

設 $\mathbf{e} \in \mathbb{R}^{K \times D}, i=1, \dots, K$ 是 VQ-VAE 中的潛在嵌入空間(也稱為“碼本”),其中 $K$ 是潛在變數類別的數量,$D$ 是嵌入大小。單個嵌入向量為 $\mathbf{e}_i \in \mathbb{R}^{D}, i=1, \dots, K$。

編碼器輸出 $E(\mathbf{x}) = \mathbf{z}_e$ 經過最近鄰查詢以匹配到 $K$ 嵌入向量之一,然後此匹配的程式碼向量成為解碼器 $D(.) 的輸入。$:

$$ \mathbf{z}_q(\mathbf{x}) = \text{Quantize}(E(\mathbf{x})) = \mathbf{e}_k \text{ 其中 } k = \arg\min_i \|E(\mathbf{x}) - \mathbf{e}_i \|_2 $$

請注意,離散潛在變數在不同的網路應用中可以具有不同的形狀;例如,1D 用於語音,2D 用於影像,3D 用於影片。

從Autoencoder to Beta-VAE圖 10.VQ-VAE的架構(圖片來源:van den Oord, et al. 2017)

由於 argmin() 在離散空間上是不可微分的,因此解碼器輸入 $\mathbf{z}_q$ 的梯度 $\nabla_z L$ 被複制到編碼器輸出 $\mathbf{z}_e$。除了重建損耗外,VQ-VAE還最佳化了:

  • VQ 損失:嵌入空間和編碼器輸出之間的 L2 誤差。
  • 承諾損失:一種鼓勵編碼器輸出靠近嵌入空間並防止其從一個程式碼向量到另一個程式碼向量過於頻繁地波動的措施。
$$ L = \underbrace{\|\mathbf{x} - D(\mathbf{e}_k)\|_2^2}_{\textrm{重建損失}} + \underbrace{\|\text{sg}[E(\mathbf{x})] - \mathbf{e}_k\|_2^2}_{\textrm{VQ loss}} + \underbrace{\beta \|E(\mathbf{x}) - \text{sg}[\mathbf{e}_k]\|_2^2}_{\textrm{承諾損失}} $$

其中 $\text{sq}[.]$ 是運算子。stop_gradient

碼本中的嵌入向量透過 EMA(指數移動平均線)進行更新。給定一個程式碼向量 $\mathbf{e}_i$,假設我們有 $n_i$ 編碼器輸出向量 $\{\mathbf{z}_{i,j}\}_{j=1}^{n_i}$,它們被量化為 $\mathbf{e}_i$:

$$ N_i^{(t)} = \gamma N_i^{(t-1)} + (1-\gamma)n_i^{(t)}\;\;\; \mathbf{m}_i^{(t)} = \gamma \mathbf{m}_i^{(t-1)} + (1-\gamma)\sum_{j=1}^{n_i^{(t)}}\mathbf{z}_{i,j}^{(t)}\;\;\; \mathbf{e}_i^{(t)} = \mathbf{m}_i^{(t)} / N_i^{(t)} $$

其中 $(t)$ 是指時間上的批處理序列。$N_i$ 和 $\mathbf{m}_i$ 分別是累積向量計數和體積。

VQ-VAE-2 (Ali Razavi, et al. 2019) 是一個結合自注意力自迴歸模型的兩級分層 VQ-VAE。

  1. 第 1 階段是訓練分層 VQ-VAE:分層潛在變數的設計旨在將區域性模式(即紋理)與全域性資訊(即物件形狀)分開。較大的底層程式碼本的訓練也以較小的頂級程式碼為條件,因此它不必從頭開始學習所有內容。
  2. 第 2 階段是學習潛在離散碼本的先驗,以便我們從中取樣並生成影像。透過這種方式,解碼器可以接收從與訓練中的分佈相似的分佈中取樣的輸入向量。使用具有多頭自注意力層增強的強大自迴歸模型來捕獲先驗分佈(如 PixelSNAIL;Chen 等人,2017 年)。

考慮到 VQ-VAE-2 依賴於在簡單的分層設定中配置的離散潛在變數,其生成的影像質量非常驚人。

從Autoencoder to Beta-VAE圖 11.分層VQ-VAE架構和多級影像生成。(圖片來源:Ali Razavi, et al. 2019)從Autoencoder to Beta-VAE圖 12.VQ-VAE-2 演算法。(圖片來源:Ali Razavi, et al. 2019)

TD-VAE型

TD-VAE(“時間差VAE”;Gregor 等人,2019 年)適用於順序資料。它依賴於三個主要思想,如下所述。

從Autoencoder to Beta-VAE圖 13.狀態空間模型作為馬爾可夫鏈模型。

1. 狀態空間模型
在(潛在)狀態空間模型中,一系列未觀測到的隱藏狀態 $\mathbf{z} = (z_1, \dots, z_T)$ 決定了觀測狀態 $\mathbf{x} = (x_1, \dots, x_T)$。圖 13 中馬爾可夫鏈模型中的每個時間步長都可以採用與圖 6 類似的方式進行訓練,其中棘手的後驗 $p(z \vert x)$ 由函式 $q(z \vert x)$ 近似。

2.信念狀態
智慧體應該學會對所有過去的狀態進行編碼,以推理未來,命名為信念狀態,$b_t = belief(x_1, \dots, x_t) = belief(b_{t-1}, x_t)$。鑑於此,以過去為條件的未來狀態的分佈可以寫成 $p(x_{t+1}, \dots, x_T \vert x_1, \dots, x_t) \approx p(x_{t+1}, \dots, x_T \vert b_t)$。在TD-VAE中,迴圈策略中的隱藏狀態被用作代理的置信狀態。因此,我們有 $b_t = \text{RNN}(b_{t-1}, x_t)$。

3.跳躍預測
此外,智慧體應該根據迄今為止收集到的所有資訊來想象遙遠的未來,這表明能夠做出跳躍預測,即預測未來幾步的狀態。

回想一下我們從上面的方差下限中學到的東西:

$$ \begin{aligned} \log p(x) &\geq \log p(x) - D_\text{KL}(q(z|x)\|p(z|x)) \\ &= \mathbb{E}_{z\sim q} \log p(x|z) - D_\text{KL}(q(z|x)\|p(z)) \\ &= \mathbb{E}_{z \sim q} \log p(x|z) - \mathbb{E}_{z \sim q} \log \frac{q(z|x)}{p(z)} \\ &= \mathbb{E}_{z \sim q}[\log p(x|z) -\log q(z|x) + \log p(z)] \\ &= \mathbb{E}_{z \sim q}[\log p(x, z) -\log q(z|x)] \\ \log p(x) &\geq \mathbb{E}_{z \sim q}[\log p(x, z) -\log q(z|x)] \end{對齊} $$

現在,讓我們將狀態 $x_t$ 的分佈建模為機率函式,該函式以當前時間步長和退一步時的所有過去狀態 $x_{<t}$ 和兩個潛在變數 $z_t$ 和 $z_{t-1}$ 為條件:

$$ \log p(x_t|x_{<{t}}) \geq \mathbb{E}_{(z_{t-1}, z_t) \sim q}[\log p(x_t, z_{t-1}, z_{t}|x_{<{t}}) -\log q(z_{t-1}, z_t|x_{\leq t})] $$

繼續擴充套件等式:

$$ \begin{aligned} & \log p(x_t|x_{<{t}}) \\ &\geq \mathbb{E}_{(z_{t-1}, z_t) \sim q}[\log p(x_t, z_{t-1}, z_{t}|x_{<{t}}) -\log q(z_{t-1}, z_t|x_{\leq t})] \\ &\geq \mathbb{E}_{(z_{t-1}, z_t) \sim q}[\log p(x_t|\color{red}{z_{t-1}}, z_{t}, \color{red}{x_{<{t}}}) + \color{blue}{\log p(z_{t-1}, z_{t}|x_{<{t}})} -\log q(z_{t-1}, z_t|x_{\leq t})] \\ &\geq \mathbb{E}_{(z_{t-1}, z_t) \sim q}[\log p(x_t|z_{t}) + \color{blue}{\log p(z_{t-1}|x_{<{t}})} + \color{blue}{\log p(z_{t}|z_{t-1})} - \color{green}{\log q(z_{t-1}, z_t|x_{\leq t})}] \\ &\geq \mathbb{E}_{(z_{t-1}, z_t) \sim q}[\log p(x_t|z_{t}) + \log p(z_{t-1}|x_{<{t}}) + \log p(z_{t}|z_{t-1}) - \color{green}{\log q(z_t|x_{\leq t})} - \color{green}{\log q(z_{t-1}|z_t, x_{\leq t})}] \end{對齊} $$

請注意以下兩點:

  • 根據馬爾可夫假設,紅色項可以忽略不計。
  • 藍色項根據馬爾可夫假設展開。
  • 綠色項被擴充套件為包括追溯到過去的一步預測,作為平滑分佈。

準確地說,有四種型別的分佈需要學習:

  1. $p_D(.)$ 是解碼器分佈:
  • $p(x_t \mid z_t)$ 是通用定義的編碼器;
  • $p(x_t \mid z_t) \to p_D(x_t \mid z_t)$;
  1. $p_T(.)$ 是轉移分佈:
  • $p(z_t \mid z_{t-1})$ 捕獲潛在變數之間的順序依賴關係;
  • $p(z_t \mid z_{t-1}) \to p_T(z_t \mid z_{t-1})$;
  1. $p_B(.)$ 是置信分佈:
  • $p(z_{t-1} \mid x_{<t})$ 和 $q(z_t \mid x_{\leq t})$ 都可以使用置信狀態來預測潛在變數;
  • $p(z_{t-1} \mid x_{<t}) \to p_B(z_{t-1} \mid b_{t-1})$;
  • $q(z_{t} \mid x_{\leq t}) \to p_B(z_t \mid b_t)$;
  1. $p_S(.)$ 是平滑分佈:
  • 回溯到過去的平滑項 $q(z_{t-1} \mid z_t, x_{\leq t})$ 也可以重寫為依賴於信念狀態;
  • $q(z_{t-1} \mid z_t, x_{\leq t}) \to p_S(z_{t-1} \mid z_t, b_{t-1}, b_t)$;

為了結合跳躍預測的思想,順序 ELBO 不僅要處理 $t t+1$,還要處理兩個遙遠的時間戳 $t_1 < t_2$。以下是要最大化的最終TD-VAE目標函式:

$$ J_{t_1, t_2} = \mathbb{E}[ \log p_D(x_{t_2}|z_{t_2}) + \log p_B(z_{t_1}|b_{t_1}) + \log p_T(z_{t_2}|z_{t_1}) - \log p_B(z_{t_2}|b_{t_2}) - \log p_S(z_{t_1}|z_{t_2}, b_{t_1}, b_{t_2})] $$

從Autoencoder to Beta-VAE圖 14.TD-VAE架構的詳細概述,做得非常好。(圖片來源:TD-VAE論文)


被引用為:

@article{weng2018VAE,
  title   = "From Autoencoder to Beta-VAE",
  author  = "Weng, Lilian",
  journal = "lilianweng.github.io",
  year    = "2018",
  url     = "https://lilianweng.github.io/posts/2018-08-12-vae/"
}

相關文章