Approximation par rollout parallèle pour la génération d'images autorégressives dans l'espace des pixels

Résumé

La génération autorégressive (AR) à jetons continus dans l’espace des pixels modélise directement les images comme des séquences de patches de pixels bruts, évitant ainsi la tokenisation discrète ou un tokenizer pré-entraîné séparé. Cependant, elle est confrontée à deux défis couplés : la génération de patches de haute dimension provoque des erreurs importantes à chaque étape, et l’apprentissage par forçage enseignant crée un écart entre l’entraînement et l’inférence, ce qui fait que ces erreurs s’accumulent au fil des étapes AR. Les correctifs existants, tels que la prédiction en x et l’injection de bruit en entrée, n’atténuent que partiellement ces problèmes. L’apprentissage par rollout exact correspond mieux aux conditions d’inférence, mais il est irréaliste en raison d’un échantillonnage séquentiel excessivement lent. Nous proposons l’Approximation de Rollout Parallèle (PRA), un cadre évolutif qui aborde conjointement ces deux défis. PRA génère des états intermédiaires de faible dimension au lieu de patches de pixels de haute dimension, puis les reconvertit en jetons dans l’espace des pixels à l’aide d’un décodeur de pixels, préservant ainsi une interface AR pixel-entrée, pixel-sortie. Il construit également des entrées pixels similaires à celles de l’inférence via le même chemin état-intermédiaire-vers-pixel utilisé en inférence, de manière indépendante pour chaque position, approximant ainsi l’interface de rétroaction pixel rencontrée lors du rollout en inférence tout en conservant un apprentissage parallèle par forçage enseignant. Sur la génération ImageNet-1K conditionnée par classe à la résolution 256×256, PRA-S avec 135M paramètres atteint un FID de 2,58, surpassant le précédent résultat AR dans l’espace des pixels à l’échelle du milliard de paramètres (3,60). La mise à l’échelle vers PRA-L avec 511M paramètres améliore encore le FID à 1,94, établissant un nouvel état de l’art parmi les modèles AR dans l’espace des pixels. Au-delà de la génération, PRA obtient une précision de sondage en classification ImageNet plus élevée que d’autres modèles AR et de diffusion de référence, suggérant son potentiel pour une génération et une compréhension unifiées d’images dans l’espace des pixels.

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.