Understanding VQ-VAE (DALL-E Explained Pt. 1)

EunYeong Park·2021년 11월 14일
2

생성모델

목록 보기
1/1
post-thumbnail

이 글은 다음 포스트(https://ml.berkeley.edu/blog/posts/vq-vae/)를 번역하고 정리한 내용입니다.

The Basics

Latent Spaces

잠재공간은 원시 데이터의 주어진 분포에 대한 숨겨진 표현

원시 데이터 x들이 n차원의 실수 범위에 있을 때, 원시 데이터들은 저차원의 z의 선형변환으로부터 noisily하게 생성된다고 생각해보자.

  • xRn,zRm (m<n)x \in R^n, z \in R^m~(m<n)
  • x=Az+vx=Az+v , vv는 n차원의 독립적으로 동등하게 분포되어있는 가우시안 노이즈, ARn×mA \in R^{n\times m}

일반적으로 우리는 원시 데이터 x만 볼 수 있고, z접근할 수 없다. (z를 데이터의 잠재적 or 숨겨진 표현이라고 하는 이유)

원시 데이터가 잠재 공간의 선형 변환이라는 이 특정 설정은 "PCA"라는 고전적인 비지도 알고리즘이 다음을 위해서 설계된 것과 정확히 같다. PCA는 기본적으로 근본적인 z 표현을 찾으려고 한다.

잠재 표현이 원시 데이터와 더 복잡하고 비선형적인 관계를 가진다면?
예를 들어 아래 그림은 더 복잡한 잠재 공간이 인코딩할 수 있는 고급 정보의 종류를 시각화 한 것이다.

이 경우 PCA는 최고의 잠재 표현을 찾지 못한다. 대신 우리는 이보다 추상적인 잠재 공간을 찾기 위해 오토 인코더를 사용할 수 있다.

참고: 잠재 공간은 연속 벡터 공간일 필요가 없으며 대신 VQ-VAE가 찾으려는 잠재 공간의 유형인 이산 변수 집합일 수 있다. 그전에 먼저 바닐라 자동 인코더를 이해하자!

Autoencoders

오토 인코더는 신경망을 사용하여 주어진 데이터 분포에 대한 비선형 잠재 표현을 찾는 비지도 학습 기술이다.

신경망은 두 부분으로 구성된다

  • 인코더 네트워크 z=f(x)z=f(x)
  • 디코더 네트워크 x^=g(z)\hat{x}=g(z)
    • xx는 입력 데이터, zz는 잠재 벡터 표현, x^\hat{x}는 잠재 공간으로부터 재구성된(reconstruct) xx
    • f(x),g(z)f(x),g(z)는 신경망
    • 두 부분을 합친 전체 모델은 x^=g(f(x))\hat{x}=g(f(x)) 관계로 표현 가능

이상적으로 디코더는 인코더의 잠재 표현에서 원시 데이터를 정확하게 재구성할 수 있어야 한다.

모델이 이러한 재구성을 학습할 수 있다면 잠재 공간이 데이터를 잘 나타낸다고 가정할 수 있다.

이 목적을 실현하기 위해 우리는 xxx^\hat{x} 사이에 재구성 손실로 모델을 훈련한다 (즉, xx^22||x-\hat{x}||^2_2 or log(P(xz))log(P(x|z)))

z는 항상 x보다 차원이 낮아야 한다는 점도 중요하다.

요점은 알고리즘이 원시 데이터의 가장 필수적인 구성 요소를 찾도록 강제로 데이터의 압축된 표현을 인코딩 하는 것이다.

오토 인코더의 이 압축된 잠재 부분은 데이터를 훨씬 더 작은 공간으로 압축하기 때문에 네트워크의 병목 현상이라고 부르기도 한다

Variational Autoencoders (VAE)

(해당 블로그에서 VAE의 수식적인 세부사항은 다루지 않으므로 , 관련내용이 잘 정리된 블로그를 추천한다)

오토 인코더의 핵심을 이해했으므로 Variational Autoencoder가 다른 점을 살펴보자.

주요 차이점은 VAE의 잠재 공간에 부과되는 구조에서 비롯된다.

이상적으로는 우리는 잠재 공간이 의미적으로 유사한 데이터 포인트를 서로 옆에 묶고 의미적으로 유사하지 않은 포인트를 멀리 떨어뜨려 배치하기를 원한다.

가급적이면 데이터 분포의 대부분은 잠재 공간에서 상당히 컴팩트한 볼륨을 구성하고 무한대로 확장되지 않아야 한다.

바닐라 오토 인코더의 문제는 학습된 잠재성이 이런 속성들을 반드시 가지고 있지는 않다는 점이다. 모델은 원하는 모든 잠재 공간을 학습할 수 있으므로 개별 데이터 포인트를 잠재 공간의 멀리 떨어진 주머니에 배치하여 기억하게 되는 경우가 많다.
아래 그림은 오토인코더 잠재 공간과 관련된 위 문제를 시각화하고 VAE와 비교한다.

VAE는 잠재 공간에 사전 확률(a probabilistic prior)을 적용하여 이 문제를 극복한다.

이를 공식화하기 위해 잠재 공간 z를 확률 변수로 간주한다.

먼저 우리의 잠재성(latents)에 사전 확률(prior) p(z)p(z)를 적용한다. 대부분의 VAE에서는 일반적으로 표준 가우시안 분포N(0,1)\mathcal{N}(0, 1)이다

원시 데이터 포인트 x가 주어지면 잠재 공간에 대한 사후 확률(posterior) p(zx)p(z|x)도 정의한다

우리의 최종 목표는 데이터에 대한 사후 값을 계산하는 것. 이는 베이즈 규칙을 사용하여 다음과 같이 표현 할 수 있다.
p(zx)=p(xz)p(z)p(x)p(z|x)=\frac{p(x|z)p(z)}{p(x)}

연속 변수를 다루기 때문에 p(x)는 계산하기 힘들다. 계산 가능하도록 만들기 위해 우리는 사후 분포의 근사치를 특정 분포인 독립 가우스 분포로 제한한다. 이 근사 분포를 q(zx)q(z|x)라고 한다

이제 우리의 목표는 실제 사후분포와 이 근사치 사이의 KL발산(KL divergence)을 최소화 하는 것이다.

수학적인 세부 사항은 생략하지만 이 목적은 손실 함수를 최소화하여 계산 가능하다.

Ezq(zx)[log(p(xz))]+KL(q(zx)p(z))−E_{z∼q(z∣x)}​[log(p(x∣z))]+KL(q(z∣x)∣∣p(z))

이 경우 q(z|x)는 인코더에 의해 근사되고 p(x|z)는 디코더에 의해 표시된다.

VAE 손실은 직관적으로 해석된다.

첫 번째 항은 본질적으로 재구성 손실이고

두 번째 항은 posterior의 정규화를 나타낸다.
posterior은 KL divergence에 의해 prior으로 당겨지고 가우시안 prior을 향해 잠재 공간을 정규화한다. 이는 잠재 분포를 0 부근에서 조밀하게 분포된 상태로 유지하는 효과가 있다.

충분히 표현 가능한 신경망이 주어지면 VAE 잠재 공간은 복잡한 데이터 분포를 매우 깔끔하게 맞출 수 있다.

일반적으로 잠재 공간에서 서로 다른 유형의 데이터 포인트 간에 부드러운 전환을 기대할 수 있다.

예를 들어, VAE가 MNIST에서 훈련된 경우 6에 대한 클러스터와 5에 대한 별도의 클러스터를 기대할 수 있다. 두 클러스터 사이를 이동하면 5와 6이 이상하게 혼합된 것처럼 보이는 숫자가 표시된다

Basic Summary

  • 잠재 공간은 압축된 데이터 표현으로, 표현에서 데이터의 가장 중요하고 의미론적으로 흥미로운 특징을 강조하는 경우가 많다.
  • 학습된 잠재 표현은 많은 다운스트림 알고리즘에 유용할 수 있다.
  • Autoencoder는 데이터의 복잡한 비선형 함수인 잠재 공간을 찾는 방법.
  • Variational autoencoders는 오토 인코더 잠재 공간에 prior을 추가한다. 이것은 학습된 잠재 공간에 아주 좋은 속성을 제공(즉, 잠재 공간을 통해 데이터 분포를 부드럽게 보간할 수 있음).

Discrete Spaces

VAE와 VQ-VAE의 근본적인 차이는 VAE는 연속 잠재 표현을 학습하는 반면 VQ-VAE는 이산 잠재 표현을 학습한다는 것이다.

지금까지 연속 벡터 공간을 사용하여 오토 인코더에서 잠재성을 나타내는 방법을 살펴보았다.그러나 잠재성은 반드시 연속 벡터일 필요는 없으며 데이터에 대한 수치적 표현만 있으면 된다.

벡터 공간에 대한 잠재적으로 바람직한 대안 중 하나는 이산 표현이다. 일반적으로 현실 세계에서 접하는 많은 데이터는 이산 표현을 선호한다.

예를 들어, 인간의 말은 별개(discrete)의 음소와 언어로 표현된다.
또한 이미지에는 일부 이산 한정자 집합(some discrete set of qualifiers)이 포함된 이산 개체가 포함되어 있다.
하나는 객체 유형, 하나는 색상, 하나는 크기, 하나는 방향, 하나는 모양, 하나는 질감, 하나는 배경색, 하나는 배경 질감, 등…

표현 외에도 이산 데이터에 대해 작동하도록 설계된 여러 알고리즘(예: transformers)이 있으므로 이러한 알고리즘이 사용할 데이터의 이산 표현이 필요하다.

  • 분명히 이산 잠재 표현이 유용할 수 있지만, 그런 표현을 어떻게 배울 수 있을까?
  • 일반적으로 이산적인 것들은 딥러닝과 어울리지 않기 때문에 어려워 보일 수 있다
  • 다행히도 VQ-VAE는 바닐라 오토인코더 패러다임을 약간만 수정하면 작업에 대한 딥 러닝 작업을 수행할 수 있다

Quantizing Autoencoders

벡터양자화란?

VQ-VAE는 네트워크에 개별 코드북 구성 요소를 추가하여 표준 오토 인코더를 확장한다.

코드북은 기본적으로 해당 인덱스와 관련된 벡터 목록이며 오토 인코더의 병목 현상을 양자화하는 데 사용된다.

인코더 네트워크의 출력은 코드북의 모든 벡터와 비교되고 유클리드 거리에서 가장 가까운 코드북 벡터가 디코더에 공급된다.

수학적으로는 다음과 같이 표기

  • zq(x)=argminize(x)ei2z_q(x)=argmin_i||z_e(x)-e_i||_2
  • ze(x)z_e(x)는 원시 데이터 x에 대한 인코더 벡터 = 연속적인 데이터
  • eie_i는 i번째 코드북 벡터 = 카테고리 = 이산적인 데이터
  • zq(x)z_q(x)는 디코더에 대한 입력으로 전달되는 양자화된 결과 벡터

실제로 이 작업을 통해 디코더 그래디언트를 인코더로 직접 전달하면 모든게 잘 작동하는 듯 보인다. (즉, 인코더 및 양자화된 코드북 벡터의 경우 기울기를 1로 설정하고 다른 모든 코드북 벡터의 경우 0으로 설정)

그런 다음 디코더는 표준 오토 인코더 공식에서와 같이 이 양자화된 벡터에서 입력을 재구성하는 작업을 수행한다.

Producing Multiple Codes

디코더에 공급할 수 있는 제한된 벡터 집합(코드북 벡터 집합)가 있다는 사실이 의아할 수도 있다.

디코더가 코드북 벡터 집합만 입력으로 받아들일 수 있는데 어떻게 가능한 이미지의 엄청난 양과 다양성을 생성할 수 있을 것이라고 예상할 수 있었을까?

모든 데이터를 재구성할 수 있으려면 각 훈련 지점에 대한 고유한 이산 값이 필요하다.

그리고 이것이 사실이라면 모델은 각 훈련 지점을 다른 개별 코드에 매핑하여 데이터를 효과적으로 기억하지 않을까?

인코더가 하나의 벡터만 출력하는 경우 매우 유효한 문제이지만 실제 VQ-VAE에서 인코더는 일반적으로 일련의 벡터를 생성한다.

예를 들어, 이미지의 경우 인코더는 32x32 벡터 그리드를 출력할 수 있으며, 이들 각각은 양자화되고 전체 그리드가 디코더에 공급된다.

또는 유사하게 오디오 처리를 위해 인코더는 긴 1차원 벡터 시퀀스를 출력할 수 있다.
모든 벡터는 동일한 코드북으로 양자화되므로 이산 값의 수는 바뀌지 않는다.

그러나 여러 코드를 출력함으로써 디코더가 구성할 수 있는 데이터 포인트의 수를 기하급수적으로 늘릴 수 있다

예를 들어 이미지 작업을 하고 있고 크기가 512인 코드북이 있으며 인코더가 32x32 벡터인 그리드를 출력한다고 생각해보자. 이 경우 디코더는 5123232=29216512^{32*32}=2^{9216} 개의 다른 이미지를 출력할 수 있다. 이 숫자는 정말 헤아릴 수 없을 정도로 큰 숫자이며 이산 공간의 크기는 더이상 여기서 문제가 되지 않는다.

물론 모델은 여전히 훈련 데이터를 기억할 수 있지만, 올바른 유도 편향(inductive biases)(즉, 이미지의 경우 conv-net 사용)과 잠재 코드에 대한 올바른 구조(즉, 이미지의 경우 32x32 그리드)를 통해 모델은 데이터를 잘 나타내는 이산 공간을 학습해야 한다.

Learning the Codebook

인코더 및 디코더 네트워크와 마찬가지로 코드북 벡터는 경사 하강법을 통해 학습된다.
이상적으로 인코더는 일반적으로 학습된 코드북 벡터 중 하나에 가까운 벡터를 출력한다.

여기에는 기본적으로 양방향 문제가 있다. 즉, 인코더 출력에 맞춰 코드북 벡터를 학습하고 코드북 벡터에 맞춰 인코더 출력을 학습하는 것이다.

이 두 문제 모두 손실 함수에 항을 추가하여 해결된다.

전체 VQ-VAE 손실 함수

L=log p(xzq(x))+sg[ze(x)]e22+ze(x)sg[e]22L = log \space p(x|z_q(x)) + ||sg[z_e(x)] - e||_2^2 + ||z_e(x) - sg[e]||_2^2

여기서 우리는 이전 절의 동일한 표기법을 사용하고, sg[x]\text{sg}[x]는 "stop gradient"를 의미하며, 기울기가 방정식의 해당 부분을 통과하는 것을 방지한다.

첫 번째 항은 단지 표준적인 재건 손실(the standard reconstruction loss)이다.

두 번째 항은 코드북 정렬 손실(codebook alignment loss)로, 선택한 코드북 벡터를 인코더 출력에 최대한 가깝게 하는 것이 목표. 이 항은 코드북만 업데이트하기 위한 것이므로 인코더 출력에는 stop gradien operator가 있다.

세 번째 항은 두 번째 항과 비슷하지만, 코드북 벡터에 stop gradient를 적용한다. 인코더 출력이 가능한 한 가장 가까운 코드북 벡터에 커밋되도록 한다. 이 항을 코드북 커밋 손실(Commitment loss)이라고 하며, 전체 손실에 대한 중요도는 조정 가능한 초 매개 변수 β\beta로 조정한다

물론 이 마지막 두 항은 두 개 이상일 경우 모델의 각 양자화된 벡터 출력에 대해 평균화된다.

이를 통해 원본 이미지와 거의 눈에 띄게 다른 다양한 이미지 세트를 재구성할 수 있는 VQ-VAE를 완전히 훈련시킬 수 있다. 또한 VQ-VAE를 오디오나 비디오와 같은 다른 양식을 재구성하도록 훈련시킬 수 있다.

Learning the Prior

VQ-VAE가 완전히 훈련되면 훈련 시간에 부과된 균일한 prior을 포기하고 latents에 대해 새롭고 업데이트된 prior p(z)를 학습할 수 있다.

이산 코드의 분포를 정확하게 나타내는 prior을 학습하면 이 prior에서 샘플링하고 샘플을 디코더에 공급하여 분포에서 새 데이터를 생성할 수 있다.(원본 VQ-VAE 샘플은 아래 그림 참조).

prior 훈련에는 압축 이점도 있다. 잠재 분포가 균일하지 않은 경우 잠재 시퀀스를 나타내는 비트는 이전 분포에 표준 허프만 또는 산술 코딩을 적용하여 추가로 압축할 수 있다.

인코더가 각 데이터 포인트에 대한 일련의 잠재 코드를 출력한다고 가정하면 prior 학습을 위해 원하는 자동 회귀 모델(예: RNN 또는 transformer)을 사용할 수 있다. 자기회귀 인수분해는 시퀀스의 모든 이전 잠재 코드가 주어지면 다음 코드를 예측한다. 이것은 잠재 분포를 p(z)로 나눈다.

p(z)=p(z1)p(z2z1)p(z3z1,z2)p(z4z1,z2,z3)...,p(z)=p(z_1)p(z_2|z_1)p(z_3|z_1,z_2)p(z_4|z_1,z_2,z_3)..., where ziz_i is the ithith latent in the sequence.

오디오와 같은 1차원 신호로 작업하는 경우 문제를 자기회귀 형식으로 변환하는 것은 매우 간단하다. 오디오에서 인코딩된 1차원 시퀀스의 다음 잠재성을 예측하기만 하면 된다.

이 접근 방식은 시퀀스의 나중 잠재성이 오디오의 나중 섹션을 나타내도록 인코더가 구성되어 있다고 가정한다. (인코딩에 컨볼루션을 사용하는 경우 이 속성이 예상되어야 함).

이미지의 경우 32x32 잠재 그리드를 먼저 1차원 시퀀스로 롤아웃하여 시퀀스가 왼쪽 위에서 오른쪽 아래로 이동하도록 문제를 비슷하게 처리 할 수 있으며, 그런 다음 이 시퀀스에 자기회귀 학습을 적용할 수 있다.

profile
사람을 위한 기술을 공부합니다

1개의 댓글

comment-user-thumbnail
2023년 6월 30일

글 잘 읽었습니다. 혹시 다른 해외 블로그를 참조하여 글을 쓰셨나요? 레퍼런스를 같이 올려주시면 감사하겠습니다.
(수정) 맨 위에 레퍼런스가 있군요 감사합니다.

답글 달기