Transfusion: Predict the Next Token and Diffuse Images with One Multi-Modal Model

Training: how do you teach one model two skills?

Transfusion trains on a carefully balanced diet:

Total: 2 trillion tokens for the 7B model.

The training loop

1. Encode images — Run each image through a frozen VAE encoder → 256 continuous patch vectors per image
2. Build the mixed sequence — Interleave text tokens and image patches: [text] [BOI] [patches] [EOI] [more text]
3. Add noise to image patches — Sample a random diffusion timestep t, add Gaussian noise to each image patch
4. Forward pass — The transformer processes the entire mixed sequence
5. Compute dual loss — Cross-entropy for text positions, MSE for image positions
6. Backpropagate — Both losses flow through the shared transformer weights

Balancing: λ = 1 (equal weighting) works well — the two objectives are surprisingly compatible.

Diffusion inside a transformer

Each image patch gets a timestep embedding added before entering the transformer — this tells the model “these patches have noise level t.”

Inference: generating a mixed document

Step 1 — Text generation. The model generates text tokens one at a time, left to right, exactly like GPT. When it predicts [BOI] (beginning of image), it switches to image mode.

Step 2 — Image generation. Start from pure noise. Run N denoising passes through the transformer, each pass predicting and subtracting noise, producing progressively cleaner patches. After N steps, decode patches to pixels via the VAE decoder.

Step 3 — Resume text. The model continues generating text autoregressively, now conditioned on both the preceding text AND the generated image. When it predicts another [BOI], it generates another image.

Scaling experiments

The scaling curve tells two stories: text quality scales similarly for both approaches, but Transfusion scales much better for images. The gap grows with compute.

The competition

Limitations the authors acknowledge

Deep dive: GenEval

GenEval tests compositional image generation — not “does the image look pretty” but “did the model actually generate what you asked for?”

Scoring uses an object detector (DDETR) to check presence, color, position, and count.

Transfusion dominates on color attribution (0.40 vs 0.23 vs 0.01) — the hardest compositional skill. Still struggles with counting and position, which remain unsolved even for dedicated image models.

Deep dive: efficiency

Chameleon converts images into 1,024 discrete tokens, each processed like a text token. Attention cost scales as O(n²) with sequence length.

Saving #1 — Fewer patches.

Saving #2 — Compression at inference. 256 → 64 patches: FID goes from 6.78 to 7.03 (barely noticeable). 256 → 16 patches: FID goes to 11.0 (still decent, 64× cheaper). You can choose your quality–cost tradeoff at inference time.

Saving #3 — No codebook bottleneck. No quantization loss, no codebook collapse, no straight-through estimator hacks.

Saving #4 — Better scaling curve. At every FLOP budget, Transfusion produces better images. To match Transfusion at 10²¹ FLOPs, Chameleon needs ~3.3× more compute.

Deep dive: architecture

Shared backbone: All transformer layers (attention + FFN) process both text and image data. The model learns shared representations across modalities.

What’s separate: Input embeddings (text uses token embeddings; images use a VAE encoder + linear projection) and output heads (text uses softmax over vocabulary; images use a U-Net-inspired decoder).

Attention masking — the critical design choice:

• Text tokens: standard causal (left-to-right) attention
• Image patches: bidirectional attention within the same image — every patch sees every other patch
• Image patches CAN attend to all preceding text (conditioned on caption)
• Text after an image CAN attend to the image patches
• Patches from different images do NOT attend to each other

U-Net skip connections: Intermediate representations from earlier transformer layers are concatenated with the final layer output and fed to the image decoder. This gives access to both low-level features (textures, edges) and high-level semantics.

Mathematical formulations

The joint objective:

L_total = L_LM + λ · L_DDPM

Language modelling loss (standard next-token cross-entropy over text positions):

L_LM = -Σ_{i ∈ T} log P_θ(x_i | x_{<i})

Diffusion loss (noise prediction over image positions):

L_DDPM = Σ_{j ∈ I} E_{t,ε} [ || ε - ε_θ(z_j^(t), t, c) ||² ]

where z_j^(t) = √(ᾱ_t) · z_j + √(1 - ᾱ_t) · ε

Here z_j is the clean latent patch vector from the VAE encoder, t is a uniformly sampled timestep, ε is the Gaussian noise added, and ᾱ_t is the cumulative noise schedule (cosine schedule from Nichol & Dhariwal 2021).

The attention mask (formally)

Given text positions T and image positions I_k for the k-th image, the attention mask M has five rules:

Training pseudocode

def train_step(batch, model, vae_encoder):
    for sequence in batch:
        patches, text_tokens = [], []
        for element in sequence:
            if element.type == "text":
                text_tokens.append(model.text_embed(element.token_id))
            elif element.type == "image":
                z = vae_encoder(element.pixels)        # [256, d_vae]
                z = model.image_proj(z)                 # [256, d_model]
                t = randint(1, T)                       # sample timestep
                eps = randn_like(z)                     # sample noise
                z_noisy = sqrt(alpha_bar[t]) * z + sqrt(1 - alpha_bar[t]) * eps
                z_noisy += model.timestep_embed(t)      # add timestep info
                patches.append((z_noisy, eps, t))

        # Build mixed sequence, construct attention mask
        hidden = model.transformer(mixed_sequence, mask)

        # Compute dual loss
        lm_loss  = cross_entropy(hidden[text_pos], next_tokens)
        dif_loss = mse(model.image_decoder(hidden[img_pos]), target_noise)
        loss = lm_loss + lambda_ * dif_loss
        loss.backward()

Inference pseudocode

def generate_document(model, vae_decoder, prompt):
    tokens = tokenize(prompt)
    while not done:
        # --- Autoregressive text ---
        next_token = model.predict_next(tokens)
        tokens.append(next_token)
        if next_token == BOI:
            # --- Iterative image denoising ---
            latents = randn(256, d_vae)         # pure noise
            for step in range(N, 0, -1):
                t_emb = model.timestep_embed(step)
                pred_noise = model.forward_image(tokens, latents + t_emb)
                latents = ddpm_update(latents, pred_noise, step)
            image = vae_decoder(latents)
            tokens.append(image_repr)
            tokens.append(EOI)
    return tokens

DDPM update step:

z^(t-1) = (1/√α_t) · (z^(t) - (1-α_t)/√(1-ᾱ_t) · ε_θ) + σ_t · w

where w ~ N(0, I) for t > 1, σ_1 = 0

Critical evaluation

Where this paper sits in the field

Transfusion received an ICLR 2025 Oral (top ~1% of submissions). Here’s what it sparked and what comes next.

What happened since August 2024

Improvement vectors

1. High resolution

256px is a toy. 1024px with 8× downsampling = 16,384 patches — impossible with standard O(n²) attention. Paths forward: hierarchical generation (generate low-res, super-resolve), windowed/sparse attention, progressive patch compression (coarse-to-fine), or latent cascade with a lightweight upsampler.

2. Multi-image consistency

Generate a 4-page comic → the fox looks different on every page. Paths: entity embeddings (persistent identity vector injected into each image), cross-image attention (let later image patches attend to earlier images), reference image conditioning. The attention mask needs a new rule allowing later images to reference earlier ones.

3. Faster inference

250 denoising steps × full transformer forward pass = very slow. Paths: DDIM (→ 50 steps), DPM-Solver++ (→ 20–25), consistency distillation (→ 1–4), flow matching (→ 10–20). Flow matching (behind SD3 and Flux) is the likely winner — mathematically simpler, straighter trajectories, fewer steps. Fully compatible with Transfusion’s framework.

4. Controllability

No way to condition on pose, depth, edges, or style references. Paths: ControlNet-style adapters, concatenating control embeddings with image patches, T2I-Adapter approach. The continuous patch representation makes this easier than adding control to discrete-token approaches.

5. Video extension

The framework naturally extends to video. Each frame → patches with bidirectional intra-frame attention, plus temporal attention across frames. Frames can be diffused jointly or autoregressively (generate frame 1, condition frame 2 on it).

6. Better VAE

Image quality is capped by VAE reconstruction quality. Low-hanging fruit: swapping to SDXL’s VAE (~20% better reconstruction), consistency decoder for sharper high-frequency detail, adversarial training (VAE-GAN), or higher-resolution latents with less spatial downsampling.

Scorecard