[Study] MMDiT를 이용한 SD3

Stable Diffusion 3 — UNet 을 버리고 Transformer 로

2024년 2월 Stability AI 는 Stable Diffusion 3 의 기술 보고서를 공개했고, 같은 해 6월 SD3 Medium(2B 파라미터) 을 정식 출시했다. SD3 는 SDXL 의 후속이지만 단순한 업그레이드가 아니다. UNet 기반 디퓨전 구조를 통째로 갈아엎고 Transformer 기반 아키텍처(MMDiT) 로 옮겼고, 학습 방식도 표준 디퓨전 대신 Rectified Flow 를 쓴다. 텍스트 인코더는 SDXL 의 두 개에 T5-XXL 을 더해 세 개로 늘렸다.

바꾼 부분이 많고 각각의 변화가 모델의 작동 방식을 근본적으로 건드린다. 각 변경이 어떻게 작동하는지를 원리 중심으로 본다.

---

1. SDXL 의 구조적 한계

SD3 가 왜 구조 자체를 바꿨는지 보려면 SDXL 의 구조적 한계를 짚어야 한다. 데이터를 더 넣거나 모델을 키우는 것만으로는 풀리지 않는 문제들이다.

• UNet 의 지역적 처리. UNet 은 컨볼루션 기반이라 기본적으로 인접한 픽셀끼리 정보를 주고받는다. 멀리 떨어진 영역끼리의 관계는 여러 단계를 거쳐야 통합되고, 그 과정에서 정보가 흐려진다.
• Cross-Attention 의 일방향성. SDXL 에서는 텍스트가 이미지 토큰의 attention 에 일방적으로 조건만 걸어준다. 이미지 정보가 텍스트 표현을 다시 다듬는 흐름은 없다.
• 텍스트 렌더링이 약하다. 위의 두 가지가 합쳐진 결과다. "Hello" 라는 글자를 이미지에 그리려면 문자 모양과 공간 배치를 정밀하게 통합해야 하는데, UNet + Cross-Attention 구조에서는 이게 잘 안 된다.
• 표준 디퓨전의 곡선 경로. 학습된 노이즈 → 이미지 경로가 직선이 아니라 휘어져 있다. 그래서 적은 step 으로 좋은 결과를 내기 어렵다.

SD3 는 이 네 가지에 각각 다른 변경으로 대응한다. 아키텍처는 MMDiT 로 바꾸고, attention 은 joint attention 으로 바꾸고, 텍스트 인코더는 T5-XXL 을 추가하고, 학습은 Rectified Flow 로 옮긴다.

---

2. MMDiT — 이미지를 토큰으로 다룬다

SD3 의 아키텍처는 MMDiT (Multimodal Diffusion Transformer) 다. 이름 그대로 멀티모달(텍스트+이미지) 디퓨전 Transformer 다. Peebles 와 Xie 가 발표한 DiT(Diffusion Transformer) 아이디어를 멀티모달로 확장한 구조다.

UNet 에서 Transformer 로 옮긴다는 게 무슨 의미인가

UNet 은 이미지를 2D 격자로 다룬다. 컨볼루션 필터가 작은 영역을 훑으면서 특성을 뽑고, 다운샘플링과 업샘플링을 반복하면서 다양한 스케일의 정보를 만든다. 공간 구조가 모델 안에 그대로 박혀 있다.

Transformer 는 다르다. 입력을 토큰의 시퀀스로 다룬다. 토큰끼리의 관계는 self-attention 으로 계산되는데, attention 은 모든 토큰 쌍을 동시에 본다. 거리에 상관없이 모든 위치가 매 블록마다 직접 연결된다. 공간 구조는 별도의 position embedding 으로 주입한다.

이미지를 토큰으로 어떻게 만드나

1. VAE 로 latent 로 인코딩. 1024×1024 이미지를 VAE 로 압축해 128×128 정도의 latent 로 만든다. 이 부분은 기존 LDM / SDXL 과 동일한 아이디어다.
2. Patch 로 자른다. latent 를 2×2 patch 로 잘라서, 각 patch 하나가 토큰 하나가 된다. 128×128 latent 라면 64×64 = 4,096 개의 이미지 토큰이 나온다.
3. Linear projection + Position Embedding. 각 patch 를 선형 변환으로 임베딩 벡터로 만든다. 여기에 2D position embedding 을 더해 공간 위치 정보를 토큰에 박는다.
4. 텍스트 토큰과 합쳐 Transformer 블록을 통과. 이미지 토큰과 텍스트 토큰이 같은 Transformer 안에서 처리된다.
5. 다시 latent 로 복원해 VAE 디코딩. Transformer 가 출력한 이미지 토큰들을 모아 다시 latent 모양으로 만든 뒤, VAE 로 이미지를 복원한다.

이게 왜 중요할까. UNet 에서는 멀리 떨어진 두 영역(예: 이미지 왼쪽 위와 오른쪽 아래) 이 정보를 주고받으려면 다운샘플링으로 해상도가 충분히 작아진 단계까지 가야 한다. MMDiT 에서는 모든 토큰이 매 attention 블록마다 직접 연결돼 있다. 첫 번째 블록부터 멀리 떨어진 영역끼리 정보를 주고받을 수 있다.

---

3. Joint Attention 

MMDiT 의 핵심은 단순히 Transformer 를 쓴다는 게 아니다. 텍스트와 이미지가 서로를 어떻게 보는지가 SDXL 과 다르다.

SDXL 의 Cross-Attention 방식

SDXL 의 UNet 안에서 텍스트는 다음과 같이 들어간다. 이미지 latent 가 self-attention 으로 자기 안에서 정보를 정리하고, 그 다음 cross-attention 블록에서 텍스트가 들어온다. cross-attention 에서는 이미지 토큰이 Query 를 만들고, 텍스트 토큰이 Key / Value 를 제공한다.

이 구조의 한계는 분명하다. 텍스트 → 이미지 방향으로만 정보가 흐른다. 텍스트는 이미지를 조건화하지만, 이미지 정보가 텍스트 표현을 다시 갱신하지는 않는다.

MMDiT 의 Joint Attention 방식

MMDiT 는 이걸 바꾼다. 이미지 토큰과 텍스트 토큰을 이어 붙여(concatenate) 같은 attention 블록에 통과시킨다.

텍스트 토큰  +  이미지 토큰  →  Concatenate  →  Joint Self-Attention

단순히 합치는 것만은 아니다. 두 modality 는 각각 독립된 QKV projection 가중치를 갖는다. 이게 이름의 "Multimodal" 이 가리키는 부분이다. 텍스트 토큰은 텍스트용 가중치로 Q, K, V 를 만들고, 이미지 토큰은 이미지용 가중치로 만든다. 그렇게 만들어진 Q, K, V 를 전체 토큰에 대해 한꺼번에 attention 연산을 수행한다.

핵심 차이는 이렇다. Cross-attention 에서는 텍스트가 이미지에 영향을 주는 일방향이었다. Joint attention 에서는 같은 attention 행렬 안에서 텍스트 토큰이 이미지 토큰을 보고, 이미지 토큰이 텍스트 토큰을 본다. 양쪽이 동시에 갱신된다. 그래서 텍스트 표현 자체가 이미지 내용을 반영해서 조정될 수 있고, 이미지도 텍스트의 의도를 더 세밀하게 반영하게 된다.

이 차이가 텍스트 렌더링에 미치는 영향

SD3 가 이미지 안에 글자를 잘 그리는 것이 자주 언급되는데, 그 원인이 여기에 있다. "Hello" 라는 글자를 그리려면 모델이 각 문자의 형태와 위치, 순서를 동시에 결정해야 한다. Cross-attention 구조에서는 텍스트가 이미지에 한 번 조건만 걸어주고 끝이니, 이미지가 생성되는 동안 텍스트가 다시 정렬되거나 정제되지 않는다.

Joint attention 은 다르다. 이미지 토큰이 글자를 그리려고 시도하는 동안, 그 정보가 텍스트 토큰 쪽으로 역류해서 텍스트 표현이 다시 정돈된다. 텍스트와 이미지가 매 블록마다 서로를 보며 협의하는 구조다.

---

4. Rectified Flow — 노이즈 경로를 직선으로

SD3 가 바꾼 또 하나의 큰 부분은 학습 방식이다. SDXL 까지는 표준 디퓨전(DDPM 계열) 을 썼지만, SD3 는 Rectified Flow 를 채택했다.

표준 디퓨전이 학습하는 경로

표준 디퓨전은 이미지에 노이즈를 단계적으로 더해가는 과정을 정의한다. x_t = √(1-β_t) x_{t-1} + √β_t · ε 같은 식이다. β_t 가 노이즈 스케줄인데, 보통 작게 시작해서 점점 커진다. 학습은 모델이 각 t 에서 "어떤 노이즈가 더해졌는지" 를 예측하게 한다.

문제는 이 식으로 그려지는 데이터-노이즈 사이의 경로가 직선이 아니라는 점이다. 노이즈 스케줄에 따라 휘어진 곡선을 따라 이동한다. 그래서 추론할 때 큰 step 으로 점프하면 곡선에서 벗어나 결과가 무너진다. 적은 step 으로 좋은 결과를 내려면 별도의 distillation 기법 같은 게 필요했다.

Rectified Flow 의 발상

Rectified Flow 는 데이터와 노이즈 사이를 직선으로 연결한다. 시간 t 는 0 부터 1 까지 균일하게 잡고, 학습 데이터 x₀ 와 노이즈 ε 를 다음과 같이 선형 보간한다.

x_t  =  (1 − t) · x₀  +  t · ε

t = 0 이면 순수 데이터, t = 1 이면 순수 노이즈. 그 사이는 두 점을 잇는 직선 위의 점이다. 모델은 노이즈를 예측하는 대신 "이 위치에서 어느 방향으로 가야 노이즈에 도달하는가" 라는 속도 벡터(velocity field) 를 학습한다. 정의상 이 속도는 v = ε − x₀, 즉 노이즈에서 데이터를 뺀 일정한 방향이다.

왜 직선 경로가 유리한가

• 적은 step 으로 추론 가능. 학습된 경로가 직선이면 큰 step 으로 점프해도 경로에서 멀리 벗어나지 않는다. 같은 품질을 더 적은 step 으로 얻을 수 있다.
• 학습이 안정적. 노이즈 스케줄을 따라 가중치가 비대칭으로 쏠리지 않고, 시간 t 에 대해 균일하게 학습된다.
• 고해상도에 더 잘 적응한다. SD3 논문에서는 해상도별로 시간 가중치를 다르게 주는 기법(logit-normal sampling) 을 함께 제안한다. 큰 해상도일수록 노이즈 제거에 더 많은 시간이 필요한 점을 반영해 t 분포를 조정한다.

주의할 점은 "직선" 은 이론적인 목표지 실제 학습된 vector field 가 완벽히 직선을 그리지는 않는다. Rectified Flow 는 직선에 가깝게 만들어 sampling 효율을 높이는 방법이지, 한 step 으로 끝낼 수 있다는 의미는 아니다. SD3 추론도 보통 20~30 step 정도를 쓴다.

---

5. 세 개의 Text Encoder — T5-XXL 의 합류

SDXL 은 CLIP 계열 인코더 두 개를 썼지만, SD3 는 여기에 T5-XXL 을 추가해 세 개를 쓴다.

인코더	특징
CLIP ViT-L (~150M)	짧고 명료한 프롬프트에 강한 OpenAI 원조 CLIP.
OpenCLIP ViT-bigG (~700M)	LAION 데이터로 학습된 더 큰 CLIP. 어휘 범위가 넓다.
T5-XXL (~4.7B)	구글의 대형 텍스트-투-텍스트 LM. 문장 구조 이해와 텍스트 렌더링의 핵심.

CLIP 과 T5 는 다르게 작동한다

CLIP 은 이미지-텍스트 매칭으로 학습됐다. 그래서 "이미지를 묘사하는 짧은 캡션" 에 강하다. 하지만 긴 문장의 문법 구조나 단어 사이의 미묘한 관계는 잘 못 잡는다.

T5 는 다르다. 텍스트-투-텍스트 변환 과제로 학습된 진짜 언어 모델이다. 긴 문장을 끝까지 읽고 토큰 사이의 관계를 풍부하게 표현한다. 특히 글자 자체를 정확히 표현하는 데 강하다. "a sign that says 'Open'" 이라는 프롬프트에서 Open 이라는 문자열 자체가 그림에 들어가야 한다고 모델이 인식하려면 T5 의 토큰 표현이 도움이 된다.

결합 방식

세 인코더는 단순히 출력을 다 합치는 게 아니라 다른 방식으로 들어간다.

• CLIP-L 과 CLIP-G — 두 모델의 pooled vector (문장 전체 요약 벡터) 를 concatenate 해서 timestep embedding 과 같은 위치에 더해진다. 이미지 토큰 처리에 글로벌 조건으로 작용한다.
• T5-XXL — 토큰별 임베딩(각 단어 / 서브워드의 표현) 을 그대로 텍스트 토큰으로 사용한다. 이게 MMDiT 의 joint attention 에 들어가는 텍스트 토큰이다.

한 가지 실용적인 포인트로, T5-XXL 은 4.7B 파라미터라 VRAM 을 많이 잡는다. 추론할 때 T5 를 제외하고 CLIP 두 개만 쓰는 옵션이 있는데, 이 경우 메모리는 크게 줄지만 텍스트 렌더링 품질과 긴 프롬프트 이해도가 떨어진다. SD3 Medium 의 공식 변형 중 하나도 T5 없는 버전이다.

---