β-VAE: KL 항의 계수는 온도 자리에 선다

VAE를 쓰다 보면 KL 항 앞에 계수 β를 붙인 변형을 만난다. Higgins 외(2017)의 β-VAE는 이 β를 1보다 크게 잡아 잠재변수가 더 잘 분리되도록 했다. 여기서 잠재변수가 분리된다는 것은 잠재변수 하나가 데이터를 만든 요인 하나, 이를테면 3차원 물체 그림의 위치나 크기, 조명, 색 가운데 하나에만 반응하고 나머지 요인에는 거의 반응하지 않는다는 뜻이다. 계수 하나를 바꿨을 뿐인데 인코더가 내놓는 q는 어떻게 달라질까? 그리고 그때 목적 함수의 최댓값은 여전히 로그우도의 하한일까?

KL 항에 붙은 계수

ELBO의 KL 항은 사전분포 몫의 에너지 평균 ⟨−ln p(h)⟩_q에서 엔트로피 H(q)를 뺀 것이다. 그러니 KL 항 앞에 붙은 β는 엔트로피 항 앞에도 붙는다. 변분 원리에서 엔트로피 항 앞에 붙은 계수가 온도였고, 온도가 높을수록 분포가 고르게 퍼지고 낮을수록 에너지가 가장 낮은 상태 하나로 몰렸다. β도 같은 자리에 서 있으니 비슷한 역할을 하리라 짐작할 수 있다. 잠재변수가 네 개뿐인 아래 위젯에서 β를 움직이며 가장 좋은 q가 어디로 옮겨 가는지 보고, 정확한 계산은 문제 16에서 해 보자.

문제 16. KL 항에 β를 곱하면

문제 15의 잠재변수 모델(사전분포 p(h) = (0.5, 0.3, 0.2), 우도 p(x | h) = (0.1, 0.4, 0.8), ln p(x) = −1.1087)에서 β-VAE처럼 KL 항에 β를 곱한 목적 함수 ⟨ln p(x | h)⟩_q − βD_KL(q‖p(h))를 가장 크게 하는 q와 그 최댓값을 β = 0.5, 2에서 구하라. 최댓값을 ln p(x)와 견주어라.

목적 함수가 ⟨ln p(x | h)⟩_q − βD_KL(q‖p(h))예요. 이걸 β로 나누면 ⟨ln p(x | h)/β⟩_q − D_KL(q‖p(h))인데, 르장드르 변환 배울 때 본 ln⟨e^z⟩_q = max_p(⟨z⟩_p − D_KL(p‖q))에서 기준 분포를 사전분포로, z를 ln p(x | h)/β로 둔 모양이에요. 그러면 답은 사전분포에 e^z = p(x | h)^(1/β)를 곱해 정규화한 분포이고, 최댓값은 β를 다시 곱해서 β ln Σ_h p(h)p(x | h)^(1/β)예요.

β = 1이면 사후분포와 ln p(x)로 돌아오네요. β = 0.5면 우도를 제곱해서 (0.028, 0.265, 0.707)이고, β = 2면 제곱근을 곱해서 (0.300, 0.360, 0.340)… 이건 거의 사전분포예요.

문제 10이랑 똑같은 구조다. β는 KL 항, 곧 사전분포를 기준으로 잰 엔트로피 쪽에 붙은 계수니까 온도야. β를 키우면 q가 사전분포 쪽으로 녹고, 0으로 보내면 h = 3 하나로 얼어붙어. β-VAE는 온도를 올려서 인코더를 사전분포 쪽에 붙잡아 두는 거네.

최댓값은 β = 2에서 −1.2821, β = 0.5에서 −0.8546이에요. β = 0.5 쪽은 ln p(x) = −1.1087보다 높으니까 β = 1일 때보다 더 좋은 하한이네요. β를 작게 잡는 게 낫겠어요.

하한이 참값보다 위에 있다고? 그럼 하한이 아니잖아.

그래요. β가 1이 아니면 그 최댓값은 더 이상 ln p(x)의 하한이 아니에요. KL 항에 β를 곱하면 엔트로피 쪽 계수, 곧 온도를 바꾼 것이니 다른 온도의 계에서 자유에너지를 재는 셈이에요.

β를 작게 잡으면 목적 함수 숫자가 좋아 보이는데 로그우도가 좋아진 건 아니네요. β가 다른 모델끼리 그 숫자를 비교하면 안 되겠어요.

결국 VAE 손실도 한 줄이었어. 평균 에너지 빼기 온도 곱하기 엔트로피.