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):

Image resolution PSNR after 1× (dB) ↑ LPIPS after 1× ↓ PSNR after 10× (dB) ↑ LPIPS after 10× ↓

1–2 MP 48.7 0.004 45.4 0.008

0.25–0.5 MP 45.1 0.005 42.2 0.010

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 :