Resumen Técnico: Aproximación de Despliegue Paralelo para la Generación de Imágenes Autorregresiva en el Espacio de Píxeles

1. Planteamiento del Problema

La generación de imágenes autorregresiva (AR) en el espacio de píxeles busca modelar imágenes como secuencias de parches de píxeles brutos sin depender de una tokenización discreta o de tokenizadores preentrenados por separado. Aunque es conceptualmente sencillo, este enfoque enfrenta dos desafíos acoplados que históricamente han limitado su rendimiento en comparación con los modelos AR en el espacio latente o los modelos de difusión:

1. Desafío del Lado de la Salida (Predicción de Alta Dimensionalidad): Predecir parches de píxeles brutos de alta dimensión (por ejemplo, vectores de 768 dimensiones para parches de 16×16) en un solo paso AR es difícil, lo que conduce a grandes errores de generación de un solo paso.
2. Desafío del Lado de la Entrada (Desajuste entre Entrenamiento e Inferencia): El entrenamiento estándar mediante teacher forcing condiciona al modelo en prefijos de verdad fundamental (ground-truth) limpios. Sin embargo, durante la inferencia, el modelo debe condicionarse en sus propios parches de píxeles previamente generados (y potencialmente erróneos). Este desajuste provoca que los errores se acumulen y se propaguen a través de los pasos AR, un fenómeno conocido como sesgo de exposición (exposure bias).

Las estrategias de mitigación existentes, como la predicción de $x$ (efectiva en modelos de difusión) o la inyección de ruido en la entrada, solo abordan parcialmente estos problemas. El entrenamiento de despliegue (rollout) exacto, que expondría al modelo a sus propios prefijos generados durante el entrenamiento, es computacionalmente prohibitivo para el AR de tokens continuos porque requiere un muestreo secuencial para cada token.

2. Metodología: Aproximación de Despliegue Paralelo (PRA)

Los autores proponen la Aproximación de Despliegue Paralelo (PRA), un marco escalable que aborda conjuntamente los desafíos de salida y entrada manteniendo una interfaz de "píxel de entrada, píxel de salida".

2.1 Lado de la Salida: Objetivos Intermedios de Extremo a Extremo

Para reducir la dificultad de la generación de un solo paso, PRA no predice directamente parches de píxeles de alta dimensión. En su lugar, genera estados intermedios de baja dimensión ($z_i \in \mathbb{R}^{d_z}$, donde $d_z \ll d$).

• Construcción del Objetivo: Durante el entrenamiento, el token de píxel de verdad fundamental $x_i$ y la representación del prefijo causal $h_{i-1}$ se mapean a un estado intermedio objetivo $z_i$ a través de una red de proyección ligera $g_\psi$. Este objetivo es sensible al contexto, incorporando información del prefijo.
• Decodificación: Un decodificador de píxeles causal $D_\phi$ mapea el estado intermedio generado $\hat{z}_i$ (y los estados previos $\hat{z}_{<i}$) de vuelta al espacio de píxeles, produciendo el token de salida $\hat{x}_i$.
• Entrenamiento: El modelo AR se entrena para predecir $z_i$ utilizando una cabeza de difusión a nivel de token (objetivo de flujo rectificado o rectified-flow), mientras que el decodificador se entrena para reconstruir $x_i$ a partir de $z_i$. Esto permite que el núcleo (backbone) AR opere en un espacio de menor dimensión mientras la interfaz externa permanece en el espacio de píxeles.

2.2 Lado de la Entrada: Entradas Similares a la Inferencia en Paralelo

Para cerrar la brecha entre el entrenamiento y la inferencia sin un despliegue secuencial, PRA construye entradas de píxeles similares a las de inferencia en paralelo.

• Mecanismo: En lugar de entrenar con píxeles de verdad fundamental limpios, PRA perturba los estados intermedios objetivo $z_i$ (usando interpolación de flujo rectificado) para crear estados ruidosos $\tilde{z}_t$. Estos se pasan a través del mismo decodificador de píxeles utilizado en la inferencia para generar tokens de píxeles reconstruidos $\bar{x}_i$.
• Paralelismo: Dado que cada $\bar{x}_i$ se construye independientemente de su correspondiente estado objetivo (en lugar de una trayectoria autorregresiva secuencial), este proceso sigue siendo totalmente paralelizable.
• Bucle de Entrenamiento: El modelo AR se entrena entonces para predecir el objetivo $z_i$ condicionado por el prefijo de estos píxeles reconstruidos y similares a la inferencia $\bar{x}_{<i}$. Esto aproxima la estructura de error del despliegue en tiempo de inferencia manteniendo la eficiencia del entrenamiento con teacher forcing.

3. Contribuciones Clave

1. Marco Unificado: PRA es el primer marco escalable que aborda conjuntamente los errores de salida de alta dimensión y el desajuste entre entrenamiento e inferencia en la generación de píxeles AR sin requerir un tokenizador preentrenado por separado.
2. Aprendizaje de Estados Intermedios: El artículo introduce un espacio de objetivos intermedios aprendido de extremo a extremo que está optimizado específicamente para el modelo AR causal, en lugar de depender de espacios latentes congelados de VAEs.
3. Aproximación de Despliegue Paralelo: El método aproxima con éxito los beneficios del entrenamiento de despliegue (exponiendo al modelo a contextos generados) mediante la construcción paralela y no secuencial de entradas de píxeles decodificadas.
4. Interfaz de Píxel de Entrada, Píxel de Salida: A diferencia de los modelos AR en el espacio latente, PRA mantiene una interfaz directa con parches de píxeles brutos, facilitando posibles aplicaciones en la comprensión y generación unificada de imágenes.

4. Resultados Experimentales

Los autores evaluaron PRA en la generación condicional de clase para ImageNet-1K a una resolución de 256×256.

• Rendimiento de Vanguardia (SOTA):
  • PRA-S (135M parámetros): Logró un FID de 2.58, superando a modelos anteriores de AR en el espacio de píxeles de escala de miles de millones (ej. FARMER-1.9B con FID 3.60).
  • PRA-L (511M parámetros): Logró un FID de 1.94, estableciendo un nuevo estado del arte para los modelos AR en el espacio de píxeles.
• Estudios de Ablación:
  • Objetivos Intermedios: Los objetivos conscientes del contexto (que incorporan la representación del prefijo) superaron a los objetivos puramente locales o a los codificadores LDM congelados. Se encontró que una dimensión intermedia de $d_z=16$ era óptima.
  • Construcción de la Entrada: El uso de entradas de píxeles decodificadas (PRA) superó significativamente el entrenamiento con píxeles de verdad fundamental limpios, estados intermedios de verdad fundamental o la simple inyección de ruido, confirmando la necesidad de igualar el bucle de retroalimentación de píxeles de la inferencia.
• Comprensión de Imágenes:
  • Al aplicarse al linear probing de ImageNet, PRA-L logró un 68.80% de precisión top-1, superando tanto a los modelos base de AR en el espacio latente (SphereAR-L: 52.19%) como a los modelos base de difusión en el espacio de píxeles (JiT-L: 42.76%). Esto sugiere que el modelo aprende representaciones visuales transferibles.

5. Significado y Reivindicaciones

El artículo afirma que PRA resuelve los cuellos de botella fundamentales que impiden que el AR en el espacio de píxeles compita con los modelos de difusión y de AR en el espacio latente. Al reducir la dificultad de la generación de un solo paso mediante estados intermedios de baja dimensión y mitigar la acumulación de errores mediante la aproximación de despliegue paralelo, PRA permite la generación de imágenes de alta calidad directamente desde píxeles brutos.

Los autores destacan que la naturaleza de "píxel de entrada, píxel de salida" de PRA, combinada con su fuerte rendimiento en el linear probing, sugiere un potencial para modelos unificados en el espacio de píxeles capaces tanto de la generación de imágenes de alta fidelidad como de una comprensión visual robusta, eliminando la necesidad de tokenizadores o autoencoders separados.

Limitaciones: Los autores señalan que PRA introduce componentes adicionales (codificador de estado intermedio, decodificador de píxeles y una pasada forward de AR paralela adicional durante el entrenamiento). Aunque estos se entrenan de extremo a extremo, se requiere trabajo futuro para validar el marco en dominios de datos y tareas de generación más amplios.