[22’ CVPR] Stable Diffusion: High-Resolution Image Synthesis with Latent Diffusion Models
카테고리: Generation
태그: Diffusion
🔍 Abstract
- Pixel space에서의 diffusion 과정은 계산량이 많다. 따라서 이를 latent space로 옮겨서 계산량을 줄이는 LDM(Latent Diffusion Model)을 제안한다.
- LDM은 complexity는 줄이면서 detail은 유지할 수 있다는 장점이 있다.
- Cross-attention layer를 통해 이론적으로는 어떠한 conditioning 정보도 사용할 수 있다.
1. Introduction
DDPM에서 언급된 내용으로, diffusion model은 imperceptible detail을 계산하는 데에 많은 시간을 소요한다는 것이 알려져 있다. 아래 그림은 DDPM에서도 분석한 그림으로, DDPM의 과정은 실제로 데이터의 semantic composition을 학습하는 semantic compression 부분과 사람이 인지하지 못할 정도의 high-frequency detail을 학습하는 perceptual compression 부분으로 나뉘어 해석할 수 있다.
따라서 LDM(Latent Diffusion Model)에서는 이 perceptual compression 과정을 diffusion model 대신 autoencoder로 대체하여 계산량을 줄이고자 한다. 즉, 실제 pixel space에서의 data와 perceptually equivalent한 latent space에서의 data를 autoencoder로 생성한 뒤, 이를 가지고 diffusion model을 학습 및 추론하면 더 효율적일 것이라는 주장이다. 이러한 방법을 통해 저자들은 고차원의 데이터를 적은 계산량으로 생성할 수 있었으며, latent space에서의 diffusion model은 훨씬 안정적으로 동작했다고 한다. 중요한 정보만을 사용하여 diffusion model을 학습했으니 더 robust하게 학습할 수 있었던 것이다.
2. Architecture
전체적인 구조는 다음과 같다. (아마 diffusion 계열에서 가장 유명한 Figure 중 하나가 아닐까 싶다.) Encoder $\mathcal{E}$는 pixel space에서 latent space로의 mapping을 담당하며, Decoder $\mathcal{D}$는 latent space에서 pixel space로의 mapping을 담당한다. Diffusion model $\epsilon_\theta$은 latent space에서의 data를 통해 계산되며, cross-attention layer를 통해 conditioning 정보를 사용할 수 있다.
2.1. Encoder & Decoder
이제 각 component를 살펴보자. 먼저 encoder $\mathcal{E}$는 pixel image $x \in \mathbb{R}^ {H \times W \times 3}$를 latent representation $z \in \mathbb{R}^ {h \times w \times c}$로 변환한다. 즉 $z = \mathcal{E}(x)$이다. 한편 decoder $\mathcal{D}$는 latent representation $z$를 pixel image $x$로 변환한다. 즉 $\tilde{x} = \mathcal{D}(z) = \mathcal{D}(\mathcal{E}(x))$이다.
이때 encoder는 다른 latent representation과는 달리 2D latent representation을 사용하여 spatial information을 보존하였고, 따라서 마치 downsampling과 유사하게 해석할 수 있다. 여기서 downsampling factor $f = H/h = W/w$이다.
Latent space가 high-variance를 가지는 것을 막기 위해 regularization을 사용한다. 두 가지 방법을 소개하는데, 첫 번째는 VAE에서 사용하는 KL-reg이다. 즉 standard normal과 latent distribution 사이의 KL-divergence를 최소화하도록 하여 latent space의 모양을 standard normal에 가깝게 만든다. 두 번째는 VQ-GAN에서 사용하는 VQ-reg로, decoder에 vector quantization layer를 사용한다.
2.2. Latent Difffusion Models
LDM(Latent Diffusion Model)은 간단히 latent space에서의 DDPM이라고 생각할 수 있다. Encoder도 2D latent representation을 생성하고, DDPM도 U-Net backbone을 사용하여 inductive bias를 최대한 유지하였다. Objective는 DDPM과 같다.
\[L_{\text{LDM}} = \mathbb{E}_ {\mathcal{E}(x), \epsilon \sim \mathcal{N}(0, 1), t} \left[ \Vert \epsilon - \epsilon_ \theta (z_t, t) \Vert ^ 2 _ 2 \right]\]이때 text, semantic map, image 등의 condition은 domain-specific encoder $\tau_ \theta$를 통해 latent space에 embedding된다. 이후 cross-attention layer를 통해 diffusion model에 주입된다. 혹은 condition의 종류에 따라(e.g. image, layout) 단순 concatenation을 통해 주입될 수도 있다. 이때 $\tau_ \theta$도 $\epsilon_ \theta$와 함께 학습된다.
\[L_{\text{LDM}} = \mathbb{E}_ {\mathcal{E}(x), \epsilon \sim \mathcal{N}(0, 1), t} \left[ \Vert \epsilon - \epsilon_ \theta (z_t, t, \tau_ \theta (y)) \Vert ^ 2 _ 2 \right]\]3. Experiments
3.1. Perceptual Compression Tradeoff
저자들은 downsampling factor $f$에 따라 모델 이름을 LDM-$f$로 표기한다. $f$가 너무 작은 경우에는 일반적인 diffusion model과 크게 다르지 않고 계산량이 많은 반면 $f$가 너무 크면 정보의 손실이 너무 커서 성능이 떨어진다. 따라서 적절한 $f$를 찾는 것이 중요하다. 그 결과 적절한 efficiency와 perceptual quality를 가진 LDM-4, 8 등을 찾을 수 있었다. 예상할 수 있듯이 데이터셋의 크기가 크고 복잡할수록 $f$가 작은 경우가 더 좋은 성능을 보였다.
3.2. Image Generation with Latent Diffusion
LDM 모델은 unconditional image synthesis, text-conditional image synthesis, layout-to-image synthesis, super-resolution, image inpainting 등 정말 다양한 task에서 준수한 성능을 보여준다.
Text-conditional image synthesis의 경우 BERT-tokenizer를 사용하고, transformer 구조의 encoder $\tau_ \theta$를 사용하여 text를 latent space에 embedding하였다. 이를 확장하여 semantic layout을 바탕으로 layout-to-image synthesis도 수행할 수 있다.
Super-resolution의 경우 단순히 low-resolution conditioning을 concatenate하는 SR3 논문의 방법을 그대로 사용하였다. 결과적으로 SR3보다 FID 면에서는 앞섰지만 IS 면에서는 뒤쳐졌다. 더 현실적인 이미지를 생성해내는 경향이 있지만 이미지의 디테일은 떨어진다는 것이다. 저자들은 LDM에서 pixel space로의 reconstruction capability가 bottleneck이 되어 성능이 떨어진다며 LDM의 한계로 지적했다.
이외에도 convolutional sampling을 사용하여 larger image를 생성할 수도 있고, image inpainting도 가능하다.
💡 Summary
- LDM(Latent Diffusion Model)을 통해 latent space에서의 diffusion을 구현하였으며, 이를 통해 계산량을 줄이고 고차원의 데이터를 더 안정적으로 다룰 수 있게 되었다.
- LDM은 spatial information을 보존하는 encoding 방식을 사용하였고, 이 덕분에 일반적인 encoding보다 pixel space의 정보를 더 잘 표현할 수 있었으며 다양한 방면에서 쉽게 확장이 가능했다.
댓글 남기기