Training Design for Text-to-Image Models: Lessons from Ablations

Summary

The second part of a series on training efficient text-to-image models from scratch documents training-design ablations for the PRX model. The authors establish a Flow Matching baseline (PRX-1.2B, 100K steps, 256×256, AdamW, GemmaT5, RoPE, 1M MidJourneyV6 synthetic images) and benchmark interventions against FID, CMMD, DINO-MMD, and throughput across four buckets: Representation Alignment, Training Objectives, Token Routing and Sparsification, and Data. REPA representation alignment boosts early convergence when used as a burn-in, while the Muon optimizer lowered FID from 18.20 to 15.55 and improved CMMD and DINO-MMD. Storing denoiser weights in BF16 instead of using BF16 autocast raised FID to 21.87, and token routing methods such as TREAD/SPRINT were minor at 256² but became major wins at higher resolutions. Long captions proved critical, synthetic versus real data shifted texture versus structure, and upcoming work releases the full training code alongside a 24-hour end-to-end speedrun.

Takeaways

1. 1

   REPA representation alignment accelerates early PRX training when used as a burn-in before being turned off, while better latents such as REPA-E and FLUX2-AE deliver a large quality jump with explicit speed trade-offs.

2. 2

   The Muon optimizer lowered FID from 18.20 to 15.55, CMMD from 0.41 to 0.36, and DINO-MMD from 0.39 to 0.35 against the AdamW baseline on PRX-1.2B at 256×256 resolution.

3. 3

   Storing denoiser weights in BF16 rather than using BF16 autocast while keeping parameters in FP32 raised FID to 21.87 and regressed CMMD and DINO-MMD, demonstrating that numerically sensitive modules such as normalization, attention, RoPE, and optimizer state must remain in higher precision.