Diffusion Transformers (DiT): Architecture Overview

Building upon the adaptations of transformers for image data, such as the Vision Transformer (ViT), the Diffusion Transformer (DiT) architecture represents a significant shift in designing the backbone network for diffusion models. Proposed by Peebles and Xie (2022), DiT replaces the commonly used convolutional U-Net with a pure transformer architecture, demonstrating strong performance and scalability, particularly for image generation tasks.

The core idea is to treat the diffusion process, which operates on noisy images $x_t$ at timestep $t$, as a sequence modeling problem suitable for transformers. Instead of relying on convolutional layers' inductive biases (like locality and translation equivariance), DiT uses the transformer's ability to model long-range dependencies across the image.

Input Processing: Preparing Images for the Transformer

Like ViT, DiT cannot directly process raw pixel grids. The input noisy image $x_t \in \mathbb{R}^{H \times W \times C}$ must first be converted into a sequence of tokens:

1. Patchification: The input image $x_t$ is divided into a grid of non-overlapping patches. For an image of height $H$ and width $W$, using a patch size $P \times P$ results in $N = (H \times W) / P^2$ patches.
2. Linear Embedding: Each patch is flattened into a vector and then linearly projected into a token embedding of dimension $D$. This creates an initial sequence of $N$ tokens: $z_0 \in \mathbb{R}^{N \times D}$.
3. Positional Embeddings: Since transformers are permutation-invariant, positional information must be explicitly added. Standard learned 1D or 2D positional embeddings $E_{pos} \in \mathbb{R}^{N \times D}$ are added to the patch embeddings: $z_0' = z_0 + E_{pos}$. This sequence $z_0'$ forms the input to the transformer blocks.

Core Transformer Blocks

The main body of the DiT consists of a series of $L$ transformer blocks. Each block processes the sequence of tokens, refining the representation. A standard DiT block typically includes:

1. Layer Normalization (LN): Applied before the attention or MLP layers to stabilize activations.
2. Multi-Head Self-Attention (MHSA): Allows each token to attend to all other tokens in the sequence, capturing global relationships between image patches.
3. Layer Normalization (LN): Applied again before the feed-forward network.
4. Feed-Forward Network (FFN): Usually a simple Multi-Layer Perceptron (MLP) with two linear layers and a non-linear activation function (e.g., GeLU), applied independently to each token.

Residual connections are used around both the MHSA and FFN sub-layers, ensuring smooth gradient flow and enabling the training of deep transformers.

$$z_l'' = \text{MHSA}(\text{LN}(z_{l-1}')) + z_{l-1}' \\ z_l' = \text{FFN}(\text{LN}(z_l'')) + z_l''$$

Where $z_{l-1}'$ is the input to block $l$, and $z_l'$ is the output.

Conditioning: Incorporating Timestep and Context

A diffusion model must be conditioned on the current timestep $t$ and potentially other contextual information $c$ (like class labels, text embeddings, etc.). DiT needs mechanisms to inject this conditioning information into the transformer blocks.

1. Time and Condition Embeddings: The timestep $t$ is first converted into a vector embedding $e_t$, often using sinusoidal embeddings followed by an MLP. Similarly, context $c$ is mapped to an embedding $e_c$.
2. Adaptive Layer Normalization (adaLN / adaLN-Zero): This is a prevalent technique in DiTs. Instead of standard Layer Normalization, adaptive normalization layers dynamically compute scale ($\gamma$) and shift ($\beta$) parameters based on the embeddings $e_t$ and $e_c$. These parameters modulate the normalized activations within each transformer block, typically before the MHSA and FFN layers. $$\text{adaLN}(h, e_t, e_c) = (\gamma) \cdot \text{LayerNorm}(h) + (\beta)$$ where $\gamma$ and $\beta$ are produced by projecting $e_t$ and $e_c$ (often concatenated or summed) through linear layers. The "adaLN-Zero" variant initializes the projection layers generating $\gamma$ and $\beta$ such that the initial output is close to an identity mapping, which can aid training stability by ensuring conditioning blocks initially behave like residual connections.
3. Other Methods: While adaLN is common, alternatives exist:
   • Adding $e_t$ and $e_c$ directly to the token sequence (e.g., concatenating them as extra tokens or adding them to each token embedding).
   • Using cross-attention mechanisms if the conditioning $c$ is a sequence itself (like text embeddings).

Output Processing: Predicting the Diffusion Target

After processing through $L$ transformer blocks, the output token sequence $z_L'$ needs to be converted back into the format required by the diffusion objective (typically predicting the noise $\epsilon$ added at step $t$, or predicting the original image $x_0$).

1. Final Modulation: The conditioning embeddings $(e_t, e_c)$ often modulate the output one last time using a final adaLN layer or similar mechanism applied to $z_L'$.
2. Linear Projection: A final linear layer projects each output token embedding back to the dimension of a flattened patch (e.g., $P \times P \times C$).
3. Unpatchification / Reshaping: The output tokens are rearranged ("unpatched") to reconstruct the final output tensor, which has the same dimensions as the input image, representing the predicted noise $\epsilon_\theta(x_t, t, c)$ or $x_0$.

Overall Architecture Diagram

The following diagram illustrates the high-level flow of information within a Diffusion Transformer:

Overall data flow in a Diffusion Transformer (DiT). Input image $x_t$ is patched and embedded. Transformer blocks process these tokens, modulated by time $t$ and context $c$ embeddings via adaLN. The final tokens are projected and reshaped to produce the diffusion target (e.g., noise $\epsilon_\theta$).

DiTs offer a powerful alternative to U-Nets. By operating on sequences of patches, they can effectively model global image structure. Their scalability, demonstrated by performance improvements with larger models, aligns well with trends observed in large language models, making them a significant architecture in the ongoing development of diffusion-based generative modeling. The next sections will compare DiTs with U-Nets and discuss practical implementation details.