VQ-VAE 리뷰

본 논문의 제목은 Neural Discrete Representation Learning [1] 이며, vector quantizing 기법, codebook 등을 사용하여 continuous한 data를 discrete하게 representation 추출이 가능한 기법을 제안한 논문이다. 본 논문에서는 clustering 기법으로 해당 벡터를 분류하였지만, 이 논문 외에도 Gumbel-Softmax [2] 등을 사용하여 더 higher 한 dimension의 codebook을 사용할 수 있음을 밝힌다. 

 

Overview

- 연속적 (continuous) 데이터 특성에서 이산 (discrete)으로 학습

- auto-encoder 혹은 VAE의 bottleneck 부분인 latent space에서 뽑힌 latent vector들을 codebook 형태로 embedding 한 후, 제한된 수의 코드북 vector를 이용하여 data를 reconstruction

 

Proposed method

 

VQ-VAE model architecture

Learning from Continuous to Discrete

- 모델은 입력 $x$로부터 받아 인코더를 통해 출력 $z_e(x)$를 생성함

- 이 때, discrete latent variable $z$는 shared embedding space인 $e$로부터 가장 가까운 neighborhood와의 distance 계산에 의해 결정됨. 즉 가장 가까운 vector로 결정

- 여기에서 $e$인 embedding space를 codebook으로 불리우며, $e\in{R^{K*D}}$이며, 여기에서 $K$는 discrete latent space의 size가 됨 (즉, $K$-way categorical)

- 그림에서 $D$는 각각의 latent embedding vector인 $e_i$의 차원 수

- 종합적으로, $K$개의 embedding vector가 존재하며, 이를 수식으로 나타내면 $e_i\in{R^D}, i\in{1,2,...,K}$가 됨

 

Figure from Dongseong Hwang

해당 내용을 풀어서 나타내면 위의 그림 처럼 될 수 있음

- $K$는 one-hot vector의 개수가 되며, $d$는 one-hot vector의 dimension이 됨

1) (e.g., 이미지) 이미지로부터 CNN을 통해 풀링되어 reshape되어 32x32가 64개가 discrete하게 됨

2) 위로부터 각 각에 인덱스가 1개씩 뽑히고, 이게 one-hot vector가 되기 위해 K-means clustering에 의해 distance 계산

3) 가장 가까운 녀석을 찾은 뒤

4) dictionary에 해당 $e_1,...,e_K$로 insertion

5) 원래 이미지로 복원하기 위해 다시 reshape 이후

6) 1번과의 reconstruction loss 계산

 

- 즉, 코드북은 dimension이 512이고, 총 64개의 one-hot vector가 있음 

- 이것은 latent vector랑 가장 가까운 녀석이 무엇인지 뽑음 (여기에서 online k-mean을 쓰고, Euclidean distance로 measuring하여 뽑음)

- 이 식으로부터 디코더가 사진을 다시 복원하는데, 즉 가장 가까운 녀석이 뽑혔기 때문에 encoder와 decoder가 비슷해 짐

 

Loss

VQ-VAE loss

 

- 모델이 잘 학습되면, $z_e(x)$랑 $z_q(x)$가 거의 비슷해짐

- loss가 타고 오면, codebook은 학습 시킬 수 있고, argmin으로 들어오는 것은 gradient가 못 흐름

- 그러므로 gradient를 복사하며, 이것이 encoder로 끌고 들어감 (stop gradient --> sg)

 

1) 전체적인 로스
- X=이미지, $z_q(x)$가 latent space, 전체 decoder (recon)

 

2) Codebook
- Stop gradient이며, $e$가 코드북

- Encoder로부터 뽑혀 나온 녀석에 stop gradient 걸고 codebook 업데이트

- 즉, encoder에서 만들어진 녀석과 비슷해지도록 함 (codebook --> encoder)

 

3) Commitment

- Encoder를 codebook 처럼 학습시킴 (encoder --> codebook)

 

 

[1] van den Oord, Aäron, Oriol Vinyals, and Koray Kavukcuoglu. "Neural Discrete Representation Learning." NIPS. 2017.