DiT - A Diffusion Model Based on the Transformer Architecture

What is DiT

DiT (Diffusion Transformers) is a novel diffusion model proposed by William Peebles (one of the heads of Sora's research and development) and Xie Saining, an assistant professor at New York University, combining Denoising Diffusion Probabilistic Models (DDPMs) and the Transformer architecture. A diffusion model is a generative model that generates new samples by simulating the gradual denoising process of data. The core idea of DiT is to use the Transformer as the backbone network of the diffusion model, instead of the traditional convolutional neural network (such as U-Net), to process the latent representation of images. Recently, along with OpenAI's video generation modelSorabecoming hugely popular, DiT has attracted widespread attention as one of the technical foundations behind Sora.

In DiT, images are first compressed by an autoencoder (such as a Variational Autoencoder, VAE) into a smaller latent representation, and then the diffusion model is trained in this latent space. The advantage of doing this is that it can reduce the amount of computation required to train the diffusion model directly in the high-resolution pixel space. The DiT model processes these latent representations through the Transformer's self-attention mechanism, which enables the model to capture long-range dependencies in images and thereby generate high-quality images.

DiT Official Website Entry

• Official project homepage:https://www.wpeebles.com/DiT
• ArXiv research paper:https://arxiv.org/pdf/2212.09748.pdf
• GitHub code repository:https://github.com/facebookresearch/DiT
• Hugging Face space:https://huggingface.co/spaces/wpeebles/DiT
• Replicate Demo：https://replicate.com/arielreplicate/scalable_diffusion_with_transformers
• Google Colab run address:http://colab.research.google.com/github/facebookresearch/DiT/blob/main/run_DiT.ipynb

Technical Principles of DiT

1. Data Preparation：
   • A pretrained Variational Autoencoder (VAE) is used to encode the input image into a representation in latent space. This latent representation is usually a low-dimensional representation of the image, for example, encoding a 256×256×3 RGB image into a 32×32×4 latent vector.
   • This latent representation is then used as the input to the DiT model.
2. Patchification：
   • The input latent representation is first divided through a process called "patchify" into a series of small patches, each corresponding to an input token of the Transformer model. This process is similar to dividing an image into small blocks so that the model can process it block by block.
   • Each patch is converted through linear embedding into a vector of fixed dimension, and then positional embeddings are added so that the model can understand the position of the patch in the image.
3. Transformer Blocks Module：
   • The input token sequence is processed through a series of Transformer blocks. These blocks include components such as self-attention layers, feedforward neural networks, and layer normalization.
   • In DiT, the researchers tried different Transformer block designs, including adaptive layer normalization (adaLN), cross-attention, and in-context conditioning, to handle conditional information such as timesteps and class labels.
4. Conditional Diffusion Process：
   • During training, the DiT model learns the reverse diffusion process, that is, recovering a clear image from noisy data. This process involves predicting the statistical properties of the noise, such as the mean and covariance.
   • The Variational Lower Bound (VLB) is used to train the model, which involves minimizing the mean squared error (MSE) between the predicted noise and the true noise.
5. Sample Generation：
   • After training is complete, new images can be generated through the DiT model. First, a latent representation is sampled from a standard normal distribution, and then the reverse diffusion process is carried out through the DiT model to gradually remove noise and finally decode back to pixel space to obtain the generated image.
6. Scalability：
   • The scalability of the DiT model is reflected in improving the model's computational load (Gflops) by increasing the number of Transformer layers, width, or the number of input tokens, thereby improving the quality of generated images. This scalability enables the DiT model to generate high-quality images at different resolutions and levels of complexity.

In this way, the DiT model leverages the powerful expressive capability of Transformers and the generative capability of diffusion models to achieve efficient and high-quality output in image generation tasks.

Key Features of DiT

• Transformer-based architecture：DiT adopts Transformer as its core architecture, which enables the model to process serialized representations of images and capture long-range dependencies in images through the self-attention mechanism.
• Latent space operation：DiT is trained in latent space, which is usually more efficient than training directly in pixel space. By using a variational autoencoder (VAE) to encode images into latent space, DiT reduces computational complexity.
• Scalability：DiT demonstrates excellent scalability. By increasing the model's computational load (measured in Gflops), the quality of generated images can be significantly improved. This scalability allows DiT to generate images at different resolutions and levels of complexity.
• Conditional generation capability：DiT supports conditional generation and can generate images of a specific category according to a given category label. This capability makes DiT very useful in image generation tasks in specific domains.
• Adaptive layer normalization (adaLN)：DiT uses adaptive layer normalization technology, which is a normalization method used in Transformer blocks that adjusts the parameters of layer normalization through learning, thereby improving the model's expressive capability and training efficiency.
• Multiple Transformer block designs：DiT explores different Transformer block designs, including adaptive layer normalization (adaLN), cross-attention, and in-context conditioning, to handle conditional information.
• Efficient training process：DiT shows a high degree of stability during training. Even without using learning rate warmup and regularization techniques, it can train stably to high performance.
• Diversity and quality of generated images：DiT can generate images with high visual quality and diversity. By adjusting the guidance strength of class-conditional generation, a trade-off can be made between the clarity and diversity of generated images.
• High computational efficiency：When generating images, DiT can achieve relatively high computational efficiency while maintaining high image quality. This makes DiT attractive in environments with limited resources.
• Application potential：DiT has broad application potential in the field of image generation, including artistic creation, game development, virtual reality, data augmentation, etc., especially in scenarios that require the generation of high-quality images.