PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling Laws For Diffusion Transformers

Zhengyang Liang, Hao He, Ceyuan Yang, Bo Dai

arXiv:2410.08184v2cs.CV

TL;DR

DiT scaling laws were insufficiently explored, limiting predictive choices of model size and data for a given compute budget. The paper fits power-law relationships across broad compute budgets and tests their extrapolation to large models and datasets. It finds that pretraining loss predicts generation quality, including FID, across datasets, while the benchmark is limited to compute-constrained settings and does not compare conditioning mechanisms.

  • Problem

    Diffusion transformers lacked explicit, accurate scaling laws linking compute with model size, data quantity, and loss, making optimal configuration selection dependent on costly heuristic searches.

    Source: §1 · p. 1, §2 · p. 3

  • Method

    The paper trains DiTs from 1M to 1B parameters across 1e17–6e18 FLOPs, extracts compute-optimal configurations from isoFLOP curves, and fits power laws for compute, model size, data, and loss.

    Source: §1 · p. 1, §3.3 · p. 5

  • Results

    Pretraining loss follows a power law with compute, and generation metrics such as FID follow the same trend across datasets; a 1B-parameter-scale extrapolation also matches the predicted loss.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Scaling laws can predict compute-optimal DiT configurations and provide a lower-cost benchmark for evaluating model and dataset quality.

    Source: §1 · p. 2, §5 · p. 10

    Experiments are limited to compute-constrained settings with sufficiently large datasets, so scaling behavior under data-limited regimes remains unstudied.

    Source: I LIMITATIONS AND FUTURE DIRECTIONS · p. 33

Abstract

from arXiv · show

Diffusion transformers (DiT) have already achieved appealing synthesis and scaling properties in content recreation, e.g., image and video generation. However, scaling laws of DiT are less explored, which usually offer precise predictions regarding optimal model size and data requirements given a specific compute budget. Therefore, experiments across a broad range of compute budgets, from 1e17 to 6e18 FLOPs are conducted to confirm the existence of scaling laws in DiT for the first time. Concretely, the loss of pretraining DiT also follows a power-law relationship with the involved compute. Based on the scaling law, we can not only determine the optimal model size and required data but also accurately predict the text-to-image generation loss given a model with 1B parameters and a compute budget of 1e21 FLOPs. Additionally, we also demonstrate that the trend of pre-training loss matches the generation performances (e.g., FID), even across various datasets, which complements the mapping from compute to synthesis quality and thus provides a predictable benchmark that assesses model performance and data quality at a reduced cost.

1 INTRODUCTION

Diffusion transformers lacked explicit scaling laws connecting compute, model size, data, and loss, despite prior evidence that larger models improve visual quality and text-image alignment. This work fits power laws across compute budgets, validates predictions at large scale, and links pretraining loss with generation quality across datasets.

1 INTRODUCTION

  • The study positions explicit DiT scaling laws as a way to allocate compute efficiently and evaluate configurations through fitted power-law coefficients.The introduction frames this as a reduced-cost assessment of model and data performance.Source: §1 · pp. 1–2
  • Earlier DiT studies found that larger models improve visual quality and text-image alignment, but did not accurately predict scaling or determine compute-optimal model and data sizes.Without explicit laws, configuration selection requires costly heuristic searches.Source: §1 · p. 1
  • A power-law relationship between compute budget and pretraining loss is established for DiT across 1e17–6e18 FLOPs, enabling predictions of optimal model and data scales.The study trains models from 1M to 1B parameters, extracts compute-optimal configurations from isoFLOP curves, and fits scaling relationships.Source: §1 · pp. 1–2
  • The predicted configuration at 1.5e21 FLOPs is approximately 1B parameters, and a trained model’s final loss matches the prediction.The reported validation demonstrates that the extrapolated scaling law remains accurate at a much larger compute budget.Source: §1 · p. 2, Figure 1 · p. 2
  • Generation quality, including FID, follows the trend of pretraining loss and the same scaling patterns transfer across datasets, including out-of-domain COCO validation data.This supports using scaling laws as a lower-cost benchmark for assessing model and dataset quality.Source: §1 · p. 2

2 RELATED WORK

Prior work established power-law scaling in language models, while diffusion-transformer research mainly demonstrated qualitative scalability. This paper addresses the missing systematic account of how DiT compute, model size, data, and loss scale together.

2 RELATED WORK

  • Language-model studies established that performance scales predictably with model size, dataset quantity, and compute, motivating analogous analysis for DiTs.The related work cites Hestness et al. (2017), Kaplan et al. (2020), Henighan et al. (2020), and Hoffmann et al. (2022).Source: §2 · p. 2
  • Diffusion-model studies showed that larger DiT models improve visual quality and text-image alignment, but did not provide accurate predictive scaling laws.The unresolved gap concerns explicit relationships among compute budget, model size, data quantity, and loss.Source: §2 · p. 3
  • The paper systematically investigates DiT scaling behavior to provide a more comprehensive account of these relationships.Source: §2 · p. 3

3 METHOD

The method establishes scaling laws for DiT by fitting compute-dependent relationships across model, data, loss, and generation metrics. It uses a vanilla transformer setup, evaluates multiple loss and likelihood measures, and validates predictions across in-domain and out-of-domain data.

3 METHOD

  • Training uses Rectified Flow with v-prediction, Logit-Normal timestep sampling, a vanilla transformer, and 108 million image-text pairs from Laion-Aesthetic.The experimental settings are intended to keep architecture and training choices consistent while studying scaling behavior.Source: §3.1 · p. 3
  • Loss is the primary scaling metric because it directly reflects the diffusion model’s time-conditioned velocity-field objective and is observable during training.Training loss is estimated by Monte Carlo sampling of timesteps and noise, with batch size 1024 and EMA smoothing using αEMA = 0.9.Source: §3.2.1 · p. 4, §3.2.2 · p. 5
  • All four evaluated metrics—training loss, validation loss, offset VLB, and offset exact likelihood—show similar trends as compute increases, supporting their use for monitoring scaling laws.Training loss is prioritized because it avoids additional evaluation steps and can be tracked in real time.Source: §3.2.2 · p. 5
  • Compute-optimal model size and token count are obtained by fitting parabolas at each budget, with fits aligning closely with empirical results except at 1e17 FLOPs.The fitted relationships imply that model and data sizes must grow together, with model size increasing slightly faster than data size.Source: §3.3 · p. 6
  • The fitted loss and FID power laws are tested by extrapolating to 1.5e21 FLOPs, where the 958.3M-parameter model’s loss and the predicted FID closely match observations.FID decreases consistently as compute increases, and the large-budget prediction is reported as highly accurate.Source: §3.3 · p. 6, §3.4 · p. 7
  • Across Laion5B and COCO validation data, loss, VLB, exact likelihood, and FID retain consistent compute trends despite dataset-dependent vertical offsets.The COCO metrics are higher in absolute terms, while the FID gap widens with budget; the FID-budget relationship remains power-law predictable.Source: §3.5 · pp. 7–8

4 SCALING LAWS AS A PREDICTABLE SCALABILITY BENCHMARK

Scaling laws provide a benchmark for comparing how efficiently DiT architectures and datasets benefit from additional compute. In the matched comparison, Cross-Attention Transformers improve loss faster than Vanilla In-Context Transformers, while the benchmark is intended to assess scalability within an architecture rather than declare a universally superior conditioning mechanism.

Benchmark utility

  • Scaling exponents can benchmark model and dataset design quality by measuring how efficiently performance improves as compute and data increase.The analysis is explicitly not intended to prove that In-Context Conditioning is inherently superior to Cross-Attention.Source: §4 · pp. 8, 10

Evaluation principles

  • With a fixed dataset, a more efficient model has a lower model scaling exponent and higher data scaling exponent, indicating more effective use of additional data.The benchmark interprets these exponents as indicators of how model and data resources should be balanced.Source: §4 · pp. 8–9
  • With a fixed model, a higher-quality dataset has a lower data scaling exponent and higher model scaling exponent, implying more efficient scaling with fewer resources.Across model and data changes, a smaller loss/FID exponent indicates better performance relative to compute.Source: §4 · p. 9

Architecture comparison

  • The Cross-Attention Transformer shows a steeper loss decline and better performance at the same compute than the Vanilla In-Context Transformer.The comparison uses compute budgets spanning 1e17 to 1e19 FLOPs, and the fitted curves support the observed difference.Source: §4 · p. 9
  • A smaller model exponent for Cross-Attention indicates that increasing compute should allocate relatively more resources to dataset scaling.Its smaller loss exponent also indicates a more rapid decline in loss as compute increases.Source: §4 · p. 9, Table 1 · p. 10

5 CONCLUSION

The paper concludes that DiT training loss follows a power law with compute and that this relationship predicts model, data, and generation-performance behavior across datasets. It presents scaling laws as a benchmark for optimizing model and data design, while the supplied conclusion materials also document the paper’s supporting appendices and disclosure about LLM use.

Conclusion

  • DiT pretraining loss follows a power-law relationship with compute across budgets from 1e17 to 6e18 FLOPs.The relationship supports predictions of optimal model size, data requirements, and performance at larger budgets such as 1e21 FLOPs.Source: §5 · p. 10
  • Generation quality, measured with metrics such as FID, tracks compute-based scaling and remains robust across different datasets.The paper reports that training budgets can predict visual quality beyond a specific data distribution.Source: §5 · p. 10
  • Scaling laws can serve as a predictable benchmark for evaluating and optimizing both DiT model architectures and data design.The paper validates this potential using In-Context and Cross-Attention Transformers.Source: §5 · p. 10

Appendix organization

  • The appendix provides notation, related work, experimental details, likelihood derivations, FLOPs counting, extra experiments, limitations, and mechanism discussions.It also includes a notation-and-abbreviation summary in Table 2.Source: APPENDIX · p. 15

Disclosure

  • The authors state that LLMs were used only to polish wording and readability, not for research design, data analysis, or result interpretation.They attribute the scientific content and conclusions entirely to the authors.Source: C LLM USAGE · p. 17

E.1 DATA

The experiments use JourneyDB for formulation and model-design ablations and a curated, re-captioned Laion-Aesthetic subset for large-scale training. The evaluated architectures include PixArt and In-Context Transformers with normalization and patch-size choices specified for training.

Datasets

  • JourneyDB is used for several ablation studies on formulation and model design.The paper also curates a 108 million image-text subset from Laion-Aesthetic using an aesthetic-score threshold of 5.Source: §E.1 · p. 18
  • The Laion-Aesthetic subset is re-captioned with LLAVA 1.5 using the LLaVA-Lightning-MPT-7B-preview model.The supplied passage describes this preprocessing after collection.Source: §E.1 · p. 18

Model configurations

  • The evaluation includes PixArt and In-Context Transformer architectures, with QK-Norm, RMSNorm, and patch size 2 used for large-scale model training.The In-Context Transformer design follows the block described in Peebles & Xie (2023).Source: §E.2 · p. 18

F DERIVATION OF THE LIKELIHOOD

The paper estimates likelihood in latent diffusion models using a variational lower bound and a normalizing-flow formulation. Its FLOPs accounting includes context-dependent attention terms, while EMA smoothing stabilizes noisy training-loss estimates and affects fitted scaling curves.

Likelihood estimation

  • VLB estimation uses a latent-space variational lower bound, but the fixed Stable Diffusion 1.5 VAE makes the surrogate differ from the true VLB by a constant factor.The reconstruction and encoder-entropy terms remain constant across models, so the offset VLB is estimated directly in latent space.Source: F DERIVATION OF THE LIKELIHOOD · pp. 18–19
  • The latent-space VLB treats noise-schedule terms as irreducible losses and reduces the remaining Gaussian KL terms to weighted L2 distances.The model predicts velocity, which is used to estimate the clean sample before computing the VLB.Source: F DERIVATION OF THE LIKELIHOOD · pp. 19–20
  • Normalizing-flow likelihood integrates probability-density change along the diffusion ODE, using reverse-time sampling to convert clean samples into Gaussian noise.The experiments use 500 reverse-sampling steps for a relatively accurate estimate.Source: F DERIVATION OF THE LIKELIHOOD · p. 20

FLOPs calculation

  • FLOPs counting follows attention-block computation while retaining context-dependent terms because model dimension is comparable to context length.The analysis reports FLOPs per sample and details separate in-context and cross-attention transformer calculations.Source: G SCALING FLOPS COUNTING · pp. 20–21

Training-loss and generation evaluation

  • EMA smoothing reduces stochastic training-loss fluctuations, but larger coefficients can overestimate early loss and bias scaling-curve fitting; αEMA = 0.9 provides a balance.EMA models also exhibit a power-law scaling trend, and FID scaling remains linear across tested CFG values.Source: §H.2 · p. 24, §H.3 · p. 26, §H.4 · p. 26

H.5 SCALING LAWS FOR MORE ARCHITECTURES & DATA

Scaling-law behavior extends beyond the primary diffusion-transformer setup. The reported experiments find consistent trends across architectures, image resolutions, and evaluation settings involving CFG and EMA models.

More architectures

  • PixArt and Flux both exhibit clear linear loss-scaling behavior, suggesting consistency across diffusion-transformer architectures.The corresponding results are shown in Fig. 13(b) and Fig. 12(b).Source: §H.5 · p. 26, Figure 12 · p. 28, Figure 13 · p. 28

More data resolutions

  • Scaling laws remain linear for 512 × 512 images trained with compute budgets from 1e17 to 6e18 FLOPs.This indicates that the observed scaling behavior holds across the tested data resolutions.Source: §H.5 · p. 26, Figure 14 · p. 29

CFG scale

  • FID scaling curves remain largely invariant to CFG scale, with linear scaling behavior across all tested CFG values.The comparison is presented in Fig. 10.Source: Figure 10 · p. 27, Figure 9 · p. 27

H.6 DATA ASSESSMENT

Scaling laws can distinguish captioning quality when image content is held constant. Dense captions produce faster FID improvement than sparse tag captions in the reported comparison.

Data quality assessment

  • The dense-caption dataset achieves an FID scaling exponent of -0.234, compared with -0.216 for sparse tag captions.The authors interpret the steeper dense-caption trend as evidence of superior data quality.Source: §H.6 · p. 29

H.7 TRANSFER RESULTS ON MORE DATASETS

The scaling-law findings transfer to new evaluation datasets. On Flickr30k and JourneyDB, both loss and FID retain clear linear relationships with compute.

Transfer results

  • On Flickr30k and JourneyDB, pretrained models preserve clear linear loss and FID relationships across compute budgets.The transfer experiments follow the same setup as COCO and are reported in Fig. 16 and Fig. 17.Source: §H.7 · pp. 29–30, Figure 16 · p. 30, Figure 17 · p. 30

H.8 GENERATED SAMPLES

Generated quality improves with compute, while scaling relationships remain consistent across evaluation metrics, sampling schedules, datasets, and conditioning architectures. These results support predictable compute-to-quality trends rather than state-of-the-art generation claims.

Generated samples

  • Sample quality improves as compute increases, although the experiments target scaling behavior rather than state-of-the-art image generation.Samples use fixed prompts and seeds.Source: §H.8 · p. 30, Figure 18 · p. 31

ImageNet scaling

  • Loss curves retain clear power-law trends in the finite-data ImageNet setting across compute budgets from 2 × 1017 to 6 × 1018.Source: §H.9 · p. 30, Figure 19 · p. 31

Evaluation metrics

  • GenEval and human-preference rewards also follow consistent power-law scaling across compute-optimal models.GenEval measures text-image alignment, while HPSv2.1 and ImageReward capture human preference.Source: §H.10 · p. 32, §H.11 · p. 32, Figure 20 · p. 32

Sampling schedule

  • Changing diffusion sampling steps shifts overall performance, but results still form a clear straight line in log–log space.The comparison uses 35 and 50 sampling steps with classifier-free guidance scale fixed to 10.0.Source: §H.12 · p. 32, Figure 22 · p. 33

Conditioning architecture

  • Cross-attention has a steeper FID scaling slope and greater scalability than the vanilla in-context architecture in the reported comparison.This matches the corresponding loss-curve trend.Source: §H.13 · p. 33, Figure 23 · p. 33

I LIMITATIONS AND FUTURE DIRECTIONS

The study’s conclusions are bounded by its text-to-image focus, compute-limited regime, and shared hyperparameter settings. Its conditioning comparison is also limited to simplified, carefully matched architectures rather than modern systems with different recipes.

Limited modalities and downstream tasks

  • Observed scaling laws are evaluated only for text-to-image generation and a small set of aggregate metrics, leaving broader modalities and downstream tasks open.Unexamined settings include text-to-video, image editing, and inpainting.Source: I LIMITATIONS AND FUTURE DIRECTIONS · p. 33

Compute-limited regime

  • Experiments assume sufficiently large datasets and vary compute mainly through model size and training steps, so data-limited scaling may deviate.The paper leaves systematic corrections for finite-data regimes to future work.Source: I LIMITATIONS AND FUTURE DIRECTIONS · p. 33

More accurate hyperparameters

  • Reusing one tuned hyperparameter set across compute budgets may leave scaling curves biased relative to fully optimized operating points.Learning-rate schedules, regularization, and optimization settings could reduce variance and systematic bias.Source: I LIMITATIONS AND FUTURE DIRECTIONS · p. 34

In-context versus cross-attention

  • The cross-attention advantage is established only against a simple in-context baseline, while modern in-context systems differ in conditioning, timestep design, data, compute, and training recipe.Therefore, the reported comparison does not establish an intrinsic scalability advantage over advanced in-context architectures.Source: J DISCUSSION: IN-CONTEXT VS.… · p. 34

Related papers

Loading 2410.08184v2…