DiT (Diffusion Transformer)

📄 관련 논문: [ICCV 2023] Scalable Diffusion Models with Transformers

DiT는 기존 디퓨전 모델에서 표준처럼 사용되던 U-Net 백본을 Transformer 아키텍처로 대체한 모델이다.

DiT 모델의 파이프라인은 다음과 같다.

DiT 또한 LDM처럼 latent 공간에서 동작하기 때문에, VAE 인코더를 거친 feature가 입력된다.

Patchify

입력으로 들어온 노이즈가 섞인 latent $x_t \in \mathbb{R}^{C \times I \times I}$를 ViT와 동일한 방식으로 $p \times p$ 크기의 패치들 $x_{t,i}$로 분할한다.

이후 각 패치들을 Flatten하고 Linear Projection을 거쳐 1차원 토큰 시퀀스로 변환한다.

\[\vx_{t,i}=\text{MLP}(x_{t,i}) \tag{1}\]

Positional Embedding

트랜스포머는 입력 토큰의 순서를 본질적으로 알 수 없기 때문에, 토큰 시퀀스에 위치 정보를 더해준다.

논문 원본에서는 2D Sine-Cosine 임베딩을 사용하였다.

Conditioning 방식

논문에서는 Noise timestep $t$와 조건 $c$ (텍스트 프롬프트, 클래스 라벨 등)를 네트워크에 가장 효과적으로 주입하기 위해 다양한 Conditioning 방법을 실험하였다.

\[\ve_{t}=\text{Embed}(t)\quad,\quad\ve_{c}=\text{Embed}(c) \tag{2}\]

논문에서 최종적으로 선택한 방법은 adaLN-Zero이다.

In-Context Conditioning

Timestep 임베딩과 Condition 임베딩을 patch token sequence에 단순히 concat하는 방식이다.

\[\begin{bmatrix} \vx_{t,1}, & \vx_{t,2}, & \cdots, & \vx_{t,T}, & \ve_{t}, & \ve_{c} \end{bmatrix} \tag{3}\]

• 장점: 표준 ViT Block을 구조적 수정 없이 그대로 사용할 수 있고, 모델에 추가되는 연산량이 매우 작다.
• 단점: 3가지 방식 중 성능이 가장 떨어진다.

Cross-Attention Block

DiT block에서 Self-Attention 뒤에 Cross-Attention layer를 추가해, Conditioning을 key와 value로, Input token을 query로 사용하는 방식이다.

• 장점: 기존 LDM에서 사용된 표준적인 Conditioing 방식과 유사하여 직관적이다.
• 단점: 3가지 방식 중 가장 많은 연산량을 추가 (대략 15%의 Gflops 오버헤드)하며, 연산 비용이 가장 큼에도 불구하고 adaLN-Zero 방식보다 FID가 낮다.

adaLN-Zero (Adaptive LayerNorm with Zero-initialization)

Adaptive LayerNorm

Transformer block 내부의 LayerNorm 연산은 다음과 같이 정의된다.

\[y=\gamma\odot\frac{x-\mu}{\sqrt{\sigma^2+\epsilon}}+\beta \tag{4}\]

Adaptive LayerNorm은 표준 Layer Norm을 조건부로 만들기 위해, 조건 임베딩의 합으로부터 scale parameter $\gamma$, $\alpha$와 shift parameter $\beta$를 회귀해내는 방식이다.

\[\gamma,\beta,\alpha=\text{MLP}(\ve_{t}+\ve_{c}) \tag{5}\]

즉, 이 MLP는 입력된 $t$나 $c$에 맞춰, 각 토큰이 어떻게 정규화되어야 하는지를 네트워크 스스로 조정하도록 만든다.

여기서 각 scale parameter $\gamma$와 $\alpha$는 적용되는 위치와 목적이 다르다.

• $\gamma$ : Sub-Block으로 들어가기 전에 LayerNorm 연산 내부에서 정규화된 데이터에 직접 곱해지며, Condition을 데이터에 주입하는 역할을 한다.
• $\alpha$ : Sub-Block에서 처리된 데이터와 원본 데이터의 Residual Connection 직전에 곱해지며, 연산 블록의 출력값을 원본에 얼만큼 반영할지 조절하는 역할을 한다.

\[x_{l+1} = x_{l} + \alpha \odot \text{SubBlock}(\text{adaLN}(x_{l}, \gamma, \beta)) \tag{6}\]

Zero-initialization

Zero-initialization은 Residual Connection 직전에 곱해지는 scale parameter $\alpha$를 0으로 초기화하는 방식이다.

이 덕분에 학습 극초기에는 Sub-Block 연산 결과가 0이 되어 무시되고, 네트워크 전체가 데이터 원본을 그대로 통과시키는 Identity function처럼 작동하게 만들어 학습 불안정성을 막는다.