GAN, VAE, Flow, Diffusion 무작정 코딩해보기 | 생성모델 공부 (1)

출처: MS Bing Image Creator, getimg.ai

 

  릴리안 웡의 블로그 글을 포함해서 생성모델에 대한 유용한 글은 널리고 널렸지만, 아무리 읽어봐도 쉽게 체득이 되지를 않는다. 그나마 그림으로 보면 '아 그런갑다' 하겠는데, 여전히 수식 보는 눈은 까막눈이고, 최신 생성모델 구현 코드들을 봐도 이해가 잘 되질 않는다.

 

생성모델 정리의 국룰 다이어그램, GOAT, 교과서. (출처: Lil'Log)

 

  위 네가지 모델들(+ 그리고 GPT와 같은 언어모델들)이 2010년대 후반부터 등장한 생성모델들의 기초 모델들이다. 최근 생성모델들은 위 네 가지 모델들을 각각 발전시키거나, 서로 합치고 개조해서 만들어졌다고 보면 된다. 때문에 이 네 가지 모델들에 대한 이해는 필수적이다.

 

	GAN	VAE	Normalizing Flow	Diffusion
생성 품질	⭕	❌	🔺	⭕
생성 속도	⭕	⭕	🔺	❌
생성 다양성	❌	⭕	⭕	⭕

생성모델 종류별 특징. (출처: 엔비디아)

 

  네가지 모델들을 간단히 설명하면 다음과 같다.

 

• GAN
  • 가짜 이미지를 만들어내는 위조지폐범(Generator)과 이미지가 가짜인지 진짜인지 판별하는 경찰(Discriminator)을 동시에 학습시켜서, 최종적으로 위조지폐범(Generator)이 더 좋은 품질의 이미지를 생성하도록 한다.
  • 비교적 다양한 샘플을 생성해내지 못한다는 단점을 갖고 있으며, 특히 잘못 훈련될 경우 Mode Collapse라는 현상에 빠져 입력과 상관없이 동일한 샘플만을 생성하게 될 수도 있어 훈련과정에서 주의가 필요하다.

 

• VAE
  • 이미지를 책에 비유 했을 때, 여러 가지 책들의 요약본을 만들어 이 요약본들의 대략적인 경향성을 파악하여 랜덤한 요약본으로부터 전체 책을 복원해내는 작가(Decoder)를 학습시키는 것과 비슷하다.
  • 생성 샘플의 품질이 좋지 못한 것으로 알려져 있다.

 

• Normalizing Flow
  • 이미지를 요리에 비유 했을 때, 요리를 바탕으로 필요한 여러 가지 단순한 재료들의 비율과 양을 추정하고(역변환), 다시 여러가지 재료들을 섞고 조리하여 해당 요리를 복원(정변환)해내는 요리사를 학습시킨다.
  • 변환 과정에 반드시 역변환이 가능한 함수가 쓰여야 하기 때문에, 사용 가능한 함수에 제한이 있고, 전체 모델의 속도도 어떤 변환 함수가 쓰였는지에 따라 달라질 수 있다.

 

• Diffusion
  • 이미지를 별에 비유하자면, 실험실에 가상 우주를 만드는 것과 비슷하다. 별의 탄생 이전 단계로 시간을 되돌려서 우주의 먼지로 바꿨다가, 다시 시간을 앞당기면서 비슷한 공간에 비슷한 별을 만드는 식이다.
  • 생성 품질과 다양성에서 강점을 가지지만, 시간이 굉장히 오래 걸린다는 것으로 유명하다. 

 

  실제로는 어떨까? 구현하기 쉽도록 만만한 MNIST 데이터셋을 학습시켜 숫자 손글씨를 만들어주는 생성모델을 만들어보고 실험해보기로 계획하고 바로 실천에 옮겨보았다. (와중에 코딩마저 ChatGPT + MS Bing Chat의 힘을 빌려가며 진행했다. 생성모델 코드 마저 생성모델의 힘을 빌려야만 하는 수준...)

 

 

  PyTorch 1.11.0을 사용해서 GAN, VAE, Normalizing Flow, Diffusion 모델을 만들되, 크기와 구조가 서로 최대한 비슷한 아키텍처를 갖도록 했다. 옵티마이져는 동일하게 Adam을 사용하고, Learning Rate Scheduler도 동일하게 ExponentialLR을 사용했다. 손실함수를 어떻게 구성하고 모델을 어떻게 학습시킬지를 달리하며 계속 실험해보면서 각 모델 종류에 맞는 실험 구성을 찾기 위해 삽질을 해보았다.

 

GAN (Generative Adversarial Network)

• 구글 코랩에서 실행:
• 모델 코드: shhommychon/mnist_generative_practice

  • import torch
    import torch.nn as nn
    
    
    class MNISTGan(nn.Module):
        def __init__(self, input_dim=128, output_dim=28*28, label_dim=4, n_classes=10):
            super(MNISTGan, self).__init__()
    
            # label embedding as condition for GAN
            if n_classes > 0:
                self.label_emb = nn.Embedding(n_classes, label_dim)
            else:
                label_dim = 0
    
            self.conv_1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1)
            self.conv_2 = nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1)
            self.conv_3 = nn.Conv2d(in_channels=32, out_channels=1, kernel_size=3, stride=1, padding=1)
    
            self.hidden_layer = nn.Linear(input_dim+label_dim, output_dim)
    
            self.relu = nn.ReLU()
            self.tanh = nn.Tanh()
    
        def forward(self, x, labels=None):
            if labels is not None:
                condition = self.label_emb(labels)
                x = torch.cat([x, condition], 1)
    
            x = self.hidden_layer(x)
    
            x = x.view(x.size(0), 1, 28, 28)
    
            x = self.conv_1(x)
            x = self.relu(x)
    
            x_ = self.conv_2(x)
            x_ = self.relu(x_)
            x = x + x_
    
            x = self.conv_3(x)
    
            x = self.tanh(x)
    
            return x #.type(torch.float16)
    
    
    class Discriminator(nn.Module):
        def __init__(self):
            super(Discriminator, self).__init__()
            self.hidden_layer = nn.Linear(28*28, 256)
            self.output_layer = nn.Linear(256, 1)
            self.relu = nn.ReLU()
            self.sigmoid = nn.Sigmoid()
    
        def forward(self, y):
            y = y.view(y.size(0), -1)
            y = self.hidden_layer(y)
            y = self.relu(y)
            y = self.output_layer(y)
            real_fake = self.sigmoid(y)
    
            return real_fake

 

  학습시켜봤을 때, Mode Collapse에 의도적으로 빠지게 하는 게 훨씬 쉬울 정도로 Mode Collapse에 잘 빠졌다. 원래는 Adversarial Loss + Reconstruction Loss의 조합에 Reconstruction Loss를 Pixelwise Loss, 즉 생성된 이미지와 정답 이미지 간 각 픽셀의 차이로 설정했었는데, Reconstruction Loss를 Classification Loss(간단한 MNIST 분류기(코랩, 모델 코드)를 학습시킨 뒤 얼리고 생성된 이미지를 얼마나 잘 분류할 수 있는지 측정)로 교체해 보니 Mode Collapse에서 쉽게 벗어날 수 있었다.

  추가로, Discriminator의 구조를 너무 복잡하게 만들어 놓으면 GAN 모델이 잘 학습되지 않았다.

 

• 생성 결과:

실제 MNIST 데이터(맨 윗줄)와 내 MNIST Gan의 생성 결과(2~5번째 행)

 

VAE (Variational AutoEncoder)

• 구글 코랩에서 실행:
• 모델 코드: shhommychon/mnist_generative_practice

  • import torch
    import torch.nn as nn
    
    
    class MNISTVae(nn.Module):
        def __init__(self, latent_dim=96, image_dim=28*28, label_dim=16, n_classes=10):
            super(MNISTVae, self).__init__()
    
            # label embedding as condition for VAE
            if n_classes > 0:
                self.label_emb = nn.Embedding(n_classes, label_dim)
            else:
                label_dim = 0
    
            # Encoder
            self.fc_mu = nn.Linear(image_dim, latent_dim)  # mu layer
            self.fc_logvar = nn.Linear(image_dim, latent_dim)  # logvariance layer
    
            # Decoder
            self.linear = nn.Linear(latent_dim+label_dim, image_dim)
    
            self.conv_1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1)
            self.conv_2 = nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1)
            self.conv_3 = nn.Conv2d(in_channels=32, out_channels=1, kernel_size=3, stride=1, padding=1)
    
            self.relu = nn.ReLU()
            self.tanh = nn.Tanh()
    
        def encode(self, x):
            return self.fc_mu(x), self.fc_logvar(x)
    
        def reparameterize(self, mu, logvar):
            std = torch.exp(0.5*logvar)
            eps = torch.randn_like(std)
            return mu + eps*std
    
        def forward(self, x, labels=None):
            x = x.view(x.size(0), -1)
            mu, logvar = self.encode(x)
            z = self.reparameterize(mu, logvar)
            return self.decode(z, labels), mu, logvar
    
        def decode(self, z, labels=None):
            if labels is not None:
                condition = self.label_emb(labels)
                z = torch.cat([z, condition], 1)
    
            x = self.linear(z)
    
            x = x.view(x.size(0), 1, 28, 28)
    
            x = self.conv_1(x)
            x = self.relu(x)
    
            x_ = self.conv_2(x)
            x_ = self.relu(x_)
            x = x + x_
    
            x = self.conv_3(x)
    
            x = self.tanh(x)
    
            return x #.type(torch.float16)

 

  Kullback–Leibler divergence의 이해가 어려운 게 특징인 VAE다. 학습 초기에는 안개 낀듯한 뿌연 숫자 이미지를 내뱉고, Pixelwise Loss를 추가로 넣어주지 않는다면 그 특유의 안개 낀듯한 이미지에서 벗어나질 못한다. 제대로 된 학습이 진행 될 수록 점점 선명해지는 숫자 이미지를 얻을 수 있었고, 학습 과정 중간마다 생성되는 이미지를 확인하면서 신기해했다.

 

• 생성 결과:

실제 MNIST 데이터(맨 윗줄)와 내 MNIST VAE의 생성 결과(2~5번째 행)

 

Generative Flow model

• 구글 코랩에서 실행:
• 모델 코드: shhommychon/mnist_generative_practice

  • import torch
    import torch.nn as nn
    import torch.nn.functional as F
    
    
    class PlanarFlow(nn.Module):
        def __init__(self, latent_dim):
            super(PlanarFlow, self).__init__()
            self.weight = nn.Parameter(torch.Tensor(1, latent_dim))
            self.scale = nn.Parameter(torch.Tensor(1, latent_dim))
            self.bias = nn.Parameter(torch.Tensor(1))
            self.reset_parameters()
    
        def reset_parameters(self):
            self.weight.data.uniform_(-0.01, 0.01)
            self.scale.data.uniform_(-0.01, 0.01)
            self.bias.data.uniform_(-0.01, 0.01)
    
        def forward(self, z):
            psi = torch.tanh(F.linear(z, self.weight, self.bias))
            return z + self.scale * psi
    
        def inverse(self, z):
            # This is an approximation to the inverse of the tanh function
            psi = torch.atanh((z - self.bias) / self.weight)
            return z - self.scale * psi
    
    
    class MNISTFlow(nn.Module):
        def __init__(self, input_dim=28*28, output_dim=28*28, hidden_dim=128, n_classes=10):
            super(MNISTFlow, self).__init__()
    
            self.hidden_dim = hidden_dim
    
            # label embedding as condition for generative flow model
            if n_classes > 0:
                self.label_emb = nn.Embedding(n_classes, hidden_dim)
            else:
                hidden_dim = 0
            self.n_classes = n_classes
    
            self.enc = nn.Linear(input_dim, hidden_dim)
            self.dec = nn.Linear(hidden_dim, input_dim)
    
            # Flow
            self.flow = PlanarFlow(hidden_dim)
    
            self.conv_1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1)
            self.conv_2 = nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1)
            self.conv_3 = nn.Conv2d(in_channels=32, out_channels=1, kernel_size=3, stride=1, padding=1)
    
            self.relu = nn.ReLU()
            self.tanh = nn.Tanh()
    
        def encode(self, x):
            x = x.view(x.size(0), -1)
            x_enc = self.enc(x)
            return x_enc
    
        def decode(self, x_enc):
            x = self.dec(x_enc)
            return x
    
        def generate_embedding(self, c=None, batch_num=1):
            # class embedding as condition
            if self.n_classes > 0:
                batch_num = c.size(0)
                c = self.label_emb(c)
            else:
                c = torch.zeros(batch_num, self.hidden_dim, requires_grad=False)
    
            # random
            z = torch.randn_like(c)
            z_enc = z + c
            return z_enc
    
        def transformation(self, input, reverse=False):
            if not reverse:
                x_enc = self.flow(input)
                return x_enc
            else:
                z_enc = self.flow.inverse(input)
                return z_enc
    
        def conversion(self, z_enc, c0, c1):
            c0 = self.label_emb(c0)
            c1 = self.label_emb(c1)
    
            z = z_enc - c0
            z_enc = z + c1
            return z_enc
    
        def forward(self, z_enc):
            x = self.decode(z_enc)
            x = x.view(x.size(0), 1, 28, 28)
    
            x = self.conv_1(x)
            x = self.relu(x)
    
            x_ = self.conv_2(x)
            x_ = self.relu(x_)
            x = x + x_
    
            x = self.conv_3(x)
    
            x = self.tanh(x)
    
            return x #.type(torch.float16)

 

  Flow 모델의 변환 과정에는 ⓐ NICE (Non-linear Independent Components Estimation) (Dinh et al., 2014, doi:10.48550/arXiv.1410.8516), ⓑ RealNVP (Real-valued Non-Volume Preserving) (Dinh et al., 2016, doi:10.48550/arXiv.1605.08803), ⓒ Glow (Kingma&Dhariwal, 2018, doi:10.48550/arXiv.1807.03039), ⓓ MAF (Masked Autoregressive Flow) (Papamakarios et al., 2017, doi:10.48550/arXiv.1705.07057), ⓔ IAF (Inverse Autoregressive Flow) (Kingma et al., 2016, doi:10.48550/arXiv.1606.04934), ⓕ Neural Spline Flows (Durkan et al., 2019, doi:10.48550/arXiv.1906.04032)와 같은 여러 가지 변환 알고리즘들이 쓰일 수 있다. 난 Bing이 추천해 준 대로 Planar Flow (Rezende&Mohamed, 2015, doi:10.48550/arXiv.1505.05770)를 사용했다.

  여담으로, (잘 기억은 안 나지만) 전 회사에서 VITS 모델의 ONNX화와 관련해서 조사할 때, Flow 모델 부분 때문에 ONNX화가 어려웠던 걸로 기억한다.

 

• 생성 결과:

실제 MNIST 데이터(맨 윗줄)와 내 MNIST Flow model의 생성 결과(2~5번째 행)

 

  그리고 내가 처음 Flow 모델을 공부한 게 VITS 때문이었어서, 이 MNIST Flow 모델도 VITS 마냥 역함수를 이용한 변환이 가능하도록 구성했다. 근본 없이 코딩했다고 생각했는데 의외로 뭐가 되는 거 같아서 놀라웠다.

 

• 변환 결과:

실제 MNIST 데이터(맨 윗줄)와 각 이미지들의 변환 결과(각 열 참고)

 

Diffusion model

• 구글 코랩에서 실행:
• 모델 코드: shhommychon/mnist_generative_practice

  • import torch
    import torch.nn as nn
    
    
    class MNISTDiffusion(nn.Module):
        def __init__(self, img_height=28, img_width=28, label_dim=32, n_classes=10, num_steps=10, noise_std=0.25): 
            super(MNISTDiffusion, self).__init__()
            self.img_height = img_height
            self.img_width = img_width
    
            if n_classes > 0:
                self.label_emb = nn.Embedding(n_classes, label_dim)
            else:
                label_dim = 0
            self.n_classes = n_classes
    
            self.num_steps = num_steps
            self.noise_std = noise_std
    
            self.linear = nn.Linear(img_height*img_width+label_dim, img_height*img_width)
    
            self.conv_1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1)
            self.conv_2 = nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1)
            self.conv_3 = nn.Conv2d(in_channels=32, out_channels=1, kernel_size=3, stride=1, padding=1)
    
            self.relu = nn.ReLU()
    
        def forward(self, x, labels=None, backward=True):
            if backward:
                # Backward diffusion process
                x = x.view(x.size(0), -1)
                if self.n_classes > 0:
                    c = self.label_emb(labels)
    
                x = self.linear(torch.cat([x, c], dim=1) if self.n_classes > 0 else x)
                x = x.view(x.size(0), 1, self.img_height, self.img_width)
    
                x = self.conv_1(x)
                x = self.relu(x)
    
                x_ = self.conv_2(x)
                x_ = self.relu(x_)
                x = x + x_
    
                x = self.conv_3(x)
    
                x = x.view(x.size(0), -1)
                
                return x.view(x.size(0), 1, self.img_height, self.img_width)
            
            else:
                # Forward diffusion process
                x = x.view(x.size(0), -1)
     
                noise = torch.randn_like(x) * self.noise_std
                x = x * (1 - self.noise_std) + noise
    
                return x.view(x.size(0), 1, self.img_height, self.img_width)

 

  위의 GAN, VAE, Flow 모델과 구조가 비슷하더라도, Diffusion 모델은 num_step회만큼 모델을 돌려야 한다. 때문에 학습 과정이 드럽게 오래 걸렸는데, Diffusion 모델에 대한 경험이 없어서 그런지 긴 시간 기다려서 나온 결과물의 퀄조차 안 좋아서 제일 골치 아팠던 모델이다.

  내가 뭘 잘못한 것 일수도 있지만, Diffusion 모델 또한 GAN과 유사하게 Mode Collapse와 비슷한 현상이 나타나서 문제였다. 해결하기 위해 GAN 때 사용했던 MNIST 사전학습 분류기(코랩, 모델 코드)를 재활용했지만, GAN 학습 때처럼 드라마틱하게 좋아지지는 않았다.

  결국 적당히 뭐가 나온다 싶을 때 관둬버려서 살짝 아쉽다. Diffusion이 제일 생성물 퀄리티가 좋은 것으로 알려져 있어서 기대를 좀 했는데, 아래에 이미지에 나와 있듯이 내 Diffusion 모델에서는 가장 특이한 창작물(?)이 나오긴 했지만 이 샘플들 퀄리티가 좋다고 말하긴 다소 어려울 것 같다.

 

• 생성 결과:

실제 MNIST 데이터(맨 윗줄)와 내 MNIST Diffusion model의 생성 결과(2~5번째 행)

 

  아래 이미지에서는 각 Backward Diffusion step별 생성 결과를 확인할 수 있다. (맨 윗줄은 비교용 실제 MNIST 데이터들이고) 두 번째 행부터 시작해서 점점 아래 이미지처럼 변해간다는 것을 확인할 수 있는데, 확실히 Diffusion이 많이 진행될수록 이미지 품질은 좋아진다는 것을 직접 확인할 수 있었다.

 

• step별 생성 결과:

내 MNIST Diffusion model의 step별 생성 결과(2~11번째 행)

 

후기

  내심 MNIST 숫자 이미지 데이터 생성 정도야 쉽지 않을까 생각했지만, 마음처럼 되지 않아서 아쉬웠다. 각 모델 주요 수식들에 대한 이해 없이 코딩해서 그런 걸까? 그래도 한번 직접 만들어봤으니 다른 생성모델 코드를 봐도 해당 코드가 대략 어떤 계열의 모델이고 어떤 식으로 돌아가는지 파악하기 훨씬 쉽지 않을까 생각한다. 다음에는 MNIST GPT를 한번 만들어보고 싶다.