UNITE: Single-Stage Joint Training of Image Tokenization and Diffusion Generation

Paper: 2603.22283 Authors: Shivam Duggal, Xingjian Bai, Zongze Wu, Richard Zhang, Eli Shechtman, Antonio Torralba, Phillip Isola, William T. Freeman Categories: cs.CV, cs.AI, cs.GR, cs.LG

The Gap

Latent diffusion models (LDMs) like Stable Diffusion achieve impressive image synthesis, but they’re stuck in a two-stage training trap. First, you train an autoencoder to compress images into latent codes. Then you freeze it and train a separate diffusion model in that latent space. This staging creates a chicken-and-egg problem: the tokenizer doesn’t know what latent structure would help generation, and the generator can’t influence the latent space it’s forced to work in.

Prior work (LDM, DALL-E 2, Imagen) accepts this split as necessary. The assumption: tokenization (encoding real images) and generation (sampling from noise) are fundamentally different tasks requiring separate training. But what if they’re actually two sides of the same coin?

Problem: Two-stage training creates misaligned latent spaces
   |
   v
Assumption: Tokenization and generation are the same inference problem
            under different conditioning (full image vs noise+text)
   |
   v
Method: Share encoder weights between both tasks, train jointly
   |
   v
Evidence: FID 2.12 on ImageNet 256x256 without adversarial losses
   |
   v
Conclusion: Single-stage joint training is feasible and competitive

The Increment

One sentence: Before UNITE, you trained tokenizers and diffusion models separately in sequence; after UNITE, you train them simultaneously through a shared encoder that learns a “common latent language.”

Core Mechanism

UNITE introduces a Generative Encoder that plays two roles through weight sharing. In the tokenization pass, it takes a real image and encodes it to latents (like a standard VAE encoder). In the generation pass, it takes noisy latents plus text conditioning and predicts the denoised latents (like a diffusion model). The decoder is always the same—it reconstructs images from latents.

The training alternates between two forward passes per batch. First pass: encode real images to latents, decode them back, compute reconstruction loss. Second pass: add noise to those latents, feed noisy latents + text to the same encoder (now acting as denoiser), predict clean latents, compute diffusion loss. Gradients from both tasks flow through the shared encoder weights.

Training Loop (per batch):

Pass 1 (Tokenization):
  Image --> [Generative Encoder] --> Latent --> [Decoder] --> Reconstructed Image
              (shared weights)                                        |
                                                                      v
                                                              Reconstruction Loss

Pass 2 (Generation):
  Noisy Latent + Text --> [Generative Encoder] --> Predicted Clean Latent
                           (same shared weights)              |
                                                              v
                                                        Diffusion Loss

Both losses backprop through shared encoder --> joint latent space optimization

Think of it like training a bilingual translator. Most systems train two separate translators: English→French, then French→English. UNITE trains one person who learns both directions simultaneously. When translating English→French, they learn what French structure makes sense. When going French→English, they learn what English structure is natural. The shared brain (encoder weights) develops an internal “meaning space” that works for both directions. The tokenization task teaches the encoder what latent codes faithfully represent images. The generation task teaches it what latent codes are easy to denoise from noise. The shared weights force a compromise—a latent space that’s both reconstructable and generatable.

Key Concepts

• Latent Space as Common Language: In traditional LDMs, the latent space is designed only for reconstruction—compress the image, decompress it, minimize pixel error. The diffusion model is then forced to generate in this space that wasn’t designed for generation. UNITE’s insight: if both tasks share weights, the latent space must satisfy both reconstruction fidelity and generation ease. It’s like designing a file format that’s both human-readable and machine-parseable—the constraints from both uses shape the format. Here, gradients from reconstruction push for information preservation, while gradients from diffusion push for smooth, denoising-friendly structure.

• Conditioning as Task Selector: The Generative Encoder doesn’t have separate code paths for tokenization vs generation. Instead, conditioning signals tell it which task to perform. For tokenization: condition on the full image (via cross-attention or concatenation). For generation: condition on text and noise level. Same weights, different inputs. This is like a Swiss Army knife—one tool, multiple functions selected by how you hold it. The network learns to route information differently based on what’s being conditioned on.

• Weight Sharing vs Parameter Efficiency: This isn’t about saving memory (though it does). It’s about forcing the encoder to learn representations that serve both tasks. If tokenization and generation had separate encoders, tokenization could learn a latent space that’s great for reconstruction but terrible for denoising. Weight sharing creates pressure: the latent space must be a compromise that works for both. It’s architectural inductive bias—the structure of the model enforces a constraint that the latent space should be “generation-aware.”

Framework Shift

Before (LDM two-stage):                After (UNITE single-stage):

Stage 1: Train Tokenizer              Joint Training:
  Image --> [Encoder] --> Latent         Image -----> [Generative Encoder] --> Latent
              |                                         (shared weights)         |
              v                                              ^                   v
          [Decoder] --> Reconstructed                       |               [Decoder]
              |                                    Noisy Latent + Text          |
              v                                              |                   v
      Reconstruction Loss                                    |           Reconstruction
         (freeze encoder)                                    |               Loss
                                                             |                   +
Stage 2: Train Diffusion                                     |            Diffusion
  Noisy Latent + Text --> [Diffusion Model]                 |               Loss
                               |                             |                   |
                               v                             +-------------------+
                       Predicted Clean Latent                  (both backprop)
                               |
                               v
                        Diffusion Loss

Latent space shaped by:                Latent space shaped by:
- Reconstruction only                  - Reconstruction + Generation jointly

One sentence: From sequential specialization (train tokenizer, then generator) to simultaneous co-optimization (train both through shared weights), the core shift is treating tokenization and generation as dual views of the same latent inference problem.

Expert Assessment

Problem choice: This is a real gap, not manufactured. The two-stage training of LDMs is an acknowledged pain point—you can’t fine-tune the latent space after the diffusion model reveals what structure would help. The problem sits at a natural inflection point: LDMs are mature enough that people are questioning their training inefficiencies, but not so entrenched that alternatives seem impossible.

Method maturity: The core insight (tokenization and generation are both latent inference under different conditioning) is elegant, but the execution is straightforward weight sharing plus alternating training. There’s no deep algorithmic innovation—it’s more of an architectural rearrangement. That said, sometimes the simple idea is the right idea. The paper doesn’t explore why this works theoretically (e.g., what properties of the loss landscape enable joint convergence), which feels like a missed opportunity.

Experimental integrity: The baselines are fair—they compare against LDM and DiT trained with the same compute budget. FID 2.12 on ImageNet 256x256 is competitive but not state-of-the-art (DiT-XL gets 2.27, but uses more parameters). The claim “near state of the art” is defensible but slightly generous. One red flag: they don’t show what happens if you train the two-stage baseline longer to match UNITE’s total training time. Is UNITE better, or just more compute-efficient? The ablations are solid—they show that weight sharing matters more than just joint training with separate encoders.

Writing quality: The paper front-loads motivation well, but the method section is dense. The “Generative Encoder” name is confusing—it’s not generating images, it’s encoding to latents (sometimes from images, sometimes from noise). A clearer name would help. The analysis sections (representation alignment, compression) feel tacked on—they don’t build a coherent story about why joint training works. Rewriting Section 4 to unify these analyses under a single narrative would elevate the paper.

Verdict: Weak accept — solid empirical result with a clean idea, but lacks theoretical depth and the gains are incremental rather than transformative.

Takeaways

Steal the framing: When you have two tasks that operate on the same representation (here, latents), ask if they’re actually the same task under different conditioning. This reframing can collapse multi-stage pipelines into single-stage training. Applies beyond vision—think NLP (encoding and decoding), RL (policy and value), or audio (analysis and synthesis).

Weight sharing as regularization: Forcing two tasks to share parameters isn’t just about efficiency—it’s a way to regularize the learned representation. If your representation works for both tasks, it’s probably capturing something fundamental rather than task-specific quirks. Use this when you suspect your representation is overfitting to one task’s idiosyncrasies.

Alternating training for multi-objective optimization: The two-pass-per-batch training is a simple way to handle multiple loss functions that operate on the same parameters. Instead of weighting losses (which requires tuning), alternate between them. Each pass gets clean gradients from its own objective. This is cleaner than multi-task learning with weighted sums when the tasks have different data requirements.