PRX Part 4: Our Data Strategy

Welcome back! This is Part 4 of the PRX series. Parts 1 to 3 covered model architectures, training design, and a 24-hour speedrun. This time we're pulling back the curtain on the part that quietly underpins all of it: the data. Of all the things that shaped PRX's quality, the data pipeline was one of the least glamorous parts to build but nevertheless an important piece to get right. Here's what we did, what we'd do differently, and a few things we only learned the slow way.

In one sentence: we assemble training data from a mix of public and internal datasets, re-caption the images with a VLM, and turn the result into the streamable corpus we trained PRX on.

At a high level, the data pipeline looks like this:

In the following we will dive into it in detail.

1. Guiding principles

A diverse dataset for pre-training

The goal was to assemble a large, diverse dataset for pre-training. At this stage the model is learning how the world looks: the visual concepts, the objects and scenes, how things are composed and lit, and the sheer range of what images can contain. That is a problem of coverage and diversity, not of per-image perfection. A broad, representative corpus teaches the model far more about the structure of the visual world than a smaller, prettier one would, even if many of the individual images are ordinary snapshots or slightly compressed. Over-filtering for aesthetics at this stage would actually hurt, narrowing the distribution and costing the model concepts and compositional variety it cannot recover later. Making generations look polished is a separate, later concern, which we leave to fine-tuning and preference alignment on small, ruthlessly curated sets. Pre-training is for breadth; fine-tuning is for taste.

A mix of data sources

We assemble our pre-training data from a mix of public and internal datasets. The priority at this stage is breadth, diversity, and standing on curation that already exists rather than redoing it ourselves. Where a source already comes quality-filtered, deduplicated, and filtered for NSFW content and personal information, we build on that work instead of repeating it at scale. Sources arrive in different shapes: some come with the image data itself, others as metadata plus baseline captions that we bring into a common form. We took a pragmatic approach: rather than build a corpus entirely from scratch, we leaned on existing datasets and our own tooling to assemble one quickly. In hindsight it is not necessarily the absolute best dataset one could build, but it was a solid and lightweight starting point for pre-training a 7B model.

Our captions philosophy

In our experience, what matters most for pre-training is to use long captions that accurately describe everything in the image. We saw this directly in Part 2, where switching from short captions to long ones substantially improved sample quality. If captioning is faithful, we don't need to worry about the occasional screenshot, advertisement, logo, or bit of text in an image, because those things get described in the caption too, so the model learns them as conditioned, controllable attributes rather than reproducing them unconditionally. Accurate captioning turns "noise" into something you can prompt for, or prompt away. This is precisely why the filtering we do later is deliberately light. We remove what is genuinely unusable, not everything that is imperfect.

Data formats

We have been using Mosaic Streaming and Mosaic Data Shards (MDS) as a dataset format for distributed training for a while now. In combination with Mosaic Composer we found it to be a very low-maintenance, flexible, and well-performing framework for distributed training. Furthermore, MDS datasets can easily and effectively be mixed and shuffled and also allow for distributed training directly from object storage like S3 or GCS.

However, MDS datasets are very rigid. Adding a column or creating a subset for a given filter basically means that one has to scan and rewrite the whole dataset. That's why we use Lance for this kind of feature engineering and dataset curation. Lance is a columnar data format with cheap predicate pushdown, scalar indexes, and vector search, the right tool for building and exploring datasets with billions of rows.

These two formats play off each other throughout this post and the PRX data pipeline: Lance to build, MDS to stream.

On text latents

For previous training runs, we used T5Gemma as the text encoder and pre-computed the text latents, storing them in MDS as bytes. This time around, after switching the text encoder to Qwen3-VL, we decided to compute text latents on the fly during training instead. Running the text encoder inside the training loop costs throughput, but how much depends on the model: for a small denoiser it can be significant, whereas at PRX's 7B scale the text encoder's compute is negligible next to the denoiser, and we measured only a roughly 3–4% throughput cost (about 1 extra day on a 30-day run). In return we get two things. Skipping pre-computation keeps our MDS shards much smaller, small enough to store the full pre-training dataset on the SSD-backed shared filesystem of our SLURM cluster rather than streaming it over the network from object storage. And it keeps us free to change the text encoder later without rewriting terabytes of stored latents, which is exactly the kind of switch we made when moving to Qwen3-VL.

On image encoding

We encoded all images as JPEG at quality 92, instead of a lossless format such as PNG. We did not just assume quality 92 was safe, we measured it. Real-world images have usually already been JPEG-compressed several times, so the real question is whether one more re-encode hurts. Across repeated decode/encode cycles on both high-resolution (1–2 MP) and lower-resolution (0.25–0.5 MP) real images, the first re-encode at quality 92 is essentially imperceptible. Each further cycle adds almost nothing, because JPEG quickly converges to a stable state, and even after 10 cycles the images stay well within the imperceptible range, while PNG would be 3–10× larger for no perceptual gain. The numbers, averaged over 100 images at quality 92 and measured against the original (higher PSNR and lower LPIPS are better):

Since most source images already arrive JPEG-compressed, storing them losslessly as PNG buys little, so we convert everything to JPEG at high quality (quality 92).

We also checked what matters most: whether training on JPEGs changes what the model produces. Our corpus is mostly JPEG anyway, so the only advantage PNG could offer is not introducing new artifacts. We deliberately drew the comparison from high-resolution sources, on the expectation that even originally-JPEG images carry few artifacts at high resolution. We trained two identical PRX models at 1024px on the same images, one stored as PNG and one as JPEG at quality 92. Both trained equally well on our metrics, and their generations were practically indistinguishable, including when we estimated how JPEG-compressed the outputs look using a known technique for estimating JPEG quality by matching quantization tables:

The detection rate is the fraction of generations in which any quantization structure is found at all. The mean and median estimated JPEG quality are taken over only those flagged images. The two models are practically indistinguishable: only about one generation in ten from either shows any detectable quantization structure, and the differences are small enough that we are confident the training image format has a minimal effect on output quality. So high-quality JPEG storage adds nothing measurable on top of what the sources already carry. For large-scale text-to-image training that is more than good enough. For artifact-sensitive data used to train other Photoroom models, such as our custom AI Shadows model, we still rely on PNGs.

2. Building the dataset in Lance

You cannot interactively explore Parquet tables with hundreds of millions or even billions of rows. We therefore store the data in Lance, where we can index, query, and browse it. We run the ingest with Ray Data, reading the source tables in parallel and writing many fragments of a Lance table across the cluster.

One lesson worth sharing is about fragmentation. A Lance table is split into fragments, and several operations scale with the total number of fragments rather than the row count: a scan opens the files of every fragment, and some metadata operations are O(number of fragments). A table broken into too many small fragments therefore queries slowly no matter its size. The Lance performance guidance is to keep that count low (the rule of thumb is on the order of a hundred fragments even for a billion-row table) and to run compaction periodically to merge small fragments into larger ones.

We learned this the slow way. Our first ingest targeted only 100,000 rows per fragment, which left thousands of tiny fragments and made even simple filters and full-text queries crawl. After compacting up to roughly a million rows per fragment, which for our tables of several hundred million rows meant on the order of a thousand fragments, queries were fast and we stopped worrying about it. The right target depends on row width, query patterns, and how the data is written, and in hindsight fewer, larger fragments might have been a better default. Lance supports distributed compaction (for example via the Lance-Ray integration), so it is easy to adjust after the fact, but it pays not to over-fragment in the first place.

Existing captions and embeddings for exploration

We chose to re-caption every image ourselves rather than rely on the captions some of the datasets happened to ship with. The reason is consistency: across a mix of sources, caption length, style, and quality vary a lot, and we wanted a single uniform standard across the whole corpus (more on the captioner in Section 3).

The metadata a dataset already carries is still valuable, though, just for a different purpose. Pre-existing captions, and the vector embeddings that sometimes come alongside them (for example CLIP embeddings), are what make a dataset explorable in Lance right away: full-text search over the captions and nearest-neighbor search over an embedding column let us navigate the data and quickly judge its quality and the filtering it will need, well before we have run our own captioning over it. So we keep both, joining them in when they live in a separate table, and lean on them for the browsing we describe next.

Profiling and exploring the data

With the data queryable, we could quickly profile it. The resolution distribution, for example, told us where to set cutoffs. Our smallest training bucket is 512² pixels, and we cap upscaling into a bucket at 4/3 (about 33%), so the effective cutoff is 384² ≈ 147k pixels plus an aspect ratio in [0.5, 2.0]. Everything below that we discard.

Lance tables support full-text search on text columns as well as nearest-neighbor similarity search on vector embedding columns out of the box, and both can be accelerated by building indexes. We did exactly that for the caption and embedding columns, then built a small UI to browse the dataset: real-time text search over the captions as well as navigating clusters of visually similar images via vector search.

Browsing the data this way let us assess its diversity and quality qualitatively and spot issues early, simply by looking at it. We found uninformative baseline captions, a surprising amount of non-photographic content (screenshots, slides, documents, infographics), and some near-duplicate images. These observations shaped the rest of the pipeline directly: re-caption everything (longer, accurate captions describe even imperfect images well, and guarantee that the caption matches the exact image we train on), and add light filtering and deduplication passes later. In general, we believe this sort of exploration tool is invaluable for getting a feel for a dataset before committing to it.

3. Re-captioning everything with a VLM

We found that long, detailed captions are a very strong lever on output quality. In an early benchmark we trained a small diffusion model twice on the same images, once on captions we generated with Qwen2.5-VL-7B and once on much shorter captions (produced with LLaVA-1.5-LLaMA3-8B), and scored both over training. The long captions generated with Qwen2.5-VL-7B won across the board, with lower FID, CMMD, and DINO-MMD at every checkpoint (lower is better on all three).

A small diffusion model trained on Qwen2.5-VL-7B captions (red) versus the baseline LLaVA-1.5-LLaMA3-8B captions (blue), scored every 10k steps up to 100k. Lower is better on all three metrics.

Final metrics (around 100k steps):

Here is an example image with both our long, VLM-generated caption as well as the shorter caption used as the baseline:

Our captioning runs as a streaming pipeline built on Ray Data: it reads the data from Lance, prepares each image, captions it on the GPU, and writes the caption back as a new column we can query and filter on. We prompt for a single dense, visually grounded paragraph rather than a one-line caption.

Choosing the captioner

A tight schedule and limited resources meant that we couldn't spend too much time on ablations of different captioning models and system prompts.

We used the following prompts to generate the captions.

System prompt:

You are an expert image captioning model specialized in producing highly detailed, visually grounded descriptions. Write a single, flowing paragraph of approximately 100–200 words that precisely describes the given image. The caption should be written in natural prose, without bullet points, headings, or labeled sections. Describe exactly and only what is visible in the image, without speculation, interpretation of intent, or assumptions beyond visual evidence. Your description should naturally integrate the following aspects: the type of image; the main subjects and objects with appearance, materials, colors, shapes, and details; positions and spatial relationships (left/right, foreground/background, depth, overlaps, scale, framing); composition and layout (centering, balance, negative space, perspective, viewpoint); lighting and color palette; style/aesthetic qualities when supported by visual cues; and any visible text transcribed exactly as it appears, stating whether it is part of the depicted scene or a graphic overlay (watermark/UI). Do not translate or interpret the text. Maintain a neutral, precise, and descriptive tone. Avoid repetition, metaphor, emotional projection, or narrative embellishment. The goal is a dense, accurate, visually faithful caption.

User prompt:

Write the caption for this image.

We shortlisted three VLMs as captioner candidates:

• Qwen2.5-VL-7B-Captioner-Relaxed (fine-tuned for captioning)
• Qwen3-VL-8B
• Qwen3.5-9B

We benchmarked caption quality indirectly: we trained a small diffusion model for 100k steps on each caption variant and scored its generations with FID, CMMD, and DINO-MMD. Alongside the three candidates we scored two reference captioners, to connect this comparison to the first benchmark: the base Qwen2.5-VL-7B (the same model as in the benchmark above) and a set of Gemini 1.5 Flash captions we had generated earlier for the 1M-image internal test set used here. The Gemini captions are a similar length but were produced with a different prompt (all others use the system prompt above), so its curve is indicative rather than a like-for-like comparison.

FID, CMMD, and DINO-MMD over training for the three captioner candidates and two reference captioners, scored every 10k steps up to 100k. Lower is better on all three.

Final metrics (around 100k steps):