Read, Allocate, Route: Role-Coupled Tokenization for Pixel-Space Generative Transformers

Core Insight§

Pixel-space generative Transformers typically use one token per image patch, so finer detail requires more tokens, making global attention expensive. The authors observe that high-frequency detail can be modeled locally, while low-frequency structure benefits from global exchange. This motivates SpecSlot, a pixel-space rectified-flow Transformer that decouples output detail from the length of the global attention sequence.

Methodology§

SpecSlot expands each spatial anchor into multiple slots, each assigned a disjoint band of DCT coefficients. Formally, for an image patch $P$, we compute its 2D DCT coefficients $C = \text{DCT}(P)$. The coefficients are partitioned into $K$ frequency bands $\{B_1, \dots, B_K\}$ (e.g., by zigzag order). Each band $B_k$ is embedded into a slot token $s_k = \text{Embed}(C_{B_k})$. Thus, an anchor produces $K$ slot tokens, with the lowest-frequency slot $s_1$ capturing global structure and higher-frequency slots $s_2,\dots,s_K$ capturing local detail.

Attention Scheme§

• Global attention: Only the lowest-frequency slot $s_1$ from each anchor participates in global self-attention, keeping the global sequence length equal to the number of anchors.
• Local attention: All slots interact within a local neighborhood (e.g., within the same anchor and adjacent anchors) via local attention or convolution.

This design allows increasing spectral detail (by adding more slots) without enlarging the global sequence.

Spatial Head§

A spatial head adds a local correction to the spectral prediction. The generative state, target, and loss remain in pixel space. The model predicts the DCT coefficients for each slot, which are inverse-transformed to pixels, and the spatial head refines the result with a local convolutional correction.

Training§

The model is trained with a rectified flow objective in pixel space. Given a target image $x$, the loss is:

$$\mathcal{L} = \mathbb{E}_{t, \epsilon} \left[ \| (x - \epsilon) - v_\theta(x_t, t) \|^2 \right]$$

where $x_t = (1-t)x + t\epsilon$ is the interpolated state, and $v_\theta$ is the predicted velocity. The velocity is computed from the slot representations via inverse DCT and the spatial head.

Implementation Details§

• Slot expansion: Each anchor is expanded into $K$ slots (e.g., $K=4$), each assigned a DCT band of size $d$.
• Global attention: Only the first slot per anchor is used, so the global sequence length equals the number of anchors.
• Local attention: All slots attend to each other within a local window (e.g., 3x3 anchors).
• Spatial head: A lightweight convolutional network (e.g., 2-3 layers) predicts a residual correction to the spectral prediction.
• Rectified flow: The model is trained with the standard flow matching objective.

PyTorch Code Snippet§

import torch
import torch.nn as nn
import torch.nn.functional as F

class SpecSlotBlock(nn.Module):
    def __init__(self, dim, num_slots=4, window_size=3):
        super().__init__()
        self.num_slots = num_slots
        self.window_size = window_size
        self.global_attn = nn.MultiheadAttention(dim, num_heads=8, batch_first=True)
        self.local_attn = nn.MultiheadAttention(dim, num_heads=8, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(dim, dim*4), nn.GELU(), nn.Linear(dim*4, dim))
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.norm3 = nn.LayerNorm(dim)

    def forward(self, slots):
        # slots: (B, num_anchors, num_slots, dim)
        B, N, K, D = slots.shape
        # Global attention on lowest-frequency slot
        global_tokens = slots[:, :, 0, :]  # (B, N, D)
        global_out, _ = self.global_attn(global_tokens, global_tokens, global_tokens)
        global_out = self.norm1(global_tokens + global_out)
        # Local attention on all slots within windows
        # Reshape to windows: (B, num_windows, window_size*window_size*K, D)
        # ... (implementation of window partition)
        local_out = ...  # local attention on all slots within windows
        local_out = self.norm2(slots + local_out)
        # MLP
        out = self.norm3(local_out + self.mlp(local_out))
        return out

class SpatialHead(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, in_channels, 3, padding=1)
        self.conv2 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
    def forward(self, x):
        return self.conv2(F.gelu(self.conv1(x)))

Results§

• ImageNet-256: Comparable final FID to a FLOPs-matched JiT control while converging faster.
• ImageNet-512: Improves FID by 2.84 over the control with 3.9% fewer FLOPs, and comes within 0.69 FID of JiT-B/16 at 38.2% lower cost.
• Text-to-image: Surpasses JiT-B/8 on DPG-Bench by 1.06 at 40.6% lower forward cost.

SpecSlot demonstrates that decoupling output detail from global attention length is an effective strategy for efficient high-resolution image generation.