Scaling Laws For Diffusion Transformers

TL;DR

Liang et al. (Shanghai AI Lab) run the first systematic Chinchilla-style isoFLOP sweep on text-to-image Diffusion Transformers (DiT), covering compute budgets from 1e17 to 6e18 FLOPs with models spanning 1M–1B parameters, and confirm a clean power-law relationship between compute, optimal model size, optimal token count, and pretraining loss. Fitted exponents yield $N^* \propto C^{0.56}$ and $D^* \propto C^{0.43}$ under an in-context DiT, with model size scaling slightly faster than data — mirroring but not matching LLM Chinchilla scaling. Crucially, the same power-law extends to FID (Fréchet Inception Distance), so generation quality itself becomes predictable from compute, and the relationships transfer to out-of-domain COCO validation with only a constant vertical offset. A 1B-parameter model trained at 1.5e21 FLOPs matches the extrapolated loss and FID, validating a ~250× compute extrapolation from the fit.

Key claims

DiT pretraining loss follows a Chinchilla-style power law in compute; a 1B / 1.5e21 FLOP extrapolation from fits at 1e17–6e18 FLOPs matches the predicted loss to within tight tolerance [§3.2, Fig. 3].
The compute-optimal ratio is $N^* \propto C^{0.56}$ , $D^* \propto C^{0.43}$ — model size grows slightly faster than data with compute, and the loss exponent is $-0.0273$ for in-context DiT [§3.2 Eqs. 5–7, Tab. 1].
FID scales as a clean power law in compute for the compute-optimal model at each budget, and the fit accurately predicts FID at the 1.5e21-FLOP extrapolation point [§3.3, Fig. 3 (left), Eq. 8].
Scaling laws are robust to out-of-domain evaluation: on COCO 2014 validation, all four metrics (validation loss, VLB, exact likelihood, FID) preserve the power-law trend with only a constant vertical offset from Laion5B [§3.4, Fig. 4].
Four training-time metrics — training loss, validation loss, offset VLB, and offset exact likelihood — track each other under scaling and can each be used to observe scaling behavior; training loss is the practical default [§3.1, Fig. 2].
Cross-Attention DiT has a smaller loss exponent (−0.0385 vs −0.0273) and a slightly larger data exponent than In-context DiT, indicating better compute efficiency and a prescription to shift compute toward data as budgets grow [§4, Tab. 1, Fig. 5].
Scaling exponents themselves are proposed as a benchmark: a better model architecture drops the model exponent and raises the data exponent under fixed data; a better dataset does the reverse under fixed model; and an improved training pipeline lowers the loss exponent [§4].

Method

The pretraining setup uses a vanilla in-context DiT (image, text, and timestep tokens concatenated), Rectified Flow with $v$ -prediction, and Logit-Normal timestep sampling per SD3. Training data is a 108M subset of Laion-Aesthetic recaptioned with LLaVA-1.5, with a 1M-pair validation set. All runs use AdamW (β=(0.9, 0.95), wd=0.01, ε=1e-15), batch size 1024, gradient clip 1.0, BF16, and CFG dropout 0.1. For each compute budget in {1e17, 3e17, 6e17, 1e18, 3e18, 6e18} FLOPs, the authors sweep transformer depth from 2 to 15 layers, fit an isoFLOP parabola in log-LR to identify the compute-optimal (N, D), and then fit power laws across budgets. To validate, a 958.3M-parameter model is trained at 1.5e21 FLOPs and compared to the extrapolated loss and FID. Compute is estimated as $C = 6ND$ .

Results

IsoFLOP parabolas at each budget align cleanly with empirical loss curves for budgets ≥3e17 FLOPs; the 1e17 fit is noisier [§3.2, Fig. 1].
Extrapolated 1B / 1.5e21 FLOP model: training loss matches the fitted power-law prediction, and FID matches the extrapolated FID trend [§3.2–3.3, Fig. 3].
On out-of-domain COCO validation, validation loss, VLB, and exact likelihood show a constant vertical offset from Laion5B across budgets; FID’s offset widens with budget but still follows a power law in compute [§3.4, Fig. 4].
Cross-Attention DiT achieves a steeper loss decrease per unit compute than In-context DiT at matched budgets, consistent with the DiT paper’s original architecture ranking [§4, Fig. 5, Tab. 1].

Why it’s interesting

This paper is the closest thing the wiki has to a canonical “Chinchilla for DiT” citation, and it predates most of the current wave of diffusion-side scaling-law work. It provides the isoFLOP methodology and the specific exponents that subsequent papers either replicate, refine, or contest — most directly Chimera: Designing and Chinchilla-Scaling Hybrid Visual Diffusion Transformers, which runs a three-axis Chinchilla fit (activated params × tokens × image/video ratio) on a hybrid KDA/MLA/MoE DiT with HeteroP hyperparameter transfer, and reports compute-optimal splits close to even for image but modestly size-favoring for video. It also complements The Design Space of Tri-Modal Masked Diffusion Models‘s finding that tri-modal masked diffusion has $D^* \propto N^{0.476}$ (larger models more data-efficient) — the inverse asymmetry of what Liang et al. find for continuous DiT, an interesting objective-dependent contrast worth revisiting. The proposal to use scaling exponents themselves as a benchmark for architecture/data quality is methodologically influential, and predates the “isoFLOP fit tells you if your change is good” framing now common in Hyperparameter scaling laws and Diffusion training efficiency.