5장. 해부: 변분 오토인코더

변분 오토인코더(VAE)는 2013년 논문 “Auto-Encoding Variational Bayes”(Kingma, Welling)에서 나왔고, 여기서 정리된 ELBO(evidence lower bound)는 확산 모델, 잠재 확산 모델, 세계 모델, 베이즈 딥러닝의 공통 언어가 되었습니다. 이 장에서는 ELBO를 처음부터 유도하고, 논문에 나오는 정규분포 KL의 닫힌 꼴과 재매개변수화 기법을 해부합니다. 필요한 것은 베이즈 정리(확률과 통계 10장), KL 발산(11장), 옌센 부등식입니다.

5.1 잠재 변수 모형

관측 $\mathbf x$(그림)가 보이지 않는 잠재 변수(latent variable) $\mathbf z$(그림의 “요인”: 모양, 색, 각도)에서 만들어진다고 봅니다.

• $p(\mathbf z)$: 사전분포(prior).
• $p_\theta(\mathbf x \mid \mathbf z)$: 디코더(decoder), 신경망.

문제. 최대우도(확률과 통계 7장)로 $\theta$를 학습하려면 $\ln p_\theta(\mathbf x)$가 필요한데, 고차원 $\mathbf z$에 대한 적분이라 계산할 수 없습니다. 사후분포 $p_\theta(\mathbf z \mid \mathbf x) = \frac{p_\theta(\mathbf x \mid \mathbf z)p(\mathbf z)}{p_\theta(\mathbf x)}$도 분모 때문에 계산할 수 없습니다(10장에서 말한 “증거 적분”의 어려움).

5.2 근사 사후분포

정의 5.1. 다룰 수 있는 분포족으로 사후분포를 흉내 내는 $q_\phi(\mathbf z \mid \mathbf x) = \mathcal N(\boldsymbol\mu_\phi(\mathbf x), \mathrm{diag}(\boldsymbol\sigma^2_\phi(\mathbf x)))$를 근사 사후분포라 하고, 이를 내는 신경망을 인코더(encoder)라 합니다.

이렇게 적분을 최적화로 바꾸는 방법을 변분 추론(variational inference)이라 합니다. 이름의 “변분”이 여기서 왔습니다.

5.3 ELBO 유도

정리 5.2. 모든 $q$에 대해

증명. 적분 안에 $q/q$를 곱하면

부등호는 $\ln$이 오목함수라서 성립하는 옌센 부등식(Jensen's inequality) $\ln\mathbb E[Y] \ge \mathbb E[\ln Y]$입니다. 오른쪽의 로그를 나누면 $\mathbb E_q[\ln p_\theta(\mathbf x \mid \mathbf z)] + \mathbb E_q[\ln\frac{p(\mathbf z)}{q}]$이고, 둘째 항이 $-D_{\mathrm{KL}}(q\Vert p)$입니다. □

정리 5.3 (간격). $\ln p_\theta(\mathbf x) - \mathrm{ELBO} = D_{\mathrm{KL}}\!\left(q_\phi(\mathbf z \mid \mathbf x)\,\Vert\,p_\theta(\mathbf z \mid \mathbf x)\right) \ge 0$.

이 하한을 증거 하한(evidence lower bound, ELBO)이라 합니다. 그러므로 ELBO를 $\phi$에 대해 키우는 것은 근사 사후분포를 진짜 사후분포에 가깝게 하는 것이고, $\theta$에 대해 키우는 것은 가능도의 하한을 키우는 것입니다. 둘을 함께 경사상승법으로 학습합니다.

5.4 두 항의 뜻

• 재구성 항 $\mathbb E_q[\ln p_\theta(\mathbf x \mid \mathbf z)]$: 인코딩한 $\mathbf z$로 원래 $\mathbf x$를 얼마나 잘 되살리는가. 디코더가 정규분포이면 음의 제곱오차(확률과 통계 7장), 베르누이이면 음의 이진 교차 엔트로피.
• KL 항 $D_{\mathrm{KL}}(q\Vert p)$: 인코딩한 분포가 사전분포 $\mathcal N(\mathbf 0, I)$에서 얼마나 벗어났는가. 잠재 공간을 정돈해 두어, 나중에 $\mathbf z \sim \mathcal N(\mathbf 0, I)$를 뽑아 디코더에 넣으면 그럴듯한 그림이 나오게 합니다.

질문 3 (극단값). KL 항이 없으면 보통의 오토인코더처럼 잠재 공간이 엉망이 되고, KL 항만 있으면 $q = p$가 되어 $\mathbf z$에 정보가 없습니다(사후분포 붕괴, posterior collapse). β-VAE는 KL 항에 $\beta$를 곱해 이 균형을 조절합니다.

5.5 정규분포 KL의 닫힌 꼴

논문 부록의 식을 1차원에서 유도합니다.

정리 5.4. $D_{\mathrm{KL}}\!\left(\mathcal N(\mu, \sigma^2)\,\Vert\,\mathcal N(0, 1)\right) = \frac12\left(\mu^2 + \sigma^2 - \ln\sigma^2 - 1\right)$.

증명. $\ln\frac{q}{p} = -\ln\sigma - \frac{(z - \mu)^2}{2\sigma^2} + \frac{z^2}{2}$. $q$에 대한 기댓값을 취하면 $\mathbb E[(z - \mu)^2] = \sigma^2$, $\mathbb E[z^2] = \mu^2 + \sigma^2$이므로 $-\ln\sigma - \frac12 + \frac{\mu^2 + \sigma^2}{2}$. □

잠재 변수가 $d$차원이고 성분이 독립이면 성분별 KL을 더합니다. 질문 4 (특수한 경우). $\mu = 0$, $\sigma = 1$이면 $0$: 같은 분포끼리의 KL은 0입니다.

5.6 재매개변수화 기법

ELBO의 기댓값은 $q_\phi$에서 뽑은 표본으로 어림하는데, “뽑는다”는 연산은 $\phi$로 미분할 수 없습니다. 그래서 무작위성을 밖으로 빼냅니다.

이제 $\mathbf z$는 $\phi$의 미분가능한 함수이고, 무작위성은 $\phi$와 무관한 $\boldsymbol\epsilon$에만 있습니다. 그래디언트가 연쇄법칙(다변수 미적분 3장)으로 인코더까지 흐릅니다. 이를 재매개변수화 기법(reparameterization trick)이라 하며, 4장의 $\mathbf x_t = \sqrt{\bar\alpha_t}\mathbf x_0 + \sqrt{1 - \bar\alpha_t}\boldsymbol\epsilon$도 같은 꼴입니다.

5.7 이 장의 요약

• 잠재 변수 모형의 가능도는 적분이라 계산 불가 → 근사 사후분포 $q_\phi$.
• 옌센 부등식으로 ELBO $= \mathbb E_q[\ln p(\mathbf x \mid \mathbf z)] - D_{\mathrm{KL}}(q\Vert p) \le \ln p(\mathbf x)$. 간격은 진짜 사후분포와의 KL.
• 정규분포 KL: $\frac12(\mu^2 + \sigma^2 - \ln\sigma^2 - 1)$.
• 재매개변수화: $\mathbf z = \boldsymbol\mu + \boldsymbol\sigma \odot \boldsymbol\epsilon$로 미분 가능하게.

쓰이는 곳. 잠재 확산 모델(Stable Diffusion은 VAE의 잠재 공간에서 확산), 확산 모델의 학습 목표 유도(4장의 손실은 ELBO를 정리한 것), 세계 모델과 강화학습, 베이즈 신경망, 음성·분자 생성, 표현 학습의 분리(disentanglement) 연구가 이 장의 식을 씁니다.

문제

0.5

0.8069

1.4

1