Transformer Autoregressive Flow (TARFlow)

Updated 30 June 2025

TARFlow is a model class that merges autoregressive Transformers with normalizing flows to enable tractable likelihood computation and expressive density modeling.
It employs a deep-shallow architecture where a deep Transformer block captures high-level semantics while shallow blocks efficiently refine local details.
STARFlow, a scaled instantiation of TARFlow, demonstrates state-of-the-art performance in high-resolution conditional image synthesis.

Transformer Autoregressive Flow (TARFlow) refers to a class of models that combine the theoretical expressivity of normalizing flows with the structured, scalable modeling capability of autoregressive Transformers. TARFlow and its recently scaled instantiation, STARFlow, integrate autoregressive flow transformations—parameterized by Transformers—within an invertible, end-to-end normalizing flow framework. This design enables tractable likelihood computation, efficient and expressive density modeling, and competitive sample quality, particularly for high-dimensional image and conditional generative modeling tasks.

1. Universality and Theoretical Expressivity

The STARFlow work establishes rigorous universality guarantees for TARFlow. Stacked autoregressive flows with $T \geq 3$ blocks, each with $D$ autoregressive steps (where $D$ is the data dimension) and alternating variable orderings, form a universal approximator for continuous densities $p \in L^1(\mathbb{R}^D)$ .

The proof leverages the property that with sufficient block depth and order reversals, the flow can represent any continuous density using mixtures of Gaussians. For $T=2$ , most conditionals are infinite mixtures of Gaussians (dense in $L^1$ ), but the final coordinate $x_D$ is only modeled as a single Gaussian, limiting expressivity. With $T \geq 3$ , this limitation is lifted, yielding full universality for modeling continuous distributions in high dimensions. This result (see equations (2)-(3) and the "Why TARFlows are Capable Generative Models?" section) is foundational for the application of TARFlow to high-fidelity generative modeling of complex data.

2. Architectural Design: Deep-Shallow Transformer Blocks

STARFlow introduces a "deep-shallow" architecture to maximize modeling capacity and computational efficiency:

Deep Transformer Block: The majority of model parameters are allocated to a single deep Transformer block. This block captures most of the modeling power, learning the high-level semantic structure and effectively serving as a "Gaussian LLM" on the underlying noise or latent representation.
Shallow Transformer Blocks: Two to five post-deep blocks, each much smaller in parameter count and depth, follow the deep block. These shallow blocks refine local or low-level structure (such as texture or subtle details) at minimal computational cost.
Block Allocation: For a total of $L$ layers and $T$ blocks, architectures are instantiated as one large $l$ -layer block and $T-1$ shallow blocks (each of two layers), so $L = l + 2(T-1)$ .
Conditional Guidance: Control signals (e.g., class labels, captions) are injected only into the deep block, which streamlines the conditional generation pipeline and localizes the high-impact guidance to the most influential part of the model.

This architecture reflects the empirical finding that "effective compute concentrates in just the top few AF blocks." Most representation capacity is thus devoted to high-level generative structure, while fine-scale conditioning is performed efficiently.

3. Modeling in Latent Space

STARFlow, like recent advances in diffusion-based generation, operates in the latent space of a pretrained autoencoder rather than directly in pixel space. This approach exhibits several advantages:

Dimensionality Reduction: The encoder compresses images (e.g., from $256 \times 256$ pixels to $32 \times 32$ latent representations), making the learning task tractable for high-resolution images.
Semantic Abstraction: The latent space captures higher-level image semantics, allowing the AF/Transformer to model global coherence over fewer, more informative variables.
Sample Quality and Training Stability: Directly modeling latents is less prone to local pixel artifacts, and training is stabilized through proper noise injection in the latent encoding process.
Objective: The following objective is adopted, supporting exact likelihoods for autoencoding flows:

$\max_{\theta, \phi} \mathbb{E}_{\tilde{z} \sim q(\tilde{z} | x), x \sim p(x)} \big[ \log p(\tilde{z}; \theta) + \log p(x | \tilde{z}; \phi) - \log q(\tilde{z}|x) \big]$

The use of a frozen encoder/decoder further allows STARFlow to focus on learning rich generative priors in the abstract space with the full expressivity and invertibility of normalizing flows.

4. Algorithmic Enhancements: Guidance and Exact Flow

STARFlow advances both guidance and likelihood computation within the TARFlow framework:

Classifier-Free Guidance (CFG) for Flows: Previous AR/flow-based models naively extended CFG by linearly interpolating means and variances between conditional and unconditional predictions, which was unstable at high guidance weights. STARFlow reformulates CFG from a score-function perspective:

$\nabla_x \log \tilde{p}_c(x) = \nabla_x \log p_c(x) + \omega \big( \nabla_x \log p_c(x) - \nabla_x \log p_u(x) \big)$

For Gaussian distributions, this yields

$\tilde{\mu}_c = \mu_c + \frac{\omega s}{1 + \omega - \omega s} (\mu_c - \mu_u), \qquad \tilde{\sigma}_c = \frac{1}{\sqrt{1 + \omega - \omega s}} \sigma_c$

where $s = \sigma_c^2 / \sigma_u^2$ . This method robustly supports aggressive guidance, essential for sharp conditional generation (e.g., in text-to-image settings).

End-to-End Exact Likelihood, No Discretization: STARFlow maintains pure normalizing flow semantics throughout, enabling:
- Invertible Architecture: Every component from encoder (pretrained) through the stacked AF/Transformer flow is invertible.
- Tractable Log-Likelihoods: The data likelihood is computed via the change-of-variables formula:
$\log p(x;\theta) = \log p_0(f_\theta(x)) + \log \left| \det \left( \frac{\partial f_\theta(x)}{\partial x} \right)\right|$

optimizing likelihood directly in continuous space. - No Binning/Quantization: Unlike discrete AR or VAEs, STARFlow never resorts to discretization, preserving information throughput and invertibility.

5. Performance, Scalability, and Sample Quality

Comprehensive evaluation demonstrates that STARFlow closes the gap between normalizing flows and state-of-the-art diffusion or autoregressive models on both class- and text-conditional high-resolution image synthesis benchmarks:

ImageNet 256×256, class-conditional (FID):
- STARFlow: 2.40
- DiT: 2.27
- MaskDiT-G: 2.50
- TARFlow (previous flow): 5.56
ImageNet 512×512, class-conditional (FID):
- STARFlow: 3.00
- DiT-XL/2: 3.04
MS-COCO, text-to-image (FID, 256×256, 3.8B params):
- STARFlow-FullData: 9.1
- Imagen: 7.3
- Parti-20B: 7.2

STARFlow thus matches or matches the best diffusion and discrete AR models of similar (or higher) parameter count, and dramatically outperforms previous flows or AR approaches at comparable size. Inference and training are tractable at scale due to the latent-space modeling and deep-shallow block structure; full AR-style generation is performed in a small number of flow blocks.

6. End-to-End Flow for Continuous-Space Synthesis

A central feature of STARFlow is its status as the first demonstration of end-to-end normalizing flows operating effectively at scale for high-resolution image synthesis. The flow is invertible from data through latent space to noise, supporting:

Exact likelihood computation
Efficient and flexible conditional generation
Bidirectional use cases (inpainting, editing, likelihood evaluation, etc.)
Theoretical performance guarantees via universality

Unlike diffusion, which is fundamentally score- or path-based, and unlike discrete AR models that rely on non-invertible quantization or binning, STARFlow's framework supports both efficient training and sampling, and strong theoretical guarantees.

Summary Table: STARFlow Key Features and Results

Attribute	STARFlow Contribution	Comparison/Impact
Universality	Stacked AFs with $T \geq 3$ are universal	Guarantees expressive fit
Architecture	Deep-shallow AF/Transformer block stack	Parameter efficiency, scalability
Latent-space modeling	Pretrained autoencoder latent domain	High-res, efficient, semantically rich
Guidance	Robust score-informed CFG for AFs	Enables strong conditioning
Likelihood	End-to-end normalizing flow, exact NLL	Outperforms prior AR/flow models
Sample quality (FID)	2.40 (INet256), 3.00 (INet512), 9.1 (COCO)	SOTA or comparable to best diffusion

Transformer Autoregressive Flow, as instantiated in STARFlow, represents a theoretically well-founded, practically scalable, and empirically competitive approach to high-resolution, conditional image generation. By leveraging universal expressivity, architectural concentration of capacity, efficient latent-space modeling, robust algorithmic guidance, and maintaining full normalizing flow invertibility, STARFlow demonstrates the viability of autoregressive flow models for modern generative modeling at the highest scale and fidelity.

PDF Markdown Chat (Pro)