Aproximação Paralela de Rollout para Geração Autorregressiva de Imagens no Espaço de Pixels

Resumo

A geração autoregressiva (AR) de tokens contínuos no espaço de pixels modela diretamente imagens como sequências de patches de pixels brutos, evitando tokenização discreta ou um tokenizador pré-treinado separado. No entanto, enfrenta desafios acoplados: a geração de patches de alta dimensionalidade causa grandes erros de etapa única, e o treinamento forçado pelo professor cria uma lacuna treinamento-inferência que faz com que esses erros se acumulem ao longo das etapas AR. Correções existentes, como predição-x e injeção de ruído na entrada, apenas mitigam parcialmente esses problemas. O treinamento exato com rollout corresponde melhor às condições de inferência, mas é impraticável devido à amostragem sequencial proibitivamente lenta. Propomos a Aproximação Paralela de Rollout (PRA), uma estrutura escalável que aborda ambos os desafios em conjunto. A PRA gera estados intermediários de baixa dimensionalidade em vez de patches de pixels de alta dimensionalidade e, em seguida, os mapeia de volta para tokens no espaço de pixels com um decodificador de pixels, preservando uma interface AR pixel-in, pixel-out. Ela também constrói entradas de pixels no estilo de inferência por meio do mesmo caminho de estado intermediário para pixel usado na inferência, independentemente entre posições, aproximando a interface de feedback de pixels encontrada durante o rollout em inferência, enquanto mantém o treinamento paralelo forçado pelo professor. Na geração condicional de classe ImageNet-1K na resolução 256×256, o PRA-S com 135M parâmetros atinge um FID de 2,58, superando o resultado AR no espaço de pixels anterior de escala de bilhões de parâmetros (3,60). A escalada para PRA-L com 511M parâmetros melhora ainda mais o FID para 1,94, estabelecendo um novo estado da arte entre modelos AR no espaço de pixels. Além da geração, o PRA atinge maior precisão de sondagem de classificação ImageNet do que outras linhas de base AR e de difusão, sugerindo seu potencial para geração e compreensão unificadas de imagens no espaço de 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.