[논문 리뷰] Stable Diffusion (Latent Diffusion Model)

Abstract

이미지 형성 프로세스를 노이즈를 제거하는 오토인코더 순차적인 어플리캐이션으로 분해하면서 diffusion model은 이미지에서 SOTA 합성 결과를 얻었습니다. 기존에는 픽셀 공간에서 연산을 하여 막대한 GPU 연산이 필요했지만 사전학습된 오토인코더의 latent space를 활용하여 제한된 컴퓨팅 리소스로도 재학습할 수 있게 했습니다. LDM(Latent Diffusion Model)은 pixel 기반의 DM보다 연산량을 상당히 줄이면서, image inpainting, class-conditional image synthesis, 그리고 text-to-image 합성, unconditional image generation, super-resolution과 같은 다양한 분야에서 SOTA를 달성했습니다.

1. Introduction

• 확산(Diffusion)은 시간이 지남에 따라 분자들이 퍼져나가는 현상을 의미합니다. 이런 현상에 착안하여 이미지에 노이즈를 점차적으로 추가하여 최종적으로 노이즈를 만들고 이를 다시 복원하도록 학습하는 것이 Diffusion Model의 핵심입니다.

2. Related Work

Generated Models for Image Synthesis

• GAN 모델은 고해상도의 이미지를 좋은 품질로 잘 생성하지만 모델을 최적화하기 어렵고, 데이터 분포를 잘 잡아내는 것이 힘듭니다.
• 반면 Likelihood 기반의 모델인 VAE(Variational Auto Encoder)나 flow 기반의 모델은 최적화가 잘 이루어지도록 하는 밀도 추정을 잘하지만, GAN보다 성능이 뛰어나지는 않습니다.
• ARM(Auto Regressive Model)은 밀도 추정에서 좋은 성능을 내지만 저해상도 이미지에서만 사용할 수 있다는 제한이 있습니다.

Diffusion Probabilistic Models

• DM은 밀도 추정과 샘플의 품질에서 훌륭한 성능을 보여주며 노이즈를 제거하는 과정에서 U-Net이 활용되어 이미지와 같은 데이터에서 inductive bias를 학습할 수 있습니다.
• DM은 압축 능력과 샘플의 품질이 서로 trade off의 관계를 가지며, pixel 공간에서 모델을 평가하고 최적화하는 것은 추론 속도 저하와 많은 컴퓨팅 비용이 발생합니다.
• LDM은 이러한 단점들을 해결하며 이미지 품질 저하 없이 연산량을 줄이면서 추론 속도를 높일 수 있습니다.

Two-Stage Image Synthesis

• two stage 접근을 통해 단점들을 극복하려는 많은 연구들이 있었습니다.
• VQ-VQE는 이산화된 latent space로 prior를 학습하는 auto regressive model을 사용합니다.
• dall-e는 VQ-VAE에서 더 나아가 이산화된 image와 text의 representation의 결합 분포를 학습하여 text-to-image를 수행합니다.
• net2net은 가역 네트워크를 사용하여 다양한 도메인의 latent spaces들 간에 일반적인 transfer를 제공합니다.
• VQ-VQE와 VQ-GAN은 큰 이미지에 auto regressive transformer를 스케일링하기 위해 first stage로 adversarial과 perceptual objective를 사용합니다.
• 높은 압축률은 ARM이 학습을 가능하게 하지만 성능에 제한이 생기고, 낮은 압축률은 연산 비용이 매우 커지게 됩니다.
• LDM은 좋은 품질의 이미지를 생성하는 것을 보장하면서 first stage에서 perceptual한 압축을 너무 많이 하지 않고 강력한 첫 번째 단계 학습 사이를 최적으로 중재하는 압축 수준을 자유롭게 선택할 수 있다.

3. Method

Fig2는 디지털 이미지에서 대부분의 bits가 사람이 인지할 수 없는 것들이라는 것을 나타냅니다. DM이 의미 상 필요 없는 정보를 제거하여 불필요한 계산을 최소화하려 해도 학습이나 추론 시에 모든 픽셀에 대해 평가를 해야 하기 때문에 불필요한 계산을 해야합니다. LDM은 효과적인 생성 모델로, 인간이 지각할 수 없는 영역을 제거하는 압축 단계를 분리합니다.

3.1 Perceptual Image Compression

• perceptual loss와 patch-based loss adversarial objective로 오토인코더를 학습하여 local realism을 강제하고 L2와 L1 objectives와 같은 픽셀 공간에 의존함으로써 생기는 블러를 피하게 해줍니다.
• 디코더는 latent로부터 이미지를 재현하고, 인코더는 $2^m$ factor로 이미지를 다양하게 downsample합니다.
• latent spaces의 고분산을 피하기 위하여 KL-reg를 적용하며, 디코더에서 VQ-reg를 적용합니다.
• 기존의 DM은 임의의 1차원 latent를 사용하지만 LDM은 2차원의 latent space를 사용하여 가벼운 압축률과 detail한 정보를 보존하여 매우 좋은 재현성을 달성할 수 있었습니다.

3.2 Latent Diffusion Models

Diffusion Models

Diffusion Model은 길이가 T인 fixed Markov Chain의 역과정에 상응하는 정규 분포 변수의 노이즈를 점진적으로 제거하여 원래의 데이터 분포 $p(x)$를 학습하도록 고안된 확률 모델입니다. 이러한 모델들은 노이즈가 있는 $x_t$의 노이즈를 제거하여 원래의 입력값 x를 예측하도록 학습합니다. 이 때, x는 이미지 원본입니다.

Forward diffusion process

원본 이미지에 노이즈를 점진적으로 더해서 최종적으로 노이즈를 만드는 과정

Reverse diffusion process

노이즈로부터 점차 원본 이미지로 되돌리는 과정

그림 출처 :  https://arxiv.org/abs/2006.11239

 

Generative Modeling of Latent Representations

LDM은 high-frequency, imperceptible details를 추상화한 정보가 담긴 저차원의 latent space에 노이즈를 추가하고 복원하는 과정을 거칩니다. LDM은 DM과 비교했을 때 여러가지 장점이 있습니다. 이 과정은 생성 모델에서 사용되는(데이터 분포의 확률을 최대화하는) likelihood에 DM보다 훨씬 적합하고 중요하고 의미가 있는 bits에 집중하며 저차원에서 효과적으로 전체 연산량을 줄일 수 있습니다.

U-Net 구조로 노이즈로부터 이미지를 복원하는 과정을 거치게 됩니다.

3.3 Conditioning Mechanisms

• cross-attention mechanism을 통해 다양한 modality로 conditioning을 할 수 있게 되어 되었습니다.
• 텍스트가 conditioning으로 주어지면 $τ_Θ$는 트랜스포머와 같은 도메인별 전문가를 통해 매개변수화할 수 있습니다.
• Q는 U-Net representation $(z)를 K, V는 conditioning y을 프로젝션한 representation을 활용하여 프로젝션 행렬 W_Q, W_K, W_V를 학습합니다.

• 최종적인 LDM은 다음과 같은 식을 통해 학습합니다. 노이즈, 반복 횟수 t, condition

4. Experiments

4.1 On Perceptual Compression Tradeoffs

• NVIDIA A100 1개로 down sampling factor {1,2,4,8,16,32} 실험을 진행하였고 factor가 1인 경우는 pixel-based DMs와 같은 경우로 생각합니다.
• Fig6는 샘플의 품질을 factor가 1,2인 경우는 학습 시간이 길어지는 반면 계수가 훨씬 더 높은 값들은 학습한 시간이 얼마 지나지 않아 성능의 정체가 오게 됩니다.
• Fig2에서 보듯이 너무 많은 perceptual compress는 정보의 손실을 야기해 좋은 품질에 대한 제한이 걸립니다.
• LDM-4 ~ 16 사이가 성능이 가장 좋으며 pixel-based인 LDM-1과 비교할 때 FID가 2M step에서 38의 차이가 납니다.

• Fig7은 CelebA-HQ와 ImageNet에서 학습한 모델들을 DDIM 샘플러로 다양한 denoise step으로 실험하여 샘플링 속도와 FID를 나타낸 것입니다.
• LDM 4-8이 가장 좋은 품질을 산출한다는 것을 알 수 있습니다.

4.2 Image Generation with Latent Diffusion

• Table1에서 보듯이 CelebA-HQ, FFHQ, LSUN-Churches, LSUN-Bedrooms 데이터셋에서 FID, Precision, Recall이 매우 좋은 성능을 내고 있음을 알 수 있다.

4.3 Conditional Latent Diffusion

4.3.1 Transformer Encoders for LDMs

• text-to-image 모델링은 BERT-Tokenizer를 활용하여 LAION-400M의 Language prompts conditioning으로 14억 5천만 개의 파라미터를 학습합니다.
• multi-head cross-attention으로 U-Net에 매핑된 latent code를 추론하기 위해 transformer로 $τ_Θ$를 구현합니다.
• 언어와 비전의 도메인별 전문가 조합은 복잡하고, 사용자 기반의 text prompt를 잘 일반화하여 좋은 결과를 산출할 수 있습니다.
• Table2는 분류기가 없는 LDM-KL-8-G이 다른 모델들보다 파라미터 수가 적음에도 더 좋은 성능을 낼 수 있는 것을 볼 수 있습니다.

4.3.2 Convolutional Sampling Beyond

• 공간적인 conditioning 정보를 인코딩 입력 값과 concatenating하여 LDM은 image-to-image 모델링을 설계할 수 있습니다.
• semantic synthesis, super-resolution, inpainting에 대해 모델들을 학습시킬 수 있습니다.

4.4 Super-Resolution with Latent Diffusion

• super-resolution 학습을 위해 LDM은 저해상도의 이미지를 concatenation을 통해 conditioning 합니다.
• Fig 10에서 보듯이, 이미지를 bicubic 방식으로 품질 저하를 시킨 후에 이를 원상 복구시키는 SR3 데이터 처리 파이프라인을 따라 학습을 시킵니다.
• Table 5에서 보듯이, FID에서는 LDM-SR이 우위이지만, IS에서는 SR3가 우위임을 볼 수 있습니다.

4.5 Inpainting with Latent Diffusion

• inpainting은 이미지에서 망가지거나 없애고 싶은 영역을 자연스러운 내용으로 대체하는 기술입니다. Fig 11은 LDM의 결과물입니다.
• pixel 기반의 DM과 latent 기반의 DM은 속도 면에서 최소 2.7배, FID scores 면에서 최소 1.6배의 차이가 나는 것을 알 수 있습니다.

5. Limitations & Societal Impacts

Limitations

• LDM은 pixel 기반의 DM 보다는 계산 필요량이 상당히 줄어들지만 GAN의 순차적인 샘플링 과정보다는 느립니다.
• pixel 공간에서 fine-grained 정확도를 요구하는 task들에서는 재현성이 약간 떨어집니다.
• super-resolution model은 다소 성능에 제한이 있습니다.

Societal Impact

• 이미지 생성 모델은 양날의 검입니다. 훈련, 추론 비용을 줄여 기술에 대한 접근성이 높아지고 많은 사람들이 연구할 수 있는 반면에 조작된 데이터나 잘못된 정보가 쉽게 퍼져 나갈 수 있습니다.
• 특히 deep fake와 같은 정교한 이미지 조작과 특히 여성과 관련된 불균형한 영향을 끼칠 수 있습니다.
• 또한 생성 모델은 민감하거나 개인적인 정보를 노출할 수 있고, 명확한 동의 없이 데이터가 수집될 수 있다는 문제점이 있습니다.
• 딥러닝 모듈들은 데이터에 담긴 bias를 재생산하거나 가중화시키는 경향이 있습니다.

6. Conclusion

LDM은 품질을 떨어뜨리지 않은 채로 노이즈를 제거하는 간단하고 효율적인 diffusion 방식으로 학습과 샘플링 효율을 향상시켰습니다. 이것들과 cross attention conditioning 방식으로 task-specific한 구조를 갖지 않고도 폭넓은 이미지 합성 분야에서 SOTA 방식들과 비교했을 때 좋은 결과를 성취할 수 있었습니다.

참고 )

https://pitas.tistory.com/9

https://lilianweng.github.io/posts/2021-07-11-diffusion-models/

 

 

Stable Diffusion으로 colab을 이용해 이미지 생성해보는 작업을 해볼 수 있습니다.

https://journey-to-serendipity.tistory.com/34

[토이 프로젝트] Stable Diffusion 텍스트로 원하는 이미지 생성해보기(초간단 colab)