Stable Diffusion의 아키텍처는 왜 이렇게 생겼나

Stable Diffusion을 처음 보면 낯선 선택들이 눈에 띈다. 왜 픽셀이 아니라 잠재 공간에서 diffusion을 하는가? 왜 의료 영상 분할 모델이었던 UNet을 사용하는가? 왜 최신 모델은 UNet을 버리고 Transformer로 갔는가? 이 질문들은 서로 연결되어 있다 — “어디서 계산할 것인가, 무엇이 bottleneck인가” 라는 단 하나의 문제를 각 시대의 설계자들이 다르게 풀어낸 역사다.

VAE: 계산 공간을 바꾼다

DDPM은 픽셀 공간 $(512 \times 512 \times 3 \approx 786\text{K dim})$에서 diffusion을 수행한다. 수백 스텝 동안 고해상도 UNet을 반복 통과시키면 GPU 메모리와 학습 시간이 기하급수적으로 늘어난다.

Rombach et al. (2022)의 핵심 아이디어는 단순하다 — diffusion을 픽셀이 아닌 잠재 공간에서 하자. Pretrained VAE encoder $\mathcal{E}$가 이미지를

$z = \mathcal{E}(x), \quad z \in \mathbb{R}^{64 \times 64 \times 4}$

로 압축하면 $(786\text{K} \to 16\text{K})$ 약 48배 압축이 이루어진다. Diffusion은 $z$ 위에서만 작동하고, 최종 이미지는 $x = \mathcal{D}(z_0)$로 복원한다.

VAE 학습 목표는 픽셀 L1 복원 손실 + LPIPS perceptual 손실 + adversarial 손실의 조합이다.

$\mathcal{L}_{\text{VAE}} = \mathbb{E}[\|x - \mathcal{D}(z)\|_1 + \lambda_{\text{lpips}} \text{LPIPS}(x, \hat{x}) + \lambda_{\text{adv}} \mathcal{L}_{\text{adv}}] + \beta\,\text{KL}(q(z|x) \| p(z))$

MSE만으로는 인지적으로 같은 이미지를 보장하지 못한다. LPIPS는 deep feature 기반으로 perceptual 유사도를 측정하므로, VAE가 “사람 눈에 같아 보이는” 잠재 표현을 학습한다. 결과적으로 잠재 공간은 픽셀 공간보다 intrinsic dimension이 훨씬 낮고 smooth한 manifold를 형성한다 — diffusion에 유리한 환경이다.

UNet: 다중 스케일이 denoising에 맞는 이유

Diffusion의 denoising network로 왜 UNet인가? Ronneberger et al. (2015)가 의료 영상 분할용으로 설계한 이 구조는 세 가지 속성을 갖는다.

Multi-scale feature: Encoder가 해상도를 낮추면서 의미론적 정보를 쌓고, Decoder가 복원하면서 세부 정보를 얹는다. Denoising도 정확히 이 계층적 처리가 필요하다 — 저해상도에서 구조를 잡고, 고해상도에서 texture를 복원한다.

Skip connection의 gradient 보존: Skip이 없으면 깊은 네트워크에서 gradient가 곱셈적으로 감쇠한다. Residual path가 추가되면

$\frac{\partial x_\ell'}{\partial x_\ell} = I + \frac{\partial f}{\partial x_\ell}$

identity $I$가 항상 gradient 1을 보장한다.

Timestep 주입 — FiLM: Diffusion step $t$는 sinusoidal encoding → MLP를 거쳐 affine 파라미터 $(\gamma_t, \beta_t)$가 된다. 각 ResBlock의 GroupNorm 출력에

$y = \gamma_t \cdot \text{GroupNorm}(x) + \beta_t$

가 적용된다. Timestep 정보가 모든 계층에 broadcast되므로 네트워크는 “지금이 몇 번째 denoising 스텝인지”를 항상 알고 있다.

DiT: inductive bias를 버리면 scaling law가 나타난다

Peebles & Xie (2023)는 UNet의 한계를 지적한다. Skip connection과 hierarchical structure는 특정 이미지 크기를 가정하는 inductive bias다. Scale up하면 이 bias가 오히려 병목이 된다.

DiT는 ViT처럼 이미지를 patch로 tokenize하고, N개의 identical Transformer block을 쌓는다. Timestep과 conditioning은 AdaLN-Zero로 주입된다.

$y = (1 - \alpha) x + \alpha [\gamma \cdot \text{LN}(x) + \beta]$

$\alpha = 0$으로 초기화하면 학습 초반에 각 block이 identity에 가깝게 동작한다. 네트워크가 안정적으로 warmup된 후 $\alpha$가 자라면서 점진적으로 modulation을 학습한다. 별도의 learning rate schedule 없이 파라미터 레벨에서 자동 warmup이 이루어지는 셈이다.

핵심 결과는 FID가 GFLOPs에 단조 감소한다는 것이다 — UNet에서는 볼 수 없었던 clean한 scaling law. 동일한 compute 예산에서 모델 크기를 키우면 품질이 예측 가능하게 개선된다. Sora(비디오)와 SD3(멀티모달)가 DiT를 기본 backbone으로 채택한 이유다.

MM-DiT: 단방향 cross-attention의 정보 비대칭을 해결한다

SD1-2는 text conditioning을 cross-attention으로 구현했다.

$\text{CrossAttn}(\mathbf{x}_{\text{img}}, \mathbf{x}_{\text{text}}) = \text{softmax}\left(\frac{Q_{\text{img}} K_{\text{text}}^T}{\sqrt{d_k}}\right) V_{\text{text}}$

이미지 query만 text와 interact한다. Text는 이미지를 직접 보지 못하고, 정보 흐름이 단방향이다.

MM-DiT (Esser et al., 2024)는 image patch tokens와 text tokens를 하나의 sequence로 concatenate하고 joint self-attention을 적용한다. 각 modality는 자신만의 Q, K, V projection을 갖되, attention 자체는 전체 sequence에 걸쳐 계산된다.

$A = \begin{bmatrix} A_{ii} & A_{it} \\ A_{ti} & A_{tt} \end{bmatrix}$

$A_{it}$는 이미지가 text를 보는 흐름, $A_{ti}$는 text가 이미지를 보는 흐름이다. 양방향 정보 흐름이 동시에 학습된다.

Cascaded Diffusion: 해상도 문제를 분업으로 푼다

고해상도를 직접 생성하는 대신 Imagen (Saharia et al., 2022)은 단계를 나눈다.

noise → [Base 64×64] → [SR 4× → 256×256] → [SR 4× → 1024×1024]
         (semantic)       (detail 1)            (detail 2)

각 stage는 독립적인 diffusion model이다. Stage $i$는 이전 stage의 출력을 upsampling한 뒤 conditioning으로 받아, 자신이 담당하는 해상도의 고주파 detail만 추가한다. Base 모델은 “무엇을 그릴지”를, SR 모델은 “얼마나 선명하게”를 담당한다.

오류 전파가 걱정될 수 있다. Stage 1의 실수가 이후 stage로 증폭될 것처럼 보인다. 그러나 SR 모델은 구조적으로 저주파 perturbation에 둔감하다 — 저주파는 upsampling으로 smooth되고, SR 모델은 고주파 detail만 학습하도록 최적화되기 때문이다.

정리

• VAE 잠재 공간은 픽셀 공간 대비 48배 압축으로 diffusion 계산을 현실적으로 만든다. LPIPS 손실이 perceptual하게 faithful한 잠재 표현을 보장한다.
• UNet + FiLM은 multi-scale feature 보존과 timestep modulation을 결합한다. Skip connection의 identity path가 gradient 소실을 막는다.
• DiT의 AdaLN-Zero는 inductive bias를 제거하고 clean한 scaling law를 얻는다. Identical block 구조가 compute 효율을 parameter 수와 FLOPs에만 의존하게 만든다.
• MM-DiT의 joint attention은 text↔image 양방향 정보 흐름을 허용한다. Cross-attention의 단방향 비대칭을 해소하는 대가로 $O(N^2)$ complexity가 증가한다.
• Cascaded Diffusion은 semantic vs detail을 분업으로 분리한다. SR 모델의 저주파 robustness가 오류 전파를 제한한다.

공통 주제는 하나다 — 어디서 계산하고, 어떻게 조건화하며, 무엇을 분리할 것인가. 아키텍처의 진화는 이 세 질문에 대한 답이 바뀌어 온 역사다.