Spanning Tree Autoregressive Visual Generation

AI summary

Overview

Research area: Computer vision, specifically autoregressive (AR) image generation with Decoder-only Transformers, and its extension to region-level image editing (inpainting).

Technical level: Advanced. The paper assumes familiarity with autoregressive likelihood factorization, causal attention, tokenizers/VQ autoencoders, spanning trees, and generative evaluation metrics (FID, IS, Precision, Recall).

Scope: The paper proposes a structured sequence-order randomization strategy for autoregressive image generation based on breadth-first search (BFS) traversal of uniform spanning trees on a patch lattice, and evaluates it for class-conditional ImageNet-1k generation at 256×256 and for image inpainting.

What This Paper Is About

Conventional autoregressive image models flatten a 2D image into a 1D sequence using a fixed raster-scan order, which forces a single unidirectional prediction dependency and makes native image editing (region-level modification) difficult. The alternative of training with random permutation exposes the model to arbitrary sequence orders, but it degrades sampling performance, and prior work tends to trade either performance or flexibility in the sequence order available at inference time. STAR aims to keep the sampling performance of raster-scan AR while retaining enough flexibility in sequence order to complete masked regions natively as postfix completion.

Key Contributions

1. A structured sequence-order randomization strategy for AR visual generation based on BFS traversal orders of uniformly sampled spanning trees over the lattice of patch positions, intended to preserve both sequence-order flexibility and sampling performance.

2. A demonstration that BFS traversal orders of uniform spanning trees are sufficient to capture the benefits that random permutation provides for diverse sequence orders in AR visual generation, without increasing training complexity.

3. A new building block (a rejection-sampling procedure for choosing training/inference sequence orders) that supports region-level modification capabilities for image editing with minimal changes to the conventional Decoder-only Transformer.

4. Empirical evidence across class-conditional ImageNet-1k generation and image inpainting, including analysis of conditional model entropy, locality, center bias, traversal choice (BFS vs. DFS), and sampling overhead.

Main Findings

• Generation quality comparable to state of the art with a simple change: On class-conditional ImageNet-1k at 256×256, STAR-B (261M) reaches FID 2.24 / IS 295.3, STAR-L (461M) FID 1.98 / IS 322.0, STAR-XL (955M) FID 1.65 / IS 333.2, and STAR-XXL (1.5B) FID 1.55 / IS 338.8, with Precision 0.82, 0.82, 0.80, 0.81 and Recall 0.57, 0.58, 0.62, 0.62. The paper reports that RAR-XXL (1.5B) achieves FID 1.48 / IS 326.0 / P 0.80 / R 0.63, so RAR-XXL has a lower FID while STAR-XXL has a higher IS.

• Outperforms random-permutation AR at matched scales: STAR-B (FID 2.24, IS 295.3) is reported better than RandAR-L (FID 2.55, IS 288.8), and STAR-XL (FID 1.65, IS 333.2) better than RandAR-XL (FID 2.25, IS 314.2).

• Much stronger image inpainting than RAR and DiT: Averaged over masking ratios 0.1 to 0.9 in 0.1 increments, STAR-L obtains inpainting FID 2.20 / IS 110.1 and STAR-XL FID 2.07 / IS 111.3, versus RandAR-L 2.66 / 58.8, RandAR-XL 2.58 / 60.3, RAR-L 3.56 / 86.7, RAR-XL 3.40 / 88.2, and DiT-XL/2 4.58 / 50.4. The paper attributes RAR's and DiT's drop to exposing bidirectional dependency only during training, or only through noisy partial observations in reverse-time diffusion.

• Parameter scaling helps inpainting at mid-to-high masking ratios: The paper reports that scaling parameter size significantly improves FID at 0.4–0.6 masking ratios, and speculates that sufficient sample diversity (reflected in the Recall trend) helps the model use ambiguous but informative partial observations.

• Structured randomization beats both permutation and order annealing at the same budget: In ablations at the STAR-B budget, raster-scan three-stage training gives generation FID 2.04 / IS 266.5 and inpainting FID 3.91 / IS 82.7; full permutation gives generation FID 3.57 / IS 291.4 and inpainting FID 2.57 / IS 108.7; permutation training annealed to spanning-tree orders gives generation FID 2.31 / IS 285.2 and inpainting FID 2.35 / IS 107.9; the default spanning-tree strategy gives generation FID 2.24 / IS 295.3 and inpainting FID 2.39 / IS 108.8.

• Prior knowledge is visible in token representations: Measuring conditional model entropy on ImageNet-1k validation images using 30 sampled sequence orders per image for each AR model, the authors find entropy increases with distance from predicted prefix tokens and is higher in central regions, suggesting that connected and adjacent context matters most.

• BFS is far more efficient than DFS for postfix completion: Testing 50k random connected masks at masking ratios 0.1–0.9 (trials bounded at 100, exceeding which counts as failure), BFS with a farthest-root strategy requires on average 4.41, 3.03, 2.50, 1.95, 1.57, 1.51, 1.43, 1.27, and 1.15 trials with 0.0% failure at every ratio. DFS with random root requires 60.45 to 2.22 trials with failure rates from 50.1% down to 0.0%, and DFS with farthest root is worse (82.3% to 0.0% failure).

• Center bias behavior: Dividing ImageNet classes into four quartiles by strength of center bias, STAR is reported comparable to raster-scan under high center bias and consistently better than random permutation at all levels, with performance degrading as center bias weakens for all sequence-order choices.

• Negligible sampling overhead: Sampling a uniform spanning tree and running BFS costs 0.56 ms per instance on average, which the paper states is less than 0.0004% of total inference time across the four STAR configurations.

• Root-cause analysis of sequence order quality: Conditional model entropy across token positions is relatively constant for raster-scan and for spanning-tree BFS orders, whereas random permutation shows high entropy at initial positions that decays quickly after half the tokens; both the smoother cases perform better, which the paper links to uniform prediction difficulty aiding convergence.

Methodology in Plain English

The researchers treat each patch position of an image as a vertex in a regular two-dimensional grid, where neighboring positions are connected by edges. Instead of either fixing one ordering (raster-scan) or shuffling tokens completely at random, they randomly draw a spanning tree of that grid — a connected acyclic subgraph that touches every vertex — and then linearize it by breadth-first search starting from a randomly chosen corner. This gives a family of orderings that always predicts tokens adjacent to what has already been seen and always starts at a corner, which matches known image statistics such as locality and center bias, while still providing many different orders.

Training is standard maximum likelihood next-token prediction, with one addition shared by permutation-based approaches: the model is conditioned on the position of the next token to be predicted (here, using learnable positional embeddings). The architecture otherwise stays a Decoder-only Transformer with a causal mask.

For inference-time editing, the key requirement is that the known (unmasked) part of the image must appear as a prefix and the missing region as a postfix. Because BFS orders group tokens by depth, the authors enforce this by sampling spanning trees for the unmasked region until the maximum-depth vertices lie on every boundary next to a masked region (rejection sampling), then attaching spanning trees inside each masked region. They choose the corner farthest from the nearest mask boundary as the root to increase acceptance, and a fixed tie-breaking rule during BFS.

Experiments use a MaskGIT-VQGAN tokenizer (reported as a 1024-codebook size and a 16 downsampling rate), a ViT backbone with causal attention, QK LayerNorm, and class conditioning via adaLN, following RAR's configuration. Models are trained for 250k steps with batch size 2048 (approximately 400 epochs) using AdamW, a cosine learning rate scheduler with linear warmup, Wilson's algorithm for spanning-tree sampling, class dropout of 0.1 for classifier-free guidance, and a power-cosine CFG scheduler without top-p or top-k sampling. Evaluation follows the ADM protocol with FID, IS, Precision, and Recall, extended to inpainting by dropping a random connected set of tokens at given masking ratios.

Why This Matters

Impact on research: The paper expands the trade-off frontier between sampling quality and inference-time sequence-order flexibility in autoregressive visual generation, showing that restricting the space of random orderings to a locality-preserving, exponentially sized family (the number of spanning trees grows as exp(N·z_G) with z_G ≈ 1.166 for the regular taxicab lattice, versus N! permutations) can recover raster-scan-level performance while enabling native inpainting. It also offers a diagnostic lens — conditional model entropy across token positions and distances — for explaining why some orderings converge faster than others.

Real-world applications:

• Interactive photo editing with object removal and content fill, where a user marks a region and the model completes it consistently with the surrounding context.
• Image restoration and retouching workflows (for example, repairing damaged scans or removing unwanted objects) that require coherence between retained and generated pixels.
• Unified multimodal systems that share one Decoder-only Transformer for text and images, where editing must be expressible as a single generation call rather than a separate architecture.
• Product, advertising, and e-commerce pipelines needing class- or condition-controlled image synthesis at 256×256 scale with region-level adjustments.

Industry relevance: The work is co-authored across LG AI Research, KIST, Yonsei University, Seoul National University, and the University of Michigan, Ann Arbor, and the code is released publicly. Because it keeps the standard Decoder-only Transformer and adds only next-token position conditioning, it is directly relevant to teams already training language-style autoregressive stacks and to those building early-fusion multimodal models, where architectural simplicity is a practical constraint. The measured 0.56 ms per-instance overhead and the 0.0% postfix-completion failure rate for BFS make the editing path cheap to deploy in practice.

Future Directions

• Combining STAR's spanning-tree sequence orders with a diffusion loss and a continuous VAE tokenizer, which the paper explicitly flags as a promising extension given that MAR uses random permutation with diffusion loss.
• Testing alternative ways of conditioning on the next-token position, such as RoPE or adaLN, which the authors state they did not evaluate (they used learnable positional embeddings).
• Handling masking scenarios that violate the paper's assumptions — masks that cover all corners or that disconnect the partial observation — which the authors note are discussed in the supplementary material.
• Extending evaluation beyond 256×256 ImageNet-1k class-conditional generation, since all reported generation and inpainting experiments in the main text use that setting; no other resolutions, datasets, or modalities are reported.

Target Audience

Researchers and engineers working on autoregressive image generation, Decoder-only Transformer architectures, multimodal and early-fusion models, and image editing/inpainting. It is also useful for readers interested in how sequence order choices interact with convergence behavior and conditional entropy in generative models. The paper is best suited to readers already comfortable with generative modeling metrics and AR likelihood training; newcomers will find the spanning-tree and rejection-sampling portions challenging without background reading.

Authors’ abstract

We present Spanning Tree Autoregressive (STAR) modeling, which can incorporate prior knowledge of images, such as center bias and locality, to maintain sampling performance while also providing sufficiently flexible sequence orders to accommodate image editing at inference time. Approaches that expose conventional autoregressive (AR) models in visual generation to arbitrary sequence orders via random permutation suffer from degraded sampling performance or compromise the flexibility in sequence order choice at inference time. Instead, STAR utilizes traversal orders of uniform spanning trees in a lattice defined by the positions of image patches. Traversal orders are obtained via breadth-first search, allowing us to efficiently construct a spanning tree via rejection sampling whose traversal order ensures that the connected partial observation of the image appears as a prefix for native image inpainting support. Through the tailored yet structured sequence order randomization strategy, STAR preserves the capability of postfix completion while maintaining sampling performance, without any significant changes to the model architecture widely adopted in language AR modeling.