Accelerate image generation by generating multiple patches in parallel instead of sequentially. Uses locality-aware ordering that generates spatially-close tokens while keeping concurrent tokens far apart, reducing steps from 256 to 20 (12× speedup) while maintaining quality.
Accelerate image generation by generating multiple patches in parallel instead of sequentially. Uses locality-aware ordering that generates spatially-close tokens while keeping concurrent tokens far apart, reducing steps from 256 to 20 (12× speedup) while maintaining quality.
Autoregressive image generation processes tokens sequentially: generate position 1, then 2, then 3... up to 256 positions for 256×256 images. This is memory-bound—computation finishes fast but bandwidth limits throughput. Generating 256 sequential steps creates 256 network latencies that compound into unacceptable delay for interactive applications.
Locality-aware Parallel Decoding generates multiple positions simultaneously by recognizing that spatial locality dominates attention patterns. Tokens far apart in the image don't interact strongly. By generating spatially-distant tokens in parallel, we reduce sequential steps from 256 to 20 while maintaining quality—a 12× speedup with competitive visual output.
Core Concept
Autoregressive generation must balance two objectives:
Strong Conditioning: Generate from strong context (nearby already-generated tokens)
Low Dependencies: Concurrent tokens shouldn't depend on each other (parallelizable)
These objectives conflict: nearby tokens have strong dependencies; distant tokens are independent but poorly conditioned. Locality-aware Parallel Decoding optimizes this trade-off by:
Locality-Aware Ordering: Generation schedule prioritizes tokens spatially close to already-generated context (strong conditioning) while keeping concurrent tokens far apart (minimal mutual dependency)
Position Query Tokens: Learnable tokens representing "what to generate at position X" decouple generation targets from conditioning role
Flexible Attention Masks: During both training and inference, special attention patterns enable tokens to see context but not each other
This enables parallel generation of, say, 12 tokens per step instead of 1, reducing 256 steps to 20.
Architecture Overview
The system consists of these components:
Position Query Token Generator: Learnable embeddings representing "generate at this position"
Flexible Autoregressive Attention: Attention masks allowing query tokens to see context but not other concurrent queries
Locality-Aware Scheduling Algorithm: Computes optimal generation order balancing conditioning strength and parallelizability
Proximity Threshold System: Dynamically adjusts which tokens can be generated in parallel
Farthest-Point Sampling: Selects maximally-separated tokens for concurrent generation
Benchmark Testing: Evaluation on token-to-image and patch-to-image generation
Implementation
This section demonstrates how to implement locality-aware parallel decoding.
Step 1: Design position query tokens for flexible generation
This code implements learnable tokens for parallel position generation:
import torch
import torch.nn as nn
import math
classPositionQueryTokens(nn.Module):
"""
Learnable tokens representing "generate at position X".
Decouples what positions to generate from how to condition generation.
"""def__init__(self, max_positions=256, embed_dim=768):
super().__init__()
self.max_positions = max_positions
self.embed_dim = embed_dim
# Learnable position queriesself.position_queries = nn.Parameter(torch.randn(max_positions, embed_dim))
# Positional encoding for spatial awarenessself.spatial_pos_encoding = self._create_spatial_encoding(max_positions, embed_dim)
def_create_spatial_encoding(self, max_positions, embed_dim):
"""Create 2D spatial positional encodings."""# Assume square grid: sqrt(max_positions) × sqrt(max_positions)
grid_size = int(math.sqrt(max_positions))
# Create spatial position embeddings
spatial_encoding = torch.zeros(max_positions, embed_dim)
for pos inrange(max_positions):
row = pos // grid_size
col = pos % grid_size
# Standard sinusoidal encoding but on 2D coordinatesfor d inrange(embed_dim // 4):
spatial_encoding[pos, 2*d] = math.sin(row / (10000 ** (2*d / embed_dim)))
spatial_encoding[pos, 2*d + 1] = math.cos(col / (10000 ** ((2*d + 1) / embed_dim)))
return spatial_encoding
defget_position_queries(self, positions: torch.Tensor) -> torch.Tensor:
"""
Retrieve position queries for given positions.
positions: (B, num_concurrent) indices of positions to generate
Returns: (B, num_concurrent, embed_dim) learnable queries
"""
B, num_concurrent = positions.shape
# Get learnable queries for these positions
queries = self.position_queries[positions] # (B, num_concurrent, D)# Add spatial encoding to make queries position-aware
spatial_emb = self.spatial_pos_encoding[positions]
queries = queries + spatial_emb.unsqueeze(0).expand(B, -1, -1)
return queries
# Test position query tokens
pos_queries = PositionQueryTokens(max_positions=256, embed_dim=768)
# Generate 12 positions concurrently
batch_size = 1
positions = torch.tensor([[10, 25, 40, 55, 70, 85, 100, 115, 130, 145, 160, 175]])
queries = pos_queries.get_position_queries(positions)
print(f"Position query shape: {queries.shape}")
print(f"Queries learn what to generate at each position")
This implements learnable tokens for flexible parallel generation.
Step 2: Define flexible autoregressive attention with custom masks
This code enables tokens to see context but not each other:
classFlexibleAutoregressiveAttention(nn.Module):
"""
Attention mechanism allowing parallel generation while maintaining autoregressive properties.
- Queries see all context
- Queries don't see each other
- New tokens can't see future positions
"""def__init__(self, embed_dim=768, num_heads=8):
super().__init__()
self.attention = nn.MultiheadAttention(
embed_dim=embed_dim,
num_heads=num_heads,
batch_first=True,
dropout=0.1
)
self.embed_dim = embed_dim
defcreate_attention_mask(
self,
context_positions: torch.Tensor,
query_positions: torch.Tensor,
grid_size: int) -> torch.Tensor:
"""
Create attention mask for parallel generation.
context_positions: indices of already-generated tokens
query_positions: indices of tokens being generated now
grid_size: sqrt(total_positions) for 2D layout
Mask properties:
- Queries attend to all context positions (autoregressive)
- Queries don't attend to other queries (parallel-safe)
- Queries attend to "self" (their own position representations)
"""
max_pos = grid_size * grid_size
num_context = len(context_positions)
num_queries = len(query_positions)
# Attention matrix: (num_queries, num_context + num_queries)
attention_mask = torch.zeros(num_queries, num_context + num_queries)
# Queries can attend to all context
attention_mask[:, :num_context] = 1.0# Queries can only attend to their own "self" query, not othersfor q_idx inrange(num_queries):
# Each query attends to its own representation
attention_mask[q_idx, num_context + q_idx] = 1.0# Convert to additive mask format for transformer (0 = attend, -inf = mask)
attention_mask = (1.0 - attention_mask) * -1e9return attention_mask
defforward(
self,
context_features: torch.Tensor,
position_queries: torch.Tensor,
context_positions: torch.Tensor,
query_positions: torch.Tensor,
grid_size: int) -> torch.Tensor:
"""
Generate features for query positions using context.
context_features: (B, num_context, D) - already-generated features
position_queries: (B, num_queries, D) - learnable position queries
"""
B = context_features.shape[0]
num_context = context_features.shape[1]
num_queries = position_queries.shape[1]
# Concatenate context and queries as KV# Queries only use context as Key/Value
kv = context_features # (B, num_context, D)# Create attention mask
attn_mask = self.create_attention_mask(
context_positions,
query_positions,
grid_size
).to(context_features.device)
# Apply attention: queries attend to context
output, _ = self.attention(
position_queries, # Query: what to generate
kv, # Key/Value: conditioning context
kv,
attn_mask=attn_mask
)
return output # (B, num_queries, D) - features for new positions# Test flexible attention
flex_attn = FlexibleAutoregressiveAttention()
context_feats = torch.randn(1, 50, 768) # 50 context tokens
pos_queries = torch.randn(1, 12, 768) # 12 concurrent queries
context_pos = torch.arange(50)
query_pos = torch.tensor([100, 110, 120, 130, 140, 150, 160, 170, 180, 190, 200, 210])
output = flex_attn(context_feats, pos_queries, context_pos, query_pos, grid_size=16)
print(f"Attention output shape: {output.shape}")
This implements attention allowing parallel generation with autoregressive guarantees.
Batch processing where throughput matters more than latency
Autoregressive patch/token-based image models
Applications tolerating slight quality trade-offs for speed
Scenarios with consistent grid-like generation patterns
When NOT to use:
Quality-critical applications where sequential generation is necessary
Non-grid-based generation (causal structures without spatial locality)
Models where token dependencies don't follow spatial patterns
Real-time systems with extreme latency requirements (other optimizations needed)
Scenarios where position queries add excessive overhead
Hyperparameters and Configuration:
Parameter
Recommended
Notes
Tokens per Step
12-20
Balance between parallelism and conditioning quality
Proximity Threshold
4-6 (in pixel distance)
Tokens too close compete; too far breaks conditioning
Grid Size
16×16 (256 tokens)
Standard for efficient transformers; larger grids → more parallelism possible
Position Query Dimension
768
Match model embedding dimension
Attention Heads
8-12
Standard for multi-head attention
Learning Rate
1e-4
Fine-tune from pretrained models
Spatial Encoding Frequency
10000
Standard sinusoidal encoding scale
Common Pitfalls:
Generating too many tokens in parallel (breaks conditioning, quality suffers)
Setting proximity threshold too high (wastes context information)
Not training position queries end-to-end (random initialization underutilizes parallel benefits)
Ignoring spatial locality in attention masks (defeats the purpose)
Using grid-incompatible models (method assumes 2D spatial structure)
Over-relying on greedy scheduling (suboptimal token selection)
Key Design Decisions:
Position query tokens decouple "what positions to generate" from "how to condition generation," enabling flexible parallelism. Locality-aware scheduling balances two competing objectives: tokens close to context are well-conditioned, while tokens far apart can be generated in parallel. Flexible attention masks implement this by allowing queries to see all context but not each other. The method achieves 12× speedup by reducing 256 sequential steps to 20 parallel batches while maintaining quality through intelligent ordering.
Reference
Shen, Y., Cai, B., Jiao, X., Zhang, Y., Zhang, T., Zhu, K., ... & Yan, H. (2025). Locality-aware Parallel Decoding for Efficient Autoregressive Image Generation. arXiv preprint arXiv:2507.01957. https://arxiv.org/abs/2507.01957
