VQ-VAE Neural Discrete Representation Learning

논문리뷰

VQ-VAE Neural Discrete Representation Learning

AI바라기 2025. 3. 15. 17:33

Overall Summary

VQ-VAE는 vector quantisation을 VAE framework에 결합하여 discrete latent representation을 효과적으로 학습하는 새로운 generative model이다. "Posterior collapse" 문제를 해결하고, continuous latent variable을 사용하는 VAE와 동등한 성능을 달성하며, image, speech, video 등 다양한 modalities에서 유의미한 representation을 unsupervised 방식으로 학습할 수 있음을 보여주었다. 이 논문은 discrete latent variable model의 가능성을 제시하고, representation learning, generative modeling 분야에 중요한 기여를 하였다.

쉬운 설명

VQ-VAE는 이미지를 "단어" (codebook의 embedding vector)로 표현하는 것과 유사하다. 이미지를 encoder에 넣어 얻은 feature (ze(x))를 가장 가까운 "단어" (ek)로 mapping (vector quantisation)하고, 이 "단어"를 decoder에 넣어 이미지를 reconstruct 한다. VQ-VAE는 이미지를 구성하는 핵심 "단어" (discrete latent representation)를 unsupervised 방식으로 학습하며, 이 "단어"들을 조합하여 새로운 이미지를 생성할 수도 있다. 마치 단어를 조합하여 문장을 만들듯이, VQ-VAE는 latent space에서 "단어"들을 조합하여 다양한 이미지를 표현하고 생성한다.

Neural Discrete Representation Learning (VQ-VAE) 학습 노트

Purpose of the Paper

기존 Variational AutoEncoder (VAE)의 continuous latent variable의 한계 극복.
Discrete latent representation을 사용, "posterior collapse" 문제 (latent variable 무시 현상) 해결 및 다양한 modality (image, speech, video)에 적합한 representation learning.
Unsupervised learning 방식으로 유용한 discrete representation 학습하는 generative model 제시.

Key Contributions

VQ-VAE (Vector Quantised-Variational AutoEncoder) Model:
- Encoder network는 discrete code를 출력.
- Prior는 static 대신 learned.
- Vector Quantisation (VQ) 도입, discrete latent representation 학습.
- "Posterior collapse" 방지, variance 문제 없음.
Discrete Latent VAE의 성능 입증:
- Discrete latent model (VQ-VAE)이 continuous model counterparts와 동등한 log-likelihood 성능 달성.
Autoregressive Prior와의 결합:
* Discrete representation과 autoregressive prior (PixelCNN, WaveNet)를 결합, high-quality image, speech, video 생성.
Unsupervised Learning of Meaningful Representations:
- Raw speech에서 supervision 없이 language structure 학습 (phoneme-level language model).
- Unsupervised speaker conversion 가능성 제시.

Novelty

VQ method를 VAE에 통합, "posterior collapse" 없이 discrete latent variable을 효과적으로 학습하는 최초의 모델.
Continuous latent variable을 사용하는 VAE와 동등한 성능을 달성하는 최초의 discrete latent VAE.
다양한 modalities (image, speech, video)에서 유의미한 discrete representation을 unsupervised 방식으로 학습.

Experimental Highlights

ImageNet, CIFAR10, DeepMind Lab, VCTK datasets 사용.
Reconstruction error, log-likelihood, qualitative sample evaluation 사용.
CIFAR10: VQ-VAE (4.67 bits/dim)가 continuous VAE (4.51 bits/dim) 및 VIMCO (5.14 bits/dim)와 comparable한 성능 달성.
ImageNet: 128x128x3 image를 32x32x1 discrete space (K=512)로 압축, high-quality reconstruction.
DeepMind Lab: PixelCNN prior를 활용, long-term structure 학습 및 coherent video generation.
VCTK (Speech):
- Discrete latent representation (64x smaller)으로 content (text) 보존, prosody는 변경된 reconstruction.
- Unsupervised phoneme-level language model 학습.
- Speaker conversion (source speaker의 content는 유지, target speaker의 voice로 변환) 성공.

Limitations and Future Work

Soft-to-hard relaxation training의 어려움: Decoder가 continuous relaxation을 쉽게 invert, quantisation이 제대로 이루어지지 않음.
Joint training of prior and VQ-VAE 미실시: Potential performance improvement 가능성.
Perceptual loss function 미사용: MSE 대신 GAN loss 등 사용 시 reconstruction quality 향상 가능성.
Encoder/decoder structure의 영향: Discrete latent의 의미가 이전 latent에 dependent할 수 있음 (bi/tri-grams 등).
더욱 발전된 mapping to phonemes를 통해, higher accuracy 확보 가능성.

용어 설명

Posterior collapse: VAE에서 powerful decoder 사용 시, latent variable z가 무시되고 decoder가 input x를 perfectly modeling하는 현상.
Continuous latent variable: VAE에서 Gaussian distribution과 같이 continuous value를 갖는 latent variable.
Discrete latent variable model: VAE에서 categorical distribution과 같이 discrete value (code)를 갖는 latent variable을 사용하는 model.

Abstract

Supervision 없이 유용한 representations을 학습하는 것은 machine learning의 주요 과제로 남아 있습니다. 이 논문에서는 이러한 discrete representations을 학습하는 간단하면서도 강력한 generative model을 제안합니다. 우리의 model, Vector Quantised-Variational AutoEncoder (VQ-VAE)는 두 가지 핵심적인 면에서 VAE와 다릅니다.

encoder network는 continuous가 아닌 discrete codes를 출력합니다.
prior는 정적이지 않고 학습됩니다.

discrete latent representation을 학습하기 위해, 우리는 vector quantisation (VQ)의 idea를 통합합니다. VQ method를 사용하면 model이 "posterior collapse" 문제를 회피할 수 있습니다. "posterior collapse"는 powerful autoregressive decoder와 짝을 이룰 때 latents가 무시되는 현상으로, 일반적으로 VAE framework에서 관찰됩니다. 이러한 representations을 autoregressive prior와 결합함으로써, model은 고품질 image, video, speech를 생성할 수 있을 뿐만 아니라 고품질 speaker conversion 및 phonemes의 unsupervised learning을 수행하여 학습된 representations의 유용성에 대한 추가적인 증거를 제공합니다.

원래 VAE에서는 prior distribution (p(z))를 보통 고정된 표준 정규 분포로 사용합니다. 하지만 VQ-VAE에서는 이 prior distribution도 학습을 통해 업데이트

Autoregressive prior는 이전 시점의 latent variable들을 조건부로 하여 현재 시점의 latent variable의 분포를 모델링하는, 순차적으로 학습되는 prior distribution입니다.

1. Introduction

image, audio, video의 generative modelling에서의 최근 발전은 인상적인 샘플과 응용 프로그램을 산출했습니다. 동시에, few-shot learning, domain adaptation, reinforcement learning과 같은 도전적인 작업은 raw data로부터 학습된 representations에 크게 의존하지만, unsupervised 방식으로 trained된 일반적인 representations의 유용성은 여전히 지배적인 접근 방식과는 거리가 멉니다.

Maximum likelihood와 reconstruction error는 pixel domain에서 unsupervised models을 train하기 위해 사용되는 두 가지 일반적인 목표이지만, 그것들의 유용성은 features가 사용되는 특정 응용 프로그램에 따라 달라집니다. 우리의 목표는 maximum likelihood를 위해 최적화하면서 latent space에서 데이터의 중요한 features를 보존하는 model을 달성하는 것입니다. [7]의 연구에서 제안하듯이, 최고의 generative models (log-likelihood로 측정)은 latents가 없지만 powerful decoder (예: PixelCNN)를 가진 모델일 것입니다. 그러나 이 논문에서는 discrete하고 유용한 latent variables를 학습하는 것을 주장하며, 다양한 domain에서 이를 증명합니다.

continuous features를 사용한 representations 학습은 이전의 많은 연구의 초점이었지만, 우리는 discrete representations에 집중합니다. discrete representations는 우리가 관심을 갖는 많은 modality에 잠재적으로 더 자연스럽게 들어 맞습니다. Language는 본질적으로 discrete하며, speech도 일반적으로 symbols의 sequence로 표현됩니다. Image는 종종 language로 간결하게 설명될 수 있습니다. 또한, discrete representations는 복잡한 reasoning, planning 및 predictive learning (예: 비가 오면 우산을 쓸 것이다)에 자연스럽게 적합합니다. Deep learning에서 discrete latent variables를 사용하는 것은 어려운 일이지만, discrete variables에 대한 분포를 모델링하기 위해 강력한 autoregressive models이 개발되었습니다.

우리의 연구에서, 우리는 observation이 주어졌을 때 (discrete) latents의 posterior distribution에 대한 새로운 parameterisation을 통해 variational autoencoder (VAE) framework와 discrete latent representations를 성공적으로 결합하는 새로운 generative models family를 소개합니다. Vector quantization (VQ)에 의존하는 우리의 model은 train하기 간단하고, 큰 variance를 겪지 않으며, powerful decoder를 가진 많은 VAE models에서 문제가 되는 "posterior collapse" 문제를 피합니다. "posterior collapse"는 종종 latents가 무시되어 발생합니다. 또한, 이 모델은 discrete latent VAE model 중 처음으로 continuous counterparts와 유사한 성능을 얻으면서 discrete distributions의 유연성을 제공합니다. 우리는 우리의 model을 VQ-VAE라고 명명합니다.

VQ-VAE는 latent space를 효과적으로 사용할 수 있기 때문에, data space에서 일반적으로 많은 차원에 걸쳐 있는 중요한 features (예: image의 object는 많은 pixel에 걸쳐 있고, speech의 phonemes, text fragment의 message 등)를 성공적으로 모델링할 수 있습니다. 반면, 종종 local한 noise와 감지할 수 없는 detail에 집중하거나 capacity를 소비하는 것과는 대조적입니다.

마지막으로, VQ-VAE를 통해 modality의 좋은 discrete latent structure가 발견되면, 이러한 discrete random variables에 대한 powerful prior를 train하여 흥미로운 샘플과 유용한 응용 프로그램을 생성합니다. 예를 들어, speech에 대해 trained될 때, 우리는 phonemes나 words에 대한 어떠한 supervision이나 사전 지식 없이 language의 latent structure를 발견합니다. 또한, decoder에 speaker identity를 부여하여 speaker conversion, 즉 내용 변경 없이 한 speaker의 목소리를 다른 speaker에게 전달할 수 있습니다. 우리는 또한 RL을 위한 environments의 long term structure를 학습하는 데 유망한 결과를 보여줍니다.

따라서 우리의 기여는 다음과 같이 요약할 수 있습니다.

VQ-VAE model을 소개합니다. 이 모델은 간단하고, discrete latents를 사용하며, "posterior collapse"를 겪지 않고 variance 문제가 없습니다.
Discrete latent model (VQ-VAE)이 log-likelihood에서 continuous model counterparts만큼 잘 수행됨을 보여줍니다.
Powerful prior와 결합했을 때, 우리의 샘플은 speech, video generation과 같은 다양한 응용 프로그램에서 일관성 있고 high quality를 가집니다.
Supervision 없이 raw speech를 통해 language를 학습하는 증거를 제시하고, unsupervised speaker conversion의 응용을 보여줍니다.

VQ-VAE Introduction 정리 노트 (AI 연구자 대상)

핵심

Discrete Latent Representations: 기존 VAE의 continuous latent space 대신, discrete latent representations를 사용.
- 장점:
  - Language, speech 등 본질적으로 discrete한 modality에 더 적합.
  - Complex reasoning, planning, predictive learning에 유리.
  - "Posterior collapse" 문제 (powerful decoder 사용 시 latent variable 무시되는 현상) 해결.
Vector Quantization (VQ): Discrete latent representation 학습을 위해 VQ 도입.
- Encoder output을 nearest codebook vector로 quantization.
- Codebook vector (embedding)도 함께 학습.
Learned Prior: Latent variable의 prior distribution (p(z))을 고정된 표준 정규 분포 대신, 데이터로부터 학습 (autoregressive model 사용).
성능:
- Continuous latent variable을 사용하는 VAE와 동등한 성능 (log-likelihood 기준).
- Image, speech, video 등 다양한 modality에서 high-quality generation.
- Unsupervised learning of language structure from raw speech (phoneme 발견).
- Speaker conversion 등 응용 가능.

기존 VAE와의 차별점

Discrete vs. Continuous Latent: VAE는 보통 continuous latent space를 가정하지만, VQ-VAE는 discrete representation 사용.
Learned vs. Fixed Prior: VAE는 prior를 보통 고정하지만, VQ-VAE는 prior도 학습.
Posterior Collapse 방지: VQ를 사용하여 powerful decoder 사용 시에도 posterior collapse 문제 해결.

한계점(명시적으로 언급되진 않았지만 추론 가능)

Discrete representation 자체의 한계: Continuous representation만큼 미세한 표현은 어려울 수 있음.
Autoregressive prior 학습: prior를 학습하기 위해 autoregressive model (PixelCNN, Transformer)이 필요하며, 이는 계산 비용이 많이 들 수 있음.

쉬운 설명 :

VQ-VAE는 기존 VAE (Variational AutoEncoder)의 업그레이드 버전이라고 생각할 수 있습니다.

기존 VAE의 문제점:

VAE는 데이터를 압축해서 표현하는 잠재 변수(latent variable)를 사용하는데, 이 잠재 변수가 연속적인 값(continuous)을 가집니다.
이 잠재 변수의 분포를 미리 정해놓고 (보통 표준 정규 분포) 사용하는데, 이게 실제 데이터의 잠재 변수 분포와 잘 안 맞을 수 있습니다.
생성 모델(generative model)에서 이미지를 생성할 때, 너무 강력한 생성 모델(decoder)을 사용하면 잠재 변수를 무시하고 이미지를 생성하는 "posterior collapse" 문제가 발생할 수 있습니다.

VQ-VAE의 해결책:

잠재 변수를 이산적인 값(discrete)으로 표현: 연속적인 값 대신, 딱딱 끊어지는 값(예: 단어, 음소)으로 잠재 변수를 표현합니다.
Vector Quantization (VQ): 데이터를 가장 가까운 "대표 벡터" (codebook vector)로 매핑하는 방식으로 이산적인 잠재 변수를 학습합니다.
잠재 변수의 분포도 학습: 잠재 변수가 어떤 값을 가질 확률 분포(prior)를 미리 정해놓지 않고, 데이터로부터 학습합니다. (이전 잠재 변수들을 보고 다음 잠재 변수를 예측하는 방식 - autoregressive)

결과:

이미지, 음성, 비디오 등 다양한 데이터를 더 잘 생성할 수 있습니다.
음성 데이터에서 자동으로 음소를 학습하거나, 화자를 바꾸는 등 재미있는 응용도 가능합니다.

간단 비유:

기존 VAE: 그림을 그릴 때, 물감의 색깔(잠재 변수)을 연속적인 값으로 표현하고, 어떤 색을 쓸지는 미리 정해진 규칙(prior)을 따름.
VQ-VAE: 그림을 그릴 때, 정해진 색깔 팔레트(codebook)에서 색깔(잠재 변수)을 골라서 사용하고, 어떤 색을 쓸지는 그림을 그려가면서 배움(prior 학습).

2. Related Work

이 연구에서 우리는 discrete latent variables를 가진 variational autoencoders (VAE)를 training하는 새로운 방법을 제시합니다. Deep learning에서 discrete variables를 사용하는 것은, 현재 대부분의 연구에서 continuous latent variables가 지배적이라는 점(심지어 underlying modality가 본질적으로 discrete한 경우에도)에서 알 수 있듯이, 어려운 문제로 여겨져 왔습니다.

Discrete VAE를 training하기 위한 많은 대안들이 존재합니다. NVIL estimator는 variational lower bound를 최적화하기 위해 single-sample objective를 사용하고, training 속도를 높이기 위해 다양한 variance-reduction techniques를 사용합니다. VIMCO는 multi-sample objective를 최적화하며, 이는 inference network에서 여러 sample을 사용하여 수렴 속도를 더욱 가속화합니다.

최근 몇몇 저자들은 Concrete 또는 Gumbel-softmax distribution이라고 불리는 새로운 continuous reparemetrisation의 사용을 제안했습니다. 이 distribution은 continuous distribution이며, training 중에 annealing될 수 있는 temperature 상수를 가지고 있어 극한에서 discrete distribution으로 수렴합니다. Training 초기에는 gradients의 variance가 낮지만 biased되고, training이 끝날 무렵에는 variance가 높아지지만 unbiased됩니다.

그러나 위의 방법 중 어느 것도 Gaussian reparameterisation trick을 사용할 수 있는 continuous latent variables를 가진 VAE의 성능 격차를 좁히지 못합니다. Gaussian reparameterisation trick은 gradients에서 훨씬 낮은 variance를 갖는 이점이 있습니다. 게다가, 이러한 기술의 대부분은 일반적으로 MNIST와 같은 비교적 작은 datasets에서 평가되며, latent distributions의 dimensionality는 작습니다(예: 8 미만). 이 연구에서는 세 가지 복잡한 image datasets (CIFAR10, ImageNet, DeepMind Lab)와 raw speech dataset (VCTK)를 사용합니다.

우리의 연구는 또한 autoregressive distributions가 VAE의 decoder 및/또는 prior에서 사용되는 연구 라인을 확장합니다. 이는 LSTM decoders를 사용한 language modelling과, 최근에는 dilated convolutional decoders를 사용하여 수행되었습니다. PixelCNN은 convolutional autoregressive models이며, VAE의 decoder에서 distribution으로 사용되기도 했습니다.

마지막으로, 우리의 접근 방식은 neural networks를 사용한 image compression 연구와도 관련이 있습니다. Theis et. al.은 arithmetic encoding 전에 lossy image compression을 위해 activations을 압축하기 위해 scalar quantisation을 사용합니다. 다른 저자들은 vector quantisation을 사용하는 유사한 compression model을 위한 방법을 제안합니다. 그들은 시간이 지남에 따라 annealing되어 hard clustering을 얻는 vector quantisation의 continuous relaxation을 제안합니다. 그들의 실험에서, 그들은 먼저 autoencoder를 train한 다음, vector quantisation을 encoder의 activations에 적용하고, 마지막으로 전체 network를 soft-to-hard relaxation을 사용하여 작은 learning rate로 fine-tuning합니다. 우리의 실험에서 우리는 scratch부터 soft-to-hard relaxation 접근 방식을 사용하여 training할 수 없었습니다. 왜냐하면 decoder가 training 중에 항상 continuous relaxation을 invert할 수 있어서 실제 quantisation이 일어나지 않았기 때문입니다.

VQ-VAE Related Work 정리 노트 (AI 연구자 대상)

기존 Discrete VAE 연구들의 한계

Continuous Relaxation 기반 방법들의 문제점 (Gumbel-Softmax 등):
- Training 초반: Gradients variance는 낮지만 biased.
- Training 후반: Variance는 높아지지만 unbiased.
- Gaussian reparameterization trick을 쓰는 continuous latent variable VAE만큼 성능이 안 나옴.
- 주로 작은 datasets (MNIST) 및 낮은 차원의 latent space에서 평가됨.
Soft-to-Hard Vector Quantization의 문제점:
- Continuous relaxation을 사용하지만, decoder가 relaxation을 invert시켜 quantization 효과가 사라지는 문제 발생 (본 논문 저자들의 실험에서 확인).

VQ-VAE 연구의 차별성 (기존 연구 대비)

더 복잡한 Datasets 사용: CIFAR10, ImageNet, DeepMind Lab, VCTK (raw speech) 등.
Gaussian Reparameterization Trick을 쓰는 Continuous VAE와 대등한 성능: Discrete latent variable을 쓰면서도 성능 저하 없음.
Autoregressive Model 활용: Decoder와 prior 모두에 autoregressive model (PixelCNN 등) 활용 연구 확장.

연관 연구

Neural Network Image Compression: Lossy image compression을 위한 scalar/vector quantization 연구들과 관련 있음. (Theis et al., 다른 저자들)

쉬운 설명:

VQ-VAE 이전에도 discrete latent variable을 사용하는 VAE 연구들이 있었습니다. 하지만 대부분 다음과 같은 문제점들이 있었습니다.

Continuous Relaxation의 한계: Discrete variable을 사용하기 위해 continuous approximation (Gumbel-Softmax 등)을 사용했는데, 이게 학습이 불안정하고 성능도 continuous latent variable을 쓰는 VAE보다 좋지 않았습니다.
Soft-to-Hard Vector Quantization의 실패: Vector quantization을 사용하되, 처음에는 "soft"하게 (뭉뚱그려서) quantization 했다가 점점 "hard"하게 (정확하게) 만드는 방법을 시도했지만, decoder가 "soft"한 상태를 원래대로 되돌려 버려서 quantization 효과가 없었습니다.

VQ-VAE는 이러한 문제들을 해결하고, 다음과 같은 특징을 가집니다.

더 크고 복잡한 데이터 (이미지, 음성)에서 실험했습니다.
Continuous latent variable을 쓰는 VAE와 성능이 비슷합니다.
이미지 생성 등에 사용되는 autoregressive model (PixelCNN 등)을 VAE의 decoder와 prior에 잘 활용했습니다.

결론적으로, VQ-VAE는 discrete latent variable을 사용하는 VAE의 새로운 가능성을 제시한 연구입니다.

3 VQ-VAE

아마도 우리 접근 방식과 가장 관련이 있는 연구는 VAE일 것입니다. VAE는 다음 부분으로 구성됩니다.

Encoder network: 입력 데이터 x가 주어졌을 때 discrete latent random variables z의 posterior distribution q(z|x)를 parameterise합니다.
Prior distribution p(z).
Decoder: 입력 데이터에 대한 distribution p(x|z)를 가집니다.

일반적으로 VAE의 posteriors와 priors는 diagonal covariance를 갖는 정규 분포로 가정되며, 이는 Gaussian reparametrisation trick을 사용할 수 있게 합니다. 확장에는 autoregressive prior 및 posterior models, normalising flows, inverse autoregressive posteriors가 포함됩니다.

이 연구에서 우리는 vector quantisation (VQ)에서 영감을 받은 새로운 training 방법을 사용하여 discrete latent variables를 사용하는 VQ-VAE를 소개합니다. Posterior 및 prior distributions는 categorical이며, 이러한 distributions에서 추출된 samples은 embedding table을 index합니다. 그런 다음 이러한 embeddings은 decoder network에 대한 input으로 사용됩니다.

3.1 Discrete Latent Variables

우리는 latent embedding space e ∈ R^(K×D)를 정의합니다. 여기서 K는 discrete latent space의 크기 (즉, K-way categorical)이고, D는 각 latent embedding vector e_i의 dimensionality입니다. K개의 embedding vectors e_i ∈ R^D, i ∈ {1, 2, ..., K}가 있습니다.

Figure 1에서 볼 수 있듯이, model은 입력 x를 받아 encoder를 통과시켜 output z_e(x)를 생성합니다. 그런 다음 discrete latent variables z는 equation 1과 같이 공유 embedding space e를 사용한 nearest neighbour look-up으로 계산됩니다. Decoder에 대한 input은 equation 2에 주어진 해당 embedding vector e_k입니다.

이 forward computation pipeline은 latents를 1-of-K embedding vectors로 mapping하는 특정한 non-linearity를 가진 regular autoencoder로 볼 수 있습니다. Model의 전체 parameters 집합은 encoder, decoder 및 embedding space e의 parameters의 합집합입니다. 간단하게 하기 위해 이 Section에서는 단일 random variable z를 사용하여 discrete latent variables를 나타내지만, speech, image 및 video의 경우 실제로 각각 1D, 2D 및 3D latent feature spaces를 추출합니다.

Posterior categorical distribution q(z|x) probabilities는 다음과 같이 one-hot으로 정의됩니다.

q(z = k|x) = {
    1,  if k = argmin_j ||z_e(x) - e_j||_2
    0,  otherwise
}

(equation 1)

여기서 z_e(x)는 encoder network의 output입니다. 우리는 이 model을 ELBO를 사용하여 log p(x)를 bound할 수 있는 VAE로 간주합니다. 우리의 proposal distribution q(z = k|x)는 deterministic하며, z에 대해 simple uniform prior를 정의함으로써 KL divergence는 상수가 되고 log K와 같아집니다.

Representation z_e(x)는 discretisation bottleneck을 통과한 후 equation 1과 2에 주어진 대로 embedding e의 가장 가까운 element에 mapping됩니다.

z_q(x) = e_k,  where k = argmin_j ||z_e(x) - e_j||_2

(equation 2)

3.2 Learning

Equation 2에는 실제 gradient가 정의되어 있지 않지만, straight-through estimator와 유사하게 gradient를 근사화하고 decoder input z_q(x)에서 encoder output z_e(x)로 gradients를 복사합니다. Quantisation operation을 통해 subgradient를 사용할 수도 있지만, 이 간단한 estimator가 이 논문의 초기 실험에서 잘 작동했습니다.

Forward computation 동안에는 가장 가까운 embedding z_q(x) (equation 2)가 decoder로 전달되고, backward pass 동안에는 gradient ∇_zL이 변경되지 않고 encoder로 전달됩니다. Encoder의 output representation과 decoder의 input은 동일한 D 차원 공간을 공유하기 때문에, gradients는 encoder가 reconstruction loss를 낮추기 위해 output을 어떻게 변경해야 하는지에 대한 유용한 정보를 포함합니다.

Figure 1 (오른쪽)에서 볼 수 있듯이, gradient는 다음 forward pass에서 encoder의 output이 다르게 discretise되도록 밀어낼 수 있습니다. 왜냐하면 equation 1의 할당이 달라지기 때문입니다.

Equation 3은 전체 loss function을 지정합니다. 이 loss function은 VQ-VAE의 다른 부분을 train하는 데 사용되는 세 가지 components를 가집니다.

Reconstruction Loss (Data Term): Decoder와 encoder를 최적화합니다 (위에서 설명한 estimator를 통해).
VQ Objective: Embedding space를 학습하기 위해 사용됩니다. l2 error를 사용하여 embedding vectors e_i를 encoder outputs z_e(x)쪽으로 이동시킵니다(equation 3의 두 번째 항). 이 loss 항은 dictionary를 업데이트하는 데만 사용되므로, z_e(x)의 moving averages의 함수로 dictionary items를 업데이트할 수도 있습니다 (이 연구의 실험에서는 사용되지 않음). 자세한 내용은 Appendix A.1을 참조하십시오.
Commitment Loss: Embedding space의 volume이 dimensionless이기 때문에, embeddings e_i가 encoder parameters만큼 빠르게 train되지 않으면 임의로 커질 수 있습니다. Encoder가 embedding에 commit하고 output이 커지지 않도록 하기 위해 commitment loss를 추가합니다 (equation 3의 세 번째 항).

따라서 전체 training objective는 다음과 같습니다.

L = log p(x|z_q(x)) + ||sg[z_e(x)] - e_k||_2^2 + β||z_e(x) - sg[e]||_2^2

(equation 3)

여기서 sg는 stopgradient operator를 나타내며, forward computation 시간에는 identity로 정의되고 partial derivatives가 0이므로, 사실상 피연산자가 업데이트되지 않는 상수가 되도록 제한합니다.

Decoder는 첫 번째 loss 항만 최적화하고, encoder는 첫 번째와 마지막 loss 항을 최적화하며, embeddings은 중간 loss 항에 의해 최적화됩니다. 결과 algorithm은 β 값(0.1에서 2.0 사이)에 대해 결과가 변하지 않아 β에 대해 상당히 robust하다는 것을 발견했습니다. 우리는 모든 실험에서 β = 0.25를 사용하지만, 일반적으로 이것은 reconstruction loss의 scale에 따라 달라집니다.

z에 대해 uniform prior를 가정하기 때문에 ELBO에 일반적으로 나타나는 KL 항은 encoder parameters에 대해 일정하므로 training에서 무시할 수 있습니다.

실험에서 N개의 discrete latents를 정의합니다 (예: ImageNet의 경우 32 x 32 latents, CIFAR10의 경우 8 x 8 x 10). 결과 loss L은 k-means 및 commitment loss에 대해 N개의 항에 대한 평균을 얻는다는 점을 제외하고는 동일합니다(각 latent에 대해 하나씩).

Complete model의 log-likelihood log p(x)는 다음과 같이 평가할 수 있습니다.

log p(x) = log Σ_k p(x|z_k)p(z_k)

Decoder p(x|z)는 MAP-inference에서 z = z_q(x)로 trained되기 때문에, decoder는 완전히 수렴되면 z != z_q(x)에 대해 p(x|z)에 probability mass를 할당해서는 안 됩니다. 따라서 log p(x) ≈ log p(x|z_q(x))p(z_q(x))로 쓸 수 있습니다. 우리는 이 근사치를 section 4에서 경험적으로 평가합니다. Jensen의 부등식에서 log p(x) ≥ log p(x|z_q(x))p(z_q(x))라고 쓸 수도 있습니다.

3.3 Prior

Discrete latents p(z)에 대한 prior distribution은 categorical distribution이며, feature map의 다른 z에 의존하게 함으로써 autoregressive하게 만들 수 있습니다. VQ-VAE를 training하는 동안 prior는 일정하고 uniform하게 유지됩니다. Training 후, z에 대한 autoregressive distribution p(z)를 fitting하여 ancestral sampling을 통해 x를 생성할 수 있습니다. Image에는 discrete latents에 대해 PixelCNN을 사용하고, raw audio에는 WaveNet을 사용합니다. Prior와 VQ-VAE를 jointly training하는 것은 향후 연구로 남겨두며, 이는 결과를 강화할 수 있습니다.

3. VQ-VAE 정리 노트 (AI 연구자 대상)

핵심 구조

Encoder: 입력 x를 받아 z_e(x) (continuous) 출력.
Vector Quantization (VQ): z_e(x)를 embedding space e에서 가장 가까운 embedding vector e_k로 mapping (discrete latent variable z_q(x) 생성).
Decoder: z_q(x) ( = e_k)를 입력받아 x reconstruction.
Learned Prior: Discrete latent variable z에 대한 prior distribution p(z)를 autoregressive model (PixelCNN, WaveNet)로 학습.

학습 방법

Loss Function:
1. Reconstruction Loss: Decoder와 encoder 학습 (standard autoencoder loss).
2. VQ Loss: Embedding vectors e_i를 encoder outputs z_e(x) 쪽으로 이동 (dictionary learning).
3. Commitment Loss: Encoder가 특정 embedding에 "commit"하도록 유도 (embedding space volume exploding 방지).
4. ```
L = log p(x|z_q(x)) + ||sg[z_e(x)] - e_k||_2^2 + β||z_e(x) - sg[e]||_2^2
```
Gradient Estimation: z_e(x)에서 z_q(x)로의 mapping은 미분 불가능하지만, "straight-through estimator" 사용 (decoder input의 gradient를 encoder output으로 복사).
Prior Learning: VQ-VAE 학습 후, latent space에 autoregressive prior (PixelCNN, WaveNet) 별도 학습.

VQ-VAE의 특징 (요약)

Discrete Latent Variable: Continuous 대신 discrete latent variable 사용.
Vector Quantization: Nearest neighbor look-up으로 discrete representation 학습.
Learned Prior: Autoregressive model로 prior distribution 학습.
"Straight-Through" Gradient: 미분 불가능한 VQ 부분을 gradient "복사"로 해결.
세 가지 Loss Term: Reconstruction, VQ, commitment loss로 구성된 loss function.

쉬운 설명:

VQ-VAE는 이미지를 압축해서 표현하는 방법을 배우는 모델입니다. 기존 VAE와 다른 점은, 이미지를 연속적인 숫자로 표현하는 대신, "단어"처럼 딱딱 떨어지는 이산적인(discrete) 표현으로 압축한다는 것입니다.

VQ-VAE의 구성 요소:

Encoder: 이미지를 입력받아 "특징 벡터" (z_e(x))를 만듭니다. (이 벡터는 아직 연속적인 값입니다.)
Codebook (Embedding Space): 여러 개의 "대표 벡터" (e_i)들을 가지고 있습니다. (마치 단어 사전처럼)
Vector Quantization: Encoder가 만든 "특징 벡터"를 codebook에서 가장 가까운 "대표 벡터"로 바꿉니다. (이것이 이산적인 표현 z_q(x)가 됩니다.)
Decoder: 이산적인 표현 (z_q(x))을 받아서 다시 이미지를 복원합니다.
Prior: 이산적인 표현들이 어떤 순서로 나타나는것이 자연스러운지를 학습합니다. (e.g. PixelCNN)

학습 방법:

Loss Function (세 가지):
1. Reconstruction Loss: 복원된 이미지가 원래 이미지와 얼마나 비슷한지 측정.
2. VQ Loss: "특징 벡터"와 "대표 벡터" 사이의 거리를 줄임. (codebook 학습)
3. Commitment Loss: Encoder가 "대표 벡터"를 골고루 사용하도록 유도.
Gradient Estimation: "특징 벡터"를 "대표 벡터"로 바꾸는 과정은 미분이 불가능하지만, "straight-through estimator"라는 트릭을 써서 학습합니다. (decoder의 gradient를 encoder로 그냥 복사)

핵심 아이디어:

이미지를 연속적인 숫자가 아닌, "단어"와 같은 이산적인 표현으로 압축하고, 그 "단어"들의 분포(prior)도 학습해서 더 강력한 생성 모델을 만드는 것입니다.

1. 이미지를 인코딩함 (학습되는 부분)

2. 코드북 초기화 (코드북도 학습되는 부분)

3. 이미지의 피쳐를 코드북에서 가까운 피쳐로 변환시킴

4. 그게 z인데 z를 받아서 디코더에서 이미지를 다시 복원 (학습되는 부분)

5. 이미지와 정답간의 loss를 통해

업데이트가 이루어짐.