PRX Part 4: Data Strategy for Diffusion Model Training

Photoroom’s PRX Part 4 details the data pipeline behind their diffusion model training. The guiding principle is that pre-training needs breadth and diversity, not per-image perfection; fine-tuning later handles taste. They assemble a mix of public and internal datasets, then re-caption every image with a VLM for consistency. Captions are long, faithful paragraphs that describe everything visible, turning potentially undesirable elements into controllable attributes. Images are stored as JPEG at quality 92 after verifying that repeated compression causes negligible degradation (after one re-encode: PSNR 48.7 dB for high-res images, LPIPS 0.004), and that training on JPEG vs PNG yields indistinguishable output quality. The data pipeline uses two complementary formats: Lance for interactive exploration, indexing, and filtering (with a caution against over-fragmentation; they learned to keep fragment count low), and Mosaic Data Shards (MDS) for streaming during training. MDS provides deterministic shuffling, elastic mid-epoch resume, and weighted mixing. Text latents are computed on the fly instead of pre-computed, costing only ~3–4% throughput (about one extra day per 30-day run) but keeping shards small and allowing encoder changes without rewriting. Aspect-ratio bucketing snaps images to 13 shapes per resolution tier (512, 1024, 2048, 4096 px) under a pixel budget of ~256 patches, keeping compute per image flat.

For captioning, they benchmarked three VLMs (Qwen2.5-VL-7B-Captioner-Relaxed, Qwen3-VL-8B, Qwen3.5-9B) plus two references by training small diffusion models for 100k steps. Qwen3.5-9B led on all metrics (FID 10.51, CMMD 0.278, DINO-MMD 0.162), but they chose Qwen3-VL-8B for its balance of quality (FID 10.98, DINO-MMD 0.182) and throughput (20 images/second per H200) with stable vLLM support. Captioning runs as a Ray Data streaming pipeline.

Data filtering is deliberately light. After re-captioning, they use Qwen3-8B in text-only mode (200 captions/s per GPU) to label samples as visual, text, or NSFW. Rather than rewriting the corpus, they maintain per-shard skip lists that the loader unions at training time, a flexible mechanism that also supports user opt-outs and ablation studies. Deduplication uses a perceptual hash (DCT-based) with zero Hamming distance threshold, removing only near-pixel-identical copies; cross-resolution duplicates keep the highest-resolution version. Deduplication removed a few percent of images, the text filter a few percent more, and NSFW a tiny fraction. Overall, the pipeline builds a “solid, lightweight” pre-training corpus for a 7B model, with the understanding that future fine-tuning and preference alignment will use smaller, ruthlessly curated sets.

PRX Part 4: Our Data Strategy

View Original