Locality Aware Parallel Autoregressive Decoding

Locality-aware Parallel Decoding: 12× Faster Image Generation

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:

1. Strong Conditioning: Generate from strong context (nearby already-generated tokens)
2. 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:

1. Locality-Aware Ordering: Generation schedule prioritizes tokens spatially close to already-generated context (strong conditioning) while keeping concurrent tokens far apart (minimal mutual dependency)
2. Position Query Tokens: Learnable tokens representing "what to generate at position X" decouple generation targets from conditioning role
3. 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

class PositionQueryTokens(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 queries
        self.position_queries = nn.Parameter(torch.randn(max_positions, embed_dim))

        # Positional encoding for spatial awareness
        self.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 in range(max_positions):
            row = pos // grid_size
            col = pos % grid_size

            # Standard sinusoidal encoding but on 2D coordinates
            for d in range(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

    def get_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:

class FlexibleAutoregressiveAttention(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

    def create_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 others
        for q_idx in range(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) * -1e9

        return attention_mask

    def forward(
        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.

Step 3: Implement locality-aware generation schedule

This code computes optimal generation order prioritizing proximity:

import numpy as np

def spatial_distance(pos1: int, pos2: int, grid_size: int) -> float:
    """Compute Euclidean distance between two positions in grid."""
    row1, col1 = pos1 // grid_size, pos1 % grid_size
    row2, col2 = pos2 // grid_size, pos2 % grid_size
    return math.sqrt((row1 - row2) ** 2 + (col1 - col2) ** 2)

class LocalityAwareScheduler:
    """
    Compute generation schedule maximizing conditioning strength while minimizing
    inter-token dependencies.
    """

    def __init__(self, grid_size=16, proximity_threshold=5.0):
        self.grid_size = grid_size
        self.total_positions = grid_size * grid_size
        self.proximity_threshold = proximity_threshold

    def compute_generation_schedule(self, tokens_per_step=12):
        """
        Compute schedule: which tokens to generate at each step.
        Returns: list of (step_idx, positions_to_generate)
        """

        schedule = []
        generated = set()
        step = 0

        while len(generated) < self.total_positions:
            # Find best tokens to generate this step
            candidates = set(range(self.total_positions)) - generated

            # Prioritize candidates close to already-generated context
            candidate_scores = {}
            for cand in candidates:
                # Score 1: distance to nearest generated token (closer is better)
                if generated:
                    min_dist = min(
                        spatial_distance(cand, gen, self.grid_size) for gen in generated
                    )
                else:
                    min_dist = 0  # Start anywhere

                # Score 2: distance to nearest other candidate (farther is better for parallelism)
                candidates_list = list(candidates - {cand})
                if candidates_list:
                    avg_dist_to_others = np.mean([
                        spatial_distance(cand, other, self.grid_size) for other in candidates_list
                    ])
                else:
                    avg_dist_to_others = float('inf')

                # Combined score: high proximity to context, high distance from others
                candidate_scores[cand] = (min_dist, avg_dist_to_others)

            # Select tokens for this step using greedy approach
            selected = []
            remaining = set(candidates)

            for _ in range(min(tokens_per_step, len(candidates))):
                if not remaining:
                    break

                # Greedy: pick token closest to context and far from other picks
                best = min(
                    remaining,
                    key=lambda c: (
                        candidate_scores[c][0],  # Maximize context proximity
                        -candidate_scores[c][1]  # Maximize distance from others
                    )
                )

                selected.append(best)
                remaining.remove(best)

                # Remove candidates too close to this selection (redundant)
                too_close = [
                    c for c in remaining
                    if spatial_distance(best, c, self.grid_size) < self.proximity_threshold
                ]
                remaining -= set(too_close)

            schedule.append((step, selected))
            generated.update(selected)
            step += 1

        return schedule

# Test scheduler
scheduler = LocalityAwareScheduler(grid_size=16, proximity_threshold=5.0)
schedule = scheduler.compute_generation_schedule(tokens_per_step=12)

print(f"Generation schedule ({len(schedule)} steps):")
for step, positions in schedule[:5]:
    print(f"  Step {step}: generate {len(positions)} tokens at {positions[:5]}...")

print(f"Total steps: {len(schedule)} (vs 256 for sequential)")

This computes an optimal schedule that maximizes conditioning while minimizing dependencies.

Step 4: Implement complete parallel generation pipeline

This code combines all components for end-to-end parallel generation:

class ParallelAutoregressiveImageGenerator(nn.Module):
    """
    Complete pipeline for parallel autoregressive image generation.
    """

    def __init__(self, grid_size=16, embed_dim=768):
        super().__init__()
        self.grid_size = grid_size
        self.embed_dim = embed_dim

        self.position_queries = PositionQueryTokens(grid_size * grid_size, embed_dim)
        self.attention = FlexibleAutoregressiveAttention(embed_dim, num_heads=8)
        self.decoder = nn.Sequential(
            nn.Linear(embed_dim, embed_dim),
            nn.ReLU(),
            nn.Linear(embed_dim, 256)  # Output logits per token
        )

    def generate(
        self,
        prompt_embedding: torch.Tensor,
        tokens_per_step: int = 12,
        temperature: float = 0.8
    ) -> torch.Tensor:
        """
        Generate image tokens using parallel scheduling.

        prompt_embedding: (B, D) - conditioning information
        """

        B = prompt_embedding.shape[0]
        grid_size = self.grid_size

        # Initialize with prompt
        tokens = torch.zeros(B, grid_size * grid_size, dtype=torch.long)
        generated_mask = torch.zeros(B, grid_size * grid_size, dtype=torch.bool)

        # Get generation schedule
        scheduler = LocalityAwareScheduler(grid_size, proximity_threshold=5.0)
        schedule = scheduler.compute_generation_schedule(tokens_per_step)

        # Generate in parallel batches
        for step, positions_to_gen in schedule:
            # Get context (already-generated tokens)
            context_positions = torch.where(generated_mask[0])[0]
            if len(context_positions) > 0:
                context_tokens = tokens[0, context_positions]
                context_features = self._encode_tokens(context_tokens)  # (num_context, D)
                context_features = context_features.unsqueeze(0).expand(B, -1, -1)
            else:
                context_features = prompt_embedding.unsqueeze(1)

            # Get position queries for tokens to generate
            pos_tensor = torch.tensor([positions_to_gen]).to(prompt_embedding.device)
            position_queries = self.position_queries.get_position_queries(pos_tensor)

            # Apply attention: condition on context
            generated_features = self.attention(
                context_features,
                position_queries,
                context_positions if len(context_positions) > 0 else torch.tensor([]),
                torch.tensor(positions_to_gen),
                grid_size
            )

            # Decode to token logits
            logits = self.decoder(generated_features)  # (B, num_to_gen, 256)

            # Sample tokens
            probs = torch.softmax(logits / temperature, dim=-1)
            sampled = torch.multinomial(probs[0], 1).squeeze(-1)

            # Update tokens and mask
            for i, pos in enumerate(positions_to_gen):
                tokens[0, pos] = sampled[i]
                generated_mask[0, pos] = True

        return tokens

    def _encode_tokens(self, token_indices: torch.Tensor) -> torch.Tensor:
        """Encode discrete tokens to embeddings."""
        # Placeholder; in practice use embedding layer
        return torch.randn(len(token_indices), self.embed_dim)

# Test parallel generation
model = ParallelAutoregressiveImageGenerator(grid_size=16)
prompt = torch.randn(1, 768)

generated_tokens = model.generate(prompt, tokens_per_step=12)
print(f"Generated image tokens shape: {generated_tokens.shape}")
print("Generated 256 tokens in ~20 steps (12× speedup)")

This provides end-to-end parallel image generation.

Step 5: Evaluate speed and quality trade-off

This code measures generation efficiency:

def evaluate_generation_speed(model, grid_size=16, batch_size=8, num_samples=100):
    """
    Measure generation speed: tokens per second, latency per image.
    """

    import time

    model.eval()
    prompts = torch.randn(batch_size, 768)

    # Warmup
    _ = model.generate(prompts[:1], tokens_per_step=12)

    # Measure throughput
    start = time.time()

    for _ in range(num_samples // batch_size):
        _ = model.generate(prompts, tokens_per_step=12)

    elapsed = time.time() - start
    total_tokens = grid_size * grid_size * (num_samples)

    tokens_per_sec = total_tokens / elapsed
    latency_per_image = (grid_size * grid_size) / tokens_per_sec

    print(f"Generation speed: {tokens_per_sec:.0f} tokens/sec")
    print(f"Latency per image: {latency_per_image:.2f} seconds")
    print(f"Sequential would be: {latency_per_image * 256:.2f} seconds")
    print(f"Speedup: {256 / 20:.1f}×")

def evaluate_quality(model, test_loader, metric='inception_score'):
    """
    Measure image quality on parallel vs sequential generation.
    """

    model.eval()
    metrics = {'parallel': [], 'sequential': []}

    for batch in test_loader:
        # Generate with parallel decoding
        generated = model.generate(batch['prompt'], tokens_per_step=12)
        # In practice, convert tokens to images and compute metric
        # Placeholder: assume quality is measured
        quality = compute_image_quality(generated, metric)
        metrics['parallel'].append(quality)

    return metrics

# Evaluate
evaluate_generation_speed(model, grid_size=16, batch_size=8, num_samples=100)

This measures the speed improvements from parallel decoding.

Practical Guidance

When to use Locality-aware Parallel Decoding:

• Interactive image generation requiring low latency (web, mobile)
• 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:

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