Transfusion — Autoregressive + Diffusion

Chameleon과 Emu3는 모든 것을 이산 토큰(discrete token)에 걸었습니다. 동작은 하지만 양자화 병목(quantization bottleneck)이 눈에 보입니다. 이미지 품질이 연속 공간(continuous-space) 디퓨전 모델(diffusion model)보다 낮은 지점에서 정체(plateau)됩니다. Transfusion(Meta, Zhou et al., 2024년 8월)은 반대 방향에 베팅합니다. 이미지는 연속(continuous)으로 유지하고, VQ-VAE를 완전히 제거하며, 하나의 트랜스포머(transformer)를 두 개의 손실(loss)로 학습시킵니다. 텍스트 토큰(text token)에는 다음 토큰 예측(Next-Token Prediction; NTP)을, 이미지 패치(image patch)에는 흐름 정합(flow matching) 또는 디퓨전 손실(diffusion loss)을 적용합니다. 두 목적함수(objective)는 같은 가중치(weight)를 최적화합니다. Stable Diffusion 3의 기반 아키텍처(architecture)인 MMDiT는 이와 가까운 사촌입니다. 이번 강의에서는 Transfusion 논문의 핵심 주장(thesis)을 읽고, 토이(toy) 규모의 두-손실(two-loss) 학습기를 만들며, 하나의 트랜스포머가 두 작업을 모두 수행하게 해주는 어텐션 마스크(attention mask)를 추적해봅니다.

유형: Build 언어: Python (표준 라이브러리, MNIST 규모 토이 위의 two-loss trainer) 선수 지식: Phase 12 · 11 (Chameleon), Phase 8 (Generative AI) 예상 시간: 약 180분

학습 목표

• 하나의 백본(backbone) 위에서 두 손실을 함께 실행하는 트랜스포머를 연결합니다. 텍스트 토큰에는 NTP를, 이미지 패치에는 디퓨전 MSE를 적용합니다.
• 이미지 패치끼리는 양방향 어텐션(bidirectional attention)을 허용하고 텍스트 토큰에는 인과 어텐션(causal attention)을 적용하는 마스크가 왜 올바른 선택인지 설명합니다.
• Transfusion 방식(연속 이미지, 디퓨전 손실)과 Chameleon 방식(이산 이미지, NTP)을 연산량(compute), 품질(quality), 코드 복잡도(code complexity) 관점에서 비교합니다.
• MMDiT의 기여를 한 줄로 정리합니다. 각 블록에서 모달리티별(modality-specific) 가중치를 사용하고, 잔차 스트림(residual stream)에서는 공동 어텐션(joint attention)을 수행한다는 점입니다.

문제

이산 이미지 토큰과 연속 이미지 토큰을 둘러싼 논쟁은 LLM보다 오래되었습니다. 연속 표현(continuous representation)인 원본 픽셀(raw pixel)이나 VAE 잠재 변수(latent)는 디테일을 보존합니다. 이산 토큰(VQ index)은 트랜스포머의 본래 어휘집(native vocabulary)에 잘 맞지만, 양자화 단계에서 디테일을 잃습니다.

Chameleon과 Emu3는 이산 쪽을 선택했습니다. 손실 하나와 아키텍처 하나로 학습 구조는 단순해졌지만, 이미지 충실도(image fidelity)는 토크나이저(tokenizer) 품질에 의해 제한되었습니다.

디퓨전 모델은 연속 쪽을 선택했습니다. 이미지 품질은 뛰어나지만 LLM과는 별도의 모델로 존재하며, 잡음 스케줄(noise schedule) 엔지니어링이 복잡하고 텍스트 생성과 깔끔하게 통합되지 않습니다.

Transfusion은 이 둘을 모두 가질 수 있는지 묻습니다. 이미지는 연속으로 유지하면서도, 여전히 하나의 모델만 학습하고, 하나의 경사 단계(gradient step) 안에 두 손실을 함께 꿰매 넣는 방식입니다.

개념

두-손실 아키텍처(two-loss architecture)

하나의 디코더 전용(decoder-only) 트랜스포머가 다음 요소를 포함하는 시퀀스(sequence)를 처리합니다.

• 텍스트 토큰: BPE 어휘집에서 온 이산 토큰입니다.
• 이미지 패치: 연속 패치입니다. 16x16 픽셀 블록을 선형 임베딩(linear embedding)으로 은닉 차원(hidden dim)에 사영합니다. ViT 인코더의 입력 레이어와 같은 방식입니다.
• <image>와 </image> 태그: 연속 패치가 어느 구간에 들어 있는지 표시합니다.

순전파(forward pass)는 한 번만 실행됩니다. 손실은 토큰 종류에 따라 두 개의 헤드(head) 중 하나를 고릅니다.

• 텍스트 토큰: 어휘 로짓 헤드(vocab-logits head)에 대해 표준 교차 엔트로피(cross-entropy)를 사용합니다.
• 이미지 패치: 연속 패치에 대한 디퓨전 손실을 사용합니다. 각 패치에 더해진 잡음을 예측합니다.

경사(gradient)는 공유된 트랜스포머 본체(shared transformer body)를 통해 흐릅니다. 두 손실이 공유 가중치를 동시에 개선합니다.

어텐션 마스크: 텍스트는 인과, 이미지는 양방향

텍스트 토큰은 반드시 인과적(causal)이어야 합니다. 텍스트 토큰이 미래의 텍스트를 보게 하면 교사 강요(teacher forcing)가 깨집니다. 반면 이미지 패치는 하나의 스냅샷을 표현하므로, 같은 이미지 블록 안에서는 서로 양방향으로 어텐드(bidirectionally attend)해야 합니다.

마스크는 다음과 같습니다.

M[i, j] = 1 if:
  (i is text and j is text and j <= i)   # 텍스트끼리는 인과적
  OR (i is image and j is image and same_image_block(i, j))   # 이미지 내부에서는 양방향
  OR (i is text and j is image and j < i_image_end)   # 텍스트는 이전 이미지를 봄
  OR (i is image and j is text and j < i_image_start)   # 이미지는 앞선 텍스트를 봄

학습과 추론 모두에서 블록 삼각(block-triangular) 마스크 형태로 구현됩니다.

트랜스포머 안의 디퓨전 손실

디퓨전 손실 자체는 표준적입니다. 이미지 패치에 잡음을 더한 뒤, 모델에게 그 잡음을 예측하게 합니다. 동등하게, 깨끗한(clean) 패치를 예측하도록 학습시킬 수도 있습니다. Transfusion은 흐름 정합(flow matching)을 사용합니다. 잡음이 섞인(noisy) 상태에서 깨끗한 상태로 향하는 속도장(velocity field)을 예측합니다.

학습 중에는 다음을 수행합니다.

1. 각 이미지 패치 x0에 대해 무작위 시점(random timestep) t를 표본 추출(sample)합니다.
2. 잡음 ε를 표본 추출하고, 흐름 정합을 위해 선형 보간 xt = (1-t) * x0 + t * ε를 계산합니다.
3. 트랜스포머가 v_theta(xt, t)를 예측합니다. 손실은 MSE(v_theta(xt, t), ε - x0)입니다.
4. 같은 시퀀스에서 나온 텍스트 NTP 손실과 함께 역전파(backprop)합니다.

추론 시 생성은 다음과 같이 진행됩니다.

• 텍스트 토큰: 표준 자기회귀(autoregressive) 샘플링입니다.
• 이미지 패치: 앞선 텍스트 토큰을 조건으로 하는 디퓨전 샘플링 루프입니다. 일반적으로 10–30 단계(step)를 사용합니다.

MMDiT: Stable Diffusion 3의 변형

Stable Diffusion 3(Esser et al., 2024년 3월)은 Transfusion과 비슷한 시기에 MMDiT(Multimodal Diffusion Transformer)를 공개했습니다. 두 아키텍처는 형제(sibling)에 가깝습니다.

MMDiT의 핵심 차이는 다음과 같습니다.

• 블록별 모달리티별 가중치(modality-specific weight). 각 트랜스포머 블록은 텍스트 토큰과 이미지 패치에 대해 서로 다른 Q, K, V, MLP 가중치를 가집니다. 어텐션은 공동(joint), 즉 모달리티 교차(cross-modality)로 수행하지만, 그 외의 부분은 모달리티별로 분리되어 있습니다.
• 정류 흐름(rectified flow) 학습. DDPM보다 샘플링과 수학이 더 단순한 특정 흐름 정합 변형(flow-matching variant)입니다.
• 규모(scale). MMDiT는 SD3의 백본입니다. 2B와 8B 파라미터(parameter) 변형이 있습니다. Transfusion 논문은 7B까지 규모를 확장합니다.

두 모델 모두 같은 핵심 아이디어에 도달합니다. 하나의 트랜스포머가 텍스트에는 NTP를, 연속 이미지 표현에는 디퓨전을 수행한다는 것입니다.

Chameleon 방식보다 나은 이유

이미지 생성에서 연속 디퓨전(continuous-diffusion)과 이산 NTP(discrete-NTP) 사이의 품질 격차는 측정 가능합니다. Transfusion 논문은 다음을 보고합니다.

• 7B 파라미터 규모에서 같은 크기의 Chameleon 방식 모델보다 FID 점수가 3~5점 더 좋습니다.
• 토크나이저 학습이 필요 없습니다. 이미지 인코더가 더 단순해집니다. ViT의 입력 레이어와 동일하게, 은닉 차원으로의 선형 사영(linear projection)만 사용합니다.
• 자기회귀 이미지 토큰과 달리, 추론 단계에서 이미지 패치 디노이징(denoising)을 병렬화(parallelize)할 수 있습니다.

단점도 있습니다. Transfusion은 이중 손실(dual-loss) 모델이므로 학습 동역학(training dynamics)이 더 까다롭습니다. 손실 가중치를 따로 조정해야 합니다. NTP와 디퓨전 사이의 스케줄 불일치(schedule mismatch)가 생기면 한쪽 헤드가 다른 쪽 헤드를 압도(dominate)할 수 있습니다.

다운스트림에 놓인 것들

Janus-Pro(Lesson 12.15)는 Transfusion의 아이디어를 더 정교하게 다듬어, 이해(understanding)와 생성(generation)을 위한 비전 인코더(vision encoder)를 분리합니다. 한쪽에는 SigLIP을, 다른 쪽에는 VQ를 사용하면서 트랜스포머 본체는 공유합니다. Show-o(Lesson 12.14)는 디퓨전을 이산 디퓨전(discrete-diffusion), 즉 마스크 예측(masked prediction)으로 바꿉니다. Transfusion 이후 통합 생성(unified-generation) 계열은 빠르게 갈라져 나갑니다.

2026년 운영(production) VLM 중 이미지를 출력하는 Gemini 3 Pro, GPT-5, Claude Opus 4.7의 이미지 생성 경로(path)는 거의 확실히 이 계열의 어떤 후손(descendant)을 사용합니다. 세부 사항은 비공개(proprietary)입니다.

사용해보기

code/main.py는 작은 MNIST 풍의 문제 위에서 토이 Transfusion을 만듭니다.

• 텍스트 캡션(caption)은 숫자(digit) 0–9를 묘사하는 짧은 정수 시퀀스입니다.
• 이미지는 바이트(byte)로 이루어진 4x4 격자(grid)입니다.
• 공유 가중치(shared-weight) 선형 사영 한 쌍이 트랜스포머의 대역(stand-in) 역할을 합니다. 텍스트에는 NTP 손실을, 잡음이 섞인 패치에는 MSE 손실을 적용합니다.
• 학습 루프는 두 손실을 번갈아 계산하고, 어텐션 마스크를 명시적으로 구성합니다.
• 생성 단계는 하나의 순전파 안에서 텍스트 캡션과 4x4 이미지를 함께 만들어냅니다.

트랜스포머 자체는 토이 수준입니다. 실제로 배워야 할 산출물은 두-손실 배관(two-loss plumbing), 어텐션 마스크 구성, 추론 루프입니다.

산출물 만들기

이번 강의는 outputs/skill-two-loss-trainer-designer.md를 만듭니다. 새로운 멀티모달 학습 과제(text + image, text + audio, text + video)가 주어졌을 때, 두-손실 스케줄(손실 가중치, 마스크 형상, 공유 블록 대 모달리티별 블록 비율)을 설계하고 구현상의 위험(implementation risk)을 표시합니다.

연습문제

1. 쉬움: Transfusion 방식 모델이 70%의 텍스트 토큰과 30%의 이미지 패치로 학습됩니다. 이미지 디퓨전 손실의 크기가 텍스트 NTP 손실보다 약 10배 큽니다. 두 손실을 균형 있게(balance) 맞추려면 어떤 손실 가중치를 사용해야 합니까?

2. 중간: 시퀀스 [T, T, <image>, P, P, P, P, </image>, T]에 대한 블록 삼각 마스크(block-triangular mask)를 구현하세요. 각 항(entry)을 0 또는 1로 표시합니다.

3. 중간: MMDiT는 모달리티별 QKV 가중치를 가집니다. 완전 공유(fully-shared) Transfusion 트랜스포머 대비 파라미터 수 오버헤드(parameter count overhead)는 얼마나 됩니까? 7B 파라미터 규모에서 그 비용을 감수할 가치가 있습니까?

4. 어려움: 생성 비용을 계산해보세요. 텍스트 프롬프트가 주어졌고, 모델이 50개 토큰 동안 NTP를 실행한 뒤 <image>를 만나고, 이후 256개 패치에 대해 20번의 디노이즈 단계로 디퓨전을 실행합니다. 전체 순전파는 몇 번입니까?

5. 어려움: SD3 논문 Section 3을 읽으세요. 정류 흐름(rectified flow)이 무엇이고, DDPM보다 더 적은 추론 단계로 수렴하는 이유를 설명하세요.

핵심 용어

더 읽을거리

• Zhou et al. — Transfusion (arXiv:2408.11039)
• Esser et al. — Stable Diffusion 3 / MMDiT (arXiv:2403.03206)
• Peebles & Xie — DiT (arXiv:2212.09748)
• Zhao et al. — MonoFormer (arXiv:2409.16280)
• Xie et al. — Show-o (arXiv:2408.12528)