이번 글에서는 Stable Diffusion 논문리뷰를 진행하겠습니다 :)
※ 해당 논문을 이해하기 위해서는 Diffusion Model에 대한 배경지식이 필요합니다.
✅ Key Idea
1. 기존 픽셀공간에서의 Diffusion Model 연산을 Latent Space로 매핑하여 기존 모델에 요구되는 자원보다 훨씬 적은 자원만으로 고퀄리티의 이미지 생성가능
2. 모델이 매우 light: 기존 모델과는 다르게 8GM VRAM만으로 이미지 생성가능
Introduction
기존 Diffusion Model(GLIDE, Imagen, DALLE2 등 ...)들은 이미지 생성분야에서 좋은 성능을 보여왔습니다. 하지만 이러한 기존 Diffusion 모델들은 직접적으로 픽셀 공간에서 동작하기 때문에 데이터의 감지할 수 없는 세부정보를 모델링하는데 너무 많은 컴퓨팅 자원을 소비하게 되는 한계를 가지고 있습니다.
위의 그림은 학습단계에 따른 Rate-Distortion의 trade-off를 나타내는 그림입니다.
각 학습단계를 살펴보면 Perceptual Compression 단계는 이미지의 시각적인 세부사항을 줄이면서 의미(semantic)는 거의 학습하지 않는 단계이고, Semantic Compression은 이미지의 의미론적(semantic) 구성과 개념적(conceptual) 구성을 학습하는 단계입니다.
Bit Rate가 작아질수록 사람의 눈에는 인지가 잘 되지 않는 부분인데, GAN과 Autoencoder는 Bit Rate가 비교적 큰 Perceptual한 부분을 학습하는데 주력하고 있습니다. 반면에 Diffusion Model은 Bit Rate가 작은 Non-perceptual한 부분을 학습하는데 초점을 맞추고 있는 것을 확인할 수 있습니다. 즉, Diffusion Model은 사람 눈에 잘 인지되지 않는 부분에 중점을 두고 학습을 진행하면서 과도한 자원을 소비하고 있는 것입니다.
따라서 저자들은 이러한 문제를 해결하기 위한 방법으로 원본 이미지 데이터와 perceptual하게 동등하지만 낮은 차원의 잠재공간에서 Diffusion Model의 학습을 진행하는 Latent Diffusion Model을 제안합니다. 우선, 원본 이미지 데이터와 perceptual하게 동일한 저차원 표현 공간을 찾기 위해 Autoencoder를 학습시키고, 그 후 학습된 잠재공간(latent space)에서 Diffusion Model을 학습을 진행합니다.
이러한 방법을 사용함에 따라 컴퓨팅 cost를 크게 낮추면서 이미지생성의 다양한 task에서 SOTA를 달성하는 성과를 얻을 수 있었습니다. 또한 더 높은 차원의 데이터들도 자연스러운 이미지 합성이 가능합니다. 또한 이러한 방식의 추가적인 장점은 한 번 학습한 Autoencoding 단계를 재사용할 수 있다는 점입니다. 현재 이 모델은 오픈소스로 공개되어 있으며 GitHub에서 찾아볼 수 있습니다.
Method
Stable Diffusion에 대한 전체적인 아키텍처에 대해 먼저 살펴보도록 하겠습니다.
Stable Diffusion의 전체적인 아키텍처는 다음과 같이 VAE와 Diffusion Model을 합쳐놓은듯한 구조로 이루어져 있습니다.
위의 그림과 같이 인코더를 이용해 데이터를 저차원 latent space에 매핑한 후에 diffusion process를 latent space에서 진행하고 디코더를 이용해서 원본이미지의 크기로 재구성하게 됩니다. 이 과정을 자세하게 나타낸 그림은 아래와 같습니다.
복잡해보이지만 크게는 1. Encoder, 2. Diffusion Process, 3. Decoder, (+ 4. Conditioning)으로 구성되어 있는 것을 확인할 수 있습니다. 보다 자세한 내부구조에 대해서 하나씩 살펴보도록 하겠습니다!
1. Perceptual Image Compression
원본 이미지 $x$를 지각적으로 동등한 latent space로 매핑하는 Autoencoder를 학습시키는 과정을 Perceptual Image Compression(지각적 이미지 압축)이라고 표현합니다.
Autoencoder는 1. Perceptual loss와 2. patch-based Adversarial Objective 두 가지 목표를 조합해서 학습이 이루어지게 됩니다. Perceptual loss는 지각적인 측면을 현실적으로 유지한채로 압축할 수 있도록 하는 역할이고, patch-based Adversarial Objective는 이미지의 로컬한 특징을 보존하도록 하는 역할을 하게 됩니다. 따라서 이러한 목표로 훈련된 Autoencoder는 이미지 매니폴드에 집중하여 원본 이미지를 재구성하게 되고 픽셀 공간 손실로 인한 흐림을 방지할 수 있게 됩니다.
2. Latent Diffusion Models
지금까지 살펴본 내용은 Autoencoder로 구성된 Perceptual Image Compression 모델을 이용해서 픽셀로 표현된 이미지를 Latent 공간에서 imperceptible한 세부정보를 추상화하는 모델링을 진행하는 것입니다.
그래서 기존 Diffusion Model의 Loss func.과 Latent Diffusion Model의 Loss func.을 비교해보면 RGB픽셀 변수($x_t$)가 latent 변수($z_t$)로 바뀐 것을 확인할 수 있습니다.
이렇게 도달한 저차원의 Latent space는 고차원 픽셀공간보다 다음과 같은 두 가지 측면에서 장점을 가지고 있습니다.
1. Autoencoder를 통해 도달한 Latent space에는 마치 PCA(주성분 분석)의 느낌처럼 데이터의 중요한 Semantic 비트만 남아 이 부분에 더 집중할 수 있습니다.
2. 차원이 낮기 때문에 더 효율적인 계산이 가능합니다.
3. Conditioning Mechanisms
다른 유형의 생성모델과 마찬가지로 Diffusion Model은 기본적으로 $p(x|y)$ 형태의 조건부 분포를 모델링 할 수 있습니다.
여기서 조건화(Conditioning)한다는 것은 텍스트 프롬프트에 맞는 이미지를 생성하는 것과 같이 모델이 이미지를 생성할 때 조건(semantic map, text, representation, Images ..)를 추가로 전제화하는 것을 의미합니다.
기존의 이미지 합성 Diffusion Model들은 입력 이미지에 대한 클래스 레이블로 조건화 하거나 이미지를 흐릿하게 변형시키는 연구들은 진행되어왔지만 다른 유형의 컨디셔닝 조건을 결합하는 것은 연구되지 않았습니다. 본 논문에서는 Condition을 text, semantic map ... 등 다른 유형의 컨디셔닝 조건을 이미지 합성프로세스에 결합시켰습니다. 아래의 그림의 가장 오른쪽에서 Conditioning Mechanism 도식을 확인할 수 있습니다.
그렇다면 다양한 형태의 Condition을 입력으로 받기 위해서는 각 Condition별 인코더 $\tau_\epsilon$이 필요합니다. 그리고 각 Condition에 적합한 인코더를 통해 변환된 정보가 Latent Space에서의 Diffusion Process 과정의 $z_t$와 매핑하기 위해서는 Cross-attention 연산을 사용하게 됩니다.
이러한 Conditioning 과정을 수식으로 표현하면 잠재변수(latent variable, $z$), 입력조건(Input condition, $z$)에 대해서 조건화된 분포는 $p(z|y)$라고 나타낼 수 있으며 Conditional Denoising Autoencoder는 $\epsilon_{\theta}(z_t, t, y)$로 구현될 수 있습니다. 또한 다양한 종류의 $y$를 전처리하기 위한 도메인별 인코더를 $\tau_\epsilon$이라고 한다면 이미지와 Conditioning 쌍을 기반으로 하는 Latent Diffusion Model의 Loss func.은 다음과 같이 작성할 수 있습니다.
Experiments
Latent Diffusion Model과 픽셀기반 Diffusion Model의 훈련과 추론의 결과를 비교해보겠습니다.
모델은 downsampling factor $f \in \{1, 2, 4, 8, 16, 32\}$에 따라서 LDM 성능을 비교하였으며 LDM-f 약칭으로 표현됩니다. 이 때 $f=1$은 픽셀기반 Diffusion Model에 해당합니다.
On Perceptual Compression Tradeoffs
LDM-f의 하이퍼파라미터에 따른 재구성 능력
각 step에 따른 샘플품질(FID, IS)
1. 작은 downsampling factor는 학습을 느리게 한다.
2. 지나치게 큰 f는 비교적 적은 step에서 샘플품질의 정체를 유발
3. LDM-4 ~ LDM-16까지는 효율성과 perceptual하게 충실한 결과 사이에서 적절한 균형을 유지한다.
Image Generation with Latent Diffusion
Uncoditional Latent Diffusion Model의 하이퍼파라미터는 다음과 같습니다.
CelebA-HQ, FFHQ, LSUN-Churches, -Bedrooms로 학습한 모델의 FID, Precision, Recall 정량적 결과
정성적 결과
Conditional Latent Diffusion
Conditional LDM의 하이퍼 파라미터
LAION 데이터셋 LDM text-to-image 생성 정성적 결과
MS-COCO 데이터셋 LDM text-to-image 생성 정량적 결과
COCO 데이터셋 LDM text-to-image 생성 정성적 결과
class-conditional LDM 정량적 결과 비교
풍경이미지 semantic synthesis를 위해 256^2 크기로 학습한 LDM이 512 x 1024 크기를 생성하도록 한 결과
Super-Resolution with Latent Diffusion
이미지넷-Val. LDM-SR의 64 → 256 super-resolution 샘플
이미지넷-Val.에 대한 64 → 256 super resolution 정량적 결과
Inpainting with Latent Diffusion
super-resolution, inpainting 정량적 성능 비교 표
impainting model로 object removal을 진행한 결과
impainting 속도 정량적 비교
512 ^ 2크기에서 impainting 성능 정량적 비교
긴 글 읽어주셔서 감사합니다.
잘못된 부분에 대한 피드백과 조언은 언제나 환영입니다 :)
