ELBO: 못 구하는 로그가능도를 아래에서 받치기

VAE의 손실 함수는 언제나 두 항의 합입니다.

recon = F.mse_loss(x_hat, x, reduction='sum')
kld = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
loss = recon + kld

첫 항은 복원이 잘 됐는지를 재고, 둘째 항은 잠재 분포가 표준정규분포에서 얼마나 벗어났는지를 잽니다. 처음 보면 정규화 항을 하나 얹은 것처럼 읽힙니다 — 재구성만 하면 과적합할 테니 벌점을 붙였다는 식으로요.

그렇게 읽으면 답할 수 없는 질문이 생깁니다. 왜 하필 계수가 1인가요. 왜 recon + 0.3 * kld가 아닌가요. 그리고 확산 모델의 손실은 그냥 잡음의 MSE 하나뿐인데, 그것과 이 두 항짜리 손실은 무슨 관계인가요.

두 질문의 답은 같습니다. 저 합은 정규화가 붙은 재구성 손실이 아니라, 구할 수 없는 값 하나를 아래에서 받치는 하한입니다. 계수가 1인 것은 취향이 아니라 유도의 결과입니다.

지난 글에서 표집을 계산 그래프 안에 남기는 장치를 봤다면, 이 글은 그 장치가 무엇을 위해 필요했는지를 봅니다. 목적지는 한 줄입니다.

$$\log p(x) = \mathrm{ELBO}(q) + \mathrm{KL}\big(q(z) \,\|\, p(z \mid x)\big)$$

구할 수 없는 값

잠재변수 모형은 데이터를 이렇게 설명합니다 — 먼저 보이지 않는 $z$ 를 뽑고, 그것으로부터 $x$ 를 만든다. 그러면 관측된 $x$ 의 확률은 모든 $z$ 를 훑어 더한 값입니다.

$$p(x) = \int p(x \mid z)\, p(z)\, dz$$

최대가능도가 시키는 대로 $\log p(x)$ 를 최대화하고 싶은데, 이 적분이 막혀 있습니다. $z$ 가 256차원 연속 벡터이고 $p(x \mid z)$ 가 신경망이면 닫힌 형태로 적을 방법이 없고, 표본으로 근사하려 해도 $p(z)$ 에서 뽑은 $z$ 는 거의 전부 $p(x \mid z) \approx 0$ 인 자리라 아무것도 얻지 못합니다.

그래서 우회합니다. 아무 분포 $q(z)$ 를 하나 들여옵니다 — 변분 분포(variational distribution)라 부르고, 우리가 마음대로 고를 수 있는 후보입니다. 이제 적분 안에 $q$ 를 곱하고 나눕니다.

$$\log p(x) = \log \int q(z)\, \frac{p(x, z)}{q(z)}\, dz = \log\, \mathbb{E}_{z \sim q}\!\left[\frac{p(x,z)}{q(z)}\right]$$

여기까지는 항등식입니다. 그리고 여기서 두 갈래가 열립니다.

경로 A · 젠센 부등식

로그는 위로 볼록한 함수이고, 젠센 부등식은 그런 함수에 대해 기댓값과 함수의 순서를 바꿀 때 부등호가 생긴다고 말합니다. $\log \mathbb{E}[Y] \ge \mathbb{E}[\log Y]$ 입니다.

그대로 적용합니다.

$$\log p(x) = \log\, \mathbb{E}_{q}\!\left[\frac{p(x,z)}{q(z)}\right] \;\ge\; \mathbb{E}_{q}\!\left[\log \frac{p(x,z)}{q(z)}\right]$$

오른쪽이 ELBO(evidence lower bound), 우리말로 증거 하한입니다. 관측된 $x$ 의 로그가능도를 통계학에서 증거(evidence)라 부르고, 그것을 아래에서 받치는 값이라는 뜻입니다.

$$\mathrm{ELBO}(q) = \mathbb{E}_{q}\big[\log p(x, z)\big] - \mathbb{E}_{q}\big[\log q(z)\big]$$

세 줄로 끝났습니다. 그런데 이 경로가 알려 주는 것은 "하한이다"까지입니다. 얼마나 헐거운 하한인지, 언제 등호가 성립하는지는 말해 주지 않습니다.

경로 B · KL을 더하고 빼기

두 번째 길은 부등식을 아예 쓰지 않습니다. 베이즈 정리 $p(x, z) = p(z \mid x)\,p(x)$ 를 ELBO 안에 넣어 봅니다.

$$\begin{aligned} \mathrm{ELBO}(q) &= \mathbb{E}_{q}\!\left[\log \frac{p(z \mid x)\, p(x)}{q(z)}\right] \\ &= \mathbb{E}_{q}\!\left[\log \frac{p(z \mid x)}{q(z)}\right] + \mathbb{E}_{q}\big[\log p(x)\big] \\ &= -\,\mathrm{KL}\big(q(z) \,\|\, p(z \mid x)\big) + \log p(x) \end{aligned}$$

둘째 줄에서 $\log p(x)$ 는 $z$ 와 무관하므로 기댓값을 그냥 통과합니다. 그리고 남은 항은 정확히 KL 발산의 정의에 마이너스를 붙인 것입니다. 옮겨 적으면

$$\boxed{\;\log p(x) = \mathrm{ELBO}(q) + \mathrm{KL}\big(q(z) \,\|\, p(z \mid x)\big)\;}$$

이것이 이 글의 한 줄입니다. 부등식이 아니라 등식이고, 세 가지를 한꺼번에 알려 줍니다.

1. KL은 항상 0 이상이므로 ELBO는 하한이다 — 경로 A의 결론이 여기서 따라 나옵니다.
2. 간극의 정체가 KL이다. 하한이 헐거운 정확한 이유는 $q$ 가 진짜 사후분포 $p(z \mid x)$ 에서 떨어져 있기 때문입니다.
3. 왼쪽은 $q$ 와 무관합니다. $q$ 를 아무리 바꿔도 $\log p(x)$ 는 그대로이므로, ELBO를 키우는 것은 KL을 정확히 그만큼 줄이는 것입니다.

세 번째가 이 방법의 정당성 전체입니다. 못 구하는 값을 최대화하는 대신 구할 수 있는 하한을 최대화하는데, 그 하한을 올리는 일이 곧 하한과 참값 사이의 거리를 좁히는 일입니다.

숫자로 확인한다

식이 맞는지 잠재변수가 둘뿐인 장난감 모형에서 직접 재 봅니다. $p(z=1) = 0.3$, $x \mid z \sim \mathcal{N}(\mu_z, 1)$ 에 $\mu_0 = 0, \mu_1 = 3$ 이고 관측값은 $x = 2$ 입니다. $z$ 가 둘뿐이니 $\log p(x)$ 도 사후분포도 손으로 정확히 구할 수 있습니다.

$$\log p(x) = -2.203782, \qquad p(z=1 \mid x) = 0.657619$$

이제 $q(z=1)$ 을 여러 값으로 바꿔 가며 ELBO와 KL을 잽니다.

$q(z=1)$	ELBO	KL(q‖사후분포)	합
0.05	−3.044463	0.840681	−2.203782
0.20	−2.644671	0.440889	−2.203782
0.50	−2.256115	0.052333	−2.203782
0.657619	−2.203782	0.000000	−2.203782
0.80	−2.253049	0.049267	−2.203782
0.95	−2.457031	0.253249	−2.203782

오른쪽 열이 소수 여섯째 자리까지 한 값으로 고정돼 있습니다. $q$ 를 어디에 두든 두 조각의 합은 $\log p(x)$ 이고, 경계만 움직입니다.

곡선이 천장에 닿는 자리는 $q$ 가 사후분포와 같아지는 한 점뿐입니다. 그 점에서 하한이 참값과 정확히 같아집니다 — 근사가 아니라 등식입니다.

두 항으로 갈라 읽기

이제 처음의 코드로 돌아갑니다. ELBO를 다시 묶어 봅니다. $p(x, z) = p(x \mid z) p(z)$ 를 넣으면

$$\begin{aligned} \mathrm{ELBO}(q) &= \mathbb{E}_q\!\left[\log \frac{p(x \mid z)\,p(z)}{q(z)}\right] \\ &= \underbrace{\mathbb{E}_q\big[\log p(x \mid z)\big]}_{\text{재구성 항}} \;-\; \underbrace{\mathrm{KL}\big(q(z) \,\|\, p(z)\big)}_{\text{정규화 항}} \end{aligned}$$

같은 ELBO를 두 번 다르게 묶었을 뿐인데 나오는 그림이 다릅니다. 앞의 묶음은 $q$ 를 사후분포와 비교했고, 이 묶음은 $q$ 를 사전분포와 비교합니다.

• 재구성 항은 $q$ 가 내놓는 $z$ 로 $x$ 를 얼마나 잘 되살리는지입니다. 가우시안 디코더를 쓰면 이것이 곧 MSE입니다.
• 정규화 항은 $q$ 가 사전분포에서 얼마나 벗어났는지입니다. 표준정규 사전분포에 대각 가우시안 $q$ 를 쓰면 닫힌 형태로 나오고, 그것이 위 코드의 kld 한 줄입니다.

그러니 계수 1은 선택이 아닙니다. 저 둘은 원래 한 식이었고, 나눠 적으면 우연히 재구성과 정규화처럼 보이는 것뿐입니다. 계수를 바꾸는 순간 그것은 더 이상 $\log p(x)$ 의 하한이 아닙니다 — 실무에서 쓰는 $\beta$-VAE는 그 사실을 알고 일부러 하한을 포기하는 쪽입니다.

두 항의 최적점이 다르다는 것도 눈으로 보입니다. 재구성만 보면 $x = 2$ 를 잘 설명하는 $z = 1$ 쪽으로 $q$ 를 끝까지 밀고, KL만 보면 사전분포 0.3으로 끌어당깁니다.

$q(z=1)$	재구성 항	KL(q‖사전분포)	ELBO
0.05	−2.8439	0.2005	−3.0445
0.30	−2.4689	0.0000	−2.4689
0.657619	−1.9325	0.2713	−2.2038
0.95	−1.4939	0.9631	−2.4570

$q = 0.3$ 에서 KL이 0이고, $q = 0.95$ 에서 재구성이 가장 좋습니다. 그런데 ELBO의 최고점은 그 사이 0.658입니다. 사후분포는 정확히 그 줄다리기의 균형점입니다. 손실에 두 항이 나란히 있는 이유가 이것입니다 — 어느 하나를 이기게 하면 사후분포에서 멀어지고, 그만큼 하한이 헐거워집니다.

하한이 얼마나 헐거운지는 무엇이 정하는가

간극이 $\mathrm{KL}(q \| p(z \mid x))$ 라는 것을 알았으니, 헐거움의 원인도 그 KL을 크게 만드는 것들입니다. 실제 모형에서는 셋이 겹쳐 있습니다.

• $q$ 를 고르는 족이 좁다. 대각 공분산 가우시안으로만 $q$ 를 잡으면, 진짜 사후분포가 축이 기운 모양이거나 봉우리가 둘이면 아무리 잘 맞춰도 KL이 0이 되지 않습니다. 위 그림에서 곡선의 최고점이 천장 아래에 머무는 상황입니다.
• 입력마다 $q$ 를 최적화하지 않는다. VAE는 $x$ 를 받아 $q$ 의 파라미터를 뱉는 인코더 하나를 학습합니다. 표본마다 따로 최적화하는 대신 함수 하나로 때우는 것이라, 그 차이만큼 간극이 더 벌어집니다.
• 모형 자체가 데이터를 설명하지 못한다. 이것은 하한의 문제가 아니라 천장이 낮은 것이라 $q$ 를 고쳐도 안 낫습니다.

앞의 둘은 $q$ 쪽 문제이고 마지막 하나는 $p$ 쪽 문제인데, 학습 중에는 ELBO 값 하나만 보이므로 셋이 구별되지 않습니다. ELBO가 낮다는 사실만으로 모델이 나쁘다고 말할 수 없는 이유가 여기 있습니다.

같은 뼈대 위의 셋

이 분해가 확산 모델까지 이어집니다. 확산의 잠재변수는 $z_1, \dots, z_T$ 가 아니라 $x_1, \dots, x_T$ 라는 이름을 달고 있을 뿐이고, 변분 분포 $q$ 는 학습하는 것이 아니라 forward 과정으로 이미 정해져 있습니다. 그래서 같은 ELBO를 적고 항을 정리하면 스텝마다의 KL 항들이 나오고, 그것이 결국 잡음 예측의 MSE 하나로 줄어듭니다. 확산 손실이 한 줄인 것은 뼈대가 달라서가 아니라 $q$ 가 고정이라 인코더를 학습할 필요가 없기 때문입니다.

EM 알고리즘도 같은 식의 다른 읽기입니다. $\log p(x) = \mathrm{ELBO}(q) + \mathrm{KL}(q \| p(z \mid x))$ 에서

• E 단계는 $\theta$ 를 고정하고 $q$ 를 사후분포로 놓습니다 — KL을 0으로 만들어 하한을 천장에 붙입니다.
• M 단계는 $q$ 를 고정하고 $\theta$ 로 ELBO를 최대화합니다 — 천장 자체를 올립니다.

두 단계를 번갈아 하는 것이 왜 $\log p(x)$ 를 단조증가시키는지가 이 한 줄에서 곧바로 나옵니다. E 단계 직후 하한과 참값이 같으므로, M 단계가 하한을 올리면 참값도 최소한 그만큼 올라갑니다.

VAE는 사후분포를 정확히 구할 수 없어 E 단계를 인코더로 근사한 것이고, 확산은 $q$ 를 아예 손으로 정해 놓은 것입니다. 셋의 차이는 $q$ 를 어떻게 조달하느냐 하나뿐입니다.

recon + kld로 돌아가면 이제 계수가 1인 이유를 한 문장으로 말할 수 있습니다 — 저 둘은 서로 다른 목적의 두 손실이 아니라, $\log p(x)$ 를 아래에서 받치는 한 개의 식을 두 조각으로 적은 것이기 때문입니다.

정리

• 잠재변수 모형의 $\log p(x) = \log \int p(x \mid z)p(z)dz$ 는 적분이 막혀 직접 최대화할 수 없다. 그래서 마음대로 고를 수 있는 변분 분포 $q(z)$ 를 들여온다.
• 젠센 부등식 경로는 $\log \mathbb{E}[\cdot] \ge \mathbb{E}[\log \cdot]$ 로 세 줄 만에 $\mathrm{ELBO} = \mathbb{E}_q[\log p(x,z)] - \mathbb{E}_q[\log q(z)]$ 를 얻지만, 간극에 대해서는 아무것도 말하지 않는다.
• KL 경로는 베이즈 정리를 넣어 $\log p(x) = \mathrm{ELBO}(q) + \mathrm{KL}(q \| p(z \mid x))$ 라는 등식을 얻는다. 하한이라는 사실과 간극의 정체가 함께 나온다.
• 왼쪽이 $q$ 와 무관하므로 ELBO를 키우는 것은 KL을 정확히 그만큼 줄이는 것이다. 이것이 하한을 최대화하는 일이 정당한 이유 전체다.
• 장난감 모형에서 $q$ 를 여섯 값으로 바꿔 재 보니 ELBO + KL이 소수 여섯째 자리까지 $-2.203782$ 로 고정이었고, $q$ 가 사후분포 0.657619 일 때 간극이 정확히 0이었다.
• 같은 ELBO를 다르게 묶으면 재구성 항 − KL(q‖사전분포) 가 된다. 두 항의 계수가 1인 것은 선택이 아니라 유도의 결과이고, 계수를 바꾸면 더 이상 하한이 아니다.
• 재구성은 $q$ 를 데이터 쪽으로, KL은 사전분포 쪽으로 당긴다. 장난감 모형에서 각각의 최적점이 0.95와 0.3인데 ELBO의 최고점은 그 사이 0.658 — 사후분포가 그 줄다리기의 균형점이다.
• 하한이 헐거운 원인은 셋이다 — $q$ 의 족이 좁은 것, 입력마다 최적화하지 않고 인코더 하나로 때우는 것, 그리고 모형 자체가 데이터를 설명하지 못하는 것. ELBO 값 하나로는 셋이 구별되지 않는다.
• VAE·확산·EM은 같은 한 줄 위에 있고 차이는 $q$ 의 조달 방식뿐이다 — 인코더로 근사하거나, forward 과정으로 고정하거나, E 단계에서 사후분포로 정확히 놓는다.

---

읽어주셔서 감사합니다. 😊