LDM 논문 리뷰(High-Resolution Image Synthesis with Latent Diffusion Models)

Diffusion model 에서 LDM이 유명하다고 해서 논문 리뷰를 해보고자 한다.

Introduction

LDM은 perceptually 동등하지만 계산적으로 더 적합한 공간을 찾아 그 공간에서 diffusion model을 훈련하고자 한다. 따라서 LDM은 perceptual compression을 통해 계산 비용을 줄이면서도, semantic compression을 통해 중요한 의미적 정보를 유지하고자 한다.

 

* Semantic Compression : 이미지의 중요한 의미나 개념적 정보를 유지하면서 불필요한 세부사항을 제거하는 과정

* Perceptual Compression : 인간의 시각적 지각에 크게 영향을 미치지 않는 세부정보를 제거하여 데이터를 압축하는 과정

 

우선 LDM은 autoencoder를 훈련하여 데이터 공간과 지각적으로 동등한 저차원의 latent space를 생성한다. 이렇게 생성된 잠재 공간에서 single network pass만으로도 효율적으로 이미지 생성이 가능하게 한다. latent space에서 DM을 학습하기 때문에 우리는 이 모델을 Latent Diffusion Models(LDM)이라고 부르는 것이다. 

이 방법의 장점은 autoencoder를 재활용할 수 있다는 점이다. 이를 이용해 많은 종류의 diffusion models에 적용 가능하다. 특히 text-to-image작업은 트랜스포머와 Unet 백본을 연결하여 설계할 수 있다.

 

LDM의 장점

1) 순수하게 transformer만 사용했을 때 보다 고해상도에서 유리

2) 여러 작업(unconditional image synthesis, inpainting, stochastic super-resolution)과 데이터셋에서 좋은 성능을 보이면서도 계산 비용↓

3) reconstruction이 잘 되며, latent space에 대한 제약도 거의 필요하지 않다.

4) cross-attention 기반의 general-purpose conditioning mechanism을 사용하여 multi-modal training이 용이하게 한다. 

 

related work

VQ-VAEs : Autoregressive models을 사용하여 discretized latent space에서 표현력이 풍부한 사전 분포(expressive prior)를 학습한다. 이 방법은 text-to-image 생성에서 discretied image와 text 표현에 대한 공동 분포를 학습한다. 

VQGans : adversarial objective과 perceptual objective을 사용하여 autoregressive transformer를 더 큰 이미지로 확장한다.

* VQ : Vector Quantized

 

Method

LDM에서는 압축 학습 단계와 생성 학습 단계를 명확히 분리한다. 이를 위해 이미지 공간과 지각적으로 동등하지만 계산 복잡도가 크게 줄어든 공간을 학습하는 오토인코딩 모델을 활용한다. 이러한 접근 방식은 (1) 저차원 공간에서 샘플링이 수행되기 때문에 계산적으로 효율적이며 (2) Unet 구조에서 유래한 DMs의 inductive biase을 활용할 수 있고(따라서 과거 연구들처럼 공격적으로 품질을 저하시킬 필요가 없다.) (3) 잠재 공간을 사용하여 여러 생성 모델을 훈련할 수 있는 범용 압축 모델을 얻을 수 있다.

 

Perceptual Image Compression

LDM에서는 autoencoder 학습에 perceptual loss, patch-based adversarial objective를 사용하고 있으며, 이는 재구성이 image manifold를 제한하도록 하여 지역적 현실감을 강화하고 L1이나 L2 목표 함수 등에만 의존할 때의 흐릿함을 방지한다.

* perceptual loss : L1이나 L2 손실과 같은 픽셀 기반 손실 함수 대신, 인간의 시각적 지각과 유사하게 이미지를 평가하는 손실함수로, 생성된 이미지의 시각적 품질을 향상시킨다.

* patch-based adversarial objective : GAN에서 사용되는 적대적 학습 방식을 패치 단위로 적용하여 이미지의 지역적 현실감을 보장하는 방법

 

아래 그림에서 image $ \mathbf{x} \in \mathbb{R}^{H \times W \times 3}$, \mathbf{z} \in \mathbb{R}^{h \times w \times c}$

이며 인코더가 image를 $f = H/h = W/w$로 downsample하였다. 이 논문에서는 $f = 2^m$으로 놓고 $m$을 바꾸면서 실험할 예정이다.

 

또한 잠재 공간의 변동성이 커지는 것을 막기 위해 2가지 regularization을 실험해본다.

(1) KL-reg : 학습된 잠재 공간에 대해 표준 정규분포로의 약간의 KL-penalty를 부과

(2) VQ-reg : 디코더 내에 VQ layer를 사용한다. 이 모델은 VQGAN과 흡사하지만, 여기서는 VQ layer가 디코더에 이미 흡수된 형태로 있다. 

 

LDM은 학습된 잠재 공간 $z$의 이차원 구조를 활용하도록 설계되어 있기 때문에, 비교적 낮은 압축률을 사용할 수 있으며, 매우 우수한 재구성을 달성할 수 있다.

 

Latent Diffusion Models

DM은 일렬의 weighted sequence of denoising autoencoders ${\epsilon}_{\theta}(x_t, t)$로 해석할 수 있다. DM의 objective는 아래와 같다.

 

우리는 이제 인코더와 디코더로 이루어진 perceptuals compression model을 통해 저차원의 잠재 공간을 다룰 수 있으며, 고차원 공간과 비교했을 때 likelihood-based generaitve model에 유리하다. 이는 (1) 중요한 semantic 정보를 집중적으로 다루며, (2) 저차원에서의 학습이 계산량이 적기 때문이다.

 

LDM은 이미지 특유의 inductive bias를 활용하고 주로 2D convolution layer로 Unet을 구축하여 reweighted bound를 사용해 지각적으로 중요한 비트에 목표를 집중시킨다. 목적함수는 아래와 같다.

여기서 신경망의 backbone ${\epsilon}_{\theta}(  , t)$는 시간 조건부 Unet으로 구현된다. forward process는 고정되어 있으므로 훈련 중에 $z_t$를 인코더로부터 효율적으로 얻을 수 있으며, $p(z)$에서 샘플을 추출하여 디코더를 한 번 통과시키면 이미지공간으로 디코딩할 수 있다.

 

Conditioning Mechanisms

DM 역시 conditional distribution $p(z|y)$가 모델링이 가능하다. 이는 conditional denoising autoencoder  ${\epsilon}_{\theta}(z_t, t, y)$를 구현하고, condition $y$를 input으로 넣어주면 된다. 하지만 conditional DM은 아직 연구가 많이 이뤄지지 않은 분야로 LDM에서는 유연한 조건부 이미지 생성기를 만들기 위해 기본 Unet 백본에 cross-attention mechanism을 추가했다. cross-attention mechanism은 다양한 입력 형태에 대해 attention 기반 모델을 학습하는데 효과적이다.

$y$의 전처리를 위해 LDM에서는 도메인에 특화된 인코더 ${\tau}_{\theta}$를 도입하여 $y$를 중간 표현 $ \mathbf{ {\tau}_{\theta}(y) } \in \mathbb{R}^{M \times d_{\tau} }$로 투영한다. 그런 다음 cross-attention layer를 통해 Unet의 중간 layer로 매핑한다. 여기서 cross-attention layer는 아래와 같이 구현된다.

여기서 $\varphi_i(z_t) \in \mathbb{R}^{N \times d_{e}^{i}}$는 Unet 구조 내에서 생성된 중간 표현이며, $Q$, $K$, $V$와 곱해지는 $W$ 행렬은 모두 학습 가능하다. conditional LDM의 objective는 다음과 같이 표현 가능하다.

 

Experiments

- On Perceptual Compression Tradeoffs

downsampling factors $f \in {{1, 2, 4, 8, 16, 32}}$에 대해 LDM-f를 생각해보자. 실험 결과를 분석해보면 i) LDM-{1, 2}의 작은 downsampling factor는 훈련이 느리고 ii) f의 값이 지나치게 큰 경우 비교적 적은 훈련 단계 후에 fedelity가 정체된다. 즉, i)에서 대부분의 perceptual compression을 확산 모델이 맡느라 시간이 오래 걸리고, ii)에서는 첫 번째 단계의 압축이 너무 강해 정보 손실을 초래해 달성 가능한 품질을 제한한다고 생각할 수 있다. LDM-{4, 16}은 효율성과 지각적으로 충실한 결과 사이에 좋은 균형을 이룬다. 

 

- Image Generation with Latent Diffusion

LDM은 Celeb-HQ에서 기존 GAN과 likelihood-based model보다 좋은 FID score를 얻었으며, LSGM(Latent Score Generative Model)보다 더 나은 성능을 보였다. LSUN-Bedrooms 데이셋에서는 다른 데이터셋보다 성능은 약간 떨어지지만 매개변수의 수와 훈련 차원이 크게 늘었음에도 ADM과 비슷한 점수를 기록했다.

또한 LDM은 Precision과 Recall 측면에서도 GAN보다 좋았는데, 이는 적대적 접근법에 비해 mode covering likelihood-based training objective의 이점을 보여준다. 

* precision : 생성된 이미지 중 실제 데이터와 유사한 이미지의 비율

* recall : 실제 데이터의 다양성을 얼마나 잘 반영하는지

 

※ Conditional LDM의 실험 결과는 논문을 참고하세요.

 

Limitation

LDM은 pixel-based approach에 비해 계산 요구 사항이 크게 줄었지만, 그 순차적 샘플링 과정은 여전히 GAN보다 느리다. 또한 높은 정밀도가 요구되는 경우 LDM의 사용이 문제가 될 수 있다.