Diffusion Transformer(DiT)

1. DiT란?

DiT(Diffusion Transformer)는 확산 모델에서 기존 U-Net 대신 Transformer를 사용하는 구조다.

확산 모델은 원본 이미지에 노이즈를 추가하고, 모델이 그 노이즈를 예측하도록 학습한다.

원본 이미지 x₀
→ 노이즈 추가
→ noisy image xₜ
→ DiT
→ 예측 노이즈 ε̂

학습 손실은 보통 실제 노이즈와 예측 노이즈의 MSE를 사용한다.

loss = F.mse_loss(predicted_noise, noise)

즉, DiT는 확산 과정 전체가 아니라 노이즈를 예측하는 Backbone이다.

---

2. U-Net과 DiT 차이

구분	U-Net	DiT
구조	CNN 기반	Transformer 기반
입력 처리	Feature Map	Patch Token
전역 정보	계층적으로 처리	Self-Attention
확장 방식	채널과 Block 증가	Depth, Width, Head 증가

DiT는 Transformer의 확장성을 활용하기 쉽지만, token 수가 많아지면 연산량도 크게 증가한다.

---

3. DiT 구조

Noisy Image 또는 Latent
        ↓
Patch Embedding
        ↓
Position Embedding 추가
        ↓
Timestep Embedding 생성
        ↓
DiT Transformer Blocks
        ↓
Patch 복원
        ↓
Predicted Noise

Patch Embedding

이미지를 작은 패치로 나누고 token으로 변환한다.

예를 들어 32×32 이미지에 4×4 패치를 사용하면:

(32 / 4) × (32 / 4)
= 64 Tokens

PyTorch에서는 Conv2d로 구현할 수 있다.

self.patch_embed = nn.Conv2d(
    in_channels=3,
    out_channels=hidden_size,
    kernel_size=patch_size,
    stride=patch_size,
)

Timestep Embedding

현재 이미지에 노이즈가 얼마나 섞였는지 모델에 알려준다.

작은 timestep → 노이즈가 적음
큰 timestep   → 노이즈가 많음

Self-Attention

각 패치가 이미지의 다른 패치를 참고한다.

예를 들어 머리카락 패치가 얼굴, 눈, 배경 패치를 함께 참고할 수 있다.

다만 Attention 연산량은 일반적으로 token 수에 대해 다음과 같이 증가한다.

O(N²)

---

4. adaLN-Zero

DiT는 timestep 조건을 Transformer Block에 넣기 위해 adaLN-Zero를 사용한다.

기본 형태는 다음과 같다.

output = normalized_x × (1 + scale) + shift

Attention과 MLP 출력에는 gate를 적용한다.

x = x + gate * attention_output

초기 gate를 0으로 설정해 학습 초기에 Block이 입력을 크게 변경하지 않도록 한다.

---

5. PyTorch 간단 구현

import math

import torch
import torch.nn.functional as F
from torch import nn


def timestep_embedding(t, dim):
    half = dim // 2

    frequencies = torch.exp(
        -math.log(10000)
        * torch.arange(
            half,
            device=t.device,
            dtype=torch.float32,
        )
        / half
    )

    angles = t.float()[:, None] * frequencies[None]

    return torch.cat(
        [
            torch.cos(angles),
            torch.sin(angles),
        ],
        dim=-1,
    )


class DiTBlock(nn.Module):
    def __init__(self, hidden_size, num_heads):
        super().__init__()

        self.norm1 = nn.LayerNorm(hidden_size)

        self.attention = nn.MultiheadAttention(
            hidden_size,
            num_heads,
            batch_first=True,
        )

        self.norm2 = nn.LayerNorm(hidden_size)

        self.mlp = nn.Sequential(
            nn.Linear(hidden_size, hidden_size * 4),
            nn.GELU(),
            nn.Linear(hidden_size * 4, hidden_size),
        )

    def forward(self, x):
        hidden = self.norm1(x)

        hidden, _ = self.attention(
            hidden,
            hidden,
            hidden,
            need_weights=False,
        )

        x = x + hidden
        x = x + self.mlp(self.norm2(x))

        return x


class MiniDiT(nn.Module):
    def __init__(
        self,
        image_size=32,
        patch_size=4,
        hidden_size=256,
        depth=4,
        num_heads=4,
    ):
        super().__init__()

        self.patch_size = patch_size
        self.grid_size = image_size // patch_size

        num_patches = self.grid_size ** 2

        self.patch_embed = nn.Conv2d(
            3,
            hidden_size,
            kernel_size=patch_size,
            stride=patch_size,
        )

        self.position_embedding = nn.Parameter(
            torch.zeros(
                1,
                num_patches,
                hidden_size,
            )
        )

        self.time_mlp = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size),
        )

        self.blocks = nn.ModuleList(
            [
                DiTBlock(
                    hidden_size,
                    num_heads,
                )
                for _ in range(depth)
            ]
        )

        self.output = nn.Linear(
            hidden_size,
            patch_size * patch_size * 3,
        )

    def forward(self, x, timestep):
        batch_size = x.shape[0]

        x = self.patch_embed(x)
        x = x.flatten(2).transpose(1, 2)

        x = x + self.position_embedding

        time_condition = timestep_embedding(
            timestep,
            x.shape[-1],
        )

        time_condition = self.time_mlp(
            time_condition
        )

        x = x + time_condition[:, None, :]

        for block in self.blocks:
            x = block(x)

        x = self.output(x)

        x = x.view(
            batch_size,
            self.grid_size,
            self.grid_size,
            self.patch_size,
            self.patch_size,
            3,
        )

        x = torch.einsum(
            "bhwpqc->bchpwq",
            x,
        )

        return x.reshape(
            batch_size,
            3,
            self.grid_size * self.patch_size,
            self.grid_size * self.patch_size,
        )

간단한 학습 흐름은 다음과 같다.

model = MiniDiT()

x0 = torch.randn(2, 3, 32, 32)
noise = torch.randn_like(x0)

timestep = torch.randint(
    0,
    1000,
    (2,),
)

noisy_image = x0 + noise

predicted_noise = model(
    noisy_image,
    timestep,
)

loss = F.mse_loss(
    predicted_noise,
    noise,
)

loss.backward()

실제 diffusion 학습에서는 x0 + noise 대신 timestep에 맞는 noise scheduler를 사용해야 한다.

---

6. Latent DiT

실제 대규모 모델은 RGB 이미지를 직접 처리하기보다 VAE latent에서 DiT를 실행한다.

Image
→ VAE Encoder
→ Latent
→ DiT
→ VAE Decoder
→ Generated Image

latent를 사용하면 이미지 크기와 token 수를 줄일 수 있다.

---

7. 영상 DiT

영상에서는 시간축이 추가된다.

이미지: [B, C, H, W]
영상:   [B, C, T, H, W]

일반적으로 공간 방향의 Spatial Attention과 시간 방향의 Temporal Attention을 함께 사용한다.

Video
→ Spatial Patch
→ Spatial Attention
→ Temporal Attention
→ Output Video

영상은 프레임 수만큼 token이 증가하기 때문에 이미지 DiT보다 훨씬 많은 GPU 메모리가 필요하다.

---

정리

DiT는 확산 모델의 노이즈 예측 Backbone을 Transformer로 구현한 구조다.

DiT
=
Diffusion
+
Transformer

핵심 과정은 다음과 같다.

이미지 또는 latent
→ Patch Token 변환
→ Timestep 조건 추가
→ Transformer 처리
→ 이미지 형태 복원
→ 노이즈 예측

Transformer의 확장성과 전역 정보 처리 능력 때문에 이미지뿐 아니라 최근 영상 생성 모델에서도 많이 사용되고 있다.