DiT: Diffusion Transformer.

Implements the DiT architecture from "Scalable Diffusion Models with Transformers" (Peebles & Xie, ICCV 2023). Replaces the traditional U-Net backbone in diffusion models with a Transformer, using Adaptive Layer Normalization (AdaLN-Zero) for timestep and class conditioning.

Key Innovation: AdaLN-Zero Conditioning

Instead of cross-attention for conditioning (expensive), DiT modulates LayerNorm parameters based on the conditioning signal:

# Standard LayerNorm:
y = gamma * normalize(x) + beta

# AdaLN-Zero:
gamma, beta, alpha = MLP(condition)    # Learned modulation
y = gamma * normalize(x) + beta       # Modulated norm
y = alpha * y                         # Scale (initialized to zero)

Initializing alpha to zero means each DiT block starts as an identity function, enabling stable deep training.

Architecture

Input [batch, input_dim]
      |
      v
+--------------------------+
| Patchify + Position Embed|
+--------------------------+
      |
      v
+--------------------------+
| DiT Block x depth        |
|  AdaLN-Zero(cond)        |
|  Self-Attention          |
|  Residual                |
|  AdaLN-Zero(cond)        |
|  MLP                     |
|  Residual                |
+--------------------------+
      |
      v
| Final AdaLN + Linear    |
      |
      v
Output [batch, input_dim]  (predicted noise or v-prediction)

Conditioning

Timestep t -----> Sinusoidal Embed --> MLP --+
                                             |--> condition vector
Class label c --> Embedding ----------> MLP --+

Usage

model = DiT.build(
  input_dim: 64,
  hidden_size: 256,
  depth: 6,
  num_heads: 4
)

Reference

• Paper: "Scalable Diffusion Models with Transformers"
• arXiv: https://arxiv.org/abs/2212.09748