VAE(Variational AutoEncoder)의 원리 · 설명 with MNIST Pytorch

안준영·2024년 1월 22일

안녕하세요, Justin4AI입니다. 정말 바빴던 2학년 2학기를 마치고, 방학 중에 1) NLP 공부 2) Spring 공부 3) 개인 토이프로젝트 4) Kaggle 참가 등 여러 계획이 있었는데, 저는 결국 generative AI를 계속 공부하고 있습니다. ~~겨울동안 SNU의 ExploreCSR 프로그램에 참가하게 되어 NLP 공부는 강제 달성!~~

오늘 포스트에서 다룰 것은 생성형 AI와 함께 각광받았고, 이제는 해당 분야에서 필수 지식이 되어버린 VAE(Variational AutoEncoder)의 원리입니다. 어떻게 보면, AutoEncoder까지의 모델들은 수학적 지식을 크게 요구하지는 않으나, VAE와 Diffusion models를 포함한 SOTA들을 제대로 이해하고 활용하려면 다음 두 가지를 동시에 만족해야 합니다.

1) 수학적 지식을 요구하는 수식을 보고, 해당 모델의 아이디어를 읽어낼 수 있어야 합니다.
2) 동시에, 그 아이디어를 code로 implement할 수 있어야 합니다.

기존 1)에 집중했던 포스트들과는 다르게, 이번에는 Pytorch code를 함께 보며 원리를 설명하고, 가능한 수학적 원리를 자세하게 다루되 가볍게 이해하실 수 있도록 설명해보겠습니다.

이번 포스트의 실습 코드는 제 깃허브의 justin4ai/pytorch-mnist-vae repository를 통해서도 확인하실 수 있습니다.

코드는 MNIST generation을 기준으로 설명하지만, 원리만 이해하면 input data에 따라 유동적으로 구조를 변경할 수 있게 됩니다. 이를 위해 각 parameter의 역할까지 자세히 설명할 예정입니다!

이 글의 목표는 generative AI에 처음 관심을 가지고 VAE를 공부해보시는 분들이, 여러 사이트를 돌아다니실 필요 없이 포스트 단 하나로 완벽히 이해할 수 있도록 돕는 것입니다.

만약 VAE에 대해 대략적으로 알고 계시지만 가장 어려운 부분들만 선택적으로 이해하고 싶으신 분들께서는, Reparameterization Trick과 Loss chapter로 바로 넘어가셔도 됩니다!

Prerequisites

이번 포스트는 아쉽게도 비전공자분들을 위한 포스트는 아닙니다.

1) Backpropagation(오차역전파)에 대한 기본적인 지식이 있고,
2) Linear/convolutional layers를 이용한 Neural Network의 구성에 능숙하셔야 하며,
3) Entropy이론에 대한 기본적인 이해가 있으셔야 합니다.

Note 1) Autoencoder에 대한 사전 이해가 있으면 이해하시는 데에 좋습니다.
Note 2) 하지만 위의 3가지만 만족하신다면, 생성형 AI에 대한 지식 없이도 이해하실 수 있습니다.

이 정도를 가정하고, 해당하는 독자 분들도 수학이나 복잡한 코드 implementation 부분에서 막힘이 없이 이 포스트 하나로 VAE를 졸업하실 수 있도록 잘 설명해보겠습니다.

Introduction

이전 제 Deepfake의 원리와 Stable Diffusion의 원리 포스트를 보셨다면 아시겠지만, 우리가 막연하게 신기술로만 생각했던 generative AI는, 사실 여러분이 한 번쯤 들어보셨을 법한 모델들을 backbone으로 여러 모듈과 결합하여 개발되는 기술입니다.

그러한 모델들의 대표적인 예시로, GAN, VAE, Flow-based models 그리고 최근 가장 활발히 연구되고 있는 Diffusion models가 있습니다.

이 중에서, VAE와 Diffusion model은 수학적으로 연결 고리가 있기 때문에, 이번 포스트를 통해 VAE를 수학적으로도 이해한다면 추후 generative AI의 SOTA들을 이해하시는 데에 강점이 생기실겁니다.

보시면 VAE : Maximize ELBO( $=$ variational lower bound)라고 나와있습니다. 말 그대로, 아직 무엇인진 모르시더라도 VAE는 ELBO라는 것을 최대화하는 학습이 특징임을 기억해주세요!

Background

본격적으로 VAE를 분석하기 전, 바로 이전 세대에 해당하는 AE(AutoEncoder)의 무엇이 부족해 VAE가 개발되었을 지를 살펴보겠습니다.

AE(AutoEncoder)

AE는 근본적으로 generative model의 역할이라기 보다는, input의 정보를 잘 담아낸 compact(=latent) code $z$ 를 압축하는 encoder와 이를 다시 input과 유사한 data로 복호화하는 decoder를 구성하도록 설계되었습니다.

AE는 결국 $x \rightarrow z \rightarrow x'$ 의 구조 때문에, loss로 MSE를 사용하던 BCE를 사용하던 결국 reconstruction(= $x' \sim x$ )만 잘 되도록 설계되어 있습니다. 따라서 이러한 단순한 구조는, latent space를 아름답게 구성하기엔 부족합니다.

Slide credit : Understanding VAEs

우측의 latent space에서는 비슷한 의미를 가진 벡터는 서로 붙어있고, 다른 의미를 가진 벡터끼리는 잘 분리된 반면, 좌측의 latent space는 벡터의 의미와 상관 없이 뒤섞여 있습니다.

즉, AE의 단순한 loss로 인해 고도화된(=regular) latent space의 구성이 어려워지고, 실제 데이터를 대표하지 못하는 irregular latent space를 구성하게 됩니다.

그렇다면 input data $x$ 의 latent representation인 $z$ 의 분포 자체를 잘 학습할 수 있도록 loss를 구성한다면 어떨까요?

$x'$ 이 $x$ 와 비슷하게 만들어지도록 latent space를 제한하는 방식이 아니라, $z$ 의 분포 자체를 학습할 수 있다면, $x$ 를 reconstruct하는 방법이 아닌 sampling된 noise $z$ 로부터 decode하여 realistic한 데이터를 생성할 수 있게 될 것입니다.

Difference with AE

한 호흡으로 쭉 가져가기 전에, AE와의 결정적인 차이점을 바로 짚어보겠습니다.

그림에서 보시다시피, 가운데에 input data가 최대한으로 압축되는 부분이 존재합니다. 이 부분을 bottleneck이라고 부릅니다. AE는 input data $x$ 를 잘 압축하는 것만이 목표이지만, VAE는 latent space 상에서 input data의 representative인 $z$ 가 아름답게 구성되어 있기를 바랍니다.

그러한 latent space를 구축하도록 강제하기 위해 standard deviation과 mean를 구하고, 학습 과정의 loss에서 $KLD$ 에 포함시킵니다.

이 logvar과 mu를 어떻게 활용하여 VAE가 원하는 목표를 달성하는지는, Reparameterization Trick과 Loss chapter에서 다루겠습니다. 당장은 bottleneck 부분에서 어떤 차이가 있는 지만 아시면 됩니다!

VAE Principle with Codes

가장 핵심적인 부분이며, 양이 방대한 만큼 상대적으로 덜 중요한 부분은 최대한 담백하게 설명하겠습니다. 혹시라도 궁금한 점이 있으시면 댓글을 남겨주세요!

Input

# Data preprocessing
dataset = dset.MNIST(root=DATA_PATH, download=True,
                     transform=transforms.Compose([
                     transforms.Resize(X_DIM), # Reszie from 28x28 to 64x64
                     transforms.ToTensor(),
                     #transforms.Normalize((0.1307,), (0.3081,))
                     ]))

# Dataloader
VAEdataloader = torch.utils.data.DataLoader(dataset, batch_size=BATCH_SIZE,
                                         shuffle=True)

Input으로 들어갈 dataset인 MNIST를 다운로드하면서, 본래 (LARGE_NUM, 1, 28, 28) 크기인 LARGE_NUM개의 숫자 이미지 데이터를 (LARGE_NUM, 1, 64, 64)로 resize합니다.

ToTensor()를 사용했기 때문에, 모든 값의 범위가 0~1이 되었음을 꼭 알아두셔야 합니다.
필요에 따라 적절한 normalization을 진행하되, 꼭 생성한 tensor를 MNIST image로 바꿀 때 역연산을 취해주어야 합니다.

VAEdataloader 객체는 한 번 iterate할 때마다 (BATCH_SIZE, 1, 28, 28)의 tensor를 return합니다.

Backbone

Backbone으로는 어떠한 구조를 사용하던 상관이 없습니다. 그러나, input data의 특징을 잘 잡아낼 수 있도록 구성된 네트워크여야 합니다.

우리의 MNIST는 VAE에게 있어 생성 난이도가 낮은 편일 것이므로, 비교적 단순한 conv2d & conv2dtranspose의 조합으로 다음과 같이 구성합니다.

말씀드렸듯이, AE와 bottleneck 부분에서 차이가 있으나, encode-decode 구조는 동일합니다.
다음은 위의 architecture에 대한 Pytorch implementation입니다.

간단화를 위해 본래 적용한 BatchNorm 등을 삭제했으니, 꼭 full code와 함께 학습하시길 권장드립니다.

class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size()[0], -1).to(device) # conv2d로 구성된 encoder의 output은 입체 tensor.
                                                          # 이를 linear한 fc1, fc2 각각과 연결시키기 위해서, BATCH 차원을 제외한 모든 차원을 하나로 합친다.
class UnFlatten(nn.Module):
    def forward(self, input):
        return input.view(input.size()[0], 64, 2, 2).to(device) # Flatten class의 역연산.

class VAE(nn.Module):
    def __init__(self, image_channels= IMAGE_CHANNEL, output_channels= INITIAL_CHANNEL, h_dim=256, z_dim=16): # h_dim : 마지막 hidden dimension, z_dim : latent dimension
        super(VAE, self).__init__()
        self.z_dim = z_dim
        self.encoder = nn.Sequential(
            nn.Conv2d(image_channels, output_channels, kernel_size=3, stride=2, padding = 1),
            nn.ReLU(),
            nn.Conv2d(output_channels, output_channels*2, kernel_size=3, stride=2, padding = 1),
            ...
            nn.Conv2d(output_channels*8, output_channels*16, kernel_size=3, stride=2, padding = 1),
            nn.ReLU(),
            nn.Dropout(0.8),
            Flatten() # (BATCH_SIZE, ?, ?, ?) -> (BATCH_SIZE, ?)
        )
        
        self.fc1 = nn.Linear(h_dim, z_dim).to(device) # reparametriazation의 첫 번째 argument인 mu를 위한 마지막 fc layer
        self.fc2 = nn.Linear(h_dim, z_dim).to(device) # reparametriazation의 두 번째 argument인 logvar를 위한 또다른 마지막 fc layer 

        self.fc3 = nn.Linear(z_dim, h_dim).to(device) # decoder의 시작 부분
        
        self.decoder = nn.Sequential(
            UnFlatten(), # (BATCH_SIZE, ?) -> (BATCH_SIZE, ?, ?, ?)
            nn.ConvTranspose2d(output_channels*16, output_channels*8, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(output_channels*8, output_channels*4, kernel_size=3, stride=2, padding=1, output_padding=1),
            ...
            nn.ConvTranspose2d(output_channels*2, output_channels, kernel_size=3, stride=2, padding=1, output_padding=1),        
            nn.ReLU(),
            nn.ConvTranspose2d(output_channels, image_channels, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.Sigmoid() # input의 value range가 0~1이기 때문에 output의 range도 맞춰주기 위함
        )
        
    def reparameterize(self, mu, logvar):

        std = logvar.mul(0.5).exp_()
        esp = torch.randn(*mu.size()).to(device)
        z = mu + std * esp # N(mu, std) ~ N(0, 1) * std + mu
        return z

    def bottleneck(self, h):
        mu, logvar = self.fc1(h), self.fc2(h) # activation function이 없어야 함
        logvar = torch.clamp(logvar, min=-4, max=4) # 추후 exp()로 인한 exploding 방지
        z = self.reparameterize(mu, logvar)
        return z, mu, logvar

    def encode(self, x):
        h = self.encoder(x)
        return self.bottleneck(h)
        
    def decode(self, z):
        z = F.relu(self.fc3(z))
        z = self.decoder(z)
        return z

    def forward(self, x):
        z, mu, logvar = self.encode(x) # mu와 logvar은 encoder가 지나고 저장해두기
        z = self.decode(z) # reparameterize된 z를 decode
        return z, mu, logvar

Q. Bottleneck 부분에서 어떤 부분이 달라졌을까요?

Reparameterization Trick

    def reparameterize(self, mu, logvar):
        std = logvar.mul(0.5).exp_()
        esp = torch.randn(*mu.size()).to(device)
        z = mu + std * esp # N(mu, std) ~ N(0, 1) * std + mu
        return z

위의 코드를 보시면, reparameterize이라는 함수가 있었습니다. reparameterization은, 미분이 불가했던 변수를 다른 변수와의 조합을 통해 미분이 가능하도록 만드는 trick입니다.

갑자기 왜 등장했을까요?
우선 저희가 원하는 것을 말해보겠습니다.

저희는 VAE를 통해, noise $z \sim N(\mu, \sigma^2)$ 로부터 output을 생성해내기를 원합니다.

여기서 $\mu$ 와 $\sigma$ 는, latent space의 distribution에 대한 값입니다.

이러한 가정 하에,

Backpropagation 과정 중 위의 그림의 Summing Junction을 VAE encoder과 연결된 bottleneck의 output, 즉 $\textbf{z}$ 라고 해보겠습니다. 이를 Chain Rule을 적용시키기 위해 각 $x$ 에 대해 편미분하면,

z = \epsilon \sim N(\mu, \sigma^2) \;\;\rightarrow\;\; \frac {\partial z}{\partial x_1}, \frac {\partial z}{\partial x_2}, \frac {\partial z}{\partial x_3} ... = 0

모두 0이 되어버려 이전의 nodes를 통한 역전파가 불가능해집니다. 이러한 성질 때문에, 저희는 $z = \epsilon$ 에서 $\epsilon$ 를 stochastic(random) term이라고 합니다.

우리는 last level의 nodes로부터 first level nodes까지의 Chain Rule을 통한 loss update를 원합니다. 따라서, bottleneck으로부터 first level nodes까지 가는 길이 끊겨서는 안됩니다(= 편미분값이 0). 즉, stochastic(random) term이 아닌 deterministic term을 원합니다.

복습해보겠습니다.

$z$ 를 단순히 input data의 encoded vector를 사용하는 AE와는 달리, VAE는 본질적으로 $z$ 를 random noise로 구성하여, 이로부터 output을 생성해내도록 설계되었습니다.

AE의 $\textbf{z}$ 의 경우, 이전 layers의 nodes에 대한 식으로 나타낼 수 있으므로, Chain Rule을 통해 encoder 파트의 모든 부분에 오차를 역전파시킬 수 있습니다.
하지만, VAE의 $\textbf{z}$ 는 이전 layers의 nodes와 무관한 별개의 분포에서 random sampling된 것이므로, 편미분한 기울기가 모두 0이 되어버려 역전파를 할 수 없게 됩니다.

즉, 수학적으로 encoder 부분과 decoder 부분이 연결되어있지 않기 때문에, bottleneck에서 encoder로 향하는 방향으로의 backpropagation가 모두 불가능해진다는 것입니다.

그러나 reparameterization을 통해서,

Stochastic term으로만 이루어졌던 $\textbf{z}$ 를 stochastic term과 deterministic term의 합으로 표현할 수 있게 됩니다.

이를 수식으로 표현하면 다음과 같습니다.

z_x = \mu_x + \sigma_x*\epsilon, \;\;\;\;\;\; \epsilon \sim N(0,1)

예를 들어 위와 같은 reparameterization을 통해, 기존에는 $x$ 로 표현되는 변수가 아니기 때문에 편미분 값이 0이었던 $\epsilon$ 를, 우리가 encoder 부분에서 사용하는 nodes들(예시의 $x$ 와 같은)의 조합으로 표현함으로써 backpropagation이 가능해졌습니다.

이제, 우리는 decoder 부분만 떼어 그에 대한 input으로 단순히 $z \sim N(0,1)$ 만 주더라도, 이미 VAE는 이러한 상황에 맞춰 학습되어 있으므로 잘 구성된 latent space로부터 원하는 data를 생성할 수 있게 됩니다.

여기서 이해를 아주 잘 하신 분들이라면(!) 의문을 가지실 겁니다.

Q. $z$ 가 아니라 $\epsilon$ 가 $N(0,1)$ 를 따르며, 코드 상에선 reparameterization은 encoder에 있기 때문에, decoder만 사용하여 $z$ 로부터 data를 생성하려면 input으로 $z \sim N(\mu,\sigma)$ 를 sampling해야 하는 것이 아닌가요?

    def bottleneck(self, h):
        mu, logvar = self.fc1(h), self.fc2(h) # activation function이 없어야 함
        logvar = torch.clamp(logvar, min=-4, max=4) # 추후 exp()로 인한 exploding 방지
        z = self.reparameterize(mu, logvar)
        return z, mu, logvar

    def encode(self, x):
        h = self.encoder(x)
        return self.bottleneck(h)
        
    def decode(self, z):
        z = F.relu(self.fc3(z))
        z = self.decoder(z)
        return z

    def forward(self, x):
        z, mu, logvar = self.encode(x) # mu와 logvar은 encoder가 지나고 저장해두기
        z = self.decode(z) # !!!!!reparameterize된 z를 decode!!!!!!
        return z, mu, logvar

정답입니다. 무조건 $N(0, 1)$ 로부터 뽑는 것이 아니라, 강제된 latent space에 따라서 해당 분포로부터 추출해야 합니다.

하지만 보통의 VAE는 latent space를 $N(0, 1)$ 로 강제시키는 loss function을 사용하기 때문에, input으로 zero-mean Gaussian noise를 넣는 것입니다. 중요한 점은, $input \sim N(0, 1)$ 을 보시고 input을 $z$ 가 아닌 $\epsilon$ 으로 생각하시면 안된다는 이야기입니다.

이를 통해서, $\mu$ 와 $\sigma$ 가 각각 0과 1이 나와야 잘 학습된 것임을 알 수 있습니다.
(이 설명은 다른 포스트에서 다뤄진 것을 본 적이 없어서, 확인을 위해 제가 진행해본 재밌는 실험의 결과를 포스트 끝 부분에 첨부하였습니다!) ~~역시 AI 공부는 너무 재미있네요~~

Loss

VAE의 꽃은 두 가지가 있습니다. 하나는 reparameterization, 다른 하나는 ELBO(Evidence of Lower BOund) 입니다. ELBO는 code에서 직접 등장하진 않으나, 이와 밀접한 관계를 가진 KLD가 VAE의 loss function에서 등장하게 됩니다.

def loss_fn(recon_x, x, mu, logvar):
    BCE = F.binary_cross_entropy(recon_x, x, reduction='sum')
    
    #BCE = nn.MSELoss()
    #BCE = BCE(recon_x, x) # input/output의 range가 0~1이 아니라면,
                           # BCE를 사용할 수 없으므로 MSELoss를 사용

    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp() + 1e-10) # 안정적인 학습을 위해 아주 작은 noise를 추가
    return BCE + KLD, BCE, KLD

BCE : BCE를 사용할 경우 $x$ 와 reconstruct된 $x'$ 의 분포 차이를 줄이도록 학습할 수 있게 됩니다. 그러나 값이 0~1의 범위에 있지 않다면, 그 둘의 차이를 단순하게 MSE(Mean Squared Error)를 통해 줄이도록 학습할 수도 있습니다.
KLD : KL-Divergence를 나타내며, BCE처럼 분포에 관한 이야기이지만, latent space 상에서의 분포들에 대한 함수입니다. $\mu$ 와 $\sigma$ 로 표현되는 latent space의 distribution이 $N(0,1)$ 에 approximate되도록 구성되었습니다.

이전의 AE vs. VAE에서 latent space의 구성 문제를 다시 짚어보면, regular latent space를 구성하도록 강제하는 것을 목표로 VAE가 탄생했다고 했습니다.

이제 KLD가 그 목표를 이루도록 도울 것입니다.

KLD(KL-Divergence)

Kullback-Leibler Divergence (KLD)는 두 확률 분포 간의 차이를 측정하는 지표 중 하나로, VAE에서는 latent space의 분포를 정규분포에 가깝게 만들기 위해 사용됩니다.

다음과 같은 $B(x)$ 가 있을 때, 데이터들로부터 해당 분포를 정확하게 계산하는 것은 무척 어려운 일일 것입니다.

이제부터 수학 기호가 조금씩 등장할 예정이지만, 거부감을 꾹 참고 차근차근 읽으시면 잘 이해하실 수 있으리라 보장합니다!

만약 이것을 $A(x)$ 라는 우리가 아주 잘 알고있는 Guassian Distribution으로 근사시킬 수 있다고 한다면 얼마나 좋을까요? Latent space의 point인 $z$ 가 우리가 이미 잘 아는 분포인 $N(0, 1)$ 을 따르도록 근사되므로, regular한(= 설명 가능한) latent space를 구성했다고 말할 수 있게 됩니다.

이제부터는 정확한 notation을 사용하겠습니다. 수학 기호가 나온다고 포기하지 않으셔도 됩니다! 자세히 보시면 전부 이해하실 수 있습니다. 다음 두 가지만 잘 구분해주세요:

$z$ : Latent space를 나타내는 변수
$x$ : Input data를 나타내는 변수

$x$ 라는 input이 주어진 상태의 $z$ 의 분포인, 즉 간단하게 어떤 input이 주어졌을 때 latent space의 분포인 $p(z|x)$ 를 우리가 아주 잘 아는 Guassian Distribution을 따르는 $q(z)$ 로 잘 근사시킨다면, 불가능에 가까운 latent space의 분포 자체를 계산하는 것 대신 regular space로 대체할 수 있게 됩니다.

이것을 변분추론(Variational Inference)이라고도 하지만, 저는 앞으로 이 용어를 따로 사용하지는 않겠습니다.

KLD를 수학적으로 나타내면:

D_{KL}(q(z)||p(z|x)) = E_q[logq(z)] - E_q[logp(z|x)]\;\;\; (1)

이것은 KLD의 정의로, 헷갈리신다면 각각 $A(x)$ 와 $B(x)$ 로 대체하여 넣어보겠습니다.

D_{KL}(A(x)||B(x)) = E_A[logA(x)] - E_A[logB(x)]\;\;\;

여기서, $E$ (기댓값)의 아래 첨자에 있는 $A$ 는, probability function으로 $A$ 를 사용함을 명시해주기 위함입니다. 이산 확률 변수의 기댓값을 구할 때 $\Sigma{variable}*probability$ 와 같은 연산이 이루어지기 때문에, 어떠한 확률 함수인지를 꼭 명시해주어야 합니다.

정보 이론의 Entropy(기초이기 때문에 꼭 알고 계셔야 합니다!) notation에 의하여, 다음과 같이 쓸 수도 있습니다.

= H(q(z), p(z|x)) - H(q(z))\;\;\; (2)

즉, $KLD$ 를 loss로 사용함으로써 minimize한다는 것은, $Eq. (2)$ 를 minimize하는 것과 같습니다. 그런데, $H(q(z), p(z|x))$ term의 최선의 상황은 $H(q(z)$ )과 같아지는 것입니다. $p(z|x)$ 가 $q(z)$ 에 완벽히 근사되었음을 의미하기 때문입니다.

한마디로 KLD $= Eq. (2)$ 는 0보다 작아질 수 없다는 것입니다( $KLD \ge 0$ ). 여전히 0보다 작아질 순 없으나, loss function인 $KLD$ 를 minimize함으로써 latent space를 Guassian Distribution에 approximate한다는 사실은 변하지 않습니다.

ELBO(Evidence of Lower BOund)

위의 그림을 기억하시나요? 방금 KLD를 minimize하는 loss를 구성한다고 결론을 내렸는데, $ELBO$ 라고 불리는 variational lower bound를 maximize한다고 써있었습니다.

제 설명이 잘못되지 않은 이상, 다음과 같은 결론을 내실 수 있습니다.

VAE의 학습을 minimizing KLD와 maximizing ELBO로 interchangable하게 나타낼 수 있다는 것은, KLD + ELBO가 상수임을 의미할 '수도' 있습니다(아직은 알 수 없지만요).

$ELBO$ 를 잠깐 머리에서 지우시고, $KLD$ 를 다시 불러오겠습니다.

D_{KL}(q(z)||p(z|x)) = E_q[logq(z)] - E_q[logp(z|x)]\;\;\; (1)

두 번째 term에서, posterior probability인 $$p(z|x)$$에 **Bayes' Theorem**을 적용하면 다음과 같습니다.

= E_q[logq(z)] - E_q[log{\frac{p(x|z)p(z)}{p(x)}}]\;\;\; (3)

두 번째 term의 조건부확률을 joint distribution function으로 바꾸겠습니다.

= E_q[logq(z)] - E_q[log{\frac{p(x,z)}{p(x)}}]\;\;\; (4)

Log의 성질에 의해, 두 번째 term을 두 terms의 차로 나타낼 수 있습니다.

= E_q[logq(z)] - E_q[log{p(x,z)} - logp(x)]\;\;\; (5)

= E_q[logq(z)] - E_q[log{p(x,z)}] - E_q[ - logp(x)]\;\;\; (5.1)

= E_q[logq(z)] - E_q[log{p(x,z)}] + E_q[logp(x)]\;\;\; (5.2)

마지막 term에서 q(x)와 p(x)는 독립된 분포이기 때문에 q에 대한 기댓값 연산은 trivial하므로,

= E_q[logq(z)] - E_q[log{p(x,z)}] + logp(x)\;\;\; (6)

즉,

D_{KL}(q(z)||p(z|x)) = E_q[logq(z)] - E_q[log{p(x,z)}] + logp(x)

입니다. 이를 $logp(x)$ 에 대하여 정리하면,

logp(x) = D_{KL}(q(z)||p(z|x)) + E_q[log{p(x,z)}] - E_q[logq(z)]

이 때, $logp(x)$ 라는 것은 사실 learnable parameter인 $z$ 와 완전히 무관하게 input data의 distribution을 나타내기 때문에, $z$ 를 위한 backpropagation 과정 중에는 상수로 생각할 수 있습니다.

위의 식에서 두 번째 항과 세 번째 항을 묶어서 $ELBO$ 라고 부릅니다.

ELBO = E_q[log{p(x,z)}] - E_q[logq(z)]

$ELBO$ 의 첫 번째 항은, $q(x)$ 가 클 때 $p(x, z)$ 또한 큰 값을 가질 때 값이 커집니다. 즉, 모델이 높은 probability를 할당하는 곳에 q(x)도 focus하도록 만드는 것입니다.

$ELBO$ 의 두 번째 항은, $H(q)$ 를 의미하므로 chaos에 가까워질 수록 값이 커집니다.

따라서 첫 번째 항과 두 번째 항은 trade-off 관계에 있으며, $ELBO$ 를 maximize한다는 것은 $KLD$ 를 minimize하는 것과 같음을 알 수 있습니다.

Train

from tqdm import tqdm
from torch.optim.lr_scheduler import StepLR

model = VAE()
model = model.to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=0.0003) 
scheduler = StepLR(optimizer, step_size=25, gamma=0.9)  # 25 epoch마다 90%로 감소


from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter("runs/mnist_bce_test")
epochs = 500

for epoch in range(epochs):
    for idx, (images, label) in enumerate(tqdm(VAEdataloader, desc=f'Epoch {epoch + 1}/{epochs}')):
        optimizer.zero_grad()

        images = images.float().to(device) # train image (BATCH_SIZE, 1, 64, 64)
        recon_images, mu, logvar = model(images) # encoder-decoder를 모두 거쳐 나온 output과,
                                                 # encoder의 끝에서 저장해둔 mu와 logvar
        
        loss, bce, kld = loss_fn(recon_images, images, mu, logvar) # loss 계산
        
        if torch.isnan(loss).any(): # loss가 nan이면 안됨
            print("NaN value in loss!")
            break

        loss.backward()
        optimizer.step()

    scheduler.step()

    print(f"Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.4f} (total loss), {bce.item():.4f} (bce), {abs(kld.item()):.4f} (kld)")
    print(f"Mu range: {torch.min(mu[0])} ~ {torch.max(mu[0])}, Logvar range: {torch.min(logvar[0])} ~ {torch.max(logvar[0])}")
        
    writer.add_scalar('Training Loss', loss.item(), epoch) # tensorboard에 log 기록

writer.close()
torch.save(model, './checkpoint/mnist_vae_bce_test.pth')

Train code에 대해서는 특별히 언급할게 없는 것 같습니다. $KLD$ loss가 잘 줄어들지 않는다면, mu와 logvar 값이 각 batch에 대하여 0으로 잘 converge하고 있는지를 확인하셔야 합니다.

$N(0,1)$ 을 따르도록 하기 위해선 variance가 1이 되어야 하니, log variance는 0이 될 때가 최적인 것이 맞습니다.

Generation

이제 train이 완료된 모델로 generate를 해보겠습니다.

Ground truth images

Reconstructed images

Generated images from noises

약간의 blur된 듯한 이미지는 VAE의 단점이기도 합니다. 이러한 특성 때문에, 선명도가 중요한 image data에 대해서는 GAN을 선호하는 분들이 더 많은 것 같습니다.

그러나 GAN은 latent space에 대한 학습은 없기 때문에, 이를 절충한 VAE-GAN도 존재합니다. 최근 직접 code implementation까지 해보았기 때문에, 별도의 포스트로 다루지는 않겠지만 pytorch-mnist-vae repository에 이번 달 내에 간단한 설명과 함께 올려놓도록 하겠습니다!

Experiment

Latent space가 정말 zero-mean Guassian distribution에 잘 근사되었는지 확인하는 방법은, 전체 loss가 아니라 KLD의 추세를 보는 것입니다.

그런데, 기존과 같은 loss function을 사용하면:

시작과 함께 $KLD$ 가 떨어지기는 커녕 오히려 발산하는 것처럼 보이며, BCE loss만 떨어지고 있음을 알 수 있습니다. 이는 실제로 $\mu$ 와 $log\sigma$ 가 전혀 0으로 converge하고 있지 않음을 통해서 확신할 수 있었습니다.

해석해보면, BCE를 통해 $x \sim x'$ 가 되도록 분포를 근사시키는 데에는 성공했으나, latent space의 분포가 $N(0,1)$ 를 근사하는 데에 실패했으므로, input으로 zero-mean Gaussian noise를 주었을 때 결과가 그다지 좋지는 못한 것입니다.

def loss_fn(recon_x, x, mu, logvar):
    BCE = F.binary_cross_entropy(recon_x, x, reduction='sum')
    
    #BCE = nn.MSELoss()
    #BCE = BCE(recon_x, x) # input/output의 range가 0~1이 아니라면,
                           # BCE를 사용할 수 없으므로 MSELoss를 사용

    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + 30 * KLD, BCE, 30 * KLD

그렇다면, $KLD$ 에 30이라는 weight를 주고 다시 train해보겠습니다.

BCE loss의 개선 속도는 조금 느려졌지만, $KLD$ 가 정상적으로 떨어지면서 $\mu$ 와 $log\sigma$ 또한 0에 converge하는 방향으로 학습이 진행됐습니다.

이 모델로 generation을 새롭게 해보면:

Generated images from noises (BCE + KLD * 30)

Generated images from noises (BCE + KLD)

Reconstruction 성능은 더 떨어졌으나(blurred된 듯한 느낌), noise로부터 생성하는 능력이 조금은 개선되었다고 생각되지 않으시나요?

제가 생각했을 때에는, $KLD$ 에 20 근처의 weight를 주고, configuration을 조금 다르게 하면 reconstruction과 generation이 모두 훌륭한 VAE를 만들 수 있을 것 같습니다.

이렇듯 같은 VAE 내에서도, 여러 가지 configuration을 영리하게 조절함으로써 성능을 향상시키거나, 원하는 방향으로 학습을 유도할 수도 있습니다.

저는 포스트 작성을 위해 무작위에 가까운 configuration을 진행했기 때문에, full code로 실습하시면서 성능 개선에 도전해보시는 것도 재밌을 것 같습니다 ㅎㅎ

Outro

수학적인 원리까지 설명하는 것이 처음이고 Latex도 처음이라 생각보다 굉장한 노력이 들어갔네요..^^.. 부족한 부분도 무척 많지만, 제가 target으로 생각한 독자의 수준에는 어느 정도 맞춘 것 같아 스스로 고생했다고 칭찬해주고 싶습니다.

다음 포스트는 Diffusion Model의 원리를 다룰 예정입니다. 저는 generative AI enthusiast로서 최근 DM을 통한 연구 성과들에 감명받았고, 2년 전 제가 우리나라 최초로 분석을 다뤘던 Stable Diffusion 포스트에서보다 깊은 수학적 원리를 설명드릴 수 있는 좋은 기회가 될 것이라고 생각합니다.

당시와는 달리 이제 diffusion-based model이 generative AI의 SOTA에서 차지하는 비중이 점점 많아지게 되었고, 문제 의식을 느껴 DMs를 열심히 공부하게 되었습니다. 그래서 포스트에서 바로 다룰까 하다가, 선행지식의 개념으로 VAE를 다뤄 놓으면 흐름상 더 좋을 것 같아서 VAE를 먼저 작성하게 되었구요.

혹시라도 질문이 있으신 분께서는, 꼭 댓글을 남겨주시면 자세하게 답변해드리겠습니다. 만약 댓글이 불편하시면 이메일을 주셔도 괜찮습니다.

또한 피드백 및 오류 지적은 언제나 환영이며, 미리 감사드립니다.

항상 generative AI의 발전을 위해 열심히 공부하시는 여러분께 감사드리고, 제 포스트로써 도움이 된다면 정말 행복할 것 같습니다.

취업 시장에 겨울이 찾아온다고 하지만, 스스로를 뜨겁게 만들어 시장의 추위를 녹여버리는 사람이 되는 것을 목표로 하고 있습니다. 방학 기간에도 제 글을 끝까지 읽어주신 분들께 진심으로 감사드립니다 🎶

안준영

Dedicated my life to AI

이전 포스트

Deepfake 탐지 : 음지의 Deepfake에 맞서는 Transformer-Based AI

다음 포스트

Sora 설명 : 텍스트 한 줄이 영상을 만든다 (Feat. Stable Diffusion 3)

2개의 댓글

서기현

2024년 1월 22일

BCE Loss를 계산할 때 reduction 옵션을 sum으로 사용하신 이유가 있을까요?

1개의 답글