Skip to main content

VAE (Variational Autoencoder)

학습 목표

  • VAE가 오토인코더와 다른 점(확률적 잠재 공간)을 이해한다
  • ELBO(Evidence Lower Bound)의 두 항을 설명할 수 있다
  • 재매개변수화 트릭(Reparameterization Trick)의 필요성과 원리를 안다
  • KL Divergence의 역할을 이해한다

왜 중요한가

변분 오토인코더(Variational Autoencoder, VAE)는 잠재 공간에 확률 분포 구조를 부여합니다. 기본 오토인코더와 달리 잠재 공간에서 의미 있는 샘플링이 가능하여, 새로운 데이터를 생성할 수 있습니다. 생성 모델의 이론적 기초이며, Diffusion 모델의 Latent Space 설계에도 영향을 줍니다.

구조

수학적 배경

ELBO (Evidence Lower Bound)

VAE는 데이터의 로그 가능도 logp(x)\log p(x)의 하한(Lower Bound)을 최대화합니다. logp(x)Eqϕ(zx)[logpθ(xz)]DKL(qϕ(zx)p(z))\log p(x) \geq \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - D_{KL}(q_\phi(z|x) \| p(z))
의미역할
E[logpθ(xz)]\mathbb{E}[\log p_\theta(x\|z)]재구성 항입력을 잘 복원하도록 유도
DKL(qϕ(zx)p(z))D_{KL}(q_\phi(z\|x) \| p(z))KL 항잠재 분포를 사전 분포 N(0,I)N(0, I)에 가깝게 유도

KL Divergence (가우시안 폐쇄형)

인코더 출력이 가우시안일 때 KL Divergence의 해석적 해: DKL=12j=1d(1+logσj2μj2σj2)D_{KL} = -\frac{1}{2} \sum_{j=1}^{d} (1 + \log \sigma_j^2 - \mu_j^2 - \sigma_j^2)

재매개변수화 트릭

zqϕ(zx)=N(μ,σ2)z \sim q_\phi(z|x) = N(\mu, \sigma^2)에서 직접 샘플링하면 기울기 전파가 불가능합니다. 재매개변수화 트릭으로 확률적 노드를 우회합니다. z=μ+σϵ,ϵN(0,I)z = \mu + \sigma \odot \epsilon, \quad \epsilon \sim N(0, I) 이렇게 하면 zzμ\muσ\sigma에 대해 미분 가능해져 역전파가 가능합니다.

구현

import torch
import torch.nn as nn
import torch.nn.functional as F

class VAE(nn.Module):
    """변분 오토인코더"""
    def __init__(self, input_dim=784, hidden_dim=256, latent_dim=16):
        super().__init__()

        # 인코더: x → (μ, log σ²)
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
        )
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)       # 평균
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)    # 로그 분산

        # 디코더: z → x̂
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim),
            nn.Sigmoid(),
        )

    def encode(self, x):
        """인코더: 입력 → 평균, 로그 분산"""
        h = self.encoder(x)
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        return mu, logvar

    def reparameterize(self, mu, logvar):
        """재매개변수화 트릭: z = μ + σ · ε"""
        std = torch.exp(0.5 * logvar)       # σ = exp(0.5 * log σ²)
        eps = torch.randn_like(std)          # ε ~ N(0, I)
        return mu + std * eps

    def decode(self, z):
        """디코더: 잠재 벡터 → 재구성"""
        return self.decoder(z)

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        x_hat = self.decode(z)
        return x_hat, mu, logvar

손실 함수

def vae_loss(x, x_hat, mu, logvar, beta=1.0):
    """VAE 손실 함수 (ELBO의 부정)

    Args:
        x: 원본 입력
        x_hat: 재구성 출력
        mu: 인코더의 평균
        logvar: 인코더의 로그 분산
        beta: KL 항의 가중치 (β-VAE)
    """
    # 재구성 손실: BCE (픽셀 단위)
    recon_loss = F.binary_cross_entropy(x_hat, x, reduction='sum')

    # KL Divergence: 폐쇄형 해
    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

    return recon_loss + beta * kl_loss

학습

import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# 데이터
transform = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)

# 모델
model = VAE(input_dim=784, latent_dim=16).to('cuda')
optimizer = optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(30):
    total_loss = 0
    for images, _ in train_loader:
        images = images.view(-1, 784).to('cuda')  # (배치, 784)

        x_hat, mu, logvar = model(images)
        loss = vae_loss(images, x_hat, mu, logvar)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    avg_loss = total_loss / len(train_loader.dataset)
    print(f"Epoch {epoch+1}: 평균 ELBO 손실 = {avg_loss:.2f}")

잠재 공간에서 생성

# 잠재 공간에서 랜덤 샘플링 → 새 이미지 생성
with torch.no_grad():
    z = torch.randn(16, 16).to('cuda')  # 16개의 잠재 벡터 샘플링
    generated = model.decode(z)
    generated = generated.view(-1, 1, 28, 28)
    # 생성된 이미지 시각화 (torchvision.utils.make_grid 사용)

# 잠재 공간 보간 (두 이미지 사이의 부드러운 전환)
with torch.no_grad():
    z1 = model.encode(image1.view(-1, 784).to('cuda'))[0]  # 첫 이미지의 μ
    z2 = model.encode(image2.view(-1, 784).to('cuda'))[0]  # 두 번째 이미지의 μ

    alphas = torch.linspace(0, 1, 10)
    interpolated = []
    for alpha in alphas:
        z = (1 - alpha) * z1 + alpha * z2
        img = model.decode(z).view(1, 28, 28)
        interpolated.append(img)

β-VAE

β>1\beta > 1로 설정하면 KL 항의 비중이 커져 분리된(Disentangled) 표현을 학습합니다. 각 잠재 차원이 독립적인 의미를 가지게 됩니다.
β 값재구성 품질잠재 공간 구조특징
0.5높음약한 구조재구성 중시
1.0보통표준원래 VAE
4.0낮음분리된 표현β-VAE
VAE는 재구성 손실로 MSE 또는 BCE를 사용하기 때문에, 생성된 이미지가 흐릿한(blurry) 경향이 있습니다. 이는 평균적인 픽셀 값을 학습하기 때문입니다. 선명한 이미지를 생성하려면 GAN의 적대적 학습이 더 효과적이며, 이를 결합한 VAE-GAN 같은 하이브리드 모델도 연구되었습니다.

참고 논문

논문학회/연도핵심 기여
Auto-Encoding Variational Bayes (Kingma & Welling)ICLR 2014VAE 제안, 재매개변수화 트릭
β-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework (Higgins et al.)ICLR 2017분리된 표현 학습

체크리스트

  • VAE와 AE의 차이(확률적 잠재 공간)를 설명할 수 있다
  • ELBO의 재구성 항과 KL 항의 역할을 안다
  • 재매개변수화 트릭의 필요성을 설명할 수 있다
  • 잠재 공간에서 샘플링과 보간을 이해한다

다음 문서