기본 개념
- prior
- p(z)를 학습하는 모델
- 인코더는 “입력 x가 있을 때”만 z를 뽑을 수 있음 (압축/복원용)
- prior는 “입력 x 없이도” 그럴듯한 z를 뽑아, 새로운 x를 생성할 수 있게 해줌. (생성용)
- discrete latent variable
- 값의 집합이 연속 공간이 아니라, 유한/가산 집합인 잠재변수
- VQ-VAE에서는 정해진 숫자의 latent vector 집합으로 이루어짐
- 즉 D dimension의 K개의 embedding 집합으로 구성 됨
- continuous latent variable의 경우 해석도 어렵고 효율이 떨어짐
VQ-VAE
모델 구조
- 인코더 - 코드북 - 디코더의 구조
- 코드북
- D차원의 K개의 임베딩으로 이루어져 있음
- 인코더 → 코드북
- 인코더의 아웃풋인 z와 가장 거리가 가까운 임베딩을 코드북에서 찾음
- 즉 연속된 latent를, 유한한 심볼의 집합에서 하나를 고르도록 하는 것
- 코드북 → 디코더
- 코드북에서 찾은 임베딩을 디코더에 전달하여 foward
- z는 2차원, 3차원이 될 수도 있음
- 예를 들어 이미지넷을 학습시킬 때에는 z의 차원은 32x32였음
- 128x128x3 차원의 이미지를 32x32x1 (K=512)로 줄였다가 다시 128x128x3로 복원함
학습 방법
- straight-through gradient estimation
- forward시에는 코드북을 거쳐서 가지만, 코드북은 미분 불가이기 때문에 backward가 안 됨
- 따라서 backward시에는 decoder의 인풋에서 encoder의 output으로 곧바로 연결함
- loss
- sg는 stopgradient operator로, 순전파 시에는 값을 그대로 사용하고, 역전파시에는 gradient가 흐르지 않도록 하여 상수처럼 취급되도록 함
- 첫번째 loss항
- reconstruction 항
- 디코더가 원본을 잘 복원할 수 있도록
- 두번째 loss항
- 코드북 업데이트 항. Vector Quantisation (VQ).
- straight-through gradient estimation로 인해 코드북은 backward시 gradient가 연결되는 부분이 없음.
- 이 loss를 통해 선택된 코드북의 임베딩이, 인코더에서 나온 Ze(x)와 비슷해지도록 학습
- 세번째 loss항
- commitment loss
- 인코더 출력 Ze(x)가 선택된 코드북 임베딩과 비슷해지도록 학습 (너무 커지지 않도록 함)
왜 이렇게 구성하였는가?
- posterior collapse를 겪지 않음
- posterior collapse란?
- 현상: posterior q(z∣x)가 prior p(z)와 거의 같아진다
- 결과: pixcelCNN과 같이 autoregressive 모델에서 디코더가 latent를 무시하게 됨
- latent를 사용하지 않고 Xt-1까지의 정보만 가지고 Xt를 예측함.
- 왜 발생하는가?
- posterior collapse가 생기는 이유
- ELBO
- VAE의 maximum likelihood
- VAE가 최적화하려는 이론적 목적함수
- 따라서 ELBO를 키우려면 reconstruction은 키우고 KL은 줄여야 함
- KL을 최소화하는 최선의 방법은 q(z∣x)를 p(z)와 같게 만드는 것
- 왜 VQ-VAE는 posterior collapse가 발생하지 않는가?
- VQ-VAE 학습은 2단계로 이루어지는데,
- 1단계, prior는 유니폼으로 두고, 인코더 - 코드북 - 디코더 학습
- 2단계, 인코더 - 코드북 - 디코더를 프리징하고, z값을 뽑아서 prior 학습
- 1단계에서 KL이 이미 상수여서 z를 무시하는 쪽으로 학습되지 않는다고 함
- 기존 discrete VAE보다 gradient variance가 덜 큼
- 미분 불가능한 구간들로 인해서 사용한 근사적 추정기들로 인해 큰 분산이 발생한다고 함
- 분산이 크면 업데이트 방향이 요동치고, 수렴이 느려지며, 학습이 불안정해짐. 학습률을 낮춰야 하는 상황도 자주 생김
- VQ-VAE는 straight-through 추정기로 gradient를 직접 전달하거나, codebook을 이용해 결정론적 방식으로 분산이 적다함
- discrete latent 모델이지만 continuous latent 모델만큼의 성능이 나옴
실험 결과
- 정량 비교: CIFAR-10 bits/dim
- 목적: 이산 latent 써도 연속 VAE급 성능 나오는지 확인함
- 결과(bits/dim, 낮을수록 좋음)
- Continuous VAE: 4.51
- VQ-VAE: 4.67
- VIMCO(discrete VAE 계열): 5.14
- bits/dim
- 의미: 데이터 x를 설명하는 데 필요한 비트
- VQ-VAE가 연속 VAE에 근접함, 기존 discrete 추정기보다 훨씬 나음
- 그 외 이미지 생성/비디오/오디오 task에서도 좋은 성능이 나옴
Share article