El problema fundamental que aborda P-JiT es la ineficiencia computacional y la lentitud de convergencia en los modelos generativos de video, particularmente aquellos basados en arquitecturas de atención con complejidad O(n²). Históricamente, la compresión de datos para estos modelos se ha delegado a Variational Autoencoders (VAEs) o modelos especializados, lo que introduce un cuello de botella y un límite empírico en la agresividad de la compresión, además de posibles divergencias entre la optimización del VAE y la calidad de la generación final.
La tesis central de P-JiT es que un modelo de difusión en espacio de píxeles puede aprender un espacio latente intrínsecamente optimizado para su tarea de generación, eliminando la necesidad de un VAE separado. Esto permite una compresión más agresiva y una convergencia más rápida al integrar la compresión y la generación en un único proceso de difusión. La clave reside en la capacidad del modelo de difusión para manejar directamente la representación de píxeles, lo que facilita la predicción de versiones de baja resolución de la imagen final en etapas tempranas de la red, mejorando la señal de gradiente y acelerando el entrenamiento.
Arquitectura del Sistema
La arquitectura Pyramid-JiT (P-JiT) es un modelo de difusión basado en transformadores (DiT) que opera directamente en el espacio de píxeles. Se compone de dos bloques 'refiner' de modalidad única (uno para texto y otro para imagen) seguidos de una secuencia de 22 bloques DiT idénticos. La representación de texto se obtiene de un modelo Qwen3.5-4B congelado, extrayendo características de sus capas de atención completas (7, 15, 27) y proyectándolas en el refiner de texto. Las imágenes se dividen en parches de 32x32, se comprimen a un espacio de 1024 dimensiones y se proyectan en el refiner de imagen.
Una característica distintiva es la inclusión de 'readouts' a lo largo del tronco del DiT. Específicamente, en el bloque 5, se utiliza un readout PixelREPA de 2 capas para alinear el espacio latente del modelo con DINOv3-L. Adicionalmente, hay dos cabezas de readout lineales después de los bloques 10 y 16 que predicen la imagen limpia a un cuarto y mitad de resolución, respectivamente. Esta estrategia de readouts tempranos y de baja resolución está diseñada para reforzar la señal de gradiente en las capas iniciales de la red, acelerando la convergencia. Cada bloque DiT está condicionado por el timestep mediante AdaLN de bajo rango, siguiendo el patrón de Z-Image, con una proyección descendente del embedding del timestep a 256 dimensiones y proyecciones ascendentes individuales para escala y puerta, inicializadas a cero para que cada bloque comience como una identidad. La pérdida total combina MSE en los readouts, distancia coseno de PixelREPA, MSE en la cabeza de salida final y pérdidas perceptuales (LPIPS y P-DINO) en la imagen de resolución completa.
Flujo de Entrenamiento de Pyramid-JiT
- 1 Entrada de Texto e Imagen Texto pasa por Qwen3.5-4B (capas 7, 15, 27). Imagen se divide en parches 32x32.
- 2 Refiner de Modalidad Única Texto y imagen pasan por bloques DiT 'refiner' separados para su modalidad.
- 3 Fusión y Bloques DiT Representaciones fusionadas, pasan por 22 bloques DiT idénticos.
- 4 Readout PixelREPA (Bloque 5) Alinea espacio latente con DINOv3-L.
- 5 Readouts de Imagen (Bloques 10, 16) Predicen imagen limpia a 1/4 y 1/2 resolución.
- 6 Condicionamiento por Timestep Cada bloque DiT usa AdaLN con embedding de timestep.
- 7 Cálculo de Pérdida Total Combina MSE, PixelREPA, LPIPS y P-DINO.
- 8 Actualización de Pesos Optimizador Muon/AdamW con curriculum de ruido y EMA.
| Capa | Tecnología | Justificación |
|---|---|---|
| compute | NVIDIA H100 GPUs | Hardware principal para el entrenamiento distribuido de modelos de difusión a gran escala. 64x H100 (8 nodos), bf16, FSDP2, 32 imágenes por GPU. |
| data-processing | Qwen3.5-4B | Modelo de lenguaje congelado utilizado como codificador de texto para generar embeddings de texto que condicionan el modelo de difusión. Lectura en capas 7, 15 y 27 para concatenar y proyectar representaciones. |
| data-processing | DINOv3-L | Modelo de visión utilizado para alinear el espacio latente del modelo mediante la pérdida PixelREPA, mejorando la coherencia visual. |
| orchestration | FSDP2 (Fully Sharded Data Parallel) | Estrategia de paralelización de datos para distribuir el entrenamiento del modelo a través de múltiples GPUs y nodos, optimizando el uso de memoria y la eficiencia. vs Data Parallel, Model Parallel |
| compute | Muon Optimizer | Optimizador utilizado para las matrices de bloques 2D del modelo, combinado con AdamW para el resto de los parámetros. vs AdamW, SGD momentum 0.95, Nesterov, actualizaciones RMS-matched a AdamW. |
| compute | AdamW Optimizer | Optimizador estándar utilizado para la mayoría de los parámetros del modelo. vs SGD, Muon β = 0.9, 0.999, ε = 1e-8. |
Trade-offs
Ganancias
- ▲ Costo de inferencia
- ▲ Velocidad de entrenamiento (convergencia)
- ▲ Agresividad de compresión
Costes
- △ Capacidad de generar texto o infografías
- △ Rendimiento en escenas complejas (cuerpos completos, multitudes)
Fundamentos Teóricos
El concepto de mejorar la señal de gradiente en las capas tempranas de redes profundas, especialmente en modelos generativos, tiene raíces en la investigación sobre el problema del 'vanishing gradient'. Trabajos como los de Hochreiter y Schmidhuber (1991) sobre redes neuronales recurrentes ya señalaban la dificultad de propagar gradientes a través de muchas capas. En el contexto de los modelos de difusión y transformadores, la observación de gradientes débiles en los bloques tempranos del DiT ha sido explorada en trabajos adyacentes como REPA, RAE-v2 y Self-Flow. Estos estudios sugieren que las capas iniciales pueden no contribuir eficazmente al aprendizaje de detalles finos o a la estructura global, lo que ralentiza la convergencia.
La estrategia de P-JiT de predecir versiones de baja resolución de la imagen final en etapas tempranas de la red se alinea con principios de aprendizaje multi-escala y supervisión auxiliar, donde la adición de objetivos de pérdida intermedios ayuda a guiar el entrenamiento. Esto es análogo a técnicas como la 'deep supervision' o 'intermediate supervision' que han sido aplicadas en redes convolucionales para tareas de clasificación y segmentación, donde se añaden clasificadores o cabezas de predicción en capas intermedias para combatir el problema del gradiente desvanecido y acelerar la convergencia, como se vio en la arquitectura GoogLeNet (Szegedy et al., 2015).