21.1 - Introduction

21.2 - VAE Basics

21.3 - VAE Generalizations (간단하게 Motivation + 컨셉까지만)

• $\beta$-VAE

• InfoVAE

21.4 - Avoiding Posterior Collapse
방법들만 간단히 설명 (별거 없음)

21.5 - Hierarchical VAE

21.6 - VQ-VAE

21.1 Introduction

Generative Models / 생성모델 - 데이터셋 $\mathcal X$ 에 대한 분포를 추정하는 모델이다. 추정된 분포 (모델) 에서 랜덤 샘플링을 하게 된다면 우리의 데이터셋 에서 뽑는 원리가 되기에. Label 이 필요없기에 비지도 학습의 일종이다.

다양한 생성모델 기법들 :

그럼 What is VAE?

우리의 데이터 $\mathbf{x}$ 를 Encoder $q(\mathbf{z}\mid \mathbf x)$ 에 넣어서 $\mathrm{z}$ 으로 압축을 시켜준다. 그런 다음에 Decoder $p(\mathbf x \mid \mathbf z)$ 에 넣어서 다시 원본 이미지처럼 만들어준다.

이 장에서 $\mathrm z, \mathbf x \mid \mathrm z$ 에 대한 가정은 다음과 같다.

• $\mathbf z \sim p_{\boldsymbol \theta}(\mathbf z)$

• $\mathbf x \mid \mathbf z \sim \mathrm{Expfam}\left( \mathbf x \mid d_{\boldsymbol \theta}\left(\mathbf z\right) \right)$

이때 :

• $p(\mathbf z) :$ 우리가 알고 있는 Latent $\mathbf z$ 에 대한 가정.

• $d_{\boldsymbol \theta}(\mathbf z) :$ Decoder (NN)

• $\mathrm{Expfam}\left( \mathbf x \mid \mathbf \eta \right) :$ Exponential Family Distribution, 예시로 Gaussian

Message Passing Algorithms 개념을 빌리자면,

• $p(z \mid x) :$ 우리의 데이터 $x$ 가 $z$ 의 어느 부분에 해당되는지 모른다. Hidden 이기 때문에. 그렇기에 이것을 Encoder 측에서 모델링 해줘야 한다.

• $p(x \mid z) :$ 우리의 Latent $z$ 를 $x$ 차원으로 다시 늘려주며, 이것은 이미 우리가 $z$ 가 뭔지 알고 있다는 가정이기에 그저 MLE 를 사용하면 된다.

VI 개념을 빌리자면, $p(z \mid x) = \frac{p(x \mid z) p(z)}{p(x)}$ 이며 $p(x) = \int p(x \mid z) p(z) dz$ 이기에 Intractable 하기에 VI 를 활용하며, 이때 활용되는건 Amortized VI, 다른 분포 $q$ 으로 $p$ 를 각 데이터 마다 하는 것이 아닌 모든 데이터 포인트에 대해서 모델링 시켜줍니다.

Why is it Intractable?

1. $z$ 차원이 커지면 적분 값을 구하는게 불가능해진다.

2. $p(x\mid z)$ 를 구하는데 Neural Network 을 사용하게 된다면 Non-linear 해지며, 우리는 $p(x\mid z)$ 를 그러면 $z$ 에 대한 함수로 바라보는게 불가능해진다.

21.2 VAE Basics

Assumptions

아주 기본적인 세팅에서 VAE 의 생성모델 부분 (Generative Model) 이 정의해주는 것은 :

• $p_{\boldsymbol{\theta}}(\boldsymbol{z}) :$ 보통은 Gaussian 으로 세팅하며 $z$ 차원이 고차원이면 그저 다변량 정규분포이다.

• $p_{\boldsymbol{\theta}}(\boldsymbol{x} \mid \boldsymbol{z}) = \prod_{d=1}^{D}\mathrm{Ber}\left( x_d \mid \sigma\left( d_{\boldsymbol \theta}(\boldsymbol{z} \right) \right) :$ 예시 중 하나, 그저 Product of Exponential Family Distributions 이기만 하면 된다.

아주 기본적인 세팅에서 VAE 의 추론 부분 (Recognition Model) 이 정의해주는 것은 :

• $q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})=\mathcal{N}(\boldsymbol{z} \mid \boldsymbol{\mu}, \operatorname{diag}(\exp (\boldsymbol{\ell}))) :$ 보통은 이것도 Gaussian 이라고 가정, 즉 Latent $z$ 도 가우시안.

• $(\boldsymbol{\mu}, \boldsymbol{\ell})=e_{\boldsymbol{\phi}}(\boldsymbol{x}) :$ $\boldsymbol \ell = \log \boldsymbol \sigma$, 시그마 부분이 벡터, 이유는 Prior 를 Vanilla VAE 에서는 Mean-field Gaussian 으로 Assume 하기 때문. 그래서 아래 알고리즘은 튀어나오는 값이 조금 다르다.

Model Fitting

우리가 최대화 하고 싶은 확률은 그러면

• $p_{\boldsymbol{\theta}}(\boldsymbol{z}, \boldsymbol{x})= p_{\boldsymbol{\theta}}(\boldsymbol{z}) p_{\boldsymbol{\theta}}(\boldsymbol{x} \mid \boldsymbol{z}) :$ Decoder 부분

• $q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x}) :$ Encoder 부분

이것들을 Negative Log-Likelihood (NLL) 형태로 바꿔서 로스 형태로 표현하면 된다.

알고리즘은 다음과 같으며, Lines 5-7 부분의 경우에는 Cholesky Decomposition 을 활용해서 다변량 정규분포에서 뽑는 모습. 그러면 여기에서의 Encoder 는 사실은 $e_{\boldsymbol \phi}(\boldsymbol x)$ 이다. 추가적으로 Full Covariance Gaussian Posterior 를 Assume 함.

• 참조 for Multivariate Normal Dist. Sampling using Cholesky Decomposition

  from 통계전산 강의 (정성현 교수님 강의안)

  

  Covariance Matrix is always PD (아주 특이 케이스 제외하면), hence we can use Cholesky Decomposition.

  위의 식에서는 $L^\prime$ 이 사실은 하삼각행렬일 보장이 없기에 Diagonal + 상단 부분을 제외한 Mask $M$ 으로 씌워준 다음에 Diagonal 에 $\sigma$ 를 넣어주면서 Covariance Matrix 을 만들어줌.

  $\boldsymbol \varepsilon \sim \mathrm{N}$ 이면 Linear Combination 으로 우리가 원하는 정규분포로 샘플링 해줄 수 있다.

  $$Z \sim \mathrm{N}(0, 1), \quad aZ + b \stackrel{D}{\rightarrow}\mathrm{N}(b, a^2)$$

• Line 10 : Computing for the NLL of $q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})$

• Line 11 : Computing for the NLL of $p_{\boldsymbol \theta}(\boldsymbol z)$

• Line 12 : Computing for the NLL of $p_{\boldsymbol \theta}(\boldsymbol x \mid \boldsymbol z)$

사실 Line 13 에서 정의하는 로스가 그저 -ELBO Term 이 되어버린다. 정확하게는 (w/o considering parameters),

• About ELBO

  from 10장 (손우진 님 정리안)

  $$\mathcal{L}(\theta, \psi \mid x) = \mathbb{E}_{q_\psi(z)} \big[ -\log p_\theta(x, z) + \log q_\psi(z) \big]$$

  

How to sample in this VAE?

어느 정도 수렴이 된다면 우리가 디코더 $d_{\boldsymbol \theta}(\boldsymbol z)$ 를 Prior $p(\boldsymbol z)$ 에서 시작해서 이미지로 나오게끔 잘 만들어줬기 때문에 그저 $\boldsymbol z \sim p(\boldsymbol z)$ 에서 뽑은 다음에 이것을 디코더 $d_{\boldsymbol \theta}(\boldsymbol z)$ 에 넣어주기만 하면 된다.

Advantage of VAE over AE?

Autoencoder 는 사실은 생성모델이 아니다. 압축하는게 사실은 목표이며 더 잘 압축하기 위해서 Decoder 를 결합한게 Autoencoder 일 뿐이다. 한 마디로 Randomness 요소가 없다. 그러기에 Latent 에 대해서 Randomness 요소가 있는 VAE 가 생성 측면에서는 이점을 가진다.

• About Autoencoder

  from 생성모델 강의 (송경우 교수님 강의안)

  

Final Loss Term would be :

Alternative Expression for ELBO in VAE

• Joint Generative 부분 :
  • $p_{\boldsymbol{\theta}}(\boldsymbol{x}, \boldsymbol{z})=p_{\boldsymbol{\theta}}(\boldsymbol{z}) p_{\boldsymbol{\theta}}(\boldsymbol{x} \mid \boldsymbol{z})$
  • $p_{\boldsymbol{\theta}}(\boldsymbol{x})= \int_{\boldsymbol z} p_{\boldsymbol \theta}(\boldsymbol x, \boldsymbol z) d\boldsymbol z$
  • $p_{\boldsymbol{\theta}}(\boldsymbol{z} \mid \boldsymbol{x})=p_{\boldsymbol{\theta}}(\boldsymbol{x}, \boldsymbol{z}) / p_{\boldsymbol{\theta}}(\boldsymbol{x})$

• Joint Inference 부분 :
  • $q_{\mathcal{D}, \phi}(\boldsymbol{z}, \boldsymbol{x})=p_{\mathcal{D}}(\boldsymbol{x}) q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})$
  • $p_{\mathcal{D}}(\boldsymbol{x})=\frac{1}{N} \sum_{n=1}^N \delta\left(\boldsymbol{x}_n-\boldsymbol{x}\right)$
  • $q_{\mathcal{D}, \boldsymbol{\phi}}(\boldsymbol{z})=\int_{\boldsymbol{x}} q_{\mathcal{D}, \boldsymbol{\phi}}(\boldsymbol{x}, \boldsymbol{z}) d \boldsymbol{x}$
  • $q_{\mathcal{D}, \boldsymbol{\phi}}(\boldsymbol{x} \mid \boldsymbol{z})=q_{\mathcal{D}, \boldsymbol{\phi}}(\boldsymbol{x}, \boldsymbol{z}) / q_{\mathcal{D}, \boldsymbol{\phi}}(\boldsymbol{z})$

이것들을 이용해서 ELBO Term 을 다음과 같이 바꿔줄 수가 있습니다.

• 풀이

  Original Loss (Negative ELBO) :

  $$\begin{aligned} L & = -\mathbb{E}_q\left[\log p(x \mid z) \right] + \mathrm{KL}\left(q\left(z\mid x\right) \| p\left(z\right) \right), \quad \text{w/o params} \\ L(\theta, \phi \mid \mathcal D) &= \mathbb{E}_{p_{\mathcal D}(\boldsymbol x)}\left[-\mathbb{E}_{q_\phi(\boldsymbol z \mid \boldsymbol x)}\left[\log p_{\boldsymbol \theta}(\boldsymbol x \mid \boldsymbol z) \right] + D_{\mathbb{KL}}\left(q_\phi(\boldsymbol z \mid \boldsymbol x) \| p_{\boldsymbol \theta}\left(\boldsymbol z\right) \right) \right] \end{aligned}$$

  Positive ELBO with parametrization :

  $$\text{ELBO}(\theta, \phi \mid \mathcal D) = \mathbb{E}_{p_{\mathcal D}(\boldsymbol x)}\left[\mathbb{E}_{q_\phi(\boldsymbol z \mid \boldsymbol x)}\left[\log p_{\boldsymbol \theta}(x \mid z) \right] - D_{\mathbb{KL}}\left(q_\phi(\boldsymbol z \mid \boldsymbol x)\left(z\mid x\right) \| p_{\boldsymbol \theta}\left(z\right) \right) \right]$$

  풀이 :

  

여기에 $\stackrel{c}{=}$ 을 Mean Equal up to additive constants 으로 둔다면, 뒤에 있는 상수 부분인 $p_\mathcal{D}(x)$ Term 을 버릴 수가 있다.

• 풀이

  

21.21 과 같이 해석을 할 수도 있고 아래 보이는 식과 같이 또 다르게 해석할 수도 있다 (Marginal as z instead of x).

둘 다 달라진건 사실 없고 ELBO 에 대한 관점만 바뀌었을 뿐! 그럼 그래서 이렇게 보는 이점이 뭘까?

• 1번째 Term 에서 보이는 바와 같이 Inference & Generative Prior 의 차이를 최소화 시켜준다는 것을 알 수가 있음.

• 2번째 Term 에서 보이는 바와 같이 $z$ 를 통해서 $x$ 를 만드는 Reconstruction 과정에서 Inference & Generative Distribution 의 차이를 최소화 시켜준다는 것을 알 수가 있음.

• 참조 Figure

  

VAE Generalizations

$\beta$-VAE

기존의 VAE Loss 식은 다음과 같다.

근데 결과를 조금 살피면 Blurry 하다는 모습을 보인다는 것을 알 수가 있다 (맨 아랫 줄에 해당함). 이에 대한 이유로는 Gaussian Decoder 를 사용한다는 것이 문제이다.

Decoder 에 대해서 얘기했을 때 우리는 Decoder 가 그저 Neural Network 이라고 말만 했었지만, 사실은 그것보다 조금 더 구체적으로 들어가야된다.

이미지 $x$ 가 예시로 MNIST 이미지라서 $28\times 28$ 이며 $z$ 는 예시로 $d$-차원이라고 한다면, 우리의 생성된 이미지 $\hat{x}$ 는 정규분포에서 뽑힌 값이다.

그럼 왜 이게 Blurry 랑 연관됐을까?

실제 이미지 $x$ 랑 생성된 이미지 $\hat{x}$ 를 비교했을 때의 확률을 최대화 시키는 것은 픽셀별로 평균낸 꼴이 되어버림.

평균을 내어버리기에 Blurry 한 이미지들이 나오게 된다고 한다.

그러면 이걸 어떻게 고칠 수가 있을까? $\rightarrow$ Prior 를 더 표현력 있게 만든다 (예시: Normalizing Flows, $z_K=f_K \circ \cdots \circ f_1\left(z_0\right), \quad z_0 \sim \mathcal{N}(0, I)$)

더 간단한 방법으로는 $\beta$-VAE 가 있다. 근데 이것을 알기 전에 Disentanglement 에 대해서 알아야 한다.

원래 비전에서는 이것을 Total-Correlation Loss 으로 사용하기는 하지만,

다른 논문에서는 다음의 Modified 된 VAE 의 로스함수에서 $\beta > 1$ 을 사용하는 것이 TC 를 줄이는 것과 똑같다고 한다.

왜 이럴까?

• $\beta = 0 :$ 그저 Vanilla VAE

• $\beta < 1 :$ Encoder 가 압축하는 $z$ 에 $N(0, I)$ 를 덜 따라가게끔 만들어준다.
  • 일종의 $z$ 로 인해서 발생한 Information Bottleneck 을 줄여주게 된다고 보면 되며, $z$ 가 $x$ 에 대한 더 많은 정보를 담고 간다고 보면 됨.

• $\beta > 1 :$ 반대로 Penalize 를 더 해주며, Encoder 가 $N(0, I)$ 으로 안 가면 더 페널티를 높여주며 Disentanglement 유도가 된다.

원래 $\beta$-VAE 논문에서는 이를 통해서 특정 $z$ 를 바꾸면 이미지에 대한 특징들을 찾을 수 있다고 한다.

InfoVAE

그럼 VAE 에서 생길 수 있는 또 다른 문제점은 뭐가 있을까?

Encoder 의 목표에서는 $z\mid x$ 를 $N(0, I)$ 에 가깝게 하는 목표도 있지만 $x$ 에 관한 정보를 최대한 담는다는 것도 있다. 하지만 Vanilla ELBO 에서는 $z\mid x$ 가 $N(0, I)$ 으로 ‘Collapse’ 해버릴 수도 있다는 단점이 존재한다.

그러면 이때 Mutual Information 을 고려해보자. $I(X;Z)$ 이라고 하면 $X$ 와 $Z$ 가 공유하고 있는 정보이다.

이 값을 올리게 된다면 $X$ 와 $Z$ 는 더욱 더 많은 정보를 공유하게 될 것이다. 이것을 VAE Loss 에 포함시켜준다면 다음과 같다.

• $\alpha$ 값을 올리면 $MI$ 가 올라가면서 Encoder 가 $x$ 에 대해서 $z$ 으로 압축시킬 때 $x$ 에 대한 정보를 너무 많이 잃지 않도록 만든다.

• $\lambda$ 값의 경우 $\beta$-VAE 부분을 반영한 모습, $\alpha := 1- \lambda$ 인 경우에는 $\lambda$ 가 $\beta$ 가 된다.

Conditional VAE’s

새로운 변수 $c$ 가 Condition (Class) 일 때 VAE 를 모델링 하는 것.

Generative Part :

Inference Part :

Avoiding Posterior Collapse

Posterior Collapse 이란 우리가 Encoder 를 통해서 압축한 $z$ 가 $q_\phi (z\mid x) \approx p(z)$ 으로 되어버리는 케이스이다. 즉 $x$ 에 대해서 독립이 되어버린다는 케이스.

해결 방법으로는 :

• KL-Annealing - $\beta$-VAE 에서의 $\beta$ 값을 학습 초반에는 $0.0$ 에서 시작하다가 점차적으로 $1.0$ 으로 올려주는 Scheduling 방식.

• Free bits - 원래는 $\mathcal{L}_R=\sum_i D_{\mathbb{K} \mathbb{L}}\left(q_{\boldsymbol{\phi}}\left(z_i \mid \boldsymbol{x}\right) \| p_{\boldsymbol{\theta}}\left(z_i\right)\right)$ 이었다면 이것을
  $\mathcal{L}_R^{\prime}=\sum_i \max \left(\lambda, D_{\mathbb{K} \mathbb{L}}\left(q_{\boldsymbol{\phi}}\left(z_i \mid \boldsymbol{x}\right) \| p_{\boldsymbol{\theta}}\left(z_i\right)\right)\right)$ 으로 바꿔주는 것 (Lower bounds the rate)

VAEs with Hierarchical Structure

이 문제를 정의를 좀 해보자면, 다음과 같다.

그러면 모델의 Expressive Power 가 늘어나게 된다.

Generative Model :

Markovian 가정을 깨버릴 수가 있다. Layer 들로 나눠줄 수가 있기에.

Inference Model :

$l=2 \rightarrow L$ 과 같은 Bottom-up 이 아닌 아래와 같은 Top-down Approach 도 고려해볼 수 있음.

둘을 고려했을 때의 목적함수는 다음과 같다.

결과 및 예시 :

VQ-VAEs (Vector Quantized VAEs)

Posterior Collapse 문제를 해결하기 위한 또 다른 방법. $x$ 마다 $z$ 가 좀 다양하게 나오길 기대하지만 그렇지 않고 비슷하게 나온다는 것.

VAE 와는 다르게 Encoder 가 Continuous 가 아닌 Discrete Latent 변수를 뱉게끔 만든다.

Encoder 의 뒷 부분은 $z_e(x)$ 의 각각의 Spatial Location $(i,j)$ 에 관해서 Mapping 을 시켜주는 다음과 같은 함수라 하자.

이것에 그러면 $(i,j)$ 에 해당하는 $k$ 값이 있다는 뜻. 그러면 우리의 Codebook 을 찾아서 이 값에 해당되는 Embedding Vector $e_k$ 를 찾는다 (Embedding lookup).

그럼 다음에 이 값을 모아서 $z_q(x)$ 이라고 한다면 (Fig 에서 디코더의 Input) 이것을 Decoder 에 넣고 Output 을 뱉는다.

로스는 다음과 같이 나온다.

• 1번째 Term 은 그저 Reconstruction loss, Decoder 에 줄 입력값인 $\boldsymbol z_q(\boldsymbol x)$ 에 대해서 출력된 이미지에 대한 로스 업데이트.

• 2번째 Term 은 Encoder output $\boldsymbol z_e$ 를 고정시키고 Codebook 인 $\boldsymbol e$ 에 대해서만 업데이트 해주는 것.

• 3번째 Term 은 2번째 Term 의 반대.

Why no ELBO?

KL Term 대신에 Codebook 개념으로 바라보기에 이렇게 다르게 바라봐야 된다.

Why the stop gradient to divide the loss into two?

둘다 업데이트를 동시에 하게 된다면 Encoder 와 Codebook 둘다 똑같은 방향으로 업데이트 하게 되며 둘에 대한 거리는 계속해서 유지될 것이다 (Moving Target).

Why the $\beta$?

Encoder 의 Output 이 Codebook 과 정확하게 일치하게끔 만들기 싫기에 $\beta$ 값을 줄여주면 이것을 방지해준다.

Conclusion

• What is VAE?
  • Autoencoder 의 Latent 를 Prior 분포으로 모델링 시키면서 확률적 요소를 더해준 것이다.
  • Variational Inference 를 활용한다.

• VAE 의 문제점으로는?
  • Disentanglement: $z$ 가 특징을 잡으면서 압축이 되어야 하는데 그게 안될 수도 있다.
    • $\beta$-VAE
    • Hierarchical VAE
  • Posterior Collapse : 압축된 $z$ 가 Input $x$ 에 대한 정보를 대부분 잃어가면서 Prior $z$ 에 맞춰지게끔 하는 경우. Encoder collapses to the prior.
    • InfoVAE
    • VQ-VAE