Diffusion Transformer

A Diffusion Transformer (DiT) is a diffusion model whose denoising network is a plain transformer over a sequence of patches, in the manner of the Vision Transformer, instead of the convolutional U-Net that earlier diffusion models used. Peebles and Xie (2023) introduced the name and the design: a class-conditional latent diffusion model whose transformer reads $p \times p$ patches of Stable Diffusion’s autoencoder latents and conditions every block on the noise step and the class through adaptive layer normalization initialized to the identity (adaLN-Zero). Their main finding was that sample quality tracked the network’s forward-pass compute, measured in Gflops, whether the compute came from a deeper and wider transformer or from more tokens. The same paper reported a Fréchet Inception Distance (FID) of 2.27 on class-conditional ImageNet $256 \times 256$. U-ViT (Bao et al., 2023), published the same year, added U-Net-style long skips between shallow and deep transformer blocks; Stable Diffusion 3 (Esser et al., 2024) built its multimodal MM-DiT on DiT’s modulation, and video generators apply the same idea to spacetime patches. “Diffusion transformer” is also used loosely for any transformer denoiser; this article uses DiT for Peebles and Xie’s model and the longer name for the family.

Unless stated otherwise, DiT details below come from the ICCV 2023 open-access version and its supplementary material, and code details from the official repository at commit ed81ce2 (as of November 2026). The notation follows the latent diffusion model article: $z_0$ is a latent, $z_t = \sqrt{\bar\alpha_t}\, z_0 + \sqrt{1 - \bar\alpha_t}\, \epsilon$, and the network predicts the noise, $\epsilon_\theta(z_t, t, y)$, for a class label $y$.

Core Idea

The diffusion U-Net came from PixelCNN++ by way of Ho et al.’s DDPM; Dhariwal and Nichol had ablated choices within it, but its high-level design had stayed largely intact (§1). Peebles and Xie asked whether its inductive bias was needed at all, and argued that a standard architecture would let diffusion models inherit the scaling behaviour and training recipes of transformers in other domains (§1). They kept everything else of a latent diffusion model fixed (the pretrained autoencoder, the noise schedule, training hyperparameters taken almost entirely from ADM) and replaced only the denoiser.

Two choices shape the result. First, the network works on a latent: for a $256 \times 256 \times 3$ image the Stable Diffusion autoencoder gives $z \in \mathbb{R}^{32 \times 32 \times 4}$, small enough for a transformer to process with $2 \times 2$ patches, 256 tokens (§3.2, §4). Second, complexity is counted in forward-pass Gflops rather than parameters, because parameter counts ignore resolution and token count, which the authors argue matter for image models (§2). With the depth, width and patch size of the transformer as dials, the paper then shows that parameter counts do not uniquely determine FID, while Gflops track it closely.

Building Blocks

Patchify. The first layer cuts the $I \times I \times C$ latent into non-overlapping $p \times p$ patches and maps each linearly to a token of width $d$, giving $T = (I/p)^2$ tokens; fixed two-dimensional sine–cosine position embeddings are added, as in the ViT literature (§3.2, Fig. 4). The design space uses $p = 2, 4, 8$. Halving $p$ quadruples $T$, which the paper says at least quadruples total transformer Gflops, while hardly changing the parameter count: only the patch embedding and the output layer depend on $p$.

Conditioning inputs. The step $t$ is embedded by a 256-dimensional sinusoidal frequency embedding followed by a two-layer MLP with SiLU activations, and the class by a learned table (supplement §A). For classifier-free guidance the label is replaced during training by a learned “null” embedding, with probability 0.1 in the official code (class_dropout_prob=0.1).

Four block designs. The paper compares four ways for a transformer block to take in $t$ and $y$ (§3.2, Fig. 3):

• In-context conditioning. The embeddings of $t$ and $y$ are appended to the sequence as two extra tokens and treated like image tokens, like ViT’s class token; standard ViT blocks are used unchanged, and the two tokens are removed after the last block. The cost in Gflops is negligible.
• Cross-attention. The two embeddings form a separate length-two sequence; each block gains a multi-head cross-attention layer after its self-attention, as in the original transformer and in LDM’s class-conditional models. This adds the most compute, about 15%.
• Adaptive layer norm (adaLN). Each layer normalization loses its learned scale and shift; instead, dimension-wise $\gamma$ and $\beta$ are regressed from the sum of the embeddings of $t$ and $y$. This follows adaptive normalization in GANs and in the U-Nets of Dhariwal and Nichol (see normalization layers). It adds the least compute, and it is the only design that applies the same function to every token.
• adaLN-Zero. adaLN plus a dimension-wise scale $\alpha$, also regressed from the conditioning, applied to each residual branch just before it is added back. The regression MLP is initialized to output the zero vector for every $\alpha$, so every block starts as the identity. The authors cite precedents: zero-initializing the last batch-norm scale of each residual block, which Goyal et al. found accelerated large-scale supervised training, and the zero-initialized final convolution of diffusion U-Net blocks.

Output layer. After the last block, a final layer normalization (adaptive in the adaLN designs) and a linear layer decode each token into a $p \times p \times 2C$ tensor, which is rearranged to the input layout: $C$ channels of predicted noise and $C$ of a diagonal covariance (§3.2). The final linear layer is initialized to zero in every model (§4); in the official code the modulation layers of the blocks and of the final layer are zero-initialized too, while the other linear layers use Xavier-uniform initialization with zero biases (initialize_weights in models.py).

Model sizes. The paper follows ViT’s Small, Base and Large configurations and adds XLarge (Table 1); a name such as DiT-XL/2 gives the configuration and the patch size.

Model	Blocks $N$	Width $d$	Heads	Gflops ($I = 32$, $p = 4$)
DiT-S	12	384	6	1.4
DiT-B	12	768	12	5.6
DiT-L	24	1,024	16	19.7
DiT-XL	28	1,152	16	29.1

Across the twelve configuration and patch-size combinations, Gflops run from 0.36 (DiT-S/8) to 118.6 (DiT-XL/2), and parameters from 33 million to 676 million, without the autoencoder’s 84 million (supplement, Table 4).

How Data Flows

Training. Encode the image once with the frozen autoencoder, $z_0 = E(x)$; draw $t$ and $\epsilon$; form $z_t$; patchify and embed it; pass the tokens and the conditioning vector through $N$ blocks; decode and unpatchify into a noise and a covariance prediction; and take a gradient step on the loss below. DiT keeps ADM’s diffusion settings: $t_{\max} = 1000$ steps, a linear variance schedule from $10^{-4}$ to $2 \times 10^{-2}$, ADM’s parameterization of the covariance and its embedding of steps and labels (§4).

Sampling. Start from $z_{t_{\max}} \sim \mathcal{N}(0, I)$ at latent size, run the reverse process with the network (the paper’s results use 250 DDPM steps), optionally with classifier-free guidance, and decode once with the autoencoder’s decoder. Diffusion sampling covers faster samplers, which apply to a DiT as to any noise predictor.

Mathematical Formulation

Patchify. With $P_i(z) \in \mathbb{R}^{p^2 C}$ the $i$-th patch flattened, in raster order,

$$h^{(0)}_i = W_E\, P_i(z_t) + b_E + \mathrm{pos}_i, \qquad i = 1, \dots, T, \qquad T = (I/p)^2,$$

with $W_E \in \mathbb{R}^{d \times p^2 C}$; as in ViT, this is a convolution with kernel and stride $p$.

An adaLN-Zero block. With $c = e_t + e_y$ the summed embeddings, one linear map regresses six vectors of width $d$:

$$\big(\beta_1, \gamma_1, \alpha_1, \beta_2, \gamma_2, \alpha_2\big) = W_c\, \mathrm{SiLU}(c) + b_c, \qquad W_c = 0,\ b_c = 0 \text{ at initialization},$$ $$h' = h + \alpha_1 \odot \mathrm{MSA}\big(\mathrm{LN}(h) \odot (1 + \gamma_1) + \beta_1\big), \qquad h'' = h' + \alpha_2 \odot \mathrm{MLP}\big(\mathrm{LN}(h') \odot (1 + \gamma_2) + \beta_2\big),$$

where LN has no learned affine parameters, the vectors are broadcast over the $T$ tokens, and MSA is multi-head self-attention. The paper describes the regressed scale as $\gamma$; the official code applies it as $1 + \gamma$ (modulate in models.py), so that a zero output leaves the normalized features unchanged. adaLN is the same block without $\alpha_1, \alpha_2$, so that the map outputs four vectors, and with ordinary initialization; the supplement gives the regression output widths as $4d$ and $6d$ (§A). In our reading, the zero initialization has two consequences: at the first step every block is the identity and the output layer emits zeros, so the prediction is $\epsilon_\theta = 0$ and the loss equals the noise variance; and no gradient reaches the attention and MLP weights until the output layer and then $\alpha$ have moved away from zero.

Objective. The noise prediction is trained with the simple loss of Ho et al. and the covariance $\Sigma_\theta$ with the full variational bound, following Nichol and Dhariwal (§3.1):

$$L_{\text{simple}}(\theta) = \mathbb{E}_{z_0, y, t, \epsilon}\big\lVert \epsilon_\theta(z_t, t, y) - \epsilon \big\rVert_2^2 .$$

Compute. A pre-norm transformer block on $T$ tokens of width $d$ with MLP ratio 4 spends $12 T d^2$ multiply-adds in its linear layers and $2 T^2 d$ in the attention products ($Q K^\top$ and the weighted sum of values); the Vision Transformer article derives these counts. Adding what DiT adds (our accounting, checked in the example below),

$$\mathrm{MACs} \approx N\big(12 T d^2 + 2 T^2 d + k d^2\big) + T d\, p^2 (C + C_{\text{out}}),$$

with $k = 6$ for adaLN-Zero and 4 for adaLN, $C_{\text{out}} = 2C$ and small terms for the embeddings and the final modulation omitted. For DiT-XL/2 ($N = 28$, $d = 1152$, $T = 256$) the full count is $118.62 \times 10^9$ multiply-adds, against 118.64 “Gflops” in the paper’s Table 4, so in our reading the paper’s Gflops count multiply-adds, not separate multiplications and additions. Halving $p$ multiplies the first term by 4 and the second by 16, hence the paper’s “at least quadruple”; the per-sample terms $k d^2$ and the embeddings do not grow with $T$, so on very few tokens the total grows less (in the example below, from 1,052,672 to 3,534,848 multiply-adds, 3.4-fold, between $p = 8$ and $p = 4$). The paper estimates total training compute as Gflops × batch size × steps × 3, the factor 3 approximating a backward pass twice as costly as the forward pass (§5).

Why It Works

Gflops, not parameters, predicted FID. The authors trained all twelve models for 400,000 steps with identical hyperparameters (§5). Deeper and wider models and smaller patches each lowered FID at every stage of training (Fig. 6). Reducing the patch size leaves the parameters unchanged (slightly lower, in fact) and raises only the Gflops, yet improved FID throughout, so parameters alone do not determine quality. Plotting FID-50K at 400,000 steps against Gflops gave a strong negative correlation (Fig. 8), and different configurations with similar Gflops scored similarly: DiT-S/2 (6.06 Gflops) and DiT-B/4 (5.56) reached 68.40 and 68.38 (supplement, Table 4). The same trend held for sFID, Inception Score, precision and recall (supplement, Fig. 12), and larger models reached lower training loss sooner (supplement, Fig. 13).

FID-50K at 400K steps, no guidance	$p = 8$	$p = 4$	$p = 2$
DiT-S	153.60	100.41	68.40
DiT-B	122.74	68.38	43.47
DiT-L	118.87	45.64	23.33
DiT-XL	106.41	43.01	19.47

These scores use the ft-MSE decoder of the autoencoder (supplement, Table 4).

Larger models used training compute better. Plotted against total training compute, small models trained longer eventually fell behind larger models trained for fewer steps, and DiT-XL/4 was overtaken by DiT-XL/2 after about $10^{10}$ Gflops (Fig. 9).

Sampling compute did not substitute for model compute. With 16 to 1,000 sampling steps, small models did not catch up with large ones even when they spent more compute per image: DiT-L/2 with 1,000 steps used 80.7 Tflops per image and scored FID-10K 25.9, against 23.7 for DiT-XL/2 with 128 steps at 15.2 Tflops (§5.2, Fig. 10).

The block design mattered. Four DiT-XL/2 models, one per block design, were trained for 400,000 steps (§5, Fig. 5; supplement, Table 4):

Block	Gflops	Parameters	FID-50K, no guidance
In-context	119.37	449M	35.24
Cross-attention	137.62	598M	26.14
adaLN	118.56	600M	25.21
adaLN-Zero	118.64	675M	19.47

adaLN-Zero reached lower FID than cross-attention and in-context conditioning at every stage of training while costing the least compute, and its FID at 400,000 steps was nearly half that of in-context conditioning; it also clearly beat adaLN, so the initialization mattered as well as the mechanism (§5). All later models use adaLN-Zero. The parameter counts differ because the regression layers are large ($6d^2$ weights per block), but the Gflops barely do, since these layers run once per sample, not once per token. In our reading the comparison therefore favours adaLN-Zero on compute, not on parameters.

The paper offers no theory for the scaling; its argument is empirical. That the transformer, unlike the U-Net, has no built-in multiscale structure and still scales, the authors take as evidence that the U-Net’s inductive bias is “not crucial” (§1).

Training Considerations

• Recipe. AdamW with a constant learning rate of $10^{-4}$, no weight decay, batch size 256, horizontal flips as the only augmentation, and an exponential moving average of the weights with decay 0.9999, used for every reported result (§4). The authors used no learning-rate warm-up or regularization, observed no loss spikes, and did not tune the hyperparameters, which are almost all ADM’s.
• Scale and hardware. The models were implemented in JAX and trained on TPU v3 pods; DiT-XL/2 ran at about 5.7 iterations per second on a TPU v3-256 pod (§4). The official PyTorch reimplementation reproduces the JAX results: at 400,000 steps, DiT-XL/2 scored FID-50K 18.1 against 19.5 for the JAX model, without guidance (README, as of November 2026).
• Autoencoder. The latents come from the original LDM “f8” autoencoder; ft-MSE and ft-EMA are Stable Diffusion’s two decoders fine-tuned with the encoder unchanged, so they can be swapped without retraining the DiT. Scaling results use ft-MSE and the final results ft-EMA; with guidance scale 1.5, FID was 2.46 with the original decoder, 2.30 with ft-MSE and 2.27 with ft-EMA (supplement, §D, Table 5).
• Guidance on three channels. The guided results apply classifier-free guidance only to the first three of the four latent channels; the authors found this similar to four-channel guidance with a scale of about $1 + \tfrac34 x$ in place of $1 + x$: four-channel guidance at 1.375 gave FID-50K 2.20, against 2.27 for three channels at 1.5 (supplement, §A). The official code keeps the three-channel version for exact reproducibility (forward_with_cfg).
• Training length. For both final models the authors “never observed FID saturate” and trained as long as they could (supplement, Table 4 caption).

Results

DiT’s headline benchmark is class-conditional ImageNet, scored with FID-50K computed by ADM’s TensorFlow evaluation suite from exported samples, 250 DDPM sampling steps, and sFID, Inception Score and precision and recall as secondary metrics (§4). The image generation article explains why FID values are comparable only under one protocol.

ImageNet $256 \times 256$ (Table 2). DiT-XL/2, trained for 7 million steps, scored:

Model	FID	sFID	IS	Precision	Recall
LDM-4-G (cfg = 1.50)	3.60	–	247.67	0.87	0.48
StyleGAN-XL	2.30	4.02	265.12	0.78	0.53
DiT-XL/2	9.62	6.85	121.50	0.67	0.67
DiT-XL/2-G (cfg = 1.25)	3.22	5.28	201.77	0.76	0.62
DiT-XL/2-G (cfg = 1.50)	2.27	4.60	278.24	0.83	0.57

With guidance at scale 1.5 it improved on the previous best diffusion FID, LDM’s 3.60, and on the GAN StyleGAN-XL’s 2.30. At every guidance scale tested it had higher recall than LDM-4 and LDM-8, and after 2.35 million steps, about ADM’s training length, it already scored 2.55 (§5.1). Its forward pass costs 118.6 Gflops, against 103.6 for LDM-4, 1,120 for pixel-space ADM and 742 for ADM-U (§5.1; supplement, Table 6).

ImageNet $512 \times 512$ (Table 3). A separate DiT-XL/2 trained for 3 million steps on $64 \times 64 \times 4$ latents (1,024 tokens, 524.6 Gflops) scored FID 3.04 with guidance at 1.5, below the 3.85 of ADM-G with ADM-U (ADM-U alone uses 2,813 Gflops); StyleGAN-XL’s 2.41 in the same table remained lower. Following ADM, precision and recall at this resolution use 1,000 real samples.

Variants

U-ViT: all inputs as tokens, with long skips

U-ViT (Bao et al., 2023), published at CVPR 2023, is a ViT backbone for diffusion built on two different choices (§3, Fig. 1). It treats everything as tokens: the time, the condition and the noisy patches enter one sequence, the in-context design that DiT found weakest. And it adds long skip connections between shallow and deep blocks, $(L - 1)/2$ of them for $L$ blocks, as a U-Net joins its encoder and decoder but with no down- or up-sampling. The authors argue that noise prediction is a pixel-level task sensitive to low-level features, which the skips pass forward. An optional $3 \times 3$ convolution after the output projection is meant to suppress artefacts.

Their ablations on CIFAR-10 used FID on 10,000 samples (§3.1, Fig. 2). For combining the skip branch $h_s$ with the main branch $h_m$, a linear projection of their concatenation, $\mathrm{Linear}(\mathrm{Concat}(h_m, h_s))$, did best. Plain addition $h_m + h_s$ gave no benefit, which the authors explain by the residual connections inside each block: $h_m$ already contains $h_s$ linearly, so adding it again only changes a coefficient. Feeding time as a token beat an adaptive layer norm of the form $y_s \odot \mathrm{LN}(h) + y_b$; linear patch embedding beat a convolutional stem; a learned 1D position embedding beat 2D sinusoids, and without position embedding the model failed to generate meaningful images. Patch size 2 was needed for good results, and patch size 1 gained nothing (§3.2).

The two papers disagree on how to feed the time, but the set-ups differ: U-ViT’s comparison was on CIFAR-10 in pixel space, with long skips and its own adaptive norm, without DiT’s gates; DiT’s was on ImageNet latents with DiT-XL/2 and no long skips. In our reading neither result transfers to the other’s setting without a test.

U-ViT-H/2 (29 blocks, width 1,152, 501 million parameters, 133 Gflops) on Stable Diffusion’s $f = 8$ latents reached FID 2.29 on class-conditional ImageNet $256 \times 256$ with classifier-free guidance, sampled with 50 DPM-Solver steps after 500,000 training iterations at batch size 1,024 with the ft-EMA decoder (Table 1; appendix, Tables 1 and 3). U-ViT-H/4 scored 4.05 at $512 \times 512$. Removing the long skips hurt on ImageNet as on CIFAR-10 (Fig. 4). In text-to-image generation trained on COCO alone, with CLIP text tokens in the sequence, U-ViT-S/2 with 17 blocks reached FID 5.48 on 30,000 validation prompts, against 7.32 for the authors’ own latent U-Net with cross-attention of comparable size trained in the same way (Table 4). DiT’s 2.27 and U-ViT’s 2.29 are not a ranking: the samplers (250 DDPM steps against 50 DPM-Solver steps), guidance settings and training budgets differ, and neither paper compares with the other.

The U-ViT paper disagrees with itself in two training details. The main text gives a weight decay of 0.3 for all datasets and 300,000 iterations on ImageNet $256 \times 256$ (§5.1), while the appendix gives 0.03 (Table 1 and text) and 500,000 iterations (Table 1); the 2.29 row of appendix Table 3 is labelled “trained 500K”.

MM-DiT: joint attention in Stable Diffusion 3

Stable Diffusion 3 (Esser et al., 2024) generalizes DiT from a class label to text (§4, Fig. 2). It keeps DiT’s modulation, fed by the step embedding and a pooled text vector, but adds the token sequence of three frozen text encoders, because a pooled vector keeps only coarse information. Latent patches of size $2 \times 2$ with position embeddings and the text tokens are embedded to one width and concatenated; each modality has its own weights for the projections and MLPs, and the two sequences meet only in a joint attention over both. The authors describe this as two transformers joined for the attention operation. The text conditioning article gives the joint-attention equations, the encoders and the comparison on CC12M in which MM-DiT beat a cross-attention DiT, a “UViT” and a DiT with text and image tokens concatenated under shared weights (§5.2.3, Fig. 4). The “UViT” there is the hybrid of Hoogeboom et al. (2023), cited by Esser et al., not Bao et al.’s U-ViT; the cross-attention DiT is attributed to Chen et al. (2023).

Model size is set by one number, the depth $d$: width $64d$, MLP width $4 \cdot 64d$ and $d$ attention heads (§4). Trained on $256 \times 256$ images for 500,000 steps at batch size 4,096, validation loss fell smoothly with depth and training steps and correlated with GenEval, T2I-CompBench and human preference, with no sign of saturation (§5.3, Fig. 5); the largest model, depth 38, has 8 billion parameters and needed a learning-rate adjustment at 300,000 steps to avoid divergence. Larger models also lost less when sampled with fewer steps (Table 6). For high-resolution fine-tuning in mixed precision, attention logits grew uncontrollably in the last blocks; normalizing queries and keys with RMSNorm in both streams fixed this (Appendix C.2, Fig. 9). Position embeddings for other sizes and aspect ratios come from extended and interpolated position grids, center-cropped before embedding (Appendix C.2). On GenEval, the depth-38 model fine-tuned with DPO at $1024^2$ scored 0.74 overall, against 0.67 for DALL·E 3 (Table 5). The Stable Diffusion article covers the released SD 3 and 3.5 models, and the latent diffusion model article the rectified-flow objective.

Video diffusion transformers, in outline

A video latent can be patchified in time as well as space. Esser et al. ran a preliminary video study from their image weights with $2\times$ temporal patching, adding full attention over all spatiotemporal tokens after each spatial attention, on 16-frame clips at $256 \times 256$; validation loss again fell with model size (§5.3). The video generation article describes the large systems built this way, OpenAI’s Sora (a diffusion transformer on spacetime latent patches, without model details) and Meta’s Movie Gen (a 30-billion-parameter transformer whose 16-second clips become about 73,000 tokens), and the factorized space–time attention that keeps earlier video diffusion models tractable. Token counts grow with frames × height × width, so the patch-size and Gflops trade-off of DiT applies with one more axis.

Use in Computer Vision

Class-conditional ImageNet generation was DiT’s own task; U-ViT added unconditional generation on CIFAR-10 and CelebA $64 \times 64$ and text-to-image generation on COCO. Stable Diffusion 3 and 3.5 use MM-DiT transformers (see the Stable Diffusion article), and the large video systems in the video generation article use transformers over spacetime patches. Peebles and Xie suggested DiT as a drop-in backbone for text-to-image models such as DALL·E 2 and Stable Diffusion (§6), which is the role MM-DiT took.

Limitations

• Quadratic attention. The attention term grows with $T^2$, and $T$ grows with resolution and with frames; DiT relies on the autoencoder to keep $T$ small, so its quality and cost are bounded by the first stage, as for any latent diffusion model.
• Compute, not parameters. Smaller patches buy quality with compute; DiT-XL/2 at $512 \times 512$ costs 524.6 Gflops per forward pass, and the paper’s results use 250 sampling steps (§4, §5.1).
• Conditioning by a vector. adaLN applies the same scale and shift to every token, which suits a class label or a step but carries a caption only as a pooled vector; text models add a token sequence, by cross-attention or joint attention (see text conditioning).
• Stability at scale. DiT trained without spikes up to 675 million parameters, but Stable Diffusion 3’s 8-billion-parameter model needed a learning-rate adjustment and its high-resolution fine-tuning needed QK-normalization (Esser et al., §5.3, Appendix C.2).
• Empirical scaling claims. The DiT scaling study covers one dataset, one autoencoder and 400,000 steps; the correlation of FID with Gflops is measured, not derived.

Important Models

• DiT (Peebles & Xie, ICCV 2023): latent transformer with adaLN-Zero, class-conditional ImageNet.
• U-ViT (Bao et al., CVPR 2023): all inputs as tokens, long skip connections.
• MM-DiT in Stable Diffusion 3 (Esser et al., ICML 2024): separate weights per modality joined by joint attention, rectified flow, up to 8 billion parameters.
• Video diffusion transformers: Sora and Movie Gen among those described in the video generation article.

Example

Counting compute, and three conditioning blocks on synthetic latents. The data are $32 \times 32$ images of one shape each (disk, square, triangle or ring, at a random position and size, $+1$ on $-1$), whose class is the shape. A lossless space-to-depth step turns each $2 \times 2$ pixel block into 4 channels, so every image becomes a $16 \times 16 \times 4$ grid that stands in for an autoencoder latent; no autoencoder is trained. The PyTorch model follows the paper’s description, with the official code’s initialization of the embeddings: linear patchify, fixed 2D sine–cosine position embeddings, a 256-dimensional frequency embedding of the step through a two-layer SiLU MLP, a class table, $c = e_t + e_y$, and blocks of the chosen design, with a zero-initialized final linear layer in every variant. It predicts only the noise; DiT’s learned covariance is left out.

Part A counts multiply-adds with PyTorch’s FlopCounterMode (which reports two flops per multiply-add) on the toy model for patch sizes 1 to 8, checks the closed form of the Mathematical Formulation section and a parameter formula against it, and then applies both formulas to the paper’s configurations beside the supplement’s Table 4. Part B trains a 4-block, width-64 DiT with $p = 2$ (64 tokens) and adaLN-Zero, adaLN (ordinary Xavier initialization of the modulation, no gates) or in-context conditioning, each for the same 500 steps at batch size 128 with AdamW at a constant learning rate of $5 \times 10^{-4}$ and no weight averaging, for seeds 0, 1 and 2. Training samples $t$ uniformly over DiT’s 1,000-step linear schedule; the code indexes steps from 0. The validation loss uses 512 held-out images with fixed steps and noise, and the script first prints two reference losses on them: predicting zero, and the exact noise predictor for a Gaussian with the training latents’ mean and covariance. The script uses PyTorch 2.14.1 with NumPy 2.5.4 and four threads, and took 8 to 15 minutes per run in three runs on our CPU, which other jobs were sharing.

import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.flop_counter import FlopCounterMode

torch.set_num_threads(4)

# Data: 32x32 shape images in [-1, 1], class = shape kind. A lossless f = 2 "latent"
# (each 2x2 pixel block becomes 4 channels) stands in for an autoencoder latent: 16x16x4.
KINDS = ["disk", "square", "triangle", "ring"]
yy, xx = np.mgrid[:32, :32] + 0.5
def shape(kind, cy, cx, r):
    if kind == "disk":     return (yy - cy) ** 2 + (xx - cx) ** 2 < r ** 2
    if kind == "square":   return (abs(yy - cy) < r) & (abs(xx - cx) < r)
    if kind == "triangle": return (yy - cy < r) & (abs(xx - cx) < (yy - cy + r) / 2) & (yy - cy > -r)
    if kind == "ring":     return abs(np.hypot(yy - cy, xx - cx) - r) < 1.5

def make_latents(n, seed):
    rng = np.random.default_rng(seed)
    y = rng.integers(4, size=n)
    img = np.empty((n, 32, 32), np.float32)
    for i in range(n):
        cy, cx = rng.uniform(10, 22, size=2)
        img[i] = np.where(shape(KINDS[y[i]], cy, cx, rng.uniform(5, 9)), 1.0, -1.0)
    return img.reshape(n, 16, 2, 16, 2).transpose(0, 2, 4, 1, 3).reshape(n, 4, 16, 16), y

z_tr, y_tr = make_latents(4096, 0)
z_va, y_va = make_latents(512, 1)

# DDPM linear schedule as in DiT (T = 1000, beta from 1e-4 to 0.02); epsilon-prediction only.
TT = 1000
abar = np.cumprod(1 - np.linspace(1e-4, 0.02, TT))
t_va = np.arange(512) * TT // 512                      # validation: fixed, evenly spread steps ...
eps_va = np.random.default_rng(2).standard_normal(z_va.shape).astype(np.float32)   # ... and fixed noise
a_va, s_va = np.sqrt(abar[t_va])[:, None, None, None], np.sqrt(1 - abar[t_va])[:, None, None, None]
x_va = (a_va * z_va + s_va * eps_va).astype(np.float32)

# Reference: the exact noise predictor for a Gaussian with the training latents' mean and covariance.
D = 4 * 16 * 16
X = z_tr.reshape(-1, D).astype(np.float64)
mu = X.mean(0)
lam, U = np.linalg.eigh(np.cov(X, rowvar=False)); lam = np.clip(lam, 0, None)
a2 = abar[t_va][:, None]
proj = (x_va.reshape(512, D) - np.sqrt(a2) * mu) @ U
eps_gauss = (np.sqrt(1 - a2) * proj / (a2 * lam + 1 - a2)) @ U.T
print(f"validation loss of the zero predictor {np.mean(eps_va ** 2):.4f}, "
      f"of the Gaussian denoiser {np.mean((eps_gauss - eps_va.reshape(512, D)) ** 2):.4f}")

# ---------------- the model ----------------
def sincos_2d(d, g):           # fixed 2D sine-cosine table: half the channels for the row, half for the column
    def one(pos, dd):
        a = pos[:, None] / 10000 ** (np.arange(dd // 2) / (dd / 2))[None]
        return np.concatenate([np.sin(a), np.cos(a)], 1)
    r, c = np.meshgrid(np.arange(g), np.arange(g), indexing="ij")
    return torch.tensor(np.concatenate([one(r.ravel(), d // 2), one(c.ravel(), d // 2)], 1), dtype=torch.float32)

def t_freq(t, dim=256):        # 256-dimensional frequency embedding of the step
    f = torch.exp(-math.log(10000) * torch.arange(dim // 2) / (dim // 2))
    a = t[:, None].float() * f[None]
    return torch.cat([torch.cos(a), torch.sin(a)], 1)

class Attn(nn.Module):
    def __init__(s, d, h):
        super().__init__(); s.h = h; s.qkv = nn.Linear(d, 3 * d); s.proj = nn.Linear(d, d)
    def forward(s, x):
        B, N, d = x.shape
        q, k, v = s.qkv(x).reshape(B, N, 3, s.h, d // s.h).permute(2, 0, 3, 1, 4)
        w = (q @ k.transpose(-1, -2)) / math.sqrt(d // s.h)
        return s.proj((w.softmax(-1) @ v).transpose(1, 2).reshape(B, N, d))

def mlp(d):
    return nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(approximate="tanh"), nn.Linear(4 * d, d))

def mod(x, shift, scale):
    return x * (1 + scale[:, None]) + shift[:, None]

class Block(nn.Module):
    def __init__(s, d, h, kind):
        super().__init__(); s.kind = kind
        affine = kind == "in-context"                  # adaptive variants: no learned affine in the LN
        s.n1 = nn.LayerNorm(d, elementwise_affine=affine, eps=1e-6); s.attn = Attn(d, h)
        s.n2 = nn.LayerNorm(d, elementwise_affine=affine, eps=1e-6); s.mlp = mlp(d)
        if kind != "in-context":                       # regress shifts and scales (and gates) from c
            s.ada = nn.Sequential(nn.SiLU(), nn.Linear(d, (6 if kind == "adaLN-Zero" else 4) * d))
    def forward(s, x, c):
        if s.kind == "in-context":
            x = x + s.attn(s.n1(x))
            return x + s.mlp(s.n2(x))
        m = s.ada(c).chunk(6 if s.kind == "adaLN-Zero" else 4, dim=1)
        if s.kind == "adaLN-Zero":
            sh1, sc1, g1, sh2, sc2, g2 = m
            x = x + g1[:, None] * s.attn(mod(s.n1(x), sh1, sc1))
            return x + g2[:, None] * s.mlp(mod(s.n2(x), sh2, sc2))
        sh1, sc1, sh2, sc2 = m
        x = x + s.attn(mod(s.n1(x), sh1, sc1))
        return x + s.mlp(mod(s.n2(x), sh2, sc2))

class DiT(nn.Module):
    def __init__(s, kind="adaLN-Zero", p=2, d=64, depth=4, heads=4, C=4, Cout=4, size=16, classes=4):
        super().__init__(); s.kind, s.p, s.C, s.Cout, s.g = kind, p, C, Cout, size // p
        s.embed = nn.Linear(p * p * C, d)                                  # patchify: one linear map per patch
        s.register_buffer("pos", sincos_2d(d, s.g))
        s.temb = nn.Sequential(nn.Linear(256, d), nn.SiLU(), nn.Linear(d, d))
        s.yemb = nn.Embedding(classes, d)
        s.blocks = nn.ModuleList([Block(d, heads, kind) for _ in range(depth)])
        s.nf = nn.LayerNorm(d, elementwise_affine=(kind == "in-context"), eps=1e-6)
        if kind != "in-context":
            s.adaf = nn.Sequential(nn.SiLU(), nn.Linear(d, 2 * d))
        s.out = nn.Linear(d, p * p * Cout)
        for m in s.modules():                                              # ordinary initialization ...
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight); nn.init.zeros_(m.bias)
        nn.init.normal_(s.yemb.weight, std=0.02)
        nn.init.normal_(s.temb[0].weight, std=0.02); nn.init.normal_(s.temb[2].weight, std=0.02)
        nn.init.zeros_(s.out.weight)                                       # ... final linear layer zero in all variants
        if kind == "adaLN-Zero":                                           # ... and adaLN-Zero zeroes every modulation
            for b in s.blocks:
                nn.init.zeros_(b.ada[1].weight); nn.init.zeros_(b.ada[1].bias)
            nn.init.zeros_(s.adaf[1].weight); nn.init.zeros_(s.adaf[1].bias)

    def patchify(s, z):        # (B, C, H, W) -> (B, T, p*p*C), patches in raster order
        B, p = len(z), s.p
        return z.reshape(B, s.C, s.g, p, s.g, p).permute(0, 2, 4, 3, 5, 1).reshape(B, s.g * s.g, p * p * s.C)

    def unpatchify(s, x):      # (B, T, p*p*Cout) -> (B, Cout, H, W)
        B, p = len(x), s.p
        return x.reshape(B, s.g, s.g, p, p, s.Cout).permute(0, 5, 1, 3, 2, 4).reshape(B, s.Cout, s.g * p, s.g * p)

    def forward(s, z, t, y):
        x = s.embed(s.patchify(z)) + s.pos
        et, ey = s.temb(t_freq(t)), s.yemb(y)
        c = et + ey                                                        # one conditioning vector per sample
        if s.kind == "in-context":
            x = torch.cat([x, et[:, None], ey[:, None]], 1)                # two extra tokens
        for b in s.blocks:
            x = b(x, c)
        if s.kind == "in-context":
            x = s.nf(x[:, :-2])                                            # drop them after the last block
        else:
            sh, sc = s.adaf(c).chunk(2, dim=1)
            x = mod(s.nf(x), sh, sc)
        return s.unpatchify(s.out(x))

# ---------------- part A: patch size, tokens, compute ----------------
def macs(T, d, depth, p, C, Cout, kind="adaLN-Zero", M=2):
    """Multiply-adds of one forward pass (one sample), MLP ratio 4."""
    m = T * p * p * C * d + 256 * d + d * d                                # patch embedding, timestep MLP
    n = T + M if kind == "in-context" else T                               # tokens inside the blocks
    blk = 12 * n * d * d + 2 * n * n * d                                   # QKV, output, MLP; QK^T and AV
    if kind == "adaLN-Zero": blk += 6 * d * d
    if kind == "adaLN":      blk += 4 * d * d
    if kind == "cross-attention": blk += 2 * T * d * d + 2 * M * d * d + 2 * T * M * d
    m += depth * blk + T * d * p * p * Cout                                # blocks, linear decoder
    if kind.startswith("adaLN"): m += 2 * d * d                            # final adaptive LN
    return m

def params(d, depth, p, C, Cout, classes, kind="adaLN-Zero"):
    n = (p * p * C + 1) * d + (256 * d + d) + (d * d + d) + classes * d    # patch, timestep MLP, class table
    blk = (4 * d * d + 4 * d) + (8 * d * d + 5 * d)                        # attention, MLP
    blk += {"adaLN-Zero": 6 * d * d + 6 * d, "adaLN": 4 * d * d + 4 * d,
            "in-context": 4 * d, "cross-attention": 4 * d * d + 4 * d + 6 * d}[kind]
    n += depth * blk + (d + 1) * p * p * Cout
    n += 2 * d * d + 2 * d if kind.startswith("adaLN") else 2 * d          # final adaptive or affine LN
    return n

print("\ntoy DiT (d = 64, 4 blocks) on a 16x16x4 latent")
print(" p  tokens  params (model / formula)   MACs (counted / formula)")
for p in [1, 2, 4, 8]:
    net = DiT(p=p)
    with FlopCounterMode(display=False) as fc:
        net(torch.zeros(1, 4, 16, 16), torch.zeros(1, dtype=torch.long), torch.zeros(1, dtype=torch.long))
    T = (16 // p) ** 2
    print(f"{p:2d}{T:8d}{sum(q.numel() for q in net.parameters()):12,d} /{params(64, 4, p, 4, 4, 4):8,d}"
          f"{fc.get_total_flops() // 2:16,d} /{macs(T, 64, 4, p, 4, 4):11,d}")

print("\nDiT on 32x32x4 latents: formula against Table 4 (supplement); 8 output channels, 1,001 labels")
print("model      block            tokens  GMACs  Table 4 Gflops  params (M)  Table 4 (M)")
table4 = [("S", 8, 0.36, 33), ("S", 4, 1.41, 33), ("S", 2, 6.06, 33), ("B", 8, 1.42, 131), ("B", 4, 5.56, 130),
          ("B", 2, 23.01, 130), ("L", 8, 5.01, 459), ("L", 4, 19.70, 458), ("L", 2, 80.71, 458),
          ("XL", 8, 7.39, 676), ("XL", 4, 29.05, 675), ("XL", 2, 118.64, 675)]
variants = [("in-context", 119.37, 449), ("cross-attention", 137.62, 598), ("adaLN", 118.56, 600)]
cfg = {"S": (12, 384), "B": (12, 768), "L": (24, 1024), "XL": (28, 1152)}
rows = [(m, p, "adaLN-Zero", g, n) for m, p, g, n in table4] + [("XL", 2, k, g, n) for k, g, n in variants]
for m, p, kind, gf, pm in rows:
    depth, d = cfg[m]; T = (32 // p) ** 2
    print(f"{f'DiT-{m}/{p}':11s}{kind:16s}{T:5d}{macs(T, d, depth, p, 4, 8, kind) / 1e9:9.2f}{gf:12.2f}"
          f"{params(d, depth, p, 4, 8, 1001, kind) / 1e6:13.1f}{pm:11d}")

# ---------------- part B: three conditioning designs, fixed budget ----------------
abar_t = torch.tensor(abar, dtype=torch.float32)
Ztr, Ytr = torch.tensor(z_tr), torch.tensor(y_tr)
Xva, Tva, Eva, Yva = map(torch.tensor, (x_va, t_va, eps_va, y_va))

def train(kind, seed, steps=500, bs=128, lr=5e-4, every=100):
    torch.manual_seed(seed)
    net = DiT(kind)
    opt = torch.optim.AdamW(net.parameters(), lr=lr, weight_decay=0)
    g = torch.Generator().manual_seed(seed)
    hist, curve = [], []
    for step in range(1, steps + 1):
        i = torch.randint(len(Ztr), (bs,), generator=g)
        t = torch.randint(TT, (bs,), generator=g)
        e = torch.randn(bs, 4, 16, 16, generator=g)
        a, s = abar_t[t].sqrt()[:, None, None, None], (1 - abar_t[t]).sqrt()[:, None, None, None]
        loss = F.mse_loss(net(a * Ztr[i] + s * e, t, Ytr[i]), e)
        opt.zero_grad(); loss.backward(); opt.step()
        hist.append(loss.item())
        if step % every == 0:
            with torch.no_grad():
                curve.append(F.mse_loss(net(Xva, Tva, Yva), Eva).item())
    return np.mean(hist[-200:]), curve

print("\n500 steps, batch 128; validation loss every 100 steps; train = mean of the last 200 steps")
results = {}
for kind in ["adaLN-Zero", "adaLN", "in-context"]:
    for seed in [0, 1, 2]:
        tr, curve = train(kind, seed)
        results.setdefault(kind, []).append((tr, curve[-1]))
        print(f"{kind:11s} seed {seed}  val " + " ".join(f"{v:.4f}" for v in curve) + f"   train {tr:.4f}")
print("\nrange over seeds 0-2:  final val          train (last 200)")
for kind, r in results.items():
    tr, va = np.array(r).T
    print(f"{kind:11s}          {va.min():.4f}-{va.max():.4f}      {tr.min():.4f}-{tr.max():.4f}")

Output:

validation loss of the zero predictor 0.9992, of the Gaussian denoiser 0.1200

toy DiT (d = 64, 4 blocks) on a 16x16x4 latent
 p  tokens  params (model / formula)   MACs (counted / formula)
 1     256     328,516 / 328,516      84,144,128 / 84,144,128
 2      64     330,064 / 330,064      14,938,112 / 14,938,112
 4      16     336,256 / 336,256       3,534,848 /  3,534,848
 8       4     361,024 / 361,024       1,052,672 /  1,052,672

DiT on 32x32x4 latents: formula against Table 4 (supplement); 8 output channels, 1,001 labels
model      block            tokens  GMACs  Table 4 Gflops  params (M)  Table 4 (M)
DiT-S/8    adaLN-Zero         16     0.36        0.36         33.1         33
DiT-S/4    adaLN-Zero         64     1.41        1.41         32.9         33
DiT-S/2    adaLN-Zero        256     6.06        6.06         32.9         33
DiT-B/8    adaLN-Zero         16     1.42        1.42        130.9        131
DiT-B/4    adaLN-Zero         64     5.57        5.56        130.4        130
DiT-B/2    adaLN-Zero        256    23.01       23.01        130.3        130
DiT-L/8    adaLN-Zero         16     5.01        5.01        458.6        459
DiT-L/4    adaLN-Zero         64    19.70       19.70        458.0        458
DiT-L/2    adaLN-Zero        256    80.70       80.71        457.8        458
DiT-XL/8   adaLN-Zero         16     7.39        7.39        675.7        676
DiT-XL/4   adaLN-Zero         64    29.04       29.05        675.0        675
DiT-XL/2   adaLN-Zero        256   118.62      118.64        674.8        675
DiT-XL/2   in-context        256   119.35      119.37        449.2        449
DiT-XL/2   cross-attention   256   137.60      137.62        598.0        598
DiT-XL/2   adaLN             256   118.55      118.56        600.5        600

500 steps, batch 128; validation loss every 100 steps; train = mean of the last 200 steps
adaLN-Zero  seed 0  val 0.1005 0.0827 0.0721 0.0678 0.0628   train 0.0673
adaLN-Zero  seed 1  val 0.1076 0.0860 0.0698 0.0622 0.0639   train 0.0648
adaLN-Zero  seed 2  val 0.1025 0.0829 0.0724 0.0652 0.0602   train 0.0645
adaLN       seed 0  val 0.1234 0.1004 0.0855 0.0801 0.0708   train 0.0782
adaLN       seed 1  val 0.1212 0.0914 0.0844 0.0747 0.0761   train 0.0767
adaLN       seed 2  val 0.1232 0.0972 0.0866 0.0784 0.0747   train 0.0766
in-context  seed 0  val 0.1922 0.1159 0.0978 0.0810 0.0700   train 0.0805
in-context  seed 1  val 0.1932 0.1168 0.0946 0.0758 0.0725   train 0.0785
in-context  seed 2  val 0.2008 0.1197 0.1037 0.0818 0.0727   train 0.0837

range over seeds 0-2:  final val          train (last 200)
adaLN-Zero           0.0602-0.0639      0.0645-0.0673
adaLN                0.0708-0.0761      0.0766-0.0782
in-context           0.0700-0.0727      0.0785-0.0837

The closed form agrees exactly with the counted multiply-adds and the model’s parameters for every patch size, and on the paper’s configurations it reproduces Table 4 to within 0.02 Gflops and to the rounding of the parameter counts, including the in-context, cross-attention and adaLN blocks; the paper’s “Gflops” therefore behave as multiply-adds in our accounting. On the toy model, halving $p$ from 2 to 1 quadruples the tokens and multiplies the multiply-adds by 5.6, more than four because the attention products grow 16-fold, while the parameters change by under 1% and fall as $p$ shrinks, since only the patch embedding and the output layer scale with $p^2$. The formula shows the same pattern at full size: 675.7 million parameters for DiT-XL/8 and 674.8 million for DiT-XL/2, which Table 4 rounds to 676 and 675.

In training, all three designs pass the Gaussian denoiser’s 0.1200 within 200 steps, so each learns more than the latents’ second-order statistics. adaLN-Zero has the lowest validation loss at every checkpoint and the lowest training loss for every seed, with final validation losses of 0.060 to 0.064 against 0.070 to 0.076 for the other two, adaLN-Zero first, as in DiT’s Fig. 5. The other two order differently on the two measures: in-context conditioning starts far behind (0.19 to 0.20 at step 100, against about 0.12 for adaLN) and has the highest training loss over the last 200 steps, but its final validation loss overlaps adaLN’s range. At this scale, with 4 blocks and 500 steps, the example shows zero-initialized, gated modulation training faster than adaLN, which has neither the gates nor the zero initialization; it does not show which of the other two would end ahead after longer training, and it measures loss, not sample quality or FID.