Scalable Diffusion Models with Transformers (DiT)

Replaces convolutional U-Net backbones in diffusion models with pure transformer architectures (DiT blocks), enabling linear scaling with model capacity and improved computational efficiency. Uses standard transformer layers with adaptive layer normalization (AdaLN) to inject diffusion timestep and class conditioning directly into attention mechanisms, eliminating separate conditioning pathways and reducing architectural complexity.

Injects diffusion timestep and class information directly into transformer blocks via learned affine transformations (scale and shift) applied to layer normalization outputs, eliminating the need for separate conditioning networks or concatenation-based feature fusion. Each transformer block learns independent AdaLN parameters conditioned on timestep embeddings and optional class embeddings, enabling efficient multi-modal conditioning without architectural branching.

Analyzes how generation quality (FID/IS) scales with model size (parameters), training compute, and data, demonstrating that transformer-based diffusion models follow predictable scaling laws similar to language models. Enables principled decisions about model size, training duration, and data requirements by fitting power-law relationships between compute and quality metrics.

Converts images into sequences of flattened patch embeddings by dividing images into non-overlapping patches (e.g., 16x16 pixels), projecting each patch to a fixed embedding dimension via a linear layer, and flattening the spatial grid into a sequence. This enables transformer processing of images by converting 2D spatial data into 1D sequences compatible with standard attention mechanisms, with patch size as a tunable hyperparameter controlling sequence length and receptive field.

Encodes diffusion timestep indices (0 to T-1) into continuous embeddings using sinusoidal positional encoding (similar to transformer position embeddings) or learned embeddings, then passes these embeddings through an MLP to produce conditioning vectors injected into each transformer block. Supports standard noise schedules (linear, cosine, quadratic) that define the variance schedule σ(t) used during training and inference, enabling flexible control over the diffusion process dynamics.

Implements distributed training across multiple GPUs using PyTorch DDP or DeepSpeed, with gradient checkpointing to reduce memory usage by recomputing activations during backpropagation rather than storing them. Enables training of large DiT models (1B+ parameters) by distributing batch processing across GPUs and using activation checkpointing to trade compute for memory, critical for fitting models on 40GB+ VRAM devices.

Supports class-conditional generation by learning a class embedding table (num_classes × embedding_dim) that maps discrete class labels to continuous embeddings, which are then injected into transformer blocks via AdaLN. Enables controlled generation of specific object classes or categories by conditioning the diffusion process on class embeddings, with optional dropout of class embeddings during training for unconditional generation.

Implements classifier-free guidance at inference time by computing predictions for both conditioned and unconditional diffusion paths, then blending them with a guidance scale parameter λ: x̂ = x̂_uncond + λ(x̂_cond - x̂_uncond). This enables post-hoc control over generation quality and diversity without retraining, trading inference speed (2x forward passes) for improved sample quality and stronger adherence to conditioning signals.

Implements DDIM (Denoising Diffusion Implicit Models) sampling to reduce inference steps from 1000 (DDPM) to 50-100 steps with minimal quality loss, using a deterministic sampling procedure that skips timesteps while maintaining the diffusion trajectory. Enables fast inference by trading off some quality for speed, with configurable step counts allowing users to balance latency against sample fidelity.

Uses relative position embeddings instead of absolute position embeddings in transformer blocks, enabling the model to generalize to image resolutions not seen during training. Relative embeddings encode the distance between patches rather than absolute positions, allowing the same model to generate images at 256x256, 512x512, or 1024x1024 without retraining or position embedding interpolation.

Computes Fréchet Inception Distance (FID) and Inception Score (IS) metrics to quantitatively evaluate image generation quality by comparing generated images to real images using features from a pre-trained Inception network. FID measures the distance between feature distributions of real and generated images; IS measures the quality and diversity of generated images independently. Enables systematic comparison of model variants and hyperparameter choices.