VAE : Auto-Encoding Variational Bayes (2022.12.V11)
Diederik P. Kingma, Max Welling(Universiteit van Amsterdam)
https://arxiv.org/abs/1312.6114
오늘은 확률적 생성 모델인 변분 오토인코더를 이용해 데이터의 잠재 공간을 학습하고, 이를 통해 생성 품질을 향상시키는 방법론을 제시한 논문 VAE : Auto-Encoding Variational Bayes를 리뷰해보도록 하겠습니다! 😁
📌 Abstract
본 논문에서는 연속적인 latent variable이 있는 확률 모델에서 효율적으로 추론 및 학습을 진행하는 방법에 대해 소개합니다. 특히 사후 분포가 다루기 어렵다거나, 데이터가 매우 큰 경우에도 효과적으로 작동하는 알고리즘을 다룹니다.
Q. 사후 분포가 다루기 어렵다는 말이 무엇이냐?
사후 분포는 기본적으로 X라는 데이터가 입력으로 들어왔을 때, Z라는 잠재변수가 나올 실제 확률 분포를 의미하며, P(Z∣X)로 표현됩니다. 이때 사후 분포는 베이즈 정리를 통하여 계산할 수 있고, P(Z∣X)=P(X∣Z)P(Z)/P(X)와 같은 식으로 표현됩니다.
분모에 들어가는 X가 등장할 확률 P(X)는 다시 P(X)=∫P(X∣Z)P(Z)dZ 로 표현이 가능한데, 모든 Z에 대해서 적분을 한다는 것이 모든 Z에 대해서 무한번 샘플링하여 Z의 latent space를 빈틈 없이 다 활용한다는 것과 동일한 개념인데 사실상 이것이 현실세계에서 사실상 불가능하다는 문제점이 있습니다.
따라서 논문은 본 문제를 해결하고, 연구 내용은 크게 2가지 기여를 다룹니다.
1) 위와 같이 사후 분포를 직접적으로 다루기 어려운 경우, 이를 대신 최적화 할 수 있는 하한 추정기 ELBO를 제공합니다. SGD와 같은 일반적인 Optimizer로 하한을 쉽게 학습할 수 있어, 이를 최적화하는 것이 logP(X)를 최대화하는 것과 비슷한 효과를 내는 원리입니다. 이 과정에서 Reparameterization Trick을 활용하여 Z의 확률적 샘플링을 미분 가능하게 변환함으로써 원활한 학습이 가능하도록 조정하였습니다.
2) i.i.d 데이터셋, 즉 동일 분포 내의 독립적인 각 데이터 샘플들에 대해 연속적인 latent variable을 갖는 사후 분포를 근사할 수 있는 추론 모형을 학습합니다. 정확한 사후 분포를 계산할 수 없기 때문에, 근사를 통해 어떠한 모형의 분포를 사후 분포와 유사해지도록 학습을 통해 조정하는 과정입니다. 이 모형이 우리가 흔히 알고 있는 인코더에 해당한다고 봐주시면 되겠습니다!
📌 Introduction
연속적인 잠재 변수 z 또는 모델 파라미터를 가진 방향성 확률 모델은, 데이터가 주어졌을 때의 사후 분포를 정확히 계산하기 어렵습니다. 전통적인 mean-field Variational Bayes는 잠재 변수끼리 독립이라고 가정해 근사 분포 q(z∣x) 를 최적화하지만, 그 과정에서 각 변수의 기대값을 해석적으로 계산해야 합니다.
Stochastic Gradient Variational Bayes, SGVB는 ELBO를 재매개변수화하여, 편향은 없고 미분 가능한 기대값 추정기를 만듭니다. 예를 들어 정규 분포라면 처럼 샘플 과정을 완전히 계산 그래프 안으로 옮겨 버립니다. 덕분에 미니배치,SGD만으로도 연속 잠재 변수를 가진 거의 모든 모델에서 근사 사후 추론이 가능해집니다.
SGVB를 인코더와 결합하면 Auto-Encoding Variational Bayes(AEVB) 알고리즘이 됩니다. 인코더는 입력을 받아 잠재 공간의 평균·분산을 즉시 내놓는 신경망이고, ELBO를 SGVB로 최대화하면서 추론과 학습을 한꺼번에 끝냅니다.
📌 Method
- Problem Scenario
본 연구는 대규모 i.i.d. 데이터셋 \(X = \{x^{(i)}\}_{i=1}^N\)을 가정하고, 각 관측값 \( x^{(i)} \)가 두 단계 생성 과정을 통해 얻어진다고 설정합니다.
첫 단계에서는 잠재 연속 변수 \( z^{(i)} \) 가 사전 분포 \(p_{\theta^*}(z)\)에서 샘플링되고, 두 번째 단계에서는 조건부 분포 \(p_{\theta^*}(x \mid z)\)로부터 \( x^{(i)} \) 가 생성됩니다.
이때 진짜 파라미터 \( \theta^* \) 와 잠재 변수 \( z^{(i)} \) 는 미지수이며, 주변우도 \(p_{\theta}(x) = \int p_{\theta}(z)\, p_{\theta}(x \mid z)\, dz\) 및 사후분포 \(p_{\theta}(z \mid x) = \frac{p_{\theta}(x \mid z)\, p_{\theta}(z)}{p_{\theta}(x)}\)를 정확히 계산하거나 미분할 수 없는 상황을 전제로 합니다. 따라서 EM 알고리즘이나 전통적 mean-field 변분 추론은 적용이 불가능하고, 대규모 데이터 때문에 배치 최적화 역시 비현실적입니다.
이러한 환경에서 본 연구는 파라미터 \( \theta \) 의 근사 최대우도(ML) 또는 최대사후(MAP) 추정, 주어진 \( x \) 에 대한 잠재 변수 \( z \) 의 근사 사후 추론, 변수 \( x \) 의 근사 주변 추론 이라는 세 가지 문제를 동시에 해결하고자 합니다. 이를 위해 실제 사후분포 \( p_{\theta}(z \mid x) \) 를 근사할 인코더 \(q_{\phi}(z \mid x)\)를 도입합니다.
이 모델은 mean-field 가정을 요구하지 않으며, 파라미터 \( \phi \) 는 닫힌꼴 기대값이 아닌 학습 과정에서 \( \theta \) 와 함께 직접 최적화됩니다.
\( q_{\phi}(z \mid x) \) 는 관측 \( x \) 로부터 생성 가능한 잠재 코드 \( z \) 에 대한 분포를 출력하므로 확률적 인코더로, \(p_{\theta}(x \mid z)\)는 코드 \( z \) 로부터 가능한 관측 \( x \) 분포를 출력하므로 확률적 디코더로 해석됩니다.
- The variational bound
데이터셋 \( X = \{x^{(i)}\}_{i=1}^{N} \) 의 로그우도는 \(
\log p_{\theta}(x^{(1)}, \dots, x^{(N)}) = \sum_{i=1}^{N} \log p_{\theta}(x^{(i)})
\) 처럼 항별로 분해할 수 있습니다.
각 항은 다시 근사 사후 \( q_{\phi}(z \mid x^{(i)}) \) 와 진짜 사후 \( p_{\theta}(z \mid x^{(i)}) \) 사이의 KL-Divergence, 그리고 ELBO \( L(\theta, \phi; x^{(i)}) \)로 나뉩니다.
\[
\log p_{\theta}(x^{(i)}) = D_{\mathrm{KL}}\left( q_{\phi}(z \mid x^{(i)}) \,\|\, p_{\theta}(z \mid x^{(i)}) \right) + L(\theta, \phi; x^{(i)})
\]
KL-Divergence가 항상 0 이상이므로, 다음이 성립합니다.
\[
\log p_{\theta}(x^{(i)}) \geq L(\theta, \phi; x^{(i)}) = \mathbb{E}_{q_{\phi}(z \mid x^{(i)})} \left[ -\log q_{\phi}(z \mid x^{(i)}) + \log p_{\theta}(x^{(i)}, z) \right]
\]
이를 다시 두 항으로 정리하면,
\( L(\theta, \phi; x^{(i)}) = -D_{\mathrm{KL}}\left( q_{\phi}(z \mid x^{(i)}) \,\|\, p_{\theta}(z) \right) + \mathbb{E}_{q_{\phi}(z \mid x^{(i)})} \left[ \log p_{\theta}(x^{(i)} \mid z) \right] \)가 됩니다. 연구 목표는 \( \theta \) 와 \( \phi \) 양쪽에 대해 이 하한을 최대화하는 것입니다.
그러나 \( \phi \) 에 대한 기울기를 계산하기 위해 단순 몬테카를로 추정을 사용하면, \( \nabla_{\phi} \, \mathbb{E}_{q_{\phi}(z)}[f(z)] = \mathbb{E}_{q_{\phi}(z)} \left[ f(z) \nabla_{\phi} \log q_{\phi}(z) \right] \approx \frac{1}{L} \sum_{l=1}^{L} f(z^{(l)}) \nabla_{\phi} \log q_{\phi}(z^{(l)}) \) 여기서 \( z^{(l)} \sim q_{\phi}(z \mid x^{(i)}) \) 이며, 이 방식은 기울기 추정치의 분산이 극단적으로 커져 학습이 사실상 불가능해집니다.
이러한 고분산 문제를 해결하기 위해 재매개변수화 트릭과 SGVB 추정량이 도입되며, 이를 통해 ELBO를 미니배치 확률적 경사하강법으로 안정적으로 최적화할 수 있습니다.
- The SGVB estimator and AEVB algorithm
SGVB 추정량과 AEVB 알고리즘의 핵심은 재매개변수화 트릭으로, 인코더 \( q_{\phi}(z \mid x) \) 에서의 샘플링을 미분 가능한 형태로 바꾸는 것입니다.
먼저 \( \varepsilon \sim p(\varepsilon) \) 와 미분 가능 변환 \( g_{\phi} \) 를 이용해 샘플을 다음과 같이 재표현할 수 있습니다. \(z^{\sim} = g_{\phi}(\varepsilon, x)\) 이렇게 하면 어떤 함수 \( f(z) \) 에 대한 기대값을 다음과 같이 바꿀 수 있습니다. \(
\mathbb{E}_{q_{\phi}(z \mid x)}[f(z)] = \mathbb{E}_{p(\varepsilon)}[f(g_{\phi}(\varepsilon, x))]\) 이는 미분 가능하며 소수의 샘플로도 안정적인 몬테카를로 추정이 가능합니다.
이 트릭을 ELBO에 적용하면 \(L(\theta, \phi; x) = -D_{\mathrm{KL}}(q_{\phi}(z \mid x) \,\|\, p_{\theta}(z)) + \mathbb{E}_{q_{\phi}(z \mid x)}[\log p_{\theta}(x \mid z)]\)인데, 이때, 첫 번째 KL 항은 보통 해석적으로 계산이 가능하므로 그대로 두고, 두 번째 재구성 항만 샘플링을 통해 추정하여 다음과 같은 SGVB 추정량을 얻습니다.
\[
\tilde{L}_B(\theta, \phi; x) = -D_{\mathrm{KL}}(q_{\phi}(z \mid x) \,\|\, p_{\theta}(z)) + \frac{1}{L} \sum_{l=1}^{L} \log p_{\theta}(x \mid z^{(l)}), \quad z^{(l)} = g_{\phi}(\varepsilon^{(l)}, x)
\]
실험적으로는 \( L = 1 \) 만으로도 충분하며, 미니배치 \( X_M \) (예: \( M = 100 \))에 대해 전체 손실을 다음과 같이 추정합니다.
\[
\tilde{L}_M(\theta, \phi; X_M) = \frac{N}{M} \sum_{i=1}^{M} \tilde{L}(\theta, \phi; x^{(i)})
\]
이후 이 손실에 대해 확률적 그래디언트를 구해 최적화합니다.
\[
\nabla_{\theta, \phi} \tilde{L}_M
\]
이를 SGD, AdaGrad 등의 옵티마이저로 업데이트하면 대규모 데이터셋에서도 빠르고 안정적인 학습이 가능합니다.
이 수식들은 오토인코더 관점에서도 해석됩니다. KL 항은 인코더 \( q_{\phi} \) 가 사전 분포 \( p_{\theta}(z) \) 에 가깝도록 정규화하는 역할을 하며, 재구성 항은 디코더 \( p_{\theta}(x \mid z) \) 의 음의 재구성 오차 역할을 합니다.
- The reparameterization trick
연속 잠재변수 \( z \) 에 대해 근사 사후 \( q_{\phi}(z \mid x) \) 로부터 샘플을 뽑을 때는, 잡음 변수 \( \varepsilon \sim p(\varepsilon) \) 와 미분 가능한 변환 \( g_{\phi} \) 를 사용해 \(z = g_{\phi}(\varepsilon, x)\) 처럼 재매개변수화 하면 편리합니다.
이때 \( q_{\phi}(z \mid x)\,dz = p(\varepsilon)\,d\varepsilon \) 이므로,
\( \int q_{\phi}(z \mid x)\,f(z)\,dz = \int p(\varepsilon)\,f(g_{\phi}(\varepsilon, x))\,d\varepsilon \approx \frac{1}{L} \sum_{l=1}^{L} f(g_{\phi}(\varepsilon^{(l)}, x)), \quad \varepsilon^{(l)} \sim p(\varepsilon) \) 이 몬테카를로 추정치는 \( \phi \) 에 대해 완전히 미분 가능하므로, ELBO를 확률적 경사하강법 SGD로 직접 최적화할 수 있습니다.
대표적 예로, 단변량 정규분포 \( z \sim \mathcal{N}(\mu, \sigma^2) \) 는 \(
z = \mu + \sigma \varepsilon, \quad \varepsilon \sim \mathcal{N}(0, 1)\)로 재매개변수화되어 기대값을 쉽게 추정할 수 있습니다.
📌 Example : Variational Auto-Encoder
위는 가우시안 인코더와 베르누이 디코더 조합으로 이루어진 VAE의 한 예시이고, 본 챕터에서는 확률적 인코더 \(q_{\boldsymbol\phi}(z\mid x)\)를 다층 퍼셉트론(MLP)으로 구성하고, 인코더–디코더 쌍을 AEVB 알고리즘으로 공동 최적화하는 전형적인 VAE 구조를 설명합니다.
사전 분포는 매개변수가 없는 등방 다변량 정규분포 \(p_{\boldsymbol\theta}(z)=\mathcal N(\mathbf 0,\mathbf I)\) 입니다. 디코더 \(p_{\boldsymbol\theta}(x\mid z)\)는 실수형 데이터에 대해서는 다변량 가우시안, 이진 데이터에 대해서는 다변량 베르누이 분포를 사용하며, 두 경우 모두 분포의 파라미터를 잠재 변수 \(z\)를 입력으로 받는 MLP 한 층에서 계산합니다.
실제 사후 \(p_{\boldsymbol\theta}(z\mid x)\)가 대각 공분산 가우시안에 가깝다고 가정하고, 근사 사후를
\(
q_{\boldsymbol\phi}(z\mid x^{(i)})
=\mathcal N\!\bigl(z;\,\boldsymbol\mu^{(i)},
\operatorname{diag}(\boldsymbol\sigma^{(i)})^{2}\bigr)
\)로 둡니다.
인코더 MLP는 데이터 \(x^{(i)}\)로부터 평균 \(\boldsymbol\mu^{(i)}\)와 표준편차 \(\boldsymbol\sigma^{(i)}\)를 출력합니다. 재매개변수화 트릭을 이용하여 \(
z^{(i,l)}=\boldsymbol\mu^{(i)}
+\boldsymbol\sigma^{(i)}\odot\varepsilon^{(l)},\qquad
\varepsilon^{(l)}\sim\mathcal N(\mathbf 0,\mathbf I),
\) (\,\(\odot\)\,는 요소별 곱)과 같이 샘플을 얻습니다.
사전과 근사 사후가 모두 정규분포이므로 KL-Divergence 항은 해석적으로 계산할 수 있습니다. 데이터 포인트 \(x^{(i)}\)에 대한 ELBO의 추정치는 다음과 같습니다.
\[
\begin{aligned}
\mathcal L(\boldsymbol\theta,\boldsymbol\phi; x^{(i)})
\approx
\frac12\sum_{j=1}^{J}
\Bigl(
1+\log\bigl(\sigma^{(i)}_{j}\bigr)^{2}
-\bigl(\mu^{(i)}_{j}\bigr)^{2}
-\bigl(\sigma^{(i)}_{j}\bigr)^{2}
\Bigr)
+\frac1L\sum_{l=1}^{L}
\log p_{\boldsymbol\theta}\!\bigl(x^{(i)}\mid z^{(i,l)}\bigr),
\end{aligned}
\]
\[
z^{(i,l)}=\boldsymbol\mu^{(i)}
+\boldsymbol\sigma^{(i)}\odot\varepsilon^{(l)},\qquad
\varepsilon^{(l)}\sim\mathcal N(\mathbf 0,\mathbf I).
\]
여기서 첫 번째 합은 KL-Divergnece에 해당하며, 두 번째 합은 디코더의 로그우도, 즉 음의 재구성 오차를 평균한 값입니다.
두 항 모두 \(\boldsymbol\theta\)와 \(\boldsymbol\phi\)에 대해 완전히 미분 가능하므로, 미니배치 확률적 경사하강법(SGD, AdaGrad 등)을 통해 대규모 데이터셋에서도 빠르고 안정적으로 최적화할 수 있습니다.
📌 Related Work
웨이크-슬립 알고리즘은 인코더와 디코더를 따로 최적화해 연속·이산 잠재변수를 모두 처리할 수 있지만, 두 목적식이 ELBO를 직접 높이지 못한다는 한계가 있습니다. 반면 AEVB는 SGVB 추정량과 재파라미터화 트릭을 도입해 단일 ELBO를 미니배치 확률적 경사하강법으로 곧바로 최적화하므로, 연속 잠재변수를 지닌 생성 모델을 대규모 데이터에서도 빠르고 안정적으로 학습할 수 있습니다.
최근 SVI 연구들은 제어변수 기법이나 유사한 재파라미터화 기법들을 통해 그래디언트 분산을 줄이며 효율을 높여 왔고, 이는 AEVB와 이론적 기반을 공유합니다. AEVB의 ELBO에는 정규화 항이 자동 포함돼 있어 오토인코더 계열이 별도로 사용하는 하이퍼파라미터 없이도 유용한 표현을 학습하며, PCA, 노이즈·스파스 오토인코더, GSN, DARN 등과 개념적 연관은 있지만 연속 잠재변수와 일반 방향성 확률 모델로의 범용적 확장성에서 차별화되어집니다.
📌 Experiments
MNIST 숫자와 Frey Face 얼굴 이미지를 대상으로, 인코더와 디코더의 은닉층 크기를 동일하게 둔 대칭 구조 VAE를 학습하여 ELBO와 marginal likelihood로 성능을 비교하였습니다.
Frey Face처럼 픽셀이 연속값을 갖는 데이터에는 디코더를 가우시안 출력으로 설정하고, 평균값에 시그모이드를 적용하여 0–1 범위로 제한함으로써 안정성을 확보하였습니다.
잠재 공간의 차원을 2로 제한한 결과, 학습된 인코더는 고차원 이미지를 2차원 잠재 공간으로 부드럽게 투사하였습니다. 이로써 MNIST에서는 숫자 ‘0’과 ‘8’처럼 형태가 유사한 클래스가 인접하게 배치되었고, Frey Face에서는 웃음이나 놀람 같은 표정 변화가 하나의 연속적인 궤적으로 나타났습니다.
웨이크-슬립 알고리즘은 인코더와 디코더를 별도로 최적화하여 연속·이산 잠재변수를 모두 처리할 수 있으나, 두 목적식이 ELBO를 직접 향상시키지 못한다는 한계가 있습니다. 반면 AEVB는 SGVB 추정량과 재매개변수화 트릭을 도입하여 단일 ELBO를 미니배치 확률적 경사하강법으로 바로 최적화합니다. 실험은 이러한 AEVB 접근이 대규모 데이터에서도 빠르고 안정적으로 학습된 근사 사후와 높은 주변우도를 달성함을 확인시켜 주었습니다.
추가적으로, 다양한 학습 데이터 개수에 따른 marginal likelihood 추정 측면에서 AEVB를 wake-sleep 알고리즘 및 Monte Carlo EM과 비교했습니다.
📌 Conclusion
SGVB는 VB의 ELBO 최적화를 미니배치 확률적 경사하강법으로 수행해, 연속 잠재변수를 가진 복잡한 모델도 대용량 데이터에서 빠르게 근사 추론합니다. 이를 활용한 AEVB 알고리즘은 VAE처럼 인코더와 디코더를 동시에 학습해 효율적인 추론, 생성을 가능하게 합니다.
📌 느낀점(논문을 읽으면서..)
논문을 읽으면서 통계적 베이스가 많이 부족하구나... 느끼며 한번 더 반성을 하게 되는 경험을 했습니다.. 전반적인 VAE의 목표와 방법론, 모델의 전반적인 구조 정도를 이해한 것 같은데, 생성 분야가 매우 어려운 분야라는 것을 느꼈고, 통계적 이론들에 익숙해질 수 있게끔 연습을 많이 해야할 것 같습니다! 두서 없는 리뷰 읽어주셔서 감사합니다! 질문이나 훈수 댓글에 달아주시면 보고 저도 조금 더 고민하며 정리해보도록 할게요!
📌References
https://www.youtube.com/watch?v=GbCAwVVKaHY