PRX Part 3 — Training a Text-to-Image Model in 24h

TL;DR

Photoroom’s PRX Part 3 stacks the best ideas from Parts 1 and 2 into a 24-hour text-to-image speedrun on 32 H200s for ~$1500 total compute. The recipe is pixel-space x-prediction (no VAE) with patch-32 tokenization, LPIPS + DINOv2 perceptual losses, TREAD token routing (50% of tokens skip blocks 2 → penultimate), REPA alignment to a DINOv3 teacher, and the Muon optimizer on 2D parameters with Adam on the rest. Training is 100k steps at 512px (batch 1024) followed by 20k steps at 1024px (batch 512, REPA off). The accompanying open-source PRX repo bundles both this run and the Part 1/2 experimental framework, positioning the recipe as the foundation for Photoroom’s larger-scale future runs.

Key claims

Stacking pixel-space x-prediction (Back to Basics, arXiv:2511.13720) with patch-32 + a 256-d projection bottleneck keeps the 512px sequence at 64 tokens and 1024px at 256 tokens, making VAE-free training computationally manageable on a 24h H200 budget [§The Training Recipe → X-prediction and Training in the Pixel Space].
Predicting pixels directly lets classical perceptual losses (LPIPS at weight 0.1; DINOv2 perceptual at weight 0.01) be applied on the predicted clean image at all noise levels, on pooled full images rather than patch-wise features — both deviations from PixelGen’s exact recipe but consistently better in the authors’ setup [§Perceptual Losses].
TREAD token routing is preferred over SPRINT for this setup: 50% of tokens skip from block 2 to the penultimate block, with the extra SPRINT complexity not paying off because TREAD’s sequence-length advantage (64 vs 128 at 512px) is already small [§Token Routing with TREAD].
Routed models look worse under vanilla CFG when undertrained, so PRX uses a self-guidance scheme from Krause et al. 2025 (arXiv:2601.01608) that guides with a dense-vs-routed conditional pair instead of a conditional-vs-unconditional pair [§Token Routing with TREAD].
REPA alignment is computed only on non-routed tokens at the 8th transformer block with loss weight 0.5, using DINOv3 as the teacher because it gave the best quality in prior PRX experiments [§Representation Alignment with REPA and DINOv3].
Muon (FSDP implementation from muon_fsdp_2) on 2D parameters with lr=1e-4, momentum=0.95, nesterov=true, ns_steps=5, plus Adam on biases/norms/embeddings with lr=1e-4, betas=(0.9, 0.95), is the optimizer split, motivated by clear improvement over Adam-only in prior PRX runs [§Optimizer: Muon, table].
Training schedule: 512px for 100k steps at batch 1024, then 1024px fine-tune for 20k steps at batch 512 without REPA, with EMA smoothing=0.999, update_interval=10ba for sampling and eval [§Training Settings].
Datasets are three publicly available synthetic corpora — Flux-generated (1.7M), FLUX-Reason-6M (6M), and midjourney-v6-llava (1M) — with the midjourney captions re-captioned through Gemini 1.5 for prompt consistency [§Training Settings].
The authors characterize the resulting failure modes (texture glitches, occasional weird anatomy, hard-prompt fragility) as undertraining + limited data diversity rather than structural flaws in the recipe, framing further scaling as predictable improvement [§Results and Closing Thoughts].

Method

PRX Part 3 is a single integration recipe rather than a new method. The model is a DiT trained directly on pixels under x-prediction (predict $\hat{x}_0$ ), with input tokenization via patch-32 + a 256-d projection bottleneck to keep sequence length manageable (64 tokens at 512px, 256 at 1024px). On top of the standard flow-matching x-prediction loss the authors add LPIPS (weight 0.1) and a DINOv2-based perceptual loss (weight 0.01) on the predicted clean image, applied to pooled full images at all noise levels — a tweak vs PixelGen’s patch-wise application. To cut per-step compute, TREAD routes 50% of tokens from block 2 to the penultimate block; because vanilla CFG degrades on routed models when undertrained, the authors use a Krause-et-al.-2025 self-guidance scheme (dense vs routed conditional). REPA representation alignment to a DINOv3 teacher is applied once at the 8th block with loss weight 0.5, only on the non-routed token subset so the alignment signal is consistent. Optimization splits parameter groups: Muon (FSDP via muon_fsdp_2) for 2D matrices, Adam for biases / norms / embeddings. The schedule is 100k steps at 512px (batch 1024) → 20k steps at 1024px (batch 512, REPA disabled), with EMA-of-weights for sampling. Datasets are three public synthetic corpora, with midjourney captions re-written by Gemini 1.5.

Results

The run is qualitative rather than benchmarked: the post shows training-curve plots and a sample grid from the final EMA checkpoint, with no FID / GenEval / DPG-Bench numbers given. Headline framings:

A ~ $1500 budget (32 H200 × ~24h ×$ 2/GPU-hour) produces a “clearly usable” T2I model with strong prompt adherence and consistent aesthetic, with the 1024px stage sharpening details without breaking composition [§Results and Closing Thoughts].
Residual artifacts (texture glitches, occasional anatomy errors, fragility on hard prompts) are attributed to undertraining and limited dataset diversity rather than recipe flaws — implying scaling should be predictable [§Results and Closing Thoughts].
Sample diffusion-side compute is normalized through TREAD-50% routing and the self-guidance variant (no separate unconditional pass for guidance) [§Token Routing with TREAD].

The acknowledgements explicitly position this run inside a small cohort of speedrun-style efforts: Nitro-T (AMD ROCm blog), Speedrunning ImageNet Diffusion (Bhanded, arXiv:2512.12386), Stretching Each Dollar (Sehwag et al., arXiv:2407.15811), and Home-made Diffusion (Yeh, arXiv:2509.06068) [§Acknowledgements].

Why it’s interesting

PRX Part 3 is the first wiki-filed integration recipe for the cluster of training-efficiency ideas that have been showing up separately: it stacks PixelGen-style perceptual losses (PixelGen: Improving Pixel Diffusion with Perceptual Supervision) on top of REPA/DINO-teacher alignment (Improved Baselines with Representation Autoencoders) plus Muon optimization, then layers TREAD token routing on the same DiT — answering an open question on Pixel-space diffusion about whether these recipes compose in a single training run. The choice to deviate from PixelGen on perceptual-loss application (pooled-full-image, all-noise-levels) is a concrete data point for that page’s “where exactly to apply DINO” thread. It’s also a useful counterpart to L2P: Unlocking Latent Potential for Pixel Generation (L2P), which gets to pixel space via transfer from an LDM on synthetic data — PRX gets there from-scratch under a strict 24h budget, on three public synthetic datasets, so the two together bracket the “cheap pixel T2I” design space. The self-citation of Swayam’s own Speedrunning ImageNet Diffusion in the acknowledgements (the same researcher who posted this link) makes it likely worth filing 2512.12386 next.