Aproximación de Despliegue Paralelo para la Generación de Imágenes Autoregresiva en el Espacio de Píxeles

Resumen

La generación autoregresiva (AR) de tokens continuos en el espacio de píxeles modela directamente las imágenes como secuencias de parches de píxeles en bruto, evitando la tokenización discreta o un tokenizador preentrenado por separado. Sin embargo, enfrenta desafíos acoplados: la generación de parches de alta dimensionalidad provoca grandes errores de un solo paso, y el entrenamiento con forzamiento de profesor crea una brecha entrenamiento-inferencia que hace que estos errores se acumulen a lo largo de los pasos AR. Las soluciones existentes, como la predicción en x y la inyección de ruido en la entrada, solo mitigan parcialmente estos problemas. El entrenamiento exacto con despliegue (rollout) se adapta mejor a las condiciones de inferencia, pero es inviable debido al prohibitivamente lento muestreo secuencial. Proponemos Parallel Rollout Approximation (PRA), un marco escalable que aborda ambos desafíos de manera conjunta. PRA genera estados intermedios de baja dimensionalidad en lugar de parches de píxeles de alta dimensionalidad, y luego los mapea de vuelta a tokens en el espacio de píxeles mediante un decodificador de píxeles, preservando una interfaz AR de entrada y salida en píxeles. También construye entradas de píxeles similares a las de inferencia a través del mismo camino estado-intermedio-a-píxel utilizado en inferencia, de forma independiente entre posiciones, aproximando la interfaz de retroalimentación de píxeles encontrada durante el despliegue en inferencia, mientras retiene el entrenamiento paralelo con forzamiento de profesor. En la generación condicional por clase ImageNet-1K a resolución 256×256, PRA-S con 135 millones de parámetros alcanza un FID de 2.58, superando el resultado previo de AR en espacio de píxeles a escala de mil millones de parámetros, que era de 3.60. Escalar a PRA-L con 511 millones de parámetros mejora aún más el FID a 1.94, estableciendo un nuevo estado del arte entre los modelos AR en espacio de píxeles. Más allá de la generación, PRA logra una mayor precisión de sondeo de clasificación en ImageNet que otras líneas base AR y de difusión, lo que sugiere su potencial para la generación y comprensión unificada de imágenes en el espacio de píxeles.

English

Pixel-space continuous-token autoregressive (AR) generation directly models images as sequences of raw pixel patches, avoiding discrete tokenization or a separately pretrained tokenizer. However, it faces coupled challenges: high-dimensional patch generation causes large single-step errors, and teacher-forced training creates a train--inference gap that makes these errors accumulate across AR steps. Existing fixes such as x-prediction and input noise injection only partially mitigate these issues. Exact rollout training better matches inference-time conditions, but is impractical due to prohibitively slow sequential sampling. We propose Parallel Rollout Approximation (PRA), a scalable framework that addresses both challenges jointly. PRA generates low-dimensional intermediate states instead of high-dimensional pixel patches, then maps them back to pixel-space tokens with a pixel decoder, preserving a pixel-in, pixel-out AR interface. It also constructs inference-like pixel inputs through the same intermediate-state-to-pixel path used at inference, independently across positions, approximating the pixel-feedback interface encountered during inference-time rollout while retaining parallel teacher-forced training. On class-conditional ImageNet-1K generation at 256times256 resolution, PRA-S with 135M parameters achieves an FID of 2.58, surpassing the previous billion-scale pixel-space AR result of 3.60. Scaling to PRA-L with 511M parameters further improves FID to 1.94, establishing a new state of the art among pixel-space AR models. Beyond generation, PRA achieves higher ImageNet classification probing accuracy than other AR and diffusion baselines, suggesting its potential for unified pixel-space image generation and understanding.