How STARFlow2 Unifies Text and Image Generation

Peter Bubenik · Apple ML · · Source
Image for STARFlow2: Bridging Language Models and Normalizing Flows for Unified Multimodal Generation

Step-by-Step Teaching

Step 1: The Problem — Why Unified Multimodal Generation Is Hard

Imagine you want one model that can read text, look at images, reason about them, and generate both text and images — all together, fluidly.

Current approaches fail in one of three ways:

ApproachWhat It DoesThe Problem
Discrete tokenizationConverts images to tokens like wordsLoses visual quality/fidelity
Causal text + diffusionLLM for text, diffusion model for imagesStructurally asymmetric — two different mechanisms
Adapting VLMs for generationFine-tunes understanding models to also generateDamages the pretrained understanding capability

Key Insight: None of these approaches are truly unified — they're always patching two different systems together.


Step 2: Foundation Concept — What Is an Autoregressive Language Model?

Before understanding the solution, you need to understand how LLMs work structurally.

Autoregressive generation means:

  • Generate one token at a time
  • Each new token depends on all previous tokens
  • Direction: strictly left-to-right (causal)

Two key mechanisms:

  • Causal mask: Prevents the model from "seeing the future" — token 5 cannot attend to token 6
  • KV-Cache: Stores computed Key-Value pairs from previous tokens so they don't need to be recomputed — makes generation fast
Token 1 → Token 2 → Token 3 → Token 4 ...
   ↑ cached  ↑ cached  ↑ cached

Step 3: Foundation Concept — What Are Normalizing Flows?

Normalizing flows (NFs) are a class of generative models with specific properties:

Core Idea:

Transform a simple distribution (like Gaussian noise) into a complex distribution (like a realistic image) using a series of invertible transformations.

Simple Noise (z) → [Invertible Transformations] → Complex Data (x)

Key Properties:

  • Likelihood-based: You can compute the exact probability of any data point
  • End-to-end trainable: Single training objective
  • Continuous: Works directly on continuous data (pixel values, latent vectors) — no discretization needed
  • Single-pass: Generate in one forward pass, not iterative like diffusion

Why Does This Matter?

Unlike diffusion models (which require many denoising steps), normalizing flows generate data in one pass — much more efficient.


Step 4: The Core Insight — Flows and LLMs Are Structurally Identical

This is the central intellectual contribution of the paper.

Autoregressive normalizing flows generate data by:

  • Producing outputs left-to-right, one step at a time
  • Using a causal mask (same as LLMs)
  • Leveraging KV-cache (same as LLMs)
LLM Structure:          [Token₁] → [Token₂] → [Token₃] ...
                         causal mask + KV-cache

Autoregressive Flow:    [Patch₁] → [Patch₂] → [Patch₃] ...
                         causal mask + KV-cache

The Realization: These two systems are architecturally the same thing — both are causal Transformers. This means you don't need two separate mechanisms. You can unify them.


Step 5: The STARFlow2 Architecture — The Pretzel Design

Now that we understand the building blocks, let's understand how STARFlow2 is built.

The "Pretzel" Architecture

The name comes from how two streams are interleaved (twisted together like a pretzel).

Two streams running in parallel:

Input Sequence (text + image patches)
         ↓
┌─────────────────────────────────┐
│  Stream 1: Frozen VLM           │  ← Pretrained Vision-Language Model
│  (handles understanding)        │    FROZEN — weights don't change
└────────────┬────────────────────┘
             │ residual skip connections (vertical interleaving)
┌────────────▼────────────────────┐
│  Stream 2: TARFlow              │  ← Transformer Autoregressive Flow
│  (handles image generation)     │    TRAINABLE
└─────────────────────────────────┘
         ↓
    Text Output + Image Output

Key Design Decisions:

1. Frozen VLM Stream

  • The pretrained VLM is kept frozen (weights unchanged)
  • This preserves all the understanding and reasoning capabilities already learned
  • Solves the problem of "degrading pretrained understanding"

2. TARFlow Stream

  • A new normalizing flow model built as a Transformer
  • Learns to generate high-quality continuous images
  • Trained on top of the frozen VLM

3. Residual Skip Connections

  • Information flows vertically between the two streams at each layer
  • The TARFlow stream receives rich semantic information from the VLM stream
  • Like adding shortcuts between floors of a building

4. Same Causal Mask

  • Both streams operate under identical causal masking
  • This is what makes it truly unified — one mechanism governs everything

Step 6: Additional Components — Deep-Shallow Flow and FAE

Deep-Shallow Flow Design

Not all parts of the image need equal computational effort.

Early layers (Deep):   Complex transformations — capture structure
Later layers (Shallow): Simpler transformations — refine details

This makes the model more computationally efficient without sacrificing quality.

Unified FAE Latent Space (Flow Autoencoder)

The problem: Images are high-dimensional (millions of pixels). Working directly on pixels is too expensive.

The solution: FAE compresses images into a latent space — a compact representation.

Why "unified"?

  • Both text tokens and image latents exist in the same space
  • Both can directly enter the KV-cache without re-encoding
  • This enables smooth interleaved generation (text...image...text...image...)
Text token  → KV-cache ✓
Image patch → FAE latent → KV-cache ✓  (no extra re-encoding step needed)

Step 7: What STARFlow2 Achieves — Putting It All Together

Let's verify it solves the original three problems:

Original ProblemSTARFlow2 Solution
Discrete tokenization loses visual qualityUses continuous latent space via FAE — no discretization
Structural asymmetry (LLM + diffusion)Both text and image use same causal Transformer mechanism
Adapting VLMs degrades understandingVLM stream is frozen — understanding is fully preserved

Additional Benefits:

  • Single-pass generation: No iterative denoising like diffusion
  • Cache-friendly: KV-cache works for both modalities
  • Interleaved sequences: Can generate text...image...text naturally in one model

Step 8: Conceptual Summary — The Big Picture

BEFORE STARFlow2:
Text Generation:  [LLM — causal, single-pass]
Image Generation: [Diffusion — iterative, separate mechanism]
                  ↑ Two different systems, awkwardly combined

AFTER STARFlow2:
Text + Image:     [Causal Transformer — single-pass, unified]
                  ↑ One system, one mechanism, one causal mask

The paper's core argument is elegant:

"Since autoregressive normalizing flows are already Transformers with causal masks and KV-caches — the same as LLMs — why not just use one unified system for everything?"

STARFlow2 is the answer to that question.


Quick Knowledge Check

Test your understanding:

  1. Why does discrete tokenization hurt image quality?
  2. What structural property do autoregressive flows share with LLMs?
  3. Why is the VLM stream kept frozen in the Pretzel architecture?
  4. What problem does the FAE latent space solve?
  5. How does STARFlow2 differ from a diffusion-based multimodal model?

(Answers are all contained in the material above — try to answer from memory!)

More to study