Transformer로 확장 가능한 확산 모델 (DiT)
Digest: 확산 모델은 사실상 항상 합성곱 기반 U-Net을 노이즈 예측 백본으로 사용해 왔지만, 다른 분야(NLP·비전 인식)에서 표준이 된 Transformer의 **확장성(scalability)**이 생성 모델 백본에서도 성립하는지는 검증되지 않았다 (Context). 저자들은 잠재 확산 모델(latent diffusion)의 U-Net을 잠재 패치 위에서 작동하는 순수 Transformer(Diffusion Transformer, DiT)로 대체하고, 백본의 계산량(Gflops)을 늘리면 생성 품질(FID)이 단조롭고 예측 가능하게 향상된다는 강한 상관관계를 발견했다 (Insight). 이를 위해 (1) VAE 잠재맵을 패치화해 토큰 시퀀스로 변환하고, (2) 표준 ViT 블록에 timestep·class 조건을 주입하는 네 가지 방식을 비교해 잔차 블록을 항등사상으로 초기화하는 adaLN-Zero가 가장 우수함을 확인했으며, (3) 패치 크기·모델 폭/깊이를 체계적으로 스케일링했다 (Solution). 그 결과 가장 큰 DiT-XL/2가 ImageNet 256×256 class-conditional 생성에서 FID 2.27로 LDM(3.60)·ADM(3.94)을 능가하는 SOTA를 달성했고, 512×512에서도 FID 3.04를 기록했다 (Evidence). 다만 여전히 사전학습된 VAE 잠재공간에 의존하고, 작은 패치(p=2)가 토큰 수를 늘려 추론 비용이 크며, U-Net의 귀납적 편향(다중 스케일 skip)을 토큰 수·계산량으로 보상한다는 trade-off가 있다 (Limitations). 이 단순하고 확장 가능한 백본은 이후 Sora·Stable Diffusion 3·PixArt 등 대규모 생성 모델의 사실상 표준 아키텍처가 되었다 (Open Questions).
섹션별 요약
Introduction
- 확산 모델은 Denoising Diffusion Probabilistic Models 이래 거의 예외 없이 U-Net 백본을 사용해 왔다. U-Net은 강력하지만 그 합성곱 구조가 확산 모델 성능에 필수적인지는 검증된 적이 없다.
- 다른 도메인에서 Transformer는 모델·데이터·계산량을 키울수록 성능이 향상되는 scaling law를 보였다. 본 논문의 질문: “확산 모델 백본의 아키텍처 선택이 샘플 품질에 얼마나 중요한가? Transformer로 바꾸면 그 확장성을 물려받는가?”
- 기여: (1) U-Net을 ViT로 대체한 DiT 제안, (2) 조건화 메커니즘 4종 비교 → adaLN-Zero, (3) Gflops ↔ FID의 강한 음의 상관을 실증(아키텍처 복잡도가 아니라 계산량이 품질을 결정), (4) ImageNet 256·512 SOTA.
Methods
-
잠재 공간 작동: 입력 이미지(256×256×3)를 LDM의 사전학습 VAE로 8배 압축한 32×32×4 잠재맵에서 확산을 수행(픽셀이 아닌 잠재 확산).
-
Patchify: 잠재맵을 패치 크기
p ∈ {2,4,8}로 분할해 토큰 시퀀스로 변환. p=2 → 16×16=256 토큰, p=4 → 64 토큰, p=8 → 16 토큰. 작을수록 토큰↑·Gflops↑. ViT처럼 sine-cosine 위치 임베딩 추가. -
DiT 블록 + 조건 주입(4종 비교):
- In-context: 조건 토큰을 시퀀스에 append(ViT의 cls 토큰처럼).
- Cross-attention: 조건을 별도 key/value로 cross-attend(약 +15% Gflops).
- adaLN: LayerNorm의 scale·shift(γ, β)를 timestep+class 임베딩으로부터 회귀.
- adaLN-Zero(최종 선택): adaLN에 더해 각 잔차 연결 직전 스케일 α도 회귀하고, α를 0으로 초기화해 블록을 항등사상에서 출발시킴 → 학습 안정·품질 최고.
-
모델 구성(ViT 스케일링 규칙 차용):
Config Layers Hidden d Heads Params DiT-S 12 384 6 ~33M DiT-B 12 768 12 ~130M DiT-L 24 1024 16 ~458M DiT-XL 28 1152 16 ~675M 표기
DiT-XL/2= XL 구성 + 패치 크기 2. DiT-XL/2의 백본 계산량은 약 118.6 Gflops(256 해상도). -
학습: AdamW, 상수 lr 1e-4, 배치 256, EMA, horizontal flip 외 데이터 증강 없음. DiT-XL/2는 약 7M 스텝 학습. classifier-free guidance로 샘플링.
Results
| Model | Dataset | Params/Gflops | FID ↓ | 비고 |
|---|---|---|---|---|
| DiT-XL/2 (+cfg) | ImageNet 256² | 675M / 118.6G | 2.27 | SOTA (Table) |
| DiT-XL/2 (no guidance) | ImageNet 256² | 675M | 9.62 | guidance 효과 큼 |
| LDM-4-G | ImageNet 256² | — | 3.60 | 이전 best |
| ADM-G,U | ImageNet 256² | — | 3.94 | U-Net + classifier guidance |
| BigGAN-deep | ImageNet 256² | — | 6.95 | GAN 계열 |
| DiT-XL/2 (+cfg) | ImageNet 512² | 675M | 3.04 | SOTA |
| ADM-G,U | ImageNet 512² | — | 3.85 | — |
Discussion
- 한계: (1) 사전학습 VAE 잠재공간에 의존 — 재구성 품질 상한에 묶임, (2) 작은 패치(p=2)는 토큰 수를 늘려 추론 Gflops가 커짐, (3) U-Net의 다중 스케일 skip 같은 귀납적 편향 없이 순수 attention으로 처리 → 계산량으로 보상.
- 향후: 더 큰 DiT·더 큰 데이터로의 스케일링, 텍스트 조건화(이후 PixArt·SD3), 비디오·3D로의 확장(이후 Sora).
Insights
- 주목할 점: 샘플 품질을 결정하는 것은 “U-Net이라는 특정 아키텍처”가 아니라 **백본의 총 계산량(Gflops)**이다 — 파라미터 수보다 Gflops가 FID와 더 잘 상관.
- 연결 고리: ViT의 패치·스케일링 레시피를 DDPM·LDM 프레임워크에 이식.
- 시사점: 확산 모델도 NLP·인식 분야의 scaling law를 따른다 — “더 큰 Transformer = 더 좋은 생성”이 성립.
- 비판적 코멘트: adaLN-Zero가 in-context/cross-attention보다 나은 이유에 대한 이론적 설명은 약하며, 주로 경험적 FID 비교에 의존.
Discussion Points
- 논쟁점: 귀납적 편향(U-Net skip) vs 순수 scale — 데이터·계산이 충분하면 편향이 불필요한가?
- 검증 필요 가정: VAE 잠재공간이 병목이 아닌가? (이후 SD3는 더 나은 VAE·rectified flow로 개선.)
- 후속 연구: PixArt-α(텍스트-투-이미지 DiT), Stable Diffusion 3(MM-DiT), Sora(비디오 DiT), Latte 등.
메타데이터
| 항목 | 내용 |
|---|---|
| 제목 | Scalable Diffusion Models with Transformers |
| 저자 | William Peebles, Saining Xie |
| 소속 | UC Berkeley, New York University |
| 연도 | 2022 (arXiv) / ICCV 2023 |
| 발표 | ICCV 2023 / arXiv:2212.09748 |
| 링크 | arXiv, GitHub, Project |
| 키워드 | diffusion transformer, latent diffusion, adaLN-Zero, scalability, ImageNet |
왜 이 연구를 하는가?
핵심 질문
확산 모델의 노이즈 예측 백본으로 합성곱 U-Net이 정말 필수인가? 이를 표준 Transformer로 대체하면, Transformer가 다른 분야에서 보인 확장성(계산량을 키울수록 성능 향상)을 생성 모델에서도 물려받는가?
기존 접근법의 한계
| 한계 | 설명 |
|---|---|
| U-Net 독점 (DDPM·ADM·LDM) | 모든 SOTA 확산 모델이 U-Net을 쓰지만, 그 구조적 선택이 품질에 얼마나 기여하는지 분리 검증된 적 없음 |
| 아키텍처 비교 부재 | 백본을 통제변수로 두고 계산량 대비 품질을 측정한 체계적 연구 부재 |
| 확장 레시피 부재 | ”확산 모델을 더 크게 만들면 얼마나 좋아지는가”에 대한 깨끗한 scaling 곡선 부재 |
핵심 통찰
- 잠재 확산의 노이즈 예측기를 잠재 패치 위의 ViT로 재구성할 수 있으며, U-Net 특유의 합성곱·다중 스케일 skip 없이도 동등 이상 성능이 가능하다.
- 모델 폭/깊이를 키우거나 패치를 작게 해 백본 Gflops를 늘리면 FID가 단조 감소한다 — 즉 확산 모델은 계산량으로 스케일된다.
- 조건 주입은 **항등 초기화(adaLN-Zero)**가 결정적이다: 학습 초기에 각 블록이 항등사상이면 깊은 생성 Transformer가 안정적으로 학습된다(ResNet/Fixup 계열 통찰과 동형).
방법 (Method)
프레임워크 개요
graph LR IMG[이미지 256²×3] -- "VAE Encoder (×8)" --> Z[잠재 32×32×4] Z -- "Patchify (p=2/4/8)" --> TOK[토큰 시퀀스 + 위치임베딩] T[timestep t] --> EMB[조건 임베딩] C[class label] --> EMB TOK --> DIT[DiT 블록 ×N] EMB -- "adaLN-Zero (γ,β,α)" --> DIT DIT --> HEAD[Linear: 노이즈 ε & Σ 예측] HEAD -. "역확산 샘플링" .-> Z Z -. "VAE Decoder" .-> OUT[생성 이미지]
핵심 구성요소
- 잠재 패치화: 32×32×4 잠재 → p×p 패치 → d차원 토큰 T개. 작은 p일수록 T↑·Gflops↑.
- DiT 블록: 표준 ViT 블록(multi-head self-attention + MLP)에 adaLN-Zero 조건화. LN의 γ·β와 잔차 스케일 α를 (t+class) 임베딩의 MLP로 회귀, α는 0 초기화.
- 예측 헤드: 최종 LN + linear로 잠재 노이즈 ε와 대각 공분산 Σ를 예측(DDPM + Improved DDPM 학습 목표).
- 샘플링: classifier-free guidance(scale≈1.5)로 잠재 역확산 후 VAE 디코더로 픽셀 복원.
발견 (Findings)
주요 결과
| 실험 | 메트릭 | 값 | 비교 |
|---|---|---|---|
| ImageNet 256² class-cond | FID | 2.27 (+cfg) | LDM-4-G 3.60, ADM-G,U 3.94 |
| ImageNet 512² class-cond | FID | 3.04 (+cfg) | ADM-G,U 3.85 |
| Gflops vs FID | 상관 | 강한 음의 상관 | 백본 계산량↑ → FID↓ (핵심 그림) |
| 조건화 ablation | FID | adaLN-Zero 최저 | in-context·cross-attn보다 우수 |
| 패치 크기 p | FID | p=2 < p=4 < p=8 | 토큰 수↑ → 품질↑ |
핵심 발견
- Gflops가 품질을 예측한다: 모델 폭·깊이를 키우거나 패치를 줄여 Gflops를 올리면 FID가 일관되게 감소. 파라미터 수 단독보다 Gflops가 FID와 더 강하게 상관.
- adaLN-Zero ≫ 다른 조건화: 동일 Gflops에서 adaLN 계열이 in-context·cross-attention보다 낮은 FID, 그중 항등 초기화(Zero)가 최고.
- 작은 패치가 유리: p=8→p=2로 줄이면 파라미터는 거의 그대로지만 Gflops↑·FID↓ — 토큰 해상도가 생성 품질의 핵심 레버.
이론적 의의
아키텍처 통일: 확산 백본도 Transformer로
DiT는 인식(ViT)·언어(GPT)에서 입증된 “단순한 Transformer + 스케일” 레시피가 생성 확산 모델에서도 통한다는 것을 보였다. 도메인 특화 귀납적 편향(U-Net의 합성곱·skip)을 제거하고도 SOTA를 달성함으로써, 멀티모달·대규모로 확장하기 쉬운 단일 백본 패러다임을 열었다.
계산량 중심의 스케일링 관점
“어떤 아키텍처인가”보다 “백본에 얼마의 계산을 쓰는가(Gflops)“가 FID를 결정한다는 발견은, 확산 모델 설계를 아키텍처 탐색에서 계산 예산 배분 문제로 재프레이밍한다. 이는 LLM의 compute-optimal scaling 논의와 같은 결의 통찰이다.
항등 초기화의 보편성
adaLN-Zero(잔차 블록을 항등에서 출발)는 ResNet의 zero-init residual·Fixup·DDPM의 안정적 학습과 같은 원리로, 깊은 생성 네트워크 학습의 일반 처방으로 재확인된다.
재현성 및 신뢰도 평가
| 항목 | 등급 | 비고 |
|---|---|---|
| 코드 공개 | ✅ | facebookresearch/DiT PyTorch 공식 + 사전학습 체크포인트 |
| 데이터 공개 | ✅ | ImageNet 공개, Stable Diffusion VAE 사용 |
| 하이퍼파라미터 | ✅ | 모델 구성·lr·배치·스텝·guidance scale 명시 |
| 실험 환경 | ✅ | TPU/GPU·학습 스텝 보고 |
| 통계적 신뢰도 | ⚠️ | FID 단일 측정(50k 샘플), 표준편차 미보고 |
| 종합 등급 | A | 광범위하게 재현·후속 채택(SD3, PixArt, Sora) |
주장별 신뢰도
| # | 주장 | 근거 | 신뢰도 |
|---|---|---|---|
| 1 | DiT-XL/2 ImageNet 256² FID 2.27 SOTA | 결과 표, 공식 체크포인트 재현됨 | 🟢 |
| 2 | Gflops ↔ FID 강한 음의 상관 | 스케일링 그림, 다수 구성 측정 | 🟢 |
| 3 | adaLN-Zero가 최적 조건화 | 조건화 ablation 표 | 🟢 |
| 4 | 512² FID 3.04 | 결과 표 | 🟢 |
| 5 | U-Net 귀납편향 불필요 | 강한 경험적 증거이나 데이터·계산 규모 의존 | 🟡 |
읽기 난이도: ⭐⭐⭐
- 필요 배경: ViT, Transformer, DDPM/LDM, classifier-free guidance, FID.
- 수식보다 설계·ablation 중심이라 DDPM 대비 직관적이나, scaling 곡선 해석이 핵심.
관련 연구 비교 매트릭스
| 축 | DiT (본 논문) | LDM (Rombach 2022) | ADM (Dhariwal 2021) | U-ViT (Bao 2022) |
|---|---|---|---|---|
| 백본 | 순수 Transformer (잠재 패치) | 합성곱 U-Net | 합성곱 U-Net | Transformer + long skip |
| 작동 공간 | VAE 잠재 | VAE 잠재 | 픽셀 | 잠재/픽셀 |
| 조건화 | adaLN-Zero | cross-attention | classifier guidance | concat 토큰 |
| 256² FID | 2.27 | 3.60 | 3.94 (ADM-G,U) | ~2.29 |
| 핵심 주장 | 백본 Gflops가 품질 결정 | 잠재공간이 효율 핵심 | guidance가 품질 핵심 | skip 연결이 Transformer에 유용 |
| 확장성 | 매끄러운 scaling 곡선 | 제한적 분석 | — | 유사 시기 병렬 연구 |
| 코드 공개 | ✅ | ✅ | ✅ | ✅ |
관련 연구
- Denoising Diffusion Probabilistic Models — DiT가 계승한 확산 학습 목표(ε·Σ 예측)와 샘플링.
- An Image is Worth 16x16 Words - Transformers for Image Recognition at Scale — DiT 블록·패치화·스케일링 레시피의 직접 출처(ViT).
- Attention Is All You Need — self-attention 백본의 원류.
- LDM / Stable Diffusion (Rombach et al. 2022) — DiT가 사용하는 VAE 잠재공간·잠재 확산 프레임워크.
- ADM (Dhariwal & Nichol 2021) — U-Net + classifier guidance 기반 이전 SOTA 비교 대상.
- U-ViT (Bao et al. 2022) — 같은 시기 Transformer 확산 백본 병렬 연구.
원자적 인사이트 (Zettelkasten)
💡 확산 모델의 샘플 품질은 백본 아키텍처가 아니라 백본의 총 계산량(Gflops)이 결정한다
출처: Scalable Diffusion Models with Transformers (Peebles & Xie, 2022)
유형: 실험적
DiT의 폭·깊이를 키우거나 패치 크기를 줄여 백본 Gflops를 증가시키면 ImageNet FID가 일관되게 단조 감소한다. 파라미터 수 단독보다 Gflops가 FID와 더 강하게 상관하며, 이는 “U-Net이라는 특정 합성곱 구조”가 아니라 노이즈 예측에 투입된 계산량이 생성 품질의 1차 결정 요인임을 시사한다.
핵심 조건/맥락: class-conditional ImageNet 256/512, 사전학습 VAE 잠재공간 가정.
연결: LLM compute-optimal scaling, ViT scaling law.
활용 가능성: 생성 모델 설계를 “아키텍처 탐색”에서 “계산 예산 배분”으로 재프레이밍.
💡 잔차 블록을 항등사상으로 초기화하는 adaLN-Zero가 깊은 생성 Transformer 학습을 안정화한다
출처: Scalable Diffusion Models with Transformers (Peebles & Xie, 2022)
유형: 방법론적
LayerNorm의 scale·shift를 조건 임베딩으로 회귀하는 adaLN에 더해, 각 잔차 연결 직전의 스케일 α를 0으로 초기화하면 학습 초기에 모든 DiT 블록이 항등사상에서 출발한다. 이 조건화가 in-context·cross-attention 방식보다 낮은 FID를 안정적으로 달성한다.
핵심 조건/맥락: timestep+class 조건을 단일 임베딩으로 합산해 주입.
연결: ResNet zero-init residual, Fixup, DDPM 안정적 학습.
활용 가능성: 깊은 조건부 생성 네트워크의 일반적 초기화 처방.
💡 잠재맵의 패치 해상도(토큰 수)는 추론 비용과 생성 품질을 동시에 지배하는 핵심 레버다
출처: Scalable Diffusion Models with Transformers (Peebles & Xie, 2022)
유형: 실험적
32×32×4 잠재를 패치 크기 p로 나눌 때 p=8(16토큰)→p=2(256토큰)로 줄이면 파라미터는 거의 불변이지만 Gflops와 FID가 함께 개선된다. 즉 모델을 “더 깊게/넓게” 만드는 것 외에 “토큰을 더 잘게” 만드는 것도 독립적인 스케일링 축이다.
핵심 조건/맥락: 잠재공간 해상도 고정(VAE ×8) 하에서의 패치 분할.
연결: ViT 패치 크기 trade-off, 잠재공간 해상도 vs 픽셀공간.
활용 가능성: 추론 예산이 제약일 때 모델 크기 대신 토큰 해상도로 품질-비용 조절.
핵심 용어 정리
| 용어 | 정의 |
|---|---|
| Diffusion Transformer (DiT) | 확산 모델의 노이즈 예측 U-Net을 잠재 패치 위 ViT로 대체한 백본. |
| Latent diffusion | 픽셀이 아닌 사전학습 VAE의 압축 잠재공간에서 확산을 수행하는 방식(LDM). |
| Patchify | 잠재맵을 p×p 패치로 나눠 토큰 시퀀스로 변환하는 ViT식 입력 처리. |
| adaLN | LayerNorm의 scale·shift(γ,β)를 조건 임베딩으로부터 회귀하는 적응적 정규화. |
| adaLN-Zero | adaLN + 잔차 스케일 α를 0 초기화해 블록을 항등사상에서 출발시키는 조건화(최적). |
| Gflops | 백본 1회 순전파의 부동소수 연산량. DiT에서 FID와 강하게 음의 상관. |
| Classifier-free guidance | 조건/무조건 예측을 외삽해 조건 충실도를 높이는 샘플링 기법. |
| FID | 생성·실제 분포의 Inception feature Fréchet 거리. 낮을수록 좋음. |
| DiT-XL/2 | XL 구성(28층·d=1152) + 패치 크기 2. 본 논문 최고 성능 모델. |
태그
paper #2022 diffusion transformer generative-model architecture vision ICCV
BibTeX
@inproceedings{Peebles2023DiT,
title = {Scalable Diffusion Models with Transformers},
author = {Peebles, William and Xie, Saining},
booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
year = {2023},
url = {https://arxiv.org/abs/2212.09748}
}