PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2605.10780v2cs.CVcs.AI

TL;DR

Existing representation-based visual tokenizers use only final-layer features, leaving hierarchical low-level information underused. DRoRAE adaptively fuses encoder layers with constrained routing and incremental correction, then trains the decoder in phases. On ImageNet-256, it improves reconstruction and generation, and fusion capacity follows a log-linear reconstruction scaling law.

  • Problem

    Existing representation-based tokenizers extract only final-layer features, discarding hierarchical information that shallow layers retain more faithfully.

  • Method

    DRoRAE adaptively fuses all encoder layers using energy-constrained routing and incremental correction, with fusion trained before decoder fine-tuning.

  • Results

    0.57→0.29 rFID and 1.74→1.65 gFID with AutoGuidance on ImageNet-256, with gains transferring to text-to-image synthesis.

  • Takeaways & Limitations

    Representation richness, determined by fused-layer count and per-layer expert capacity, is a predictably scalable dimension for visual tokenizers.

  • Takeaways & Limitations

    Experiments use DINOv2-B with 12 layers, and generation evaluation is pending for the further-unfrozen-backbone configuration.

Abstract

from arXiv · show

Representation autoencoders that reuse frozen pretrained vision encoders as visual tokenizers have achieved strong reconstruction and generation quality. However, existing methods universally extract features from only the last encoder layer, discarding the rich hierarchical information distributed across intermediate layers. We show that low-level visual details survive in the last layer merely as attenuated residuals after multiple layers of semantic abstraction, and that explicitly fusing multi-layer features can substantially recover this lost information. We propose DRoRAE (Depth-Routed Representation AutoEncoder), a lightweight fusion module that adaptively aggregates all encoder layers via energy-constrained routing and incremental correction, producing an enriched latent compatible with a frozen pretrained decoder. A three-phase decoupled training strategy first learns the fusion under the implicit distributional constraint of the frozen decoder, then fine-tunes the decoder to fully exploit the enriched representation. On ImageNet-256, DRoRAE reduces rFID from 0.57 to 0.29 and improves generation FID from 1.74 to 1.65 (with AutoGuidance), with gains also transferring to text-to-image synthesis. Furthermore, we uncover a log-linear scaling law ($R^2{=}0.86$) between fusion capacity and reconstruction quality, identifying \textit{representation richness} as a new, predictably scalable dimension for visual tokenizers analogous to vocabulary size in NLP.

1 Introduction

Existing representation autoencoders rely on last-layer features, discarding hierarchical visual information that multi-layer fusion can recover. DRoRAE addresses adaptive fusion and decoder compatibility, improving reconstruction and generation while revealing scalable representation richness.

  • Motivation: Last-layer tokenizers discard shallow-layer textures, edges, and color gradients because semantic transformations progressively dilute low-level information.Low-level information remains mainly as attenuated residual leakage in the final layer, whereas shallower layers retain it more faithfully.
  • Motivation: Multi-layer fusion introduces two capacity axes—fused-layer count and per-layer expert capacity—that define representation richness and motivate testing a visual-tokenizer scaling law.The paper compares this concept with vocabulary size as a representation-richness dimension in NLP.
  • Method: DRoRAE uses energy-constrained routing, per-layer expert MLPs, active suppression through negative weights, incremental correction, and three-phase decoupled training.Fusion first learns under a frozen decoder’s distributional constraint before decoder fine-tuning exploits the enriched latent.
  • Results: 0.57→0.29 rFID and 1.74→1.65 gFID with AutoGuidance show improved reconstruction and class-conditional generation on ImageNet-256.The gains also transfer to text-to-image synthesis.
  • Results: Both expert capacity and fused-layer count follow the same log-linear scaling law, establishing representation richness as a predictably scalable tokenizer dimension.The scaling experiments treat these two variables as capacity axes for fusion.

2 Related Work

Representation-based tokenizers inherit pretrained visual representations to improve reconstruction and generation, but typically use only the final encoder layer. Related vision models aggregate hierarchical features, leaving multi-layer fusion for generative image tokenization largely unexplored.

  • Image Tokenizers for Latent Generation: Learned tokenizers use discrete or continuous latent spaces, while representation-based methods align or directly reuse pretrained visual features.REPA, VA-VAE, RAE, and RPiAE progressively incorporate pretrained representations, distillation, freezing, or channel expansion.
  • Image Tokenizers for Latent Generation: Representation-based tokenizers achieve strong reconstruction and downstream generation quality by inheriting structure from pretrained visual models.RAE directly freezes DINOv2 as the tokenizer and trains a decoder, while RPiAE expands channels with principal components.
  • Multi-Layer Feature Utilization in Vision Models: Encoder depth separates fine-grained textures and edges in shallow layers from high-level semantics in deep layers, but existing tokenizers extract only the final layer.This single-layer bottleneck discards hierarchical information potentially useful for reconstruction and generation.
  • Multi-Layer Feature Utilization in Vision Models: Feature Pyramid Networks, Dense Prediction Transformers, and hypercolumns demonstrate multi-layer aggregation for dense prediction, whereas generative image tokenization has scarcely explored it.The open questions are whether explicit fusion improves reconstruction beyond residual leakage and benefits generation.
  • Multi-Layer Feature Utilization in Vision Models: DRoRAE processes multi-layer frozen-backbone features through depth routing, energy-constrained aggregation, incremental updating, and decoding for reconstruction.Its design connects hierarchical feature utilization with representation-autoencoder tokenization.

3 Method

DRoRAE extends representation autoencoders by fusing hidden states from every encoder layer into an enriched latent, while preserving compatibility with the pretrained decoder. Its routed fusion and three-phase training recover complementary visual information and improve reconstruction quality.

  • 3.1 Preliminaries: Standard RAE uses only the final-layer latent, discarding intermediate hidden states that retain fine-grained visual information.The final representation is semantically rich, but shallower layers preserve details such as textures, edges, and color gradients.
  • 3.2 Depth-Routed Fusion Module: DRoRAE inserts a lightweight fusion module between the frozen backbone and latent space, producing a drop-in enriched representation from all encoder layers.The module adds approximately 29M parameters and replaces the original latent for decoding.
  • 3.2 Depth-Routed Fusion Module: Energy-constrained routing assigns per-token weights across layer-specific experts, including negative weights that suppress detrimental contributions while bounding output energy.Each expert is a two-layer MLP, and routing weights are normalized by their ℓ2 norm rather than softmax.
  • 3.2 Depth-Routed Fusion Module: Incremental correction adds the fused information to the baseline latent instead of replacing it, with β controlling fusion strength.When β = 0, the module reduces to the original single-layer RAE and focuses learning on complementary information.
  • 3.3 Training Strategy: Three-phase training first establishes a decoder, then trains fusion under the frozen decoder’s distributional constraint, and finally fine-tunes the decoder with the enriched latent.The staged procedure progressively introduces fusion complexity and avoids arbitrary latent-distribution drift during module training.
  • 3.4 Reconstruction Results: DRoRAE improves qualitative reconstruction of textures, thin structures, repetitive patterns, and color fidelity over the RAE baseline.These gains target high-frequency content that the last-layer representation tends to lose.

4 Experiments

Experiments evaluate DRoRAE across reconstruction, class-conditional generation, text-to-image synthesis, ablations, and scaling behavior. The results show improved reconstruction, preserved or improved generation quality, and predictable gains with fusion capacity and depth.

  • Experimental Setup: Experiments use ImageNet-1K at 256 × 256 for reconstruction and class-conditional generation, plus CC12M-LLaVA-Next for text-to-image generation.Evaluation includes 50K ImageNet validation images and 50K generated images for FID computation.
  • Reconstruction: 0.29 rFID, 24.32 dB PSNR, 0.134 LPIPS, and 0.701 SSIM improve over the RAE baseline with the same DINOv2-B encoder and ViT-XL decoder.The full three-phase model adds approximately 29M fusion parameters; Phase 2 alone reaches rFID 0.47.
  • Class-Conditional Generation: 1.65 gFID with AutoGuidance improves over RAE-B’s 1.74 while maintaining Precision 0.81 and increasing Recall from 0.60 to 0.61.Phase 2 already reaches gFID 1.70, and AutoGuidance recovers the mild distribution shift observed without guidance.
  • Text-to-Image Generation: 0.59 versus 0.56 overall GenEval shows comparable text-to-image quality between DRoRAE and the RAE baseline despite the reconstruction improvement.Both models use the same Bagel-MoT framework and CC12M-LLaVA-Next training data.
  • Ablation Studies: Energy-constrained routing improves rFID over softmax routing, while incremental correction lowers DiT loss from approximately 0.8 to 0.47 at epoch 12.Negative routing weights suppress detrimental layer contributions, and the β = 0.2 update anchors outputs near the last-layer distribution.
  • Scaling Behavior: 0.86 R^2 characterizes the log-linear relationship between expert capacity and reconstruction quality, while adding layers also yields consistent improvement without saturation at 12 layers.Across both axes, increasing expert capacity or the number of fused layers provides predictable gains; the unified parameter trend has R^2 = 0.59.

5 Conclusion

DRoRAE enriches representation-autoencoder latents by routing features from multiple depths of a frozen pretrained encoder. It improves reconstruction and generation on ImageNet-256, with representation richness emerging as a scalable tokenizer dimension.

  • Conclusion: DRoRAE aggregates multi-layer features from a frozen pretrained encoder using energy-constrained routing, incremental correction, and three-phase decoupled training.The router produces a structurally distinct fused representation rather than simply reproducing the last-layer output.
  • Conclusion: 0.29 rFID and 1.65 gFID with AutoGuidance improve over 0.57 and 1.74 respectively on ImageNet-256, with gains transferring to text-to-image synthesis.The conclusion reports both reconstruction and class-conditional generation improvements.

A Training Details

Training decouples decoder learning, fusion learning, and decoder adaptation while keeping the DINOv2-B encoder frozen. The implementation specifies the model components, phase-wise optimization, and diffusion-training setup.

  • Architecture: The tokenizer uses a DINOv2-B-reg encoder, ViT-XL decoder, and approximately 29M-parameter fusion module; DiTDH-XL is the class-conditional generator.Table 5 summarizes architecture and optimization configurations.
  • Phase 1: Decoder Training: Phase 1 trains the ViT-XL decoder from the frozen DINOv2-B last-layer representation at 256 × 256 resolution.Training runs for 100 epochs with cosine learning-rate decay and a DINO-based patch discriminator after a 30k-step warmup.
  • Phase 2: Fusion Training: Phase 2 freezes the encoder and decoder while training the approximately 29M-parameter fusion module with 12 layer-specific experts and β = 0.2 incremental correction.The energy-constrained router produces layer weights without softmax normalization, and GAN warmup is reduced to 10k steps.
  • Phase 3: Decoder Fine-Tuning: Phase 3 unfreezes the decoder for 20 additional epochs while keeping the fusion module frozen, allowing decoder co-adaptation to the enriched latent.Other hyperparameters remain identical to Phase 1.
  • Stage 2: Diffusion Training: DiTDH-XL is trained for 80 epochs on the DRoRAE latent with v-prediction, then sampled for 250 DDPM steps with AutoGuidance at scale 1.5.The generator has 839M parameters and uses DiTDH-S as its guidance model.

B Text-to-Image Training Details

The text-to-image evaluation uses a unified autoregressive architecture with decoupled text and vision processing. Tokenizers are compared under matched training settings, with only latent shape and vision-expert channel dimensions adapted.

  • Text-to-Image Evaluation: Bagel’s Mixture-of-Transformers architecture decouples text and vision processing within a unified autoregressive framework.The text-to-image evaluation uses this architecture.
  • Controlled Comparison: All tokenizers share the same training configuration for fair comparison.
  • Controlled Comparison: Tokenizer-specific adaptations are limited to latent shape and the corresponding vision-expert channel dimension.All other hyperparameters remain identical across runs.

C Scaling Experiment Details

The scaling experiments vary expert capacity and the number of fused layers under controlled Phase 2 training. Expert capacity shows a stronger log-linear relationship with reconstruction quality than layer count, whose scaling also changes the router’s available space.

  • Experimental Setup: The experiments use Phase 2 training, updating only the fusion module while freezing the backbone and decoder.Training configurations are otherwise identical except for fusion-module architecture.
  • Expert Capacity Scaling: Expert capacity scaling fixes 12 fused layers and varies the expert hidden dimension.The parameter count depends on backbone dimension C, expert dimension h, and layer count L.
  • Layer Count Scaling: Layer count scaling fixes expert hidden dimension at 3072 and selects progressively more of the deepest backbone layers.For DINOv2-B, the selected layers range from layers 13−N through 12.
  • Scaling Results: rFID = −0.058 · log10(params) + 0.97 achieves R2 = 0.86 for expert capacity scaling, while layer count has slope −6.2 × 10−3 per layer with R2 = 0.49.Layer scaling has higher variance because it changes both available information and the router’s routing space.

D Class-Conditional Generation Samples

Selected ImageNet-256 class-conditional samples generated with DRoRAE and AutoGuidance exhibit high visual fidelity, coherent global structure, and fine-grained local detail across diverse categories.

  • Sample Quality: The samples show coherent global structure and fine-grained local details, including sharp textures in animal fur, feathers, and food surfaces.The latent space also remains well-structured for generative modeling across diverse samples.

E Full Router Weight Visualization

The full router visualization shows stage-wise routing changes across all 12 layers, including suppression and later antagonistic activation. The fused representation is complementary to the base latent while incremental correction preserves decoder compatibility.

  • Routing Evolution: Routing weights evolve from mild shallow adoption through localized and strong suppression, then return to positive activation in later layers.L1–L3 are broadly positive, L4–L7 increasingly suppressive, and L10–L12 are uniformly positive with reduced selectivity.
  • Routing Evolution: L8–L9 activate positively where L6–L7 suppressed, forming antagonistic routing pairs.This stage-wise reversal produces a complementary fused representation.
  • Latent Comparison: cos(zfuse, zbase) is approximately −0.22, indicating nearly orthogonal latent directions, while ∥zfuse−zbase∥ remains uniformly high.These comparisons support that fusion adds complementary information rather than duplicating the base representation.
  • Latent Comparison: With β = 0.2, the final output remains 80% zbase and 20% zfuse, preserving decoder compatibility while injecting complementary high-frequency detail.
  • Generated Samples: Generated samples exhibit sharp textures and coherent structures across diverse ImageNet-256 categories.
Loading 2605.10780v2…