[筆記][NeurIPS2017] VQ-VAE: Neural Discrete Representation Learning

5 min readJun 5, 2020

Latent space 離散化

Paper: https://papers.nips.cc/paper/7210-neural-discrete-representation-learning.pdf

架構

模型

如同 VAE ， Encoder 是 posterior distribution ，觀測到 x 推論出 z。Decoder 則是 prior distribution ，給予資訊 z 推論出 x 內容。這篇 z 上－也就是 latent space－做離散化問題。離散化的好處是可以根據不同類別給予不同的資訊空間，而不是將所有種類大雜燴，這樣容易影響到相互資訊。因此在圖中所示，在 bottleneck layer 建立 embedding space 概念。

Embedding Space

在 space 上論文定義 e ∈ R^{K x D} ，意思是有 K 個 D 長度的 latent space組成。 K 代表有 K 個種類；D 也就是有維度 D 個。

Posterior categorical distribution probabilities

對於 Encoder 輸入進來的 z ，等式計算進來的 latent space 與 embedding space 每個 latent space 之間的 L2 距離，產生地圖 q(z|x)。將最近的 embedding space ，作為 decoder 輸入。如下面等式所示 z_{q}(x)：

在論文中的上圖裡，在 z 的每一格中找到對應最接近的 latent space，並且作為 decoder 的輸入。

為何訓練得起來?

離散化主要的問題，是在神經網路做 back propagation 。因為是不連續的，也代表它不可微分。

因此在 VQVAE 中，將 decoder 回傳的梯度，直接回流到 encoder 輸出端，不通過 embedding space 或者是 q(z|x)。然而 embedding space 如何學習，將在下方的 loss function-Vector Quantisation 提到。

不過為何可以直接將 decoder 的梯度回傳到 encoder？論文說道，因為 encoder 的目的是要試著接近 embedding space 的距離，而 decoder 也是接收 embedding space 的內容。一直以來都是學習維度 D 組成的空間，如上方圖片， decoder 回傳的 ▽zL 藏有可觀的資訊，間接推論 encoder 要如何靠近那些向量。

Loss function

Reconstruction loss + Vector Quantisation + β*Commitment loss

從上方式子左到右分有三個 loss function:

Reconstruction loss: 藉由 ▽zL 訓練 encoder 與 decoder 的恢復程度。
Vector Quantisation: 因為 ▽zL 直接略過離散化地方，embedding spaces 會訓練不到。因此我們將 encoder 產出的 z 與 embedding space 做距離上訓練。假如 encoder 一個 latent space 對應到e_{7} ，就對於兩個之間的距離當作 loss。要注意的是 z 是不會回傳到 encoder，因為單純讓 spaces 知道怎麼移動空間位置，才有 sg: stop gradient 的符號存在。
Commitment loss: 不過 embedding spaces 在學習 encoder 距離不會比 autoencoder 快，因為 embedding spaces 中每個 latent space 是沒有關係(dimensionless)，每一個 latent space 都可以隨意訓練。需要讓 encoder 認知這件事情，讓 spaces 跟上 autoencoder 。這 loss 也含有 β 這個 hyper parameter，控制 encoder 認知比例。