[평범한 학부생이 하는 논문 리뷰] Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding (Imagen)

1. Introduction

 본 논문은 Imagen을 도입하는데 이는 text-to-image 합성에서 전례없는 정도의 photorealism과 깊은 수준의 언어 이해를 가져오기 위해 transformer language models와 high-fidelity diffusion model을 결합한 text-to-image diffusion model이다. Imagen의 key finding은 text-only corpora로 기학습된 large LM으로부터 text embedding이 text-to-image 합성에서 놀라운 효과적이라는 것이다.

 Imagen은 input text를 sequence of embeddings로 매핑하기 위한 frozen T5-XXL encoder와 $64 \times 64$ image diffusion model과 그 뒤로 $256 \times 256$과 $ 1024 \times 1024$ 이미지를 생성하는 두 개의 super-resolution diffusion model로 구성되어 있다. 모든 diffusion model은 text embedding sequence로 condition되고 classifier-free guidance를 사용한다. Imagen은 이전 작업들에서 보이는 샘플 퀄리티 감소 없이 거대한 guidance weight 사용을 허용하는 새로운 샘플링 기법에 의존한다. 그 결과 higher fidelity와 better image-text alignment를 보인다.

 

 컨셉적으로 간단하고 학습하기 쉬운데 결과는 좋다고 한다. 자세한 결과들은 뒤에서 볼 예정이다. 본 논문에서는 DrawBench라는 새로운 benchmark를 도입했다.

본 논문의 contribution을 요약하면 다음과 같다.

1. 텍스트 데이터로만 학습된 large frozen language models는 text-to-image generation을 위한 매우 효과적인 text encoder라는 것을 보였고, diffusion models의 크기를 키우는 것 보다 frozen text encoder의 크기를 키우는 것이 샘플 퀄리티를 향상시켰다.
2. High guidance weights를 활용하는 새로운 diffusion sampling technique인 dynamic thresholding을 도입했고 이를 통해 더 photorealistic하고 디테일한 이미지를 생성했다.
3. 본 논문은 몇가지 중요한 diffusion 아키텍처 디자인 선택을 강조하고 더 빨리 수렴하고  더 메모리 효율적인 Efficient U-Net을 제안한다.
4. 7.27의 새로운 sota COCO FID를 달성했다. Human rater들은 이미지와 텍스트의 alignment 측면에서 Imagen이 참조 이미지와 동등한 수준이라고 평가했다.
5. 새로운 평가 benchmark인 DrawBench를 도입했다. 이를 통한 평가에서 다른 연구를 outperform했다.

2. Imagen

 Introduction에서 설명했듯이 Imagen은  text를 sequence of embeddings에 맵핑하는 text encoder와 증가하는 해상도의 이미지들을 이 embeddings에 맵핑하는 conditional diffusion models의 cascade로 이루어진다.

2.1 Pretrained text encoders

 Text-to-image model(CLIP)들은 임의의 자연어 text input의 complexity와 compositionality를 캡처하기 위해 강력한 semantic text encoders를 필요로 한다. Text encoder들은 image-text pair 데이터로 학습되고 이때 당시 text-to-image 모델들의 스탠다드였다. Image-text 학습 objective들은 이러한 text encoder들이 특히 text-to-image generation 테스크와 관련된 시각적으로 meaningful하고 semantic한 표현을 인코드할 것이라고 제안한다. Large Language Model(BERT,GPT,T5)들 또한 text-to-image generation을 위한 text를 인코드하기 위한 또다른 선택이 될 수 있다. 언어 모델은 text-image pair 데이터보다 큰 text only corpus로 학습되고 그렇기에 매우 풍부하고 넓은 text 분포에 노출된다. 이러한 모델들은 image-text model들의 text encoder보다 훨씬 크다.

 

 그러므로 text-to-image 테스크를 위한 두 text encoders군에 대해 모두 탐구할 필요가 있다. 그래서 Imagen은 다음과 같은 기학습된 text encoder에 대해 연구한다 : BERT, T5, CLIP. 유사성을 위해 본 논문에서는 이러한 text encoder들의 weight를 freeze한다.  Freezing은 embedding의 오프라인 계산과 같은 몇가지 이점을 가진다. 그리고 이는 text-to-image 모델의 학습동안 무시할 만한 계산과 메모리 footprint를 야기한다. 본 논문에서 text encode 크기를 scaling하는 것이 text-to-image generation의 퀄리티를 향상시킨다는 것을 보인다. 게다가 MS-COCO와 같은 간단한 benchmark에서는 T5-XXL과 CLIP이 유사한 성능을 내지만 사람 평가에서 DrawBench에 대해서는 T5-XXL의 결과를 더 선호한다고 한다.

2.2 Diffusion models and classifier-free guidance

Diffusion models 포스팅

https://juniboy97.tistory.com/45

[평범한 학부생이 하는 논문 리뷰] Generative Modeling by Estimating Gradients of the Data Distribution (NCSN)

https://juniboy97.tistory.com/46

[평범한 학부생이 하는 논문 리뷰] Denoising Diffusion Probabilistic Models (DDPM)

Classifier guidance 포스팅

https://juniboy97.tistory.com/54

[평범한 학부생이 하는 논문 리뷰] Diffusion Models Beat GANs on Image Synthesis

Classifier free guidance 포스팅

https://juniboy97.tistory.com/55

[평범한 학부생이 하는 논문 리뷰] Classifier-Free Diffusion Guidance

Imagen은 효과적인 text embedding을 위해 classifier-free guidance에 의존한다.

2.3 Large guidance weight samplers

 Classifier-free guidance weight를 증가시키면 image-text alignment는 증가하지만 매우 saturated되거나 unnatural 이미지들을 생성하는 등의 image fidelity에 데미지를 가한다. 본 논문에서는 이를 high guidance weights로부터 발생한 train-test mismatch때문에 발생한다는 것을 발견했다. 각 sampling step $t$에서 , x-prediction $\hat{x}^t_0$는 training data $x$와 동일한 bound내에 있어야한다. 이 bound는 $[-1,1]$이다. 하지만 경험적으로 high guidance weights는 x-prediction이 이러한 bound를 초과하는 것을 야기한다는 것을 발견했다. Diffusion model은 sampling을 통해 자기 자신의 output이 반복적으로 적용되기 때문에, sampling 프로세스는 unnatural 이미지를 생성하고 종종 발산한다. 이는 train-test mismatch이다. 이 문제를 해결하기 위해서 본 논문은 static thresholding과 dynamic thresholding을 연구했다.

Static thresholding

 Static thresholding은 x-prediction을 [-1,1]로 elementwise clipping하는 것이다. 본 논문에서는 static thresholding이 large guidance weights를 통한 sampling에서 매우 필수적이고 blank image의 생성을 막는다고 한다. 그럼에도 불구하고, static thresholding은 여전히 guidance weight를 증가시키면 over-saturated하고 less detailed한 이미지를 생성한다.

Dynamic thresholding

 Dynamic thresholding은 각 샘플링 단계에서 $s$를 $\hat{x}^ t_0$의 특정 백분위수 절대 픽셀 값으로 설정하고, $s$가 1보다 크면 $\hat{x} ^t_ 0$을 $[-s, s]$ 범위로 임계값을 설정한 다음 $s$로 나누는 방식이다. Dynamic thresholding은 포화 픽셀(-1과 1에 가까운 픽셀)을 안쪽으로 밀어내어 각 단계에서 픽셀이 포화되는 것을 적극적으로 방지한다. 본 논문에서는 dynamic thresholding이 특히 매우 큰 guidance weights를 사용할 때, 더 좋은 image-text alignment 뿐만 아니라 더 좋은 photorealism을 야기한다는 것을 발견했다.

2.4 Robust cascaded diffusion models

 Imagen은 base $64 \times 64$ 모델의 pipeline과 $64 \times 64$의 생성된 이미지를 $256 \times 256$과 그 후 $1024 \times 1024$로 upsample하기 위해 두 개의 text-conditional super-resolution diffusion models를 이용한다. 점진적으로 high-fidelity 이미지들을 생성하는데 noise conditioning augmentation을 가진 cascaded diffusion models는 매우 효과적이어왔다. 게다가 noise level conditioning을 통해 super-resolution 모델을 추가된 noise의 양을 알도록 만드는 것은 샘플 퀄리티를 향상시키고 저해상도 모델에 의해 생성된 artifacts를 다루기 위해 super-resolution 모델들의 robustness를 향상시키는 것을 돕는다. Imagen은 두 super-resolution 모델들에 대해 noise conditioning augmentation을 사용한다. 이는 높은 fidelity 이미지 생성에 중요하다.

2.5 Neural network architecture

Base model

 본 논문에서는 $64 \times 64$ text-to-image diffusion model에 대해 improved DDPM의 U-Net 구조를 적용한다. 이 네트워크는 pooled embedding vector를 통해 text embedding에 대해 condition되고 이 pooled embedding vector는 classifer guidance와 cascaded diffusion models에서 사용된 class embedding conditioning 방식과 유사하게 diffusion timestep embedding에 더해진다. 더 나아가 multiple resolution에 대한 text embedding에 대해 cross attention을 추가함으로써 text embedding의 전체 시퀀스에 대해 condition을 준다. 게다가 atttention과 pooling 레이어에서 text embeddings를 위한 Layer Normalization이 performance를 향상시키는데 상당한 도움이 된다고 한다.

Super-resolution models

 $64 \times 64 \rightarrow 256 \times 256$ super-resolution에 대해 본 논문은 메모리 효율, 추론 속도, 수렴  속도를 향상시킨 Efficient U-Net을 사용한다.

 $256 \times 256 \rightarrow 1024 \times 1024$ super-resolution 모델은 $1024 \times 1024$이미지의 $64 \times 64 \rightarrow 256 \times 256$ crops에 대해 학습이 된다.($256 \times 256$으로 이미지를 쪼개서 그걸 $64 \times 64 \rightarrow 256 \times 256$으로 학습시킨다는 말인가...  코드를 찾아봐도 모르겠다...) 이를 유용하게 하기 위해, self-attention layer를 제거하고 text cross-attention은 유지한다.

 두 super-resolution 모델 모두에서 text cross attention을 사용한다.

Efficient U-Net 구조

3. Evaluating Text-to-Image Models

 Metric으로는 image fidelity를 평가하기 위한 FID score,와 image-text alignment를 평가하기 위한 CLIP score를 사용한다. Guidance weight가 image quality와 text alignment를 컨트롤하는 중요한 재료이기 때문에, 본 논문은 guidance weights의 범위에 따른 CLIP과 FID score에 대한 trade-off curve를 보고한다.

 FID는 perceptual quality를 전부 align하지 않고 CLIP은 counting하기에 비효율적이다라는 한계를 가지기에 human evaluation을 이용한다.

DrawBench

 COCO는 모델들 사이의 차이에 대한 insight를 손쉽게 제공하는 프롬프트의 제한된 스펙트럼을 가진다. 이렇기에 본 논문은 DrawBench를 도입한다. 이는 text-to-image 모델들의 평과와 비교를 지원하는 comprehensive하고 challenging한 프롬프트의 set이다. DrawBench는 다른 색, 객체의 수, 공간적인 관계, 장면에 있는 text, 객체간의 흔치않은 상관관계를 충실히 render하는 능력과 같은 모델들의 다른 능력을 시험하는 11개의 카테고리의 프롬프트를 포함한다. 카테고리들은 또한 길고 복잡한 textual 설명, 드문 단어들. misspelled 프롬프트를 포함한 복잡한 프롬프트를 포함한다.

4. Experiments

4.1 Results on COCO

 Imagen은 COCO에 대한 zero-shot FIDscore를 7.27로 달성함으로써 기존의 COCO로 학습된 DALL-E2를 능가한다. 표 2는 Imagen에 대한 human evaluation을 나타낸다.

4.2 Results on DrawBench

 Image-fidelity와 image-text alignment 둘 다에서 human rater들은 Imagen의 결과를 선호하는 것을 볼 수 있다.

4.3 Analysis of Imagen

• Text encoder의 크기를 scaling하는 것은 매우 효과적이다.
• Text encoder의 크기를 scaling하는 것은 U-Net 크기를 scaling하는 것보다 더 중요하다.
• Dynamic thresholding은 중요하다.
• Human raters는 DrawBench에 대해 CLIP보다 T5-XXL을 더 선호한다.
• 큰 guidance weights를 사용하고 inference 시 저해상도 이미지에 noise를 추가하는 것(noise conditioning augmentation)은 super-resolution모델이 저해상도 이미지의 artifacts를 제거하면서 다양한 upsampled output을 생성하게 한다.
• Cross attention으로 text embedding의 sequence에 대해 conditioning하는 것은 mean or attention based pooling보다 좋은 성능을 보인다.
• Efficient U-Net은 메모리를 덜 쓰고, 빠르게 수렴하고, 더 빠른 추론으로 더 좋은 샘플 퀄리티를 가진다.

6. Conclusions

 Imagen은 diffusion model을 사용한 text-to-image 생성을 위한 text encoder로 frozen된 사전학습된 large language model의 효과를 보여준다. 이러한 언어 모델을 scaling하는 것이 U-Net을 scaling하는 것보다 더 효과적이라는 것 또한 보인다. 본 논문은 Imagen을 통해 classifier-free guidance의 중요성을 다시 강조하고 훨씬 높은 guidance weight의 사용을 가능하게 하는 dynamic thresholding을 도입한다.