[2025-1] 황징아이 - Scalable Diffusion Models with Transformers

Abstract

Transformer 아키텍쳐에 Diffusion을 접목 시킨 논문.
기존 LDM (Latent Diffusion Model에서는 U-Net을 사용했다. 그러나 U-Net의 Inductive Bias가 Diffusion Model의 성능에 중요하지 않아서 Standard design인 Transformer를 사용.
Diffusion Transformer (DiT)를 제시했다. 이는 기존의 Convolutional Network (e.g. ResNet)보다 시각적 인식에 더 효율적으로 확장할 수 있는 Vision Transformer(ViT)를 참고했다.
또한, 저자는 Scaling behavior of transformers with respect to network complexity vs. sample quality 결과를 연구했다.

Diffusion Transformer 개념

Diffusion Models (DDPM)

원본 이미지 $x_0$, 완전한 Gaussian 상태 $x_T$
Forward Process (학습 대상 X) : 원본 이미지 $x_0$에서 완전한 Gaussian 상태 $x_T$ 가 될때 까지 정해준 Schedule에 맞추어 순차적으로 조금씩 Noise를 가하는 과정 (정해진 Schedule는 Linear, Cosine, Sigmoid, Quad…)
Reverse Process (학습 대상 O) : Forward의 반대로 완전한 Gaussian 상태 $x_T$에서 원본 이미지 $x_0$까지 순차적으로Noise를 제거하는 과정

Latent Diffusion Models (LDM)

고해상도 픽셀 space에서 직접 diffusion model을 학습하는 것은 계산적으로 불가능 할 수 있다. Latent diffusion model (LDM)은 2단계 접근 방식으로 이 문제를 해결한다.

학습된 인코더 E를 사용하여 이미지를 더 작은 space의 표현으로 압축하는 오토인코더를 학습한다.
이미지 x의 diffusion model 대신 표현 z=E(x)의 diffusion model을 학습한다 (E는 고정).

그리고는 Diffusion model에서 표현 z를 샘플링하고 학습된 디코더 x=D(z) 를 사용하여 이미지로 디코딩하여 새 이미지를 생성할 수 있다.

Diffusion Transformer Design Space

Patchify (DiT의 첫 레이어)

기존의 Vision Transformer 아키텍처 기반 (Input : 패치 단위의 토큰)
DiT의 입력값 $z = (I\times\ I\times C)$ 를 각 $d$ 패치크기의 차원을 갖는 $T = (I/P)^2$개의 토큰으로 변환
ViT와 마찬가지로 Sine Cosine값의 Position Embedding를 모든 입력 토큰에 적용.
$P$를 절반으로 줄이변 $T$는 4배 커져 최종적으로 Gflops 또한 4배 커지게 된다. $P$값에 따라 Gflops가 영향을 받지만 전체 파라미터 수에는 큰 영향을 끼치지 않는다고 한다. 저자는 $P = 2,4,8$로 실험을 했다.

DiT Block Design

Patchify후, DiT Block는 토큰 외에 Conditional Information도 (noise timestep $t$, class labels $c$…) 임베딩을 한다
Condition을 적용하는 방법을 다양하게 비교
- In-Context Conditioning
  - Patchify가 완료된 Token 임베딩에 단순히 Time Step $t$ 와 Class Label $c$ 를 두개의 추가 토큰을 추가
  - ViT의 cls토큰와 유사하다. 기존 ViT블록을 수정없이 사용.
  - Gflops의 큰 변화가 없다.
- Cross-Attention Block
  - Multi-head Self-Attention 다음으로 Multi-head Cross Attention추가
  - 기존 LDM에서 클래스 레이블로 컨디셔닝하는 데 사용하는 것과 유사
  - 기존 대비 15% 이상의 Gflops 증가
- Adaptive Layer Norm (adaLN) Block
  - GAN의 adaptive normalization layer와 UNet backbone이 있는 diffusion model의 광범위한 사용에 따라 Transformer block의 표준 레이어를 adaptive layer norm (adaLN)으로 대체한다
  - 차원별 scaling 및 shift 파라미터 $\gamma$ (gamma)와 $\beta$(beta)를 직접 학습 하는 대신 $t$와 $c$의 임베딩 벡터합계에서 회귀한다.
  - 가장 적은 Gflops를 추가하므로 가장 컴퓨팅 효율적이다.
  - 모든 토큰에 동일한 feature를 적용하도록 제한되는 유일한 컨디셔닝 메커니즘이다.
- Adaptive Layer Norm (adaLN)-Zero Block
  - Diffusion U-Net 모델은 유사한 초기화 전략을 사용하여 residual 연결 전에 각 블록의 최종 convolution layer를 0으로 초기화한다
  - DiT에서는 $\gamma$와 $\beta$를 회귀하는 것 외에도 DiT 블록 내의 Residual Connection 전 Scaling을 진행하는 scale factor $\alpha$를 추가한다 (초기 $\alpha$가 0이다)

Decoder Design

최종 DiT블록의 결과물에 LayerNorm (adaLN) 적용 하고 각 토큰을 $p\times\ p \times2C$텐서로 선형 디코딩한다. (C는 입력채널 수) : ($32\times32\times8 \rightarrow 2\times32\times32\times4)$
원래 공간 레이아웃으로 재정렬하면 최종 결과는 Noise 예측값과 Covariance값

Architecture

Experiment

사용 데이터셋 : 256x256, 512x512 Image Net
Training Setup :
- Optimizer : AdamW
- Learning Rate : $1\times10^{-4}$ (Weight Decay X, Warm Up X)
- Batch Size : 256
VAE
- Stable Diffusion에서 사용한 VAE를 사용 (256→32)
평가 Metric
- FID, Inception Score, sFID, Precisio/Recall

Experiment Result

DiT Block Design

*XL/2 : DiT XL (Patch size = 2), FID-50K 낮을수록 좋다

DiT Block Design
- 조건부 매커니즘이 모델 품질에 중요한 영향을 미침
- adaLN-Zero이 Cross-attention, In-Context 매커니즘보다 낮은 FID를 보여줌
- 각 DiT 블록을 초기화하는 adaLN-Zero는 기본 adaLN보다 성능이 활씬 좋음
- adaLN-zero가장 효율적이다

Scaling Model and Patch Size

DiT (S, B, L, XL), Patch size (2, 4, 8) 총 12가지 케이스를 비교
모델 크기가 클수록, Patch Size를 작게 할 수록 성능이 향상

다양한 DiT모델의 크기에 따른 FID의 변화 :

위 : 모델 크기 고정

아래 : 패치사이즈 고정

모델 크기가 클수록, Patch Size를 작게 할 수록 성능이 향상

Gfloops의 중요성

Gflops에 따른 FID를 나타낸 그래프 : Gfloops와 FID-50K 간의 강한 음의 상관관계가 있음을 확인 ⇒ Gfloops증가가 모델 성능향상의 핵심

큰 모델일수록 더 효율적이다

→ 작은 DiT모델은 더 오래 훈련해도 더 적은 단계로 훈련된 큰 DiT모델에 비해 계산 효율이 떨어진다

동일한 모델이더라도 패치크기에 따라 제한된 Gfloops내 성능이 다르다 ($10^{10}$ Gfloops 히루 XL/4 모델은 XL/2 모델보다 성능이 떨어짐)

Visualizing Scaling

패치 사이즈와 Transformer 크기에 따른 이미지 품질 변화

기존 모델들과의 비교

DiT-XL/2는 이전 모든 Diffusion 모델보다 성능이 좋고, LDM이 달성한 이전 최소 FID-50K 3.6을 2.27로 낮췄다

결론

Diffusion Transformer 는 기존의 U-Net 백본 모델보다 높은 성능 및 효율성 확장성을 보인다
향후 Stable Diffusion 3, Sora와 같은 대형 생성 모델들의 기초가 되는 논문Diffusion Transformer 개념Diffusion Models (DDPM)