일반적인 오토인코더 (Autoencoder)는 각 입력을 잠재 코드 (Latent code) 내의 단일 '점 (point)'으로 압축한 뒤 다시 디코딩합니다. 이는 압축에는 효과적이지만, 학습된 잠재 공간 (Latent space)이 구멍이 뚫린 (holey) 상태가 됩니다. 즉, 학습 데이터 포인트들이 섬처럼 떨어져 있고 그 사이에는 빈 공간이 존재하여, 무작위 코드를 선택하거나 두 예시 '사이'의 점을 선택하면 디코더가 쓰레기 값을 생성하게 됩니다. 이로 인해 샘플링 (Sampling)이 불가능하며, 결과적으로 생성 (Generate)도 할 수 없습니다. VAE는 세 가지 작은 추가 요소를 통해 이 문제를 해결하며, 저는 재매개변수화 트릭 (Reparameterization trick)과 가우시안 KL (Gaussian KL)이 단순한 애니메이션이 아닌 실제 수학적 계산으로 브라우저에서 작동하는 실시간 2D VAE를 구축했습니다. 그 작동 원리는 다음과 같습니다.

인코더는 점이 아닌 분포를 방출합니다

일반적인 인코더는 하나의 선형 레이어 (Linear layer)를 거쳐 코드로 끝납니다. 반면 VAE 인코더는 동일한 은닉 특징 (Hidden features) 위에 두 개의 선형 헤드 (Linear heads)를 가집니다. 하나는 평균 μ를 생성하고, 다른 하나는 로그 분산 logσ²를 생성합니다. 로그를 예측함으로써 σ²를 양수이며 제약이 없는 상태로 유지할 수 있습니다. 이 둘은 해당 입력에 대한 대각 가우시안 사후 분포 (Diagonal Gaussian posterior) q(z|x) = N(μ, σ²)를 정의하며, 이는 하나의 점이 아닌 하나의 작은 이웃 영역 (Neighbourhood)을 형성합니다.

class Encoder(nn.Module):
    def __init__(self, in_dim, hidden, zdim):
        super().__init__()
...

재매개변수화 트릭 (Reparameterization trick)은 샘플링을 미분 가능하게 만듭니다

이제 우리는 z ~ N(μ, σ²)를 샘플링해야 합니다. 하지만 무작위 샘플을 추출하는 과정은 μ와 σ에 대한 그래디언트 (Gradient)가 없기 때문에, 역전파 (Backprop)가 중단됩니다. 이 트릭은 무작위성을 파라미터 '외부'로 이동시킵니다. 즉, ε ~ N(0,1)을 샘플링한 다음, z = μ + σ⊙ε를 형성합니다. 이제 z는 μ와 σ에 대한 매끄럽고 미분 가능한 함수가 되며, ∂z/∂μ = 1 및 ∂z/∂σ = ε이므로 그래디언트가 그대로 흐를 수 있습니다.

def reparameterize(mu, logvar):
    std = torch.exp(0.5 * logvar)      # sigma = exp(0.5 * log sigma^2)
    eps = torch.randn_like(std)        # epsilon ~ N(0,1), the randomness
...

이 한 줄의 코드가 VAE를 단순한 노이즈가 섞인 오토인코더와 구분 짓는 핵심입니다. 해당 코드를 샘플링하고 디코딩하는 과정은 μ 주변의 전체 영역이 정확한 한 점뿐만 아니라 합리적으로 디코딩되도록 강제합니다.

가우시안 KL (Gaussian KL)은 하나의 폐형식(closed-form) 수식입니다

손실 함수(loss)는 두 부분으로 나뉩니다. 첫 번째는 일반적인 재구성 항 (reconstruction term)으로, 디코딩된 출력이 입력으로부터 얼마나 떨어져 있는지를 나타냅니다 (픽셀의 경우 이진 교차 엔트로피 (binary cross-entropy), 연속형 데이터의 경우 평균 제곱 오차 (MSE)). 두 번째는 공간을 하나로 묶어주는 역할을 합니다. 사후 분포 (posterior) N(μ,σ²)와 사전 분포 (prior) N(0,1)가 모두 가우시안 (Gaussian) 분포이기 때문에, 이들의 KL 발산 (KL divergence)은 샘플링이 필요 없는 _정확한 폐형식 (exact closed form)_을 가집니다.

def kl_divergence(mu, logvar):
    # KL( N(mu, sigma^2) || N(0,1) ) = 0.5 * sum( mu^2 + sigma^2 - 1 - log sigma^2 )
    return -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) / mu.size(0)

이 값은 모든 μ=0 이고 σ=1 일 때, 즉 정확히 사전 분포와 일치할 때만 0이 되며, 사후 분포가 멀어질수록 커집니다. 따라서 이를 최소화하면 모든 덩어리(blob)를 단위 분산(unit spread)을 가진 원점으로 끌어당겨, 고립된 섬들을 하나의 연속적인 영역으로 병합하게 됩니다.

재구성(Recon) + β·KL은 음의 ELBO입니다

두 항을 더하면 L = recon + β·KL이 됩니다. 이것이 바로 _음의 ELBO (negative ELBO, evidence lower bound)_입니다. 즉, 이를 최소화하는 것은 log p(x)의 하한(lower bound)을 최대화하는 것과 같습니다. β=1은 오리지널 VAE이며, β>1 (β-VAE)은 선명도(sharpness)를 일부 희생하는 대신 더 연속적이고 얽힘이 해제된 (disentangled) 공간을 얻기 위해 사전 분포에 더 강한 가중치를 부여합니다. β=0은 일반적인 오토인코더 (autoencoder)로 회귀합니다. 데모를 통해 β를 조절하면 인코딩된 덩어리들이 단위 원(unit-circle) 사전 분포를 향해 끌려가며 음영 처리된 섬들이 융합되는 모습을 직접 확인할 수 있습니다.

결과물: 생성 및 보간 (generate and interpolate)

KL이 집합적 사후 분포 (aggregate posterior)를 N(0,1)로 밀어붙였기 때문에, 생성 (generation)은 매우 간단해집니다. 사전 분포에서 z ~ N(0,1)을 샘플링하여 디코딩하기만 하면 됩니다. 인코더도, 입력값도 필요하지 않습니다.

@torch.no_grad()
def generate(dec, n, zdim):
    z = torch.randn(n, zdim)   # z ~ prior N(0,1)
...

그리고 보간 (interpolation)이 작동합니다. 두 인코딩 (encoding) 사이를 직선으로 따라가면, KL 항 (KL term)이 두 덩어리 사이의 공간을 채워주었기 때문에 모든 중간 지점이 매끄럽게 변형되며 합리적인 무언가로 디코딩 (decode)됩니다. 일반적인 오토인코더 (autoencoder)에서는 이러한 중간 코드들이 노이즈로 디코딩되는데, 바로 그 간극이 결정적인 차이입니다. 여기서 구축한 확률적 잠재 공간 (probabilistic-latent) 개념은 조건부 VAE (conditional VAEs), VQ-VAE (현대적 이미지 토크나이저의 기반이 되는 이산 코드북 (discrete codebook)), 그리고 잠재 확산 모델 (latent diffusion)의 조상입니다. Stable Diffusion은 VAE의 잠재 공간 (latent space) — — 내부에서 확산 (diffusion) 과정을 수행합니다.

β를 드래그하고 잠재 공간 (latent space)을 클릭하여 어떤 지점이든 디코딩 (decode)해 보세요:
https://dev48v.infy.uk/dl/day42-vae.html