VQ-VAE & NSVQ (Noise Substitution in Vector Quantization)

예시로 든 논문 : LATENT ACTION PRETRAINING FROM VIDEOS

VQ-VAE

VQ-VAE Diagram
Figure 1: Left: A figure describing the VQ-VAE. Right: Visualisation of the embedding space. (Oord et al., 2017, "Neural Discrete Representation Learning").

VQ-VAE의 한계

인코더가 뽑은 벡터 d1​을 코드북의 z1, z2​ 중 하나로 바꾸는 과정을 함수 f(x)라고 할 때, 벡터 d1에 대한 함수 f(x)는 뚝뚝 끊기는 계단 함수 형태를 띈다. 계단 함수는 기본적으로 미분 불가능한 점이 존재하며 대부분 기울기가 0인 구간이라 d1을 출력하는 인코더를 업데이트할 수 없다. 이를 임시방편으로 해결하기 위해 VQ-VAE는 실제 오차(d1, z1 사이 거리)를 무시하는 Pseudo-gradient를 흘려보낸다(Straight-Through Estimator). 이는 학습을 불안정하게 만들고, 특정 코드만 계속 쓰는 Codebook Collapse 현상으로 이어진다.

미분이 불가능한 문제점은 크게 두 갈래,

  • 인코더(d1​)를 학습이 불가하다는 점
  • 코드북(z1​)을 학습이 불가하다는 점 (단어장을 업데이트하지 못함)

로 요약되는데, 이러한 한계를 보완한 방법론이 바로 NSVQ다.

NSVQ

VQ-VAE Diagram
NSVQ architecture. From "NSVQ: Improved Vector Quantization technique for Neural Networks Training", [Towards Data Science](https://towardsdatascience.com/improving-vector-quantization-in-vector-quantized-variational-autoencoders-vq-vae-915f5814b5ce/)

Codebook Collapse의 해결 : 미분 가능한 NSVQ

\[\hat{d}_1​=d_1​+노이즈\] \[\hat{d}_1 = d_1 + \|d_1 - z_1\| \frac{v}{\|v\|}\]

Stop Gradient

stop gradient는 모델의 역전파 과정에서 특정 경로로 흐르는 그라디언트를 강제로 차단하는 연산자다. (표기 : sg[x]) 일부 레이어의 그라디언트는 변하지 않게 함으로써 모델의 부분적 지식을 보존하는 효과가 있다. 논문 LATENT ACTION PRETRAINING FROM VIDEOS에서는 Representation Collapse를 방지하기 위해 d1(인코더가 출력한 임베딩)을 고정한채, 코드북을 이루는 벡터를 업데이트한다. 이렇게 하면 코드북 z1​이 실제 데이터(d1​)의 분포를 따라가게 된다.

Embedding Loss : \(\|sg[d_1] - z_1\|^2_2∥\)