β-VAE와 계층형 VAE로 분리 표현 학습 설계하기

잠재공간을 제어 가능한 표현으로 만드는 방식

변분 오토인코더(VAE)는 확률적 잠재공간을 통해 데이터 생성과 표현 학습을 함께 수행한다. 여기서 중요한 문제는 잠재 차원이 실제 데이터의 요인과 얼마나 분리되어 있는가다. 한 차원을 바꿨을 때 여러 속성이 얽혀 움직이면 편집, 해석, 전이에 쓰기 어렵다.

β-VAE, 분리 표현, 계층형 VAE는 이 문제를 서로 다른 방향에서 다룬다. β-VAE는 정보 병목의 강도를 조절하고, Total Correlation 기반 모델은 잠재변수 사이의 상관을 직접 억제한다. 계층형 구조는 전역 특성과 지역 디테일을 별도 잠재 계층에 배치한다.

VAE의 목적함수 ELBO는 재구성항과 KL 발산항으로 구성된다. β-VAE는 KL 항에 β(≥1)를 적용해 정보 병목을 강화한다. β가 커지면 잠재변수의 독립성과 분리 표현은 강화될 수 있지만, 재구성 품질과는 트레이드오프가 생긴다.

분리 표현은 서로 독립적인 잠재 요인을 개별 잠재 차원으로 학습하는 방식이다. 잠재 슬라이더로 속성을 조작할 수 있고, 요인과 차원의 대응을 해석하기 쉬우며, 전이 학습 효율성도 높일 수 있다.

계층형 VAE는 z2→z1→x 구조처럼 여러 잠재변수를 둔다. 상위 계층은 전역 스타일이나 배경을, 하위 계층은 지역 디테일이나 콘텐츠를 맡도록 설계할 수 있다. 생성은 하향식(top-down)으로, 추론은 상향식(bottom-up)으로 결합해 복잡한 데이터의 다중 스케일 구조를 반영한다.

정보 병목과 잠재변수 독립성을 다루는 선택지

β-VAE에서는 β>1로 KL 가중치를 높여 정보 흐름을 제한한다. 다만 학습 초기에 규제가 지나치면 후방 붕괴(posterior collapse)가 생길 수 있다. 용량 C를 쓰는 방식은 이를 완화하는 접근이다.

ELBO = Recon + β·|KL − C|

학습 초기에는 낮은 C에서 시작해 점진적으로 올리면서 분리 표현을 유도한다.

잠재변수 간 독립성을 더 직접적으로 다루려면 Total Correlation 제어가 필요하다. β-TCVAE와 FactorVAE는 KL을 분해하거나 TC 추정기를 사용해 이 상관을 제약한다. 이 선택은 독립성 제어를 강화하는 대신 안정성과 계산비용 사이의 균형을 요구한다.

계층형 구조에서는 추론 네트워크 q(z1,z2|x)가 상향식으로 특징을 추출한다. 생성 분포 p(x|z1)p(z1|z2)p(z2)는 하향식 조건부 생성을 담당한다. 상층 z2에 전역 스타일·배경을, 하층 z1에 지역 디테일·콘텐츠를 배치하는 전략이 여기에 해당한다.

학습 안정화에는 KL 어닐링, free-bits(차원별 최소 KL), 스톱-그라디언트(top-down path)를 사용할 수 있다. 미니배치로 TC를 추정할 때는 배치 크기와 추정 바이어스도 관리 대상이며, 옵티마이저와 학습률 스케줄의 영향도 크다.

분리 표현은 MIG, DCI, SAP, BetaVAE Score로 평가할 수 있지만 데이터셋 의존성이 있다. 재현성을 위해서는 시드를 고정하고 데이터 요인 노출을 제한해야 한다. 운영에서는 PSNR/SSIM 기반 재구성 품질, FID 기반 생성 품질, 유저 스터디 기반 조작 가능성을 함께 본다.

계층형 잠재공간의 학습 흐름

⤢✕

입력 배치 x는 상향 인코딩을 거쳐 계층 잠재변수로 샘플링되고, 하향 디코딩으로 재구성 x̂를 만든다. 이 과정에서 z1·z2와 로그·지표를 함께 기록한다.

후방 붕괴는 KL(z1), KL(z2) 평균이 임계값 아래로 지속될 때 경고 대상으로 삼는다. TC 추정이 불안정해 배치 분산이 커지면 추정기 가중치 λ를 점진적으로 높일 수 있다. 재구성 손실이 급등하는 발산 상황에서는 학습률 감쇠, β의 임시 하향, gradient clipping을 적용한다.

모델별로 달라지는 품질과 운영 부담

항목	Vanilla VAE	β-VAE	β-TCVAE/FactorVAE	Hierarchical VAE
분리표현	낮음	중~높음(β, C 의존)	높음(TC 직접 제약)	중~높음(스케일 분리)
재구성 품질	높음	중(과규제 시 저하)	중	중~높음(깊이에 비례)
확장성	높음	높음	중(추정기 비용)	중(최적화 난이도)
안정성	높음	중(붕괴 위험)	중(추정 바이어스)	중(탐색 어려움)
운영 편의	매우 높음	높음	중	중(튜닝 비용 증가)

분리 표현은 TC 제어나 계층 구조로 유도할 수 있다. 재구성 품질은 규제 강도와 디코더 용량의 함수이며, 운영 편의는 튜닝해야 할 매개변수 수와 반대 방향으로 움직인다.

제어와 분석이 필요한 모델에 적용하기

콘텐츠와 스타일, 배경과 전경, 포즈와 아이덴티티를 분리하면 콘텐츠를 보존하면서 스타일만 바꾸는 파이프라인을 만들 수 있다. 이미지 편집 UI에서는 잠재 슬라이더로 실시간 조작하는 방식도 가능하다.

로보틱스와 시뮬레이션에서는 잡음이 많은 관측으로부터 저차원 요인을 추출해 정책 학습을 안정화할 수 있다. 도메인 랜덤화에서는 요인별 견고성을 검증하고, 상태 추정의 불확실성을 정량화해 안전 제어를 강화하는 데 활용한다.

이상 탐지와 과학 데이터 분석에서는 정상 데이터의 요인을 분리해 비정상 요인 변화를 민감하게 탐지할 수 있다. 실험 변수를 분리하면 인과 가설 수립을 보조할 수 있지만, 정확한 인과성을 보장하는 것은 아니다.

계층형 VAE는 전역 레이아웃(z2)과 지역 텍스처(z1)를 나눠 생성하는 데 적합하다. 대용량 이미지나 시계열을 스케일별로 모델링할 때 이 구조를 적용할 수 있다.

분리 지표인 MIG, DCI, SAP, BetaVAE Score는 향상을 기대할 수 있으며, 동일 파라미터 대비 제어 가능 차원 비율도 증가한다. 복잡한 데이터에서는 전역-지역 분할 효과로 FID/LPIPS가 개선될 수 있으나 데이터셋과 아키텍처에 의존한다. 모델 해석 가능성, 적은 라벨에서의 다운스트림 전이 성능, 불확실성 추정과 잠재 조작을 통한 검증 용이성도 함께 얻을 수 있다.

PyTorch로 구현하는 β-VAE와 계층형 VAE

환경은 Python 3.10+, PyTorch 2.2+, CUDA 선택적으로 구성한다. 분리 평가에는 dSprites를, 데모에는 MNIST 또는 Fashion-MNIST를 사용할 수 있다. 재현성을 위해 시드를 고정하고 평가 지표를 정규화한다.

# pip install torch torchvision numpy
# PyTorch 2.2+, Python 3.10
import torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# ----------------------
# β-VAE with capacity C
# ----------------------
class BetaVAE(nn.Module):
    def __init__(self, z_dim=10):
        super().__init__()
        self.enc = nn.Sequential(
            nn.Conv2d(1,32,4,2,1), nn.ReLU(),
            nn.Conv2d(32,64,4,2,1), nn.ReLU(),
            nn.Flatten(), nn.Linear(64*7*7,256), nn.ReLU()
        )
        self.mu = nn.Linear(256, z_dim)
        self.logvar = nn.Linear(256, z_dim)
        self.dec_fc = nn.Sequential(nn.Linear(z_dim, 256), nn.ReLU(),
                                    nn.Linear(256, 64*7*7), nn.ReLU())
        self.dec = nn.Sequential(
            nn.ConvTranspose2d(64,32,4,2,1), nn.ReLU(),
            nn.ConvTranspose2d(32,1,4,2,1), nn.Sigmoid()
        )
    def encode(self, x):
        h = self.enc(x)
        return self.mu(h), self.logvar(h)
    def reparam(self, mu, logvar):
        std = torch.exp(0.5*logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
    def decode(self, z):
        h = self.dec_fc(z).view(-1,64,7,7)
        return self.dec(h)
    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparam(mu, logvar)
        xhat = self.decode(z)
        return xhat, mu, logvar

def kld_normal(mu, logvar):
    return -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = BetaVAE(z_dim=10).to(device)
opt = torch.optim.Adam(model.parameters(), lr=2e-4)

ds = datasets.MNIST(root='./data', train=True, download=True,
                    transform=transforms.ToTensor())
loader = DataLoader(ds, batch_size=128, shuffle=True, num_workers=2)

beta, C_max, C, C_delta = 4.0, 25.0, 0.0, 1e-3  # capacity schedule
for epoch in range(5):
    for i,(x,_) in enumerate(loader):
        x = x.to(device)
        xhat, mu, logvar = model(x)
        rec = F.binary_cross_entropy(xhat, x, reduction='none').sum(dim=(1,2,3)).mean()
        kl = kld_normal(mu, logvar).mean()
        C = min(C + C_delta, C_max)
        loss = rec + beta * torch.abs(kl - C)  # capacity control
        opt.zero_grad(); loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 5.0)
        opt.step()
        # collapse monitoring
        if i % 200 == 0:
            print(f"e{epoch} i{i} rec={rec.item():.2f} kl={kl.item():.2f} C={C:.2f}")

# ----------------------
# Minimal Hierarchical VAE (z2 -> z1 -> x)
# ----------------------
class HierVAE(nn.Module):
    def __init__(self, z1=16, z2=8):
        super().__init__()
        self.enc1 = nn.Sequential(nn.Conv2d(1,32,4,2,1), nn.ReLU(),
                                  nn.Conv2d(32,64,4,2,1), nn.ReLU(), nn.Flatten())
        self.hdim = 64*7*7
        self.mu2 = nn.Linear(self.hdim, z2); self.lv2 = nn.Linear(self.hdim, z2)
        self.mu1 = nn.Linear(self.hdim+z2, z1); self.lv1 = nn.Linear(self.hdim+z2, z1)

        self.p_z1_mu = nn.Linear(z2, z1); self.p_z1_lv = nn.Linear(z2, z1)
        self.dec_fc = nn.Sequential(nn.Linear(z1, 256), nn.ReLU(), nn.Linear(256, 64*7*7), nn.ReLU())
        self.dec = nn.Sequential(nn.ConvTranspose2d(64,32,4,2,1), nn.ReLU(),
                                 nn.ConvTranspose2d(32,1,4,2,1), nn.Sigmoid())

    def reparam(self, mu, lv): return mu + torch.randn_like(mu) * torch.exp(0.5*lv)

    def forward(self, x):
        h = self.enc1(x)
        mu2, lv2 = self.mu2(h), self.lv2(h)
        z2 = self.reparam(mu2, lv2)
        hcat = torch.cat([h, z2], dim=1)
        mu1, lv1 = self.mu1(hcat), self.lv1(hcat)
        z1 = self.reparam(mu1, lv1)
        xhat = self.dec(self.dec_fc(z1).view(-1,64,7,7))
        # priors
        pz1_mu, pz1_lv = self.p_z1_mu(z2), self.p_z1_lv(z2)
        kl2 = kld_normal(mu2, lv2).mean()
        kl1 = 0.5*torch.sum(
            (torch.exp(lv1) + (mu1 - pz1_mu)**2)/torch.exp(pz1_lv) - 1 + pz1_lv - lv1, dim=1
        ).mean()
        return xhat, kl1, kl2

hier = HierVAE().to(device)
opt2 = torch.optim.Adam(hier.parameters(), lr=2e-4)
for epoch in range(3):
    for i,(x,_) in enumerate(loader):
        x = x.to(device)
        xhat, kl1, kl2 = hier(x)
        rec = F.binary_cross_entropy(xhat, x, reduction='none').sum(dim=(1,2,3)).mean()
        loss = rec + 1.0*kl1 + 1.0*kl2
        opt2.zero_grad(); loss.backward(); opt2.step()
        if i % 300 == 0:
            print(f"[Hier] e{epoch} i{i} rec={rec.item():.2f} kl1={kl1.item():.2f} kl2={kl2.item():.2f}")

데이터 전처리에서는 표준화와 일관된 augment로 잠재 요인 교란을 줄인다. 디코더 용량이 과도하면 포스터리어를 무시하는 Posterior Collapse가 생길 수 있으므로 KL 어닐링과 free-bits를 병행한다. MIG와 DCI는 요인 라벨을 요구하며, 라벨이 없다면 latent traversal 조작 실험과 다운스트림 전이로 간접 평가한다.

튜닝할 때 함께 보는 균형

β와 C 스케줄링은 점진적으로 적용해 초기 안정성을 확보한 뒤 분리를 강화한다. 계층형 VAE에서는 상향·하향 신호의 균형이 필요하며, skip-connection을 최소화하고 LayerNorm 또는 GroupNorm으로 신호 스케일을 안정화할 수 있다. 미니배치 TC 추정에는 큰 배치와 moving-average 보정을 사용한다.

β↑와 TC 억제↑는 분리를 강화하지만 재구성 손실을 높일 수 있다. 계층을 깊게 만들수록 표현력은 커지지만 최적화 난이도와 수렴 시간도 증가한다. TC 추정과 다층 잠재 구조는 계산·메모리 비용을 늘리므로 운영 환경의 지연 한계까지 포함해 선택해야 한다.