Neural Discrete Representation Learning

• paper : https://arxiv.org/abs/1711.00937
• code : https://colab.research.google.com/github/zalandoresearch/pytorch-vq-vae/blob/master/vq-vae.ipynb

Preliminaries

• Variational AutoEncoder

  • dataset이 있을 때 likelihood $p_\theta(x)$를 구하고 싶음. 이 때, $x$는 어떤 unknown random variable $z$에 dependent함. generative 모델은 $p(z)$와 $p(x|z)$를 알아서 아무 $z$나 넣었을 때 $x$ 와 비슷한 샘플을 뽑고 싶음! 이 때 $z$에 대한 확률 분포 prior를 가정하고, $p_\theta(x)$를 likelihood를 최대화할 때, ELBO 형식으로 만들어 최적화! 이때 $p(z|x)$를 근사하는 encoder를 $q_\phi(z|x)$로 근사!

  

• posterior collapse

  • https://wain-together.tistory.com/8
  • approximate posterior가 prior를 그대로 "mimic"하며, model은 latent variable을 무시한 상태에서 학습이 진행되는것을 의미
  • $x$ ⇒ $z$ ⇒ $x$로 이어지는 VAE 형태에서, latent variable $z$가 $x$에 관한 정보를 잘 담지 못하는걸 의미
  • 수식으로 설명하면, $q_\phi(z|x) = p_\theta(z|x) = p(z)$
  • VAE이 가진 유명한 문제.
    • latent가 discrete할 때 더 잘 일어난다고 알려져 있음.
    • decoder가 encoder에 비해 너무 강하게 학습되면 일어남.

1. Introduction

3. VQ-VAE

우리의 방법론은 VAE와 매우 유사하다.

1. input data $x$를 받고 discrete한 latent random variable $z$를 내뱉는 posterior $q(z|x)$를 parametrize한 encoder network

2. prior distribution $p(z)$

3. $p(x|z)$를 내뱉는 decoder로 구성된다.

이때 VAE의 posterior과 prior는 diagonal covariance를 가진 normal distribution으로 가정되고, gaussian reparametrisation trick을 사용한다.

우리는 vector quantization에 영감을 받아 discrete한 latent variable을 사용하고 새로운 방법으로 학습한 VQ-VAE를 소개한다. posterior과 prior 분포는 카테고리 분포이고, 해당 카테고리 분포가 임베딩의 인덱스에 대한 분포가 되어 샘플링 된다. 이러한 임베딩은 디코더 네트워크의 인풋으로 사용된다.

• vector quantization

3.1 Discrete Latent variables

$K$는 discrete latent space의 크기(=discrete variable의 카테고리 개수)이고, $D$는 임베딩 차원일 때, latent embedding space $e\in \mathbb{R}^{K\times D}$를 정의한다.

모델은 input $x$를 받고 encoder를 통해 $z_e(x)$를 내뱉는다. 이후 embedding space $e$에서 가장 가까운 look-up을 계산해서 discrete latent variable $z$를 찾는다.