PaperlayerView exact revision on arXiv

Source-linked AI summary

Abra: Scaling Diffusion Image Training

Kyle Chickering, Wei-An Lin, Swayam Bhanded, Dan Saunders, Akshat Tripathi, Jiaming Song, Shyamal Buch, Xinchen Yan

arXiv:2608.17286v1cs.LG

TL;DR

Scaling laws for visual generation remain less established than for language models. This paper systematically studies compute-optimal scaling in controlled text-to-image diffusion transformers and finds predictable behavior with optimality at roughly 200 image tokens per parameter.

  • Problem

    Compute-optimal scaling for visual generation remains understudied despite large-scale text-to-image studies and domain-specific measurement challenges.

  • Method

    The paper trains controlled ABRA flow-matching transformers from 60M to 2B parameters across 10^19–10^22 FLOPs to fit compute-optimal scaling laws.

  • Results

    Approximately 200 image tokens per parameter achieves compute optimality, while diffusion models scale predictably across generative metrics, representation quality, and training-curve shape.

  • Takeaways & Limitations

    Practitioners should favor training smaller diffusion models on more data, while accounting for resolution-dependent token requirements and robust overtraining behavior.

  • Takeaways & Limitations

    The study uses a single text encoder, whose variation may measurably shift model output quality and the resulting scaling laws.

Abstract

from arXiv · show

Compute-optimal scaling laws guide the training of frontier language models yet remain largely unexplored for visual generation. We present a systematic scaling law study for text-to-image diffusion models using Abra, a controlled family of flow-matching transformers trained across three orders of magnitude worth of compute ($10^{19}$ to $10^{22}$ FLOPs), reaching significantly larger compute budgets than previous works. We demonstrate that diffusion models scale just as predictably as language models but require far more data to train optimally: compute optimality occurs at approximately $200$ image tokens per parameter, ten times the Chinchilla compute-optimal prescription for LLMs. We show that unlike language models, diffusion models are robust to overtraining and that practitioners should err on the side of more data rather than a larger model. Finally, we show that this predictability extends beyond training loss to generative quality metrics, optimal CFG settings, representation quality, and even the shape of the training curves, which collapse onto a universal form.

1 INTRODUCTION

This paper develops compute-optimal scaling laws for text-to-image diffusion models through a controlled ABRA study spanning three orders of magnitude in compute. It finds an actionable 200 image tokens-per-parameter rule and predictable scaling across losses, generative metrics, representations, and training behavior.

  • Visual generation scaling is challenging because images are high-dimensional, token information density depends on resolution, training and generation differ, and training curves are noisier.
  • ABRA studies dense flow-matching transformers from 60M to 2B parameters across 10^19 to 10^22 FLOPs using µP-fitted scaling laws.
  • 200 image tokens per parameter is the compute-optimal training prescription for text-to-image diffusion models, a 10× increase over the standard Chinchilla rule.
  • Evaluation loss, common generative metrics, and representation quality all exhibit predictable scaling trends.
  • Diffusion model training is robust to overtraining, while representation and generation capabilities scale heterogeneously and diffusion models exhibit scaling collapse.

2 RELATED WORK

Prior work established predictable neural scaling laws and introduced transformer- and flow-matching-based approaches for generative diffusion modeling. Existing diffusion scaling studies, downstream benchmark analyses, and concurrent work motivate this study while leaving distinct compute-optimal and foundational questions.

  • Text-to-Image Transformers and Flow-Matching: Flow matching largely replaced original diffusion formulations by expressing denoising as simple linear interpolation, following transformer-based diffusion studies.Peebles & Xie (2023) studied vision transformers for generative diffusion modeling; the paper uses v-preconditioning.
  • Neural Scaling Laws: Neural scaling-law studies show predictable performance improvements as compute and model size increase, particularly for large language models.The modern study of these laws is associated with Amodei et al., Hestness et al., Kaplan et al., and Hoffmann et al.
  • Scaling Laws for Diffusion Models: Liang et al. (2024) provide the only prior text-to-image diffusion scaling-law study, whereas this work uses an order of magnitude more compute and a controlled model family.Their finding that model size should scale faster than dataset size leads to systematic undertraining at frontier scales; this work instead finds D ∝ C^1/2.
  • Related Diffusion Scaling Studies: Other studies examine downstream scalability, inference scaling, or optimal CFG, while concurrent work agrees that diffusion training requires more data than LLM training but addresses different questions.Li et al. study downstream benchmarks without compute-optimal analysis; Mei et al. focus on inference scaling; Ge et al. study compute-optimal image and video generation concurrently.

3 METHOD

The method defines compute-optimal diffusion scaling as choosing model and dataset sizes that minimize a loss or evaluation metric under fixed compute. It studies this question with the ABRA family, standardized latent-diffusion training, and large-scale web image-text data.

  • Compute-optimal scaling: Compute optimality selects model size N and dataset size D to minimize loss or evaluation metric L under a fixed training budget C FLOPs.This follows the definition adopted from Hoffmann et al. (2022).
  • Compute-optimal scaling: Data is measured as image tokens per parameter (TPP), while N counts diffusion-transformer parameters excluding the frozen text encoder.Including or excluding embedding and unembedding layers changes the optimal-TPP estimate by less than 3%.
  • Training and sampling: Training uses a v-prediction flow-matching objective, maintains an EMA with decay 0.9999, and samples latent images with a 50-step Euler solver.The ODE is integrated backward from Gaussian noise, followed by decoding with the FLUX autoencoder; classifier-free guidance is swept per model and fixed across sampling steps.
  • ABRA model family: ABRA spans 60M–2B parameters and approximately 10^19–10^22 image-token FLOPs, varying depth from 19–33 blocks and width from 384–1792.The family uses four double-stream blocks followed by single-stream blocks, with FFN expansion rate 6 and a frozen Qwen3-4B text conditioner.
  • Data and optimization: Training uses DataComp-1B web image-text pairs, re-captioned with open-source vision-language models and mixed across dense, medium, sparse, and web-sourced captions.The study focuses exclusively on diffusion-model pre-training without mid- or post-training.

4 EMPIRICAL RESULTS

Diffusion transformers scale predictably across metrics, but reach compute optimality at substantially higher data budgets than language models. They are robust to overtraining, while optimal allocations vary by metric, representation objective, resolution, and training progress.

  • Compute-optimal scaling: 200 TPP is compute-optimal for diffusion models, roughly 10× the canonical LLM budget of 20 TPP.N denotes parameter count and D is measured in image tokens.
  • Overtraining and undertraining: Less than 0.5% loss penalty occurs under 2× overtraining, whereas undertraining is catastrophic and favors a smaller overtrained model.The practical rule is to err toward overtraining a small model rather than undertraining a large one when compute is limited.
  • Generative metrics: FID and KID favor faster dataset growth, while CLIPScore and CMMD favor faster parameter growth, so no single metric defines compute optimality.Generative metrics follow power laws in compute, but their compute-optimal allocations diverge.
  • Generative metrics: Optimal CFG decreases with both model size and training steps, indicating that stronger generative capacity needs less guidance.The feature model behind each metric also shifts the optimal CFG.
  • Image understanding: Linear-probe accuracy follows tight scaling laws and reaches compute-optimality at far lower TPP than generation, peaking in middle layers near timestep t ≈0.7.A joint generation-and-understanding double-point is estimated at about 5 × 10^22 FLOPs with a 6B-parameter model.
  • Scaling collapse: Compute-optimal training curves collapse when rescaled by relative compute, extending collapse behavior beyond MLPs and language models.The collapse is observed for ABRA-60M, ABRA-120M, and ABRA-250M, despite substantial diffusion training noise.
  • Resolution scaling: Compute-optimal TPP rises from about 165 at 256 × 256 resolution to around 247 at 768 × 768 resolution.High-resolution diffusion training is compute bound rather than data bound.

5 CONCLUSION

ABRA shows that text-to-image diffusion transformers obey tight, predictable scaling laws, with compute optimality at approximately 200 TPP. Diffusion training is robust to overtraining, while generative metrics, representation quality, and training-curve shape exhibit distinct but predictable scaling behavior.

  • Compute-optimal scaling: 200 TPP is the compute-optimal point for text-to-image diffusion transformers, an order of magnitude higher than for LLMs.ABRA spans 60M to 2B parameters and uses 10× as much compute as previous studies.
  • Scaling behavior: Diffusion model training is robust to overtraining, and FID and CLIPScore follow predictable scaling trends in compute.As resolution increases, models require more tokens to reach compute optimality.
  • Scaling behavior: Internal representation quality scales differently from generative capabilities, while compute-optimal training exhibits scaling collapse.The study reports scaling collapse for diffusion model training for the first time.

A ARCHITECTURAL DECISIONS

ABRA combines design choices drawn from recent image-generation architectures with scaling-oriented parameterization and normalization adjustments. The models remove biases and weight decay, while training uses repetition factors that diffusion models can tolerate.

  • Scaling Family: ABRA uses ungated SiLU activations, FFN expansion ratio 6, fixed head dimension 128, width scaled through head count, and RMSNorm-based normalization.The architecture draws components from FLUX.2 and PixArt-α, using parametric RMSNorm for QK normalization and non-parametric RMSNorm for the layer.
  • Removal of Bias Terms and Weight Decay: Removing all bias terms slightly improved performance in initial architectural ablations, while training uses no weight decay.No weight decay is described as standard in diffusion model training.
  • Maximal Update Parameterization (µP): ABRA applies µP to transfer optimal learning rates across model sizes, validating this transfer for ABRA-120M and ABRA-250M.The implementation was checked using spectral coordinate checking, which the authors describe as more robust to implementation errors than the original coordinate-checking formulation.
  • Softmax Scaling Factor: QK-normalization combined with µP occasionally caused training instabilities, attributed to altered softmax-argument scaling.The authors empirically compared alternative softmax scalings in the presence of QK-normalization.
  • Epoch Repetition Factors: At 400 TPP, the 2B model trains for nearly 8 epochs, reflecting diffusion models’ reported tolerance for repeated data exposure.Table 2 reports epoch repetition factors as passes over the 100M-sample mix.

B LIMITATIONS

The scaling results depend on several fixed design choices, and varying them may shift the compute-optimal TPP and broader scaling laws. The study examined neither alternative optimizers and encoders nor compute-optimal batch-size scaling.

  • General limitations: Fixed experimental axes may shift the optimal TPP, as already observed when changing resolution.The authors speculate that ablations along these axes could alter the reported optimum.
  • Optimizers: Only Adam was evaluated; Muon, Shampoo, and SOAP may scale differently, potentially shifting scaling-law intercepts.In LLMs, these optimizers appear to shift the intercept without changing compute-optimal TPP.
  • Text Encoder: Only one text encoder was used, and encoder variations are expected to measurably affect output quality and shift the scaling laws.The authors did not test alternative text encoders.
  • Vision Encoder: The vision encoder was not ablated, although its choice may materially affect text-to-image generation scaling behavior.The passage frames vision-encoder effects as likely but unmeasured in this work.
  • Batch Size Scaling: Compute-optimal batch-size scaling for diffusion models was not empirically studied because analogous LLM studies are recent and cheaper.The authors suspect analogous scaling laws exist for diffusion models but leave them beyond this work’s scope.

C THE EFFECT OF DIFFUSION TIMESTEP · D MODEL GENERATIONS

Diffusion timestep losses are heterogeneous, but each timestep bucket follows a clean scaling law with different rates of compute-optimal learning. Model generations are presented across prompt complexities and guidance scales using non-cherry-picked Parti Prompts samples.

  • C THE EFFECT OF DIFFUSION TIMESTEP: Each timestep individually obeys a scaling law despite heterogeneous flow-matching loss across diffusion timesteps.This finding concerns per-timestep behavior rather than only the aggregate training loss.
  • C THE EFFECT OF DIFFUSION TIMESTEP: The aggregate training loss combines per-timestep losses weighted by the logit-normal timestep distribution.The analysis isolates per-sample loss, its expectation, and the loss response curve.
  • C THE EFFECT OF DIFFUSION TIMESTEP: The oracle loss minimizes the loss response curve over an appropriate function space, with endpoint values determined by the data distribution and independent x and ε.At each endpoint, the model lacks information about the other quantity, so the minimizing prediction is its mean.
  • C THE EFFECT OF DIFFUSION TIMESTEP: The loss response curve is highly heterogeneous: endpoints approach oracle floors, while interior timesteps contain most reducible loss.The curve was measured by sampling the EMA model 100k times using the logit-normal training distribution.
  • C THE EFFECT OF DIFFUSION TIMESTEP: 10 uniform timestep buckets each yield well-behaved scaling-law fits, but differing exponents show heterogeneous rates of reaching compute optimality.Some timesteps saturate compute early, whereas others continue learning rapidly.
  • D MODEL GENERATIONS: Model generations use non-cherry-picked Parti Prompts sorted into simple, standard, and complex categories, with fixed guidance scale 3.0.Figure 12 presents generations at fixed guidance, while Figure 13 varies guidance by model size.

E EXTENDED SCALING ANALYSIS · F FIT PROCEDURES, ALTERNATIVES, AND THE KAPLAN/HOFFMANN RECIPES

The extended analysis compares Abra’s iso-FLOP suboptimality with several open-source LLM scaling datasets, while the fitting section documents procedures, canonical recipes, and alternative estimators. The main procedure fits per-TPP power laws to EMA-weight loss and infers optima from iso-FLOP slices.

  • E EXTENDED SCALING ANALYSIS: Abra’s iso-FLOP suboptimality is compared with open-source LLM pre-training data beyond McLeish et al. (2025).The comparisons include Gadre et al. (2025), Porian et al. (2024), and Hoffmann et al. (2022).
  • E EXTENDED SCALING ANALYSIS: The extended comparisons specifically use Gadre et al. (2025), Porian et al. (2024), and the Chinchilla paper.The corresponding analyses are presented in Figures 14, 16, and 15, respectively.
  • E EXTENDED SCALING ANALYSIS: The authors state that no openly available source is known for the Chinchilla comparison’s data.This qualification appears in the discussion of comparisons with Hoffmann et al. (2022).
  • F FIT PROCEDURES, ALTERNATIVES, AND THE KAPLAN/HOFFMANN RECIPES: The fitting section reports fitting methodology, error analysis, quality-of-fit analysis, canonical language-model recipes, and alternative estimators.All fits are collected in Table 4.
  • F FIT PROCEDURES, ALTERNATIVES, AND THE KAPLAN/HOFFMANN RECIPES: The main-body procedure fits per-TPP power laws to EMA-weight loss and reads optima from iso-FLOP slices.Each trajectory is interpolated with PCHIP, and N counts all diffusion-transformer parameters while excluding the frozen text encoder.
  • F FIT PROCEDURES, ALTERNATIVES, AND THE KAPLAN/HOFFMANN RECIPES: The per-TPP fits use trust-region non-linear least squares on raw loss values in physical space after discarding the first 5,000 warmup steps.The procedure uses PCHIP interpolation for each trajectory.

F.1 SUPPORTING-HYPERPLANE FIT … F.5 CHOICE OF INTERPOLANT

The appendix validates the compute-optimal frontier with supporting-hyperplane, Kaplan-style, and Hoffmann-style estimators, then tests robustness to pipeline choices and interpolation methods. The supporting-hyperplane fit agrees with the iso-FLOP envelope, while the parametric Hoffmann-style fit is non-identifiable on ABRA.

  • F.1 SUPPORTING-HYPERPLANE FIT: The supporting-hyperplane fit directly models the maximal power-law lower bound over measured (N, C, L) Pareto points by optimizing amplitude, offset, and exponent.It selects the exponent minimizing total slack while remaining weakly below every Pareto point.
  • F.1 SUPPORTING-HYPERPLANE FIT: <0.1%: The supporting-hyperplane fit recovers the same L*(C) as the iso-FLOP envelope on ABRA.The corresponding coefficients are reported in Table 5.
  • F.2 KAPLAN-STYLE FIT: Kaplan-style fitting regresses a pure power-law curve through compute-efficient envelope points rather than fitting a bound below them.Allocation exponents are estimated both freely and under the constraint a + b = 1.
  • F.3 HOFFMANN-STYLE (CHINCHILLA) FITS: The Hoffmann-style analysis estimates the frontier from training-curve minima and iso-FLOP profiles, with the latter used in the main body.The iso-FLOP method minimizes fitted loss-versus-TPP slices across fixed FLOP budgets and aggregates the results.
  • F.3 HOFFMANN-STYLE (CHINCHILLA) FITS: The parametric Hoffmann-style loss-surface fit uses log-space EMA loss, a Huber objective, and L-BFGS initialization grids but is non-identifiable on ABRA.The Huber threshold is δ = 10^-3, so no parametric TPP* fit is reported.
  • F.4 ALTERNATIVE PROCEDURES AND UNIFIED COMPARISON: Robustness checks vary the loss target, ladder composition, fit objective, and interpolation while recovering TPP* from identical ABRA data with 60M included or excluded.The compared choices include EMA versus raw loss, full versus reduced ladders, squared error versus Huber, and PCHIP versus piecewise-linear interpolation.
  • F.5 CHOICE OF INTERPOLANT: PCHIP provides the headline shape-preserving interpolation, benchmarked against plain piecewise-linear interpolation between checkpoints.Leave-one-out evaluation holds out each interior checkpoint and reports RMSE over the predicted losses in Table 6.

F.6 LEAVE-ONE-MODEL-OUT ROBUSTNESS · F.7 EVERY ESTIMATOR AGREES

Leave-one-model-out refits show that the iso-FLOP TPP* fit is stable except when ABRA-60M is omitted. Across estimators and fitting choices, the compute-optimal allocation consistently centers near 200 TPP, about ten times the LLM value of 20.

  • F.6 LEAVE-ONE-MODEL-OUT ROBUSTNESS: Holding out one model at a time leaves the iso-FLOP TPP* fit materially unchanged, except when ABRA-60M is omitted.This robustness check was performed for both ladder choices.
  • F.6 LEAVE-ONE-MODEL-OUT ROBUSTNESS: Both ladder choices support the same leave-one-model-out robustness conclusion.The refits bound the leverage of any individual model.
  • F.6 LEAVE-ONE-MODEL-OUT ROBUSTNESS: ABRA-60M is the sole held-out model whose omission materially affects the fit.All other single-model omissions do not materially affect the result.
  • F.7 EVERY ESTIMATOR AGREES: ≈200 TPP is a method-invariant headline estimate, with a ∼10× gap versus the LLM value of 20.The same estimate persists across estimator and fitting variations.
  • F.7 EVERY ESTIMATOR AGREES: Every identifiable estimator lands in [183, 202].This range is reported in Table 4.
  • F.7 EVERY ESTIMATOR AGREES: The balanced allocation a ≈0.47–0.49 holds under both free and constrained fits.The estimate is also unchanged under either interpolant and when holding out any single model.

F.8 GENERATIVE METRIC SCALING-LAW FITS · G COMPARISON TO LIANG ET AL. (2024)

The appendix establishes stable, CFG-robust scaling laws for generative metrics using compute-optimal and per-model-final fits, then compares allocation exponents with Liang et al. (2024) while noting tokenization differences make intercepts incomparable. It reports KID scaling with α ≈0.11 and R2 = 0.90, and relates Abra’s 200 image TPP to Liang et al.’s 287 total TPP for a 1B model.

  • F.8 GENERATIVE METRIC SCALING-LAW FITS: The iso-FLOP procedure replaces loss with CFG-optimized metrics, interpolates each size over training steps, and fits M(C) = A C^-α + FM across the model ladder.Distribution distances and FID use CFG-grid minima, while CLIPScore uses the maximum; the 60M model is excluded from the fit.
  • F.8 GENERATIVE METRIC SCALING-LAW FITS: KID follows a pure power law with α ≈0.11 and R2 = 0.90, while floor terms are not identifiable for KDD and KID.The value-law analysis fits both compute-optimal trajectories and per-model-final converged checkpoints, excluding 60M.
  • F.8 GENERATIVE METRIC SCALING-LAW FITS: Optimal CFG shifts with model size, so evaluating every model at one fixed guidance value can misrank them; for distribution metrics, optimal CFG also drifts downward with model size and compute.The workflow therefore compares models using per-model CFG optimization rather than a single fixed guidance value.
  • F.8 GENERATIVE METRIC SCALING-LAW FITS: The scaling exponent α remains stable across the swept CFG range for every metric, with only mild high-guidance declines for CMMD and CLIPScore and flat behavior for DMMD and KDD.This robustness check refits the per-model-final law at each measured CFG and drops fits with R2 < 0.5 or α pinned at the optimizer bound.
  • G COMPARISON TO LIANG ET AL. (2024): The appendix compares Abra’s compute-optimal allocation exponents with Liang et al. (2024), but their allocation multipliers and intercepts are not directly comparable because the studies count different token types.Abra measures image tokens, whereas Liang et al. measure total text-plus-image tokens; the scaling exponents remain comparable.
  • G COMPARISON TO LIANG ET AL. (2024): 200 image TPP at 512 × 512 resolution becomes 250 total TPP when 256 text tokens per sample are included.The text encoding increases estimated TPP because its tokens are averaged over samples.
  • G COMPARISON TO LIANG ET AL. (2024): For a 1B model, Liang et al. (2024) report optimal training at 287 total TPP, compared with Abra’s 250 total TPP under the stated 512 × 512 and 256-text-token conversion.Liang et al.’s reported optimal N and D concern 256 × 256 images, limiting direct comparison with Abra’s 512 × 512 setting.
Loading 2608.17286v1…