PaperlayerView exact revision on arXiv

Source-linked AI summary

Parallel Rollout Approximation for Pixel-Space Autoregressive Image Generation

Jiayi Xu, Di He, Guolin Ke

arXiv:2606.27978v1cs.CVcs.AI

TL;DR

Pixel-space autoregressive image generation struggles with high-dimensional patch errors and train–inference mismatch. PRA addresses both through low-dimensional intermediate states and parallel inference-like pixel inputs, achieving state-of-the-art pixel-space AR performance with 1.94 FID on ImageNet-1K.

  • Problem

    Pixel-space autoregressive models must predict high-dimensional pixel patches while handling train–inference mismatch that can accumulate errors across autoregressive steps.

  • Method

    PRA predicts low-dimensional intermediate states, decodes them into pixel tokens, and constructs inference-like pixel inputs in parallel through the same decoding path.

  • Results

    PRA improves pixel-space autoregressive generation and visual representation probing, with PRA-L achieving an FID of 1.94 on class-conditional ImageNet-1K generation.

  • Takeaways & Limitations

    PRA establishes a new state of the art among pixel-space autoregressive models and supports their potential for unified image generation and understanding.

  • Takeaways & Limitations

    PRA adds an intermediate state, pixel decoder, and additional parallel autoregressive forward during training, and broader validation remains future work.

Abstract

from arXiv · show

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 \emph{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 $256\times256$ 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.

1 Introduction

Pixel-space autoregressive generation avoids tokenizer dependence but couples high-dimensional patch prediction with imperfect pixel contexts. PRA addresses both bottlenecks by generating compact intermediate states, decoding them into pixel tokens, and constructing inference-like inputs in parallel.

  • Learned token spaces make autoregressive modeling easier but add an extra stage and constrain final quality through the tokenizer or autoencoder.This motivates end-to-end pixel-space autoregression directly on raw pixel patches.
  • Raw pixel patches are both current-step targets and future causal context, creating coupled output-side and input-side challenges.PRA’s framing treats generation errors and train–inference context mismatch as linked problems.
  • PRA generates low-dimensional intermediate states, decodes them into pixel tokens, and preserves pixel interaction through pixel prefixes and decoded inputs.The same intermediate-state-to-pixel path supports inference-like inputs constructed independently across positions while retaining parallel teacher-forced training.
  • 135M-parameter PRA-S surpasses prior billion-scale pixel-space AR models, while PRA-L reaches an FID of 1.94 on class-conditional ImageNet-1K at 256×256 resolution.Ablations find complementary gains from low-dimensional intermediate states and decoded pixel inputs.

2 Related Work

Prior work identifies train–inference mismatch as a source of exposure bias and error accumulation in autoregressive generation, while continuous-token models increasingly use diffusion heads and latent representations. PRA instead learns low-dimensional intermediate states end-to-end without a separately pretrained tokenizer, preserving a pixel-in, pixel-out interface amid pixel-space generative modeling dominated by diffusion.

  • Train–Inference Mismatch and Rollout-Based Training: Teacher-forced autoregressive training conditions on ground-truth prefixes, whereas inference uses generated prefixes, causing exposure bias and error accumulation.Existing remedies expose models to generated or perturbed inputs during training, including scheduled sampling and sequence-level rollout methods.
  • Continuous-Token Autoregressive Generation: Continuous-token autoregressive models extend generation from discrete vocabularies to continuous-valued tokens, including latent features and raw signal patches.Recent methods commonly use token-level diffusion heads to model continuous next-token distributions.
  • Continuous-Token Autoregressive Generation: Many strong continuous-token autoregressive methods operate in latent space using a pretrained tokenizer or autoencoder to map raw signals into lower-dimensional representations.PRA instead learns low-dimensional intermediate states end-to-end with the autoregressive model and avoids a separately pretrained tokenizer.
  • Pixel-Space Generative Modeling: Pixel-space image generation has been dominated by diffusion models that generate images through iterative denoising directly in the original pixel space.PRA differs by targeting pixel-space continuous-token autoregressive generation while retaining a pixel-in, pixel-out external interface.

3 Pixel-Space Image Generation via Continuous-Token AR Modeling

The method models images as raster-ordered continuous patch tokens with a causal Transformer and token-level diffusion head. During inference, generated tokens are sequentially fed back as inputs, unlike teacher-forced training on ground-truth prefixes.

  • Patch-wise continuous-token formulation: Images are partitioned into patches, each represented as one continuous token arranged in raster-scan order.With patch size 16^2, each token is a 16 × 16 × 3 = 768-dimensional vector.
  • Autoregressive modeling: A causal Transformer encodes ground-truth prefixes under teacher forcing, with optional class labels or text prompts prepended as conditioning.The hidden state is h_i−1 = f_θ(x_<i).
  • Continuous-token prediction: A token-level diffusion head models each continuous next-token distribution conditioned on the Transformer hidden state.The head transforms a simple prior sample into the target token and is trained with a rectified-flow objective.
  • Autoregressive inference: At inference, the model conditions on previously generated tokens and samples each next token by integrating its learned velocity field from t = 0 to t = 1.Sampling starts from Gaussian noise and uses N Euler-Maruyama steps.
  • Autoregressive inference: The generated continuous token is fed back for subsequent autoregressive steps, creating an inference interface that differs from teacher-forced ground-truth conditioning.The feedback uses the final state produced by the diffusion head as the next input token.

4 Challenges of Pixel-Space Autoregressive Generation

Pixel-space autoregressive generation faces two coupled challenges: high-dimensional pixel tokens cause large single-step errors, while train–inference mismatch amplifies errors across autoregressive steps. PRA addresses both by generating compact intermediate states and constructing decoded inference-like inputs in parallel.

  • Output-side challenge: With 256 autoregressive tokens in both settings, increasing token dimensionality from 48 to 768 causes pixel-space AR performance to fall substantially behind JiT.The matched autoregressive length rules out a longer generation horizon as the explanation, identifying high-dimensional continuous-token generation as an output-side bottleneck.
  • Output-side challenge: 9.70 to 7.68: x-prediction improves high-dimensional AR performance over v-prediction, but a substantial gap to JiT remains.This comparison corresponds to A3 versus A4.
  • Input-side challenge: Generated-token feedback creates a train–inference mismatch because training uses ground-truth prefixes while inference uses generated prefixes whose errors influence later predictions.This mismatch amplifies errors across autoregressive steps.
  • Input-side challenge: 9.94 to 7.68: input noise injection improves FID, but independent perturbations cannot fully match structured, model-dependent inference-time rollout errors.The improvement is reported for A6 to A3.
  • PRA response: PRA jointly addresses both challenges by generating compact intermediate states and constructing decoded inference-like inputs independently in parallel.This preserves parallel training while making inputs better resemble generated tokens encountered during inference-time rollout.

5 Parallel Rollout Approximation

PRA addresses pixel-space AR’s coupled output and train–inference challenges by generating low-dimensional intermediate states and decoding them into pixel tokens. It also constructs inference-like decoded pixel prefixes independently across positions, retaining parallel training and a pixel-in, pixel-out interface.

  • Output-side approximation: PRA replaces direct high-dimensional pixel-patch generation with low-dimensional intermediate states decoded back into pixel-space tokens.This reduces each AR step’s generation difficulty while preserving pixel outputs.
  • Output-side approximation: The causal pixel decoder reconstructs each pixel token from its intermediate state and preceding intermediate states, constraining the state to retain pixel-output information.For token xi, the reconstruction is xrec_i = Dϕ(zi, z<i) ≈ xi, with dz < d.
  • Target construction: PRA learns context-aware intermediate targets end-to-end from the current token and its causal-prefix representation, rather than using an external VAE or tokenizer.The target is zi = gψ(xi, hi−1), where hi−1 = fθ(x<i) and dz < d; token masking encourages causal-prefix information to contribute.
  • Input-side approximation: PRA approximates inference-time pixel inputs by perturbing intermediate targets and decoding them through the same intermediate-state-to-pixel path used during inference.The resulting decoded tokens are closer to inference-time generated inputs than clean ground-truth prefixes or independent pixel-space noise.
  • Input-side approximation: Because decoded inputs are constructed independently for each position, PRA trains in parallel while conditioning the AR model on reconstructed prefixes instead of clean prefixes.The reconstructed sequence is stop-gradient input, and the token-level diffusion head predicts the target intermediate states.
  • Inference pipeline: During inference, PRA autoregressively feeds decoded pixel tokens back into the causal context while using low-dimensional states only internally.The external interface remains pixel-in and pixel-out, despite the internal intermediate-state generation path.

6 Experiments

On class-conditional ImageNet-1K at 256×256 resolution, PRA improves pixel-space autoregressive generation quality across model scales and validates its intermediate-target and inference-like-input designs through ablations. PRA-L also achieves strong ImageNet linear-probing performance, indicating potential beyond generation.

  • Generation Results: 2.58 FID: PRA-S with 135M parameters outperforms the 1.9B-parameter FARMER-1.9B/8 baseline, which achieves 3.60 FID.PRA-B reaches 2.21 FID, while PRA-L reaches 1.94 FID, improving consistently with scale.
  • Output-Side Ablation: 2.88 FID: the prefix-aware learned target outperforms the local-only target, improving over 3.08 FID.The prefix-aware target uses the causal prefix representation z_i = gψ(x_i, h_i−1), whereas a frozen LDM encoder performs worse than the learned PRA target.
  • Output-Side Ablation: d_z = 16 gives the best intermediate-state performance, while d_z = 8 lacks reconstruction capacity and d_z = 32 or 64 makes continuous-token prediction harder.The intermediate-state dimension balances predictability against information preservation.
  • Input-Side Ablation: Clean ground-truth pixels perform poorly as AR inputs because inference instead conditions on pixels decoded from generated intermediate states.Decoded pixels can deviate substantially from real pixels, creating a strong input-side train–inference mismatch.
  • Input-Side Ablation: Noisy ground-truth intermediate-state inputs substantially improve performance over noise-free inputs, making training conditions closer to generated states.This improvement indicates that the intermediate space is easier for autoregressive modeling.
  • Linear Probing: 68.80% top-1 accuracy: PRA-L outperforms both SphereAR-L and JiT-L in ImageNet linear probing without architectural modifications or additional unsupervised pretraining.The frozen generation-trained backbone supports a linear classifier, demonstrating transfer to visual understanding.

7 Conclusion

PRA addresses high-dimensional pixel-space autoregressive generation errors and train–inference mismatch using low-dimensional intermediate states, pixel decoding, and parallel rollout approximation. It improves pixel-space AR generation, supports transferable visual representations, and leaves framework simplification and broader validation as future directions.

  • Method: PRA targets large single-step pixel-patch errors and cross-step error accumulation from train–inference mismatch.It learns low-dimensional intermediate targets end-to-end and maps them to pixel-space tokens with a pixel decoder.
  • Results: PRA substantially improves pixel-space AR generation and establishes a new state of the art among pixel-space AR models.Linear probing results also suggest transferable visual representations for unified pixel-space image generation and understanding.
  • Limitations: PRA adds an intermediate state, pixel decoder, and additional parallel AR forward during training, despite end-to-end training without a separately pretrained tokenizer.Simplifying the framework and validating it across broader data domains and generation tasks remain future directions.

A Model Details

PRA’s scale configurations specify diffusion-head depth, shared AdaLN groups, model widths, MLP structure, and training-cost reporting conditions. The diffusion-head width follows the model dimension, while the pixel-decoder width also sets the encoder width.

  • Architecture: PRA configurations report diffusion-head residual-block counts and shared AdaLN-group counts.The notation “res / AdaLN” refers to these two quantities.
  • Architecture: The diffusion-head width matches d, while pixel-decoder width dr also defines encoder gψ width.These widths determine the corresponding component dimensions.
  • Architecture: All blocks use an MLP ratio of 4 with SwiGLU.
  • Training: Training cost is reported as wall-clock days on 8×A100 GPUs for 400-epoch training.

A.1 Architecture Overview

PRA is an end-to-end trained architecture combining causal autoregression with intermediate-state encoding, pixel decoding, and rectified-flow modeling. It processes 256×256 RGB images as 256 non-overlapping 16×16 patches in raster-scan order.

  • Architecture Overview: PRA jointly trains four modules: a causal AR Transformer, intermediate-state encoder, pixel decoder, and diffusion head for rectified-flow modeling.The encoder produces intermediate states, the decoder maps noisy states to continuous pixel patches, and the diffusion head conditions their distribution on AR hidden states.
  • Architecture Overview: 256×256 RGB inputs are patchified into 256 non-overlapping 16×16 patches ordered by raster scan.Each patch corresponds to one autoregressive token.

A.2 Causal AR Transformer fθ … A.5 Diffusion Head vω

The model combines a causal Transformer backbone with intermediate-state encoding, causal pixel decoding, and a rectified-flow diffusion head. These components preserve causal conditioning while modeling low-dimensional states and reconstructing pixel patches.

  • A.2 Causal AR Transformer fθ: The causal Transformer uses pre-RMSNorm attention, SwiGLU layers, 2-D RoPE, FlashAttention, and 16 learnable class-conditioning prefix tokens.The prefix tokens are selected from a class-label embedding table and prepended to the patch sequence.
  • A.3 Intermediate-State Encoder gψ: The intermediate-state encoder maps each pixel patch and preceding hidden state to a 16-dimensional state using projected inputs, four SwiGLU residual blocks, LayerNorm, and AdaLN.Conditioning uses the normalized AR hidden state plus a learned position embedding, with AdaLN parameters shared every two residual blocks.
  • A.3 Intermediate-State Encoder gψ: Two-level masking encourages intermediate states to depend on causal context: selected examples use psample=0.9, and tokens are masked with ptoken=0.5.Masking replaces tokens with a learned mask embedding.
  • A.4 Pixel Decoder Dϕ: The causal pixel decoder maps noisy intermediate states to 16×16×3 = 768-dimensional pixel patches through projection, a separate-weight Transformer, RMSNorm, and a SiLU MLP head.Its causal design preserves KV caching.
  • A.5 Diffusion Head vω: The diffusion head models the rectified-flow distribution of each intermediate state conditioned on the normalized preceding hidden state and a learned position embedding.It is implemented as a SwiGLU MLP stack with shared AdaLN modulation.
  • A.5 Diffusion Head vω: For each flow timestep, the diffusion head modulates with a SiLU-transformed sum of sinusoidal timestep and conditioning embeddings, predicts velocity, and uses zero-initialized output and modulation layers.The modulation input is SiLU(tembed(s) + ci).

B Training Objective and Implementation Details · C Sampling Details

PRA combines parallel reconstruction and autoregressive training with decoded, inference-like pixel inputs, using intermediate-state perturbations, reconstruction objectives, and gradient controls. At inference, it generates patches sequentially with diffusion sampling, classifier-free guidance, normalization, and KV caching.

  • B Training Objective and Implementation Details: PRA alternates reconstruction and AR passes, using stop-gradient decoded pixel patches as causal Transformer inputs to approximate inference-time feedback during parallel training.Ground-truth patches are encoded, optionally perturbed, decoded, and supervised during reconstruction; decoded patches provide inputs for the AR pass.
  • B Training Objective and Implementation Details: During training, selected intermediate states receive Gaussian perturbations with ϵ_i ∼ N(0, I) and timesteps t_i ∼ U(t_min, 1).Perturbation probability is 0.9; t_min is 0.5 for the first 350 epochs and 0.7 for the remaining 50 epochs.
  • B Training Objective and Implementation Details: With encoder masking, clean no-mask encoder outputs define detached AR diffusion targets, while masked/noised paths train the decoder and construct decoded pixel inputs.The AR loss does not backpropagate through the detached target-producing encoder path.
  • B Training Objective and Implementation Details: The AR diffusion head uses rectified-flow interpolation with Gaussian noise and sigmoid-transformed timesteps, using μ=0 in the main experiments and decoded prefixes for conditioning.The conditioning prefix is computed from decoded pixels rather than clean pixels.
  • B Training Objective and Implementation Details: Reconstruction uses ℓ1 and LPIPS losses, while an auxiliary loss predicts per-patch normalized pixels; the main runs set λ_rec=1, λ_LPIPS=1, and λ_repr=1.A 0.3 gradient scale balances encoder-side reconstruction and auxiliary signals entering the shared AR backbone, and decoded inputs replace with noised ground-truth patches at probability 0.04.
  • C Sampling Details: At inference, generation proceeds sequentially in raster order: the Transformer consumes each prior generated patch, the diffusion head samples an intermediate state, and the decoder produces the next patch.KV caching is used for both the AR Transformer and causal pixel decoder.
  • C Sampling Details: 100 Euler–Maruyama steps plus a final deterministic Euler step evolve Gaussian noise from s=0 to s=1, with classifier-free guidance, linear positional scheduling, and normalized decoder inputs.Conditional and null-label predictions are evaluated in the diffusion head before applying the guidance schedule.

D Training Hyperparameters

The main ImageNet-1K models use a fixed 400-epoch AdamW training setup with scheduled optimization, EMA weights, and standard architectural defaults.

  • 400 epochs and a global batch size of 512 define the main ImageNet-1K training regime.
  • AdamW uses a peak learning rate of 3×10−4, β1=0.9, β2=0.95, weight decay 0.05, and gradient clipping at norm 1.0.
  • Optimization follows a cosine learning-rate schedule with 20K warmup steps and EMA weight decay 0.9999.
  • Unless specified otherwise, models use class dropout 0.1, 16 class prefix tokens, patch size 16, and intermediate dimension dz=16.
Loading 2606.27978v1…