En esta publicación, exploraremos el entrenamiento de un modelo de texto a imagen en solo 24 horas. Este proceso es una continuación de los enfoques discutidos en Part 1 y Part 2.
El objetivo es responder la pregunta de qué sucede al combinar todas las técnicas que han demostrado ser efectivas. En lugar de optimizar una sola dimensión, combinaremos los mejores elementos y evaluaremos el rendimiento dentro de un presupuesto de cómputo estricto.
Detalles del Entrenamiento
Realizamos un speedrun de 24 horas utilizando:
- 32 H200
- ~$1500 en total (2$/hora/GPU)
Este enfoque es un claro avance respecto a los primeros días de los modelos de difusión, donde los costos podían llegar a millones de dólares. Se busca demostrar la evolución del campo y los logros que se pueden alcanzar en un solo día.
Los resultados de este speedrun servirán como base para futuras recetas de entrenamiento a gran escala.
Además, publicaremos nuestro código en Github, que incluirá el código de entrenamiento y el marco experimental utilizado en publicaciones anteriores.
La Receta de Entrenamiento
El entrenamiento se basa en la formulación de x-prediction, permitiendo el entrenamiento directo en el espacio de píxeles y eliminando la necesidad de un VAE. Usamos un tamaño de parche de 32 y un cuello de botella de 256 dimensiones en la capa de proyección. Esto facilita el manejo del tamaño de secuencia, haciendo que el entrenamiento en el espacio de píxeles sea computacionalmente viable incluso a resoluciones más altas.
A 512px, la longitud de la secuencia es de 256, y a 1024px, se convierte en 1024. En lugar de seguir el programa habitual de 256px → 512px → 1024px, comenzamos directamente a 512px y luego ajustamos a 1024px. Con un control adecuado de los conteos de tokens y hardware moderno, el entrenamiento en el espacio de píxeles se ha vuelto más accesible y directo.
Pérdidas Perceptuales
Un beneficio adicional de predecir directamente en el espacio de píxeles es la posibilidad de reutilizar herramientas de visión por computadora clásica. Usamos pérdidas perceptuales que se aplican a imágenes completas en lugar de características por parches, mejorando así la calidad visual y la velocidad de convergencia.
En este entrenamiento, se añadieron dos pérdidas auxiliares, LPIPS y DINO, para mejorar la calidad del modelo. Estos enfoques son ligeros en comparación con el paso principal del transformador y no incrementan significativamente la carga computacional.
Optimización y Configuración
Utilizamos el optimizador Muon para parámetros 2D, lo que mostró claras mejoras en comparación con Adam. Este optimizador se aplica junto con Adam para todos los parámetros no 2D.
El entrenamiento se realizó en tres conjuntos de datos sintéticos disponibles públicamente. La planificación consistió en entrenar a 512px durante 100k pasos y luego afinar a 1024px durante 20k pasos.
Resultados y Conclusiones
Los resultados del entrenamiento muestran que, a pesar de ser un proceso de un día, el modelo ya es funcional, aunque no perfecto. Se observan algunos errores menores, pero la calidad general es prometedora. Los problemas restantes parecen más bien artefactos de subentrenamiento y diversidad limitada de datos.
Este speedrun destaca el avance en el entrenamiento de modelos de difusión, haciendo posible obtener un modelo significativo en un día y con un presupuesto accesible.
Próximos Pasos
Este entrenamiento de 24 horas es solo el comienzo. Continuaremos mejorando la misma receta y ajustando los conjuntos de datos y la capitulación. Todo el código y la configuración de este speedrun están disponibles en el repositorio PRX: https://github.com/Photoroom/PRX.
Agradecemos su interés y los invitamos a unirse a nuestra comunidad en Discord para discutir sobre difusión y modelos de texto a imagen.
