DiT: Diffusion Transformers
Sources:
- DiT 2022 paper
Also, it’s recommended to read my DDPM and Improved DDPM post to fully understand the context of diffusion models.
DiT: Diffusion Transformers
TL;DR: DiT = Diffusion framework + transformer backbone
1. Overview
Diffusion models generate samples by learning a reverse denoising process.
In DDPM:
\[ x_T \rightarrow x_{T-1} \rightarrow \cdots \rightarrow x_0 \]
DDPM learns a denoising function:
\[ \epsilon_\theta(x_t,t) \]
Early diffusion models mainly use U-Net architectures as denoising backbones.
Improved DDPM improves the reverse process by learning the variance:
\[ \Sigma_t \rightarrow \Sigma_\theta(x_t,t) \]
DiT asks a different question:
Can Transformer architectures replace U-Net as diffusion denoisers?
The key idea is:
\[ \epsilon_\theta: \text{U-Net} \rightarrow \text{Transformer} \]
DiT keeps the diffusion formulation unchanged and focuses on a scalable Transformer-based denoiser.
2. Notation
| Symbol | Meaning |
|---|---|
| \(x_0\) | Clean image sample |
| \(z_0\) | Latent representation from VAE encoder |
| \(z_t\) | Noisy latent representation |
| \(t\) | Diffusion timestep |
| \(T\) | Total diffusion timesteps |
| \(\epsilon\) | Gaussian noise sampled from \(\mathcal{N}(0,I)\) |
| \(\epsilon_\theta(z_t,t)\) | Predicted noise |
| \(y\) | Conditioning information |
| \(p\) | Patch size |
| \(N\) | Number of visual tokens |
| \(d\) | Transformer hidden dimension |
3. Relationship with DDPM and Improved DDPM
DiT does not introduce a new diffusion process.
It keeps:
- forward diffusion process;
- reverse Gaussian formulation;
- noise prediction objective;
- learned variance formulation from Improved DDPM.
The main difference is the denoising backbone.
| DDPM | Improved DDPM | DiT | |
|---|---|---|---|
| Forward diffusion | ✓ | ✓ | ✓ |
| Noise prediction | ✓ | ✓ | ✓ |
| Learned variance | ✗ | ✓ | ✓ |
| Backbone | U-Net | U-Net | Transformer |
4. Latent Diffusion Background
DiT follows the latent diffusion framework.
A VAE encoder maps:
\[ x_0 \rightarrow z_0 \]
and diffusion is performed in latent space:
\[ z_t= \sqrt{\bar{\alpha}_t}z_0+ \sqrt{1-\bar{\alpha}_t}\epsilon \]
The latent representation reduces computational cost for Transformer attention.
5. DiT Architecture
The noisy latent \(z_t\) is converted into Transformer tokens.
5.1 Patchification
For latent feature map:
\[ H\times W\times C \]
the number of patches is:
\[ N=(H/p)^2 \]
Each patch is projected into a token embedding:
\[ \mathbb{R}^{p^2C}\rightarrow\mathbb{R}^{d} \]
5.2 Transformer Backbone
DiT uses Transformer blocks similar to ViT:
LayerNorm |
Because diffusion requires timestep conditioning, DiT introduces adaptive LayerNorm.
5.3 AdaLN-Zero
Adaptive LayerNorm:
\[ AdaLN(x,c)=\gamma(c)LN(x)+\beta(c) \]
where:
\[ c=Embedding(t)+Embedding(y) \]
AdaLN-Zero additionally introduces:
\[ x'=x+\alpha F(x) \]
with zero initialization of the modulation parameters:
\[ \alpha\approx0 \]
which stabilizes large Transformer diffusion models.
6. Output Parameterization
DiT predicts the same quantities required by diffusion:
\[ \epsilon_\theta(z_t,t) \]
or, following Improved DDPM:
\[ [\epsilon_\theta,v_\theta] \]
The diffusion sampling procedure remains unchanged:
\[ \epsilon_\theta \rightarrow \mu_\theta \rightarrow z_{t-1} \]
7. Scaling Experiments
DiT studies whether diffusion models benefit from Transformer scaling.
The paper scales:
- Model size.
- Patch size.
- Computation.
Larger models achieve better generation quality.
Smaller patches create more tokens:
\[ p\downarrow,\quad N\uparrow \]
and improve spatial modeling with higher computation cost.
The main observation is:
\[ \text{GFlops}\uparrow \Rightarrow FID\downarrow \]
showing strong scaling behavior.