Аппроксимация параллельного прогона для авторегрессионной генерации изображений в пространстве пикселей

Аннотация

Авторегрессионная (AR) генерация в пространстве пикселей с непрерывными токенами напрямую моделирует изображения как последовательности сырых пиксельных патчей, избегая дискретной токенизации или отдельного предварительно обученного токенизатора. Однако она сталкивается с сопряженными трудностями: генерация многомерных патчей приводит к большим одношаговым ошибкам, а обучение с принудительным использованием учителя создает разрыв между обучением и инференсом, из-за которого эти ошибки накапливаются на AR-шагах. Существующие решения, такие как x-предсказание и внесение шума на входе, лишь частично смягчают эти проблемы. Точное обучение с развертыванием лучше соответствует условиям инференса, но непрактично из-за чрезвычайно медленной последовательной выборки. Мы предлагаем приближение параллельного развертывания (Parallel Rollout Approximation, PRA) — масштабируемую структуру, которая совместно решает обе проблемы. PRA генерирует низкоразмерные промежуточные состояния вместо многомерных пиксельных патчей, а затем преобразует их обратно в пиксельные токены с помощью пиксельного декодера, сохраняя AR-интерфейс с входом и выходом в пикселях. Кроме того, он создает похожие на инференс пиксельные входы через тот же путь от промежуточного состояния к пикселю, который используется при инференсе, независимо для разных позиций, аппроксимируя интерфейс пиксельной обратной связи, встречающийся при развертывании во время инференса, сохраняя при этом параллельное обучение с принудительным использованием учителя. На генерации ImageNet-1K с разрешением 256×256 в условиях классов модель PRA-S с 135M параметров достигает FID 2.58, превосходя предыдущий результат миллиардного масштаба в пиксельном пространстве (3.60). Масштабирование до PRA-L с 511M параметров дополнительно улучшает FID до 1.94, устанавливая новый рекорд среди моделей AR в пиксельном пространстве. Помимо генерации, PRA достигает более высокой точности зондирования классификации ImageNet, чем другие AR- и диффузионные базовые модели, что указывает на его потенциал для унифицированного понимания и генерации изображений в пиксельном пространстве.

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.