[NIPS 2017] Neural Discrete Representation Learning

오늘은 평소 리뷰했던 논문과는 다소 다른 분야의 논문을 리뷰하려고 합니다. 최근 GPT를 사용하며 그림을 그려달라는 요청을 할 때, 잘 워킹하는 것을 보고 이게 어떤 원리지? 라는 호기심이 생겨서 찾아보게 되었습니다.

상단 이미지처럼, GPT에게 Text로 어떤 그림을 그려달라고 요청하면, 잘 알아듣고 맥락에 맞는 그림을 잘 그려내는 것을 확인할 수 있습니다.

뿐만아니라, 왼쪽 그림같이 그림의 일정 영역을 선택(돼지 그림 아래에 꽃들을 중심으로 파란색으로 칠해진 부분이 선택된 영역)한 뒤 그 부분에 대해서 추가적인 명령을 할 수 있습니다. 상단 이미지는 제가 배경을 선택한 뒤 “돼지가 파도를 타고있는 거로 바꿔줘” 라고 요청했을 때의 캡처본입니다.

그 결과 파도를 타고있는 돼지로 배경이 잘 교체되어 생성된 것을 확인할 수 있습니다. 이 과정에서 어떻게 문자-이미지 사이의 변환이 잘 동작하는지, 이와 관련된 연구는 어떤 것이 있는지 궁금해졌는데요. 이와 관련되어 찾아보니.. 가장 기본이 되는 연구인 VQ-VAE 라는 논문을 알게되었습니다. 해당 연구는 DALL-E의 핵심 구성요소 중 하나로 Discrete latent space의 시초라고 할 수 있는 연구입니다.

1. Background: Discrete Latent Space

오토인코더, VAE 나 GAN 같은 생성 모델들은 학습 과정에서 latent space 라고 불리는 데이터의 내재된 특성을 포착하는 잠재 공간을 만들어냅니다. 일반적으로 latent space는 고차원의 입력을 낮은 차원의 표현으로 압축하는 데, 잠재공간이 잘 생성되었다면 이미지의 특정 속성과 일치하게 됩니다. 이게 무슨 뜻이냐면, 예를 들면 돼지 이미지를 압축할 때 한 축은 돼지의 색깔을, 다른 축은 이미지의 전체적인 스타일을 나타내는 저차원 공간으로 압축되게 됩니다.

그러나 실제로는 이렇게 특정 축이 이미지의 특성과 일치하지 않을 뿐더러, 이를 찾아내는 것도 어려운 일이라고 합니다. VAE를 통해 생성된 latent space는 인간이 해석하는 이미지의 특성 혹은 스타일과는 사뭇 다르게 설계되곤 합니다. 또한 이 잠재 공간은 일반적으로 고정된 차원의 연속적인(Continuous) 표현 공간으로 나타나는데, 이러한 표현은 해석하기 어렵고 효율성이 떨어지는 문제가 있습니다.

특히 리뷰 시작부분에서 보여드린 GPT 예시는 언어라는 이산적인 요소를 이미지의 특성과 결합한 사례인데.. 이처럼 “단어” 단위로 latent space를 구성한 다음, 이런 단어들의 조합으로 latent space를 다시 이미지로 디코딩할 수는 없을까요?

다들 아시다시피, 트랜스포머 계열의 모델은 자연어 처리에서 실제 단어를 생성하는 것처럼 잠재 공간의 “단어”를 생성합니다. 따라서 이런 모델들은 기존 VAE 모델에서 사용하는 Continuous latent space 에 적용하는 것은 오히려 어려운 일 같아 보입니다. 따라서 이미지에서도 discrete latent space를 사용하기 위한 연구가 등장하였습니다. 즉, 연구자들은 잠재 공간을 이산적인(Discrete) 표현형 공간으로 나타내는 방법을 탐구하였습니다. 이산적인 표현은 더 명확하고 구체적인 데이터 표현을 가능하게 하지만, 그러나 이 방식은 최적화 과정에서 그레디언트 계산하는 것이 어렵다는 단점이 있습니다.

그럼에도 불구하고, Discrete latent space는 설명 가능한 인공지능(XAI) 개발의 중요한 요소가 될 수 있음에 분명합니다. (최근과 같이 모델이 엄청 커지는 상황에서 경량화에도 유리하고, 일반화 성능까지 continuous에 뒤지지 않는다는 결과까지 보이기도 했습니다.)

따라서 제가 이번에 리뷰하려는 논문은 이런 문제를 해결하고 현대 AI의 기초가 된 VQ-VAE라는 논문을 리뷰하도록 하겠습니다.

---

리뷰하려는 VQ-VAE는 Vector Quantization 이라는 기법으로 latent space를 discrete 하게 만들었습니다.

• Title: [NIPS 2017] Neural Discrete Representation Learning
• Code: [Not official?] GitHub

2. Introduction

Image, audio, video 등을 생성하는 분야는 꾸준히 발전하고 있습니다. 이와 관련하여 Pixel domain에서 가장 효과적인 모델은 PixelCNN 이지만, 이는 연속적인 표현을 사용하고 있습니다. 그러나 현실 세계의 많은 데이터들은 Discrete하게 구성되어 있습니다. 일단 언어는 본질적으로 discrete 합니다. 그리고 이미지 역시 언어로 표현될 수 있습니다. 가령 image captioning 이 대표적인 예시겠죠. 이처럼 저자는 VAE와 이산표현을 결합한 새로운 생성 모델인 VQ-VAE 제안하였습니다. Vector Quanization(VQ)를 사용함으로써 분산이 너무 커서 생기는 어려움*과 posterior collapse** 문제를 회피할 수 있었다고 합니다. 그러면서, 연속 표현을 사용하는 모델과 비등한 성능을 발생할 수 있었습니다.
*discrete variable의 분산이 커서 학습이 느려지는 경우가 있음

**Posterior collapse를 알기 위해서는 VAE를 알아야 하는데 여기서 이에 대해 자세하게 다루다간 리뷰가 너무 길어질 것 같아 간단하게만 서술하겠습니다.

• 앞서 설명드린 것처럼 VAE는 입력 데이터 x를 받아 내부적으로 잠재 변수 z로 변환하는 인코더, 그리고 이 잠재 변수를 다시 원래의 데이터와 유사한 데이터 \hat{x} 로 복원하는 디코더로 구성되어 있습니다. 여기서 잠재 변수 z의 분포는 “prior” 이라는 사전 분포를 따른다고 가정합니다. 이를 통해 latent space에서의 데이터 복잡성을 간결하게 표현할 수 있으며, 모델은 이 사전분포를 기반으로 실제 데이터의 잠재 표현을 학습하게 됩니다.
• “Posterior collapse”는 VAE에서 특히 잠재 변수 z가 데이터 x의 유의미한 특성을 파악하는 데 실패하는 현상입니다. 자세하게 말하자면, 인코더가 생성하는 잠재 변수 z의 분포가 prior와 구분이 거의 되지 않는 상태를 말합니다. 이 경우, 잠재 변수 z가 모델의 재구성이나 데이터 생성 과정에서 별다른 역할을 하지 못하게 됩니다.

물론 지금 리뷰하는 논문이 discrete latent 를 학습하기 위한 최초의 논문은 아닙니다. 당시 latent space는 continuous하거나 discrete 했는데, 앞서 언급된 장점에도 불구하고 discrete latent에서는 성능이 잘 나오지 않았기에, 저자는 WaveNet을 참고하여 discrete latent와 autoregressive model을 활용하였다고 합니다.

3. VQ-VAE

본 연구에서는 VAE 활용한 연구를 다룹니다. VAE는 입력 데이터 x에 대하여 잠재 변수 z의 q(z|x)을 학습하는 encoder, 잠재변수에 대한 사전분포(prior) p(z), 그리고 입력 데이터에 대한 분포 p(x|z)를 학습하는 decoder로 이루어져 있습니다. 또한 일반적으로 사후 분포와 사전 분포를 diagonal covariance를 가진 정규 분포로 가정하며, 이를 통해 가우시안 재매개변수화 기법을 사용할 수 있습니다.

3.1 Discrete Latent variables

학습은 encoder, decoder, embedding space에서 이루어지게 됩니다. 위의 그림에서 Latent embedding space e∈R^{K×D} (K는 discrete size)가 이산표현을 나타내며, 이를 codebook이라고 합니다. 또한 K는 이산 표현 공간의 크기, D는 각 Embedding vector e_i의 차원입니다.

먼저, encoder에서 입력 이미지 x로부터 z_e​(x)를 출력합니다. 그 다음 식(1)을 통해 z_e​(x)와 embedding space e를 사용하여 최근접이웃 look-up 방식으로 embedding vector를 찾게됩니다. 그 결과, dictionary q(z|x)를 생성하고, 식(2) 방식으로 q(x|z)에 e를 mapping하여 z_q(x)를 재구성합니다.

마지막으로 디코더는 z_q(x)로부터 p(x|z_q)를 출력합니다. 그래서 저자들은 이 모델을 log(p(x))를 likelihood로 설정하고, ELBO 형태로 최적화하는 VAE라고 주장합니다. 또한 제안하는 분포 z를 uniform distribution으로 정의하여 D_{KL}​ term이 상수가 됩니다.
(VAE에서는 z를 gaussian distribution으로 정의)

지금까지의 내용을 보다 컨셉적으로 쉽게 다시한번 설명을 해보도록 하겠습니다. 위의 그림이 바로 Codebook 과 VQ-VQE의 예시입니다. 우선 인코더를 통해 이미지 x에 대한 Enc(x)가 출력되었다고 가정해보겠습니다. 이 때, 인코더는 하나의 출력 벡터가 아니라 벡터의 행렬로 구성된 3차원 텐서를 생성하게 된다고 합니다.

이후 이 잠재 벡터에 대해 잠재 공간에서 가장 가까운 코드북 벡터를 찾아서, 이를 해당 잠재 벡터로 대체합니다. 결과적으로 latent space에서는 코드북 벡터들만을 포함합니다. 위의 그림을 보면 codebook에 해당하는 7, 1, 5, 3의 벡터로 대체된 것을 볼 수 있죠. 그 다음 디코더는 인코더가 생성했던 것과 동일한 차원을 가진 벡터 텐서를 입력으로 받지만, 이제 잠재 코드는 이산적이라서 가능한 경우의 수는 유한 개로 정해지게 됩니다.

따라서 VQ-VAE는 연속적인 잠재 공간을 이산적인 코드북 벡터로 표현함으로써, 효율적인 데이터 압축 및 복원이 가능하게 되었습니다.

그러나 내용을 읽다보면 의문점이 몇 가지 생기실 건데요, 첫째, 인코더와 디코더 네트워크를 어떻게 학습할 것인지에 대한 것입니다. VAE와 유사하게, VQ-VAE는 가장 가까운 이웃을 취하는 상수 작업을 수행하는데, 이 연산은 그래디언트가 역전파되지 않는 문제가 발생하게 됩니다. 둘째, 코드북 벡터를 어떻게 학습할 것인가? 다음 파트에서 이에 대한 설명을 다뤄보겠습니다.

3.2 Learning

우선 첫번째 문제: 그라디언트 역전파 문제에 대해 다뤄보겠습니다. discrete latent 로 표현하는 과정은 역전파 계산이 불가능하여 학습에도 어려움이 따르게 됩니다. 이에 VQ-VAE는 디코더에서 인코더로 그래디언트를 복사(상단 그림에서 빨간 색으로 표시)하였습니다. 즉, 코드북 벡터에 대한 loss의 그래디언트를 Enc(x)에 대한 그래디언트로 사용합니다. 이는 근사를 취하는 방식으로, 저자들은 이를 통해 gradient 문제를 해소하고자 하였습니다.

두 번째 문제인 VQ-VAE의 학습 방식에 대해 설명하겠습니다. 이를 위한 loss function 은 다음과 같습니다:

굉장히 복잡해서 하나하나 설명드리겠습니다. VQ-VAE 모델의 Loss는 세 가지 항목으로 구성됩니다. 첫 번째 항은 재구성 loss(reconstruction loss)로, 모델이 입력 데이터를 얼마나 잘 복원하는지를 측정합니다. 두 번째 항은 코드북 loss(codebook loss)인데, 이는 인코더의 출력인 z_e​(x)와 코드북 내에서 가장 가까운 벡터 e_i​가 z_e(x)와 더 가까워지도록 유도합니다. 세 번째 항은 코드북 commitment loss(codebook commitment loss)로, 코드북 벡터가 인코더 출력에 더 잘 맞도록 만드는 역할을 합니다. 이 때, sg는 “stop gradient”를 의미하는데, 말 그래도 특정 항에 대해 가중치 업데이트를 멈추는 역할을 합니다.

Reconstruction Loss. 말 그대로 디코더와 인코더를 최적화하는 데 사용됩니다. 재구성 손실은 원본 데이터 x와 디코더를 통해 재구성된 데이터 사이의 차이를 최소화하게 됩니다.

• 바로 직전에 설명한 미분 불가능에 대한 상세 설명입니다. 인코더 출력 <em>z_e</em>(<em>x</em>)와 디코더 입력 z_q(x) 사이에서 비선형적이고 미분 불가능한 ‘argmin’ 연산으로 가장 가까운 코드북 벡터를 선택하는 과정을 수행하게 됩니다.
• 그러나 이는 그래디언트를 직접 계산할 수 없기 때문에, 이를 해결하고자 ‘straight-through estimator’와 유사한 방법으로 그래디언트를 직접 복사하는 방식을 사용하였다고 합니다.

Codebook Loss. 인코더 출력과 가장 가까운 코드북 벡터 e_i 사이의 L2 error를 최소화하는 데 사용됩니다. 이를 통해 코드북 벡터들을 적절하게 학습하여 인코더 출력에 더 잘 맞도록 설계할 수 있습니다.

Commitment loss. 이 loss는 인코더가 특정 코드북 벡터에 “약속”을 하도록 강제하는 역할을 합니다. 이는 인코더 출력이 너무 변동성이 크거나 확장되지 않도록 제한하는 데 도움을 줄 수 있다고 합니다.

• Embedding vector는 무한하기 때문에, codebook loss에서 e_i​는 encoder vector만큼 빠르게 학습되기 어려우며 임의로 커질 수도 있습니다.
• 따라서 이 loss를 encoder vector가 embedding vector와 유사해질 수 있도록 강제하는 텀이 추가되었습니다.

3.3 Prior

저자는 Discrete latents variable p(z)에 대한 사전 분포로는 categorical distribution를 사용하였으며, feautre map 내 다른 z 값에 따라 autoregressive하게 구성될 수 있습니다. VQ-VAE를 학습하는 동안에는 사전을 일정하고 균일하게 유지합니다. 또한 이미지에 대해서는 PixelCNN을, raw 오디오에 대해서는 WaveNet을 사용하였다고 합니다.

4. Experiments

여기까지가 VQ-VAE 설명입니다. 다만, 이 연구가 트랜스포머가 나오기 전에 수행되었기 때문에 트랜스포머를 잠재 코드 생성에 사용할 수 없었지만, 실험에는 당시 최고성능을 내던 PixelCNN을 사용하였습니다. (PixelCNN은 이미지를 생성하는 모델로 시작되었지만, 고해상도 이미지를 픽셀 단위로 자동 회귀적으로 생성하는 게 너무 느리다는 단점이 있다고 합니다)

4.1 Comparison with continuous variables

저자는 성능 평가를 위해 VQ-VQE, VAE(continuous variables), VIMCO(independent Gaussian or categorical priors)를 비교하였습니다. 모든 모델은 CIFAR10 데이터셋을 사용하여 동일한 VAE 구조로 학습되고, 잠재 변수의 수(연속 또는 이산)와 이산 공간의 차원 K를 조정해가며 실험을 진행하였다고 합니다. 즉, 실험을 통해 각 모델의 성능과 잠재 공간 구성이 데이터 재구성 및 생성에 미치는 영향을 비교 분석하고자 하였습니다.

그 결과 discrete latent space인 VQ-VAE는 continuous latent space를 사용했을 때와 비슷한 성능을 보ㅇ였습니다.

4.2 Images

우선 저자는 실험을 통해 이미지는 픽셀 수준에서 많은 중복 정보를 포함하고, 대부분의 픽셀이 상관관계를 가지며 노이즈가 많기 때문에, 픽셀 수준에서 모델을 학습하는 것이 낭비일 수 있음을 보여주고자 하였습니다. 따라서 128×128×3 이미지를 32×32×1 크기의 이산 공간으로 압축하여p(x|z)의 purely deconvolutional을 통해 모델링할 수 있음을 보여주었습니다. (비트 단위로 약 42.6배의 정보량 감소를 의미한다고 합니다).

저자는 당시 가장 좋은 성능을 내던 PixelCNN을 학습함으로써 제안하는 방법론을 실험하였습니다. 상단 이미지가 바로 VQ-VAE로부터 생성된 이미지입니다. (즉, PixelCNN을 32 × 32 × 1 discrete latent space에서 학습시킨 다음 생성한 이미지임)

동일한 실험을 DeepMind Lab 환경에서 추출한 84x84x3 프레임에 대해서도 진행하였고, 만들어진 이미지는 원본과 거의 동일하게 보였다고 합니다. 21x21x1 잠재 공간에서 학습된 PixelCNN 사전에서 추출된 샘플은 디코더를 사용하여 픽셀 공간으로 디코딩되며, 상단 그림과 같습니다.

마지막으로, 첫 번째 VQ-VAE에서 얻은 DM-LAB 프레임의 21x21x1 잠재 공간 상에서 PixelCNN 디코더가 있는 두 번째 VQ-VAE를 학습하였습니다. 이 설정은 일반적으로 VAE에서 ‘posterior collapse’ 문제를 유발하지만, 저자가 제안하는 모델에서는 이를 피할 수 있게 되었습니다. 두 번째 단계에서 전체 이미지를 모델링하기 위해 단 3개의 잠재 변수만을 사용하기 때문에 존채 but눈 3*9=27 bit로 float32보다 적은 양으로 압축하는 결과 보이게 됩니다. 그 결과는 상단 그림에서 확인할 수 있습니다.

---

(리뷰를 다 작성하고 보니 모든 예시가 돼지군요ㅋㅋ.. 오랜만에 직접 그림을 그려가며 열심히 리뷰를 작성한 보상으로 삼겹살을 먹어야겠습니다 ㅎㅋ )

이런 분야가 있다는 것을 또 새삼스럽게 알게되었습니다. 다소 예전 방법론임에도 현재까지 토대가 될만큼 역시 simple is best임을 깨닫게 되는 논문인 듯 합ㄴ디ㅏ.