[Paper Reivew] Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding (Imagen)

NeurIPS 2022. [Paper] [Github] [Demo] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S. Sara Mahdavi, Rapha Gontijo Lopes, Tim Salimans, Jonathan Ho, David J Fleet, Mohammad Norouzi
Google Research | Brain Team 23 May 2022

TL;DR

Text-only pretrained large language model(e.g. T5-XXL)의 텍스트 임베딩을 활용해, 이미지 합성에서 fidelity와 image-text alignment를 크게 향상시킨 T2I model Imagen을 제시한 논문입니다.

또한 DranBench라는 text-to-image model을 위한 새로운 benchmark를 제안합니다.

1. Introduction

저자들은 Imagen이라는 Text-to-Image 생성 모델을 제안합니다. Text-image data만 사용한 이전 연구들과는 다르게, text-only corpora로 pretraining한 large LM의 text embeeding 능력을 이용해 언어 이해능력을 크게 높일 수 있었다고 합니다.

Imagen은 frozen T5-XXL 언어 모델을 이용해 입력 텍스트를 텍스트 embedding으로 변환하고, 64x64크기의 diffusion model과 256×256, 1024×1024 크기의 super-resolution 모델로 구성되어 있습니다. CFG, 이전포스트기법을 적용하며, 추가적으로 저자들은 Imagen에서 새로운 샘플링 방법(Dynamic Thresholding)을 제시하여 guidance scale을 키울 때 생기는 품질 저하 문제를 해결했다고 합니다.

Imagen은 COCO 데이터셋에서 기존 모델보다 훨씬 더 우수한 성능을 보입니다. Imagen의 FID-30K 점수는 7.27로, GLIDE의 12.4나 DALL-E 2의 10.4보다 크게 앞서며, COCO 데이터로 학습하지 않았음에도 불구하고 COCO 데이터로 학습한 다른 모델들(예: Make-A-Scene의 7.6)보다 더 나은 성능을 기록했습니다.

또한 저자들은 새로운 benchmark DrawBench를 도입하여, Text-to-Image의 새로운 평가 지표를 제안합니다. 이는 모델의 구문성, 공간 관계, 복잡한 텍스트 프롬프트 처리 능력을 평가할 수 있다고 합니다.

저자들의 주요 기여를 정리하면 다음과 같습니다.

1. Text-only data를 이용해 pretraining한 frozen large 언어 모델(T5 등)의 텍스트 인코딩이 매우 효과적임을 발견했습니다. 텍스트 인코더의 크기를 키우는 것이 이미지 생성 모델의 크기를 키우는 것보다 샘플 품질을 더 크게 향상시킵니다.

2. Dynamic Thresholding이라는 새로운 확산 샘플링 기법을 도입하여, 더 높은 가이드 가중치를 사용하더라도 샘플 품질이 저하되지 않고 더 사실적인 이미지를 생성할 수 있게 했습니다.
3. 효율적인 U-Net 아키텍처(Efficient U-Net)를 제안하여, 모델의 메모리 효율성을 높이고 빠르게 수렴하게 했습니다.
4. COCO 데이터셋에서 7.27의 새로운 최첨단 FID 점수를 기록했으며, 인간 평가에서 참조 이미지와 동등한 수준의 이미지-텍스트 정합도를 달성했습니다.
5. DrawBench라는 포괄적이고 도전적인 평가 벤치마크를 도입해, Imagen이 다른 최신 텍스트-이미지 생성 모델들보다 뛰어나다는 사실을 입증했습니다.

2. Imagen

2.1. Pretrained text encoders

T2I 모델은 자연어 텍스트 입력의 복잡성과 구문성을 처리할 수 있는 강력한 의미적 텍스트 인코더가 필요합니다.

• 기존 standard : Image-Text pair로 텍스트 인코더를 학습하거나, CLIP등의 pretrained model 사용.
• Imagen: 더 넓고 다양한 텍스트 분포를 학습한 pretrained대규모 언어 모델(예: BERT, GPT, T5)를 freezing하여 사용(더 적은 연산 비용, 메모리 사용의 이점). 또한 텍스트 인코더의 크기를 키우면 T2I 성능이 좋아지는 명백한 경향을 확인.

저자들의 실험에서 T5-XXL과 CLIP 모두 MS-COCO 벤치마크에서는 유사한 성능을 보였으나, DrawBench라는 더 복잡한 프롬프트 모음에서 T5-XXL이 CLIP보다 텍스트-이미지 정합도와 이미지 품질 측면에서 더 우수하다는 인간 평가 결과를 얻었습니다.

2.2. Diffusion models and classifier-free guidance

Diffusion Basic, 이전 포스트 와 CFG, 이전포스트를 참고해주세요!

확산 모델(diffusion models)은 가우시안 노이즈를 학습된 데이터 분포의 샘플로 변환하는 반복적인 노이즈 제거 과정을 거칩니다. 이 모델은 클래스 레이블이나 텍스트 등 다양한 조건에 따라 작동할 수 있습니다. Imagen의 확산 모델은 denoising objective를 사용해 노이즈 제거 학습을 진행하며, 손실 함수는 다음과 같이 정의됩니다:

\[\begin{equation} \mathbb{E}_{x, c, \epsilon, t} [ w_t \| \hat{x}_\theta (\alpha_t x + \sigma_t \epsilon, c) - x \|_2^2 ] \\ t \sim \mathcal{U} ([0, 1]), \quad \epsilon \sim \mathcal{N} (0, I) \end{equation}\]

여기서 \(x\)는 데이터, \(c\)는 조건, \(\epsilon\)은 가우시안 노이즈, \(t\)는 노이즈 강도를 결정하는 매개변수입니다.

classifier guidance(CG, 이전포스트) 는 사전 학습된 calssifier 모델의 gradient를 사용해 샘플 품질을 향상시키는 방법이지만, classifier-free guidance(CFG, 이전포스트) 는 사전 학습된 모델 없이 모델을 조건부 및 비조건부 목표로 동시에 학습하는 방법입니다. CFG는 conditional, unconditional 예측을 혼합하여 샘플링 중 가중치 \(w\)로 조절된 결과를 생성합니다. \(w\)가 클수록 텍스트 조건화가 강화되는 경향이 있습니다.

2.3. Large guidance weight samplers

최근 연구에서 CFG의 guidance 가중치 \(w\)를 높이면 이미지-텍스트 alignment가 개선되는 경향을 확인했지만, 이로인해 이미지 품질이 저하되거나, 부자연스러운 이미지가 생성되는 문제가 있습니다. 저자들은 이런 문제의 원인을 높은 guidance가 train-test mismatch를 유발하기 때문이라고 주장합니다.

샘플링 단계에서, \(\hat {x}_{t0}\)​는 학습 데이터의 범위 \([-1, 1]\) 안에 있어야 하지만, 높은 가이드 가중치로 인해 \(\hat{x}_{t0}\)​가 이 범위를 초과하는 것이 관찰되었다고 합니다. 이는 모델이 반복적으로 자체 출력을 사용하여 샘플링하는 과정에서 unnatural image를 생성하거나 심지어는 발산하게 된다고 합니다.

이런 문제를 해결하기 위해서 저자들은 Static thresholding, Dynamic thresholding 방법을 실험했습니다.

• Static thresholding \(\mathbf x\)의 예측값을 \([-1, 1]\)로 clipping하는 방법입니다. 이전에도 사용된 방법이지만 그 중요성은 크게 강조되지 않았다고 합니다. 이 방법이 효과가 있기는 하지만, 가중치가 커질수록 over-saturated 혹은 less detailed image의 문제가 여전히 발생합니다.

• Dynamic thresholding 저자들은 새로운 Dynamic thresholding 방법론을 제시합니다. 각 샘플링 단계에서 \(s\)라는 값을 \(\hat {x}_{t0}\)​ 픽셀 값의 특정 percentage로 설정하고, 만약 이 값이 1보다 크다면, \(\hat {x}_{t0}\)​를 \([-s, s]\)사이로 clipping한 뒤, 다시 s로 나누어 주는 방법입니다. 이 방법은 saturated 픽셀을 내부로 밀어 넣어 픽셀이 포화 상태에 빠지지 않도록 적극적으로 방지합니다. 저자들은 동적 임계값이 매우 큰 가이드 가중치를 사용할 때 photorealism과 image-text alignment이 크게 개선된다는 것을 확인했습니다.

2.4. Robust cascaded diffusion models

Imagen은 64×64 크기의 기본 diffusion 모델과, 이를 점진적으로 업샘플링하여 256×256과 1024×1024 해상도의 이미지를 생성하는 text-conditional super-resolution diffusion 모델 두 개로 구성된 파이프라인을 사용합니다. 이런 cascaded diffusion 모델을 noise conditioning augmentation와 함께 사용하면 이미지 품질이 매우 좋아짐이 이전 연구에서 밝혀졌습니다. 이런 방식으로 추가된 노이즈의 크기를 super-resolution 모델에 전달함으로써, 저해상도에서 발생하는 aritifac문제를 해결하고 robust한 모델을 만들 수 있었다고 합니다.

주어진 저해상도 이미지와 augmentation level(aug_level)에 따라, 저해상도 이미지를 해당 증가 수준에 맞게 노이즈나 블러를 사용해 왜곡하고, diffusion 모델이 이를 aug_level에 맞춰 처리하도록 조건화합니다. training에는 aug_level을 무작위로 선택하며, inference에는 다양한 값을 적용해 가장 좋은 샘플 품질을 찾습니다.

2.5. Neural network architecture

• Base model 저자들은 Imagen의 기본 구조로 64×64 T2I diffusion model의 구조로 U-Net구조를 사용했다고 합니다. 이 네트워크는 diffusion timestep embedding에 text embedding을 추가하는 방법으로 conditioning 된다고 합니다. 또한 text embedding의 전체 시퀀스에대해 cross attention을 적용해 추가적인 conditioning을 했다고 합니다.

• Super-resolution models

  (1) \(64 \times 64 \rightarrow 256 \times 256\) 해상도로 업샘플링하는 모델에도 U-Net 아키텍처를 적용하되, 기존 모델에서 메모리 효율성, 추론 속도, 수렴 속도를 개선하는 여러 가지 수정 사항을 적용했습니다. 이 수정된 모델을 Efficient U-Net이라고 부릅니다.

  (2) \(256 \times 256 \rightarrow 1024 \times 1024\) 해상도로 업샘플링하는 모델은 \(1024 \times 1024\) 이미지의 \(64 \times 64 \rightarrow 256 \times 256\) crop을 이용해 학습합니다. 이를 위해 self-attention 레이어를 제거하였지만, 텍스트 cross attention 레이어는 유지했다고 합니다.

3. Evaluating Text-to-Image Models

COCO validation dataset은 T2I 모델을 평가하는 표준 벤치마크로, supervised 및 zero-shot 모두에서 사용됩니다. 주요 성능 지표로는 FID score 가 사용되며, 이는 이미지 품질을 측정합니다. CLIP score는 이미지-텍스트 alignment를 측정합니다.

저자들은 대부분의 이전 연구와 일치하게 FID-30K를 사용했습니다. 검증 세트에서 가이드 가중치(guidance weight)는 이미지 품질과 텍스트 정렬을 제어하는 중요한 요소이므로, CLIP 및 FID 점수 간의 상호 관계를 설명하는 trade-off(pareto curve)로 대부분의 실험 결과를 제시합니다.

하지만 저자들은 FID, CLIP score 모두 한계가 있으며, 이를 보완하기 위해 human evaluation을 도입했다고 합니다. (COCO dataset에서 무작위로 선택된 200개의 image-caption pair를 이용)

1. 이미지 품질 평가 : “어떤 이미지가 더 사실적(실제로 보이는가)인가?”
2. 정렬 평가 : “캡션이 위의 이미지를 정확하게 설명하는가?”

DrawBench: COCO는 유용한 벤치마크이지만, 저자들은 T2I 모델의 종합적인 언어 이해 능력을 평가하기위해 DrawBench라는 새로운 평가지표를 제안합니다. DrawBench는 11개의 카테고리로 나누어진 프롬프트로 구성되며, 색상 표현, 객체 개수, 공간적 관계, 장면 내 텍스트, 그리고 드문 상호작용과 같은 다양한 모델 기능을 테스트합니다. 또한 복잡한 텍스트 설명, 희귀 단어, 철자가 잘못된 프롬프트 등의 세트를 포함한다고 합니다.

4. Experiments

4.1. Training details

저자들은 64×64 Base diffusion 모델에는 2B parameter, 64×64 → 256×256 및 256×256 → 1024×1024 super-resolution model은 각각 600M 및 400M의 parameter 모델을 사용했습니다.

4.2. Results on COCO

Imagen은 COCO 검증 세트에서 FID score를 기준으로 평가한 결과, FID 7.27이라는 SOTA를 기록했습니다. 흥미로운 점은 DALL-E 2와 같은 최신 모델은 물론, COCO에서 직접 학습된 모델들도 능가한 결과입니다.

4.3. Results on DrawBench

저자들은 DrawBench를 사용해 Imagen과 DALL-E 2, GLIDE, Latent Diffusion, CLIP 기반 VQ-GAN 등의 최신 모델들을 비교했습니다.

4.4. Analysis of Imagen (Ablations)

1. 텍스트 인코더 크기 확장은 매우 효과적: 텍스트 인코더 크기를 확장하면 이미지-텍스트 alignmnet와 이미지 품질이 일관되게 개선되었습니다. (Figure 4a).

2. 텍스트 인코더 크기가 U-Net 크기보다 중요: U-Net 크기를 확장하면 샘플 품질이 개선되지만, 텍스트 인코더 크기를 확장하는 것이 U-Net 크기 확장보다 훨씬 더 큰 영향을 미쳤습니다 (Figure 4b).

3. Dynamic thresholding 적용이 중요: Dynamic thresholding을 적용하면 static thresholding 또는 no thresholding대비 훨씬 더 사실적이고 텍스트와 잘 맞는 샘플이 생성되었습니다. 특히, 큰 CFG guidance weight 사용 시 더욱 두드러진 효과가 나타났습니다 (그림 4c).

4. Noise conditioning augmentation이 중요: Noise conditioning augmentation을 적용해 super-resolution model을 학습하면 CLIP과 FID 점수가 향상되었습니다. 추론 시 저해상도 이미지에 노이즈를 추가하면, 큰 CFG guidance를 사용할 때 super-resolution model이 저해상도 이미지의 아티팩트를 제거하면서 다양한 업샘플링 출력을 생성할 수 있었습니다.

5. 텍스트 조건화 방법이 매우 중요: 텍스트 임베딩 시퀀스에 대한 cross attention을 사용해 조건화하면 단순한 평균 풀링이나 어텐션 기반 풀링보다 샘플 품질과 이미지-텍스트 정렬에서 훨씬 더 뛰어난 성능을 보였습니다.

6. Efficient U-Net의 중요성: Efficient U-Net은 메모리 사용량을 줄이고, 학습 속도가 더 빠르며, 샘플 품질이 개선되었고, 추론 속도도 더 빠르게 작동합니다.

5. Related Work

GLIDE와 달리, Imagen은 대형 frozen pretrained 언어 모델을 사용하며, 이는 fidelity과 텍스트-이미지 alignment에 매우 중요한 역할을 했습니다.

XMC-GAN 역시 BERT를 텍스트 인코더로 사용하지만, Imagen은 더 큰 텍스트 인코더를 사용하여 그 효과를 입증했습니다.

6. Conclusions

저자들의 Imagen은 frozen large pretrained language model이 텍스트-이미지 생성에서 강력한 텍스트 인코더로서 효과적임을 보여줍니다. 언어 모델의 크기를 확장하는 것이 U-Net 크기 확장보다 성능에 더 큰 영향을 미친다는 관찰은 향후 연구에서 더 큰 언어 모델을 텍스트 인코더로 탐색하는 방향을 제시하기도 합니다.

또한, Imagen도 다른 연구들과 마찬가지로 CFG의 중요성을 강조하고, dynamic thresholding을 도입해 더 높은 guidance weight 사용하면서도 샘플 품질을 유지하는 데 성공했습니다.