Distilling Specialized Orders for Visual Generation

1. Problem Statement

Autoregressive (AR) models have become dominant in language generation and are increasingly applied to image generation due to their scalability and high-quality output. However, traditional AR image generators suffer from a fundamental architectural constraint: they are locked into a fixed generation order (typically a raster-scan from top-left to bottom-right).

• Limitation: This fixed order acts as a rigid inductive bias, preventing the model from natively supporting flexible tasks like inpainting, outpainting, or region-based editing without task-specific retraining.
• The Trade-off: "Any-order" AR models attempt to solve this by learning to generate under arbitrary patch permutations, enabling multi-task flexibility. However, this comes at a significant cost: the model must distribute its finite capacity across all $N!$ possible permutations, often resulting in lower generation quality compared to specialized raster-scan models.
• Core Question: Can a single AR model achieve both flexible generation capabilities (like any-order models) and high generation quality (like specialized order models) simultaneously?

2. Methodology: Ordered Autoregressive (OAR)

The authors propose Ordered Autoregressive (OAR), a self-distillation pipeline that extracts specialized generation orders from a pre-trained any-order model. The process involves three main stages:

A. Training an Any-Order AR Model (AO-AR)

• Foundation: The authors first train a standard decoder-only AR model (e.g., RandAR) to predict tokens under arbitrary orders.
• Mechanism: The model is trained using random permutations of image patches. To handle arbitrary orders, the model receives explicit positional information for the next token to be generated (using either RandAR-style position tokens or $\sigma$-GPT-style double positional encoding).
• Goal: This stage creates a robust "teacher" model capable of understanding conditional probabilities for any subset of patches, regardless of spatial arrangement.

B. Specialized Order Extraction (Greedy Distillation)

• Concept: Instead of learning all $N!$ permutations, the method seeks a single, data-driven "specialized order" for each image that maximizes the likelihood of generation.
• Process:
  1. The trained AO-AR model acts as a teacher.
  2. For each training image, the model performs a greedy extraction of the generation order.
  3. At each step $i$, given the ground-truth prefix, the model evaluates all remaining patch locations. It selects the location $l^*_i$ that yields the highest maximum likelihood over the vocabulary:
     $$l^*_i = \arg \max_{l \in L_i} \left[ \max_{v \in V} p_\theta(v | l, x_{<i}, l_{<i}) \right]$$
  4. This creates a pseudo-labeled dataset of image-order pairs $(x, l^*)$, where $l^*$ is the specialized order derived from the model's own confidence.

C. Self-Distillation Fine-Tuning

• Fine-Tuning: The AO-AR model is fine-tuned on the extracted image-order pairs.
• Capacity Reallocation: By focusing on a single specialized trajectory per image, the model reallocates capacity previously spent on learning the vast permutation space to refine generation quality along this specific path.
• Preserving Flexibility: To prevent "catastrophic forgetting" of the any-order capabilities, the fine-tuning uses a reduced learning rate and an order-annealing schedule (gradually transitioning from mixed/random orders to purely specialized orders). This ensures the model retains the ability to handle arbitrary inputs (for inpainting/editing) while optimizing for quality on the specialized path.

D. Inference Optimization

• Parallel Querying: During inference, the model evaluates all possible next locations in parallel within a single forward pass using modified attention masks (preventing queries from attending to each other).
• Optimized KV-Caching: A cache overwriting strategy is employed to maintain $O(N)$ memory complexity, avoiding the quadratic growth typical of evaluating multiple candidates.
• Result: Despite evaluating multiple candidates, the inference time remains comparable to standard raster-scan AR models due to GPU parallelization.

3. Key Contributions

1. OAR Framework: A novel self-distillation pipeline that extracts specialized generation orders from an any-order AR model without architectural changes or additional annotations.
2. Theoretical Insight: Demonstrates that order specialization acts as capacity reallocation. By shifting focus from fitting $N!$ permutations to a single high-confidence trajectory, generation quality improves without sacrificing the multi-task flexibility inherent in any-order models.
3. Zero-Shot Flexibility: The method supports zero-shot inpainting, outpainting, and editing without retraining, as the underlying any-order conditionals are preserved.
4. Efficiency: Introduces parallel querying and optimized KV-caching techniques to make the greedy order selection computationally feasible.

4. Experimental Results

The method was validated on ImageNet (256x256), Fashion Products, and Multimodal CelebA-HQ.

• ImageNet Performance:
  • FID Improvement: OAR improved the Fréchet Inception Distance (FID) from 2.39 (Any-order baseline) to 2.17 (with annealing).
  • Comparison: It outperforms the standard Any-order baseline and is competitive with state-of-the-art raster-based AR models (e.g., RAR-XL) and diffusion models, while retaining the latter's flexibility.
  • Inpainting: Achieved an FID of 2.10 for inpainting (25-75% mask), slightly better than the baseline (2.14), proving the specialized order aids reconstruction.
• Human Evaluation:
  • In a forced-choice study across 21 ImageNet classes, OAR was preferred over the baseline in 64.33% of cases.
  • 16 out of 21 classes showed a preference for OAR, with 12 showing "strong preference" (>65% win rate).
• Other Datasets:
  • Fashion Products: FID improved from 4.07 (Random) to 2.56 (Fine-tuned Ordered).
  • CelebA-HQ: FID improved from 1.94 (Raster) to 1.41 (Fine-tuned Ordered).
  • Qualitative Analysis: The specialized orders learned by the model are semantically structured (e.g., generating high-confidence regions like cheeks/chin first, then fine-grained features like eyes/lips, and background last), rather than following a rigid raster scan.

5. Significance

This paper bridges the gap between flexibility and quality in autoregressive image generation.

• Paradigm Shift: It challenges the notion that high-quality generation requires a fixed raster order. Instead, it shows that the "best" order is data-dependent and can be learned.
• Practical Impact: OAR enables a single model to serve as a general-purpose vision generator capable of high-fidelity synthesis, editing, and completion without the need for multiple task-specific models or retraining.
• Scalability: The method is lightweight (requiring only fine-tuning) and compatible with existing AR optimizations (KV caching, parallel decoding), making it a viable path for scaling AR models in vision.

In summary, OAR demonstrates that by distilling the "wisdom" of an any-order model into specialized, confidence-driven generation paths, one can achieve state-of-the-art image quality while maintaining the versatile, zero-shot capabilities of any-order generation.