Transfusion: Unified Autoregressive Text + Diffusion Image Generation.

A single transformer model that jointly handles discrete text tokens (autoregressive next-token prediction) and continuous image patches (denoising diffusion) in one shared backbone.

Key Innovation: Mixed Attention Mask

Text tokens and image patches share the same transformer layers, but attend with different masks:

• Text positions (causal): each token sees only preceding tokens
• Image positions (bidirectional within image): each patch sees all other patches in the same image, plus all preceding text context

Combined rule:

mask[i, j] = 1  if  j ≤ i                        # causal for text
             OR  (image[i] AND image[j])            # bidir within image

Architecture

Inputs: sequence [batch, seq_len, embed_dim]  (text embeddings + image patches)
        modality_mask [batch, seq_len]          (0=text, 1=image)
        timestep [batch]                        (diffusion step for image)
        |
        v
Modality type embedding  (learnable TEXT / IMAGE vectors added to tokens)
        |
        v
Input projection  →  hidden_size
        |
        v
+----------------------------------------------+
|  Transfusion Block  ×  num_layers             |
|                                               |
|  Add time_embed at image positions            |
|  LayerNorm  →  Mixed Attention  →  Residual   |
|  LayerNorm  →  FFN (GELU)       →  Residual   |
+----------------------------------------------+
        |
Final LayerNorm
        |
     ┌──┴──┐
text_head    image_head
[b,s,V]     [b,s,P]

Dual Loss

• Text tokens: cross-entropy against next-token targets
• Image patches: MSE between predicted and target noise/velocity
• Total: text_weight * L_CE + image_weight * L_MSE

Usage

model = Transfusion.build(
  embed_dim: 64,
  hidden_size: 256,
  num_heads: 8,
  num_layers: 6,
  vocab_size: 32_000,
  patch_dim: 64
)

# Build the mixed attention mask for a 20-token text + 16-patch image
mask = Transfusion.build_mixed_mask(20, 16)

# Compute training loss
loss = Transfusion.transfusion_loss(text_logits, image_pred, %{
  text_targets:  token_ids,      # [batch, seq_len] integer indices
  image_targets: noise_targets,  # [batch, seq_len, patch_dim]
  text_mask:     text_positions, # [batch, seq_len] float, 1 at text positions
  image_mask:    image_positions # [batch, seq_len] float, 1 at image positions
})

References

• Paper: "Transfusion: Predict the Next Token and Diffuse Images with One Multi-Modal Model"
• Authors: Chunting Zhou et al., Meta (2024)
• arXiv: https://arxiv.org/abs/2408.11039