Paper Notes 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).

DAC (Descript Audio Codec) 정리

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

Descript Audio Codec: Neural Audio Codec (오디오를 작게 압축하고, 원본에 가까운 음질로 복원). 오디오 생성 모델들이 Latent를 만들 때 자주 쓰는 코덱이라 구조를 짧게 정리해 둔다. 논문은 Kumar et al., High-Fidelity Audio Compression with Improved RVQGAN (NeurIPS 2023). 숫자로 요약하면 44.1 kHz 오디오를 8 kbps, 약 90배 압축하면서 음성·음악·환경음을 모델 하나로 다룬다.

구조

Encoder: Waveform을 CNN에 태우기 → Continuous Latent Representation 생성
RVQ (Residual Vector Quantization)
- 연속적인 값을 Discrete한 Token으로 바꾸기 → Quantization
- 첫 번째 층: 전체적 뼈대만 잡아서 토큰화 → Residual이 크다
- 두 번째 ~ n번째 층: 이전 층의 오차(Residual)만 받아와서 Token화 → 반복…
- 오차가 0에 수렴 → 적은 Token으로 디테일한 소리 정보 저장.
Decoder: RVQ의 Discrete Token을 입력받아, Waveform으로 만들기
→ 훈련을 GAN으로 함. (그래서 이름이 RVQGAN)

RVQ 한 장 요약. 각 단계는 “남은 오차"만 양자화하므로, 뒤 코드북일수록 미세한 디테일을 담당한다.

구조적 특징

Fully Convolutional Network(FCN, 1차원 CNN)을 Encoder/Decoder에 사용
Snake Activation Function을 전 Layer에 사용 (주기가 있는 함수에 good)
Dilated Convolution으로 듬성듬성 넘어가서 Receptive Field 넓히기

Snake: $x+\frac{1}{\alpha}\sin^2(\alpha x)$ . 항등함수에 주기적인 굴곡이 들어가 있어서, ReLU로는 잘 못 배우는 주기 신호(= 소리)를 표현하기 좋다. 논문에서는 ReLU를 Snake로 바꾸는 것만으로 SI-SDR이 눈에 띄게 올랐다.

다음 이야기

→ Discrete Token을 생성하는 RVQ를 버리고, Continuous한 VAE를 쓰자 → DAC-VAE (Semantic-VAE)

참고 자료

R. Kumar, P. Seetharaman, A. Luebs, I. Kumar, K. Kumar, High-Fidelity Audio Compression with Improved RVQGAN, NeurIPS 2023. 코드/가중치: descriptinc/descript-audio-codec.
N. Zeghidour et al., SoundStream: An End-to-End Neural Audio Codec, 2021 — RVQ를 오디오 코덱에 처음 쓴 논문. A. Défossez et al., High Fidelity Neural Audio Compression (EnCodec), 2022 — DAC가 주로 비교하는 상대.
L. Ziyin, T. Hartwig, M. Ueda, Neural networks fail to learn periodic functions and how to fix it, NeurIPS 2020 — Snake activation의 출처. BigVGAN도 같은 활성함수를 쓴다.

방식	어떻게	비용	결과
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의 거의 절반)