Unlocking Any-Order Generation in Pretrained Autoregressive Image Models

AI summary

Overview

• Research area: Computer vision and generative modeling, specifically autoregressive (AR) image generation and any-order (permutation-based) sequence modeling.
• Technical level: Advanced. The paper assumes familiarity with autoregressive factorization, positional embeddings, KV caching, and discrete-diffusion baselines.
• Scope: The paper proposes a two-part adaptation recipe that converts an already-pretrained raster-scan AR image generator into an any-order generator capable of inpainting, outpainting, and region-conditioned editing, evaluated on ImageNet-1k with LlamaGen-L/XL and RAR-L/XL.

What This Paper Is About

Standard autoregressive image models generate tokens in a fixed raster-scan order (left to right, top to bottom), which hardcodes a single factorization of the image distribution and prevents the model from conditioning on non-contiguous context such as an inpainting mask or an outpainting boundary. Existing any-order alternatives, notably discrete diffusion and permutation-trained AR models, must be trained from scratch at roughly 3 times the compute of raster-scan training and often lose compatibility with the AR inference stack. The paper's goal is to obtain any-order generation by lightly adapting existing pretrained raster-scan checkpoints instead of training new models from scratch.

Key Contributions

1. A diagnosis of two distinct obstacles. The authors separate an identifiability problem (under varying order, identical context maps to different prediction targets, producing conflicting gradients) from an optimization problem (naively training on uniform random permutations at the pretrained initialization destroys the pretrained representations before order-flexible routing develops).
2. Future-aware positional embeddings. A learned additive table E of shape n × d that injects the identity of the next target position into the input representation at each step, at the input layer only, adding (n+1)·d parameters — 0.74M (0.095%) for LlamaGen-XL — while leaving the single-stream causal transformer, KV caching, and parallel decoding intact.
3. An ordering curriculum. Permutations are sampled from a mixture P_α = (1−α)·δ_raster + α·Uniform(S_n) with a linear schedule α(s) = min(s/S_ramp, 1), tracing a continuation path from the pretrained solution to a good any-order basin.
4. An adaptation recipe validated across four base models. The recipe is applied to LlamaGen-L/XL and RAR-L/XL, producing models that perform inpainting, outpainting, and region completion as zero-shot instances of conditional generation, with training cost of 1.13–1.33× the pretraining budget.

Main Findings

• Inpainting beats MaskGIT. On ImageNet 256×256 with 50,000 samples, the adapted model reaches FID 6.21 versus 10.44 for MaskGIT (cosine) and 7.82 for MaskGIT (Halton), and is lower than MAR-L's 6.63.
• Outpainting beats MaskGIT. FID 5.84 versus 12.17 (MaskGIT cosine) and 7.91 (MaskGIT Halton), against MAR-L's 6.86.
• Competitive but not dominant against MAR-L. MAR-L has higher outpainting Precision (0.736 vs 0.720) and Recall (0.649 vs 0.637); the authors describe the two as competitive in the absence of run-level uncertainty.
• Reference-based completion metrics favor the adapted model. LPIPS 0.125, PSNR 24.65 ± 3.34, SSIM 0.682 ± 0.17, against MAR-L's 0.174, 20.15 ± 4.60, and 0.550 ± 0.20.
• A measured quality–flexibility trade-off, not a lossless conversion. Class-conditional FID rises by between 1.17 and 1.93 across the four base models. On LlamaGen-XL, FID moves from 2.63 to 3.91 while IS (252.94), Precision (0.81), and Recall (0.53) stay close to the pretrained values. LlamaGen-L goes from 3.07 to 4.24, RAR-L from 1.70 to 3.63, and RAR-XL from 1.50 to 2.68.
• Pretraining matters at matched compute. With the same architecture and curriculum at matched total epochs, training from random initialization is worse: LlamaGen-L at 400 epochs (300 + 100) gives 4.24 for adapted versus 5.60 for random init; LlamaGen-XL at 125 epochs (100 + 25) gives 3.91 versus 4.81.
• Both recipe components are necessary. Removing future-aware embeddings makes the loss climb rather than descend as the curriculum ramps in. Removing the curriculum (α = 1 from step 0 on LlamaGen-L) gives FID 5.86, IS 191.27, Precision 0.78, Recall 0.51 versus 4.24, 231.62, 0.83, 0.52 with it — a fidelity collapse with Recall essentially unchanged.
• The embedding needs a higher learning rate. Scaling the future-aware embedding's learning rate by 10× relative to the rest of the network is reported as necessary for stable convergence.
• Teacher regularization backfires. Adding a KL term (λ = 1.0) toward the frozen raster-scan teacher raises LlamaGen-L FID from 4.24 to 9.62.
• Cost is a small fraction of from-scratch training. Adaptation costs 1.25× the pretraining budget on LlamaGen-XL (1.13–1.33× across models), against roughly 3× for from-scratch any-order training — about 2.4× lower end-to-end and over 12× lower as a marginal cost for practitioners who already hold a checkpoint.
• Parallel decoding comes for free. On LlamaGen-XL, dropping from 576 to 64 generation steps raises FID from 3.91 to 5.03 (192 steps: 4.12; 128 steps: 4.81); 9× fewer forward passes cost 1.12 FID.
• Inference latency is mixed. At maximum native-setting throughput on one RTX PRO 6000, the model takes 35.6 ms per image (batch 64), versus 56.3 ms for MAR-L and 23–24 ms for MaskGIT, whose 12-step decoders are faster but reach worse inpainting FID — not a matched-FID comparison.

Methodology in Plain English

The authors start from the observation that a raster-scan AR model has already learned the joint distribution over image tokens; the chain rule of probability says the same joint can be factored under any token ordering. So instead of teaching the model new knowledge, they teach it to express existing knowledge through arbitrary conditional pathways.

Two changes are made during a short continued-training phase on the ImageNet-1k training split (1.28M images, pre-tokenized with each model's own VQ tokenizer). First, at each step the model is told which spatial position it is about to predict, via an additive embedding added at the input layer; the same row is shared across sequence positions, and position indices used for RoPE (LlamaGen) or lookup tables (RAR) follow the spatial index rather than the sequence step, so the model's spatial-relational priors carry over unchanged. Second, rather than immediately training on fully random token orderings, the training mixture starts at pure raster scan and linearly ramps toward uniform random permutations, so earlier stages present small, low-variance gradients that keep the model near its pretrained solution while order-flexible routing gradually forms.

At inference, the user supplies a generation order: raster scan recovers the original behavior, random shuffles it, and task-induced orders place observed tokens (inpainting mask, outpainting boundary, partial edit) first — the remaining tokens are then sampled as ordinary conditional generation, with no task-specific training, sampler modification, or auxiliary network, and with KV caching applying unchanged.

Why This Matters

The paper reframes any-order capability as something to be unlocked from existing checkpoints rather than added by retraining, and it does so without abandoning the causal single-stream transformer. That makes flexible image generation reachable for groups that already have AR pretraining pipelines and AR serving infrastructure, rather than requiring a separate diffusion-style stack.

• Photo editing and restoration tools: filling in removed objects, extending a cropped photo to a wider frame, and regenerating a user-selected region, all from one model with observed pixels treated as context.
• Creative and design workflows: outpainting a subject into a larger scene while preserving identity and palette, as shown in the qualitative comparisons.
• Content pipelines with existing AR serving: organizations already running vLLM-based AR serving can add spatial position indices and the future-aware embedding via a custom model registration, keeping KV caching and tensor/pipeline parallelism.
• Latency-sensitive deployment: parallel decoding admits joint sampling of near-independent tokens in one forward pass, trading a small FID change for substantially fewer forward passes.

For industry, the practical lever is cost: adaptation at 1.13–1.33× the pretraining budget versus roughly 3× for from-scratch any-order training, and over 12× lower marginal cost for anyone who already holds a raster-scan checkpoint.

Future Directions

• Beyond class-conditional ImageNet at 256×256. The authors state explicitly that whether the recipe extends to text-to-image AR models or higher resolutions remains untested; this is the largest open question.
• Non-token-based AR variants. The identifiability diagnosis is claimed to apply to any causal any-order generator, but models like MAR, which predict continuous tokens via a diffusion head, lie outside the set of models actually adapted.
• Broader editing evaluation. The current editing study covers region completion only; instruction following and identity preservation are left as qualitative observations without metrics.
• A cleaner from-scratch control. The random-initialization comparison matches compute but not schedule, since only the pretrained pathway begins with raster-only training — a gap the authors flag themselves.

Target Audience

Researchers and engineers working on autoregressive image generation, multimodal generation stacks, or inference infrastructure for generative models. It is most useful to readers who already understand transformer-based AR modeling and want to know whether any-order or editing capability requires a new model family or can be retrofitted onto checkpoints they already have. Readers primarily interested in diffusion or flow-matching methods will find the comparison useful but the recipe itself out of scope by the authors' own framing.

Authors’ abstract

Autoregressive (AR) image generators are powerful in principle but inflexible in practice: training with a fixed raster-scan order hardcodes a single factorization of the image distribution into the model, leaving it unable to inpaint, outpaint, or edit at inference time. Existing solutions, including the modern discrete-diffusion family, train from scratch on a mixture of token orderings, paying roughly 3x the compute of standard raster-scan training; discrete diffusion further sacrifices compatibility with the AR inference ecosystem (KV caching, vLLM) by construction. We argue that this cost is unnecessary. A pretrained raster-scan AR model has already learned the joint distribution over image tokens; what it lacks is not knowledge but the ability to express that knowledge through non-sequential conditional pathways. We propose a recipe for adapting pretrained raster-scan models into any-order generators using two components: future-aware positional embeddings that resolve a target-position identifiability problem at negligible overhead, and an ordering curriculum that traces a stable optimization path from the pretrained solution to a good any-order basin. Applied to four base models across two architectures and two scales (LlamaGen-L/XL, RAR-L/XL), the recipe substantially outperforms MaskGIT on ImageNet inpainting (FID 6.21 vs. 10.44) and outpainting, is competitive with MAR-L, trades a 1.2-1.9 increase in class-conditional FID for this flexibility, and costs 13-33% additional compute beyond pretraining, while leaving the standard AR inference stack untouched. Inpainting, outpainting, and editing, our results suggest, need not be added to autoregressive image models. They can be unlocked from the ones that exist. Code is available at https://github.com/Ordered-AR/unlockedadaptation.