ABRA: Scaling Diffusion Image Training

Summary (Overview)

• Core Finding: Text-to-image diffusion transformers achieve compute-optimal training at approximately 200 image tokens per parameter (TPP), which is 10× higher than the Chinchilla compute-optimal prescription for LLMs (20 TPP).

• Comprehensive Scaling Study: The authors introduce ABRA, a controlled family of dense flow-matching transformers (60M to 2B parameters), trained across three orders of magnitude of compute ($10^{19}$ to $10^{22}$ FLOPs)—an order of magnitude more compute than prior work.

• Overtraining Robustness: Diffusion models are remarkably robust to overtraining (less than 0.5% loss penalty when 2× overtrained), unlike LLMs, suggesting practitioners should err on the side of more data rather than larger models.

• Predictable Scaling Beyond Loss: Scaling laws extend to generative quality metrics (FID, KID, CLIPScore, CMMD), representation quality (linear probe accuracy), and training curve shapes, which exhibit scaling collapse—demonstrated for the first time in diffusion models.

• Resolution Effects: Higher resolution training requires more image tokens per parameter to reach compute optimality (165 TPP at 256×256 → 247 TPP at 768×768), though the number of images required actually decreases.

---

Introduction and Theoretical Foundation

Background

Compute-optimal scaling laws describe how a fixed training compute budget $C$ (in FLOPs) should be allocated between model size $N$ (in parameters) and dataset size $D$ (in training tokens) to minimize loss. For LLMs, the "Chinchilla rule" prescribes:

$$N \propto C^{1/2}, \quad D \propto C^{1/2}$$

resulting in ~20 tokens per parameter (TPP). This framework has become standard industrial practice for frontier language models.

Challenges for Visual Generation

Extending scaling law analysis to visual generation presents unique challenges:

• Higher dimensionality: 2D images have resolution-dependent token information density
• Separate training/generation paradigms: Training loss doesn't directly measure generative quality
• Noisier training curves: Batch-to-batch noise dominates per-step loss improvements
• Prior work limitations: Liang et al. (2024) only fit iso-FLOP profiles below $10^{19}$ FLOPs

Theoretical Foundation

The paper follows Hoffmann et al. (2022)'s definition of compute optimality:

$$N_{\text{opt}}(C), D_{\text{opt}}(C) = \arg\min_{N, D} \mathcal{L}(N, D; C) \tag{1}$$

where $N$ is the total parameter count (excluding frozen text encoder) and $D$ is the number of image tokens seen during training. The paper uses v-prediction flow-matching (Albergo & Vanden-Eijnden, 2023; Ma et al., 2024) with a latent diffusion formulation, where the network $v_\theta(x_t, t; \text{text})$ regresses the velocity of the linear interpolant between a clean image latent and Gaussian noise.

---

Methodology

Model Family (ABRA)

• Architecture: Dense flow-matching transformers from 60M to 2B parameters
• Structure: Four double-stream blocks followed by single-stream blocks, FFN expansion rate of 6
• Text conditioning: Frozen Qwen3-4B model (excluded from FLOP/parameter counts)
• Scaling strategy: Monotonically non-decreasing width-to-depth ratio; batch size doubles with parameter count while keeping training steps fixed

Training Setup

• Base resolution: $512 \times 512$ with five aspect-ratio buckets
• VAE: Frozen FLUX VAE (f=8) encoding images into $32 \times 32$ latent grids
• Text conditioning dropout: 10% probability for unconditional estimates at inference
• Optimizer: Adam with $\mu$P (Yang et al., 2021) using Complete-P parameterization
• Learning rate: $4 \times 10^{-4}$ (zero-shot transferred across model sizes)
• Initialization: $\sigma = 0.01$ for all layers
• Timestep sampler: Logit-normal with mean $p_\mu = 1.9$ and standard deviation $p_\sigma = 1.0$
• EMA: Decay parameter 0.9999, kept constant across model sizes

Dataset

• Primary: DataComp-1B (Gadre et al., 2023), a large-scale web image-text collection
• Re-captioning: Images re-captioned using open-source vision-language models with a mix of dense, medium, sparse, and web-sourced captions
• Evaluation: Held-out mixture more diverse than training data (DataComp-1B is ~20% subset)

Scaling Law Fitting

The authors interpolate evaluation checkpoints into continuous loss-versus-TPP trajectories using PCHIP interpolation, then fit power laws:

$$L(C; T) = A_T C^{-\alpha_T} + \beta_T$$

Fixed-C slices give iso-FLOP profiles; the arg min over $T$ gives compute-optimal TPP.

Sampling

• ODE integration: $\frac{dx_t}{dt} = v_\theta(x_t, t; \text{text})$ integrated backwards from $\varepsilon \sim \mathcal{N}(0, 1)$ at $t = 1$
• Solver: 50-step Euler solver
• CFG: Classifier-free guidance with scale swept per model and held fixed across sampling steps

---

Empirical Validation / Results

4.1 Compute-Optimal Tokens per Parameter

Key Finding: Compute-optimality occurs at ~200 TPP, roughly 10× higher than the LLM prescription of 20 TPP.

• The minima of iso-FLOP curves tightly cluster around 200 TPP at all but the smallest training budgets
• ABRA-60M was excluded as an outlier (scaling laws hold only in appropriate regimes)
• The result is robust: across every fit procedure tested, optimal TPP lies within ±17 TPP

4.2 Diffusion Overtraining is Forgiving

The loss penalty for non-optimal training is defined as:

$$\Delta L(N, C) := L(N, C) - L(N_{\text{opt}}(C), C) \tag{2}$$

Key Findings:

• 2× overtraining: Less than 0.5% loss penalty (nearly "free")
• Undertraining: Catastrophic for diffusion models
• Practical rule: In compute-limited settings, err on the side of overtraining a small model rather than undertraining a large one

4.3 Generative Metrics Scale Predictably

All tested metrics follow power laws: $\mathcal{M}(C) = A_M C^{-\alpha} + F_M$ (sign flipped for CLIPScore), but their compute-optimal allocations diverge:

Key Findings:

• No single metric can define compute optimality
• Optimal CFG decreases with both model size and training steps (corroborating Mei et al., 2025)

4.4 Linear Probe Accuracy (Representation Quality)

• Peak accuracy: Middle layers, moderately high diffusion timestep $t \approx 0.7$
• Scaling: Linear-probe accuracy follows tight scaling laws at fixed TPP
• Compute-optimal TPP for understanding is much lower than for generation
• "Double-point": Both generation and understanding are jointly compute-optimal at ~$5 \times 10^{22}$ FLOPs with a 6B parameter model

4.5 Scaling Collapse

Following Qiu et al. (2025), the rescaled loss is:

$$\ell(x, p; \theta) = \frac{L(xt^*(p), p) - \hat{L}}{L(t^*(p), p) - \hat{L}} \tag{3}$$

with collapse deviation $\Delta$ and noise floor $\sigma$:

$$\Delta(x) := \frac{\mathbb{V}_{p, \theta(p)}[\ell(x, p; \theta)]^{1/2}}{\mathbb{E}_{p, \theta(p)}[\ell(x, p; \theta)]}, \quad \sigma(x; p) := \frac{\mathbb{V}_{\theta(p)}[\mathcal{L}(xt^*(p), p, \theta(p))]^{1/2}}{\mathbb{E}_{\theta(p)}[\mathcal{L}(xt^*(p), p, \theta(p))]}$$

Key Finding: Training curves collapse onto a universal form within the first 15% of training, with cross-size deviation $\Delta$ at or below the noise floor $\sigma$—demonstrating scaling collapse for the first time in diffusion models.

4.6 Effect of Resolution

Key Findings:

• Higher resolution requires more image tokens per parameter
• However, the number of images required actually decreases with resolution
• High-resolution training is compute-bound, not data-bound

---

Theoretical and Practical Implications

Theoretical Implications

1. Data-Intensive Nature of Visual Generation: The 10× higher TPP requirement (200 vs 20) suggests that image generation models have fundamentally different data efficiency properties than language models, possibly due to the higher information density and continuous nature of visual data.

2. Heterogeneous Scaling Across Capabilities: The divergence in compute-optimal allocations between generation metrics (FID/KID vs CLIPScore/CMMD) and between generation vs. understanding capabilities challenges the notion of a single "optimal" training configuration.

3. Scaling Collapse Generalization: The demonstration of scaling collapse in diffusion models extends this phenomenon beyond MLPs and LLMs, suggesting it may be a universal property of compute-optimal training across modalities.

Practical Implications

1. Actionable Training Rule: The 200 TPP rule provides practitioners with a concrete guideline for allocating compute between model size and data volume when training text-to-image diffusion models.

2. Risk Management: The robustness to overtraining means practitioners can safely train smaller models on more data, reducing inference costs without significant quality penalties.

3. Resolution-Specific Guidance: The finding that higher resolutions require more tokens per parameter (but fewer images) informs data collection strategies for high-resolution training pipelines.

4. CFG Scheduling: The discovery that optimal CFG decreases with model size and training steps provides guidance for inference-time hyperparameter selection.

---

Conclusion

This paper presents the most comprehensive scaling law study for text-to-image diffusion models to date, spanning three orders of magnitude of compute ($10^{19}$ to $10^{22}$ FLOPs) with a controlled model family (ABRA, 60M–2B parameters).

Main takeaways:

1. Text-to-image diffusion transformers scale predictably with compute, following tight power laws across loss, generative metrics, and representation quality
2. Compute optimality occurs at ~200 TPP, an order of magnitude higher than LLMs
3. Overtraining is forgiving for diffusion models, making it safe to prioritize data over model size
4. Scaling collapse extends to diffusion models, providing a diagnostic tool for compute-optimal training
5. Resolution affects optimal allocation: higher resolutions require more tokens per parameter but fewer images

Future directions suggested by the authors include:

• Investigating how generative metrics correlate with human preferences
• Understanding the heterogeneous capacity demands of different metrics
• Exploring the "double-point" where generation and understanding are jointly compute-optimal
• Extending these analyses to video generation and other modalities