PRX revela su estrategia de datos para preentrenamiento

PRX arma su corpus mezclando datos públicos e internos, reescribiendo descripciones con un VLM y convirtiendo el resultado en el flujo que alimentó el preentrenamiento del modelo PRX. ¿Por qué tanto énfasis en cantidad y consistencia en lugar de perfección por imagen? Porque en esta etapa el modelo aprende cobertura visual y variedad, no estilo pulido.

Resumen de la pipeline

A alto nivel la receta fue simple pero poderosa: juntar muchas fuentes, estandarizar las descripciones con un visual language model (VLM), indexar y explorar con Lance para curación y finalmente transformar todo a MDS para el entrenamiento distribuido.

El objetivo del preentrenamiento es la amplitud: que el modelo aprenda objetos, composición, iluminación y la inmensa diversidad de imágenes reales. Filtrar demasiado por estética aquí empobrece la distribución y cuesta conceptos que no se recuperan luego. Lo "bonito" llega después, en fine-tuning y alignment sobre conjuntos pequeños y muy cuidados.

Datos y formatos: por qué usar Lance para construir y MDS para entrenar

• Lance es un formato columnar pensado para explorar y hacer feature engineering a escala: predicate pushdown, índices escalares y búsqueda por vectores. Ideal para tablas de cientos de millones o miles de millones de filas.
• MDS (Mosaic Data Shards) es rígido pero eficiente para streaming en entrenamiento distribuido: shards de ~128 MB, mezcla ponderada de streams, reanudable y determinista.

La estrategia fue usar Lance para construir, filtrar y explorar, y luego convertir los datos finalizados a MDS para el trainer. Esa combinación entrega flexibilidad en la etapa de curación y bajo coste operativo en la etapa de entrenamiento.

Texto y latentes: precomputar versus calcular on-the-fly

Antes precomputaban latentes textuales con T5Gemma y los almacenaban en MDS como bytes. Para PRX 7B cambiaron a Qwen3-VL y optaron por calcular las latentes dentro del loop de entrenamiento. ¿Por qué?

• Coste de throughput: correr el encoder en línea costó solo alrededor de 3-4% de throughput (aprox. 1 día extra en una corrida de 30 días) porque el denoiser domina el cómputo a escala 7B.
• Ventajas: shards mucho más pequeños (pueden caber en SSD compartido del cluster SLURM) y libertad para cambiar el text encoder sin reescribir terabytes.

Lección práctica: si el text encoder es ligero comparado con el modelo visual, calcular latentes en tiempo real compensa por la flexibilidad y el ahorro en I/O y almacenamiento.

Imágenes y compresión: JPEG quality 92 comprobado

Guardaron todas las imágenes como JPEG calidad 92 tras medir el impacto. La razón: la mayoría de fuentes ya están JPEG-comprimidas; PNG sería 3-10x más grande sin ganancia perceptual.

Resultados sobre 100 imágenes (PSNR más alto y LPIPS más bajo son mejores):

También entrenaron dos PRX iguales (1024px), una con PNG y otra con JPEG 92. En métricas y apariencia las generaciones fueron prácticamente indistinguibles:

Conclusión: para preentrenamiento a gran escala, JPEG calidad alta es suficiente. Para tareas sensibles a artefactos (p. ej. modelos de sombras de Photoroom) siguen usando PNG.

Captioning: por qué los captions largos importan y cómo los generaron

Descubrieron que captions largos y fieles mejoran las generaciones significativamente. Un experimento con un pequeño modelo de difusión mostró que usar captions generadas por Qwen2.5-VL-7B (largo y detallado) reduce FID, CMMD y DINO-MMD frente a captions cortas.

Benchmark final (cerca de 100k pasos) comparando captioners:

Tuning de pipeline y throughput importaron: Qwen3-VL-8B rindió ~20 img/s por H200 en vLLM y fue la elección por calidad, velocidad y soporte estable. Qwen3.5-9B daba mejores métricas pero iba ~6.5 img/s y requería builds inestables en ese momento.

El prompt de sistema para el captioner pedía un párrafo de 100-200 palabras, descriptivo, sin interpretaciones ni metáforas, con posición, materiales, colores, texto transcrito literalmente y marcado si era overlay.

Indexado y exploración con Lance

Guardar todo en Lance permitió búsquedas textuales y nearest-neighbor con embeddings, y montar una UI para navegar el corpus. Eso hizo evidente la presencia de muchos screenshots, infographics y captions pobres.

Fragmentación: un punto crítico. Inicialmente escribieron 100k filas por fragmento y terminaron con miles de fragmentos pequeños que ralentizaban consultas. Compactar a ~1M filas por fragmento (para tablas de cientos de millones de filas) dejó el número de fragmentos en una escala manejable y aceleró las consultas.

Regla práctica: mantener el número de fragments bajo (orden de cientos para tablas de mil millones) y compactar periódicamente. Lance soporta compaction distribuida.

Filtrado, deduplicación y skip-lists

• Filtrado por texto: hicieron una clasificación desde las captions con Qwen3-8B en modo texto que etiqueta muestras como visual, texto o nsfw. Es rápido (~200 captions/s por GPU) y barato porque usa trabajo ya hecho.
• Skip-lists: en lugar de reescribir shards, añadieron archivos laterales con índices a saltar. El loader une los skip-lists y omite esos ítems en entrenamiento. Flexible para opt-outs o para experimentar con ablations sin duplicar datos.
• Límite práctico: si el porcentaje de muestras saltadas crece mucho (estimado ~10%), conviene reescribir el dataset.

Deduplicación: usaron hashes perceptuales DCT-based (thumbnail, 2D DCT, threshold low-freq coefficients) y consideraron iguales dos huellas con Hamming distance cero. Esto atrapa reencodes y resizes; para balance conceptual habría que deduplicar en embeddings.

Resultados agregados: dedup removió unos pocos por ciento, el filtro por texto algunos por ciento más, y NSFW una fracción pequeña.

Bucketing y streaming para difusión

• Buckets por tier: 512, 1024, 2048, 4096 px. Una imagen entra en el mayor tier sin upscaling >4/3 (floor del tier 512 ≈ 384^2). Todo debajo del floor se descarta.
• Buckets por aspect ratio: AR en [0.5, 2.0], 13 formas por tier (patch-aligned, múltiplos de 16) para mantener un budget de tokens constante por imagen.
• Resize con Lanczos y encode JPEG Q92. MDS organizado en árbol por resolución y AR para streaming eficiente.

Mosaic Streaming aporta: reordenamiento determinista y reanudable, mezcla ponderada de streams y elasticidad mid-epoch sin reprocesar datos.

Recomendaciones prácticas y lecciones

• Prioriza cobertura y diversidad en preentrenamiento; deja la pulcritud para fine-tuning.
• Usa Lance para explorar y MDS para entrenar si necesitas streaming robusto al scale.
• Calcular text latents on-the-fly es una buena opción si el encoder no domina cómputo; te da flexibilidad para reemplazar modelos de texto.
• Mantén pocos fragments grandes en Lance; compacta periódicamente.
• Prefiere captions largos y fieles: convierten ruido visual en atributos controlables.
• Implementa skip-lists para excluir ejemplos sin reescribir shards; reescribe solo si el skip crece mucho.

La pieza que relata PRX es técnica, práctica y repleta de detalles aplicables a cualquier proyecto de imagen a escala. Si estás armando un corpus grande, estas decisiones son exactamente las que te ahorran semanas de I/O, almacenamiento y dolores de cabeza de ingeniería.

Fuente original

https://huggingface.co/blog/Photoroom/prx-part4-data