[ML] VAE (Variational Autoencoder)
Autoencoder (AE) — 차원 압축의 신경망
Autoencoder는 입력을 다시 자기 자신으로 복원하도록 학습되는 신경망. 중간에 일부러 좁은 병목(bottleneck, latent space) 을 둬서 데이터의 본질적인 정보만 압축하도록 강제함.
\[x \;\xrightarrow{\text{Encoder } f_\phi}\; z \;\xrightarrow{\text{Decoder } g_\theta}\; \hat{x}\]- $x \in \mathbb{R}^D$: 입력 (예: 784차원 MNIST 이미지)
- $z \in \mathbb{R}^d$ ($d \ll D$): 잠재 표현(latent representation)
- $\hat{x}$: 복원된 출력
손실 함수
복원이 잘 되도록 단순한 재구성 손실(reconstruction loss)을 최소화.
\[\mathcal{L}_{\text{AE}}(\phi, \theta) = \| x - g_\theta(f_\phi(x)) \|^2\](이미지가 픽셀 값이면 MSE, 이진이면 BCE를 씀.)
AE의 결정적 한계
AE의 latent space $z$는 결정론적인 점(point)
학습된 latent space의 모양이 어떻게 생겼는지 알 수 없음
- 학습 데이터의 latent 위치 주변에만 의미 있는 표현이 존재
- 두 점 사이를 보간(interpolation)하면 의미 없는 결과
- 새로운 샘플 생성 불가: latent space에서 아무 점이나 뽑아 decoder에 넣으면 깨진 이미지가 나옴
- 즉, AE는 압축기이지 생성 모델이 아님
VAE (Variational Autoencoder) — 확률적 생성 모델
발상의 전환
VAE는 latent space를 점이 아니라 확률 분포로 다룸. Encoder가 입력 $x$를 받으면 하나의 $z$를 뱉는 게 아니라, $z$의 분포 $q_\phi(z|x)$ 를 출력함.
보통 이 분포는 다변량 정규분포로 가정:
\[q_\phi(z|x) = \mathcal{N}\big(z;\; \mu_\phi(x),\; \mathrm{diag}(\sigma_\phi^2(x))\big)\]즉, encoder는 두 벡터를 출력함:
- $\mu_\phi(x) \in \mathbb{R}^d$: 평균 벡터
- $\sigma_\phi(x) \in \mathbb{R}^d$: 표준편차 벡터 (실제로는 안정성을 위해 $\log \sigma^2$를 출력)
그리고 이 분포에서 샘플링해서 $z$를 얻고, 그걸 decoder에 넣어 복원.
\[x \;\to\; \text{Encoder} \;\to\; (\mu, \sigma) \;\to\; z \sim \mathcal{N}(\mu, \sigma^2) \;\to\; \text{Decoder} \;\to\; \hat{x}\]AE vs VAE 비교
| 구분 | Autoencoder | VAE |
|---|---|---|
| Encoder 출력 | 점 $z$ | 분포 파라미터 $(\mu, \sigma)$ |
| Latent space | 결정론적 | 확률적, 구조화됨 |
| 손실 함수 | reconstruction loss만 | reconstruction loss + KL 정규화 |
| 생성 능력 | ❌ | ✅ ($z \sim \mathcal{N}(0,I)$에서 샘플링) |
| 보간 | 의미 없음 | 부드럽고 의미 있음 |
$\log p(x)$ 최대화
생성 모델의 근본 목표
VAE는 단순히 “재구성을 잘 하는 것”이 목적이 아니라, 데이터의 진짜 분포 $p_{\text{data}}(x)$를 모델 $p_\theta(x)$로 흉내내는 것이 목적.
예를 들어 MNIST에서:
- 진짜 데이터 분포 $p_{\text{data}}(x)$: 실제 손글씨 숫자 이미지들의 분포
- 만들고 싶은 것: 이 분포를 모방하는 $p_\theta(x)$, 그래서 여기서 샘플을 뽑으면 손글씨 숫자처럼 생긴 새 이미지가 나오게
학습 데이터 ${x_1, x_2, \dots, x_N}$가 주어졌을 때, 이 데이터들이 모델 하에서 나올 확률을 최대화하면 모델이 데이터 분포를 잘 흉내내게 됨. 이게 바로 maximum likelihood estimation (MLE):
\[\theta^* = \arg\max_\theta \sum_{i=1}^N \log p_\theta(x_i)\]- “재구성을 잘 하고 싶다”는 결과적 효과이고, 본질적 목적은 “데이터 분포를 학습한다”
그런데 $\log p(x)$를 직접 계산할 수가 없음
VAE의 생성 과정을 보면
\[z \sim p(z) = \mathcal{N}(0, I), \qquad x \sim p_\theta(x|z)\]먼저 잠재변수 $z$를 뽑고, 거기서 $x$를 뽑음.
그러면 데이터 $x$가 나올 확률은 모든 가능한 $z$에 대해 marginalize 해야 함:
\[p_\theta(x) = \int p_\theta(x|z) \, p(z) \, dz\]이 적분은 모든 가능한 $z$에 대한 적분이라 차원이 높아지면 계산이 불가능함 (intractable). 그래서 직접 $\log p(x)$를 최대화하지 못하고, 하한선인 ELBO를 대신 최대화하는 우회 전략을 씀.
VAE의 손실 함수 — ELBO
ELBO 정의 (재구성 항 + KL 항 전체)
ELBO = Evidence Lower BOund = “Evidence(증거)의 하한선”
- Evidence: $\log p(x)$ (데이터의 로그 가능도). 베이지안 통계에서의 evidence - “이 모델 하에 데이터가 관측될 증거”
- Lower Bound: 이 evidence의 하한선
뒤의 KL 항은 항상 0 이상이므로:
\[\log p(x) \geq \text{ELBO}\]ELBO를 끌어올리면 $\log p(x)$도 따라 올라감. 그래서 ELBO를 대신 최대화.
두 항의 역할
| 항 | 이름 | 역할 |
|---|---|---|
| $\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]$ | 재구성 항 (reconstruction term) | "$z$에서 $x$를 잘 복원해라" |
| $-\mathrm{KL}(q_\phi(z|x) \,\|\, p(z))$ | 정규화 항 (KL 항) | "encoder 분포가 prior에서 너무 멀어지지 마라" |
따라서 최소화할 손실은 (ELBO를 최대화 → 부호 뒤집어 최소화):
\[\mathcal{L}_{\text{VAE}}(\phi, \theta; x) = -\mathbb{E}_{q_\phi(z|x)}\!\left[\log p_\theta(x|z)\right] + \mathrm{KL}\!\left(q_\phi(z|x)\;\|\;p(z)\right)\]여기서 prior는 보통 표준정규분포로 둠: $p(z) = \mathcal{N}(0, I)$.
KL 항은 닫힌 형식(closed form) 으로 계산됨 ($q$도 $p$도 정규분포니까):
\[\mathrm{KL}\!\left(\mathcal{N}(\mu, \sigma^2)\;\|\;\mathcal{N}(0, I)\right) = -\frac{1}{2}\sum_{i=1}^{d}\left(1 + \log \sigma_i^2 - \mu_i^2 - \sigma_i^2\right)\]- KL 항은 “encoder가 출력하는 분포가 너무 prior에서 멀리 떨어지지 말라”는 정규화 역할. 이게 있어야 latent space가 매끄러워지고 생성 모델로 쓸 수 있음.
문제 : 샘플링은 미분이 안 됨
문제 정의
재구성 항을 보면 expectation이 있음. 보통 Monte Carlo로 근사:
\[\mathbb{E}_{q_\phi(z|x)}\!\left[\log p_\theta(x|z)\right] \approx \frac{1}{L}\sum_{l=1}^{L} \log p_\theta(x|z^{(l)}), \quad z^{(l)} \sim q_\phi(z|x)\]학습하려면 이 값을 encoder 파라미터 $\phi$ (즉 $\mu, \sigma$ 출력하는 부분) 에 대해 미분해야 함.
\[\nabla_\phi \, \mathbb{E}_{q_\phi(z|x)}\!\left[\log p_\theta(x|z)\right] = \;?\]샘플링 연산 z ~ N(μ, σ²)은 확률적(stochastic)이라 미분 불가능. Backprop이 decoder까지는 잘 내려오다가, 샘플링 노드에서 막혀버림. 그러면 encoder의 $\mu, \sigma$를 어떻게 업데이트해야 할지 알 수 없음.
1
2
3
4
x ──► Encoder ──► (μ_φ, σ_φ) ──► [SAMPLE z ~ N(μ, σ²)] ──► Decoder ──► x̂ ──► Loss
▲ │
│ ▼
└──── ??? ◄── gradient ❌
왜 미분이 안 되는가
expectation을 적분으로 쓰면:
\[\mathbb{E}_{q_\phi(z|x)}\!\left[f(z)\right] = \int q_\phi(z|x) \, f(z) \, dz\]$\phi$로 미분하려면 적분의 분포 자체가 $\phi$에 의존하기 때문에 단순히 안쪽 그래디언트를 평균 내는 것으로는 안 됨.
\[\nabla_\phi \int q_\phi(z|x) f(z) \, dz \;\neq\; \int q_\phi(z|x) \nabla_\phi f(z) \, dz\]샘플 $z^{(l)}$ 자체가 $\phi$에 의존하는데, 샘플링이라는 연산은 그래디언트를 흘려보낼 통로가 없음.
해결책 : Reparameterization Trick
“확률성을 외부의 보조 변수로 빼내자.”
- 결정론과 확률을 분리
$z$를 직접 $\mathcal{N}(\mu, \sigma^2)$에서 뽑는 대신, 다음과 같이 재표현:
\[\boxed{\;z = \mu_\phi(x) + \sigma_\phi(x) \odot \varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, I)\;}\]여기서 $\odot$는 element-wise 곱셈.
이 식이 왜 동등한지 보면: 정규분포의 성질에 의해 $\varepsilon \sim \mathcal{N}(0, I)$일 때 $\mu + \sigma \odot \varepsilon$는 정확히 $\mathcal{N}(\mu, \sigma^2)$를 따름. 분포는 완전히 같음.
확률성이 $\varepsilon$ 안으로 격리됐음
- $\varepsilon$은 외부에서 뽑힌, $\phi$와 완전히 독립적인 노이즈
- $\mu, \sigma$는 encoder가 내놓는 결정론적(deterministic) 출력
- $z$는 $\mu, \sigma, \varepsilon$의 결정론적 함수
1
2
3
4
5
6
7
ε ~ N(0, I) (외부 노이즈, gradient 불필요)
│
▼
x ──► Encoder ──► (μ_φ, σ_φ) ──► z = μ + σ·ε ──► Decoder ──► x̂ ──► Loss
▲ │
│ ▼
└──────── ✅ gradient ◄────────┘
이제 $\mu$와 $\sigma$를 거치는 모든 경로가 미분 가능. Backprop이 막힘없이 흘러감.
수식
Reparameterization 이후 expectation을 다시 쓰면
\[\mathbb{E}_{q_\phi(z|x)}\!\left[f(z)\right] = \mathbb{E}_{\varepsilon \sim \mathcal{N}(0,I)}\!\left[f\big(\mu_\phi(x) + \sigma_\phi(x) \odot \varepsilon\big)\right]\]이제 기대값을 잡는 분포 $p(\varepsilon)$가 $\phi$와 무관. 따라서 그래디언트와 expectation의 순서를 바꿀 수 있음
\[\nabla_\phi \, \mathbb{E}_{\varepsilon}\!\left[f\big(\mu_\phi(x) + \sigma_\phi(x) \odot \varepsilon\big)\right] = \mathbb{E}_{\varepsilon}\!\left[\nabla_\phi \, f\big(\mu_\phi(x) + \sigma_\phi(x) \odot \varepsilon\big)\right]\]Monte Carlo로 근사하면:
\[\nabla_\phi \, \mathbb{E}_{q_\phi(z|x)}\!\left[f(z)\right] \approx \frac{1}{L}\sum_{l=1}^{L} \nabla_\phi \, f\big(\mu_\phi(x) + \sigma_\phi(x) \odot \varepsilon^{(l)}\big)\]실제 학습에선 $L=1$로도 충분함 (미니배치가 노이즈를 평균화해줌).
gradient 계산
연쇄 법칙(chain rule)을 적용하면
\[\frac{\partial \mathcal{L}}{\partial \mu_\phi} = \frac{\partial \mathcal{L}}{\partial z} \cdot \frac{\partial z}{\partial \mu_\phi} = \frac{\partial \mathcal{L}}{\partial z} \cdot 1\] \[\frac{\partial \mathcal{L}}{\partial \sigma_\phi} = \frac{\partial \mathcal{L}}{\partial z} \cdot \frac{\partial z}{\partial \sigma_\phi} = \frac{\partial \mathcal{L}}{\partial z} \cdot \varepsilon\]깔끔하게 떨어짐. $\varepsilon$은 forward pass에서 샘플링된 상수처럼 취급되고, 그래디언트는 자연스럽게 $\mu$와 $\sigma$로 흘러감.
Backpropagation은 정확히 어느 단계에서 일어나는가?
한 step의 흐름
미니배치 하나 $x$에 대해 다음 순서로 진행:
1
2
3
4
5
6
7
8
9
10
11
12
[1] Forward pass ──────────────────────────────► loss 계산
x → Encoder → (μ, log σ²) → ε ~ N(0,I) → z = μ + σ⊙ε → Decoder → x̂
│
[2] Loss 계산 │
L = ‖x - x̂‖² + KL(N(μ,σ²) ‖ N(0,I)) ◄──────┘
│
[3] Backward pass (여기서 backprop) ◄────────────────────────────────────┘
∂L/∂θ (decoder), ∂L/∂φ (encoder) 계산
[4] Optimizer step
θ ← θ - η·∂L/∂θ
φ ← φ - η·∂L/∂φ
Backprop은 [3]단계에서 일어남. PyTorch라면 loss.backward()를 호출하는 순간.
그래디언트가 흐르는 경로
loss.backward()가 호출되면, 그래디언트는 loss에서 시작해서 계산 그래프를 거꾸로 타고 흘러감:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
Loss
│
├─► 재구성 항 ‖x - x̂‖²
│ │
│ ▼
│ Decoder (θ 업데이트) ──► ∂L/∂θ
│ │
│ ▼
│ z = μ + σ⊙ε
│ ├──► μ 경로: ∂L/∂μ = ∂L/∂z · 1
│ └──► σ 경로: ∂L/∂σ = ∂L/∂z · ε
│ │
│ ▼
│ Encoder (φ 업데이트) ──► ∂L/∂φ
│
└─► KL 항 (μ, σ에 대한 닫힌 형식)
│
▼
Encoder의 μ, σ로 직접 그래디언트 흐름 ──► ∂L/∂φ
- Decoder의 파라미터 $\theta$ 는 재구성 항을 통해서만 업데이트
- Encoder의 파라미터 $\phi$ 는 두 경로로 업데이트
- 재구성 항 → decoder → $z$ → $(\mu, \sigma)$ → encoder
- KL 항 → $(\mu, \sigma)$ → encoder (직접)
- $\varepsilon$ 은 외부 노이즈라 그래디언트가 흐르지 않음 (상수 취급)
reparameterization trick이 왜 결정적인가
만약 reparameterization을 안 했다면, [3]단계에서 그래디언트가 $z$까지는 흘러오는데 $z \sim \mathcal{N}(\mu, \sigma^2)$ 샘플링 노드에서 막혀서 $\mu, \sigma$로 못 넘어감.
- Decoder $\theta$는 그나마 업데이트 가능 ($z$에서 $x$까지는 미분 가능)
- Encoder $\phi$는 재구성 항으로부터 그래디언트를 못 받음
KL 항으로 $\mu, \sigma$가 업데이트되긴 하지만, KL 항만으로는 “데이터를 어떻게 인코딩할지”를 배울 수 없음. KL은 그냥 “prior에 가깝게 가라”고만 시킬 뿐이니까. 결과적으로 encoder는 모든 입력에 대해 $\mathcal{N}(0, I)$만 뱉는 멍청한 모델이 되어버림.
Reparameterization trick은 바로 이 막힌 경로를 뚫어서, 재구성 신호가 decoder를 거쳐 encoder까지 전달되게 만드는 장치.
코드 (PyTorch)
1
2
3
4
def reparameterize(mu, logvar):
std = torch.exp(0.5 * logvar) # σ = exp(0.5 · log σ²)
eps = torch.randn_like(std) # ε ~ N(0, I)
return mu + std * eps # z = μ + σ ⊙ ε
한 줄임. torch.randn_like(std)는 $\phi$와 무관하게 외부에서 뽑힌 노이즈라서 그래디언트가 흐르지 않고, mu와 std로는 자연스럽게 backprop이 됨.
전체 학습 루프:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# [1] Forward
mu, logvar = encoder(x)
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std) # ε: 그래디언트 안 흐름
z = mu + std * eps # reparameterization
x_hat = decoder(z)
# [2] Loss 계산
recon_loss = F.mse_loss(x_hat, x, reduction='sum')
kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
loss = recon_loss + kl_loss
# [3] Backprop ← 여기!
optimizer.zero_grad()
loss.backward() # 모든 파라미터에 대한 그래디언트 자동 계산
# [4] Update
optimizer.step()
cf) 보통 $\sigma$ 대신 $\log \sigma^2$ (logvar)를 출력함. 이유는 (1) $\sigma > 0$ 제약을 자연스럽게 만족시키고, (2) 수치적으로 더 안정적이기 때문.
전체 학습 한 사이클 정리
입력 $x$ 하나에 대해:
- Forward (Encoder): $(\mu, \log \sigma^2) = \text{Encoder}_\phi(x)$
- Sample noise: $\varepsilon \sim \mathcal{N}(0, I)$
- Reparameterize: $z = \mu + \sigma \odot \varepsilon$ where $\sigma = \exp(0.5 \log \sigma^2)$
- Forward (Decoder): $\hat{x} = \text{Decoder}_\theta(z)$
- Compute loss: \(\mathcal{L} = \underbrace{\| x - \hat{x} \|^2}_{\text{재구성}} - \underbrace{\frac{1}{2}\sum_i (1 + \log \sigma_i^2 - \mu_i^2 - \sigma_i^2)}_{\text{KL (부호 주의)}}\)
- Backward: $\nabla_\phi \mathcal{L}, \nabla_\theta \mathcal{L}$ 계산 (reparameterization 덕분에 막힘 없이 흐름)
- Update: $\phi \leftarrow \phi - \eta \nabla_\phi \mathcal{L}$, $\theta \leftarrow \theta - \eta \nabla_\theta \mathcal{L}$
한 줄 요약
Autoencoder는 결정론적 압축기, VAE는 latent space를 확률 분포로 모델링한 생성 모델.
VAE의 본질적 목적은 데이터 분포 $p_{\text{data}}(x)$를 $p_\theta(x)$로 흉내내는 것, 즉 $\log p(x)$ 최대화 (MLE).
그런데 $p_\theta(x) = \int p_\theta(x|z)p(z)dz$가 intractable하므로 그 하한선인 ELBO (= 재구성 항 + KL 항 전체) 를 대신 최대화함.
학습 과정에서 $z \sim \mathcal{N}(\mu, \sigma^2)$ 샘플링이 미분 불가능하다는 문제가 생기는데, $z = \mu + \sigma \odot \varepsilon$ ($\varepsilon \sim \mathcal{N}(0, I)$)로 재표현해 확률성($\varepsilon$) 과 학습 가능한 부분($\mu, \sigma$) 을 분리하면 backprop이 decoder → $z$ → encoder까지 막힘없이 흘러감. 이게 reparameterization trick.