VAE 밑바닥 구현 — 인코더·리파라미터화·KL 발산으로 이해하는 이미지 생성

핵심 메시지

• VAE는 방대한 데이터를 학습하기 쉬운 표현으로 압축해, 현대 이미지·영상 생성 모델의 속도와 품질을 뒷받침하는 도구다.
• 인코더는 데이터 한 점을 평균·표준편차로 표현되는 정규분포로 매핑하고, 디코더는 그 표본을 다시 원래 분포로 되돌린다.
• 분포에서 무작위 표본을 뽑는 연산은 미분이 불가능해 역전파가 끊긴다 — 리파라미터화 기법으로 무작위성을 분리해 이를 해결한다.
• KL 발산 항을 손실에 더해 인코더 출력이 사전분포(단위 정규분포)에 가깝도록 제약해야 생성이 의미를 갖는다.
• 재구성 손실과 KL 항은 서로 상충하며, 이 균형(ELBO, 증거 하한)이 VAE의 본질이다. beta로 두 항의 가중치를 조절한다.

쉽게 이해하기

변분 오토인코더(VAE)는 방대한 데이터를 더 배우기 쉬운 표현으로 압축해, 최근 이미지·영상 생성 모델의 품질과 속도를 끌어올린 핵심 도구다. 허깅페이스의 최신 오픈소스 이미지 생성 모델 저장소를 보면 대부분 VAE 폴더가 들어 있을 만큼 널리 쓰인다. 발표자는 수식을 먼저 던지는 대신, 코드로 부품을 하나씩 맞춰 가며 왜 그렇게 하는지 오류와 함께 짚어 나간다.

과제는 이렇다. 어떤 데이터셋이 주어졌을 때 그와 비슷한 새 표본을 더 만들어 내고 싶지만, 데이터가 어떻게 생성됐는지 방정식이 없어 직접 표본을 뽑기 어렵다. 반면 정규분포 같은 단순한 분포에서 표본을 뽑는 일은 함수 호출 한 번이면 된다. 복잡한 분포를 단순한 정규분포로 매핑할 수 있을까? 발표자는 이삿짐 트럭에 가구를 효율적으로 채워 넣는 비유로 압축의 직관을 설명한다.

간단한 신경망 인코더는 데이터 한 점을 두 개의 평균과 두 개의 표준편차, 즉 정규분포로 매핑한다. 디코더는 그 표본을 받아 원래 데이터 점으로 되돌린다. 그런데 첫 시도는 실패한다. 분포에서 무작위로 표본을 뽑는 연산이 미분 불가능해, 역전파 시 인코더로 그래디언트가 흐르지 않기 때문이다. 실제로 디코더의 그래디언트는 정상이지만 인코더의 그래디언트는 None으로 찍힌다.

해법은 리파라미터화 기법이다. 무작위성을 별도의 노드(단위 정규분포에서 뽑은 엡실론)로 분리한 뒤, 표준편차를 곱하고 평균을 더해 원하는 정규분포의 표본으로 변환한다. 그러면 엡실론이 무엇이든 평균·표준편차가 표본에 미치는 영향을 계산할 수 있어 학습이 가능해진다. 파이토치의 rsample 메서드는 이 과정을 자동으로 처리한다.

다만 재구성만 최소화하면 모델은 항등함수를 배워 버려 생성이 실패한다. 그래서 KL 발산 항을 손실에 더해 인코더 출력을 단위 정규분포에 가깝게 제약한다. 하이퍼파라미터 beta로 두 항의 균형을 조절하는데, beta를 1에서 0.5로 낮추자 생성 데이터가 원본 분포를 닮아 갔다. 이 방식은 합성곱 층을 더해 이미지로도 확장되며, MNIST의 784차원 이미지를 64차원으로 90% 넘게 압축해 그럴듯한 이미지를 생성했다.

주요 인사이트

• 학습이 ‘너무 잘 되어’ 항등함수를 배워 버리면 생성이 실패한다 — 재구성 손실만 최소화하려는 함정에 빠지는 것이다.
• 파이토치의 rsample은 리파라미터화를 자동 처리해, 정규분포 파라미터에 대한 미분 가능성을 유지한다.
• beta를 1에서 0.5로 낮추자 재구성과 생성 데이터가 함께 깨끗해지며 원본 분포에 가까워졌다.
• 잠재공간이 전체적으로는 단위 정규분포를 따르면서도 클래스(색)별로 조직화되는 점이 VAE의 핵심이다.
• 이 손실식은 사실 ELBO(증거 하한)이며, 각 항의 역할을 코드로 먼저 체감한 뒤 수식으로 되짚을 수 있다.

자주 묻는 질문

관련 AI 소식