Source-linked AI summary
Transfusion: Predict the Next Token and Diffuse Images with One Multi-Modal Model
Chunting Zhou, Lili Yu, Arun Babu, Kushal Tirumala, Michihiro Yasunaga, Leonid Shamis, Jacob Kahn, Xuezhe Ma, Luke Zettlemoyer, Omer Levy
TL;DR
Existing multimodal models need to handle discrete text and continuous images, but common approaches combine separate modeling strategies or discretize images. Transfusion trains one transformer with language-modeling and diffusion objectives, and its experiments show better scaling than image discretization while supporting competitive text and image generation. A 7B model trained on 2T multimodal tokens reaches comparable quality to similar-scale language and diffusion models, though the broader objective space and architecture choices remain incompletely explored.
Problem
Multimodal models must model both discrete text and continuous media, while existing approaches include combining separate models or quantizing images into discrete tokens.
Method
Transfusion trains a single transformer on mixed text and image data, using next-token prediction for text and diffusion for images with modality-specific encoding and decoding.
Results
Transfusion scales better than Chameleon across modality combinations, while a 7B model trained on 2T multimodal tokens achieves competitive image and text generation performance.
Takeaways & Limitations
A shared model can generate discrete and continuous modalities using their respective modeling objectives, with U-Net layers enabling larger image patches and potentially lower serving costs.
Takeaways & Limitations
The broader combination of discrete and continuous losses, including alternatives such as flow matching, remains for future work; the study primarily uses transformers.
Abstract
from arXiv · showhide
We introduce Transfusion, a recipe for training a multi-modal model over discrete and continuous data. Transfusion combines the language modeling loss function (next token prediction) with diffusion to train a single transformer over mixed-modality sequences. We pretrain multiple Transfusion models up to 7B parameters from scratch on a mixture of text and image data, establishing scaling laws with respect to a variety of uni- and cross-modal benchmarks. Our experiments show that Transfusion scales significantly better than quantizing images and training a language model over discrete image tokens. By introducing modality-specific encoding and decoding layers, we can further improve the performance of Transfusion models, and even compress each image to just 16 patches. We further demonstrate that scaling our Transfusion recipe to 7B parameters and 2T multi-modal tokens produces a model that can generate images and text on a par with similar scale diffusion models and language models, reaping the benefits of both worlds.
1 Introduction
Transfusion integrates next-token prediction for text and diffusion for images in one transformer, targeting unified multimodal generation. It scales better than image discretization in controlled comparisons and supports competitive text and image generation at 7B parameters.
- Approach: Transfusion trains one transformer on mixed text and image data using next-token prediction for text and diffusion for images.The training mixture contains 50% text and 50% image data, with both modalities and losses present at each step.
- Architecture: Text tokens use causal attention, while image patches use bidirectional attention and modality markers separate the sequences.Inference combines language-model text generation with diffusion-based image generation.
- Scaling comparisons: Transfusion scales better than Chameleon’s discretization approach across modality combinations, reaching approximately 2× lower FID at matched FLOPs.It exceeds Chameleon in text-to-image generation at less than one-third the compute and matches it for image-to-text at 21.8% of the FLOPs.
- Ablations: U-Net image encoding and decoding preserves performance with larger patches and may reduce serving costs by up to 64×.The ablation also finds that replacing bidirectional image attention with causal attention hurts text-to-image generation.
- Large-scale demonstration: A 7B Transfusion trained on 2T multimodal tokens generates images at similar quality to diffusion models while matching Llama 1’s text-benchmark performance.On GenEval, it outperforms DALL-E 2 and SDXL; the model uses U-Net layers totaling 0.27B additional parameters.
2 Background
The background contrasts autoregressive language modeling for discrete tokens with diffusion for continuous vectors such as images. It also introduces latent image representations, which reduce the computational cost of diffusion-based image modeling.
- Language modeling: Language models predict each discrete token conditioned on the preceding prefix and optimize cross-entropy, producing the next-token or LM loss.The sequence probability is decomposed into conditional token probabilities over a fixed vocabulary.
- Language modeling: Language models generate text by sampling tokens sequentially from the model distribution, commonly using temperature and top-p truncation.This generation procedure follows the autoregressive factorization used during training.
- Diffusion: Diffusion models learn to reverse gradual Gaussian noise addition over continuous vectors, making them suited to image generation.The forward process creates increasingly noisy states, while the reverse model denoises them step by step.
- Diffusion: A single-noise reparameterization can sample a noisy state directly from the original data, and the noise schedule controls its variance.The model estimates accumulated noise from the noisy input and timestep; inference iteratively removes noise from Gaussian initialization.
- Latent image representation: VAEs encode images into lower-dimensional latent spaces so downstream latent diffusion models can operate more efficiently.Figure 3 shows conversion between images, VAE latents, and patch representations using either a linear layer or U-Net down blocks.
3 Transfusion
Transfusion trains one unified transformer with separate language-modeling and diffusion objectives for text and images. It combines modality-specific representations, attention patterns, and decoding procedures to handle mixed-modality sequences.
- Training objective: Transfusion trains a single model with language modeling for text and diffusion for images over shared data and parameters.The formulation combines a discrete distribution loss with a continuous distribution loss.
- Data representation: Text is tokenized into discrete integers, while images are encoded as continuous latent patch vectors using a VAE.Image patches are sequenced left-to-right and top-to-bottom.
- Model architecture: A shared transformer processes every modality, with lightweight modality-specific components converting text and image representations to and from its vector space.Image processing can use either a linear layer or U-Net down and up blocks.
- Attention: Transfusion applies causal attention across the sequence but bidirectional attention among patches within each image.This allows every image patch to attend to every other patch in the same image.
- Training objective: The training loss adds per-token language-modeling loss to per-image diffusion loss using balancing coefficient λ.Noise is added to latent images before patchification to compute the diffusion loss.
- Inference: During inference, the decoder switches from token-by-token language-modeling mode to diffusion mode when it samples a BOI token.Diffusion mode appends noisy image patches and follows standard diffusion decoding.
4 Experiments
Experiments evaluate Transfusion across uni-modal and cross-modal benchmarks using controlled model, data, and compute settings. The setup compares it with Chameleon’s image-tokenization approach and examines scaling, compression, and large-scale training.
- Evaluation: Transfusion is evaluated on text-to-text, text-to-image, and image-to-text benchmarks, including perplexity, task accuracy, FID, CLIP, CIDEr, and GenEval.The evaluation suite supports scaling-law analyses, ablations, and comparison with diffusion models.
- Baseline: The main baseline is Chameleon, which quantizes images into discrete tokens and models the resulting sequences with a standard language model.The baseline models are trained with controlled data and compute for direct comparison.
- Data: Most experiments use 0.5T tokens or patches sampled from text and image data at a 1:1 token ratio.The image data consists primarily of 380M licensed Shutterstock images and captions.
- Image representation: Images are represented with VAE latents, reducing each 256×256 image to a 32×32×8 tensor before patch-based processing.Depending on patch-encoder compression, images contain 1024, 256, 64, or 16 sequence elements.
- Model scaling: Models are trained at 0.16B, 0.37B, 0.76B, 1.4B, and 7B parameters to investigate scaling trends.U-Net patch encoding adds 0.27B parameters across configurations, while linear encoding adds fewer than 0.5% of total parameters.
- Large-scale training: The large-scale experiment trains on 2T tokens, comprising 1T text tokens and approximately 1T image patches or tokens.It uses 692M image-caption pairs per epoch after adding public images, upsampled Shutterstock images, and CC12M data.
4.2 Controlled Comparison with Chameleon
Controlled experiments show that Transfusion scales more efficiently than Chameleon across modalities, while architectural ablations identify bidirectional attention and U-Net layers as important design choices.
- Scaling comparison: Transfusion exhibits better scaling laws than Chameleon in every benchmark, with a particularly large compute-efficiency advantage for image generation.At equal performance, Transfusion’s FID reaches parity with Chameleon using 34× less compute.
- Scaling comparison: Training on quantized image tokens degrades text performance more than using diffusion, although Transfusion still incurs a non-zero text-performance cost.The text-only ablation compares Transfusion and Chameleon against the original Llama 2 recipe.
- Scaling comparison: Transfusion uses continuous image representations, so theoretical FLOPs are used to avoid giving it an advantage from shorter sequences and lower quadratic attention cost.The comparison uses 6ND as a proxy for compute and removes the token-length confounder.
- Architecture ablations: Intra-image bidirectional attention improves performance across benchmarks, including a FID improvement from 61.3 to 20.3 with linear encoding layers.Causal-only attention prevents later patches from informing earlier patches; U-Net blocks partly reduce this gap through their internal bidirectional attention.
- Architecture ablations: Larger patches reduce inference compute but create modality-dependent trade-offs: U-Net variants benefit on image tasks, while text performance deteriorates.With linear encoding, performance decreases as images use fewer patches; U-Net encoding changes this trend for image-related tasks.
- Architecture ablations: U-Net encoder/decoder layers retain benefits as the transformer grows, improving smaller models beyond a linear 7B model on image generation and captioning.The U-Net contribution is smaller relatively at larger scale but does not disappear.
- Image noising: Limiting diffusion noise to t = 500 when images precede captions significantly improves CIDEr while changing other benchmarks by less than 1%.This intervention targets the 20% of image-caption pairs ordered with the image first.
4.4 Comparison with Image Generation Literature
A 7B Transfusion model with U-Net encoding and decoding achieves image-generation performance comparable to high-performing diffusion models while also generating text, and it adapts to image editing after limited fine-tuning.
- Large-scale comparison: A 7B Transfusion model is trained on the equivalent of 2T tokens using U-Net encoding/decoding layers and 2×2 latent pixel patches.The training mixture contains 1T text corpus tokens and 3.5B images with captions.
- Large-scale comparison: Transfusion matches high-performing image-generation models such as DeepFloyd and surpasses previously published models including SDXL.It lags behind SD 3, whose GenEval performance benefits from synthetic caption data; Transfusion uses natural data in this setup.
- Large-scale comparison: The model performs on par with Llama models on text generation while also generating images.The Llama models use the same text-data distribution for comparison.
- Image editing: Fine-tuning on 8k image-editing examples produces edits that manual examination finds follow the instructions on the EmuEdit test set.The authors describe this as evidence that Transfusion can adapt to new modality combinations, while noting limitations of the experiment.
5 Related Work
Prior multi-modal systems often combine modality-specific components or discretize images, but these approaches can be limited in task scope or continuous-data generation quality.
- Modality-specific systems: Many existing multi-modal models attach multiple modality-specific architectures, often pretrained separately, to combine text and continuous modalities.Image and video generation systems commonly use pretrained text encoders to condition diffusion models.
- End-to-end models: Fuyu uses image patches for visual understanding, while Chameleon discretizes images into token sequences and trains over combined text-image sequences.These are cited as prior end-to-end multi-modal examples.
- End-to-end models: Prior approaches may be restricted to input-level multi-modal tasks or lag behind diffusion models in continuous-data generation.The related-work discussion presents Transfusion as an end-to-end alternative for understanding and generating multi-modal data.
- Diffusion for text: Diffusion-based approaches to discrete text generation have not yet reached the performance and scale of standard autoregressive language models.The paper identifies combining discrete and continuous modalities through this direction as future research.
6 Conclusion
Transfusion bridges discrete sequence modeling and continuous media generation by training one joint model with modality-specific objectives. The experiments report efficient scaling and generation across modalities.
- Transfusion trains a single joint model with next token prediction for discrete data and diffusion for continuous media.
A Autoencoder Details
The autoencoder objectives combine reconstruction, perceptual, adversarial, and regularization terms, with VQ-VAE replacing KL regularization by codebook commitment loss. The VAE and VQ-VAE otherwise share the same training setup.
- VAE: The VAE objective combines pixel reconstruction, LPIPS perceptual, GAN, internal-feature, and KL-regularization losses.GAN training begins after 50,000 steps, and the latent dimension is 8.
- VQ-VAE: VQ-VAE replaces the VAE’s KL loss with a codebook commitment loss encouraging encoder outputs and codebook vectors to be close.The setup uses β = 0.25 with loss weighting 1.0.
- Comparison: The VAE and VQ-VAE use the same training compute, data, and encoder/decoder architecture apart from quantization and the loss-function change.
B Examples: Image Generation
Figures 7 and 8 present generated-image examples from a 7B Transfusion model trained on 2T multi-modal tokens.
- Figures 7 and 8 show generated images from a 7B Transfusion model trained on 2T multi-modal tokens.
C Examples: Image Editing
The paper presents image-editing examples from a fine-tuned 7B Transfusion model, alongside generated-image examples and prompts illustrating diverse visual content.
- Image Editing: Figure 9 shows random image-editing examples from a fine-tuned 7B Transfusion model.
- Prompt Examples: Additional prompts depict a light bulb containing a sailboat floating through the galaxy, a horse reading a book, and other object-focused scenes.
- Image Editing: Figure 9 provides edited-image examples from the fine-tuned model.