[평범한 학부생이 하는 논문 리뷰] DiT (ICCV 2023 oral) & MM-DiT (ICML 2024)

Paper : https://arxiv.org/abs/2212.09748

Scalable Diffusion Models with Transformers

Paper : https://arxiv.org/abs/2403.03206

Scaling Rectified Flow Transformers for High-Resolution Image Synthesis

1. DiT

 DiT는 ViT의 아키텍처를 기반으로한 diffusion model이다.

Patchify

DiT Block Design

 Diffusion model은 noisy input image외에도 time, class, natural language와 같은 다른 condition을 받아야한다. DiT논문에서는 여러가지 방식을 실험한 결과 adaLN-Zero방식을 선택했다. AdaLN-Zero는 MLP가 모든 $\alpha$에 대해서는 처음에 $0$을 뱉도록 한다.

Transformer Decoder

 최종 DiT블록의 output인 sequence of image token들을 original spatial input과 동일한 크기의 output noise prediction과 output diagonal covariance prediction으로 decoding해야한다. DiT 논문에서는 standard linear decoder를 사용한다. Layer norm (adaptive if using adaLN)를 적용하고 각 토큰을 $p \times p \times 2C$의 tensor로 linearly decoding한다. 최종적으로 decoding된 token들을 original spatial layout으로 rearrange한다.

2. MM-DiT

 MM-DiT는 DiT를 기반으로 한다. 두 개의 DiT 브랜치로 구성된 느낌이고, text와 image embedding을 따로 처리한다. 그리고 attention 계산할 때에만 concat해서 joint하게 계산한다. 사실 아키텍처에 대한 내용은 논문에서도 크게 다루고 있지 않다. 논문에서는 주로 noise를 어떻게 잘 샘플링할 것인가, 어떤 forward process formulation을 사용할 것인가를 다루고 실험적으로 어떤 것이 결과가 좋았는 지를 다루는 것으로 파악했다.

3. SiT

 번외로 DiT의 variant인 SiT 또한 앞선 MM-DiT에서와 같이 모델 아키텍쳐 자체에 집중하기 보다는,

1. vector field prediction이냐 score function prediction이냐
2. continuous time로 학습하냐 discrete time으로 학습하냐
3. forward process를 어떻게 둘 것이냐

등에 대해 실험해보고 어떤 design이 성능이 좋더라를 보여준 것으로 이해했다.