🤖 AI

Beyond the Last Layer: Multi-Layer Representation Fusion for Visual Tokenizatio

This paper introduces DRoRAE, a lightweight representation autoencoder that significantly improves visual tokenization and generation quality by adaptively fusing hierarchical features from all layers of a frozen vision encoder, rather than relying solely on the final layer, thereby uncovering a scalable log-linear relationship between fusion capacity and reconstruction performance.

Original authors: Xuanyu Zhu, Yan Bai, Yang Shi, Yihang Lou, Yuanxing Zhang, Jing Jin, Yuan Zhou

Published 2026-05-12

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Xuanyu Zhu, Yan Bai, Yang Shi, Yihang Lou, Yuanxing Zhang, Jing Jin, Yuan Zhou

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Last Layer" Bottleneck

Imagine you are trying to describe a complex painting to a friend who has never seen it. You have a team of 12 art critics (the layers of a neural network) looking at the painting.

The first few critics are great at noticing tiny details: the texture of the brushstrokes, the exact shade of blue, and the sharp edges of the frame.
The last critic is an expert at the "big picture." They can tell you the painting is a "sunset over a mountain," but they have forgotten the specific texture of the clouds or the exact color of the grass because they were too busy summarizing the whole scene.

Current AI systems (like the ones used to generate images) usually only listen to that last critic. They throw away the notes from the first 11 critics.

The Result: The AI can generate a picture that looks like a sunset, but the details are blurry, the textures are muddy, and the edges are soft. It's like trying to rebuild a house using only the architect's summary, forgetting the blueprint for the bricks and mortar.

The Solution: DRoRAE (The "All-Hands" Meeting)

The authors propose a new system called DRoRAE. Instead of listening to just the last critic, DRoRAE holds a meeting with all 12 critics and combines their notes into one perfect summary.

Here is how they do it, using three clever tricks:

1. The Smart Router (The "Energy-Constrained" Mixer)

You can't just average all the notes together; the "texture" notes from the first critic might clash with the "big picture" notes from the last one.

How it works: DRoRAE uses a Smart Router. Think of this router as a DJ mixing music.
The Trick: Unlike a standard mixer that only turns volume knobs up (positive numbers), this DJ can also turn volume knobs down (negative numbers).
Why it matters: If a specific layer is giving "noisy" or bad information for a specific part of the image, the router can actively suppress it (turn the volume down to negative) rather than just ignoring it. It picks the best details from the shallow layers (textures) and the best concepts from the deep layers (semantics) and blends them perfectly.

2. The "Incremental Correction" (The Safety Net)

If you suddenly change the summary the AI uses, the "decoder" (the part that turns the summary back into a picture) might get confused and fail. It's like changing the language a translator speaks in the middle of a sentence.

The Trick: DRoRAE doesn't replace the old summary entirely. Instead, it treats the new, richer summary as a small correction to the old one.
The Analogy: Imagine you have a perfect map (the old summary). DRoRAE says, "Let's keep that map, but add a few tiny, precise notes about the potholes and street signs." This ensures the decoder doesn't get lost, but it still gets the extra details it needs.

3. The Three-Phase Training (The "Rehearsal" Strategy)

You can't just throw everyone into the final performance at once. The authors use a three-step rehearsal process:

Phase 1: Teach the decoder (the painter) to work perfectly with the old summary (just the last layer).
Phase 2: Freeze the painter. Now, train the Smart Router to create a new summary that the painter can still understand, even though it has more details. The painter acts as a strict teacher, ensuring the new summary doesn't drift too far from what the painter knows.
Phase 3: Unfreeze the painter and let them practice with the new, richer summary. Now the painter learns to use those extra details to paint even better pictures.

The Results: Sharper Pictures and Better Scaling

The paper tested this on ImageNet (a huge collection of photos).

Reconstruction: When trying to rebuild an image from the summary, DRoRAE made the picture much sharper. The "blur" (measured by a score called rFID) dropped significantly. It recovered fine details like hair strands, fabric textures, and thin lines that the old system lost.
Generation: When asking the AI to create new images, the results were also better. The images were more realistic and followed instructions more closely.
Text-to-Image: These improvements also helped when generating images from text descriptions.

The "Scaling Law" Discovery

The authors found something fascinating about how much better the system gets as they make it bigger.

In text AI (like LLMs), we know that if you increase the vocabulary size (the number of words the AI knows), the AI gets smarter in a predictable, straight-line way.
The authors discovered that for image AI, Representation Richness is the same thing.
The Analogy: Think of "Representation Richness" as the size of the AI's toolbox.
- You can make the toolbox bigger by adding more layers (more critics).
- Or you can make the toolbox bigger by making each expert smarter (more capacity per layer).
The Finding: No matter which way you make the toolbox bigger, the quality of the images improves in a predictable, log-linear way. If you double the "richness" of the information, you get a predictable boost in image quality. This means we don't have to guess how to improve these systems; we just need to keep adding more "layers of detail" to the toolbox.

Summary

DRoRAE is a new way to teach AI to see images. Instead of letting the AI forget the fine details by only looking at the "big picture" (the last layer), it forces the AI to listen to every layer of its brain. By using a smart mixer, a safety net, and a careful rehearsal process, it creates images that are sharper, more detailed, and easier to generate, all while following a predictable rule: more detailed information = better images.

Technical Summary: Beyond the Last Layer: Multi-Layer Representation Fusion for Visual Tokenization

1. Problem Statement

Current state-of-the-art visual tokenizers, particularly those leveraging frozen pretrained Vision Foundation Models (VFMs) like DINOv2 as backbones, suffer from a critical single-layer information bottleneck. Existing methods exclusively extract features from the final encoder layer to define the latent space. While these last-layer features are optimized for high-level semantics, they systematically discard low-level visual details (textures, edges, color gradients) through progressive semantic abstraction.

Recent analysis indicates that low-level information survives in the final layer only as attenuated residuals via passive residual connections, becoming increasingly lossy with each successive layer. Consequently, single-layer tokenizers discard the rich hierarchical information distributed across intermediate layers, limiting both reconstruction fidelity and the quality of downstream generation. Furthermore, it remains an open question whether increasing the "representation richness" of visual tokenizers follows a predictable scaling law analogous to vocabulary size scaling in NLP tokenizers.

2. Methodology: DRoRAE

The authors propose DRoRAE (Depth-Routed Representation AutoEncoder), a lightweight fusion module designed to aggregate features from all encoder layers while maintaining compatibility with frozen decoders and downstream generative models.

Core Architecture

DRoRAE inserts a trainable fusion module between the frozen VFM backbone and the decoder. The process involves three key mechanisms:

Layer-wise Experts: Each of the $L$ encoder layers is associated with a dedicated expert network (a two-layer MLP). These experts project heterogeneous layer features onto a common scale, normalized by the backbone's layer normalization to handle variance disparities.
Energy-Constrained Routing: Unlike standard Mixture-of-Experts (MoE) that use softmax normalization, DRoRAE employs an energy-constrained routing mechanism. A router network produces per-token weights across all layers. Crucially, this formulation allows for negative weights, enabling the router to actively suppress detrimental layer contributions (denoising) rather than just selecting winners. The output is normalized by the $\ell_2$ -norm of the weight vector to bound the output energy.
Incremental Correction: Instead of replacing the original last-layer output ( $z_{base}$ ) with the fused representation ( $z_{fuse}$ ), DRoRAE formulates the output as an incremental correction:
$z_{final} = LN_{bb}(z_{base} + \beta \cdot (z_{fuse} - z_{base}))$
Here, $\beta$ controls the fusion strength. This residual formulation ensures the fused latent remains anchored to the original distribution, preventing arbitrary drift that could break downstream generation compatibility.

Three-Phase Decoupled Training Strategy

To address the challenge of shifting the latent distribution away from what a frozen decoder expects, the authors propose a progressive training strategy:

Phase 1 (Decoder Training): Train the decoder on the standard single-layer RAE baseline to establish a strong "decoding capacity" and define the target distribution.
Phase 2 (Fusion Module Training): Freeze the backbone and the Phase 1 decoder. Train only the fusion module. The frozen decoder acts as an implicit distributional constraint, forcing the fusion module to learn complementary information that the decoder can already invert, preventing distribution drift.
Phase 3 (Decoder Fine-tuning): Freeze the optimized fusion module and unfreeze the decoder. Fine-tune the decoder to fully exploit the enriched latent space, maximizing reconstruction quality without harming generation.

3. Key Contributions

Identification of the Bottleneck: The paper identifies that existing representation autoencoders discard hierarchical visual information by relying solely on the last encoder layer, where low-level details are merely attenuated residuals.
DRoRAE Architecture: A lightweight (~29M parameters) depth-routed fusion module that adaptively aggregates all encoder layers via energy-constrained routing (allowing negative weights) and incremental correction.
Training Strategy: A three-phase decoupled training protocol that ensures generation compatibility by learning fusion under the constraint of a frozen decoder before co-adapting the decoder.
Scaling Law Discovery: The paper establishes a log-linear scaling law ( $R^2=0.86$ ) between fusion capacity (defined by expert capacity and number of fused layers) and reconstruction quality. This identifies "representation richness" as a new, predictably scalable dimension for visual tokenizers, analogous to vocabulary size in NLP.

4. Experimental Results

Experiments were conducted on ImageNet-256 (256x256) for reconstruction and class-conditional generation, and on CC12M-LLaVA-Next for text-to-image synthesis.

Reconstruction: DRoRAE significantly improves reconstruction metrics over the RAE baseline. The rFID drops from 0.57 to 0.29, PSNR increases from 18.8 to 24.32 dB, and LPIPS improves from 0.256 to 0.134.
Class-Conditional Generation: Using a DiTDH-XL generator with AutoGuidance, DRoRAE improves the generation FID (gFID) from 1.74 to 1.65 while maintaining or slightly improving Inception Score (IS) and Precision/Recall.
Text-to-Image Synthesis: The gains transfer to text-to-image tasks, with DRoRAE achieving comparable or slightly better GenEval scores (0.60 vs. 0.56 for RAE) across object, counting, and spatial dimensions.
Ablation Studies:
- Energy-Constrained Routing: Outperforms softmax routing by allowing active suppression of noisy layers.
- Incremental Correction: Essential for generation compatibility; without it, reconstruction improves slightly, but downstream diffusion training loss increases significantly, indicating distribution mismatch.
- Training Strategy: Joint training without the Phase 2 constraint fails to achieve the same generation quality, confirming the necessity of the decoupled approach.

5. Significance and Claims

The paper claims that DRoRAE demonstrates that multi-layer fusion is a practical and effective upgrade for representation-based tokenizers, recovering high-frequency details lost in single-layer extraction.

A primary significance of this work is the discovery of a predictable scaling law for visual tokenizers. By showing that increasing the "representation budget" (via more fused layers or larger expert capacities) yields log-linear improvements in reconstruction quality, the authors argue that representation richness is a scalable dimension comparable to vocabulary size in NLP. This suggests that future improvements in visual tokenization can be systematically harvested by scaling fusion capacity, rather than relying solely on architectural changes or larger datasets.

The work concludes that while the current experiments utilize DINOv2-B, the methodology is generalizable to larger encoders and potentially video tokenization, positioning multi-layer fusion as a fundamental direction for enhancing visual generation systems.