[23’ ICCV] DiT: Scalable Diffusion Models with Transformers

Date: 2024.05.06 Updated: 2024.05.06

카테고리: Generation

🔍 Abstract

Transformer 구조를 기반으로 한 diffusion model인 DiT(Diffusion Transformer)을 제안한다.
DiT는 계산량이 늘어나면 성능이 증가하는 scalable한 특징을 가진다.

1. Introduction

LDM(Latent Diffusion Model), 혹은 Stable Diffusion 논문 리뷰를 통하여 diffusion model을 latent space에서 계산하는 것의 이점을 확인할 수 있었다.

DiT(Diffusion Transformer)는 이 latent space에서의 diffusion model backbone을 일반적인 U-Net이 아닌 Transformer를 사용하여 구현한 모델이다. 단순히 U-Net을 ViT 구조로 변경한 것이지만, 이를 통해 성능이 개선되었으며 U-Net 때와는 달리 모델의 크기를 늘리거나, patch의 수를 늘리는 등 계산량을 늘리는 것이 대개 성능 향상에 도움이 되는 직관적인 특성을 가진다. 이러한 특성을 scalability라고 한다. 이제 그 구조부터 살펴보자.

2. DiT Architecture

DiT의 구조는 굉장히 직관적이다. ViT Encoder 구조를 그대로 사용한다. 눈여겨볼 부분은 timestep $t$, label $y$를 어떻게 conditioning하는지이다. 저자들은 크게 4가지 방법을 제안했다.

In-context Conditioning: timestep $t$와 label $y$를 input token에 추가한다. ViT에서 cls token을 사용하는 것과 유사하다.
Cross-attention Block: timestep $t$와 label $y$를 cross-attention layer를 추가하여 conditioning한다. LDM에서 사용한 방법이다.
adaLN Block: adaLN(Adaptive Layer Normalization)은 timestep $t$와 label $y$로부터 regression을 통해 layer의 scale 및 shift parameter $\gamma, \beta$를 얻고, 이를 통해 layer normalization을 수행하는 것이다. CG(Classifier Guidance) 논문에서 사용한 adaGN과 유사하다.
adaLN-Zero Block: 기존 adaLN에 residual connection 바로 직전에 scale parameter $\alpha$를 추가하였다. 그리고 모든 $\alpha$에 대하여 MLP vector 값이 0이 될 수 있도록 MLP를 초기화한다. 이렇게 설정하면 처음에는 DiT Block이 Identity Block과 같은 역할을 하게 되고, 이것이 학습에 도움을 준다고 한다.

이러한 4가지 conditioning strategy 중 adaLN-Zero Block이 가장 성능이 좋았다. ViT 기반의 encoder에 가장 적절한 conditioning 방법을 찾아내는 과정이라고 이해하면 되겠다.

3. Experiments

먼저 SOTA와 비교했을 때에도 성능이 향상된 모습을 보이고, 계산량이 늘어나면 성능이 향상되는 모습을 보인다. 이때 계산량을 늘리는 방법은 크게 2가지인데, 첫 번째는 모델의 크기를 키우는 것이고, 두 번째는 patch의 크기를 줄이는 것이다. Patch의 크기를 줄이면 그만큼 patch의 수가 늘어나기 때문에 계산량이 늘어나게 된다. 두 경우 모두가 Figure 2 왼쪽에 나타나 있다.

이처럼 계산량(Gflops)과 성능(FID)의 관계를 직접 그려보면 상당히 선형적인 관계를 잘 따르는 것을 보여준다(Figure 8). 그리고 크기가 큰 모델일수록 계산량을 효율적으로 사용하는 것도 알 수 있다(Figure 9). 작은 모델의 training step 수를 늘려서 큰 모델의 것과 계산량을 비슷하게 하더라도 성능에 한계가 있다는 것이다. 비슷한 방법으로 sampling step 수를 늘려서 계산량을 비슷하게 하더라도 성능에 한계가 있었다.

정리하자면, DiT의 계산량을 늘리는 방법으로 저자들은 4가지를 제시하였고, 그 중 2가지는 효과적이고 2가지는 효과적이지 않다.

효과적인 방법: model size를 키우거나 patch size를 줄이는 것
효과적이지 않은 방법: training step 수를 늘리거나 sampling step 수를 늘리는 것

결국 model architecture 차원에서의 계산량 변화가 뚜렷한 성능 향상을 가져온다는 것을 알 수 있다.

💡 Summary

DiT(Diffusion Transformer) architecture 상에서의 contribution은 transformer backbone의 diffsion model을 구현했다는 점, adaLN-Zero Block을 통한 conditioning strategy를 제안했다는 점이다.
DiT는 직관적으로 scalable하며, model size를 키우거나 patch size를 줄이는 것이 성능 향상에 도움이 된다는 것을 실험을 통해 확인했다. 반면 training step 수나 sampling step 수를 늘리는 것은 효과적이지 않다.
Opinion: 하고자 하는 것이 아주 명확한 논문이라 좋다.

📃 Reference

[23’ ICCV] DiT: Scalable Diffusion Models with Transformers

I'm rubatoyeong