---
title: "Training Design for Text-to-Image Models: Lessons from Ablations"
description: "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."
---

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

[Hugging Face](https://yomu.fyi/company/hugging-face) · David Bertoin, Roman Frigg, Jon Almazán · Feb 3, 2026

**Type:** Benchmark

## 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.

## Context

To document what actually moved the needle when training the PRX text-to-image model to train faster, converge more reliably, and learn better representations, rather than attempting an exhaustive survey of the growing list of training tricks.

## Approach / What changed

Structured as an experimental logbook that reproduces or adapts recent training-efficiency ideas in a consistent Flow Matching baseline (PRX-1.2B, 100K steps, 256×256, AdamW, GemmaT5, RoPE, 1M MidJourneyV6 synthetic images), benchmarks each technique in isolation against FID, CMMD, DINO-MMD, and throughput, and reports which interventions remain useful when combined.

## Takeaways

- 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.
- 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.
- 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.

**Tags:** [Machine Learning](https://yomu.fyi/topic/machine-learning), [Open Source](https://yomu.fyi/topic/open-source), [Performance](https://yomu.fyi/topic/performance)

- Source: [Hugging Face](https://huggingface.co/blog/Photoroom/prx-part2)
- Source URL: https://huggingface.co/blog/Photoroom/prx-part2
- Ingested by Yomu: 2026-08-28T08:00:45.618Z

[Read original post](https://huggingface.co/blog/Photoroom/prx-part2)