[논문 review] Neural Discrete Representation Learning (VQ-VAE)

공부 정리용 블로그입니다. 미숙한 실력으로 인해 표현이나 설명에서 많은 오류가 있을 수 있으니 참고 바랍니다 :)

---

vq-vae

 

이번에는 ae와 vae를 간단히 다루고 vq-vae로 넘어가겠다.

 

Background

AutoEncoder

ae 구조

 

autoencoder(이하 ae) 는 단순히 latent vector z 를 잘 추출하는 것을 목적으로 한다. 해당 network는 reconstructioon loss만을 최소화한다. 즉, 인풋과 디코더의 아웃풋 간의 비교만 진행하는 셈이다. 또한, latent vector의 길이에 영향을 많이 받는다. 길이에 따른 표현의 정도가 다르기 때문이다.

 

ae는 latent 정보를 target 값에 일대일 매핑시키는 방향으로 학습이 된다. 그러다 보니, 오버피팅의 위험도가 높고 새로운 정보에 대해서는 잘 작동하지 못한다. 생성모델로서의 역할을 수행하지 못한다고 해석할 수 있다.

 

VAE (Variational ...)

새로운 데이터를 생성하는 '생성 모델'을 위해서 제안된 것이 바로 VAE이다. 이것은 latent 정보를 target과 매핑시키는 것이 아닌 인풋 데이터의 '분포'를 학습시킨다. 학습에 사용된 데이터들의 분포를 학습하고, 해당 분포와 유사하도록 값을 조합한다면 분명 그럴듯한 데이터가 생성될 것이다! 라는 아이디어이다. 

 

VAE

 

VAE의 전체적인 흐름을 보면 다음과 같다.

1. 인코더를 통해 데이터의 평균과 표준편차를 구한다.
2. 해당 값을 이용하여 정규 분포를 형성한다.
3. 해당 분포에서 z를 샘플링한다.
4. 샘플링 값을 디코더에 넣어 새로운 데이터를 생성한다.

하지만 위의 방법에는 문제가 있다. 딥러닝의 핵심인 backprop을 사용하기 위해서는 미분이 가능해야 하지만, z 값을 랜덤으로 샘플링하는 과정은 미분이 가능하지 않다. 이를 해결하기 위해 reparameteriztion trick이 사용된다.

 

해당 개념은 정규분포에서 샘플링을 한 것은 아니지만, 샘플링한 것처럼 보여주자라는 trick이다.

1. 처음에는 모든 값들이 초기화되어 있다. 각 layer의 weight이 랜덤으로 initialize되어 있기 때문이다.
2. 표준 정규 분포에서 epsilon을 뽑는다.
3. 랜덤으로 설정된 평균, 표준편차, 입실론을 이용하여 latent vector z를 세팅한다. 처음에는 당연히 z는 임의의 분포를 가진다.
4. 우리는 z를 정규 분포로 바꾼 후, 여기에서 샘플링을 하고자 한다. 
5. 4번의 목표를 이루기 위해, 우리는 z와 정규 분포에 KL Divergence loss를 적용하여 z의 분포를 맞추어준다.
6. 학습이 완료된 network의 z는 마치 정규 분포에서 샘플링된 것 같은 결과물을 가지게 된다 (trick!)

이렇게 완벽한 줄만 알았던 VAE에도 문제가 생긴다. Posterior Collapse. 디코더가 latent를 무시하고 output을 생성하는 현상이다.

 

이것은 다음과 같은 상황에서 발생할 수 있다. 먼저 PixelCNN과 같이 디코더가 너무 강한 경우이다. 효과적인 latent를 생성하도록 모델이 학습되지 않아도 디코더가 올바른 결과를 뽑는 것이다. 이렇게 되면 인코더의 중요성이 사라지게 된다.

또는 loss 항 중에서 KL 항이 과도한 중요성을 가질 경우이다. latent를 정규 분포에 '너무' 맞추려고 노력해서 모든 값을 동일한 정규 분포로 맞추는 것이다. 이것도 latent를 학습하지 못한다.

 

분포를 학습하다가 문제가 생기니 그러면 분포 만들지마! 해서 나온 것이 vq-vae이다.

 

Methods

VQ-VAE

vq-vae

 

vq-vae는 분포를 학습하는 대신, 이산화된 값을 이용한다.

1. 인풋을 CNN 인코더에 넣어 Ze를 얻는다.
2. 각 vector와 코드북에서의 가장 가까운 벡터 인덱스를 이용하여 q(z|x) 얻는다.
3. 각 값에 해당 인덱스 코드북 벡터를 넣어 Ze와 사이즈가 같은 Zq를 얻는다.
4. Zq를 CNN 디코더에 넣어 데이터를 생성한다.

여기에서 우리가 학습해야하는 것은 인코더와 디코더, 그리고 코드북이다. (코드북은 사전에 세팅되는 것이 아니라는 점 유의하자.) 이것들을 어떻게 학습해야하는지 알아보자.

 

Loss

loss

 

Reconstruction loss

인코더의 인풋과 디코더의 아웃풋 간의 loss로 인코더와 디코더를 학습시킨다.

straight-through

이것도 코드북의 값을 매핑시키는 과정을 미분할 수 없는 문제가 생긴다. 여기에서는 straight-through라는 기법으로 문제를 해결한다. 디코더의 gradient를 인코더에 사용하여 업데이트를 하겠다는 것이다. 이게 어떻게 유의미한 결과를 가져오는지 이해는 못했지만, Ze와 Zq의 사이즈가 같아 가능하다고 논문에서는 말한다. 

recon_loss 식

 

straight-through로 인해 코드북이 학습이 진행되지 않기 때문에 추가적인 loss 설정이 필요하다.

 

Embedding loss

앞서 말한 것처럼, 해당 loss는 코드북을 학습시키는 과정이다. 인코더의 아웃풋과 코드북의 값을 유사하게 만드는 과정이다. 코드북만 학습시키기 위해 인코더의 아웃풋에는 sg (stop gradient)를 이용하여 업데이트를 막았다.

embedding_loss 식

 

코드에서 detach메소드를 이용하여 update를 막을 것을 확인할 수 있다.

 

Commitment loss

해당 loss는 인코더의 업데이트로 인해 코드북보다 인코더의 아웃풋이 더 빠르게 변하여 두 값의 차이가 커지는 것을 막고자 인코더 값을 코드북에 commit하도록 하는 역할을 한다. 코드북의 학습을 막기 위해 코드북에 sg를 부여하였다.

commit_loss 식

 

앞선 embedding loss와 달리 코드북에 detach를 적용한 모습을 확인할 수 있다.

 

Prior

이러한 과정으로 학습된 network의 q(z|x) 분포를 구하고자 하는 과정이 추가적으로 진행된다. 이는 새로운 데이터를 생성하기 위함으로 확인된다. 분포를 알아야 그 분포를 따르는 데이터의 추가 생성이 가능하기 때문이다.

process of generation

 

이미지의 경우 PixelCNN, 오디오의 경우 WaveNet을 이용하여 분포를 학습한 후, 해당 분포에서 샘플링된 값을 디코더에 넣어 데이터를 생성한다. 이것이 vq-vae가 데이터를 generate하는 과정이다.

 

Evaluation

VAE : VQ-VAE : VIMCO = 4.51 : 4.67 : 5.14 bits/dim (압축률)

해당 단위에 대한 설명이 정확하게 없지만, 얼마나 더 적은 bits로 표현이 가능하냐는 의미로 해석된다. 따라서 성능은 VAE, VQ-VAE, VIMCO 순으로 좋다고 할 수 있다.

 

image evaluation

 

VQ-VAE는 low-dimensionality를 이용하였기에 빨간색 동그라미처럼 조금 흐린 부분이 존재한다.

 

image generation

 

새로운 이미지의 생성에 있어서도 빨간 네모처럼 부자연스러운 모습도 존재한다.

 

Conclusion

본 논문은 vector quantization을 결합한 새로운 모델을 제안하였다. discrete latent 문제를 해결하기 위해 straight-through를 사용한 점은 인상적이다. 

또한, continuous latent를 사용한 기존의 모델과 유사한 성능을 보인다는 점은 충분히 가치가 있다. vq-vae는 코드북의 값으로만 인코딩을 할 수 있는 deterministic한 성격을 띠고 있기 때문에 dimension이 낮다는 장점이 있다.

 

최대한 수학적 수식을 배제하고 진행하려다 보니 내용이 아쉬운 느낌이 있다. 능력이 된다면 generative model에 대해서 수학적으로 다뤄볼 생각이다.