확산 트랜스포머

Diffusion Transformer, DiT; 확산 트랜스포머
확산 모델의 잡음 예측 신경망을 U-Net 대신 트랜스포머로 만든 구조. 잠재 이미지를 패치 토큰으로 잘라 처리한다

확산 트랜스포머(Diffusion Transformer, DiT)는 확산 모델이나 플로 매칭 모델에서 잡음(또는 속도)을 예측하는 백본을 합성곱 U-Net 대신 트랜스포머로 바꾼 구조다. 2022년 12월 UC 버클리의 피블스(William Peebles)와 뉴욕 대학교의 셰싸이닝(Saining Xie)이 발표한 논문 〈Scalable Diffusion Models with Transformers〉에서 이름이 나왔다.^[1] 이 논문은 트랜스포머의 크기나 토큰 수를 늘려 순전파 연산량(Gflops)이 커질수록 FID가 꾸준히 낮아진다는 것을 보였고, 이후 스테이블 디퓨전 3, FLUX.1, 소라, 무비 젠 같은 대형 이미지·영상 생성 모델의 표준 백본이 되었다.

"확산 트랜스포머"는 좁게는 피블스·셰싸이닝의 DiT 구조를, 넓게는 확산·플로 매칭 계열 생성 모델에 쓰는 트랜스포머 백본 전체(PixArt, MMDiT 등)를 가리킨다. 이 문서는 원조 DiT의 설계를 자세히 보고 후속 변형을 비교한다.

배경

2020~2022년의 확산 모델(DDPM, ADM, 잠재 확산 모델)은 해상도를 줄였다 늘리는 합성곱 U-Net에 셀프 어텐션 층을 섞어 썼다. 다리왈과 니컬이 U-Net 구조 선택을 상세히 실험했지만 백본 자체를 바꾸는 시도는 적었다. 한편 비전 트랜스포머는 이미지 인식에서 데이터와 모델을 키울수록 성능이 오르는 스케일링 성질을 보이고 있었다. DiT는 "U-Net의 귀납적 편향은 확산 모델 성능에 꼭 필요하지 않다"는 가설을 세우고, 잠재 확산 모델의 U-Net 자리에 ViT 형태의 트랜스포머를 넣었다.^[1]

구조

잠재 공간과 패치화

DiT는 픽셀이 아니라 스테이블 디퓨전의 사전 학습 VAE 잠재 공간에서 동작한다. 256×256×3 이미지는 가로세로 8배 줄어든 32×32×4 잠재 $z$ 가 된다.^[1]

첫 층 "패치화"(patchify)는 $I \times I \times C$ 잠재를 $p \times p$ 패치로 잘라 각각을 선형 사상으로 $d$ 차원 토큰으로 만든다. 토큰 수는

$T = (I / p)^{2}$

이다. 패치 크기 $p$ 를 반으로 줄이면 토큰이 4배가 되어 Gflops가 적어도 4배로 늘지만 매개변수 수는 거의 그대로다. 논문은 $p = 2, 4, 8$ 을 실험했고, 모든 토큰에 ViT식 2차원 사인·코사인 위치 임베딩을 더했다. 출력 쪽에서는 마지막 층 정규화 후 선형 디코더로 토큰마다 $p \times p \times 2 C$ 값을 내고, 이를 다시 공간 배치로 되돌려 잡음 예측과 대각 공분산 예측을 얻는다.

조건 주입 방식

시간 단계 $t$ 와 클래스 레이블 $c$ 를 넣는 방식 네 가지를 비교했다.

방식	내용	DiT-XL/2 Gflops
문맥 내 조건(in-context)	$t$ , $c$ 임베딩을 토큰 두 개로 시퀀스에 붙임. ViT의 cls 토큰과 비슷	119.4
교차 어텐션	$t$ , $c$ 임베딩을 별도 시퀀스로 두고 블록마다 교차 어텐션 층 추가	137.6(약 15% 증가)
adaLN	층 정규화의 스케일 γ와 이동 β를 학습 상수 대신 $t + c$ 임베딩에서 회귀	118.6
adaLN-Zero	adaLN에 잔차 연결 직전의 차원별 게이트 α를 더하고, α가 0이 되도록 초기화	118.6

adaLN-Zero가 학습 내내 가장 낮은 FID를 기록해 논문의 나머지 실험은 모두 이 블록을 썼다.^[1]

adaLN-Zero

블록 하나의 계산은 다음과 같다. $𝐜$ 는 시간 임베딩과 클래스 임베딩의 합이다.

$(β_{1}, γ_{1}, α_{1}, β_{2}, γ_{2}, α_{2}) = M L P (𝐜)$

$x \leftarrow x + α_{1} ⊙ A t t n (L N (x) (1 + γ_{1}) + β_{1})$

$x \leftarrow x + α_{2} ⊙ F F N (L N (x) (1 + γ_{2}) + β_{2})$

여기서 LN은 학습 가능한 아핀 매개변수가 없는 층 정규화다. 조건을 만드는 MLP의 마지막 선형층을 0으로 초기화하면 처음에는 $α = 0$ 이라 모든 블록이 항등 함수가 된다. ResNet의 마지막 배치 정규화 스케일을 0으로 두거나 확산 U-Net이 잔차 연결 앞 합성곱을 0으로 초기화하는 것과 같은 발상이다.^[1] adaLN은 모든 토큰에 같은 변조를 거는 방식이라 연산량이 거의 늘지 않는다. 대신 클래스 하나나 풀링된 텍스트 벡터처럼 짧은 조건에는 맞지만, 긴 텍스트 토큰열을 세밀하게 반영하려면 교차 어텐션이나 아래 MMDiT 같은 구조가 필요하다.

모델 크기

ViT의 설정을 따라 네 가지 크기를 두고 패치 크기와 조합해 이름을 붙였다. 예를 들어 DiT-XL/2는 XL 설정에 $p = 2$ 다.

설정	층 수 N	은닉 크기 d	헤드 수	Gflops(I=32, p=4)
DiT-S	12	384	6	1.4
DiT-B	12	768	12	5.6
DiT-L	24	1024	16	19.7
DiT-XL	28	1152	16	29.1

12개 모델의 연산량은 0.3~118.6 Gflops에 걸쳐 있다. 학습은 AdamW에 고정 학습률 1×10⁻⁴, 가중치 감쇠 없음, 배치 256, 가중치 EMA 감쇠 0.9999였고 데이터 증강은 좌우 반전뿐이었다. 학습률 워밍업이나 정규화 없이도 손실 급등 없이 안정적으로 학습되었다고 한다. JAX로 구현해 TPU v3-256 팟에서 학습했고 DiT-XL/2는 초당 약 5.7회 반복했다.^[1]

결과

스케일링

모델 크기를 키우거나(층·폭) 패치를 줄여 토큰을 늘리면 모든 조합에서 FID가 좋아졌다. 400K 단계 시점의 Gflops와 FID 사이에는 강한 음의 상관이 있었고, 매개변수 수가 비슷해도 Gflops가 다르면 성능이 달랐다. 즉 "매개변수 수가 아니라 연산량이 품질을 결정한다"는 것이 논문의 핵심 관찰이다. 큰 모델은 같은 총 학습 연산으로도 더 좋은 FID에 도달해 연산 효율도 높았다. 반면 작은 모델에 샘플링 단계를 많이 주어 추론 연산을 늘려도 큰 모델의 부족한 모델 연산을 메우지 못했다.^[1]

ImageNet 클래스 조건 생성

해상도	모델	FID-50K	비교
256×256	DiT-XL/2, 700만 단계, CFG 1.5	2.27	LDM 3.60, StyleGAN-XL 2.30
512×512	DiT-XL/2, 300만 단계, CFG 1.5	3.04	ADM 3.85

256×256 모델은 118.6 Gflops로 잠재 U-Net인 LDM-4(103.6 Gflops)와 비슷하고 픽셀 공간 ADM(1120 Gflops)보다 훨씬 적다. 512×512에서 DiT-XL/2는 64×64×4 잠재를 1024개 토큰으로 처리해 524.6 Gflops를 쓰며, ADM(1983 Gflops)과 ADM-U(2813 Gflops)보다 적다.^[1] 공식 코드는 재현성을 이유로 분류기 없는 가이던스를 출력 채널 중 처음 세 채널에만 적용하며, 주석으로 전 채널 적용이 표준 방식이라고 밝혀 두었다.^[2]

후속 구조

모델	연도	조건 주입	특징
DiT	2022	adaLN-Zero(시간 + 클래스)	클래스 조건 ImageNet. 스케일링 성질 입증
PixArt-α^[3]	2023	DiT 블록에 텍스트 교차 어텐션 추가, 클래스 조건 가지 간소화	학습 단계를 셋으로 나누고 VLM이 단 촘촘한 캡션 사용. 학습에 A100 675일(SD 1.5는 6,250일)
소라^[4]	2024	공개 안 됨	영상 잠재를 시공간 패치로 잘라 처리하는 확산 트랜스포머
MMDiT(스테이블 디퓨전 3)^[5]	2024	텍스트·이미지 토큰을 하나의 시퀀스로 합쳐 공동 어텐션, 모달리티별 별도 가중치	정류 흐름 학습, QK 정규화, 최대 80억 매개변수
FLUX.1^[6]	2024	멀티모달 블록과 병렬 블록을 섞은 혼합 구조	120억 매개변수, 플로 매칭, 회전 위치 임베딩, 병렬 어텐션 층
무비 젠^[7]	2024	텍스트 교차 어텐션	300억 매개변수 영상 모델, 플로 매칭

MMDiT

텍스트-이미지 생성에서는 조건이 긴 텍스트 토큰열이다. 스테이블 디퓨전 3의 MMDiT(멀티모달 DiT)는 CLIP 두 개와 T5-XXL의 텍스트 임베딩 토큰과 이미지 패치 토큰을 한 시퀀스로 이어 붙여 함께 어텐션하되, 텍스트와 이미지에 서로 다른 가중치(정규화, QKV 사영, MLP)를 둔다. 논문은 이를 "모달리티마다 독립된 트랜스포머 두 개를 두고 어텐션 연산에서만 두 시퀀스를 합친 것"과 같다고 설명했다. 그래서 정보가 텍스트에서 이미지로만 가는 교차 어텐션과 달리 양방향으로 흐른다. 풀링된 CLIP 벡터와 시간 임베딩은 여전히 adaLN 변조로 넣는다. 모델 크기는 깊이 $d$ 하나로 정해 은닉 크기 $64 d$ , 헤드 $d$ 개로 두었고, 가장 큰 $d = 38$ 모델이 80억 매개변수다. 큰 모델에서 어텐션 로짓이 커져 학습이 불안정해지는 문제는 Q와 K에 RMSNorm을 거는 QK 정규화로 막았다.^[5]

구현

adaLN-Zero 블록

공식 구현의 블록을 줄인 것이다. PyTorch 2.5에서 실행해 초기화 직후 블록 출력이 입력과 정확히 같은 것(항등 함수)을 확인했다.

import torch, torch.nn as nn

def modulate(x, shift, scale):
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

class DiTBlock(nn.Module):
    def __init__(self, d, heads, mlp_ratio=4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
        self.attn = nn.MultiheadAttention(d, heads, batch_first=True)
        self.norm2 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
        self.mlp = nn.Sequential(nn.Linear(d, int(d * mlp_ratio)), nn.GELU(approximate="tanh"),
                                 nn.Linear(int(d * mlp_ratio), d))
        # 조건 c(시간 + 클래스 임베딩)에서 shift, scale, gate를 두 벌 회귀
        self.adaLN = nn.Sequential(nn.SiLU(), nn.Linear(d, 6 * d))
        nn.init.zeros_(self.adaLN[-1].weight)   # adaLN-Zero: 처음에는 gate = 0
        nn.init.zeros_(self.adaLN[-1].bias)

    def forward(self, x, c):                    # x: (B, T, d) 패치 토큰, c: (B, d)
        sh1, sc1, g1, sh2, sc2, g2 = self.adaLN(c).chunk(6, dim=1)
        h = modulate(self.norm1(x), sh1, sc1)
        x = x + g1.unsqueeze(1) * self.attn(h, h, h, need_weights=False)[0]
        h = modulate(self.norm2(x), sh2, sc2)
        x = x + g2.unsqueeze(1) * self.mlp(h)
        return x

blk = DiTBlock(384, 6)                          # DiT-S 크기
x = torch.randn(2, 256, 384)                    # 32x32 잠재, p=2 -> 256 토큰
c = torch.randn(2, 384)
assert torch.equal(blk(x, c), x)                # 초기 상태는 항등 함수

패치화는 커널 크기와 보폭이 모두 $p$ 인 합성곱 하나(nn.Conv2d(C, d, kernel_size=p, stride=p))로 구현하는 것이 보통이며, 공식 코드는 timm의 PatchEmbed를 썼다.^[2]

diffusers로 추론

import torch
from diffusers import DiTPipeline, DPMSolverMultistepScheduler

pipe = DiTPipeline.from_pretrained("facebook/DiT-XL-2-256", torch_dtype=torch.float16).to("cuda")
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
class_ids = pipe.get_label_ids(["white shark", "umbrella"])   # ImageNet 클래스 이름 -> 번호
images = pipe(class_labels=class_ids, num_inference_steps=25, guidance_scale=4.0).images

DiTPipeline의 guidance_scale 기본값은 4.0이다.^[8]

U-Net과 비교

항목	U-Net 백본	확산 트랜스포머
귀납적 편향	국소성, 다중 해상도, 스킵 연결	거의 없음(패치와 위치 임베딩뿐)
확장 방법	채널·블록 수, 어텐션 해상도 등 손볼 곳이 많음	층 수·폭·패치 크기로 단순하게 확장
연산량	고해상도 합성곱 층에 집중	토큰 수의 제곱(어텐션)에 비례
조건 주입	교차 어텐션, 시간 임베딩 더하기	adaLN, 교차 어텐션, 공동 어텐션(MMDiT)
다른 모달리티 확장	영상은 시간 축 층 추가 필요	시공간 패치로 영상·오디오도 같은 구조
생태계	스테이블 디퓨전 1·2, SDXL	스테이블 디퓨전 3, FLUX.1, 소라, 무비 젠

트랜스포머 쪽의 이점은 대형 언어 모델에서 쌓인 기술, 곧 플래시 어텐션, 텐서 병렬화, 혼합 정밀도 학습, RoPE 같은 위치 임베딩을 그대로 가져다 쓸 수 있다는 점이다. 단점은 토큰 수에 대해 어텐션 비용이 제곱으로 늘어 고해상도·장시간 영상에서 시퀀스가 매우 길어진다는 것이다. 그래서 더 강하게 압축하는 오토인코더와 큰 패치가 함께 연구된다.

같이 보기

각주

↑ ^1.0 ^1.1 ^1.2 ^1.3 ^1.4 ^1.5 ^1.6 ^1.7 William Peebles, Saining Xie, 〈Scalable Diffusion Models with Transformers〉, ICCV 2023. arXiv:2212.09748
↑ ^2.0 ^2.1 facebookresearch, 〈DiT models.py〉. github.com
↑ Junsong Chen 외, 〈PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis〉, ICLR 2024. arXiv:2310.00426
↑ OpenAI, 〈Video generation models as world simulators〉, 2024년 2월 15일. openai.com
↑ ^5.0 ^5.1 Patrick Esser 외, 〈Scaling Rectified Flow Transformers for High-Resolution Image Synthesis〉, ICML 2024. arXiv:2403.03206
↑ Black Forest Labs, 〈Announcing Black Forest Labs〉, 2024년 8월 1일. bfl.ai
↑ Adam Polyak 외, 〈Movie Gen: A Cast of Media Foundation Models〉, 2024. arXiv:2410.13720
↑ Hugging Face, 〈pipeline_dit.py〉, diffusers 소스 코드. github.com

참고 문헌

[dit-1] 1.0 ^1.1 ^1.2 ^1.3 ^1.4 ^1.5 ^1.6 ^1.7 William Peebles, Saining Xie, 〈Scalable Diffusion Models with Transformers〉, ICCV 2023. arXiv:2212.09748

[ditcode-2] 2.0 ^2.1 facebookresearch, 〈DiT models.py〉. github.com

[pixart-3] Junsong Chen 외, 〈PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis〉, ICLR 2024. arXiv:2310.00426

[sora-4] OpenAI, 〈Video generation models as world simulators〉, 2024년 2월 15일. openai.com

[sd3-5] 5.0 ^5.1 Patrick Esser 외, 〈Scaling Rectified Flow Transformers for High-Resolution Image Synthesis〉, ICML 2024. arXiv:2403.03206

[bfl-6] Black Forest Labs, 〈Announcing Black Forest Labs〉, 2024년 8월 1일. bfl.ai

[moviegen-7] Adam Polyak 외, 〈Movie Gen: A Cast of Media Foundation Models〉, 2024. arXiv:2410.13720

[ditpipe-8] Hugging Face, 〈pipeline_dit.py〉, diffusers 소스 코드. github.com

[1]

[2]

[3]

[4]

[5]

[6]

[7]

[8]