[비전 생성 모델] VAE (Variational Auto Encoder)
AE (Auto Encoder)
AE는 인코더-디코더 구조를 가진다.
인코더는 입력 $\mathbf{x}$를 잠재 벡터 $\mathbf{z}$로 압축하며, 디코더는 이 $\mathbf{z}$를 이용해 복원한다. 이때, $\mathbf{z}$는 고정된 벡터이다.
VAE (Variational Auto Encoder)
VAE는 AE와 달리, 입력 데이터를 잠재 변수의 단일 지점이 아닌 확률 분포로 매핑한다.
즉, 같은 입력 $\mathbf{x}$에 대해서도 매번 조금씩 다른 잠재 변수 $\mathbf{z}$가 샘플링될 수 있다.
VAE의 목표는 입력 데이터 $\mathbf{x}$가 주어졌을 때, 이 데이터를 가장 잘 설명할 수 있는 잠재 변수 $\mathbf{z}$의 분포를 학습하는 것입니다.
이때, 모델 설계를 위해 잠재 변수 $\mathbf{z}$의 사전 분포를 가장 간단한 형태인 표준 가우시안 분포로 가정한다.
\[p(\mathbf{z})=\mathcal{N}(\mathbf{0},\mathbf{I})\]
하지만 이 단순한 사전 분포 $p(\mathbf{z})$에서 무작위로 $\mathbf{z}$를 샘플링하여 디코더에 통과시켜보니, 원본 이미지조차 생성하지 못하였다.
왜냐하면 고차원 공간에서 대부분의 $\mathbf{z}$는 의미 없는 잡음으로 연결되기 때문이다.
따라서 최소한 원본 데이터 $\mathbf{x}$는 잘 복원하기 위해서는, 아무 $\mathbf{z}$나 샘플링하는 것이 아니라 해당 $\mathbf{x}$를 만들어낼 확률이 높은 $\mathbf{z}$를 샘플링해야 한다.
따라서 데이터 $\mathbf{x}$가 주어졌을 때 $\mathbf{z}$의 분포인 사후 확률 (Posterior) $p(\mathbf{z}\mid\mathbf{x})$에서 $\mathbf{z}$를 샘플링한다.
\[\mathbf{z}\sim p(\mathbf{z}\mid\mathbf{x})\]이 역할을 수행하기 위해 설계된 것이 인코더이다.
인코더 (Encoder)
우리는 이상적인 사후 확률 분포인 True posterior $p(\mathbf{z}\mid\mathbf{x})$를 알고 싶지만, 이를 직접 계산하는 것은 불가능에 가깝다.
베이즈 정리에 의해 $p(\mathbf{z}\mid\mathbf{x})$는 다음과 같이 표현된다.
\[p(\mathbf{z}\mid\mathbf{x})=\frac{p(\mathbf{z},\mathbf{x})}{p(\mathbf{x})}= \frac{p(\mathbf{x}\mid\mathbf{z})p(\mathbf{z})}{\int p(\mathbf{x}\mid\mathbf{z})p(\mathbf{z})d\mathbf{z}}\]여기서 잠재 변수 $\mathbf{z}$는 고차원 벡터이며, 차원이 증가할수록 적분 계산량이 기하급수적으로 늘어난다.
예를 들어 $\mathbf{z}$가 100차원 벡터라면, 이 적분은 아래와 같은 100중 적분이 된다.
\[\int p(\mathbf{z})d\mathbf{z}=\int\cdots\int p(z_{1},z_{2},\cdots,z_{100})dz_{1}dz_{2}\cdots dz_{100}\]즉, 분모의 적분 계산이 사실상 불가능하기 때문에, True posterior를 정확히 알아낼 수 없다.
이 문제를 해결하기 위해 변분 추론을 도입한다.
계산할 수 없는 True posterior $p(\mathbf{z}\mid\mathbf{x})$를 대신하기 위해, 우리가 다루기 쉬운 가우시안 분포 형태의 근사 함수 $q_\phi(\mathbf{z}\mid\mathbf{x})$를 정의하고, 이 분포가 True posterior와 최대한 비슷해지도록 학습시키는 것이다.
이 근사 함수 $q_\phi(\mathbf{z}\mid\mathbf{x})$를 신경망으로 모델링한 것이 인코더이다.
\[q_\phi(\mathbf{z}\mid\mathbf{x})\approx p(\mathbf{z}\mid\mathbf{x})\]인코더는 입력 $\mathbf{x}$를 받아, 근사된 사후 확률 분포 $q_\phi(\mathbf{z}\mid\mathbf{x})$의 파라미터인 평균 $\boldsymbol{\mu}$와 분산 $\boldsymbol{\sigma}^2$을 출력한다.
Decoder
디코더는 샘플링된 $\mathbf{z}$를 입력 받아, 원본 데이터 $\mathbf{x}$를 복원하는 역할을 수행한다.
인코더와 마찬가지로 신경망으로 모델링된다.
\[p_\theta(\mathbf{x}\mid\mathbf{z})\]디코더 신경망은 $\mathbf{z}$를 입력받아 원본 데이터의 분포 파라미터를 출력한다.
전체 모델 구조는 아래와 같다.
재매개변수화 트릭 (Reparameterization Trick)
VAE 학습 과정에서 인코더가 추정한 분포 $q_\phi(\mathbf{z}\mid\mathbf{x})$로부터 잠재 변수 $\mathbf{z}$를 샘플링해야 한다.
그러나 단순한 샘플링 연산은 미분이 불가능하여 역전파를 통한 그라디언트 전달이 끊기는 문제가 발생한다.
이를 해결하기 위해 재매개변수화 트릭을 도입한다.
이는 확률적인 샘플링 과정을 확정적인(Deterministic) 연산과 별도의 노이즈 항으로 분리하는 기법이다.
\[\mathbf{z}=\boldsymbol{\mu}+\boldsymbol{\sigma}\odot\boldsymbol{\epsilon}~~,~~\boldsymbol{\epsilon}\sim\mathcal{N}(\mathbf{0},\mathbf{I})\]여기서 $\odot$은 요소별 곱 (Element-wise product)을 의미한다.
이 트릭을 통해 모델의 파라미터인 $\boldsymbol{\mu}$와 $\boldsymbol{\sigma}$에 대한 미분이 가능해져, 역전파를 통해 인코더를 정상적으로 학습시킬 수 있게 된다.
Hierarchical VAE (HVAE)
- VAE의 latent variable을 여러 층으로 쌓은 모델
$t$시점의 latent는 $t+1$시점의 latent의 영향만 받는다는 Markov 가정 하에 진행
$p(z_t z_{t+1},z_{t+2},\dots)$는 매우 복잡하기 때문에 풀어쓸 수 없음
Encoder
\[q_\phi(z_{1:T}|x)=q_\phi(z_{1}|x)\prod_{t=2}^Tq_\phi(z_t|z_{t-1})\]- 처음 step에는 $x$가 포함되어 있으므로, $\prod$ 안에 넣을 수 없음
Decoder
\[p(x,z_{1:T})=p(z_T)p_\theta(x|z_1)\prod_{t=2}^Tp_\theta(z_{t-1}|z_t)\]- Latent가 여러 층이므로 VAE와 달리 $z_{1:T}$로 표현
- 우변은 Markov 정의에 의해 나온 식임