[Paper Review] All are Worth Words: A ViT Backbone for Diffusion Models

Paper Review

제목 : All are Worth Words: A ViT Backbone for Diffusion Models
저널 : CVPR 2023
저자 : Fan Bao et al.
일시 : 25 Mar 2023
연구 : Dept. of Comp. Sci. & Tech., Institute for AI, BNRist Center

 

해당 논문 리뷰는 위 링크의 논문을 리뷰한 것입니다.

 

오류나 내용 정정 요청에 대한 문의는 언제나 환영입니다.

 

본 논문 리뷰를 읽기 전 꼭 아래 background를 먼저 숙지하시기 바랍니다.

 

 

Background

Diffusion models은 최근 수많은 연구로 엄청난 성능을 자랑하고 있습니다.

그리고 ViT 기술을 경량화하여 나온 Apple의 FastViT 등,

NLP 분야에서 성공한 Transformer 구조를 Computer Vision 분야에 적용하여 큰 인기를 끌었습니다.

한편, Diffusion models의 reverse process에서 항상 log-likelihood estimation에 있어서 CNN base의 U-Net 구조만 사용해 왔습니다.

이것에 더 특화되어 해당 논문에서는 ViT를 Diffusion models의 reverse process에 적용하는 방법을 제시합니다.

또한, inductive bias가 부족한 ViT를 어떻게 작은 양의 데이터에 적용하여 좋은 성능을 내는지 (FID 2.29 SOTA performance) 살펴봅시다.

 

 

 

1. Vision Transformers (ViT): Vision Transformers (ViT)는 최근에 주목받는 딥 러닝 아키텍처로, 원래는 자연어 처리에서 많은 성공을 거둔 Transformer 아키텍처에서 발전된 형태입니다. ViT는 이미지 처리 분야에서도 그 유용성을 입증하며 다양한 시각 작업에서 뛰어난 성과를 보이고 있습니다. 이 아키텍처의 핵심적인 아이디어는 이미지를 일정한 크기의 패치로 분할하고, 이러한 패치들을 토큰으로 간주하여 Transformer 아키텍처에 입력으로 사용하는 것입니다.

 

2. U-Net & CNN: U-Net은 이미지 세그멘테이션과 같은 작업에서 주로 사용되는, Convolutional Neural Network (CNN) 기반의 특별한 아키텍처입니다. 이 아키텍처는 크게 두 부분으로 구성되어 있습니다: 이미지를 점차적으로 축소하는 down-sampling 부분과 축소된 이미지를 다시 원래 크기로 복원하는 up-sampling 부분. 이 두 부분 사이에는 long skip connections이라는 특별한 연결 방식이 있어, 깊은 레이어의 정보를 얕은 레이어로 전달하게 됩니다.

 

3. Diffusion Models: Diffusion models은 이미지나 다른 형태의 데이터를 생성하는 데 특히 유용한 딥 러닝 모델로, 최근에 주목받기 시작했습니다. 이 모델들은 다양한 generative task, 즉 새로운 데이터 포인트를 생성하는 작업에서 높은 품질의 결과를 가져다 주며, 이미지 생성 뿐만 아니라 텍스트에서 이미지로의 변환, 비디오 생성 등 다양한 분야에서 활용되고 있습니다.

 

4. Tokens in U-ViT: 본 논문에서 제안하는 U-ViT 아키텍처는 특히 ViT의 토큰 기반 접근 방식을 활용하고 있습니다. 이 아키텍처에서는 입력 데이터(시간, 조건, 노이즈가 섞인 이미지 패치)를 모두 토큰으로 간주하게 됩니다. 이러한 토큰들은 U-Net에서 영감을 받아 얕은 레이어와 깊은 레이어 사이의 긴 skip 연결을 통해 처리됩니다.

 

5. Generative Models & FID: generative models은 데이터의 분포를 학습하여 새로운 데이터 포인트를 생성하는 데 사용되는 딥 러닝 모델의 한 유형입니다. 이러한 모델의 성능 평가에는 FID (Fréchet Inception Distance) 점수가 주로 사용됩니다. FID는 생성된 이미지의 품질과 다양성을 측정하는 메트릭으로, 낮은 FID 점수는 높은 이미지 품질을 나타냅니다.

 

 

 

Introduction

 

• Diffusion models은 high quality image generation을 성취하는 강력한 deep generative models입니다.

 

• 또한 여러 분야에서 사용될 수 있습니다.
  • text-to-image generation
  • image-to-image generation
  • video generation
  • speech synthesis
  • 3D synthesis

 

• 알고리즘의 발달에 있어, backbones은 diffusion models에서 중요한 역할을 합니다.

 

• 대표적인 예는 CNN 기반의 U-Net model입니다.

 

• 이 U-Net 구조는 diffusion models으로의 image generation tasks에서 지배적인 역할을 하고 있습니다. (log-likelihood estimation)

 

• 한편, Vision Transoformer(ViT)는 Computer Vision 분야에서 CNN-based models보다 우수한 성능을 보여주고 있씁니다.

 

이 상황에서,

 

본 논문의 저자들은 다음과 같은 질문을 던집니다.

 

 

CNN-based U-Net의 의존성이 diffusion models에서 필요할까?

 

 

 

 

여기서, 본 논문의 저자들은 간단하고 일반적인 ViT-based architecture인 U-ViT를 제안합니다.

 

---

• Transformer의 방법론에 따라, U-ViT는 모든 input에 있어 time, condition, noisy image patches를 tokens으로 처리합니다.

 

• U-ViT는 얕고 깊은 layers들 사이에서 U-Net model로부터 아이디어를 얻어 long skip connections을 적용합니다.

 

• 직관적으로, diffusion models에서 pixel-level prediction objective에는 low-level feature이 중요하게 여겨지며, 이러한 connection은 해당 prediction network의 training을 용이하게 할 수 있씁니다.

 

 

• 게다가 U-ViT는 output 생성 전에 3 x 3 convolutional block을 더하여 더 나은 visual quality을 얻어냅니다.

 

• 해당 모델을 가지고 unconditional image generation, class-conditional image generation, text-to-image generation에 사용합니다.

 

• U-ViT를 사용한 latent diffusion models의 경우 record-breaking FID scores 2.29를 class-conditional image generation에서 달성하였으며 (image size : ImageNet 256 x 256), text-to-image generation에서 MS-COCO에 대하여 5.48의 FID를 기록했다.

 

• 해당 방법들에 대해서 본 모델의 training 과정에 있어 large external dataset을 사용하지 않았다.

 

• 본 논문의 결과는 long skip connection이 중요하며, CNN 기반 U-Net의 down-sampling, up-sampling operators는 image diffusion models에 항상 필요하지 않다는 것을 제안합니다.

 

• U-ViT가 diffusion models backbone에 대한 future research에 대한 insight를 제공하고 larget scale cross-modality datasets에서의 generative modeling에 이점을 제공할 수 있을 것이라고 믿는다고 말합니다.

 

 

---

<Formula>

DDPM, conditional DDPM에 대한 기본적인 수식입니다.

 

 

Diffusion Probabilistic Models

 

forward process, diffusion process, noise-injection proecss등 불리우는 이름이 많은 noise를 더하는 과정인 forward process는 다음과 같은 Markov chain에 기반합니다.

 

 

 

$ \alpha_t , \beta_t $는 noise schedule이고 이 둘의 합은 1입니다.

 

한편, reverse process에서는,

 

위와 같은 Gaussian model이 채택되어,

 

ground truth reverse transion인 위 수식에 근사하게 됩니다.

 

 

이 과정에서 최적의 mean 값을 선정하고, 

 

결국, 위에서 $ \epsilon $은 $ x_t $가 수식에 더해진 standard Gaussian noisese가 됩니다.

 

결국, 학습은 noise prediction task와 동일하게 됩니다.

 

$ \epsilon_\theta(x_t, t) $는 $ \mathbb{E}[\epsilon | x_t] $을 학습하기 위해 채택되며,

 

아래의 noise prediction objective를 최소화합니다.

 

 

한편,

 

여기서 conditional diffusion models을 학습하기 위해서,

 

condition information은 noise prediction objective에 더 더해집니다.

 

위 수식과의 차이가 보이시나요?

 

c는 여기서 condition 혹은 그것의 continuous embedding이 됩니다.

 

이전의 image modeling tasks에서, diffusion models의 성공은 주로 CNN-based U-Net에 크게 의존하였습니다.

 

이 U-Net은 down-sampling block group, up-sampling block group 및 두 그룹 사이의 long skip connection으로 특징 짓는 convolutional backbone입니다.

 

그리고 c는 adaptive group normalization 및 cross attention 과 같은 메커니즘을 통해 U-Net에 공급됩니다.

 

 

ViT

 

Vision Transformer는 pure transformer architecture입니다.

 

이 구조는 image (words)를 tokens의 sequence로 다룹니다.

 

즉 image를 flattened patches의 sequence로 다시 정렬합니다.

 

ViT는 learnable 1D positional embedding을 이러한 patch들의 linear embedding에 더해주며, 이것은 transformer encoder의 입력으로 들어가기 전에 행해집니다.

 

ViT는 CV에서 유망하지만 이것이 diffusion-based image modeling에서 잘 사용될지 아직 확실하지 않습니다.

 

---

 

 

Method

 

U-ViT는 image generation에 있어 diffusion models을 위한 간단하고 일반적인 backbone입니다.

 

특히,

 

U-ViT는 noise prediction network $ \epsilon_\theta(x_t, t, c) $를 parameterization합니다.

 

이것은 time t, condition c, noisy image $ x_t $를 inputs로 받고 $ x_t $에 더해진 noise를 예측합니다.

 

ViT의 methodology에 의해서, image는 patches로 나뉘며,

 

U-ViT는 time, condition, image patches를 tokens으로 다룹니다.

 

U-Net 구조에서 영감을 받아, U-ViT도 얕고 깊은 layers에서 long skip connection을 사용합니다.

 

직관적으로 위 Eq. (1)의 목적에 있어 pixel-level prediction task는 low-level features에 민갑합니다.

 

long skip connections은 low-level features에 지름길을 제공하며, noise prediction network의 훈련을 쉽게 합니다.

 

또한,

 

U-ViT는 선택적으로 3 x 3 convolutional block을 output 전에 추가합니다.

 

이것은 transformer에 의해 생성된 image의 잠재적 인공물을 막는 경향이 있습니다.

 

이 block은 visual quality를 향상시킵니다.

 

 

Implementation Details

 

본 내용에서는 Ablation study와 같이 각 방식에 대해서 가장 성능이 잘 나오는 것을 선택한 후, figure 2.에 비교하여 넣었다.

 

각 내용은 다음과 같다.

 

 

(a) Combine the long skip branch

 

(b) Feed time into the network

 

(c) Add an extra convolutional block

 

(d) Variants of patch embedding

 

(e) Variants of position embedding

 

 

 

(a) Combine the long skip branch

 

 

 

위 수식에 대해서,

 

각각을 main branch와 long skip branch로부터의 embeddings이라고 하자.

 

위 방식에 대해서는 

 

이 가장 성능이 좋았다.

 

즉, "concatenationg them and then performing a linear projection"

 

바로 더하는 것은 benefits이 없다.

 

transformer bloack은 skip connections을 이미 내부에 adding operator로 가지고 있기에,

 

$ h_m $은 이미 $ h_s$의 정보를 linear form에 포함하고 있다.

 

결과적으로 $ h_m + h_s $는 linear form에 있어 $ h_s $의 coefficient를 증가시킨다.

 

 

 

(b) Feed time into the network

 

 

1.  time t는 token으로 network에 주어진다.

 

2. 두 번째로 고려한 방법은 AdaLN을 지나는 방법인데, 즉, U-Net의 adaptive group normalization을 하는 방법이다.

 

그러나 time t를 token으로 주는 방법이 더 성능이 좋다.

 

 

(c) Add an extra convolutional block

 

1.  image patches에 대한 token embeddings을 mapping하는 Linear projection 이후에 3 x 3 convolutional block을 더하는 방법이고,

 

2.  linear projection 전에 더하는 것이다.

 

3.  extra convolutional block을 사용하지 않는 것이다.

 

첫 번째 방법을 본 U-ViT에서는 선택했다.

 

(d) Variants of patch embedding

 

1. original patch embedding은 patch를 token embedding에 mapping하는 linear projection을 채택한다.

 

2.  3 x3 convolutional block stack 에 1 x 1 convolutional block을 쌓는 구조이다.

 

본 논문에서는 original patch emebdding이 더 좋은 성능을 낸다고 한다.

 



 

(e) Variants of position embedding

 

2가지를 고려한다.

 

1. 1-dimensional learnable position embedding

 

2. 2-dimensional sinusoidal position embedding

 

1번 방법이 더 좋은 성능을 낸다.

 

저자들은 다른 position embedding도 시도해보았지만 모델이 generation에 있어 meaningfule images를 만들 수 없다고 말한다.

 

 

 

위 내용에대 한 Ablation design choices는 다음과 같다.

 

 

Effect of Depth, Width and Patch size

위 내용들에 대해서도 조사했습니다.

 

 

depth - i.e., the number of layers

 

width - i.e., the hidden size D

 

patch size

 

아래 figure 3에 보이듯, depth는 9부터 13까지는 향상되다가 그 이후로는 향상되지 않습니다.

 

width는 256부터 512은 향상되다가 그 이후로는 향상되지 않습니다.

 

patch size는 8부터 2까지는 향상되다가, 1은 향상되지 않습니다.

 

small patch size는 good performance를 냅니다.

 

이것은 아마 classification과 같은 high-level tasks와는 다르게,

 

diffusion models의 Noise prediction task는 low-level 이며, small patches를 필요로 하기 때문입니다.

 

여기서 high resolution images에 대해서 small patch size를 사용하는 것은 비용이 많이 들기 때문에,

 

일단 그것을 low-dimensional latent representations으로 바꾸고 U-ViT로 latent representations을 modeling합니다.

 

 

 

 

 

Experiments

 

본 논문의 저자들은 U-ViT를 unconditional and class-conditional image generation, text-to-image generation에 본 모델을 시험합니다.

 

 

Datasets

 

unconditional learning:

CIFAR10: 50K의 훈련 이미지

CelebA 64x64: 162,770의 인간 얼굴 훈련 이미지

 

class-conditional learning:

ImageNet: 64x64, 256x256 및 512x512 resolution, 1K 다른 클래스에서 1,281,167 훈련 이미지 포함

 

text-to-image learning:

 

MS-COCO 256x256: 82,783 훈련 이미지 및 40,504 검증 이미지. 각 이미지에는 5개의 캡션이 있음.

 

high resolution image generation:

 

- latent diffusion models(LDM) 를 따라 256x256 및 512x512 해상도의 이미지 사용

- 먼저 Stable Diffusion2[54]에서 제공하는 사전 훈련된 이미지 오토인코더를 사용하여 32x32 및 64x64 해상도의 잠재 표현으로 변환

- 제안된 U-ViT를 사용하여 이러한 latent representation을 모델링

 

text-to-image learning:

 

- MS-COCO에서는 Stable Diffusion을 따라 CLIP 텍스트 인코더를 사용하여 이산 텍스트를 임베딩 시퀀스로 변환

- 이러한 임베딩은 토큰의 시퀀스로 U-ViT에 공급됨

 

U-ViT setting:

 

Table 2에서 U-ViT의 여러 설정을 확인

논문의 나머지 부분에서 U-ViT 설정과 입력 패치 크기를 간단히 표시 (예: U-ViT-H/2는 2x2의 입력 패치 크기로 U-ViT-Huge 설정을 의미)

 

Training:

 

모든 데이터셋에 대해 weight decay가 0.3인 AdamW optimizer 사용

 

대부분의 데이터셋에 대해 2e-4의 학습률 사용, ImageNet 64x64에서는 3e-4 사용

 

CIFAR10 및 CelebA 64x64에서 500K 반복으로 128의 배치 크기로 훈련.

 

ImageNet 64x64 및 ImageNet 256x256에서 300K 반복으로 1024의 배치 크기로 훈련. ImageNet 512x512에서 500K 반복으로 1024의 배치 크기로 훈련

 

MS-COCO에서 1M 반복으로 256의 배치 크기로 훈련

 

ImageNet 256x256, ImageNet 512x512 및 MS-COCO에서는 classifier-free guidance 채택

 

 

 

5.2. Unconditional and Class-Conditional Image Generation

• U-ViT를 기존 U-Net 기반의 확산 모델과 비교함.
• 더 작은 ViT인 GenViT와도 비교. GenViT는 긴 스킵 연결을 사용하지 않음.
• 이미지 품질을 측정하기 위해 50K 생성 샘플의 FID 점수를 보고함.
• U-ViT는 무조건부 CIFAR10 및 CelebA 64×64에서 U-Net과 비교할 수 있으며, GenViT보다 훨씬 나음.
• 클래스-조건부 ImageNet 64×64에서 U-ViT는 더 나은 FID를 얻고 성능을 향상시키기 위해 다른 구성을 사용함.
• U-ViT는 latent sapce 에서 특히 잘 수행되며, 클래스-조건부 ImageNet 256×256에서 최첨단 FID를 얻음.
• 여러 샘플링 단계에서의 성능도 나타냄. U-ViT는 U-Net보다 더 나은 성능을 보임.

5.3. Text-to-Image Generation on MS-COCO

• 표준 벤치마크 데이터셋인 MS-COCO에서 텍스트-이미지 생성을 위해 U-ViT 평가.
• 잠재 공간의 이미지에 대한 U-ViT를 훈련시키며, U-ViT-S 크기와 비교할 수 있는 U-Net도 훈련시킴.
• FID 점수를 사용하여 이미지 품질 측정.
• U-ViT-S는 이미 훈련 중에 큰 외부 데이터셋을 액세스하지 않는 방법 중에서 최첨단 FID를 달성함.
• U-ViT와 U-Net을 동일한 무작위 시드를 사용하여 비교한 결과, U-ViT가 더 높은 품질의 샘플을 생성하고 텍스트와 더 잘 일치함.

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

< 정리 >

CNN base의 U-Net 모델을 이용한 Diffusion models이 지배적인 영향을 주는 와중에,

Vision transformers(ViT) 다양한 Computer vision 분야에서 유망함을 보이고 있습니다.

이때,

diffusion models을 사용한 image generation에 간단하고 일밙벅은 ViT-based architecture인 U-ViT를 본 논문에서 제안합니다.

U-ViT의 input은 time, condition, noisy image patches의 tokens로 이루어져 있고, 얕고 깊은 long skip connection을 사용합니다.

본 논문애서는 unconditional, class-conditional image generation, text-to-image generation tasks에 대해서,

U-ViT는 유사한 크기의 CNN 기반 U-Net에 비해 비슷하거나 우월한 성능을 보입니다.

U-ViT 기반의 latent diffusion models은 record-breaking FID scores 2.29를 성취합니다.

이러한 여러 방법들은 generative models의 training 중 large external dataset을 사용하지 않고 이뤄냅니다.

 

 

 

Contirbution은 다음과 같이 정리됩니다.

 

Contirbution

• simple and general ViT-based architecture for image generation with diffusion models
• U-ViT treats all inputs including the time, condition and noisy image patches as tokens and employs long skip connections between shallow and deep layers.
• long skip connection is crucial while the down-sampling and up-sampling operators in CNN-based U-Net are not always necessary
  
• U-ViT can provide insights for future research on backbones

 

 

 

---

< 본 논문에 대한 본 포스팅 저자의 생각 >

 

해당 내용은 저의 개인적인 의견이니 참고용으로만 바라봐주세요 🌈

 

Diffusion model은 굉장히 많은 분야에 좋은 성능으로 사용되고 있다.

CV분야에 큰 파장을 일으킬만한 이 시점에 쏟아지는 최신 논문들을 읽고 있자니 감개 무량할 정도이다.

내가 이 game-changer와 같은 models이 나온 지 불과 3년 만에 알게되다니 말이다. (Since 2020 DDPM)

본 모델이 열역학 개념에서 파생된 만큼,

해당 열역학 법칙을 지키며 혹은 깨부시며 diffusion model의 정통을 이어가려는 노력들이 많이 보이며 그에 따른 변형 모델도 많이 출시되고 있는 추세이다.

그러한 관점에서,

본 논문에서 시사하는 바는 굉장히 크다.

Diffusion models에서의 log-likelihood estimation에서 reverse process의 다음 step을 예측한다는 측면에서 너무나 당연하게도 사용되던 U-Net 구조를 최신 핫한 ViT로 바꾸어 시도해봤다는 건 엄청난 시도라고 생각된다.

또한, 그러한 구조를 실용화 하는 데 있어 해당 내용이 future research에 backbone이 된다고 언급한 것 또한 큰 의미를 가진다고 생각한다.

한편, Vision transformer가 가지고 있는 약점인 inductive bias를 어떻게 해결했을지를 유심히 보았다.

 

 

 

Reference

 

https://arxiv.org/abs/2209.12152

All are Worth Words: A ViT Backbone for Diffusion Models

 

 

https://github.com/baofff/U-ViT

GitHub - baofff/U-ViT: A PyTorch implementation of the paper "All are Worth Words: A ViT Backbone for Diffusion Models".