像素空间自回归图像生成的并行部署近似方法
Parallel Rollout Approximation for Pixel-Space Autoregressive Image Generation
摘要
像素空间连续标记自回归生成方法直接将图像建模为原始像素块的序列,从而避免了离散标记化或独立预训练的标记器。不过,该方法也面临一些挑战:高维像素块的生成会导致较大的单次错误;而由教师强制式训练则会导致训练与推理之间的差距,使得这些错误在多个生成步骤中不断累积。现有的解决方案如x预测和输入噪声注入只能部分缓解这些问题。精确展开训练虽然能更好地模拟推理时的条件,但由于顺序采样速度过慢,因此并不实用。我们提出了并行展开近似方法,这是一种可扩展的框架,能够同时解决上述两个问题。该方法生成低维中间状态而非高维像素块,然后通过像素解码器将其转换回像素空间的标记,从而保持像素输入与输出的对应关系。它还通过相同的中间状态到像素的路径来生成类似推理时的像素输入,从而近似于推理过程中遇到的像素反馈机制,同时保留教师强制式训练的功能。在256×256分辨率下对ImageNet-1K数据集进行类条件生成时,拥有135百万参数的PRA-S模型的FID值为2.58,超过了之前10亿像素级别像素空间自回归模型的3.60值。当使用拥有511百万参数的PRA-L模型时,FID值进一步降低至1.94,这标志着该模型在像素空间自回归模型中达到了新的最佳性能水平。除了生成功能外,PRA还比其他自回归和扩散模型具有更高的ImageNet分类准确率,这表明其具有在像素空间范围内实现统一图像生成与理解的能力。
English Abstract
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.