StableDiffusion

High-Resolution Image Synthesis with Latent Diffusion Models - https://arxiv.org/abs/2112.10752 

High-Resolution Image Synthesis with Latent Diffusion Models

보니까 알게된게, VAE에서 결국 decoder만 떼서 사용할텐데, 이게 그럼 encoder를 학습할 필요가 있나? 해서 확산모델이 나오지 않았나.

Abstract(논문에 대한 전체적인 요약)

이전 Diffusion Model(DM)의 장점은 재훈련 없이 가이딩 메커니즘을 통해 이미지 생성 과정을 제어할 수 있다는 점이다. 하지만, 이러한 모델은 픽셀 공간에서 작동하므로, 학습에 막대한 시간과 계산 자원이 필요하며, 추론 역시 순차적인 방식으로 진행되어 효율성이 떨어지는 단점이 있다.

 이를 해결하기 위해, 이 논문에서는 사전 학습된 오토인코더(auto-encoder)를 활용한 잠재 공간(latent space)에서의 모델 훈련을 제안한다. 이 방식은 복잡도를 감소시키면서도 세부 사항을 보존하고, 시각적 충실도(자연스러운 이미지 생성)를 유지하며, 첫 시도에서도 거의 최적점에 도달할 수 있는 효과를 제공한다.

 또한, 크로스 어텐션(cross-attention) 레이어를 모델 구조에 도입하여, Diffusion Model이 텍스트, Bounding Box와 같은 일반적인 조건 입력을 처리할 수 있도록 확장했다. 이를 통해, 컨볼루션 연산(convolution)을 기반으로 고해상도 이미지 생성이 가능해졌다.

결과적으로, 제안된 Latent Diffusion Model(LDM)은 기존 DM보다 자원 소모가 적으면서도 이미지 인페인팅(image inpainting), 클래스 조건부 이미지 합성(class-conditional image synthesis) 등 다양한 과업에서 최고 수준의 성능을 달성했다.

 

기존 방식 문제점

•  픽셀 공간에서 작동 → 학습 및 추론 시간이 오래 걸림
•  계산 자원 소모가 큼

새로운 방식 제안

• Auto-encoder 기반의 latent space
• Cross-attention 레이어 도입 → 일반적인 조건 입력 처리 가능
• Conv 기반 연산 → 고해상도의 이미지 생성 가능

결과

• 복잡도 감소 및 자원 효율성
• 시각적 충실도 및 세부 사항 보존
• 다양한 과업에서 최고 성능

Introduction

이미지 합성은 CV 분야에서 큰 발전을 이뤘지만, 동시에 가장 많은 계산 자원을 필요로 하기도 한다. 복잡하고 고해상도의 이미지를 생성하기 위해서는 Auto Regressive의 수십억 개의 파라미터를 포함하는 모델을 사용하거나, Gan도 좋은 성능을 보이지만, 복잡하고 다중 모드의 분포를 모델링 하는데는 한계가 있음이 들어났다. 최근 DM이 나타나며 이미지 합성 뿐 아니라 다양한 응용 분야에서도 조은 결과를 보이고 있고, 클래스 조건부 합성과 초 고해상도의 경우의 state-or-art로 여겨지고 있다. 당연하게 DM은 Gan의 고질적인 mode collapse 문제를 격지 않으며, 파라미터 공유를 통해 AR 처럼 많은 가중치를 사용하지 않고도 고해상도의 이미지를 합성할 수 있다.

Democratizing High-Resolution Image Synthesis

DM은 mode-covering 을 위해 데이터의 세부 사항을 모델링 하는데 과한 컴퓨터 자원을 소비하는 경향이 있는 likelihood 기반의 모델이다. reweigt variational objective 는 denosing process에서 undersampling 하여 이 문제를 해결하려고 하지만, 여전히 컴퓨터 연산량이 높다. 예로 DM을 학습하는데는 단일 A-100 GPU의 50만 개의 샘플을 생성하는데 약 5일이 걸리며, 여러 연구 커뮤니티와 사용자에게 2가지 영향을 미치는데, 1. 거대한 리소스를 감당할 수 있는 극히 일부만이 모델을 훈련할 수 있고, 높은 탄소 발자국을 남긴다. 2. 학습을 마친 모델의 경우에도, 평가할 때, 각 절차를 순차적 실행해야 하기 때문에, 시간과 메모리 비용이 많이 든다.

-> 목적은 DM의 성능을 유지하면서 보다 training sampling 모두에서 적은 리소스를 사용해 접근성을 늘려보자!

Departure to Latent Space

 일반적인 우도 기반의 모델들은 학습을 두 가지 단계로 나눌 수 있는데, 1. perceptual compression stage : 이 단계에서 고 해상도의 세부 정보를 제거하지만, 의미론적 변화를 거의 학습하지 않는다. 2. sementic copression stage : 이 단계에서 실제 생성 모델이 데이터의 의미론적이고 개념적인 구성을 학습한다. 이 논문은 1번에 초점을 맞춰 개념적으로 동일하지만, 계산적으로 더 적합한 공간으로 압축하는 것을 목표로 한다. 우선, 데이터 공간과 개념적으로 동일한 저 차원으로 갈 수 있도록 auto encoder를 학습시키는데, 이 때, 기존 연구들과 달리, 지나치게 강력한 공간 압축을 요구할 필요가 없다. 

 이 방식의 눈에 띄는 장점은 autoencoder를 한 번만 학습시켜 두고, 이 auto encoder를 여러 DM을 학습시키는데 사용할 수 있다는 것이며, text2img의 경우 transformer를 DM의 UNet 백본에 연결하고, 임의의 유형의 토큰 기반 conditioning 매커니즘을 가능하게 하는 아키텍처를 설계한다.

InSUM

1. 전체 구조를 모두 Transformer로 구현한 모델과 달리 고차원의 데이터에 대해서도 잘 작동할 수 있다.

2. 이미지 합성, 고해상도 이미지등 여러 task에서 좋은 성능을 냈다 더불어 추론 비용도 절감한다.

3. encoder/decoder 구조와 점수 기반 사전확률을 동시에 학습하는 모델들과 달리, 재구성 및 생성능력에 세세한 가중치를 부여할 필요가 없다.

4. 초해상도 등 고밀도의 조건부 작업의 경우, 컨볼루션 방식을 적용해 1024x1024의 이미지를 렌더링 할 수 있다.

5. cross-attention에 기반한 범용 컨디셔닝 메커니즘을 설계해, 멀티 모델 학습이 가능하다.

6. encoder 가중치를 깃헙에 공개하고 있다.

요약

기존 DM의 문제가 학습은 물론 추론에도 너무 많은 시간이 든다는 점인데, 이를 해결하기 위해 이 논문이 나왔으며, 기본 아이디어는 모델들이 학습 앞단에서는 low-pixel 정보 뒷단에서는 high-pixel 정보를 학습하는데, 이 때, low-pixel 정보에 의미론적 학습보다는 고해상도를 저해상도로 바꾸는 역할을 하게 되고, 이를 encoder를 통해 대체함으로서, 더 적은 파라미터와 더 적은 추론시간에 동일한 생성물을 내는 것이 이 논문의 idea이다.

Related Work

GAN은 고해상도 이미지를 효율적 생성 및 시각적 품질을 제공하지만, 최적화가 어렵고, 데이터 분포 전체를 포착하는데 어렵다. 이에 반해 likelihood-based 방법은 안정적인 학습이 가능하다. 대조적으로 우도 기반 방법의 VAE와 같은 접근법은 효율적으로 고해상도 이미지를 합성할 수 있지만, GAN과 동일한 품질을 갖지는 못합니다. autoregressive model은 밀도 추정에서는 강력하지만, 연산비용과 순차적이라는 특징 때문에 저해상도 이미지에 제한된다. 이미지의 픽셀 기반 표현은 인간이 모르는 고해상도의 세부 정보를 표현하기 때문에, 자원 소모와 학습시간이 길어지는데, 높은 해상도를 위해서 몇몇의 모델들은 원시 픽셀 대신 잠재 이미지 공간에서 압축한 것을 사용한다.

 최근 DM은 density estimation의 최신모델로서 떠올랐고, 이러한 모델의 생성능력은 UNet으로 구현된 신경망 백본이 이미지와 같은 데이터의 귀납적 편향에 자연스럽게 맞아 떨어지는 데에서 비롯된다. 최상의 합성 품질은 일반적으로 reweighted objective를 사용하여 학습될 때 달성되며, 이 경우 DM은 손실 압축기로 작동하여, 이미지 품질과 압축 성능간 균형을 조절할 수 있다. 다만 픽셀공간을 사용하기에 느린 추론 속도, 매우 높은 학습 비용이라는 단점을 가진다. 느린 추론 속도는 고급 샘플링이나 계층적 접근법을 통해 부분적인 해결이 가능하지만, 이 또한 비용이 많이 든다. 이를 이 논문에서는 LDM으로 해결하고자 한다.

 개별 생성 접근법의 단점을 완화하기 위해 많은 연구에서 2단계 접근법을 통해 서로 다른 방법의 장점을 결합하여 보다 효율적이고 성능이 우수한 모델로 만드는 방법을 연구하고 있다. 그 방식으로 VQ-VAE는 양자화된 잠재 공간 위에 표현력이 좋은 사전분포를 학습하는 방식이고, VQGAN은 첫 단계에서 adversarial 및 perceptual 목적함수를 사용해 autoregressive transformers 를 더 큰 이미지로 확장한다. 그러나 모델의 현실성에 의해 압축률과 성능이 trade-off 관계를 가지게 되는 문제가 발생하는데, 

 이 논문은 그것을 방지한다. LDM은 conv backbone을 통해 더 높은 차원의 잠재공간으로 부드럽게 확장할 수 있다. 따라서, 1단계에서 강력한 학습이 가능하고, 적당한 압축 수준을 선택할 수 있다. 

-> 본인들의 연구가 뛰어난 이유로, 적당한 압축 수준을 선택할 수 있다는 것을 강조하고 있는데, 그 방법에 대해서는 봐야 알 것 같다.

Method

이미 DDPM에서 손실 함수에서 특정 부분을 under sampling 하여 지각적으로 중요하지 않은 세부사항을 무시하는 방법으로 계산량을 줄이려고 노력했지만, 여전히 픽셀공간에서 사용하기에 문제가 있었고, 이 단점의 해결을 위해 압축 학습과 생성 학습의 단계를 명시적으로 구분하는 방식을 제안한다. 이를 위해 오토인코더를 사용하는데, 이는 3가지 장점을 부여한다.

1. 저차원 공간에서의 샘플링을 통해 계산적으로 효율적인 확산 모델을 얻는다.

2. Unet architecture에서의 inductive bias를 활용해 공간적 구조를 가진 데이터에서 효율적이며, 이전 다른 연구들과 달리 고압축이 필요하지 않다.

3. 압축 모델을 따로 설계했기 때문에, 이 오토인코더를 여러 생성모델을 학습하는데 사용이 가능하며, 단일 이미지 기반의 CLIP 유도 합성과 같은 곳에도 응용할 수 있다.

 

3.1 Perceptual Image Compression

이 논문에서 사용할 perceptual compression model 은 perceptual loss, patch based, adversarial objective의 결합을 통한 오토 인코더로 구성된다. image manifold 상으로의 재구성을 보장한다. 

→ perceptual loss(출력 feature map 간의 L2), patch based(지역적인 부분에 대한 재현률을 보장)

x ← 원본이미지로 부터(H x W x 3)

E ← 인코더

z ← E(x)

D ← 디코더

x_h ← D(z) ← D(E(x)) 이때 z는 (h x w x e)

압축률은 f = H / h = W / w 로 정의될 수 있다. 그리고 2의 배수 기반의 압축 비율을 실험한다.

임의의 높은 분산의 잠재 공간을 피하기 위해 두 종류의 정규화 방법을 연구한다.

첫째는, KL-reguralization 인데, 이는 VAE에서 처럼 학습된 잠재 표현에 대해 표준 정규 분포에 약간의 KL 페널티를 추가한다.

둘째는, VQ-reguralization 인데, 디코더 내부로 양자화 계층이 흡수된 형태로 해석이 가능하다.

LDM은 학습된 잠재 공간의 2차원 구조를 활용하도록 설계되어, 비교적 낮은 압축률로도 좋은 재구성을 할 수 있다. 이는 Z의 분포를 autoregresively 모델링 하기 위해 임의의 1차원 순서에 의존했던 다른 모델들과 달리, Z의 본질적인 구조를 보존한다는 것이 차이점를 가진다.

 

3.2 Latent Diffusion Models

Diffusion Model은 확률적인 모델로 데이터의 분포 p(x)를 학습하도록 디자인 되어있고, p(x)는 마르코프 체인의 reverse process를 학습하여, 정규 분포된 변수들을 denoising 하는 과정을 학습하도록 설계되었다. 이미지 합성에 가장 성공적인 모델인[30, 15, 72 diffusion 기반 모델] 는 p(x) 의 reweighted variant of the variational lower bound 에 의존한다. 그리고 이것은 denoising score-matching의 과정으로 여겨질 수 있다. 이런 모델들은 weighted sequence of denoising autoencoder로서 해석될 수 있다. 

auto encoder를 통과한 e(x)에 대한 loss 식을 사용하면 e_theta 의 backbone network는 UNet이다.

Generative Modeling of Latent Representations

훈련된 지각 압축 모델은 ℰ 과 𝒟 로 구성되었으며, 우리는 이제 고해상도의 지각할 수 없는 세부사항들을 효율적으로 저차원의 잠재공간에 넣을 수 있게 되었다. 고차원의 pixel 공간과 비교해 이 공간은 우도 기반의 생성모델에 더 적합하며 이제 중요하고 의미있는 데이터에 집중하고 저차원의 계산적으로 더 효율적인 공간에서 학습할 수 있게 되었다.
고수준으로 압축된 잠재 공간에서 auto regressive, attention 기반 transformer 모델에 의존했던 것과 달리, 모델은 이미지의 inductive bias를 사용할 수 있고, 여기에는 주로 2D Conv에서 기본 UNet을 구축하는 기능과 reweighted bound 를 통해 목표에 더 잘 집중할 수 있다.

백본은 UNet(time-conditional UNet) 으로 구현되며, 순방향 과정이 고정되어 있어 학습중에 e_theta는 효율적으로 얻을 수 있다. 더불어 p_theta(z)의 샘플링 결과는 D를 통한 single pass로 디코딩 될 수 있다.

 

3.3 Conditioning Mechanisms

다른 생성 모델들 처럼, diffusion 모델도 원칙적으로 조건부 분포 모델링이 가능한데, 클래스 레이블이나 흐릿한 변형을 넘어서는 조건부 생성 모델은 아직 나타나지 않았는데, Unet 백본에 cross-attention을 추가해 확산 모델을 더 유연한 조건(c)부 이미지 생성기로 만들 수 있다. (즉 text와 같은 여러 정보를 보다 쉽게 받아드릴 수 있다) 이를 위해 domain specfic encoder  τθ​
를 추가하는데,  이 인코더는 c를 중간 표현으로 투영하며, 이는 Unet의 중간 레이어에 cross-attention을 통해 매핑된다. 이때 조건부 이미지 생성기의 backborn으로는 CLIP을 사용하는데, fine-tuneing 할 때는 같이 업데이트 할 수 있고, 대부분의 경우 얼려서 사용한다고 한다.

이 식을 통해 인코더와 tao를 같이 최적화 한다.