Representation Forcing: Eliminating VAE Bottlenecks in Unified Multimodal Models

Paper: 2605.31604 Authors: Yuqing Wang, Zhijie Lin, Ceyuan Yang, Yang Zhao, Fei Xiao, Hao He, Qi Zhao, Zihan Ding, Fuyun Wang, Shuai Wang Categories: cs.CV

The Gap

Unified multimodal models (UMMs) like Chameleon and Transfusion aim to handle both understanding and generation in one architecture. But they cheat: they use a frozen, separately pretrained VAE to compress images into latent codes for generation. This creates a structural bottleneck—the VAE’s latent space becomes a mandatory checkpoint that limits what the unified model can learn. Remove the VAE and generate pixels directly? Quality collapses. The model drowns trying to learn both high-level semantics (what to draw) and low-level details (how to render textures) simultaneously from raw pixels.

Problem: VAE bottleneck in UMMs
    |
    v
Assumption: Direct pixel generation fails because
            model must learn structure + details together
    |
    v
Method: Representation Forcing (RF)
        - Predict visual representations first (structure)
        - Then predict pixels conditioned on those (details)
    |
    v
Evidence: Pixel-space RF matches VAE-based quality
          + improves understanding performance
    |
    v
Conclusion: Can build end-to-end UMMs without VAE

The Increment

One sentence: Before this paper, unified multimodal models needed a frozen VAE as a crutch for image generation; after, they can generate directly from pixels at comparable quality by learning to predict visual representations as an intermediate step.

Core Mechanism

Representation Forcing (RF) works by splitting the generation task into two sequential stages within the same autoregressive backbone. First, the model predicts visual representation tokens—these are discrete codes from a learned codebook that capture high-level visual structure. Second, these representation tokens remain in the context window and guide a diffusion process that generates the actual pixels.

The key insight: by making representation prediction a native capability of the model (not an external VAE), you force the model to learn a structured latent space as part of its autoregressive training. The model learns to compress visual information into representations during understanding tasks, then reuses that same compression skill during generation. The representations act as scaffolding—they provide high-level guidance so the pixel diffusion doesn’t have to figure out both “what” and “how” simultaneously.

Concretely, during training, the model sees images encoded into representation tokens via a learned VQ-VAE. It learns to predict these tokens autoregressively, just like text tokens. During generation, after predicting representation tokens, the model switches to a diffusion mode where it iteratively denoises pixels, conditioned on the representation tokens it just produced. Both stages use the same transformer backbone—no separate networks.

Input text --> Transformer --> Representation tokens (discrete)
                                        |
                                        v
                              [tokens stay in context]
                                        |
                                        v
                              Diffusion decoder (same backbone)
                                        |
                                        v
                                   Output pixels

Think of it like building a house. The old VAE-based approach is like hiring a separate architect (frozen VAE) who hands you blueprints, then you build from those blueprints. You’re stuck with whatever the architect designed—you can’t learn to make better blueprints yourself. Representation Forcing is like learning to be both architect and builder. First, you sketch the floor plan (representation tokens)—where rooms go, rough proportions. Then, with that sketch visible on your desk, you build the actual structure (pixels). The sketch guides your hands, but you drew the sketch yourself, so over time you learn what kinds of sketches lead to good houses. The sketch and the building happen in the same workspace (transformer backbone), using the same tools. You’re not dependent on an external architect anymore.

Key Concepts

• Representation tokens: These are discrete codes from a learned codebook (like VQ-VAE) that compress an image into a sequence of high-level visual features. Imagine taking a photo and describing it with 256 carefully chosen “visual words” from a fixed vocabulary of 8192 words. Each word captures a local pattern—“blue sky texture,” “sharp edge at 45 degrees,” “smooth gradient.” The model learns this vocabulary during training. Unlike VAE latents (continuous vectors), these are discrete tokens, so they fit naturally into the autoregressive next-token prediction framework that language models use. The model predicts them just like it predicts text tokens, making them a native part of the unified architecture.

• Autoregressive representation prediction: The model generates representation tokens one at a time, left to right, each token conditioned on all previous tokens (both text and representation). This is the same mechanism used for text generation. The crucial difference from VAE-based models: the model learns to produce these representations itself, rather than receiving them from a frozen external encoder. This means the model can optimize the representation space end-to-end for both understanding and generation tasks, rather than being stuck with whatever a separately trained VAE learned.

• Contextual diffusion: After predicting representation tokens, those tokens stay in the transformer’s context window. The model then runs a diffusion process to generate pixels, but at each denoising step, it attends to the representation tokens. Think of it as painting with a reference sketch always visible. The diffusion decoder is the same transformer backbone—no separate network. This is different from typical text-to-image diffusion models where text conditioning is injected via cross-attention to a separate encoder. Here, representations and pixels share the same architectural space, allowing tighter integration.

Framework Shift

Before (VAE-based UMMs):              After (Representation Forcing):

Text input                            Text input
    |                                     |
    v                                     v
Transformer (understanding)           Transformer (unified)
    |                                     |
    +---> Text output                     +---> Text output
    |                                     |
    v                                     +---> Representation tokens
[Frozen VAE encoder]                              |
    |                                             v
    v                                    [tokens in context]
Latent codes                                      |
    |                                             v
    v                                    Diffusion (same backbone)
Transformer (generation)                          |
    |                                             v
    v                                        Pixel output
Pixel output

Bottleneck: VAE latent space          No bottleneck: end-to-end learning

From external compression to internal scaffolding, the core shift is making representation learning a native capability of the unified model rather than outsourcing it to a frozen component.

Expert Assessment

Problem choice: Real gap. The VAE bottleneck in UMMs is not manufactured—it’s a genuine architectural compromise. Models like Chameleon achieve impressive results but are fundamentally limited by their frozen VAE’s latent space. The field has been stuck here because naive pixel-space generation fails. This paper attacks the right constraint.

Method maturity: Clever insight with solid execution. The core idea—predict representations before pixels—is simple but non-obvious. It elegantly reuses the autoregressive framework for both stages. However, the method still requires a learned VQ-VAE to define the representation space, so it’s not fully “bottleneck-free” in the sense of requiring zero external components during training. The VQ-VAE is learned jointly, which is better than frozen, but it’s still a separate module. A truly end-to-end approach would learn representations implicitly without discrete tokenization, but that’s likely much harder.

Experimental integrity: Baselines are fair. The paper compares against strong VAE-based UMMs (Chameleon, Transfusion) and shows competitive generation quality (FID, CLIP scores) while improving understanding benchmarks. The ablations are thorough—they test different representation granularities, diffusion steps, and training strategies. One minor concern: the understanding improvements are modest (1-2% on most benchmarks), which could be within noise for some tasks. The generation results are more convincing. No major red flags, but I’d want to see human evaluation on generation quality, not just automated metrics.

Writing quality: The paper is clear and well-structured. The motivation is crisp, the method description is detailed, and the experiments are comprehensive. The weakest section is the related work—it’s a bit perfunctory and doesn’t deeply engage with why prior pixel-space attempts failed. A more thorough analysis of what makes RF succeed where others didn’t would strengthen the contribution. The ablation section is strong but could be condensed; some results feel redundant.

Verdict: weak accept — Solid contribution that removes a real architectural limitation, with convincing experimental results, though the method still relies on learned discrete representations rather than being fully end-to-end.

Takeaways

Staged generation with shared context: When generating complex outputs (images, audio, video), consider splitting the task into high-level structure prediction followed by detail refinement, both using the same model architecture with the structure tokens remaining in context. This pattern could transfer to other domains—e.g., generating code by first predicting an AST sketch, then filling in syntax.

Representation as intermediate supervision: If your model struggles to learn a complex mapping end-to-end, don’t just add more capacity—add intermediate targets that force the model to learn useful internal representations. Make the hidden layer explicit. This is a general principle beyond multimodal models.

Unifying perception and generation: The key to unified models isn’t just training on both tasks—it’s making the representations learned during perception directly useful for generation. If your model has separate pathways for understanding and creation, look for ways to make them share structure, not just parameters.