PRX Data Strategy: Scaling Pre-training with VLM Re-captioning and Mosaic Streaming

Photoroom has released a detailed breakdown of the data strategy used to train PRX, a 7B parameter model. The core approach involves assembling a diverse corpus from public and internal sources, re-captioning all images using a Vision Language Model (VLM) to ensure high-fidelity descriptions, and utilizing a dual-format storage system to balance dataset exploration with training throughput.

Pre-training Principles: Breadth Over Perfection

For the pre-training phase, Photoroom prioritized coverage and diversity over individual image aesthetics. The goal was to teach the model visual concepts, composition, and lighting across a broad distribution of the visual world.

• Diversity vs. Taste: The team argues that over-filtering for aesthetics during pre-training narrows the distribution and costs the model essential compositional variety. Polishing generations is deferred to fine-tuning and preference alignment on smaller, curated sets.
• Pragmatic Sourcing: Data was assembled from a mix of public and internal datasets, leveraging existing curation (NSFW and PII filtering) where available to accelerate the process.
• Captioning Philosophy: Long, accurate captions are critical. By describing everything in an image—including logos, text, or artifacts—these elements become conditioned, controllable attributes rather than unconditional noise.

Data Infrastructure: Lance and Mosaic Data Shards (MDS)

Photoroom employs two distinct data formats to handle the lifecycle of the dataset: Lance for curation and MDS for streaming.

Lance for Exploration

Lance is used as a columnar data format for building and exploring datasets with billions of rows. It supports cheap predicate pushdown, scalar indexes, and vector search, allowing the team to:

• Profile Data: Analyze resolution distributions to set cutoffs (e.g., discarding images below 384² pixels).
• Interactive Browsing: Use full-text search and nearest-neighbor similarity search via a custom UI to qualitatively assess diversity and spot issues like non-photographic content.
• Fragment Management: The team noted that over-fragmenting Lance tables (e.g., 100k rows per fragment) slows queries. Compacting to roughly one million rows per fragment significantly improved performance.

MDS for Training

Mosaic Data Shards (MDS) are used for distributed training, providing deterministic resumable shuffling, elastic mid-epoch checkpointing, and weighted mixing of multiple streams. To optimize the pipeline:

• On-the-fly Latents: Instead of pre-computing text latents, Photoroom computes them during training using the Qwen3-VL encoder. This reduced MDS shard size and allowed for encoder swaps without rewriting terabytes of data, costing only a 3–4% throughput hit at the 7B scale.
• JPEG Encoding: All images are stored as JPEG at quality 92. Empirical tests showed that this is perceptually indistinguishable from PNG and does not negatively impact the quality of generated images, while reducing storage requirements by 3–10×.

VLM Re-captioning Strategy

Photoroom found that long, detailed captions significantly improve sample quality. In benchmarks, a small diffusion model trained on Qwen2.5-VL-7B captions outperformed one trained on shorter LLaVA-1.5-LLaMA3-8B captions across FID, CMMD, and DINO-MMD metrics.

Captioner Selection

After benchmarking several candidates, Photoroom selected Qwen3-VL-8B as the primary captioner. The selection was based on a balance of quality and throughput:

• Performance: Qwen3.5-9B performed best on metrics, but Qwen3-VL-8B was close and significantly faster (20 img/s per H200 GPU).
• Prompting: The team used a dense system prompt requiring a single flowing paragraph of 100–200 words, focusing on precise visual grounding without speculation or interpretation.

Dataset Refinement and Deduplication

Once VLM captions were generated, Photoroom applied light filtering and deduplication to refine the corpus without rewriting the entire dataset.

Caption-Based Filtering

Instead of using expensive pixel-level classifiers, Photoroom used Qwen3-8B in text-only mode to classify captions as "visual," "text," or "nsfw." This approach is computationally efficient (″200 captions/s per GPU) and effectively removes screenshots, documents, and infographics.

Perceptual Hash Deduplication

To remove near-duplicates, the team implemented a DCT-based perceptual hash. If two images have a Hamming distance of zero, they are considered duplicates. In cases where a duplicate exists across different resolutions, the highest-resolution copy is retained.

The Skip-List Mechanism

To avoid rewriting the massive MDS corpus for every filter or deduplication pass, Photoroom implemented a skip-list feature in the data loader. A sidecar file per shard lists indices to skip, allowing the team to exclude data (e.g., for user opt-outs or quality issues) dynamically at load time.

Aspect-Ratio Bucketing

To maintain constant compute per image and avoid wasteful padding or cropping, Photoroom uses aspect-ratio bucketing:

1. Resolution Tiers: Images are assigned to tiers (512px, 1024px, 2048px, 4096px) based on pixel count.
2. Aspect Ratio Buckets: Within each tier, images are snapped to one of 13 patch-aligned shapes (AR 0.5 to 2.0) that maintain a constant patch token budget.
3. Processing: Images are resized using a Lanczos filter and stored in a directory tree keyed by resolution and aspect ratio for separate streaming.