[논문 리뷰] stable diffusion paper 뜯어보기

오늘의 논문리뷰는 21년도에 발표된 stable diffusion, 'High-Resolution Image Synthesis with Latent Diffusion Models' 이다.

stable diffusion은 이미지 생성 모델로 널리 쓰이는 모델이다. 오랜만에 논문 리뷰를 해보자 :) 

 

---

 

0. 사전 지식

ELBO

diffusion계열에서 자주 나오는 ELBO의 개념은 내블로그의 다른 포스트에서 자세히 다룬 적 있다. 가볍게 복습을 해보자. 

 

log likelihood를 maximize하는 것이 모델 학습에서의 핵심이다. 그러나 이 함수를 maximize하자니 식이 복잡해져, 이 함수의 '최대치'로써 근사될 수 있는 바운더리를 찾아, 대신해서 maximum log likelihood의 대상으로 삼는다. 이를 bound optimization이라고 한다. 

 

bound optimization은 최대한 작아야한다. 그래야 log likelihood의 maximize 값과 차이가 근소할 것이기 때문이다. 그러면 우린 bound optimization의 'lower bound'를 찾아야 한다. 이를 ELBO라고 한다. 

 

이제 log likelihood를 maximize 하는 것이 아닌, ELBO를 maximize하는 문제로 바뀌었다. 

 

 

Diffusion Model

이미지 생성에 주로 사용되는 diffusion계열 모델은 noise를 학습하는 모델이라 볼 수 있다. 이미지에서 (보통 gaussian distribution을 따르는) noise를 마구마구 추가한 다음에, 이 노이즈 낀 이미지를 다시 복원하면서 학습을 진행한다. 

 

그래서 실제 infer로 사용할 땐 노이즈에서 시작해서 denoising을 하는 형태로 이미지를 생성한다. 아래 식의 epsilon은 noise, epsilon_theta는 t시점에서 학습한 노이즈를 의미한다. 

diffusion objective

 

이제 논문 리뷰를 시작해보자! 

 

 

1. Introduction

Diffusion Model(DM)은 이미지 생성을 잘하지만 pixel space에서 직접 작동한다는 한계점이 있다. 또하나의 한계점은 연산량과 메모리가 너무 많이 든다. 5만장 샘플 생성하는 데에 A100으로 5일이 걸린다. (25년도 입장에서 생각해보면 엄청나게 커보이진 않긴 하다ㅎ) 

 

그래서 본 논문에서는 DM의 성능을 저해하지 않으면서도 complexity를 줄이고자 했다. 그래서 탄생한게 Latent Diffusion Model, 현재 stable diffusion이라 불리는 모델이다. 

 

LDM은 크게 2단계(two stage) 로 분리해서 볼 수 있다. 

• Perceptual Compression Stage
  • Auto Encoder(AE)를 사용해서 이미지 space를 latent space로 매핑한다. 

• Generative Modeling Stage
  • Diffusion Model을 픽셀 space가 아닌 위에서 뽑아낸 latent space에서 학습한다. 기존 diffusion의 학습에서 prior가 주요 역할을 했던 것과의 차별점이 되겠다. 모델이 데이터 자체의 semantic한 부분을 조금 더 학습할 수 있는  컨셉이다.

여기에 cross attention layer를 도입해서 텍스트 생성 등에 조금 더 유연하게 되었다. 

 

메인 컨트리뷰션은 다음과 같다. 

1. 확장성이 좀 더 좋아졌다. DALL-E, VQGAN에 비해서 이미지의 고해상도 합성에도 보다 효율적이다.
   1. 참고로 DALL-E, VQGAN 등은 Auto Regressive하게 학습하는데 이게 discrete하다. 
2. pixel based DM에 비해 연산량이 줄었다. 
3. latent space의 regularization이 필요하지 않게 되었다. 즉, prior가 학습에 거의 관여하지 않는 듯 하다. 
4. cross-attention을 도입해서 텍스트와 이미지 쌍의 학습이 가능해졌다.  

 

2. Related Works 

 

이 섹션에서는 친절하게도 GAN과 DM 모델에 대한 설명이 잘되어있다.  먼저 GAN부터 살펴보자. 

 

GAN

GAN의 단점은 mode collapse가 있을 수 있단 거다. (여기서의 mode는 클래스로 생각해도 무방해보인다.) 익히 알다시피 GAN은 discriminator와 generator로 구성되어 있다. 여러가지 클래스의 이미지로 학습했다 하더라도, 특정 클래스로 생성 편향이 생길 수 있다. 어떻게 보면 discriminator가 잘 모를법한 이미지로 계속해서 생성하도록 잘못 학습되는 거라 볼 수 있다. 

 

discriminator는 minimize, generator는 maximize되어야 하기 때문에 이 둘의 균형점을 찾으며 학습하는 것에도 난이도가 있다. 

 

 

Likelihood based Methods

VAE나 auto regressive 모델이 여기에 속한다. density estimation 기반이라 데이터를 밀도 기반으로 잘 추정하는 특성이 있다. 효율적으로 합성은 가능하지만, GAN보다는 퀄리티가 떨어진다. 

 

 

Diffusion Probalistic Models 

디퓨전 모델들은 Likelihood based의 일종이다. 보통 U net 아키텍처를 사용한다. 거기다가 noise를 제거하는 AE layer로 인해서 inductive bias(한국어로 귀납적 편향)에 강한 편이다. 게다가 이미지 퀄리티를 GAN 수준으로 끌어올렸다. 

 

다만 위의 Likelihood based Methods와 같이, pixel기반으로 동작한다. 즉, 픽셀별로 학습하기 때문에 연산량이 급격히 늘어나게 된다. 

 

 

 

3. Methods

논문에서 제시한 LDM이 기존 diffusion model에서 어떻게 발전했는지를 기술한 섹션이다. 인트로에서 언급했던 것처럼 two stage로 나뉘어진다. 이미지가 참 직관적이다. 

 

출처: stable diffusion paper fig 3

 

1. Perceptual Image Compression 

위 figure3 에서의 빨간색 부분이다. encode에선 이미지 x가 latent space로, decoder에선 latent space에서의 계산 결과를 디코딩하는 것을 확인할 수 있다. 

 

이 스텝의 목적은 diffusion에서의 불필요한 computation을 줄이는 것이다. 어디에서 과한 연산이 발생했는지 확인해보니, (위에서 계속 언급됐다시피) RGB 고차원 채널에서의 pixel기반 function evaluation이 문제였다. 그래서 아예 latent 공간으로 보내버렸다. 그래서 불필요하게 학습되던 픽셀들은 아예 학습 과정에서 빠져버리게 된다. 

 

latent space로의 매핑은 Auto Encoder(AE)를 사용한다. 기존 DM 모델들이 모든 픽셀을 대상으로 학습하기 때문에 compression rate가 빡셌던(?) 것을 이걸로 인해 완화할 수 있다. 

 

figure 1에서 VQVAE의 압축률 f가 16이었던 것을 보면 직관적으로 와닿을 수 있겠다. 

 

 

latent space의 high variance를 막기 위해 KL divergence를 추가한다. VAE에서 KL reg term을 주는 것과 유사하다. 

 

2. Latent Diffusion Models

먼저 Diffusion step. Diffusion 모델이 사용하는 개념인 Markov Chain이 여기서도 사용된다. 이는 denoising autoencoder로도 해석될 수 있다! 총 T번의 denoising step이 있다고 가정하면 다음과 같은 수식을 따르게 된다. (디퓨전은 학습된 noise epsilon에 따라 생성하는 모델이다.) 

 

latent한 디멘젼에서 학습되기 때문에 (Intro에서 나왔던 것처럼) 상대적으로 중요한 semantic bit에 더 집중할 수 있게 됐다. 거기다 기본 DM보다는 low dimension에서 학습되기 때문에 cost efficiency까지. 추가로 LDM에선 backbone으로 U-Net을 사용하는데 Conv layer로 이루어져 있다. 

 

LDM 수식) latent diffusion objective

 

 

3. conditioning mechanism

stable diffusion은 어떻게 추가로 text 정보를 학습에 이용할 수 있을까? denoised된 결과 θ(zt, t, y) 에서 y는 텍스트일 수도 있고 다른 signal일수도 있을 것이다. 이 y가 text라면 BERT encoder를 태우는 cross attention을 수행하는 컨셉이다. 아래 수식의 τθ 가 specific domain encoder이다. 

 

U-Net에서의 중간중간 결과를 flatten한 것인 z_t를 Query, cross domain y를 Key, Value로 받는다. 

 

 

그래서 LDM 수식은 이렇게 진화한다.

LDM 수식(진)

 

 

 

4. Experiment

(어떻게 보면 가장 중요한) downsampling 인자인 f에 대한 실험을 먼저 보자. f가 너무 크면(ex. f = 32) 압축률이 너무 커져 정보 손실이 생긴다. 반면 f가 너무 작으면 (f= 1,2) training 속도가 느리다. OpenImages 데이터셋으로 학습했을 때 최적은 4인듯 하다. (참고: FID는 원본이미지와 생성이미지간 얼마나 유사한지 distance를 계산한 결과다.) 

 

다른 실험들에서도 k = 4 또는 8을 권장하고 있다. 

 

 

---

 

"prior을 계속해서 학습에 사용하는 Diffusion 모델이 효율성을 위해 앞단에서 아예 latent space로 보낸다" 가 stable diffusion의 핵심으로 보인다:)