PaperlayerView exact revision on arXiv

Source-linked AI summary

Scalable Diffusion Models with Transformers

William Peebles, Saining Xie

arXiv:2212.09748v2cs.CVcs.LG

TL;DR

Diffusion models have largely relied on convolutional U-Nets, leaving the role of transformer backbones and their scaling behavior less established. This paper introduces latent Diffusion Transformers and shows that scaling Gflops improves FID, reaching 2.27 FID on class-conditional 256×256 ImageNet.

  • Problem

    Diffusion models have predominantly used convolutional U-Net backbones, while the role and scalability of transformer architectures remained less established.

  • Method

    The paper introduces Diffusion Transformers that replace the U-Net with a transformer operating on latent patches and evaluates scaling through Gflops and FID.

  • Results

    Higher DiT Gflops consistently correspond to lower FID, with DiT-XL/2 achieving a state-of-the-art FID of 2.27 on class-conditional 256×256 ImageNet.

  • Takeaways & Limitations

    DiTs show that diffusion models can use transformer backbones while retaining strong scaling properties and outperforming prior diffusion models.

Abstract

from arXiv · show

We explore a new class of diffusion models based on the transformer architecture. We train latent diffusion models of images, replacing the commonly-used U-Net backbone with a transformer that operates on latent patches. We analyze the scalability of our Diffusion Transformers (DiTs) through the lens of forward pass complexity as measured by Gflops. We find that DiTs with higher Gflops -- through increased transformer depth/width or increased number of input tokens -- consistently have lower FID. In addition to possessing good scalability properties, our largest DiT-XL/2 models outperform all prior diffusion models on the class-conditional ImageNet 512x512 and 256x256 benchmarks, achieving a state-of-the-art FID of 2.27 on the latter.

1. Introduction

The paper introduces Diffusion Transformers (DiTs), replacing diffusion models’ U-Net backbone with a transformer operating in a latent-space framework. It investigates how transformer network complexity relates to diffusion sample quality and argues that diffusion models can benefit from more unified architectural designs.

  • Motivation: Transformers have broadly displaced other neural architectures in language and vision, but remain less adopted in several image-level generative-modeling frameworks.The introduction positions diffusion models among these remaining holdouts despite transformers’ widespread use in autoregressive image models.
  • Empirical scaling: Higher-compute DiT models show steadily improving FID-50K, while DiT-XL/2 outperforms prior U-Net-based diffusion models.The comparison uses models trained for 400K iterations; lower FID-50K is better.
  • Architectural motivation: The authors argue that the U-Net inductive bias is not essential, because standard transformer designs can readily replace it in diffusion models.This supports using diffusion models to investigate broader architectural unification and shared best practices.
  • Diffusion Transformers: DiTs replace the conventional U-Net backbone with a transformer within the Latent Diffusion Models framework.The models operate in a VAE’s latent space and follow Vision Transformer design practices.
  • Scaling behavior: The study examines transformer scaling by relating network complexity to diffusion sample quality.The DiT design space is constructed and benchmarked to study this relationship.

2. Related Work

Transformers have become broadly applicable across language, vision, reinforcement learning, and meta-learning, with strong scaling behavior in language and vision. Diffusion and score-based models have achieved major success in image generation, often surpassing GANs, while DDPM progress has been driven largely by improved sampling.

  • Transformers: Transformers have replaced domain-specific architectures across language, vision, reinforcement learning, and meta-learning.
  • Transformers: Transformers exhibit strong scaling properties as model size, training compute, and data increase in language models and vision transformers.
  • Transformers: Transformers have also been used to autoregressively predict pixels and to model discrete codebooks.
  • Denoising diffusion probabilistic models (DDPMs): Diffusion and score-based generative models have been particularly successful for image generation, often outperforming generative adversarial networks.
  • Denoising diffusion probabilistic models (DDPMs): Recent DDPM improvements have largely come from better sampling techniques, notably classifier-free guidance.

3. Diffusion Transformers

DiTs adapt the Vision Transformer architecture to latent diffusion by patchifying noised latent representations into token sequences, processing them with transformer blocks, and decoding tokens into noise and covariance predictions. The design explores conditioning mechanisms while emphasizing compute-efficient latent-space operation.

  • Latent diffusion models: Latent diffusion compresses images with a frozen autoencoder before training the diffusion model, reducing computation relative to pixel-space diffusion.LDMs use an encoder E to produce spatial representations z, then generate images by sampling representations and decoding them.
  • Patchify: DiT uses a Vision Transformer backbone that patchifies noised latent representations into token sequences and adds sine-cosine positional embeddings.For 256 × 256 × 3 images, the latent input has shape 32 × 32 × 4; patch size p determines the number of tokens.
  • Patchify: A smaller patch size produces longer token sequences and more Gflops, making patch size a direct lever for DiT computation.The sequence length is T = (I/p)^2, so reducing p increases T.
  • DiT block design: 15% overhead: cross-attention adds roughly this much computation, whereas in-context conditioning introduces negligible new Gflops.The explored conditioning variants process timestep and class information either as additional tokens, through cross-attention, or with adaptive normalization.
  • Transformer decoder: The transformer decoder applies final normalization and linearly maps each token to p×p×2C outputs before rearranging them into spatial noise and covariance predictions.Both decoded outputs match the original spatial input shape.

4. Experimental Setup

The experiments study DiT scaling on class-conditional ImageNet at 256 × 256 and 512 × 512 resolutions, using latent diffusion models trained with standardized optimization and diffusion settings. Scaling is evaluated primarily with FID-50K under conventional sampling and implementation protocols.

  • Model design: DiT models are named by transformer configuration and latent patch size, with DiT-XL/2 denoting the XLarge configuration and p = 2.The experiments explore the DiT design space and scaling properties.
  • Training: Class-conditional latent DiT models are trained on ImageNet at 256 × 256 and 512 × 512 resolutions using AdamW.The final linear layer is zero-initialized, while other weights use standard ViT initialization.
  • Training: 1 × 10−4 is the constant learning rate, with no weight decay, batch size 256, and horizontal flips as the only augmentation.Training is reported as stable across configurations without learning-rate warmup or regularization.
  • Evaluation metrics: FID-50K is reported using 250 DDPM sampling steps and ADM’s TensorFlow evaluation suite for comparisons with prior work.Classifier-free guidance is excluded unless otherwise stated; Inception Score, sFID, and Precision/Recall are also reported.
  • Compute: 5.7 iterations/second is the approximate DiT-XL/2 training speed on a TPU v3-256 pod with global batch size 256.All models are implemented in JAX and trained using TPU-v3 pods.

5. Experiments

Experiments show that increasing DiT model compute—through depth, width, or token count—consistently improves FID and visual quality, with larger models using training compute more efficiently. DiT-XL/2 achieves state-of-the-art ImageNet results at both 256×256 and 512×512 resolutions.

  • Scaling model size and patch size: Increasing model size and decreasing patch size considerably improves diffusion-model quality across DiT configurations.Deeper and wider transformers improve FID, as does processing more tokens by decreasing patch size.
  • DiT Gflops are critical to improving performance: Scaling DiT Gflops is key to improving performance even when parameter counts remain approximately unchanged.Decreasing patch size increases Gflops while leaving total parameters effectively unchanged or slightly reduced.
  • Larger DiT models are more compute-efficient: Larger DiT models are more compute-efficient than smaller models trained longer, achieving better FID with fewer training steps.Training compute is estimated as model Gflops · batch size · training steps · 3.
  • 256×256 ImageNet: 2.27 FID-50K: DiT-XL/2 outperforms all prior diffusion models on class-conditional ImageNet 256×256, improving LDM’s 3.60.The model is continued for 7M steps and uses classifier-free guidance.
  • 512×512 ImageNet: 3.04 FID: DiT-XL/2 outperforms all prior diffusion models on class-conditional ImageNet 512×512, improving ADM’s 3.85.The 512×512 model processes 1024 tokens and uses 524.6 Gflops.

6. Conclusion

The paper introduces Diffusion Transformers (DiTs), a transformer-based diffusion backbone that outperforms prior U-Net models while retaining strong transformer scaling properties. Future work should scale DiTs further and explore them as backbones for text-to-image models.

  • 6. Conclusion: DiTs are a simple transformer-based diffusion backbone that outperforms prior U-Net models and inherits the transformer class’s scaling properties.The conclusion identifies DiTs as a backbone for diffusion models with promising scaling results.
  • 6. Conclusion: Future work should scale DiTs to larger models and token counts.The proposed direction follows the paper’s reported scaling results.
  • 6. Conclusion: DiTs could serve as drop-in backbones for text-to-image models such as DALL·E 2 and Stable Diffusion.The conclusion names these models as examples for future exploration.

A. Additional Implementation Details

The section provides implementation details for DiT models, including conditioning, transformer nonlinearities, model reporting, and training diagnostics. It also shows that classifier-free guidance on only three latent channels can achieve similar FID to four-channel guidance after adjusting the scale.

  • Table 4 reports DiT model details for both 256 × 256 and 512 × 512 models, while Figure 13 provides training loss curves and Table 6 lists Gflop counts for ADM and LDM DDPM U-Nets.
  • DiT timestep conditioning uses a 256-dimensional frequency embedding followed by a two-layer MLP, while adaLN layers combine timestep and class embeddings before projecting to 4× or 6× hidden size.The 4× output applies to adaLN, and the 6× output applies to adaLN-Zero.
  • The core transformer uses GELU nonlinearities approximated with tanh.
  • 2.27 FID-50K is obtained with three-channel guidance at scale 1.5, compared with 2.20 FID-50K for four-channel guidance at scale 1.375.Guidance was applied to the first three latent channels instead of all four, and adjusting the scale factor produced similar FID results.

B. Model Samples

The section presents DiT-XL/2 samples at 512 × 512 and 256 × 256 resolutions, including uncurated generations across classifier-free guidance scales and class labels. Larger guidance scales increase visual fidelity while decreasing sample diversity, and training continued without observed FID saturation for either model.

  • Model Samples: DiT-XL/2 samples are shown at 512 × 512 and 256 × 256 resolutions after 3M and 7M training steps, respectively.Figures 1 and 11 show selected samples from both models.
  • Model Samples: Uncurated samples span classifier-free guidance scales and input class labels, using 250 DDPM sampling steps and the ft-EMA VAE decoder.Figures 14 through 33 provide samples across these settings.
  • Model Samples: Larger classifier-free guidance scales increase visual fidelity and decrease sample diversity.This observation is reported for the displayed samples, consistent with prior work using guidance.
  • Model Samples: FID did not saturate for either the 256 × 256 or 512 × 512 DiT-XL/2 model during continued training.Table 4 reports FID-50K without classifier-free guidance and uses the ft-MSE VAE decoder.

C. Additional Scaling Results

Additional scaling results show that larger DiT models are more compute-efficient across evaluation metrics and achieve faster, lower-saturating training losses as Gflops increase.

  • Impact of scaling on metrics beyond FID: Across FID, sFID, Inception Score, Precision, and Recall, scaled-up DiTs are more compute-efficient, with model Gflops highly correlated with performance.Inception Score and Precision benefit especially strongly from increased model scale.
  • Impact of scaling on training loss: Increased DiT Gflops, through larger transformers or more input tokens, makes training loss decrease faster and saturate at a lower value.This trend is consistent with observations for scaled-up transformers in language models.

D. VAE Decoder Ablations

The decoder ablation compares the original LDM decoder with two Stable Diffusion fine-tuned decoders, which can be swapped without retraining because the encoders are identical. Different pre-trained decoder weights produce comparable ImageNet 256 × 256 results, while XL/2 still outperforms prior diffusion models with the LDM decoder.

  • Ablation result: Different pre-trained decoder weights yield comparable results on ImageNet 256 × 256.This is the main result reported in the decoder ablation table.
  • Decoder choices: Three VAE decoder choices were tested: the original LDM decoder and two fine-tuned Stable Diffusion decoders.The ft-MSE and ft-EMA models fine-tune only decoder weights from the original LDM “f8” model.
  • Decoder choices: Because the encoders are identical, decoders can be swapped in without retraining the diffusion model.This isolates decoder effects while keeping the diffusion model fixed.
  • Ablation result: XL/2 continues to outperform all prior diffusion models when using the LDM decoder.The decoder swap therefore does not remove XL/2’s reported advantage in this comparison.
Loading 2212.09748v2…