帶你認識Vector-Quantized Variational AutoEncoder - 理論篇
說到近年來最火紅以深度學習為主的生成模型，大家必定會想到生成對抗網路(Generative Adversarial Network, GAN)，然而在GAN(2014)還沒被提出來之前，有另外一個同樣屬於生成模型的Variational AutoEnoder (VAE)常被大家所使用，很可惜的是當時GAN在許多任務上所產生的圖片清晰度較高，因此VAE類型的模型相對而言就勢弱了一些(當然GAN在訓練的特性上有一些難以克服的問題至今也尚未完全解決)。
故事總不會就這樣結束，2017年DeepMind在NIPS研討會上提出了Vector-Quantized Variational AutoEncoder模型，雖然在效果上仍然是先與VAE做比較，但VQ-VAE提出的概念讓它擁有比其它生成模型更獨特的地方，甚至在後續2019年6月提出的VQ-VAE2甚至宣稱在生成1024*1024的高解析度人臉時與當時效果最佳的BigGAN可作比擬。如果你開始對VQ-VAE感到好奇，就跟著我們一起看下去吧。
註1：如果你對Variational AutoEncoder甚至是AutoEncoder的概念還沒那麼熟的話，可以參考此篇AutoEncoder介紹、此篇VAE介紹、或是尋找其他資源唷。
VQ-VAE的獨特之處
我們可以這樣解讀AutoEncoder家族在做的事情，Encoder試圖找出輸入圖片x在潛在空間上的表徵(representation)，在大多數的狀況中，大家使用連續型的分布去模擬z的樣貌(e.g. AE將輸入x投影至潛在空間的一個點；VAE則改為使用高斯分布模擬輸入x在潛在空間的樣貌)，然而VQVAE的作者提到離散的潛在表徵在很多情境上也許才是比較適合的，例如語言概念，因此VQ-VAE主要的突破就是試圖讓Encoder產出離散的表徵代表每一筆輸入資料，而Decoder則需要在接收這樣離散的表徵後還原原本的資料。
當然我們可以直接從模型的名稱感受到它究竟是如何做到這件事情。Vector Quantization 向量量化的技巧在訊號處理領域中已經發展了一段時間，主要的做法是將影像/音訊切割不同群組並取得每個群組的代表向量(Figure 1)，另外維護一份有K個編碼向量的編碼簿(codebook)，針對每個群組，以編碼簿中最接近群組代表向量的編碼向量索引作為這個群組的代表，這樣我們就可以將原始的資料轉換為n個索引(n為群組數量)，再加上儲存編碼簿本身，就可以達到資料壓縮與特徵離散化的目的。由於資料壓縮必定會產生資訊的遺失，這個演算法最主要的任務就是在橫跨不同的影像/音訊中找到能讓資訊遺失最少的K個編碼向量。
接著來看一下VQ-VAE的架構，以及它是如何在原本AutoEncoder的框架下加入Vector Quantization的技巧。
從Figure 2 中可以看到VQ-VAE同樣維持著Encoder-Decoder的架構，然而這邊所提取的特徵保留了多維的結構，以圖中所使用的影像資料為例，Encoder最後輸出的潛在表徵Z_e(x)大小將為(h_hidden, w_hidden, D)，其實就是在CNN中我們熟知的Feature map。接著會進入到Vector Quantization的部分，同樣我們會有K個編碼向量(Figure 2 中 Embedding Space的部分)，每一個編碼向量同樣有D個維度，根據Feature Map中(h_hidden, w_hidden)的每個點位比對D維的特徵向量與Codebook中K個編碼向量的相似程度，並且以最接近的編碼向量索引作取代(Figure 2中央藍色的Feature Map部分)，這樣就達到了將原圖轉換為離散表徵的步驟(最後的表徵為(h_hidden, w_hidden, 1)的形狀)。
在Decode階段，只要將Codebook中相對應的編碼向量置於表徵的相對應位置，就可以得到同樣為(h_hidden, w_hidden, D)的轉換後表徵Z_q(x)，接著同樣通過Decoder後就可得到還原後的原始資料。
訓練VQ-VAE網路
上述的流程事實上並不算複雜(如果有人沒有完全理解的話一定是我說明的不清楚Orz)，但在深度學習中我們期待的是end-to-end的訓練架構，因此如何設計損失函數讓我們可以同時訓練Encoder、Decoder，以及找到合適的Embedding Vectors同樣是需要思考的事情。
先來聊聊Encoder和Decoder的部分，我們都知道目前深度學習模型的訓練高度依賴倒傳遞(back-propagation)方法，也就是使用微分的方式計算梯度後以此更新模型權重(參數)，這部分在AE/VAE上也相同。但是修但幾勒，在VQ-VAE的其中一個步驟中，我們使用了argmin (尋找與Z_e(x)最接近的codebook向量並進行取代)，這個操作本身是無法計算梯度的，因此單純地使用原始的倒傳遞沒辦法更新到Encoder的參數。
- Encoding step : X → Z_e(x) (透過NN/CNN)
- Vector Quantization step : Z_e(x) → Z_q(x) (透過argmin)
- Decoding step : Z_q(x) → X' (透過NN/CNN)
若我們將VQ-VAE模型表示成上面的三個步驟，主要無法微分的是第二個步驟。那可怎麼辦呢？在此作者提出了一個Straight-through estimator的作法，在概念上也並無甚麼高深的地方，就是直接將Z_q(x)所造成的誤差直接當作Z_e(x)的誤差作倒傳遞，藉此就能讓decoder輸出與原圖的差異同時更新decoder與encoder(在Figure 2 中為紅色線條)。
另外一個部分要考量的則是codebook中的編碼向量，我們同樣也希望手中的K個編碼向量能夠盡量涵蓋並且代表輸入資料的特徵，因此在訓練VQ-VAE的模型時，同樣也應調整codebook中的參數。在此作者使用l2 norm計算Z_e(x)的每個pixel上的向量與相對應的codebook編碼向量差異，並以此計算梯度修正codebook。
Figure 3 呈現了VQ-VAE的損失函數，總共包含了三個部分，讓我們一一來看一下。第一項的loss為reconstruction loss，也就是在通過了整個VQ-VAE網路後所還原出來的資料與原始輸入資料的差距(在計算上仍然以pixel-wise mean square error估計)，值得提醒的是由於使用了Straight-through estimator方法，這一項的loss並不會調整到codebook中的內容。第二與第三項則與codebook有關，如同上一段所說的，我們計算Z_e(x)與codebook中相對應向量的l2 norm，並試圖同時更新Z_e(x)與codebook向量，然而相對而言我們希望codebook更新的幅度多一些，因此在這邊作者導入了stop gradient的符號並將l2 norm拆成兩項。在stop gradient(sg)符號內的內容在計算loss時會正常存在，但在更新時此項並不會提供任何梯度，也代表著不會被更新。第二項的代表的是計算Z_e(x)與相對應codebook向量(e)的l2 norm，但只更新codebook，第三項則計算相同的l2 norm但只針對Z_e(x)作梯度計算與更新，由於我們希望Z_e(x)不用那麼在意與codebook接近，因此第三項的前面會加上一個commitment cost作為權重(在原始論文中提的數值為0.25)。
VQ-VAE 模型效果
說了那麼多，VQ-VAE在還原資料的效果上如何呢？在原始論文中作者除了在cifar10這種較小的資料集上作測試外，也應用到了ImageNet等較大型的資料集上。從原始論文中的圖片(Figure 4)可以看到，即使在ImageNet這種有差異度極高的資料集上，VQ-VAE的還原效果也不錯。
在上述的模型架構中我們主要以圖片作為示範，然而VQ-VAE的架構在Encoder與Decoder的選擇上是非常彈性的，因此除了圖片之外，作者也應用VQ-VAE到音訊甚至是影片資料上。由於VQ-VAE針對資料做壓縮後再還原將導致部分資訊會有遺失，但在音訊資料上，實驗發現VQ-VAE所還原的資料會保留講者的內容資訊而排除聲調或語氣的部分，這也證明了VQ-VAE後續可能的發展潛力。
另外，上述提到的內容只描述到AutoEncoder系列作的資料壓縮與還原能力，但身為生成模型的一員，VQ-VAE同樣也能產生訓練資料中未見過的資料。在生成新資料時，我們只需要生成潛在表徵的codebook索引(同樣以Figure 1為例，只需生成中間藍色部分的feature map)，再使用decoder還原成原圖即可。原文中作者使用了另一類的自回歸生成模型PixelCNN作潛在表徵結構的生成，但由於篇幅安排跟主題聚焦關係，關於PixelCNN的模型介紹以及結合VQ-VAE進行圖像生成的部分請大家期待後續的系列文章。
給個總結
在GAN佔據大部分研究者的目光時，VQ-VAE試圖拓展大家對生成模型領域發展的想像。VQ-VAE擁有AutoEncoder家族維度縮減的能力，能夠應用於資料壓縮或是特徵提取，另外也能結合自回歸類的生成模型(ex. PixelRNN, PixelCNN)進行資料生成(甚至補足自回歸生成模型生成高解析度影像時效能不佳的問題)，非常期待後續能看到這類型的模型在不同場域有更加優秀的表現。
這次簡介中，我盡量不使用貝氏模型的概念作說明，若大家對一些細節有興趣，以下我也整理了一些曾看過不錯的相關資源供大家參考。在這個系列的下一篇文章中，我將會實際帶大家訓練一個VQ-VAE模型，請大家敬請期待。
