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

Welcome back! This is the second part of our series on training efficient text-to-image models from scratch.

In the first post of this series, we introduced our goal: training a competitive text-to-image foundation model entirely from scratch, in the open, and at scale. We focused primarily on architectural choices and motivated the core design decisions behind our model PRX. We also released an early, small (1.2B parameters) version of the model as a preview of what we are building (go try it if you haven't already 😉).

In this post, we shift our focus from architecture to training. The goal is to document what actually moved the needle for us when trying to make models train faster, converge more reliably, and learn better representations. The field is moving quickly and the list of “training tricks” keeps growing, so rather than attempting an exhaustive survey, we structured this as an experimental logbook: we reproduce (or adapt) a set of recent ideas, implement them in a consistent setup, and report how they affect optimization and convergence in practice. Finally, we do not only report these techniques in isolation; we also explore which ones remain useful when combined.

In the next post, we will publish the full training recipe as code, including the experiments in this post. We will also run and report on a public "speedrun" where we put the best pieces together into a single configuration and stress-test it end-to-end. This exercise will serve both as a stress test of our current training pipeline and as a concrete demonstration of how far careful training design can go under tight constraints. If you haven’t already, we invite you to join our Discord to continue the discussion. A significant part of this project has been shaped by exchanges with community members, and we place a high value on external feedback, ablations, and alternative interpretations of the results.

Before introducing any training-efficiency techniques, we first establish a clean reference run. This baseline is intentionally simple. It uses standard components, avoids auxiliary objectives, and does not rely on architectural shortcuts or tricks to save compute resources. Its role is to serve as a stable point of comparison for all subsequent experiments.

Concretely, this is a pure Flow Matching (Lipman et al., 2022) training setup (as introduced in Part 1) with no extra objectives and no architectural speed hacks. We will use the small PRX-1.2B model we presented in the first post of this series (single stream architecture with global attention for the image tokens and text tokens) as our baseline and train it in Flux VAE latent space, keeping the configuration fixed across all comparisons unless stated otherwise.

The baseline training setup is as follows:

| Setting | Value | |---|---| | Steps | 100k | | Dataset | Public 1M synthetic image generated with MidJourneyV6 | | Resolution | 256×256 | | Global batch size | 256 | | Optimizer | AdamW | | lr | 1e-4 | | weight_decay | 0.0 | | eps | 1e-15 | | betas | (0.9, 0.95) | | Text encoder | GemmaT5 | | Positional encoding | Rotary (RoPE) | | Attention mask | Padding mask | | EMA | Disabled |

This baseline configuration provides a transparent and reproducible anchor. It allows us to attribute observed improvements and regressions to specific training interventions, rather than to shifting hyperparameters or hidden setup changes. Throughout the remainder of this post, every technique is evaluated against this reference with a single guiding question in mind:

Does this modification improve convergence or training efficiency relative to the baseline?

Examples of baseline model generations after 100K training steps.

To keep this post grounded, we rely on a small set of metrics to monitor checkpoints over time. None of them is a perfect proxy for perceived image quality, but together they provide a practical scoreboard while we iterate.

Fréchet Inception Distance (FID):(Heusel et al., 2017) Measures how close the distributions of generated and real images are, using Inception-v3 feature statistics (mean and covariance). Lower values typically correlate with higher sample fidelity.CLIP Maximum Mean Discrepancy (CMMD):(Jayasumana et al., 2024) Measures the distance between real and generated image distributions using CLIP image embeddings and Maximum Mean Discrepancy (MMD). Unlike FID, CMMD does not assume Gaussian feature distributions and can be more sample-efficient; in practice it often tracks perceptual quality better than FID, though it is still an imperfect proxy.**DINOv2 Maximum Mean Discrepancy (DINO-MMD):**Same MMD-based distance as CMMD, but computed on DINOv2 (Oquab et al. 2023) image embeddings instead of CLIP. This provides a complementary view of distribution shift under a self-supervised vision backbone.**Network throughput:**Average number of samples processed per second (samples/s), as a measure of end-to-end training efficiency.

With the scoreboard defined, we can now dive into the methods we explored, grouped into four buckets: Representation Alignment, Training Objectives, Token Routing and Sparsification, and Data.

Diffusion and flow models are typically trained with a single objective: predict a noise-like target (or vector field) from a corrupted input. Early in training, that one objective is doing two jobs at once: it must build a useful internal representation and learn to denoise on top of it. Representation alignment makes this explicit by keeping the denoising objective and adding an auxiliary loss that directly supervises intermediate features using a strong, frozen vision encoder. This tends to speed up early learning and bring the model’s features closer to those of modern self-supervised encoders. As a result, you often need less compute to hit the same quality.

A useful way to view it is to decompose the denoiser into an implicit encoder that produces intermediate hidden states, and a decoder that maps those states to the denoising target. The claim is that representation learning is the bottleneck: diffusion and flow transformers do learn discriminative features, but they lag behind foundation vision encoders when training is compute-limited. Therefore, borrowing a powerful representation space can make the denoising problem easier.

REPA (Yu et al., 2024)

REPA adds a representation matching term on top of the base flow-matching objective.

Let be a clean sample and be the noise sample. The model is trained on an interpolated state (for ) and predicts a vector field . In REPA, a pretrained vision encoder processes the clean sample to produce patch embeddings , where is the number of patch tokens and is the teacher embedding dimension. In parallel, the denoiser processes and produces intermediate hidden tokens (one token per patch). A small projection head maps these student hidden tokens into the teacher embedding space, and an auxiliary loss maximizes patch-wise similarity between corresponding teacher and student tokens: Here indexes patch tokens, is the teacher embedding for patch , is the corresponding student hidden token at time , and is typically cosine similarity.

This term is combined with the main flow-matching loss:

with controlling the trade-off.

In practice, the student is trained to produce noise-robust, data-consistent patch representations from , so later layers can focus on predicting the vector field and generating details rather than rediscovering a semantic scaffold from scratch.

We ran REPA on top of our baseline PRX training, using two frozen teachers: DINOv2 and DINOv3 (Siméoni et al., 2025). The pattern was very consistent: adding alignment improves quality metrics, and the stronger teacher helps more, at the cost of a bit of speed.

| Method | FID ↓ | CMMD ↓ | DINO-MMD ↓ | batches/sec ↑ | |---|---|---|---|---| | Baseline | 18.2 | 0.41 | 0.39 | 3.95 | | REPA-Dinov3 | 14.64 | 0.35 | 0.3 | 3.72 |