Technical Summary: Spanning Tree Autoregressive Visual Generation (STAR)

Problem Statement

Autoregressive (AR) models based on Decoder-only Transformers have achieved remarkable success in language generation and are increasingly applied to visual generation via next-patch prediction. However, conventional AR visual models typically rely on a fixed raster-scan order (unidirectional dependency). While this order yields high sampling performance, it imposes a structural limitation: the model cannot condition on tokens that do not appear as prefixes in the predefined sequence. This rigidity hinders native image editing (e.g., inpainting), where arbitrary regions of an image must be masked and completed based on the remaining context.

To address this, recent approaches employ random permutation of token sequences during training to expose models to diverse orders. However, the paper notes that training with arbitrary sequence orders often leads to degraded sampling performance compared to fixed raster-scan models. This creates a trade-off: fixed orders offer performance but lack flexibility, while random permutations offer flexibility but sacrifice generation quality and convergence speed.

Methodology: Spanning Tree Autoregressive (STAR)

The authors propose STAR, a structured sequence order randomization strategy that bridges this gap by utilizing the traversal orders of Uniform Spanning Trees (UST) on the lattice defined by image patch positions.

1. Structural Randomization via Uniform Spanning Trees

Instead of sampling from the space of all $N!$ permutations, STAR samples from the space of spanning trees on a regular taxicab lattice $G=(V, E)$, where vertices represent patch positions.

• Root Selection: A random corner of the image is selected as the root node $r$. This choice is motivated by the observation that raster-scan orders (starting from corners) often outperform others, likely due to center bias (salient objects are less likely to appear in corners) and locality (adjacent patches are easier to predict).
• Traversal Order: The sequence order $\tau$ is determined by a Breadth-First Search (BFS) traversal of the sampled spanning tree.
• Theoretical Basis: The authors hypothesize that the advantage of predefined sequences stems from preserving prior knowledge (center bias and locality). They verify this by measuring conditional model entropy, showing that tokens closer to previously predicted tokens (in terms of tree distance) exhibit lower entropy, similar to raster-scan behavior. The number of spanning trees grows exponentially ($\sim \exp(N \cdot z_G)$) but remains a much smaller, more structured subset of the full permutation space.

2. Inference and Postfix Completion (Image Inpainting)

A critical challenge for AR models in inpainting is ensuring that a given partial observation (the non-masked region) appears as a prefix in the sequence order, allowing the masked region to be generated as a postfix.

• Rejection Sampling: Unlike random permutation, where prefix/postfix separation is trivial, STAR must sample a spanning tree where the BFS traversal naturally places the non-masked region before the masked region.
• Algorithm:
  1. Identify the boundaries of the non-masked region.
  2. Select a root node $r$ that is the farthest (in Manhattan distance) from the closest boundary of the non-masked region to maximize the likelihood of the non-masked region being visited first.
  3. Sample a spanning tree in the non-masked region. If the maximum depth of the tree does not satisfy the condition that all boundary vertices are visited before any masked vertex, the sample is rejected (rejection sampling).
  4. Once a valid tree is found for the non-masked region, spanning trees are sampled for the masked regions and connected to the boundary vertices to complete the global tree.
• Efficiency: The authors demonstrate that using BFS with "farthest root" selection significantly reduces the number of rejection trials compared to Depth-First Search (DFS), making inference efficient.

3. Model Architecture

STAR requires minimal changes to the standard Decoder-only Transformer architecture. It incorporates:

• Causal attention masks.
• Additional features to condition on the next-token position (using learnable positional embeddings), similar to previous permutation-based methods.
• No complex architectural changes (e.g., bidirectional attention or specialized tokenizers) are required.

Key Contributions

1. Structured Randomization Strategy: The proposal of using BFS traversal of uniform spanning trees as a sequence order for AR visual generation. This strategy structurally preserves image priors (center bias, locality) while introducing sufficient randomness for flexibility.
2. Performance-Flexibility Trade-off: The demonstration that BFS traversal of USTs is sufficient to achieve the benefits of random permutation (diverse sequence orders) without the performance degradation typically associated with arbitrary permutations. It achieves this without increasing training complexity.
3. Native Image Editing Support: The introduction of a mechanism (rejection sampling with BFS) that enables region-level modification (inpainting) with minimal architectural changes, addressing a key limitation of conventional AR models.

Experimental Results

The authors evaluate STAR on ImageNet-1k (256x256 resolution) using class-conditional generation and image inpainting tasks.

• Image Generation: STAR models (STAR-B, L, XL, XXL) achieve performance comparable to state-of-the-art models (e.g., RAR, DiT, LlamaGen) in terms of Fréchet Inception Distance (FID) and Inception Score (IS). Notably, STAR-XXL achieves an FID of 1.55, outperforming RAR-XXL (1.48) and RandAR-XXL (2.15), while maintaining competitive scores.
• Image Inpainting: STAR significantly outperforms both fixed-order models (RAR) and random-permutation models (RandAR) in inpainting tasks.
  • At various masking ratios (0.1 to 0.9), STAR maintains high FID and IS scores.
  • In contrast, RAR (fixed raster-scan) shows a significant drop in inpainting performance because it cannot naturally handle arbitrary masks as prefixes.
  • RandAR maintains inpainting capability but suffers from lower generation quality compared to STAR.
• Ablation Studies:
  • Locality: Random permutation fails to outperform STAR, suggesting that sequence orders violating locality are detrimental to convergence.
  • Center Bias: Randomly selecting corners as roots works well for images with strong center bias and generalizes better than random permutation for images with weaker center bias.
  • Entropy Analysis: Models trained with STAR exhibit a more uniform scale of conditional model entropy across token positions compared to random permutation, which correlates with better convergence.
  • Overhead: The computational overhead for sampling spanning trees is negligible (< 0.0004% of total inference time). Rejection sampling for inpainting is highly efficient, typically requiring fewer than 5 trials with BFS.

Significance and Claims

The paper claims that STAR offers a simple yet effective solution to the long-standing trade-off between sampling performance and sequence order flexibility in AR visual generation. By leveraging the structural properties of uniform spanning trees, STAR:

• Preserves the Decoder-only Transformer architecture, facilitating integration into early-fusion multimodal models.
• Avoids the need for complex training strategies like order annealing (switching from random to fixed orders during training) or specialized tokenizers.
• Provides native support for image editing (inpainting) without sacrificing generation quality.

The authors position STAR as a "basic building block" for future AR visual modeling, emphasizing that its simplicity allows it to be easily adopted and extended to multimodal generative models. The work suggests that the success of AR models relies not just on the model architecture, but on the structured randomization of the sequence order to align with the intrinsic properties of visual data.