[평범한 학부생이 하는 논문 리뷰] High-Resolution Image Synthesis with Latent Diffusion Models (Stable Diffusion)

부스트캠프 내에서 내가 진행하고 있는 diffusion 스터디에서는 아직 SDE diffusion 논문을 읽고 있고, DDIM까지 읽고 나서 Stable Diffusion을 읽으려 했지만, 최종 프로젝트 때문에 미리 읽게 되었다. 다행히 NCSN, DDPM, SDE diffusion처럼 수식이 엄청 많고 그런 논문이 아니라서 이 세 논문처럼 힘들지는 않았던 것 같다! 다들 stable diffusion에 관심이 많고 부스트캠프에서도 생성모델을 통한 프로젝트를 여러 팀에서 하는 것 같아서 부스트캠프 내에서 진행되는 논문 세미나에서도 이 논문을 발표하기로 했다. 그렇기에 최대한 꼼꼼하게 공부하고 리뷰를 하려고 한다.

https://arxiv.org/abs/2112.10752

High-Resolution Image Synthesis with Latent Diffusion Models

1. Introduction

Stable diffusion논문이 나오기 이전에 diffusion model(DM)들은 pixel 공간에서 이미지를 생성했다. 그렇기에 pixel 하나하나의 gradient를 계산하고 그 gradient를 가지고 noise에서 이미지를 생성했다. 그래서 엄청난 리소스를 필요로 했고, 큰 이미지를 생성하기에 부적합했다. Inference하는 데에도 많은 시간이 걸린다.

본 논문에서는 이러한 문제점을 극복하기 위해 stable diffusion을 제시한다. Stable Diffusion은 어떻게 이러한 문제점을 극복했을까?

 

위의 이미지를 보면 semantic compression과 perceptual compression이 나와있다.

• Semantic Compression : 생성 모델이 데이터의 의미적 및 개념적 구성을 학습하는 단계
• Perceptual Compression : high-frequency 디테일을 제거하지만 의미적 변화는 거의 학습하지 않는 단계

본 논문에서는 이 두 단계를 나눠서 진행하는 방식을 제안한다. 첫 단계에서는 autoencoder를 훈련하여 동일한 데이터 공간과 perceptual하게 동등한 저차원 representational 공간을 제공하고 그 다음 단계에서 그 공간에 diffusion model을 학습한다. 이 방식을 이용하면 universal autoencoding 단계를 한번만 훈련하면 된다고 한다. 그래서 다른 작업에서도 다시 사용이 가능하다는 장점이 있다.

2. Method

본 논문은 generative learning 단계에서 compression을 명시적으로 분리하여 이러한 단점을 피할 것을 제안한다. 이를 위해 이미지 공간과 지각적으로 동등한 공간을 학습하지만 계산 복잡성이 크게 줄어든 autoencoding 모델을 활용한다.

이를 통해 아래와 같은 advantage를 얻는다고 한다.

1. 저차원 공간에서 샘플링을 수행하기 때문에 계산이 훨씬 효율적이다.
2. UNet 아키텍쳐를 통해 얻은 inductive bias를 통해 공간 구조를 가진 데이터에 효과적이며, 품질 저하를 일으키는 수준의 압축의 필요성을 완화시킨다.
3. 여러 generative 모델을 훈련하는 데 latent space를 사용할 수 있고 다른 downstream application에도 활용할 수 있는 general-purpose compression models을 얻는다.

2.1 Perceptual Image Compression

본 논문에서 소개하는 perceptual compression model은 perceptual loss와 patch-based adversarial loss의 조합으로 훈련된 autoencoder로 구성된다. 이를 통해 국소적 사실성을 강화하여 reconstruction이 이미지 manifold에 국한되도록 하고 $L_2$ 또는 $L_1$ objective와 같은 pixel-space loss에만 의존하여 발생하는 흐릿함을 방지할 수 있습니다. 

---

Perceptual Loss 란?

VGG의 hidden layer를 통해 두 이미지의 high-level feature를 추출해서 그 feature의 오차($L_2$ norm)에 대한 Loss.

 

Patch-based adversarial Loss란?

Pix2Pix논문에서 나왔듯이 이미지 전체를 비교하는 것이 아닌 patch 단위로 이미지가 실제인지 아닌지를 판별하는 Loss

---

2.2 Latent Diffusion Model

Diffuion Models

기존의 diffusion model은 이미지에 점차적으로 gaussian noise를 추가해서 tractable한 prior distribution(표준 정규 분포)으로 만들고 desnoising process를 거쳐 다시 복원하는 과정을 학습해서 이미지를 생성하는 모델이다. 이는 동일하게 가중치가 부여된 denoising autoencoder $θ(x_t, t); \ t = 1, \cdots ,T$로, 입력 $x_t$의 denoising variant를 예측하도록 훈련된 것으로 해석할 수 있다. 이 때의 objective는 아래와 같다.

$$L_{DM} = E_{x,\epsilon∼\mathcal{N}(0,1),t}[ \Vert\epsilon - \epsilon_θ(x_t, t)\Vert^2_2] , \quad (1)$$

 

Generative Modeling of Latent Representations

앞서 autoencoder를 이용해서 저차원의 latent space에 접근할 수 있게 됩니다. 이렇게 encoder $\mathcal{E}$를 통해 먼저 input $x$를 latent space의 latent vector $z$로 encoding을 시키고 이 $z$를 이용해서 diffusion process를 진행시킨다고 보면될 것 같다. 이렇게 되면 (i) 데이터 중요하고 의미 있는 부분에 집중할 수 있고, (ii) 저차원의 계산으로 훨씬 효율적으로 훈련이 가능하다고 한다.

Latent diffusion models의 objective는 아래와 같다.

$$L_{LDM} := \mathbb{E}_{\mathcal{E}(x),\epsilon∼\mathcal{N}(0,1),t}[\Vert \epsilon - \epsilon_θ(z_t, t)\Vert^ 2_2]. \quad (2)$$

이전 방식들과 다르게 논문에서 소개하는 LDM은 이미지별 inductive biases를 활용할 수도 있다고 합니다. 이는 아마도 2D convolutional layer를 활용해서 UNet을 구축하기 때문일 것 같아요.

2.3 Conditioning Mechanisms

Stable Diffuion의 Architecture

본 논문에서는 또한 cross-attention 매커니즘을 이용한 conditional generation을 제안한다. 그렇다면 cross-attention을 어떻게 이용했을까?

위 이미지를 보면 먼저 도메인별 인코더 $τ_θ$를 도입해서 condition $y$가 주어지면 intermediate representation인 $τ_θ(y) ∈ \mathbb{R}^{M×d_τ}$ 로 projection시킨다. 이 인코더의 output은 cross-attention을 통해 UNet의 중간 layer에 매핑된다.

$${\rm Attention}(Q,K,V) = {\rm softmax}\left( \frac{QK^T}{\sqrt{d}} \right) \cdot V$$

$$Q = W^{(i)}_Q \cdot \varphi_i(z_t), \ K = W^{(i)}_K \cdot \tau_{\theta}(y), \ V = W^{(i)}_V \cdot \tau_{\theta}(y)$$

여기서 나오는 $W$들은 학습 가능한 projection matrix들이다. 이를 바탕으로 conditional LDM은 아래의 objective를 바탕으로 $\tau_{\theta}$와 $\epsilon_{\theta}$를 jointly하게 최적화한다.

$$L_{LDM} := \mathbb{E}_{\mathcal{E}(x),y,\epsilon ∼\mathcal{N}(0,1),t}[ \Vert \epsilon - \epsilon_θ(z_t, t, τ_θ(y))\Vert^2_2 ] , \quad (3)$$

3. Experiments

실험들은 pixel space에서 생성하는 모델들과 비교한다.

3.1 On Perceptual Compression Tradeoffs

이 섹션에서는 다운샘플링 계수 $f \in \{1,2,4,8,16,32 \}$에 따른 LDM의 동작을 분석한다. 여기서 다운샘플링 계수 $f$는 1단계에서 autoencoder를 통해 input image를 얼마나 줄이는 지를 나타낸 계수라고 보면 된다. ($f = H/h = W/w$)

train step에 따른 sample quality를 보여주는 이미지

1. LDM-{1,2} (다운샘플링 계수가 작은 경우) : perceptual compression을 대부분 diffusion 모델이 맡아서 훈련 진행 속도가 느리다.
2. LDM-32 (다운샘플링 계수가 너무 큰 경우) : perceptual compression이 지나치게 돼서 정보 손실이 잃어나서 적은 train step이후, fidelity가 정체된다.
3. LDM-{4-16} : 균형이 잘 맞아서 픽셀 기반 diffusion과 상당한 FID 차이를 보인다.

샘플링 속도를 보여주는 image

• LDM-{4-8} : 픽셀 기반 LDM-1에 비해 훨씬 낮은 FID 점수를 달성하는 동시에 샘플 처리량을 크게 증가시켰다.

3.2 Image Generation with Latent Diffusion

Sample quality와  FID와 Precision-and-recall을 이용해 data manifold에 대한 coverage를 평가한다.

결과 요약 table

• GAN과 latent diffusion모델과 첫 번째 단계(autoencoder)를 함께 학습하는 LSGM보다 성능이 뛰어난 것을 확인할 수 있다.
• 더 적은 리소스로 다른 diffusion model이상의 성능을 냈다.
• GAN보다 향상된 Precision-and-Recall을 보이며 training objective에 대한 이점을 보였다.

3.3 Conditional Latent Diffusion

3.3.1 Transformer Encoders for LDMs

LDM에 cross-attention based conditioning을 도입함으로써 다양한 conditioning 양식을 사용할 수 있게 됐다. 

$\tau_{\theta}$를 multi-head cross-attention통해 UNet에 매핑되는 latent code를 추론하는 transformer로 구현한다. 이를 통해 복잡한 사용자 정의 text prompt에 잘 일반화되는 강력한 모델을 생성한다. 아래의 이미지들은 이러한 모델로 생성한 결과를 보여준다.

(좌) text2image (우) layout-to-image

3.3.2 Convolutional Sampling Beyond $256^2$

Spatially aligned conditioning 정보를 $\epsilon_{\theta}$의 input에 concat함으로써 위와 같은 합성이 가능해진다.

위는 semantic map과 그와 쌍을 이루는 풍경이미지를 concat해서 이미지를 생성한 예시를 보여준다.

3.4 Super-Resolution with Latent Diffusion

LDM은 concat을 통해 저해상도 이미지를 직접 conditioning함으로써 super-resolution을 효율적으로 training할 수 있다.

위의 결과 뿐만 아니라 pixel 기반 diffusion model들 보다 빠른 샘플링으로 더 좋은 결과를 보였다고 한다.

3.5 Inpainting with Latent Diffusion

Inpainting의 결과는 아래와 같다.

많은 부분을 가렸지만 inpainting이 잘 되는 것을 확인할 수 있다.

4. Limitations

• 픽셀 기반의 diffusion model들 보다는 확실히 빠르지만 여전히 GAN에 비해 느리다.
• Super-resolution과 같은 높은 정밀도가 필요한 작업에서는 여전히 한계를 드러낸다.

5. Conclusion

픽셀 space가 아닌 latent space에서 diffuion model을 적용시킨다는 아이디어가 매우 참신했고 결과 또한 놀라웠다. 게다가 cross attention을 이용한 conditional generation도 매우 신기했다. 앞으로 읽어나갈 stable diffusion을 이용한 다양한 논문들이 궁금해졌다.