Deep Learning on dirac042

MM-DiT — 두 모달리티를 동등하게 다루는 Diffusion Transformer

Mon, 06 Jul 2026 10:30:00 +0900

이전 노트에서 남긴 찝찝함 — 이미지/오디오가 주(主), 텍스트가 부(副)인 구조 — 을 정면으로 다루는 MM-DiT 정리. Stable Diffusion 3 논문(Esser et al., Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, 2024)에서 제안됐다.

기존의 DiT 구조는 2가지의 모달리티에 계층이 존재했다.

이미지, 오디오를 메인으로 두고,
텍스트를 보조 힌트로 사용.

이렇게 되면 2개의 모달리티의 깊이가 달라서, Deeper Level의 Alignment가 잘 안 되는 문제가 생긴다.

그러나 MM-DiT는 2개의 모달리티를 완벽하게 동등한 비중으로 다루게 된다.

Architecture

Stream을 2개로 쪼개자!
- Track A: Noise Token이 지나가는 길 (이미지/오디오)
- Track B: Text Token이 지나가는 길 (텍스트)
  - 이때 텍스트는 힌트로 쓰이고 버려지지 않고, 이미지랑 똑같이 블록을 통과할 때마다 스스로 update된다.
Joint Attention: 서로 다른 도로를 달리는 2개를 합치자!
1. 이미지/텍스트는 성질이 달라서, 독립적인 Linear 신경망을 거쳐 QKV 제작
2. 이미지 QKV와 텍스트 QKV를 Sequential하게 이어붙이기 (매우 큰 공간)
  → 이 공간 안에서 토큰들이 서로 Attention (픽셀, 텍스트와의 관계 학습)
  “사람들 4000명을 앉혀놓고, 전부 떠들게 만들기”
3. 정보를 가지게 된 토큰들을 다시 2개의 갈래로 나누기
4. 분리된 토큰들이 독립적인 MLP를 거쳐 최종적인 결과를 내리고, 다음 블럭으로 넘어가기

한 블록의 흐름. 두 트랙은 가중치(adaLN, QKV, MLP)를 따로 갖고, attention 한 번만 같은 방에서 한다.

SD3에서는 실제로 어떻게 조건을 만드나

텍스트 인코더를 세 개 쓴다: CLIP L/14, OpenCLIP bigG/14, T5-XXL. CLIP 두 개의 pooled 출력을 이어 붙여 하나의 벡터로 만들고, 여기에 timestep 임베딩을 더해 adaLN의 조건 벡터(위 그림의 “t + pooled text”)로 쓴다. 반면 토큰 단위 임베딩(CLIP의 penultimate hidden state + T5 토큰)은 Track B의 시퀀스가 되어 joint attention에 들어간다. 즉 “문장의 요약"은 modulation으로, “문장의 단어 하나하나"는 attention으로 들어가는 셈이다.

그리고 이게 정말 나은지도 재봤다. 같은 조건에서 vanilla DiT < UViT < cross-attention DiT(CrossDiT) < MM-DiT 순서였고, 논문은 “MM-DiT가 cross-attention과 vanilla 변형을 크게 앞선다"고 쓴다. 학습 목표는 DDPM 대신 rectified flow이고, 중간 timestep을 더 자주 뽑는 logit-normal 샘플링을 쓴다. 깊이 38(약 8B 파라미터)까지 키워도 validation loss가 포화 없이 계속 내려갔다.

Problems

뭔가 데이터가 엄청나게 커지게 되면, Joint Attention이 엄청나게 느려질 것 같은 예감.

Transformer Attention 연산은 $\mathcal{O}(N^2)$, 즉 데이터 길이의 제곱에 비례해 계산량이 늘어난다.
- 기존 방식은 Noise Token이 Text를 참고만 하니까, $\mathcal{O}(N^2)$ .
- MM-DiT는 전부 다 몰아넣고 Attention을 실행하니까, $\mathcal{O}((N_{text}+N_{image})^2)$
- 계산량이 매우매우 커진다.
Attention 계산에서 OOM 에러
Modality Imbalance하면 효율이 떨어진다.
ex. 텍스트 100개, 이미지 4만 개 → 이미지-이미지 간의 상호작용만 여러 개.

아직까지도 안 풀린 숙제로 남아 있다.

Flash Attention 같은 최적화 기술을 영끌해서 적용하거나,
3개의 블럭에서만 Joint Attention하고, 나머지는 분리해서 연산하는 등…

실제로 나온 절충안들

FLUX.1(Black Forest Labs, 2024)은 앞쪽에 두 트랙이 분리된 MM-DiT식 double-stream 블록을 두고, 뒤쪽은 두 시퀀스를 아예 하나로 합쳐 처리하는 single-stream 블록으로 바꿔 파라미터를 아낀다 — 위 2번 아이디어와 같은 방향이다.
오디오 쪽에서는 MMAudio(Cheng et al., 2024)가 비디오·텍스트·오디오 세 스트림을 joint attention으로 묶고, 비디오-오디오 프레임을 시간축에 맞춰 정렬한 위치 임베딩(aligned RoPE)을 얹어 동기화 문제를 푼다.
계산량 자체는 FlashAttention 계열이 메모리 병목을 크게 완화하지만 $\mathcal{O}(N^2)$ 자체가 사라지진 않는다. 그래서 latent를 더 압축하거나(패치 크기, VAE 다운샘플 비율), 윈도우/희소 attention을 섞는 연구가 이어진다.

참고 자료

P. Esser et al., Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, 2024 — MM-DiT 원 논문(Stable Diffusion 3).
W. Peebles, S. Xie, Scalable Diffusion Models with Transformers, ICCV 2023 — 출발점인 DiT.
H. K. Cheng et al., MMAudio: Taming Multimodal Joint Training for High-Quality Video-to-Audio Synthesis, 2024 — 오디오에서의 multimodal joint attention.
black-forest-labs/flux — double/single-stream 블록 구현.

DiT (Diffusion Transformer) 정리

Mon, 06 Jul 2026 10:20:00 +0900

Diffusion의 백본을 U-Net에서 Transformer로 바꾼 DiT(Peebles & Xie, Scalable Diffusion Models with Transformers, ICCV 2023)를, 오디오 생성 관점에서 읽고 정리한 노트.

이전에 사용하던 U-Net 기반의 Diffusion은 이미지 생성 Task에서 많이 활용되어 왔다.

하지만, U-Net의 구조적인 한계 때문에, 시간 간격이 긴 데이터에 약하다.
즉, 앞의 데이터와 뒤의 데이터를 서로 연결할 수 있는 힘이 없다.

그래서 Transformer의 Attention 구조가 가진 Long-Term Dependency를 이용해 Transformer 기반의 Diffusion Method를 고안하게 되는데, 이것이 DiT라고 볼 수 있겠다.

Architecture

결국 Diffusion이 메인이라, Gaussian Noise에서 GT로 가는 방향을 예측하는 건 동일하다.

즉, DiT의 Block은 현재 Timestep에서 입력된 Noisy Latent를 어떤 방향으로 Denoise할 것인지를 구하는 게 목표다.

Patchify: latent를 토큰으로

Transformer는 토큰의 시퀀스를 먹으니, 먼저 latent를 잘라야 한다. $I\times I$ latent를 $p\times p$ 패치로 자르면 토큰 수는 $$T=(I/p)^2$$ . 논문은 $p\in\{2,4,8\}$ 을 실험했는데, $$p$$ 를 반으로 줄이면 토큰이 4배가 되고 Gflops도 최소 4배가 된다(파라미터 수는 거의 그대로). 그리고 그만큼 성능이 좋아진다 — 이것도 “계산량이 곧 성능"의 한 사례. 위치 정보는 ViT처럼 sine-cosine positional embedding으로 넣는다. 오디오라면 $I\times I$ 대신 (시간 × 주파수/채널) latent를 같은 식으로 자르면 된다.

블록 안에서 일어나는 일

AdaLN (Adaptive Layer Normalization)
- 먼저 시간 $$t$$ 를 주입받아, 각각의 토큰들에게 ‘얼마나 깎아야 하는지’를 전달한다.
Self-Attention
- 1에서 나온 정보들을 바탕으로, 각 토큰들은 자기 자신(Waveform) 안에서 패턴을 분석한다.
Cross-Attention
- 이때 외부의 Condition(Video/Text)이 들어와 Denoising을 위한 힌트를 제공하고, 소리 생성의 방향성이 생긴다.
MLP
- 최종적으로 “이 소리는 이 방향으로 이동해야 한다!“를 구하게 된다.

원 논문에서 가장 좋았던 adaLN-Zero 블록. 조건(t, c)은 attention이 아니라 LayerNorm의 scale/shift와 residual 앞의 gate로 들어간다.

조건을 넣는 네 가지 방법 — 그리고 승자

위 1~4번은 “cross-attention을 쓰는 DiT"의 그림이다(오디오/텍스트 조건 모델들이 실제로 이렇게 쓴다). 그런데 원 논문은 조건(timestep $$t$$ , class label $$c$$ )을 주입하는 방법을 네 가지 비교했다.

방식	어떻게	비용	결과
In-context	$$t$$ , $$c$$ 임베딩을 토큰 두 개로 시퀀스 앞에 붙임	거의 0	가장 나쁨
Cross-attention	조건 시퀀스에 대해 별도 cross-attn 층	약 +15% Gflops	중간
adaLN	조건 벡터에서 LayerNorm의 scale γ, shift β를 회귀	가장 적음	좋음
adaLN-Zero	adaLN + residual 직전의 dimension-wise gate α를 0으로 초기화	가장 적음	최고 (400K step에서 FID가 in-context의 거의 절반)

adaLN-Zero의 아이디어는 학습 초기에 각 블록이 항등함수(identity)에서 출발하게 만드는 것. 그래서 위 그림의 α가 “0으로 초기화"다. 정리하면 — 조건이 “짧은 벡터”(timestep, class, pooled text)라면 adaLN 계열이, 조건이 “긴 시퀀스”(문장 토큰, 비디오 프레임)라면 cross-attention이 자연스럽다. Stable Audio Open 같은 오디오 DiT가 텍스트 토큰은 cross-attention으로, timestep/타이밍 조건은 adaLN으로 넣는 이유가 여기 있다.

Commentary

찝찝한 부분이 있다.

현재 DiT에서는 Self-Attention에서 Waveform의 패턴을 분석한 다음,
부가적으로 Condition이 들어와 Denoising을 시행한다.

그 말인즉슨, Waveform이 주(主), Video/Text가 부(副)가 되는 형태라고 볼 수 있겠다.

이런 식으로 이미지/오디오가 메인, 텍스트가 보조가 되면 정보 불균형이 생기지 않을까?

이걸 Hierarchical하지 않고, Equal하게 다룰 수는 없을까?

→ 다음 노트: MM-DiT

참고 자료

W. Peebles, S. Xie, Scalable Diffusion Models with Transformers, ICCV 2023. 코드: facebookresearch/DiT.
P. Esser et al., Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, 2024 (Stable Diffusion 3) — DiT를 두 모달리티 동등 구조로 확장한 MM-DiT.
Z. Evans et al., Stable Audio Open, 2024 — 오디오 latent 위에서 돌아가는 DiT의 실제 예 (T5 텍스트는 cross-attention, 타이밍 조건은 adaLN).

방식	어떻게	비용	결과
In-context	$\(t\)$ , $\(c\)$ 임베딩을 토큰 두 개로 시퀀스 앞에 붙임	거의 0	가장 나쁨
Cross-attention	조건 시퀀스에 대해 별도 cross-attn 층	약 +15% Gflops	중간
adaLN	조건 벡터에서 LayerNorm의 scale γ, shift β를 회귀	가장 적음	좋음
adaLN-Zero	adaLN + residual 직전의 dimension-wise gate α를 0으로 초기화	가장 적음	최고 (400K step에서 FID가 in-context의 거의 절반)