[CVPR 2026] UNITE: Breaking the Two-Stage Paradigm with Unified Tokenization and Latent Denoising

TL;DR

The dominant recipe for Generative AI has long been "Tokenize first, Diffuse later." UNITE (Unifying Tokenization & Latent Generation) collapses this barrier. By using a single Generative Encoder that shares weights between the tokenizer and the denoiser, the model learns a "common latent language" from scratch. It hits an impressive 1.73 FID on ImageNet without any GAN losses or pretrained DINO features, while proving its versatility on complex molecular datasets.

The "Staging" Problem: Why Separate Training Is Holding Us Back

For years, Latent Diffusion Models (LDMs) have treated the latent space as a "frozen" intermediary. You train a VAE to compress images, lock it, and then ask a Diffusion Transformer (DiT) to learn the distribution of those static points.

The technical cost of this separation is significant:

1. Semantic Mismatch: The tokenizer is optimized for reconstruction, not for being easy to "denoise."
2. Dependency: Modern SOTA often requires "Anchoring" (e.g., REPA) using massive pretrained models like DINOv2, which are computationally expensive to produce.
3. Complexity: Managing multi-stage pipelines increases the risk of error and limits the model's ability to adapt to new domains (like scientific data) where no "DINO" exists.

Methodology: One Encoder to Rule Both Modes

The core innovation is the Generative Encoder (GE). The authors realize that tokenization (Image $o$ Latent) and generation (Noise $o$ Latent) are essentially the same inference problem, just with different amounts of information available ("High Observability" vs. "Low Observability").

The Unified Architecture

UNITE uses a Vision Transformer (ViT) backbone. It operates in two modes within a single training iteration:

• Tokenization Mode: It takes image patches concatenated with $K$ "register tokens." Through self-attention, the registers soak up the image information to become the clean latents ($z_0$).
• Generation Mode: These same registers are initialized with noise ($z_t$). The GE then acts as a denoiser to predict the clean target.

By sharing weights ($heta$) between these passes, the gradients from reconstructing pixels and denoising latents simultaneously "sculpt" the weights. The result? A latent space that is naturally robust to noise because the encoder was forced to make it so.

Results: SOTA Performance Without the "Bloat"

UNITE doesn't just simplify the pipeline; it beats established two-stage models.

Image Generation

Scaling from Base to Large and XL, UNITE demonstrates excellent scaling laws. It surpasses DiT-XL/2 and even many models that use auxiliary supervision.

Moving Beyond Pixels: Molecules and Crystals

Because UNITE learns from scratch, it thrives in domains like QM9 (Molecules) and MP20 (Crystals). Traditional methods struggle here because there are no billion-parameter pretrained "DINO-Molecules" to provide features. UNITE achieves a 99.37% match rate on molecular reconstruction, proving that end-to-end joint training is a universal solution for generative modeling.

Deep Insight: Shared Features vs. Separate Norms

The most fascinating part of the paper is the CKA (Centered Kernel Alignment) analysis. The authors found that even if you don't share weights, the encoder and denoiser naturally evolve toward the same feature representations (high alignment).

However, parameter tying (weight sharing) forces the model to use the same MLPs and Attention layers, whereas the two modes differentiate themselves primarily through Normalization layers. This suggests that the "logic" of understanding an image and the "logic" of denoising a latent are fundamentally the same, differing only in signal scale.

Conclusion & Future Outlook

UNITE marks a shift toward "General Generative Models." By removing the crutches of pretrained teachers and complex multi-stage training, the authors have provided a blueprint for more efficient, flexible, and robust models.

Key Takeaways for Practitioners:

• Single Stage is Enough: You don't need a 2-stage pipeline to get SOTA FID.
• Weight Sharing is a Regularizer: Tying the encoder and denoiser makes the latent space more robust and easier to sample.
• Domain Independence: UNITE is a go-to choice for specialized scientific domains where you lack the luxury of pretrained vision foundations.

Stay tuned as the team explores how this unified logic can be applied to Video and Robotics world models!