U-Net의 독주는 끝났습니다. 2025년 생성형 AI 혁명을 이끄는 새로운 아키텍처를 소개합니다.
수년 동안 U-Net은 디퓨전(Diffusion) 모델의 절대적인 제왕이었습니다. Stable Diffusion 1.5부터 초기 비디오 모델에 이르기까지, 컨볼루션(Convolutional) 네트워크가 이 분야를 지배했습니다. 하지만 Sora 2와 FLUX.2의 등장은 새로운 현실을 확정지었습니다. 미래는 이제 Diffusion Transformer (DiT)의 것입니다.
우리는 지금 AI의 “두뇌 이식” 현장을 목격하고 있습니다. 컨볼루션 백본을 트랜스포머(Transformer)로 교체함으로써, 모델들은 컴퓨팅 자원에 비례해 예측 가능한 확장을 할 수 있게 되었고, 물리학을 이해하며, 비디오 내에서 전례 없는 시간적 일관성(temporal consistency)을 유지하게 되었습니다.
왜 U-Net은 실패했고, DiT가 승리했는가
DiT를 이해하려면 먼저 무엇이 대체되고 있는지 알아야 합니다. U-Net 아키텍처(SDXL 및 구형 모델에서 사용)는 컨볼루션(Convolutions)에 의존합니다. 컨볼루션은 국소적인 디테일(가장자리, 질감)을 처리하는 데는 뛰어나지만, 전역적인 문맥(Global Context)을 파악하는 데는 어려움을 겪습니다. 이미지를 한 번에 “통으로” 보는 능력이 부족하기 때문에, 물체가 분리되거나 비디오에서 형태가 일그러지는 현상이 발생하곤 했습니다.
Diffusion Transformer (DiT)는 GPT-4의 아면 텍처(Transformer)를 시각적 생성에 적용한 것입니다. 다음 단어를 예측하는 대신, 이미지나 비디오 프레임의 다음 “패치(patch)”를 예측하는 방식입니다.
DiT의 결정적 우위:
- 확장성 (Scalability): CNN과 달리, 트랜스포머는 파라미터와 데이터를 늘릴수록 성능이 선형적으로 똑똑해집니다 (스케일링 법칙).
- 전역 어텐션 (Global Attention): 이미지의 모든 패치가 다른 모든 패치와 즉각적으로 “소통”할 수 있습니다. 왼쪽 끝에 있는 손이 오른쪽 끝에 있는 팔이 무엇을 하고 있는지 정확히 알고 있는 셈입니다.
- 문맥 인식 (Context Awareness): 예를 들어 FLUX.2는 24B 규모의 비전-언어 모델(Mistral-3)을 DiT와 결합하여, 단순한 키워드가 아니라 복잡한 프롬프트와 물리학적 상황을 이해합니다.
코드 분석: 간단한 DiT 블록 구현
Sora 2와 FLUX.2의 핵심은 마법이 아닙니다. 바로 이 블록들이 쌓여있는 적층 구조입니다. 아래는 단일 DiT 블록의 PyTorch 구현 예시입니다. 여기서 AdaLN (Adaptive Layer Norm)의 사용을 주목하세요. 이것은 타임스텝(timestep)과 텍스트 프롬프트를 생성 과정에 주입하는 “조종간” 역할을 합니다.
import torch
import torch.nn as nn
class DiTBlock(nn.Module):
"""
단일 Diffusion Transformer 블록.
인자:
hidden_size (int): 트랜스포머 토큰의 차원 크기.
num_heads (int): 어텐션 헤드의 개수.
mlp_ratio (float): MLP 내 히든 차원의 배수.
"""
def __init__(self, hidden_size, num_heads, mlp_ratio=4.0):
super().__init__()
self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.attn = nn.MultiheadAttention(hidden_size, num_heads, batch_first=True)
self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
# 피드 포워드 네트워크 (MLP)
mlp_hidden_dim = int(hidden_size * mlp_ratio)
self.mlp = nn.Sequential(
nn.Linear(hidden_size, mlp_hidden_dim),
nn.GELU(),
nn.Linear(mlp_hidden_dim, hidden_size),
)
# Adaptive Layer Norm (AdaLN) 변조
# 컨디셔닝(시간/라벨)에 기반하여 shift(감마)와 scale(베타)를 예측
self.adaLN_modulation = nn.Sequential(
nn.SiLU(),
nn.Linear(hidden_size, 6 * hidden_size, bias=True)
)
def forward(self, x, c):
"""
x: 입력 토큰 [Batch, Sequence_Length, Hidden_Size]
c: 컨디셔닝 임베딩 (시간 + 텍스트) [Batch, Hidden_Size]
"""
# 1. 컨디셔닝으로부터 변조 파라미터 회귀(Regression)
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
self.adaLN_modulation(c).chunk(6, dim=1)
)
# 2. AdaLN을 적용한 Self-Attention 블록
x_norm1 = (1 + scale_msa.unsqueeze(1)) * self.norm1(x) + shift_msa.unsqueeze(1)
attn_output, _ = self.attn(x_norm1, x_norm1, x_norm1)
x = x + gate_msa.unsqueeze(1) * attn_output
# 3. AdaLN을 적용한 MLP 블록
x_norm2 = (1 + scale_mlp.unsqueeze(1)) * self.norm2(x) + shift_mlp.unsqueeze(1)
mlp_output = self.mlp(x_norm2)
x = x + gate_mlp.unsqueeze(1) * mlp_output
return x
# 사용 예시
# batch_size=4, seq_len=256 (패치 수), dim=1152
x = torch.randn(4, 256, 1152)
c = torch.randn(4, 1152) # 컨디셔닝 벡터 (시간 + 텍스트)
model = DiTBlock(hidden_size=1152, num_heads=16)
output = model(x, c)
print(f"Output Shape: {output.shape}")
아키텍처 시각화
flowchart TD
subgraph "전처리 (Preprocessing)"
A["입력 이미지/비디오"] --> B["VAE 인코더"]
B --> C["Latent(잠재) 표현"]
C --> D["Patchify (패치 단위로 자르기)"]
end
subgraph "Diffusion Transformer (DiT)"
D --> E["선형 임베딩 + 위치 인코딩"]
E --> F["DiT 블록 1 (Attn + MLP)"]
F --> G["DiT 블록 2 (Attn + MLP)"]
G --> H["... DiT 블록 N ..."]
H --> I["Depatchify (원래 형태로 재조립)"]
end
subgraph "컨디셔닝 (Conditioning)"
J["텍스트 프롬프트"] --> K["T5/CLIP/Mistral 인코더"]
L["타임스텝 (노이즈 레벨)"] --> M["MLP"]
K --> N["컨디셔닝 벡터 (c)"]
M --> N
N --"AdaLN을 통해 주입"--> F
N --"AdaLN을 통해 주입"--> G
end
subgraph "출력 (Output)"
I --> O["예측된 노이즈 / Latent"]
O --> P["VAE 디코더"]
P --> Q["최종 생성 미디어"]
end
style A fill:#333,stroke:#fff,color:#fff
style Q fill:#333,stroke:#fff,color:#fff
style F fill:#000,stroke:#0f0,color:#fff
style G fill:#000,stroke:#0f0,color:#fff
따라해 보기: 로컬에서 FLUX.2 실행하기
Sora 2는 비공개 소스이므로, 고성능 DiT 모델을 경험해 볼 수 있는 가장 좋은 진입점은 FLUX.2입니다. 상당한 VRAM을 요구하는(32B 파라미터의 괴물급 모델입니다) 모델이므로, 최소 24GB VRAM을 갖추거나 8-bit 양자화를 사용해야 합니다.
- 환경 설정:
FLUX.2 아키텍처를 지원하는 최신diffusers라이브러리가 필요합니다.pip install -U diffusers transformers accelerate sentencepiece - Hugging Face 인증:
black-forest-labs/FLUX.2-dev에 접근하려면 Hugging Face에서 라이선스 동의가 필요합니다.huggingface-cli login # HF 토큰 입력 - 파이프라인 실행 (FP8 양자화):
일반 소비자용 RTX 4090 등에 맞추기 위해 8-bit로 로드합니다.import torch from diffusers import FluxPipeline # FP8 최적화로 모델 로드 pipe = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.2-dev", torch_dtype=torch.bfloat16 ) # VRAM 절약을 위해 오프로딩 활성화 pipe.enable_model_cpu_offload() prompt = "A cinematic shot of a futuristic datacenter, glowing blue lights, 8k resolution, photorealistic" image = pipe( prompt, height=1024, width=1024, guidance_scale=3.5, num_inference_steps=50, max_sequence_length=512 ).images[0] image.save("flux2_output.png") - DiT를 위한 프롬프팅 전략:
- 서술적으로 작성하세요: “단어 샐러드” 방식(예: masterpiece, best quality 등 태그 나열)이 필요했던 초기 SD 모델과 달리, FLUX.2 같은 DiT 모델은 강력한 LLM 인코더(Mistral/T5)를 사용합니다. 자연어를 잘 이해하므로 태그가 아닌 문장으로 작성하세요.
- 물리와 조명을 언급하세요: DiT는 빛의 이동(light transport)을 렌더링하는 데 탁월합니다. 조명에 대해 구체적으로 묘사하세요 (예: “창문으로 들어오는 볼류메트릭 라이팅”).
Diffusion Transformer로의 전환은 단순한 점진적 업데이트가 아닙니다. 이는 향후 5년 동안 생성형 AI의 표준이 될 기술입니다. Sora 2 앱을 사용하든 로컬에서 FLUX.2를 돌리든, 여러분은 이제 픽셀과 컨볼루션이 아니라 패치와 어텐션으로 “생각하는” 모델과 상호작용하고 있는 것입니다. 이 아키텍처를 마스터하는 것이 곧 디지털 창작의 미래를 마스터하는 것입니다.