[논문 리뷰] Latent Diffusion : High-Resolution Image Synthesis with Latent Diffusion Models

paper : [2112.10752] High-Resolution Image Synthesis with Latent Diffusion Models (arxiv.org)

High-Resolution Image Synthesis with Latent Diffusion Models

---

 

Stabe Diffusion은 2022년 발표된 text-to-image 딥러닝 모델이다. Stable Diffusion은 Latent Diffusion model(LDMs)의 일종이고, 위의 연구를 기반으로 만들어졌다. LDMs가 일반적으로 latent space에서의 확산 과정을 통해 이미지를 생성하는 방식을 사용한다면, Stable Diffusion은 이러한 기본 아이디어에서 추가적으로 사용자 친화적인 기능, 향상된 품질 및 효율성 등의 차별점을 가진다. 

 

+) Stable diffusion과 Latent diffusion은 엄밀히 말하면 다른 개념이다! Latent diffusion은 diffusion model의 한 종류로, latent space에서 denoising diffusion process를 수행하는 방법론을 의미하고, Stable diffusion은 Latent diffusion을 기반으로 학습된 text-to-image 모델의 이름이다. 즉, Latent diffusion은 방법론의 이름이고, Stable diffusion은 Latent diffusion 방법론을 활용해 만든 특정 모델 이름이라고 볼 수 있다. 

 

- LDMs의 핵심 아이디어 : pixel space에서 latent space로의 전환

기존의 Diffusion Model(DM)은 데이터 그 자체에 대해 노이즈를 주고 빼는 것으로 구성되어 있다. 이러한 모델들은 이미지 합성에 큰 성과를 얻었으나, pixel space 내에서 직접 작동하기에 훈련 과정과 inference 과정에서 계산 비용이 많이 발생한다는 단점을 가지고 있었다. 이에 제한된 컴퓨터 자원 안에서 퀄리티와 다양성을 유지하며 훈련시키기 위한 방법으로 사전 학습 되어있는 autoencoder들의 latent space(저차원 공간) 내에서 DM을 적용시켜보자는 주장이 대두되었고, latent diffusion model의 기본 아이디어로 자리잡혔다. 

• 머신러닝에서의 autoencoder : 입력 차원과 출력 차원이 동일한 모델 구조
• LDMs에서의 autoencoder : 입력 데이터를 낮은 차원의 잠재 공간(latent space)으로 인코딩하고, 이 잠재 공간에서 원래 입력 데이터를 디코딩해 복원하는 인공 신경망 구조

+) DM은 이미지 공간에서 작동하는데, 이미지 공간은 크기가 엄청나게 크다. 3개의 RGB 채널이 있는 512 * 512 이미지라며느 786,462차원을 필요로한다. Stable Diffusion은 이러한 문제를 latent space로 이미지를 압축한 뒤 연산을 수행해 이러한 문제를 해결한 것이다. latent space는 image space에 비해 48배나 작은 4 * 64 * 64(16,384)이기 때문에 연산이 빨라진다. stable diffusion에서 순방향 디퓨전 및 역방향 디퓨전이 모두 이 latent space상에서 이루어진다. 

 

---

 

- 일반 생성모델 : 이미지를 두 단계로 학습한다. 

Step 1) Perceptual Compression(지각적 압축) - 인간의 지각에 중점을 둔다

high-frequency details을 제거하지만, 적은 의미적 변동을 학습하는 단계이다. Semantic Compression은 유지되는 구간이다. 주로 이미지나 비디오를 저장하거나 전송할 때, 인간의 시각 체계에 기반해 중요한 정보는 보존하고 데이터를 압축하는 기술을 의미한다. 일반적으로 DM은 이 단계에서 과도한 시간이 걸린다는 특징이 있다. 

Step 2) Semantic Compression(의미적 압축) - 데이터의 의미적 정보에 중점을 둔다

이미지나 데이터의 의미적 및 개념적 구성을 학습하는 단계이다. 이미지나 데이터의 이미적 정보를 보존하면서 압축하는 기술을 의미한다. 

 

두 단계를 거쳐 지각적으로는 동등하지만 계산적으로 더 적합한 공간을 찾아내 DMs를 훈련시킨다. 정리하면, 생성 모델은 처음에 학습을 할 때 Perceptual Compression에서 이미지에 대한 전반적인 형채에 대해 학습하게 되고, 다음으로 Semantic Compression에서 실제 데이터의 본질이 학습되는 것이다. 

 

bit rate에 따른 loss를 바탕으로 GAN, Autoencoder와 LDM을 비교한 그래프

 

 

Bit Rate가 작아질수록 사람의 눈에는 인지가 안되는 부분을 의미한다. GAN과 Autoencoder는 Bit Rate가 비교적 큰 Perceptual한 부분을 학습하는 데에 주력하며, Latent Diffusion Model은 Bit Rate가 작은 Non Perceptual한 부분을 학습하는데에 초점을 맞춘다. 사람 눈에 이 부분은 그다지 중요하지 않은 부분으로, 따라서 Stable Diffusion은 이러한 부분을 개선하여 Perceptual 한 부분에 초점을 맞춰 학습하는 Diffusion Model을 제안한다. 

 

본 논문에서 제안된 모델의 전체적인 구조는

1) Autoencoder + GAN을 통한 원본 이미지의 Perceptual Compression 과정

2) LDM을 통한 Semantic Compression 과정

으로 구성된다. 

1)과 2)는 각각 학습된다. 이는 이전 연구들이 encoder / decoder 아키텍처를 diffusion에 사용했을 때 둘을 score-based 평가 지표를 통해 함께 최적화하던 것에 비해 효율적이다. 

 

디지털 이미지에서 대부분의 bits는 사람이 인지할 수 없는 것들로, DM이 의미 상 필요없는 정보를 제거해 불필요한 계산을 최소화하려고 해도 학습이나 추론 시에 모든 픽셀에 대한 평가를 해야하기 때문에 불필요한 계산이 요구된다. LDMs은 이때 인간이 지각할 수 없는 영역을 제거하는 압축 단계(compressive)를 분리한다. LDMs는 눈에 띄지 않는 세부 사항만을 제거하는 별도의 약한 압축 단계를 포함한다. 

 

- Method

고해상도 이미지 합성에서의 DM의 training 계산 요구를 위해 DM이 loss term을 다운샘플링해 지각적으로 관련이 없는 디테일들을 무시할 수 있게 하더라도, 픽셀 공간에서 비용이 많이 드는 함수의 평가를 필요로 하며, 이는 큰 계산 시간과 에너지 자원이 요구된다. 

• 고해상도 이미지 = pixel 수 많음 -> 손실 함수 계산에 높은 비용 요구됨
• 관련이 없는 디테일 무시 & 다운샘플링으로 loss 계산 -> 계산 비용 줄이기 가능
  
  • 다운샘플링 : 이미지를 축소해 pixel 수를 감소시키는 작업. 이 문장에서의 다운샘플링은 손실 함수 계산에 적용되는 것

위 단점을 피하기 위해서, LDMs에서는 compressive와 generative learning 단계를 명시적으로 분리하는 것을 제안한다. 이미지 공간과 지각적으로 동등한 공간을 학습하면서도 계산 복잡성이 크게 감소한 오토인코딩 모델을 활용한다. 

• '지각적으로 동등' : 원본 이미지의 주요한 시각적 특성이나 구조를 유사하게 보존한다는 것

 

이런 구조를 통해 얻을 수 있는 장점은 아래와 같다. 

1. 고차원 원본 이미지 공간을 떠남으로써 훨씬 더 효율적인 모델을 만들 수 있다. (이미지 생성(sampling)이 저차원 공간에서 수행되기 때문)

2. DMs이 UNet 아키텍처를 사용하는 것으로부터 얻는 inductive bias를 활용해 공간적인 구조(spartial structure)를 더 잘 활용할 수 있다. 이전에 1차원으로 압축했던 시도들에 비해 좋은 퀄리티를 얻을 수 있다. 

•  inductive bias : 주어지지 않은 입력의 출력을 예측하는 것. 일반화의 성능을 높이기 위해 만약의 상황에 대한 추가적인 가정(Additional Assumptions)라고 볼 수 있음. 
• UNet 아키텍처는 이미지 내의 공간적인 구조를 효과적으로 학습하는 inductive bias를 내재하고 있음.
  • UNet : 주로 이미지 분할(image segmentation) 작업에서 사용되는 신경망 구조. 대칭적인 형태를 가지며, 네트워크의 앞부분(encoder)에서는 입력 이미지의 차원을 줄이며 중요한 특징을 추출하고, 뒷부분(decoder)에서는 이 특징들을 확장해 원본 이미지의 크기로 복원함

3. 다른 generative model들을 훈련시키거나 다른 종류의 downstream application에도 활용할 수 있다. (autoencoder의 재사용이 가능하다.)

 

1. Perceptual Image Compression :  autoencoder를 학습하는 과정

perceptual compression model은 perceptual loss(지각적 손실)와 patch-based loss adversarial objective로 훈련된 오토인코더를 사용한다. 이는 local realism을 강화하고 L2와 L1와 같은 pixel space의 손실만을 사용하는 것의 한계를 극복해 블러 현상 피하게 해준다. 

• 기존의 DMs에서는 loss term(손실항)을 최소화함으로써 의미 없는 정보를 억제할 수 있었는데, gradient와 neural network backbone은 여전히 모든 픽셀에서 계산되어야 했다. 이는 불필요한 계산과 많은 비용이 드는 최적화(optimizer)와 추론(inference)로 이어졌다. 

latent space의 고분산을 피하기 위해 KL-reg를 적용하고, decoder에서 VQ-reg를 적용한다. 기존의 DM은 임의의 1D latent을 사용하지만 LDM은 2D의 latent space를 사용해 가벼운 압축률과 detail한 정보를 보존해 매우 좋은 재현성을 달성할 수 있다. 

 

2. Latent Diffusion Models

Diffusion Models

DMs은 데이터 분포 p(x)를 학습하기 위해 설계된 확률 모델로, 정규 분포된 변수를 점진적으로 잡음 제거함으로써 작동한다. 이는 길이가 T인 fixed Markov Chain의 역과정을 학습하는 것과 동일하다. 

즉, DMs은 데이터의 잡음이 있는 버전에서부터 점차적으로 잡음을 제거함으로써 원본 데이터 분포를 학습하는 방식으로 작동한다. 이 과정은 denoising autoencoder의 sequence로 구성되고, 노이즈가 있는 $x_t$의 노이즈를 제거해 원래의 입력값 x를 예측하도록 학습한다.

기존 DM의 loss function

• x : 이미지 원본
• 각 오토인코더 : 시간 단계 t에 따라 입력의 잡음을 점차 줄여나가는 역할
• Forward diffusion process : 원본 이미지에 노이즈를 점진적으로 더해서 최종적으로 노이즈를 만드는 과정
• Reverse diffusion process : 노이즈로부터 점차 원본 이미지로 되돌리는 과정

Stable Diffusion은 Latent Diffusion 모델이므로 당연히 Latent Diffusion의 구조를 따른다. Stable Diffusion은 이미지에 노이즈를 주고 이를 다시 역산해 신경망에 노이즈 이미지에서 일반 이미지를 복원시키도록 학습시킨 후, 이 신경망에 노이즈만 들어있는 이미지와 적절한 힌트를 주면 해당 힌트에 맞는 이미지를 출력한다. 이는 Stable Diffusion, 즉 Latent Diffusion 모델에 대한 설명이라기 보다는 Diffusion 모델에 대한 설명인데, Latent Diffusion의 기본적인 이미지 생성 원리도 이와 같지만 더 발전했다.

Diffusion Model Process

Latent Diffusion도 마찬가지로 노이즈에서 복원하는 것은 같지만, 이미지를 바로 복원하는 것이 아니라 Latent vector를 복원한 후, 이 Latent vector가 다시 VAE를 거쳐야 이미지가 되는 것이다. 

• θ : denoising 과정을 수행하는 신경망의 파라미터
  • forward diffusion process는 신경망 파라미터와 직접적인 관련이 존재하지 않기에 θ와 같은 파라미터가 조건부확률 식에 존재하지 않는 것. forward diffusion process는 일반적으로 미리 정의된 잡음 스케줄에 따라 진행되며, 학습 가능한 파라미터 없이 고정된 함수나 규칙을 사용함.

 

Generative Modeling of Latent Representaions

이미지 생성 중심으로 표현한 Latent Diffusion 모델 구조

• 빨간색 부분(pixel space) : autoencoder
  • ε : 데이터 x를 latent space로 보내는 encoder. z라는 latent representation으로 보냄
  • D : latent z에서 이미지 복원하는 decoder
  • autoencoder로 정보를 압축해준 덕분에 high frequency와 사람 눈으로 인식이 안되는 noise성 정보들은 모두 제거되었다고 볼 수 있음
  • Stable Diffusion에서는 디코더로 VAE(Variational AutoEncoder) 모델을 사용함
• 초록색 부분(latent space) : latent diffusion model
  • z : Latent Embedding
  • diffusion process : 압축된 z에 noise를 점차 주입해 거의 완전한 noise를 만들고, 랜덤한 noise z_T로부터 이를 복원
• 회색 부분(conditioning) : 어떠한 condition에 어울릴만한 데이터를 생성하기 위해 추가한 모듈. 입력된 condition은 diffusion model 내부에서 cross-attention 연산을 사용해 적용됨

더 간단한 구조도로 latent diffusion의 main component 3가지에 대해 설명하면 아래와 같다. 

 

Latent Diffusion의 더 간단한 전체 구조도

Part 1. Text Encoder

latent diffusion에서는 prompt 텍스트를 입력받고 그 텍스트에 맞는 이미지를 출력한다. 이를 수행하기 위해 CLIP 이라는 text encoder를 사용한다. text 형태인 prompt를 UNet이 이해할 수 있는 형태, 즉 token으로 변환하는 것이다. text encoder는 Tokenizer를 이용해서 문장에서 단어를 추출해 숫자로 변환하고(tokenize), 이 숫자를 latent vector의 형태인 text embedding으로 만든다. text embedding으로 변환하는 이 과정을 거쳐야 text가 이미지를 생성하는 UNet에 conditioning 할 수 있게 되는 것이다.

• CLIP (Contrastive Language-Image Pre-Training) : OPEN AI에서 개발한 딥러닝 모델로, 이미지에 대한 텍스트 설명을 설명하기 위해 사용됨.

 

Part 2. UNet

text encoder에서 만들어진 text embedding은 UNet으로 전달된다. UNet에서는 text embedding에 따라 조건화(conditioning)된 채로 random latent vector를 n번 반복해 denoise하는 과정을 거친다. (여기서 랜덤한 노이즈에 이미지를 생성하는 부분이 diffusion 모델의 원리라 할 수 있다. ) 이 반복 전에 어떤 방식(노이즈의 세기, 종류, 확률 편미분 방정식 이용 등)으로 처리하고 반복하느냐를 결정하는 것이 scheduler의 역할이다. scheduler의 종류로는 DDPM, DDIM, PNDM, Euler, Euler a, DPM++ 등이 있다. 

그리고 이 모델은 text embedding에 의해 조건화되었기 때문에 텍스트로 우리가 원하는 내용을 입력해 이미지를 출력할 수 있는 것이다. 이 conditioning 과정은 attention 등의 방법을 사용해 이루어진다. 

위 과정을 모두 거치면 저해상도의 (위 예제에서는 64 * 64) latent vector 값이 생성되게 된다. 이 값은 마지막으로 VAE로 전달된다. 

 

Part3. VAE 

데이터 분포의 latent space를 학습하는데 사용되는 기술이다. encoder는 어떤 값을 수학적 원리를 통해 그 값의 특징을 추출하여 학습하고 decoder는 임의의 값 z(특징에 대한 latent vector)가 주어지면 그 값(latent reresentation)을 바탕으로 원래 데이터(image)로 복원하는 역할을 한다. encoder는 낮은 dimensional latent로 변환하고, 이 변환된 내용은 UNet의 input으로 들어가게 된다. VAE는 그 자체로 하나의 모델이며, 미리 학습되어 준비된 것(Pretrained Model)을 사용한다. 

Latent diffusion 모델에서는 UNet에서 학습되고 이를 바탕으로 생성된 값 자체가 diffusion 모델처럼 이미지의 픽셀과 같은 원본 값이 들어간 게 아니라 VAE에 의하여 encode되어 latent vector화 된 값을 UNet에 학습시킨 것이기에 UNet에서 복원되어 나온 저해상도의 latent vector를 VAE로 다시 Decode하여 고해상도의 그림으로 만들어 준다. 

 

 

 

x를 512 * 512 * 3의 RGB 원본 이미지라고 가정하면, x는 512 * 512 * 3의 pixel space상에 있다고 할 수 있다. 이를 encoder ε를 거쳐 128 * 128 * 16 이라는 보다 저차원의 latent space상의 z로 projection 한다. (latent embedding)

• z는 h * w * c의 차원을 가지는데, 이는 f라는 하이퍼파라미터로 조정된다.
• H / f = h의 식으로 조절할 수 있다.
  • f = 8이고 H = 1024라면 h = 1024/8 = 128이 된다.

 

• f의 값을 너무 높이면 샘플링의 퀄리티가 떨어지는 trade off 관계가 있으나 4~16으로 설정했을 때 좋은 밸런스를 보여준다고 한다. 

이 z에 noise를 점차 주입하는 diffusion process를 따르게 하여 원래 diffusion model의 구조를 사용하게 한다. encoder로 압축된 z에 noise를 점차 주입해 T번의 step 후에 거의 완전한 noise z_T를 만들고, denoising autoencoder의 연결로 z_T를 z로 복원한다. 이때의 denoising autoencoder은 UNet의 구조를 따르고, 이 안에는 transformer도 포함되어 있다. 즉, UNet 구조와 cross-attention을 활용해 z_T를 T번의 denoising step을 거쳐 z로 복원한다. 

• denoising autoencoder : autoencoder에 의도적인 noise를 삽입한 데이터를 다시 복원하는 autoencoder
• transformer를 통하면 text to image, image to image, super-resolution 등의 각기 다른 객체들을 attention을 통해 관계를 알 수 있게 되어 condition을 추가해줄 수 있게 된다. 

이 과정들을 거쳐 random noise z_T를 Z로 복원한다. 이 때 z는 아직 latent space 내에 있으므로 이를 decoder D로 h * w * c의 RGB 이미지로 다시 decoding 해주면 우리가 원하는 생성 이미지(~x)를 얻을 수 있게 된다. 

LDM의 loss function

 

loss function을 기존의 DMs와 비교해보면, LDM에서 픽셀값 x가 아닌 latent embedding z가 들어간다는 차이를 볼 수 있다. 이렇게 autoencoder로 정보를 압축해준 덕분에 high frequency와 사람 눈으로 인식이 안되는 noise성 정보들이 제거되었다고 볼 수 있다. 따라서 좀 더 semantic한 부분에 집중할 수 있으면서 계산 복잡도는 줄어든다. 

 

 

3. Conditioning Mechanisms

다른 생성모델과 비슷하게 DMs도 원칙적으로 p(z|y) 형태의 조건부 분포를 모델링할 수 있다. 이는 conditional denoising autoencoder(조건부 잡음 제거 오토인코더) ε_ θ(z_t, t, y)로 구현된다. 

LDMs은 UNet 구조에 cross attention mechanism을 추가함으로써 DMs를 더 유연한 조건부 이미지 생성기로 전환시킬 수 있다. LDMs에서는 다양한 condition을 입력받을 수 있지만 각자에 맞는 적절한 encoder 𝜏가 필요하다. 

입력된 condition은 diffusion model 내에서 연산 중인 이미지 정보 z_T와 연산되어야 한다. 이때 이 둘의 상관관계를 잘 고려해 이미지 정보 z_T에 반영해주어야 한다. 이 과정에서 cross-attention이 사용된다.

• 서로 다른 정보의 상관 관계를 고려하기 위해 가장 많이 사용되는 방법이 cross-attention이다. 
• self-attention : 하나의 정보에 대한 상관관계 계산
• cross-attention : 동일한 메커니즘을 사용해 두 정보의 상관관계를 고려하는 방법

LDMs에 적용된 cross-attention을 수식으로 표현하면 아래와 같다. 

LDMs에 적용된 cross-attention

query, key, value 수식

 

• z_t : 이미지 정보
• y : condition 정보

이는 이미지와 condition의 상관관계를 고려해 condition 정보에 가중치를 반영하는 것이라고 해석할 수 있다. 이렇게 가중치가 반영된 condition 정보는 최종적으로 z_t에 다시 더해진다.

 

이렇게 cross-attention까지 적용된 최종 LDMs의 loss function은 아래와 같다.

condition이 추가된 LDMs의 loss function

 

 

 

---

 

 

 

Stable Diffusion - Diffusion Model 논문 리뷰 (ffighting.net)

Stable Diffusion - Diffusion Model 논문 리뷰

스테이블 디퓨전 원리 알아보기, High-Re.. : 네이버블로그 (naver.com)

스테이블 디퓨전 원리 알아보기, High-Resolution Image Synthesis with Latent Diffusion Models 논문 리뷰

Stable Diffusion이 대체 무엇일까?(Latent Diffusion의 작동 원리) (tistory.com)

Stable Diffusion이 대체 무엇일까?(Latent Diffusion의 작동 원리)