Scalable Diffusion Models with Transformers

Written for the Research library (claude-opus-5-5) from the paper text. Original — arxiv.org/pdf/2212.09748

Summary

Image diffusion models have almost always used a convolutional U-Net as their denoising network, and it was unclear whether that design is actually needed. Peebles and Xie test this by building Diffusion Transformers (DiTs) within the Latent Diffusion Models (LDMs) framework. A frozen, off-the-shelf VAE compresses each image into a latent. The latent is cut into patches, and a ViT-style transformer processes those patches as tokens to predict noise. The authors compare several ways of feeding in the timestep and class label: extra in-context tokens, cross-attention and adaptive layer norm. The adaptive norm variant works best. They train a grid of model sizes and patch sizes on class-conditional ImageNet. They measure complexity in Gflops instead of parameter count, and find that sample quality improves steadily as per-pass compute grows. Extra compute helps whether it comes from a deeper or wider network or from more tokens. The largest model, DiT-XL/2, sets a new state of the art on ImageNet 256×256 and needs far less compute than pixel-space U-Nets. The work suggests diffusion can adopt the same scalable transformer backbone used in other fields.

Key points

Models are named by size (S, B, L, XL) and patch size (/8, /4, /2), giving 12 configurations in total. Halving the patch size makes the token sequence four times longer and raises Gflops accordingly. Model Gflops correlate strongly and negatively with FID.
On ImageNet 256×256 with classifier-free guidance (cfg=1.50), DiT-XL/2 reaches FID 2.27 after 7M steps. This beats the previous diffusion best of 3.60 (LDM-4-G) and also StyleGAN-XL (2.30). With only 2.35M steps, it still reaches 2.55.
On ImageNet 512×512, DiT-XL/2 (1024 tokens, 524.6 Gflops) reaches FID 3.04, improving on ADM's 3.85. ADM needs 1983 Gflops for this resolution and ADM-U needs 2813 Gflops.
Compute efficiency at 256×256: DiT-XL/2 uses 118.6 Gflops, compared with 103.6 for LDM-4, 1120 for ADM and 742 for ADM-U.
Extra sampling steps cannot replace model capacity. DiT-XL/2 with 128 steps (15.2 Tflops per image) gets FID-10K 23.7. DiT-L/2 with 1000 steps (80.7 Tflops) only gets 25.9.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields: Other

Cited by

412 papers of our library cite it (of the 140673 whose bibliographies we read); 5 of them have a page here, 407 not yet.

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2212.09748 · arXiv 2212.09748

@misc{peebles2022scalable,
  title = {{Scalable Diffusion Models with Transformers}},
  author = {Peebles, William and Xie, Saining},
  year = {2022},
  doi = {10.48550/arxiv.2212.09748},
  eprint = {2212.09748},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2212.09748}
}

References

15 references, 1 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. 6 find
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale GAN training for high fidelity natural image synthesis. In ICLR, 2019. 5, 9 find
Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In NeurIPS, 2020. 1 find
Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In CVPR, pages 11315–11325, 2022. 2 find
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In NeurIPS, 2021. 2 find
Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Ilya Sutskever. Generative pre-training from pixels. In ICML, 2020. 1, 2 find
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019. 2 find
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In NAACL-HCT, 2019. 1 find
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. In NeurIPS, 2021. 1, 2, 3, 5, 6, 9, 12 find
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2020. 1, 2, 4, 5 find
Patrick Esser, Robin Rombach, and Björn Ommer. Taming transformers for high-resolution image synthesis, 2020. 2 find
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In NIPS, 2014. 3 find
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv:1706.02677, 2017. 5 find
Shuyang Gu, Dong Chen, Jianmin Bao, Fang Wen, Bo Zhang, Dongdong Chen, Lu Yuan, and Baining Guo. Vector quantized diffusion model for text-to-image synthesis. In CVPR, pages 10696–10706, 2022. 2 find
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In CVPR, 2016. 2 [50] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S. Sara Mahdavi, Rapha Gontijo Lopes, Tim Sa… find our card