[CV] 생성형 모델(Generative models) | VAE부터 Diffusion, Stable Diffusion까지

VAE

Variational Autoenconder(VAE)는 데이터의 잠재 표현(latent representation)을 확률적으로 인코딩하고, 이를 통해 원본을 재구성하거나 새로운 데이터를 생성하는 모델을 말합니다. 

VAE 구조

VAE는 크게 인코더 - 잠재 공간 - 디코더로 구분된 구조를 가지고 있습니다.

인코더(Encoder)

- 입력 데이터 x를 받아서 잠재 변수 z의 확률 분포 매개변수(평균, 분산)을 출력합니다.

- 인코더는 신경망을 통해 구현합니다.

잠재 공간(Latent Space)

- 가우시안 분포를 따르는 잠재 변수 z를 샘플링합니다.

- 샘플링할 때 reparametrization을 사용합니다. 

디코더(Decoder) 

- 샘플링된 잠재 변수 z를 받아서 원본 데이터 공간의 재구성 x'을 출력합니다. 

- 디코더는 신경망을 통해 구현합니다. 

 

여기서 구조에 따라 두가지 손실이 존재합니다. 하나는 재구성 손실(reconstruction loss)이 있고 다른 하나는 쿨백-라이블러 발산(KL divergence)입니다. 재구성 손실은 입력 x와 재구성 x'와 차이이며, KL 발산은 잠재 공간의 분포와 가우시안 부포의 차이를 의미합니다. VAE에서는 두가지를 합쳐서 총 손실(재구성 손실 + KL 발산)을 계산합니다. [1]

 

Auto-Encoding Variational Bayes

적용 방법

주어진 훈련 데이터셋을 미니 배치로 나눠 앞서 구성한 인코더와 디코더에 학습시키고, 역전파를 통해 네트워크의 가중치를 업데이트 합니다.

 

이렇게 학습된 인코더와 디코더는 데이터의 추론 및 생성에 활용됩니다. 먼저, 인코더에 입력된 값들은 평균과 분산을 활용해 reparametrization을 통해 z을 샘플링해 잠재 공간을 생성하게 됩니다. 그리고 샘플링된 데이터를 디코더에 넣어 새로운 데이터를 생성하게 됩니다.

 

아래 예제 코드를 통해 적용 방법에 대해 이해할 수 있습니다. 

GitHub - Jackson-Kang/Pytorch-VAE-tutorial: A simple tutorial of Variational AutoEncoders with Pytorch

DDPM

Denoising Diffusion Probablistic Model(DDPM)은 최근에 많이 사용되는 이미지 생성 애플리케이션의 근간이 되는 디퓨전 모델을 처음 제시한 모델입니다. 디퓨전 모델은 데이터에 점진적으로 노이즈를 추가하는 과정(forward process)과 노이즈를 제거하고 데이터를 복원하는 과정(reverse process)을 학습합니다.

 

DDPM 구조[2]

Forward process

최종 단계(T)까지 나눈 단계(time step)마다 노이즈를 추가하면서 최종에서는 완전한 가우시안 노이즈에 가깝게 변환합니다. 가우시안 노이즈를 만드는데 필요한 확률분포는 아래와 같이 모델링하고 있습니다. (여기서 beta 값은 하이퍼파라미터입니다.)

 

노이즈가 추가된 데이터와 타임스텝을 신경망에 입력해 신경망에서 노이즈 예측값을 계산하고, 예측 노이즈와 실제 노이즈 사이의 차이를 계산합니다. 이를 손실함수로 훈련을 진행합니다. 

 

Reverse process

가우시안 노이즈인 T에서 점진적으로 노이즈를 제거하면서 원본 데이터 분포의 샘플을 생성하는 것을 목표로 합니다. 현재 상태와 타임스텝을 신경망에 입력해 계산된 노이즈 예측값을 제거하는 과정을 거칩니다. 

 

간단한 것처럼 설명했지만, 실제로 보면 꽤 복잡한 구조를 갖고 있습니다. 구현된 코드를 살펴보면서 이해가 필요하니 참고용으로 링크를 남깁니다.

ML-Paper-Review/ComputerVision/Diffusion/DDPM/DDPM.ipynb at main · Glanceyes/ML-Paper-Review

Latent Diffusion (Stable Diffusion)

Latent Diffsuion(LDM) 또는 Satble Diffusion이라 불리는 기술은 단연컨데, 최근에 생성형 이미지에 가장 많이 사용되는 기술일 것입니다. 그만큼 일반적으로 활용이 가능하고 결과물의 퀄리티가 높은 편이라 꼭 알아두어야 하는 기술이라고 생각합니다. LDM은 이미 학습된 오토 인코더를 활용해 latent space에 있는 DDPM을 학습하는 형태로 모델의 훈련이 진행됩니다.

LDM 구조

LDM의 구조는 아래 그림과 같습니다. 크게 구분할 수 있습니다.

 

1) 픽셀 공간(Pixel Space) : 입력 및 출력을 담당하는 인코더와 디코더를 포함한 오토인코더 존재

2) 잠재 공간(Latent Space) : Diffusion 프로세스와 Denoising하는 Unet 구조를 포함해 DDPM 학습 수행

3) 조건부 메커니즘(Conditioning Mechanism) : 다른 입력 값(텍스트 등) 처리

 

LDM 구조 [3]

 

각 구조별 손실함수는 다음과 같습니다. 

픽셀 공간 손실함수 [3]

잠재 공간 손실함수 [3]

조건부 메커니즘 손실함수 [3]

적용 방법

앞서 VAE와 DDPM을 이해했다면, 적용 방법 자체는 비교적 간단합니다. 왜냐하면 DDPM과 가장 큰 차이점이 DDPM은 확산 및 학습 과정이 픽셀 공간에서 이뤄졌다면, LDM은 픽셀 공간에서는 이미 학습된 오토 인코더를 가져와서 인코딩/디코딩을 수행하고 잠재 공간에서 확산 및 학습 과정이 이뤄지기 때문입니다. 즉, 여러 기능들을 모듈화하여 조합하는 형태라고 이해할 수 있을 것 같습니다. 참고할만한 github 파이썬 파일 링크를 아래와 같이 기록합니다. 

 

latent-diffusion/ldm/models/autoencoder.py at a506df5756472e2ebaf9078affdde2c4f1502cd4 · CompVis/latent-diffusion

 

latent-diffusion/ldm/models/diffusion/ddpm.py at main · CompVis/latent-diffusion

응용 사례

Stable Diffusion의 추가적인 활용 사례에 대한 연구도 지속적으로 이뤄지는 것으로 보입니다. 기존 모델에서 제어 가능한 입력을 추가해 이미지 생성 과정을 정밀하게 조절할 수 있도록 제어하는 ControlNet[4], 사전 학습된 모델에 추가적인 작은 모듈만 더해 모델을 조정하는 LoRA[5]와 같은 모델들이 있습니다.

참고자료

[1] Diederik P Kingma, Max Welling. "Auto-Encoding Variational Bayes". https://arxiv.org/abs/1312.6114

[2] Jonathan Ho, Ajay Jain, Pieter Abbeel. "Denoising Diffusion Probablistic Models". https://arxiv.org/abs/2006.11239

[3] Rombach et al. "High-Resolution Image Synthesis with Latent Diffusion Models". https://arxiv.org/abs/2112.10752

[4] Lvmin Zhang, Anyi Rao, Maneesh Agrawala, "Adding Control Control to Text-to-Image Diffusion Models". https://arxiv.org/abs/2302.05543

[4] Hu et al., "LoRA : Low-Rank Adaptation of Large Language Models". https://arxiv.org/abs/2106.09685