Stable Diffusion ; High-Resolution Image Synthesis with Latent Diffusion Models 논문 리뷰

오늘 리뷰할 논문은 Stable Diffusion입니다. Stable Diffusion은 2022년 공개된 Latent Diffusion Models(LDM) 기반의 텍스트-이미지 생성 모델로, 기존 픽셀 단위의 Diffusion Model(확산 모델) 대비 97% 감소한 계산 비용으로 고해상도 이미지 합성 가능성을 입증한 획기적인 연구입니다. 핵심 아이디어는 이미지를 잠재 공간(latent space)에서 처리하는 것으로, VAE 기반의 압축/복원 구조와 결합해 효율성을 극대화합니다.

논문 : High-Resolution Image Synthesis with Latent Diffusion Models
github : https://github.com/Stability-AI/stablediffusion

1. Introduction

논문에서는 Diffusion 모델을 설명하기에 앞서 이전까지 생성 모델 Task에서 인기있었던 GAN 모델에 대해 먼저 이야기합니다. 특히 GAN 모델이 가지고 있던 문제점에 대해 이야기 하며 Diffusion 기반의 모델이기에 기존에 GAN 모델이 가지고 있던 Mode Collapse나 Training instability와 같은 문제가 발생하지 않는다는 점을 이야기합니다.

Diffusion Models

Diffusion Models(DM)는 이미지 데이터 합성에서 최첨단 성능을 달성했습니다. 그러나 이러한 모델들은 주로 픽셀 공간에서 직접 작동하기 때문에, 강력한 DM의 최적화에는 수백 GPU일의 계산 자원이 소모되며 추론 과정도 순차적 평가로 인해 비용이 많이 듭니다.
Stable Diffusion의 핵심 아이디어는 이러한 계산적 한계를 극복하기 위해 확산 과정을 이미지의 잠재 공간(latent space)에서 수행하는 것입니다. 이를 통해 다음과 같은 이점을 얻을 수 있습니다:

1. 계산 효율성 향상: 잠재 공간에서의 연산은 픽셀 공간보다 훨씬 적은 차원에서 이루어집니다.
2. 품질 유지: 적절한 압축률을 선택함으로써 시각적 품질을 유지하면서도 계산량을 크게 줄일 수 있습니다.
3. 유연성: 다양한 조건부 생성 작업에 적용할 수 있습니다.
   저자들은 이러한 접근법을 “Latent Diffusion Models (LDMs)“라고 명명했습니다. 추가로 Transformer를 Diffusion 모델 U-Net의 backbone으로 사용할 수 있게 아키텍처를 구상하였습니다.

위의 내용들을 정리하면
1. 전체적인 구조를 모두 Transformer로 구성한 모델들과 달리 고차원의 데이터에 대해서도 잘 적용됩니다.
2. Unconditional Image Synthesis, Inpainting, Stochastic Super-Resolution 등 여러 Task에서 좋은 성능을 냈습니다.
3. Score-based 모델과는 다르게 Reconstruction에 대한 Delicate Weighting이나 Generative Ability가 필요하지 않습니다.
4. 1024x1024 pixel을 가지는 큰 이미지에 적용할 수 있습니다.

2. Related Work

Stable Diffusion은 여러 선행 연구를 기반으로 합니다. 주요 관련 연구들은 다음과 같습니다:
1. DM(Diffusion Models): 점진적인 노이즈 제거 과정을 통해 이미지를 생성하는 방법입니다. DDPM(Denoising Diffusion Probabilistic Models)이 대표적입니다.
2. VAE(Variational Autoencoders): 데이터를 저차원 잠재 공간으로 인코딩하고 다시 원본 공간으로 디코딩하는 방법을 학습합니다.
3. GAN(Generative Adversarial Networks): 생성자와 판별자의 경쟁적 학습을 통해 고품질 이미지를 생성합니다.
4. VQ-VAE(Vector Quantized-VAE): 이산적인 잠재 표현을 학습하는 방법으로, 이미지 압축에 효과적입니다.
Stable Diffusion은 이러한 선행 연구들의 장점을 결합하여 새로운 접근법을 제시했습니다.

3. Method

Stable Diffusion의 핵심 아이디어는 두 가지 단계로 구성됩니다:

1. 오토인코더를 사용한 Perceptual Compression(지각적 압축)
2. 잠재 공간에서의 확산 모델링

이러한 접근 방법은 다음과 같은 장점이 존재합니다.

1. Efficient : 저차원 공간에서 샘플링이 이루어집니다.
2. High Inductive Bias : U-Net 구조를 사용하여 Compression이 잘 이루어지기 때문에 높은 Inductive Bias를 가집니다.
3. General-Purpose : Latent Space를 사용하여 한번의 학습으로 만들어진 Compression Model을 여러 Downstream task에 활용할 수 있습니다.

Perceptual Image Compression

Perceptual Compression이란 우리가 Autoencoder에서 Latent space를 학습하는 것을 의미합니다. 이 때 Latent Space의 분산이 크면 Latent Space가 가지고 있는 정보가 이질적이므로 작은 분산을 가지도록 Regularization을 가지도록 하였습니다.

이 과정은 다음과 같은 수식으로 표현됩니다:
$z = E(x)x̂ = D(z)$
여기서 E는 인코더, D는 디코더, x는 원본 이미지, z는 잠재 표현입니다.
VAE의 목적 함수는 다음과 같습니다:
$L_{VAE} = L_{rec} + L_{reg}$
$L_{rec}$는 재구성 손실, $L_{reg}$는 정규화 항입니다.

Regularization의 종류

KL-reg : 학습된 Latent에 약간의 KL-penalty를 줍니다.
VQ-reg : Decoder안에 Vector Quantization을 사용합니다.

이 논문에서 사용한 Diffusion 모델은 Latent Space를 2차원 구조로 설계하여 Latent Space가 1차원인 모델보다 Compression과 Reconstruction 성능이 좋았다고 합니다.

Latent Diffusion Models

Diffusion Model은 점진적으로 Denoising을 해가며 $p(x)$의 데이터 분포를 학습하도록 디자인된 확률 모델입니다(고정된 길이의 Marcov Chain 역과정(reverse process)로 학습합니다). 이미지 합성에서 성능이 좋은 대부분의 모델들은 $p(x)$에 대한 ELBO(변분 하한;Variational Lower Bound)를 재가중치하여 변형하여 사용하였으며 이는 Denoising score-matching과 유사합니다. 결국 이 모델들은 동일한 가중치를 가지는 Denoising AutoEncoder의 시퀀스로 해석될 수 있으며 이는 $\epsilon_{\theta}(x_t, t)$로 표현됩니다.

이는 해당 목적함수입니다. 여기서 $ε$은 Gaussian noise, $ε_θ$는 noise 예측 네트워크입니다. $x_t$는 Input x의 noise가 추가된 버전입니다.

Generative Modeling of Latent Representations

학습된 Perceptual Compression 모델을 사용하면 눈에 띄지 않는 디테일들이 효율적이고 저차원인 Latent Space에 접근할 수 있습니다. 고차원인 Pixel Space와 대비하여 Latent Space를 사용하는 것은 두가지 장점이 있습니다.
1. 데이터의 중요하고 Semantic한 정보에 더 초점을 맞출 수 있으습니다.
2. 저차원에서 학습을 진행하기 때문에 계산이 더 효율적이라는 장점을 가집니다.

이러한 구조는 Time-conditional U-Net을 Backbone으로 사용하여 기존 AutoRegressive나 Attention-based Transformer 구조보다 더 풍부한 Inductive Bias를 가지게 만듭니다.

Conditioning Mechanisms

마지막으로 Conditioning Mechanisms을 알아보겠습니다. Stable Diffusion 이전의 Diffusion 모델들에는 다양한 Conditioning 과의 결합에 대한 연구가 없었다고 합니다. Stable Diffusion은 cross-attention을 사용하여 다양한 조건을 사용할 수 있게끔 만들었습니다. Conditioning이란 텍스트나 이미지 등의 추가적인 표현을 의미한다고 볼 수 있습니다. 예로, 텍스트가 조건으로 들어갈 경우가 요즘 자주 보이는 text-to-image 모델입니다.

먼저 각각의 새로운 조건들을 $y$라고 했을 때, 이를 각 단계에 매핑하기 위해 $\tau_\epsilon(y)\in \R^{M\times d_\tau}$ 로 보내는 인코더 $\tau_\epsilon$을 설정하고, 이를 cross-attention layer를 통해 U-Net의 중간 레이어에 매핑시킵니다.
여기서의 크로스 어텐션 매커니즘은 Query는 기존 $z_t$의 값을, Key와 Value는 $\tau_\epsilon(y)$의 값을 가져오는 전형적인 방법입니다. 즉 식으로 표현하면 다음과 같이 표현할 수 있습니다.

여기서 $\phi_i(z_t)$는 U-Net의 중간 단계를 표현한다고 생각하셔도 됩니다.
따라 conditional LDM의 Loss는 다음과 같이 표현할 수 있습니다.

4. Experiments

위의 표는 이 섹션에서 비교한 LDM에 사용된 first stage model의 hyper-parameter 및 재구성 성능을 보여줍니다.

위 그래프는 ImageNet으로 클래스 조건부 모델을 200만 step 학습할 때 step에 대한 샘플 품질(FID, IS)을 보여줍니다.

이를 통해 2가지를 확인할 수 있습니다.

1. 작은 downsampling factor는 학습을 느리게 한다. (LDM-1, LDM-2)
2. 지나치게 큰 f 값은 비교적 적은 step에서 샘플 품질의 정체를 유발한다. (LDM-32)

LDM-4부터 LDM-16까지는 효율성과 perceptual하게 충실한 결과 사이에서 적절한 균형을 유지합니다.

또한 기존 Diffusion 모델, 생성 모델들과 비교해보았을 때, Stable Diffusion 모델은 다른 모델들 대비 좋은 성능을 내었으며, 특히 아래의 Task에서 좋은 성능을 보였습니다.

1. 고해상도 이미지 생성에 더 적은 자원과 시간
2. 여러 조건, 데이터 셋에 대한 모델 유연성 평가

5. Limitations & Societal Impact

Limitations

Stable Diffusion의 한계점은 1. Latent Space 기반의 방법을 사용하여 Pixel 기반 방법론 보다는 빠르지만 Sampling Process가 여전히 GAN보다 느리다고 합니다.또한 2.LDM모델은 높은 정밀도가 요구되는 작업에서 적합하지 않을 수 있다고 합니다.

Societal Impact

이미지 생성 Task에 있어 Train, Inference 비용을 절감하여 기술 접근성을 높임과 동시에, 그만큼 조작된 데이터를 쉽게 배포할 수 있다는 문제가 존재하고 데이터에 대한 윤리적인 문제가 존재할 수 있다고 합니다.

마지막으로

Stable Diffusion은 앞으로도 컴퓨터 비전과 인공지능 분야의 발전을 이끌어갈 것으로 기대됩니다. 이 모델을 기반으로 한 다양한 응용 프로그램과 서비스들이 계속해서 등장하고 있으며, 새로운 표현의 도구를 제공하고 있습니다.
하지만 논문에서 언급되었다시피 기술의 정밀도가 높아감에 따라 딥페이크, 조작 등의 윤리적인 부분에 대해서도 생각해 보아야 할 것 같다는 생각이 듭니다.