[Diffusion] ELBO

Diffusion 이해를 위한 사전 지식 3

들어가기 전

생성 모델이 하고 싶은 일은 결국 실제 데이터분포 $p_{\text{data}}(x)$를 알아내는 것이다. 하지만 이 분포는 직접 볼 수 없고, 실제 다룰 수 있는 것은 그 분포에서 뽑힌 샘플(데이터셋)이다. 따라서 파라미터 $\theta$를 가진 모델 $p_\theta(x)$를 세우고 이 모델이 실제 분포를 최대한 닮도록 만드는 방향으로 접근한다.

앞서 본 KL divergence로 두 분포의 차이를 측정할 수 있고, 실제 데이터 분포와 모델 사이의 KL을 줄이는 것은 데이터에 대한 로그 우도를 최대화하는 것과 같다. 그래서 학습의 목표는 데이터 하나에 대해 $\log p_\theta(x)$를 크게 만드는 것이 된다.

문제는 복잡한 데이터를 표현하기 위해 모델에 눈에 보이지 않는 잠재변수 $z$를 도입하면, $p_\theta(x)$는 가능한 모든 $z$에 대해 적분해야만 얻을 수 있는 값이 된다. $z$가 고차원이면 이 적분은 계산할 수 없고 최대화도 할 수 없다.

이때 앞서 살펴본 importance sampling을 사용한다. 직접 계산하기 어려운 적분을 샘플링하기 쉬운 다른 분포 $q(z\mid x)$에 대한 기댓값으로 바꿔 쓰는 방법이다. 이를 통해 $p_\theta(x)$를 q에서 뽑은 샘플들로 추정할 수 있는 형태가 된다. 다만 우리가 원하는 것은 $p_\theta(x)$이 아닌 $\log p_\theta(x)$이고 log가 기댓값 바깥에 있으면 다루기 어렵다. 따라서 log의 오목성을 이용한 Jensen 부등식으로 log를 기댓값 안에 넣는다. 그 대가로 등식은 부등식이 되고, 정확한 값 대신 항상 그보다 작거나 같은 값(하한)을 얻는다.

이 하한이 ELBO이다. 로그 우도 함수를 직접 구할 수 없을 때 대신 최대화할 수 있도록 만든 계산 가능한 하한이다. 이 하한이 $\log p_\theta(x)$와 얼마나 차이 나는지는 KL divergence로 정확하게 표현된다. 여기서 KL은 실제 데이터 분포와의 차이가 아닌, 우리가 도입한 근사 분포 $q(z\mid x)$가 모델이 정의하는 사후분포 $p_\theta(z\mid x)$와 얼마나 다른지를 나타낸다.

정의

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

위에서 언급한 것 처럼 생성 모델의 목표는 로그 우도 $\log p_\theta(x)$를 최대화하는 것이다. 잠재변수 $z$를 도입한 모델에서는 위와 같은 적분으로 표현된다($p_\theta(x\mid z)$의 가중평균). 잠재변수 $z$가 고차원이면 위 식을 계산할 수 없기 때문에 계산 가능한 하한을 만들어 그것을 최대화한다.

ELBO(Evidence Lower Bound)이 하한이다. 여기서 “evidence”는 $\log p_\theta(x)$를 가리키는 말이고, ELBO는 그 아래에 놓이는 값이라는 뜻이다.

유도(Jensen 부등식)

1.importance sampling

$$
p_\theta(x) = \int p_\theta(x, z)\,dz =\int q(z\mid x)\,\frac{p_\theta(x,z)}{q(z\mid x)}\,dz = \mathbb{E}_{q(z\mid x)}\left[\frac{p_{\theta}(x,z)}{q(z\mid x)}\right]
$$

적분을 직접 못 하면 샘플링으로 근사하는 것이 자연스러운 방법인데, 이를 위해서는 적분이 어떤 분포에 대한 기댓값 형태로 쓰여 있어야 한다. 그래서 샘플링하기 쉬운 분포 분포 $q(z\mid x)$를 곱하고 다시 나눈다. 1을 곱한 것과 같으므로 값은 전혀 변하지 않는다.(importance sampling)

$p_\theta(x)$은 $p$에서 $z$를 뽑고, 각 $z$에서 $\frac{p_\theta(x,z)}{q(z\mid x)}$를 계산해 평균 낸 값이 된다. 이 비율이 importance weight $w$이고 q가 원래 분포와 다른 곳에서 샘플을 뽑았기 때문에, 그 차이만큼 각 샘플의 기여도를 보정해 주는 역할을 한다.

2. log를 씌우면 생기는 문제

$$
\log p_\theta(x) = \log \mathbb{E}_{q(z\mid x)}\big[w(z)\big]
$$

학습의 목표는 $p_\theta(x)$가 아니라 $\log p_\theta(x)$를 최대화하는 것이므로 양변에 log를 씌운다. 여기서 문제는 log가 기댓값 바깥에 있다. 기댓값은 샘플 평균으로 추정이 가능하지만 샘플에 log를 씌운값을 평균적으로 정확히 맞히지 못하는 편향된 추정량이 된다. 고차원에서는 importance weight는 극단적으로 크거나 작은 값을 갖기 쉬워 수치적으로도 불안정하다.

반대로 log가 기댓값 안에 있으면 $\log w(z)$를 샘플마다 계산해 평균 내기만 하면 되므로 추정과 미분이 모두 간단해진다.

3. Jensen 부등식

$$
\log \mathbb{E}[Y] \;\ge\; \mathbb{E}[\log Y]
$$

$\log$가 오목함수이므로 Jensen 부등식으로 위 식이 항상 성립한다. 직관적으로 보면 평균을 낸 뒤 log를 씌우는 것이 log를 씌운 뒤 평균을 내는 것보다 크다.

$$
\log p_{\theta}(x) \;\ge\; \mathbb{E}_{q(z\mid x)}\left[\log\frac{p_{\theta}(x,z)}{q(z\mid x)}\right] = \text{ELBO}
$$

결국 1.importance sampling의 식을 Jensen 부등식에 적용하면 위와 같은 식을 얻을 수 있다.

즉 ELBO는 importance weight의 로그를 $q$에 대해 평균 낸 값이다. 계산 가능한 형태를 얻은 대신 등식이 부등식이 되어 정확한 값이 아닌 하한을 얻게 된다.

4. KL 분해

$$
\log p_\theta(x) = \text{ELBO} + D_{KL}\big(q(z\mid x)\,\|\,p_\theta(z\mid x)\big)
$$

Jensen 부등식으로는 “하한이다”라는 사실만 알 수 있다. 위 항등식을 쓰면 그 차이가 정확히 무엇인지 알 수 있다.

KL은 항상 0 이상이므로 ELBO는 하한이다. 또한 $q$가 실제 사후분포 $p_\theta(z\mid x)$와 같아질수록 하한이 실제 값과 일치(tight)해진다. 따라서 ELBO를 최대화하면 두 가지 효과가 동시에 생긴다.

• $\theta$에 대해 최대화하면 데이터 우도가 커진다.
• $q$에 대해 최대화하면 $q$가 실제 사후분포에 가까워진다.

해석

$p_\theta(x,z) = p_\theta(x\mid z)p(z)$로 풀어 쓰면 ELBO는 두 항으로 나뉜다.

• 유도

• ELBO에 대입

$$
p_\theta(x,z) = p_\theta(x\mid z)\,p(z), \; \text{ELBO} = \mathbb{E}_{q(z\mid x)}\left[\log\frac{p_{\theta}(x,z)}{q(z\mid x)}\right]
$$

$$
\text{ELBO} = \mathbb{E}_{q(z\mid x)}\left[\log\frac{p_{\theta}(x\mid z)\,p(z)}{q(z\mid x)}\right]
$$

• log를 사용해 두 항으로 나눔

$$
\text{ELBO} = \mathbb{E}_{q(z\mid x)}\big[\log p_{\theta}(x\mid z)\big] + \mathbb{E}_{q(z\mid x)}\left[\log\frac{p(z)}{q(z\mid x)}\right]
$$

• 두번째 항을 KL로 정리

$$
D_{KL}\big(q\,\|\,p\big) = \mathbb{E}_{q}\left[\log\frac{q}{p}\right]
$$
$$
\mathbb{E}_{q(z\mid x)}\left[\log\frac{p(z)}{q(z\mid x)}\right] = -\mathbb{E}_{q(z\mid x)}\left[\log\frac{q(z\mid x)}{p(z)}\right] = -D_{KL}\big(q(z\mid x)\Vert p(z)\big)
$$
$$
\text{ELBO} = \underbrace{\mathbb{E}{q(z\mid x)}[\log p\theta(x\mid z)]}_{\text{재구성 (reconstruction)}} - \underbrace{D{KL}\big(q(z\mid x)\,\|\,p(z)\big)}_{\text{prior와의 정합 (regularization)}}
$$

$$
\text{ELBO} = \underbrace{\mathbb{E}_{q(z\mid x)}[\log p_{\theta}(x\mid z)]}_{\text{재구성 (reconstruction)}} - \underbrace{D_{KL}\big(q(z\mid x)\,\|\,p(z)\big)}_{\text{prior와의 정합 (regularization)}}
$$

첫 번째 항은 잠재변수 $z$로부터 $x$를 잘 복구하기 위한 항이고, 두 번째 항은 $q$가 prior 분포에서 너무 벗어나지 않기 위한 항이다. VAE의 loss가 위 형태를 따른다.