디퓨전 트랜스포머(Diffusion Transformers)와 정류 흐름(Rectified Flow)

디퓨전(diffusion)의 비밀은 U-Net이 아닙니다. U-Net을 트랜스포머(transformer)로 바꾸고, 노이즈 스케줄(noise schedule)을 직선 흐름(straight-line flow)으로 바꾸면 SD3, FLUX, 그리고 2026년 텍스트-투-이미지(text-to-image) 모델의 대부분이 그 자리에서 모습을 드러냅니다.

유형: Learn + Build 언어: Python 선수 조건: Phase 4 Lesson 10(Diffusion DDPM), Phase 4 Lesson 14(ViT), Phase 7 Lesson 02(Self-Attention) 소요 시간: 약 75분

학습 목표

• U-Net DDPM(Lesson 10)에서 디퓨전 트랜스포머(Diffusion Transformer; DiT), MMDiT(SD3), 그리고 단일/이중 스트림(single+double-stream) DiT(FLUX)로 이어지는 변화를 추적합니다.
• 정류 흐름(rectified flow)을 설명하고, 노이즈와 데이터 사이의 직선 궤적(straight-line trajectory)이 왜 1000 스텝(step) 대신 20 스텝 샘플링(sampling)을 가능하게 하는지 이해합니다.
• 100줄 미만의 작은 DiT 블록(block)과 정류 흐름 학습 루프(training loop)를 구현합니다.
• SD3, FLUX.1-dev, FLUX.1-schnell, Z-Image, Qwen-Image 같은 모델 변형(model variant)을 아키텍처(architecture), 파라미터 수(parameter count), 라이선스(license) 기준으로 구분합니다.

문제

Lesson 10에서는 U-Net 디노이저(denoiser)를 가진 DDPM을 만들었습니다. 이 레시피(recipe)는 2020-2023년을 지배했습니다. U-Net + 베타 스케줄(beta schedule) + 노이즈 예측 손실(noise-prediction loss) 조합입니다. Stable Diffusion 1.5와 2.1, 그리고 DALL-E 2도 모두 이 계열입니다.

2026년의 모든 최신(state-of-the-art) 텍스트-투-이미지 모델은 이 지점을 지나왔습니다. Stable Diffusion 3, FLUX, SD4, Z-Image, Qwen-Image, Hunyuan-Image는 U-Net을 사용하지 않고 디퓨전 트랜스포머(DiT)를 사용합니다. SD3와 FLUX는 DDPM 노이즈 스케줄까지 정류 흐름(rectified flow)으로 바꿉니다. 이 방식은 노이즈에서 데이터로 가는 경로(path)를 곧게 펴 주기 때문에, 일관성(consistency) 모델이나 증류된(distilled) 변형에서는 1-4 스텝 추론(inference)까지 가능해집니다.

이 변화가 중요한 이유는 디퓨전 기반 이미지 생성(image generation)이 더 잘 제어 가능(controllable)하고, 프롬프트 정확도(prompt-accurate)도 높으며, 특히 SD3/SD4가 텍스트 렌더링(text rendering) 문제를 해결한 비결, 그리고 프로덕션(production)에서도 빠르게 동작하게 된 비결이 모두 여기에 있기 때문입니다. DiT와 정류 흐름을 이해한다는 것은 곧 2026년의 생성형 이미지 스택(generative-image stack)을 이해한다는 뜻입니다.

개념

U-Net에서 트랜스포머로

flowchart LR
    subgraph UNET["DDPM U-Net (2020)"]
        U1["Conv encoder"] --> U2["Conv bottleneck"] --> U3["Conv decoder"]
    end
    subgraph DIT["DiT (2023)"]
        D1["Patch embed"] --> D2["Transformer blocks"] --> D3["Unpatchify"]
    end
    subgraph MMDIT["MMDiT (SD3, 2024)"]
        M1["Text stream"] --> M3["Joint attention<br/>(modality별 별도 weight)"]
        M2["Image stream"] --> M3
    end
    subgraph FLUX["FLUX (2024)"]
        F1["Double-stream blocks<br/>(text + image 분리)"] --> F2["Single-stream blocks<br/>(concat + shared weights)"]
    end

    style UNET fill:#e5e7eb,stroke:#6b7280
    style DIT fill:#dbeafe,stroke:#2563eb
    style MMDIT fill:#fef3c7,stroke:#d97706
    style FLUX fill:#dcfce7,stroke:#16a34a

• DiT(Peebles & Xie, 2023): U-Net을 잠재 패치(latent patch) 위에서 동작하는 ViT 형태의 트랜스포머로 대체합니다. 조건 부여(conditioning)는 적응형 레이어 정규화(adaptive layer norm; AdaLN)를 사용합니다.
• MMDiT(SD3, Esser et al., 2024): 텍스트 토큰(text token)과 이미지 토큰(image token)에 별도 가중치 스트림(weight stream)을 두되, 공동 어텐션(joint attention)을 공유합니다.
• FLUX(Black Forest Labs, 2024): 앞쪽 N개의 블록은 SD3처럼 이중 스트림(double-stream)이고, 뒤쪽 블록은 토큰을 이어 붙여(concatenate) 가중치를 공유하는(shared weight) 단일 스트림(single-stream)으로 효율을 높입니다.
• Z-Image(2025): 6B 파라미터로 동작하는 효율적인 단일 스트림 DiT입니다. "무조건 규모 확장"이라는 접근에 정면으로 도전합니다.

정류 흐름을 한 문단으로 정리

DDPM은 순방향 과정(forward process)을 x_t가 점점 오염되는 잡음 섞인 SDE(noisy SDE)로 정의합니다. 학습된 역방향(reverse) 과정 역시 또 하나의 SDE이며, 1000개의 작은 스텝으로 풀어냅니다.

정류 흐름은 깨끗한 데이터와 순수 노이즈 사이를 잇는 직선(straight-line) 보간(interpolation)을 정의합니다.

x_t = (1 - t) * x_0 + t * epsilon,     t in [0, 1]

신경망(network)은 속도(velocity) v_theta(x_t, t) = epsilon - x_0를 예측하도록 학습합니다. 이 값은 깨끗한 데이터에서 노이즈로 향하는 직선 경로의 순방향(forward direction), 다시 말해 dx_t/dt에 해당합니다. 샘플링(sampling) 단계에서는 이 속도를 반대로 적분(integrate)하여 노이즈에서 데이터 방향으로 이동합니다. 결과로 얻어지는 ODE(ordinary differential equation)는 훨씬 직선에 가깝기 때문에 샘플을 만들 때 필요한 적분 스텝(integration step)이 크게 줄어듭니다.

SD3는 이를 정류 흐름 매칭(Rectified Flow Matching)이라고 부릅니다. FLUX, Z-Image를 비롯한 2026년의 많은 모델이 같은 목적 함수(objective)를 사용합니다. 일반적인 추론에서는 기존 DDPM 방식의 50개 이상의 DDIM 스텝 대신 결정적(deterministic) 오일러(Euler) 적분 20-30 스텝이면 충분합니다. 증류된(distilled) / 터보(turbo) / 슈넬(schnell) / LCM 변형은 이를 1-4 스텝까지 줄입니다.

AdaLN 조건 부여

DiT는 타임스텝(timestep)과 클래스/텍스트 조건 부여(class/text conditioning)를 적응형 레이어 정규화(adaptive layer norm)로 처리합니다. 조건 벡터(conditioning vector)에서 scale과 shift를 예측한 뒤 LayerNorm 뒤에 적용합니다. U-Net에서 자주 쓰던 FiLM 스타일의 변조(modulation)보다 훨씬 깔끔하며, 오늘날 모든 현대 DiT의 표준 설계입니다.

cond -> MLP -> (scale, shift, gate)
norm(x) * (1 + scale) + shift, then residual add * gate

SD3와 FLUX의 텍스트 인코더(text encoder)

• SD3는 세 개의 텍스트 인코더를 사용합니다. 두 개의 CLIP 모델과 T5-XXL입니다. 임베딩(embedding)은 이어 붙여진 뒤 텍스트 조건(text conditioning)으로 이미지 스트림에 들어갑니다.
• FLUX는 CLIP-L 하나와 T5-XXL을 사용합니다.
• Qwen-Image / Z-Image 변형은 자신의 기반 LLM과 정렬(aligned)된 자체 텍스트 인코더를 사용합니다.

텍스트 인코더는 SD3/FLUX가 SD1.5보다 프롬프트를 훨씬 잘 이해하는 큰 이유입니다. T5-XXL 하나만 해도 47억(4.7B) 파라미터에 달합니다.

분류기 없는 가이던스(classifier-free guidance)는 그대로 유지됩니다

정류 흐름은 샘플러(sampler)를 바꾸는 것이지 조건 부여 방식을 바꾸는 것이 아닙니다. 분류기 없는 가이던스(classifier-free guidance; CFG), 즉 학습 중 10% 확률로 텍스트를 드롭(drop)하고 추론에서 조건부 예측(conditional prediction)과 비조건부 예측(unconditional prediction)을 섞는 방식은 정류 흐름에서도 그대로 동작합니다. 2026년 모델은 가이던스 스케일(guidance scale) 3.5-5 정도를 많이 사용합니다. SD1.5의 7.5보다 낮은 이유는 정류 흐름 모델이 기본적으로 프롬프트를 훨씬 강하게 따르기 때문입니다.

Consistency, Turbo, Schnell, LCM

네 이름 모두 같은 아이디어를 가리킵니다. 느리고 많은 스텝이 필요한(many-step) 모델을 빠르고 적은 스텝이면 충분한(few-step) 모델로 증류(distil)하는 작업입니다.

• LCM(Latent Consistency Model; 잠재 일관성 모델): 임의의 중간(intermediate) x_t에서도 최종 x_0를 한 스텝에 예측하는 학생(student) 모델을 학습합니다.
• SDXL Turbo / FLUX schnell: 적대적 디퓨전 증류(adversarial diffusion distillation)로 학습된 1-4 스텝 모델입니다.
• SD Turbo: 잠재 디퓨전(latent diffusion)에 맞게 변형한 OpenAI 계열의 일관성 모델(Consistency Models)입니다.

새 모델을 프로덕션에서 서빙(serving)할 때는 "최고 품질(full quality)" 체크포인트(checkpoint)와 "터보 / 슈넬(turbo / schnell)" 변형을 함께 제공합니다. 슈넬(schnell)은 독일어로 "빠름"을 뜻하며 Black Forest Labs의 명명 규칙(convention)입니다. 1-4 스텝으로 실행되어 실시간(real-time) 파이프라인(pipeline)에 적합합니다.

2026년 모델 지형도(model landscape)

FLUX.1-schnell은 2026년 오픈소스 분야의 기본 선택지(open-source default)입니다. Z-Image는 효율성에서 앞서가는 모델(efficiency leader)이고, FLUX.2와 SD4는 현재 품질 측면에서 가장 앞선 모델(quality tip)입니다.

왜 이 패러다임 전환(phase shift)이 중요한가

DDPM + U-Net 조합도 충분히 잘 동작했습니다. 하지만 DiT + 정류 흐름은 더 잘, 더 빠르게, 더 깔끔하게 규모를 키울(scale) 수 있습니다. 이 전환은 NLP에서 RNN이 트랜스포머로 대체된 변화와 비슷합니다. 두 아키텍처 모두 같은 문제를 풀었지만, 결국 트랜스포머가 규모 확장에 성공하면서 주류를 차지했습니다. 2026년에 발표되는 이미지·비디오·3D 생성 논문은 거의 모두 DiT 형태의 디노이저(DiT-shaped denoiser)를 사용하고, 대부분 정류 흐름 목적 함수를 채택합니다. U-Net DDPM은 이제 주로 학습용 교재(Lesson 10)로 남았습니다.

만들어 보기

Step 1: AdaLN을 갖춘 DiT 블록

import torch
import torch.nn as nn


class AdaLNZero(nn.Module):
    """
    게이트(gate)가 포함된 적응형 레이어 정규화입니다.
    조건 벡터에서 (scale, shift, gate)를 예측합니다.
    전체 블록이 항등(identity) 매핑으로 시작하도록 가중치를 0으로 초기화합니다("zero init").
    """

    def __init__(self, dim, cond_dim):
        super().__init__()
        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
        self.mlp = nn.Linear(cond_dim, dim * 3)
        nn.init.zeros_(self.mlp.weight)
        nn.init.zeros_(self.mlp.bias)

    def forward(self, x, cond):
        scale, shift, gate = self.mlp(cond).chunk(3, dim=-1)
        h = self.norm(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
        return h, gate.unsqueeze(1)


class DiTBlock(nn.Module):
    def __init__(self, dim=192, heads=3, mlp_ratio=4, cond_dim=192):
        super().__init__()
        self.adaln1 = AdaLNZero(dim, cond_dim)
        self.attn = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.adaln2 = AdaLNZero(dim, cond_dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, dim * mlp_ratio),
            nn.GELU(),
            nn.Linear(dim * mlp_ratio, dim),
        )

    def forward(self, x, cond):
        h, gate1 = self.adaln1(x, cond)
        a, _ = self.attn(h, h, h, need_weights=False)
        x = x + gate1 * a
        h, gate2 = self.adaln2(x, cond)
        x = x + gate2 * self.mlp(h)
        return x

AdaLNZero는 MLP 가중치가 0으로 초기화되어 있어 항등 매핑으로 학습을 시작합니다. 학습이 진행되면서 블록은 항등 매핑에서 서서히 멀어집니다. 이 방식은 깊은 트랜스포머 기반 디퓨전 모델을 크게 안정화해 줍니다.

Step 2: 작은 DiT(TinyDiT)

def timestep_embedding(t, dim):
    import math
    half = dim // 2
    freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half)
    args = t[:, None].float() * freqs[None]
    return torch.cat([args.sin(), args.cos()], dim=-1)


class TinyDiT(nn.Module):
    def __init__(self, image_size=16, patch_size=2, in_channels=3, dim=96, depth=4, heads=3):
        super().__init__()
        self.patch_size = patch_size
        self.num_patches = (image_size // patch_size) ** 2
        self.patch = nn.Conv2d(in_channels, dim, kernel_size=patch_size, stride=patch_size)
        self.pos = nn.Parameter(torch.zeros(1, self.num_patches, dim))
        self.time_mlp = nn.Sequential(
            nn.Linear(dim, dim * 2),
            nn.SiLU(),
            nn.Linear(dim * 2, dim),
        )
        self.blocks = nn.ModuleList([DiTBlock(dim, heads, cond_dim=dim) for _ in range(depth)])
        self.norm_out = nn.LayerNorm(dim, elementwise_affine=False)
        self.head = nn.Linear(dim, patch_size * patch_size * in_channels)

    def forward(self, x, t):
        n = x.size(0)
        x = self.patch(x)
        x = x.flatten(2).transpose(1, 2) + self.pos
        t_emb = self.time_mlp(timestep_embedding(t, self.pos.size(-1)))
        for blk in self.blocks:
            x = blk(x, t_emb)
        x = self.norm_out(x)
        x = self.head(x)
        return self._unpatchify(x, n)

    def _unpatchify(self, x, n):
        p = self.patch_size
        h = w = int(self.num_patches ** 0.5)
        x = x.view(n, h, w, p, p, -1).permute(0, 5, 1, 3, 2, 4).reshape(n, -1, h * p, w * p)
        return x

Step 3: 정류 흐름 학습(rectified flow training)

import torch.nn.functional as F

def rectified_flow_train_step(model, x0, optimizer, device):
    model.train()
    x0 = x0.to(device)
    n = x0.size(0)
    t = torch.rand(n, device=device)
    epsilon = torch.randn_like(x0)
    x_t = (1 - t[:, None, None, None]) * x0 + t[:, None, None, None] * epsilon

    target_velocity = epsilon - x0
    pred_velocity = model(x_t, t)

    loss = F.mse_loss(pred_velocity, target_velocity)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

DDPM의 노이즈 예측 손실(Lesson 10)과 비교하면 구조는 같고 학습 목표(target)만 다릅니다. 노이즈 epsilon을 예측하는 대신, 직선 보간을 따라 데이터에서 노이즈로 향하는 속도 epsilon - x_0를 예측합니다.

Step 4: 오일러 샘플러(Euler sampler)

정류 흐름은 ODE입니다. 오일러 방법(Euler method)은 가장 단순한 풀이법이며, 잘 학습된 정류 흐름 모델에서는 20 스텝 이상에서 고차(higher-order) 솔버(solver)와 거의 비슷한 정확도를 보여 줍니다.

@torch.no_grad()
def rectified_flow_sample(model, shape, steps=20, device="cpu"):
    model.eval()
    x = torch.randn(shape, device=device)
    dt = 1.0 / steps
    t = torch.ones(shape[0], device=device)
    for _ in range(steps):
        v = model(x, t)
        x = x - dt * v
        t = t - dt
    return x

단 20 스텝입니다. 학습된 모델에서는 1000 스텝짜리 DDPM과 비교할 만한 샘플을 만들어 냅니다.

Step 5: 엔드 투 엔드 동작 점검(end-to-end smoke test)

import numpy as np

def synthetic_blobs(num=200, size=16, seed=0):
    rng = np.random.default_rng(seed)
    out = np.zeros((num, 3, size, size), dtype=np.float32)
    yy, xx = np.meshgrid(np.arange(size), np.arange(size), indexing="ij")
    for i in range(num):
        cx, cy = rng.uniform(4, size - 4, size=2)
        r = rng.uniform(2, 4)
        mask = (xx - cx) ** 2 + (yy - cy) ** 2 < r ** 2
        colour = rng.uniform(-1, 1, size=3)
        for c in range(3):
            out[i, c][mask] = colour[c]
    return torch.from_numpy(out)

이 합성 데이터에 TinyDiT를 정류 흐름으로 학습시킵니다. 500 스텝이 지난 뒤 샘플 출력은 희미한 색 얼룩(colour blob) 형태로 나와야 합니다.

사용하기

실제 FLUX / SD3 / Z-Image 이미지 생성에는 diffusers 라이브러리가 통합 API를 제공합니다.

from diffusers import FluxPipeline, StableDiffusion3Pipeline
import torch

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    torch_dtype=torch.bfloat16,
).to("cuda")

out = pipe(
    prompt="a golden retriever surfing a tsunami, hyperrealistic, studio lighting",
    guidance_scale=0.0,           # schnell은 CFG 없이 학습되었습니다.
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]
out.save("surf.png")

단 세 줄입니다. FLUX.1-schnell을 네 스텝으로 실행합니다. 더 높은 품질이 필요하다면 모델 id를 black-forest-labs/FLUX.1-dev로 바꾸고 CFG를 켠 채 20-30 스텝을 사용합니다.

SD3는 다음과 같이 사용합니다.

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-large",
    torch_dtype=torch.bfloat16,
).to("cuda")
out = pipe(prompt, guidance_scale=3.5, num_inference_steps=28).images[0]

산출물 만들기

이 레슨에서는 다음을 만듭니다.

• outputs/prompt-dit-model-picker.md: 품질(quality), 지연 시간(latency), 라이선스(license) 제약에 따라 SD3, FLUX.1-dev, FLUX.1-schnell, Z-Image, SD4 Turbo 중 하나를 골라 주는 프롬프트입니다.
• outputs/skill-rectified-flow-trainer.md: AdaLN DiT와 오일러 샘플링(Euler sampling)을 포함한 정류 흐름 학습 루프를 작성하는 스킬(skill)입니다.

연습문제

1. (쉬움) 위에서 만든 TinyDiT를 합성 얼룩(synthetic blob) 데이터셋으로 500 스텝 학습합니다. 10, 20, 50 오일러 스텝으로 생성한 샘플을 비교합니다.
2. (중간) 학습된 클래스 임베딩(learned class embedding)을 시간 임베딩(time embedding)에 이어 붙여 텍스트 조건 부여를 추가합니다. 색상 기준으로 10개의 얼룩 "클래스"를 정의한 뒤 클래스 0, 5, 9로 샘플링하여 색상이 일치하는지 확인합니다.
3. (어려움) 같은 크기의 신경망을 정류 흐름 버전과 DDPM 버전으로 같은 데이터, 같은 스텝 수로 학습합니다. 생성된 샘플 사이의 프레셰 거리(Fréchet distance; FID 근사)를 계산하고 어느 쪽이 더 빨리 수렴하는지 보고합니다.

핵심 용어

더 읽을거리

• Scalable Diffusion Models with Transformers (Peebles & Xie, 2023): DiT 원 논문입니다.
• Scaling Rectified Flow Transformers (Esser et al., SD3 paper): MMDiT와 정류 흐름을 대규모로 끌어올린 논문입니다.
• FLUX.1 model card and technical report (Black Forest Labs): 이중/단일 스트림 구조의 세부 사항을 확인할 수 있습니다.
• Z-Image: Efficient Image Generation Foundation Model (2025): 6B 규모의 단일 스트림 DiT 소개 자료입니다.
• Elucidating the Design Space of Diffusion (Karras et al., 2022): 디퓨전 설계의 교환 관계(trade-off)를 다룬 기준 문헌입니다.
• Latent Consistency Models (Luo et al., 2023): LCM-LoRA로 4 스텝 추론을 구성하는 방법을 설명합니다.