Stratégie de données de PRX pour l'entraînement IA

Q: Quel est l'impact d'un corpus diversifié sur l'entraînement des modèles IA ?

Un corpus diversifié permet au modèle d'apprendre une gamme plus large de concepts visuels, améliorant ainsi la qualité et la flexibilité des générations.

Back to Articles 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 n

Accueil

Outils

Annuaire

Apprendre