Abra: Scaling Diffusion Image Training

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

cs.LGAug 18, 2026

Abra: Scaling Diffusion Image Training

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

Organizations: Luma AI

Abstract

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.

Figures & tables

Figure 1: Text-to-image diffusion transformers follow predictable scaling laws and require ten times as much data as a comparably sized language model to reach compute optimality.

Figure 2: $512\times 512$ generations of the prompt ‘‘a rabbit’’ across model scales and CFG.

Figure 3: Left: The raw training loss for Abra runs. The bold curves are a simple moving average of the scattered loss data. The solid black line is a qualitative fit of this data. Right: The noise decomposition for the Abra -250M run, illustrating that the per-step differences in the loss curve are dominated by the batch-to-batch noise.

Figure 4: Overtraining compute-optimal models incurs a cost at fixed compute. The compute-optimal point occurs when the model loss curve is tangent to the scaling law curve.

Figure 5: Normalized iso-FLOP curves. (a) EMA loss

L/L^{*}(C)

and (b) linear-probe accuracy

\mathrm{acc}/\mathrm{acc}^{*}(C)

versus TPP for several compute budgets, each normalized to its per-budget optimum (star). The near-flat curves show the compute-optimal TPP is largely budget-insensitive; representational quality (b) is optimal at a substantially lower TPP than the loss (a).

Figure 6: Relative iso-FLOP loss penalty

\Delta L/L^{*}

(Eq. ( 2 )) as a function of TPP for (left) the Abra text-to-image diffusion family and (right) the Gemstones LLM suite ( McLeish et al., 2025 ) . Abra markers denote evaluated checkpoints and Gemstones curves are the raw per-checkpoint fineweb-edu validation-loss penalties. The shaded bands mark the

0.2\%

penalty threshold (white below, grey above).

Figure 7: Standard generative metrics scale predictably. Top: generative quality vs training compute $C$ for FID, KID, CLIPScore, and CMMD. The bold red curve is its compute-optimal frontier. Bottom: compute-optimal $\mathrm{TPP}^{*}$ as a function of model size $N$ , and the compute-optimal $N_{\text{opt}}$ and $D_{\text{opt}}$ as functions of compute $C$ .

Figure 8: (a) Best top-1 accuracy as a function of relative layer depth. (b) Best top-1 accuracy as a function of diffusion timestep. (c) Top-1 accuracy follows predictable scaling laws at different TPP. (d) Top-1 accuracy as a function of training compute.

Figure 9: (a) Training loss curves for six models from the Abra family rescaled with Equation ( 3 ). (b) The deviation $\Delta$ between models collapses beneath the noise floor $\sigma$ for the three seed-replicated models ( Abra -60M, Abra -120M, and Abra -250M).

Figure 10: (Top Row) Training curves and loss scaling laws at different resolutions. (Bottom Row) TPP as a function of resolution, compute-optimal model size, compute-optimal number of images , compute-optimal dataset size in number of image tokens. For these data we obtained estimated compute-optimal TPP of

165

202

235

, and

247

respectively for the image sizes

256

384

512

, and

768

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Params	Double	Single	Width	Heads	Batch (img.)	Steps	Image-tokens
Abra -60M	$60.1$ M	$4$	$15$	$384$	$3$	$128$	$183$ k	$2.4\times 10^{10}$
Abra -120M	$118.0$ M	$4$	$18$	$512$	$4$	$256$	$180$ k	$4.7\times 10^{10}$
Abra -250M	$262.3$ M	$4$	$18$	$768$	$6$	$512$	$200$ k	$1.0\times 10^{11}$
Abra -500M	$497.0$ M	$4$	$20$	$1024$	$8$	$1024$	$190$ k	$2.0\times 10^{11}$
Abra -1B	$998.8$ M	$4$	$22$	$1408$	$11$	$2048$	$191$ k	$4.0\times 10^{11}$
Abra -2B	$1.97$ B	$4$	$29$	$1792$	$14$	$4032$	$191$ k	$7.9\times 10^{11}$

Appendix

Table 1: The Abra model family and its training configuration. “Double” and “Single” count the double- and single-stream MMDiT blocks ( Esser et al., 2024 ; Black Forest Labs, 2025 ) ; “Batch” is the global batch size in images. Parameter counts exclude the frozen text encoder.

Model	Images $D_{I}$	Epochs at 200 TPP	Epochs at 400 TPP
Abra -60M	$23.4$ M	0.12	$0.23$
Abra -120M	$46.1$ M	0.23	$0.46$
Abra -250M	$102.4$ M	0.51	$1.02$
Abra -500M	$194.6$ M	0.98	$1.95$
Abra -1B	$391.2$ M	1.95	$3.91$
Abra -2B	$770.1$ M	3.85	$7.70$

Appendix

Table 2: Epoch repetition factor (number of passes over the

100

M-sample mix) for each model,

D_{I}/10^{8}

, where

D_{I}=B\cdot\text{steps}

is the total images processed at the final (

400

TPP) checkpoint.

Figure 11: Timestep loss response and per-bucket scaling. Left: the per-timestep loss response curves

\mathcal{L}^{*}(v_{\theta},t)

(mean, top; std, bottom) at the final checkpoint for the Abra family. Right: each timestep bucket individually obeys a clean scaling law. Top two rows: scaling laws in EMA model loss for 10 timestep buckets. Bottom: the exponents

\alpha_{b}

and biases

\log A_{b}

of the supporting-hyperplane fits per bucket.

Figure 12: Non-cherry-picked generations from the Abra family on Parti prompts ( Yu et al., 2022 ) at a fixed guidance scale of $3.0$ . Columns are increasing model size ( Abra -60M $\to$ Abra -2B); rows are three prompts each of simple, standard, and complex complexity (word-count terciles). Each cell is a single $512\times 512$ generation.

Figure 13: Model size vs. CFG at the compute-optimal ( $200$ -TPP) checkpoint. Panels are two Parti prompts ( Yu et al., 2022 ) each of simple, standard, and complex complexity (rows). Within each panel, columns are increasing model size ( Abra -60M $\to$ Abra -2B) and rows are the guidance scale (CFG $\in\{2,3,6,10\}$ ). Each cell is a single $512\times 512$ generation.

Figure 14: Iso-FLOP relative loss penalty

\Delta L/L^{*}

versus TPP for (left) the Abra family and (right) the overtrained LLM grid of Gadre et al. (2025) . Markers denote the measured points ( Abra evaluated checkpoints; Gadre et al.’s per-run final losses); the connecting curves interpolate through them. Doubling the optimal TPP for the largest model incurs a

{\sim}4\%

loss penalty — nearly

20\times

that of the diffusion model.

Figure 15: Iso-FLOP suboptimality penalty

(L_{N}-L^{*})/L^{*}

versus TPP for the Abra family (left) and the Chinchilla / Hoffmann et al. data (right) , the latter figure-digitized by Besiroglu et al. (2024) and binned into log-

N

model-size bands. Markers are the measured points ( Abra evaluated checkpoints; digitized Chinchilla points); curves interpolate through them. The green/red bands mark the

0.2\%

penalty threshold.

Figure 16: Iso-FLOP suboptimality penalty

(L_{N}-L^{*})/L^{*}

versus TPP for the Abra family (left) and the LLM runs of Porian et al. (2024) on RefinedWeb (right) . Markers are the measured points ( Abra evaluated checkpoints; Porian et al.’s per-step logged validation losses along each size’s longest run); curves interpolate through them. The green/red bands mark the

0.2\%

penalty threshold.

Ladder	$a$ (constr.)	$a$ (free)	$b$ (free)	$a+b$ (free)
incl. 60M	$0.467$	$0.484$	$0.550$	$1.035$
excl. 60M	$0.470$	$0.487$	$0.547$	$1.033$

Appendix

Table 3: Compute-optimal allocation exponents (

N_{\text{opt}}\propto C^{a}

D_{\text{opt}}\propto C^{b}

) for the Abra envelope. The free fit regresses

\log N_{\text{opt}}

and

\log D_{\text{opt}}

\log C

independently; the constrained fit imposes

a+b=1

. The free

a+b>1

reflects the per-model transformer FLOP count exceeding the

6ND

estimate by a size-dependent factor; the balanced split is robust either way.

Estimator / variation	TPP* (incl. 60M)	TPP* (excl. 60M)	Notes
Iso-FLOP (main body)	$186$	$\mathbf{199}$	headline; $68\%$ eval-noise $[180,270]$
Kaplan (envelope OLS)	$185$	$192$	balanced allocation $a\approx 0.47$
Supporting hyperplane (LP)	$185$	$192$	same envelope allocation as Kaplan
Hoffmann A1 (curve minima)	$185$	$192$	noisy minima; wide p10–p90
Raw-loss target	$189$	$227$	only material mover; not used
Huber fit objective	$186$	$199$	identical to least squares

Appendix

Table 4: Compute-optimal

\mathrm{TPP}^{*}

under every estimator and procedure variation on identical Abra data, with 60M included and excluded (holding the other choice fixed); the headline (iso-FLOP, 60M excluded) is bold .

Method	Ladder	Alloc.	$A$	$\alpha$	$\beta$	$A_{N}$	$a$	$A_{D}$	$b$
Iso-FLOP	incl. 60M	free	$1.37$	$0.0371$	$0.3874$	$0.0211$	$0.503$	$0.978$	$0.533$
		$a{+}b{=}1$				$0.049$	$0.485$	$2.27$	$0.515$
		$a{=}b{=}0.5$				$0.0243$	$0.500$	$4.57$	$0.500$
	excl. 60M	free	$1.2$	$0.0140$	$0.0052$	$0.0126$	$0.514$	$1.58$	$0.523$
		$a{+}b{=}1$				$0.0301$	$0.495$	$3.78$	$0.505$
		$a{=}b{=}0.5$				$0.0239$	$0.500$	$4.75$	$0.500$

Appendix

Table 5: Full fit-coefficient comparison, the companion to Table 4 (which reports only the recovered

\mathrm{TPP}^{*}

). For each of the seven estimators and variations we report the compute-optimal loss frontier

L^{*}(C)=A\,C^{-\alpha}+\beta

together with the optimal-allocation laws

N_{\text{opt}}(C)=A_{N}\,C^{a}

and

D_{\text{opt}}(C)=A_{D}\,C^{b}

, each computed with the 60M model included and excluded, and under three allocation conventions: free (independent log–log fits,

a{+}b\approx 1.04

), constrained (

a{+}b{=}1

), and balanced (

a{=}b{=}0.5

Interpolant	Abra -60M	Abra -120M	Abra -250M	Abra -500M	Abra -1B
PCHIP (main body)	$0.0167$	$0.0171$	$0.0233$	$0.0317$	$0.0430$
Piecewise-linear	$0.0379$	$0.0406$	$0.0513$	$0.0672$	$0.0882$

Appendix

Table 6: Leave-one-out interpolation RMSE per model, for PCHIP and piecewise-linear interpolation; PCHIP is the more accurate and both are small.

Held out	TPP* (incl. 60M base)	TPP* (excl. 60M base)
none (full)	$186$	$199$
Abra -60M	$199$	—
Abra -120M	$170$	$188$
Abra -250M	$189$	$207$
Abra -500M	$185$	$198$
Abra -1B	$184$	$203$

Appendix

Table 7: Leave-one-model-out iso-FLOP

\mathrm{TPP}^{*}

: the estimate with each model held out of the ladder, for the 60M-included and 60M-excluded bases. With 60M excluded, every held-out estimate stays in

[188,207]

Metric	Backbone	$\mathrm{TPP}^{*}$	$\mathrm{TPP}^{*}(C)$ exponent $p$
KDD	DINOv2	$121$	$-0.06$
DMMD	DINOv2	$133$	$-0.13$
CLIPScore	CLIP	$153$	$-0.05$
FID	Inception	$182$	$+0.12$
KID	Inception	$257$	$+0.12$
CMMD	CLIP	$294$	$-0.14$

Appendix

Table 8: Compute-optimal tokens per parameter

\mathrm{TPP}^{*}

per generative metric, recovered with the headline iso-FLOP procedure (Section 4.1 ) applied to each metric’s best-over-CFG trajectory.

Figure 17: Each generative metric as a function of the classifier-free guidance scale, swept per model size. The optimum (a minimum for the distribution distances, a maximum for CLIPScore) shifts with model size, so comparing models at a single fixed guidance value would misrank them; this motivates the per-model CFG optimization we use throughout.

Figure 18: DINOv2-family generative metrics scale predictably. Top: generation quality vs training compute $C$ for the three DINOv2-based distances—FDD (Fréchet), KDD (kernel), and DMMD (MMD). At each guidance scale the evaluation trajectories are fit with the iso-FLOP compute-optimal procedure of Section 4.1 (curves colored by CFG); dark markers and curves are the best-over-CFG trajectory $R(t)$ per model size, and the bold red curve is its compute-optimal frontier. Bottom: the resulting compute-optimal $\mathrm{TPP}^{*}$ as a function of model size $N$ (linear fit through 120M–2B, 60M extrapolated), and the compute-optimal $N_{\text{opt}}$ and $D_{\text{opt}}$ as functions of compute $C$ . All three DINOv2 distances reach compute optimality at a substantially lower TPP than the loss.

Figure 19: Optimal CFG versus training compute

C=6ND

, one panel per metric, with each model size a different colour and every swept checkpoint plotted (not just the final one). The optimal CFG drifts downward with both model size and training compute for the distribution metrics.

Allocation law	Ours ( Abra )	Liang et al. (2024)
$N_{\text{opt}}(C)$	$0.0301\,C^{0.4951}$	$0.0009\,C^{0.5681}$
$D_{\text{opt}}(C)$	$3.776\,C^{0.5049}$	$186.85\,C^{0.4319}$

Appendix

Table 9: Compute-optimal allocation laws for Abra , Abra -60M excluded, versus Liang et al. (2024) .

C

in FLOPs,

N

in parameters,

D

in image tokens; both fits impose

a+b=1

. Note that our study measures compute in image tokens while Liang et al. (2024) measure in total tokens. While this makes the intercepts incomparable, the scaling exponents are comparable between the two studies.

T2I Diffusion Models

Text-to-image (T2I) diffusion models generate images from prompts by learning to reverse noise under text conditioning.

See all 403 papers

Scaling Beyond Masked Diffusion Language Models
Feb 16, 2026Subham Sekhar Sahoo, Jean-Marie Lemercier, Zhihan Yang +4Diffusion Language Models Diffusion Models
NVIDIA, Santa Clara, USA · Department of Computer Science, Cornell Tech, NYC, USA · Department of Computer Science, Cornell University, Ithaca, USA +2
i1: A Simple and Fully Open Recipe for Strong Text-to-Image Models
Jun 9, 2026Boya Zeng, Tianze Luo, Shu Pu +4Training Data Curation T2I Generation
Princeton University
Chimera: Designing and Chinchilla-Scaling Hybrid Visual Diffusion Transformers
Jul 30, 2026Chongjian Ge, Hanwen Jiang, Tianyu Wang +9Video Diffusion Models Diffusion Transformer
Adobe Research

Feb 16, 2026J/K move · Enter open · S save

Scaling Beyond Masked Diffusion Language Models

Subham Sekhar Sahoo, Jean-Marie Lemercier, Zhihan Yang, Justin Deschenaux+3

NVIDIA, Santa Clara, USA · Department of Computer Science, Cornell Tech, NYC, USA · Department of Computer Science, Cornell University, Ithaca, USA+2

Diffusion Language Models Diffusion Models Diffusion Language Model Training Diffusion Language Model Benchmarking Masked Diffusion Models Diffusion Language Model Inference

Open paper AskPDF arXiv

Diffusion language models are a promising alternative to autoregressive models due to their potential for faster generation. Among discrete diffusion approaches, Masked diffusion currently dominates, largely driven by strong perplexity on language modeling benchmarks. In this work, we present the first scaling law study of uniform-state and interpolating discrete diffusion methods. We also show that Masked diffusion models can be made approximately 12% more FLOPs-efficient when trained with a simple cross-entropy objective. We find that perplexity is informative within a diffusion family but can be misleading across families, where models with worse likelihood scaling may be preferable due to faster and more practical sampling, as reflected by the speed-quality Pareto frontier. These results challenge the view that Masked diffusion is categorically the future of diffusion language modeling and that perplexity alone suffices for cross-algorithm comparison. Scaling all methods to 1.7B parameters, we show that uniform-state diffusion remains competitive on likelihood-based benchmarks and outperforms autoregressive and Masked diffusion models on GSM8K, despite worse validation perplexity. We provide the code, model checkpoints, and video tutorials on the project page: http://s-sahoo.github.io/scaling-dllms

Abstract

T2I Diffusion Models

Explore similar work

Scaling Beyond Masked Diffusion Language Models

i1: A Simple and Fully Open Recipe for Strong Text-to-Image Models

Chimera: Designing and Chinchilla-Scaling Hybrid Visual Diffusion Transformers