PaperlayerView exact revision on arXiv

Source-linked AI summary

dRAE: Representation Autoencoder with Hyper-Spherical Codes

Tianren Ma, Lin Long, Chuyan Chen, Mu Zhang, Junbo Zhao, Tong Zhang, Qixiang Ye

arXiv:2607.22148v1cs.CVcs.AI

TL;DR

High-dimensional visual representations are difficult to discretize because Euclidean quantization mismatches their anisotropic, orientation-driven geometry and can cause codebook collapse. dRAE introduces Hyper-Spherical Quantization, which routes codes by angular similarity while preserving magnitude, achieving scalable reconstruction and strong performance across visual understanding and generation tasks, including 0.42 rFID with a 131, 072-codebook.

  • Problem

    Euclidean codebook assignment mismatches the orientation-driven geometry of high-dimensional visual representations, contributing to magnitude bias, uneven angular distributions, and codebook collapse.

  • Method

    dRAE uses Hyper-Spherical Quantization to route codes by angular similarity while preserving magnitude information through a Euclidean commitment objective.

  • Results

    0.42 rFID with a 131, 072-codebook, alongside consistent vocabulary scaling, high code utilization, and strong performance across reconstruction, understanding, and generation tasks.

  • Takeaways & Limitations

    dRAE offers a scalable visual tokenizer with high codebook utilization, fast convergence, robust reconstruction quality, and a foundation for next-generation multimodal language models.

Abstract

from arXiv · show

In this work, we aim to discretize the high-dimensional visual representations to bridge the gap with language models - a non-trivial challenge, as existing quantization methods suffer from codebook collapse, failing to scale while preserving semantic coherence. We identify the root cause as metric mismatch: standard Euclidean codebook objectives are fundamentally misaligned with the anisotropic geometry of representation space, leading to codebook embeddings with high-variance magnitude scales and uneven angular distributions that hinder scalability. To address this, we propose Hyper-Spherical Quantization (HSQ), which decouples semantic content from feature magnitude via angular routing, preventing code assignment from being dominated by scale rather than meaning. The resulting discrete Representation Autoencoder (dRAE) achieves high-fidelity reconstruction while preserving semantic integrity and supporting scalable codebook budget. Extensive experiments demonstrate consistent performance gains as the vocabulary size scales to 131{,}072, along with 100\% codebook utilization, simplified training pipeline, and strong performance across understanding and generation tasks.

1 Introduction

Representation autoencoders bridge semantically rich visual representations and image reconstruction, but discrete tokenization of high-dimensional features suffers from codebook collapse due to metric mismatch. dRAE addresses this with Hyper-Spherical Quantization, routing codes by angular similarity while preserving magnitude, and scales to 131,072 codes across reconstruction, understanding, and generation tasks.

  • Motivation: Representation autoencoders use high-dimensional pretrained features to reconstruct images, producing visual tokens with more generalizable information than reconstruction VAEs.These representations inherit global semantic alignment from contrastive or self-supervised models.
  • Problem: Existing vector quantization methods suffer severe codebook collapse on high-dimensional features and plateau beyond 16K codes.The limited vocabulary cannot capture the latent space’s diverse semantic capacity.
  • Problem: Metric mismatch causes collapse because semantics are encoded mainly in vector orientation, whereas Euclidean assignment is sensitive to anisotropic scale.This produces codebook embeddings with high-variance magnitudes and clustered angular distributions.
  • Method: HSQ routes codes by angular similarity while preserving essential magnitude information for the decoder, decoupling routing from the quantization objective.This targets semantic orientation rather than letting feature scale dominate code assignment.
  • Results: 131,072 vocabulary codes yield continuous reconstruction gains while retaining high semantic fidelity for understanding and supporting text-to-image generation.The experiments also establish dRAE as a viable high-throughput tokenizer for visual understanding and generation.

2 Related Work

Related work contrasts semantically rich, high-dimensional visual representations with compact generative latents, motivating efforts to unify comprehension and generation. Existing approaches face substantial training costs or quantization limitations, including codebook collapse and poor scaling potential.

  • Visual Semantic and Generative Latent: Self-supervised and contrastive objectives produce semantically structured visual latents that generalize across visual understanding tasks.These methods include masked modeling, self-distillation, and image-text alignment.
  • Visual Semantic and Generative Latent: Generative models typically use heavily compressed reconstruction-trained VAE manifolds, creating dimensional incompatibility with semantic-rich embeddings.Visual representations commonly have d ∈[768, 1536], whereas generative latents commonly have d ∈[8, 64].
  • Unifying Comprehension and Generation: Unifying comprehension and generation has been pursued through joint semantic and reconstruction objectives or pretrained encoders, but end-to-end approaches demand extremely high data and training budgets.VTP explores end-to-end pretraining, while VILA-U uses residual quantization with a pretrained encoder.
  • Unifying Comprehension and Generation: RAE directly decodes high-dimensional semantic features for reconstruction, while VQRAE applies SimVQ quantization but remains prone to codebook collapse and limited scaling potential.RAE also identifies potential for visual world modeling through vectors capturing semantic density and fine-grained structure.

3 Preliminary

This section defines image reconstruction and introduces Representation Autoencoders, which use high-dimensional pretrained vision features instead of a low-dimensional bottleneck. It then formulates vector quantization and outlines self-distillation and projection mechanisms for preserving semantics and improving codebook utilization.

  • Image Reconstruction: Image reconstruction maps an input image to latent visual tokens with an encoder and reconstructs pixels through a decoder.The encoder produces N = hw/p^2 tokens, while the decoder reconstructs the image from the latent set.
  • Image Reconstruction: The reconstruction objective combines pixel-level fidelity through L1 distance with perceptual quality and discriminator losses.The losses are denoted L1, Lperc, and Ldisc.
  • Representation Autoencoder: RAE removes the VAE bottleneck by using high-dimensional, semantically structured embeddings from pretrained Vision Foundation Models as latents.RAE uses a frozen VFM encoder and a trainable ViT decoder.
  • Vector-Quantizing Semantics: Vector quantization maps each image feature to its nearest learnable codebook entry, producing quantized features for discrete modeling.VQRAE first projects image features to 1536 dimensions, and codebook capacity is denoted K.
  • Vector-Quantizing Semantics: VQRAE combines codebook and commitment losses, then uses teacher supervision and a learnable projection to preserve semantics and improve codebook utilization.The teacher is initialized from the original VFM, while SimVQ transforms codebook entries with Wci.

4 Methodology

The methodology identifies a mismatch between hyper-spherical visual representations and Euclidean quantization, then introduces HSQ to separate angular semantic routing from magnitude-aware reconstruction. This design preserves both directional and magnitude information while stabilizing scalable codebook learning.

  • Representation Geometry: High-dimensional visual features concentrate near a hyper-spherical shell, while contrastive and self-supervised encoders encourage regular, directionally distributed representations.The motivation draws on thin-shell concentration in Gaussian spaces and observed regularity in CLIP, DINO, and JEPA representations.
  • Diagnostic Hypotheses: Normalized features preserve semantic understanding but impair image reconstruction, indicating that semantics and structural fidelity depend on different feature properties.The diagnostic compares raw and ℓ2-normalized features in MLLM understanding and ViT pixel reconstruction settings.
  • The Unreliable Euclidean: Euclidean VQ assignment is sensitive to feature norms, biasing tessellation toward radial differences rather than angular structure and potentially causing codebook collapse.Samples with larger or varying norms can disproportionately influence the partitioning under the standard ℓ2 objective.
  • Objective Decoupling: HSQ replaces Euclidean assignment with angular routing and uses a spherical codebook loss that constrains codebook updates to the tangent space.A learnable linear projection is applied to raw image features before angular assignment.
  • Objective Decoupling: HSQ retains Euclidean commitment loss to provide magnitude-aware guidance, because fully angular objectives leave magnitude unsupervised and destabilize convergence.Unlike SVQ, HSQ uses spherical geometry only for nearest-neighbor assignment and does not impose a unit-sphere prior on decoder inputs.

5 Experiment

Experiments show that dRAE simplifies training through one-stage end-to-end distillation while achieving scalable code utilization, reconstruction improvements, competitive multimodal understanding, and efficient visual generation.

  • Implementation: dRAE trains end-to-end with distillation loss in one stage, without embedding initialization, anti-collapse tricks, or curriculum learning.This simplifies the training pipeline relative to prior methods using two-stage training.
  • Scaling Behavior: Increasing vocabulary size consistently raises code utilization and reconstruction fidelity, unlike the VQRAE baseline’s marginal gains in active codes and PSNR.The reported scaling behavior contrasts with prior observations of limited benefits from larger vocabularies.
  • Codebook Collapse: At K = 65,536, VQRAE activates ∼6K codes per forward pass, whereas HSQ consistently activates more than ∼12K while maintaining > 90% global code utilization.This result uses batch size 64 per GPU and requires no stochastic sampling tricks.
  • Image Understanding: dRAE demonstrates competitive performance on multimodal understanding benchmarks while using an encoder optimized for reconstruction throughout MLLLM training and testing.The evaluation follows VQRAE’s tuned-image-encoder protocol, while dRAE remains consistent with other baselines in encoder usage.
  • Visual Generation: With ∼700M trainable parameters and relatively modest training scale, dRAE achieves competitive visual generation capabilities against baselines trained on significantly larger datasets.The experiments evaluate text-to-image generation and also consider class-to-image generation with HSQ-derived tokens.

6 Ablative Study

The ablations show that directional routing improves latent-space quantization, while retaining an Euclidean commitment term is crucial for preserving magnitude information and convergence. HSQ also improves multimodal understanding as codebook size increases, unlike VQ-based methods’ stagnant performance.

  • Loss Design: Replacing Euclidean routing with cosine similarity improves performance, indicating that directional matching better fits the latent space.The baseline uses Euclidean routing with ℓ2-based codebook and commitment loss.
  • Loss Design: Defining both losses in cosine space significantly degrades performance, showing that magnitude information remains important for convergence.This is the full-spherical variant, which applies cosine metrics to all components.
  • Loss Design: The best-performing design combines a cosine-based Lcodebook with an ℓ2 commitment loss.Cosine codebook alignment captures directional structure, while Euclidean commitment preserves feature magnitude and guides projection toward the codebook distribution.
  • Impact on Semantics: HSQ performance consistently improves as codebook size increases, whereas VQ-based methods remain relatively stagnant on downstream multimodal understanding.The comparison uses different codebook designs under multimodal model training.

7 Beyond Image Reconstruction

This section tests HSQ beyond image reconstruction through feature reconstruction and joint image-semantic modeling. HSQ outperforms VQ on feature reconstruction and achieves the fastest convergence and highest codebook utilization in joint tasks.

  • Alternative Proxy Tasks: HSQ’s quantization mechanism is evaluated with alternative proxy tasks rather than only image reconstruction.The paper argues that HSQ should apply regardless of the proxy task and reports two preliminary experiments.
  • Feature Reconstruction: With the encoder frozen, a codebook and ViT decoder reconstruct encoded features using an L2 objective after 100K training steps.Performance is evaluated with PSNR and cosine similarity between the original and reconstructed features.
  • Joint Image-Semantic Modeling: A joint model trains HSQ end-to-end between a frozen visual encoder and Qwen-3.5, combining language-modeling and image-reconstruction objectives.A trainable ViT decoder reconstructs images in parallel while the language backbone performs image understanding.
  • Joint Image-Semantic Modeling: K = 131,072 and τ = 0.01 are used for fair comparison with VQ, IBQ, and ℓ2-normalized IBQ under entropy regularization.Entropy regularization is reported as essential for effective codebook utilization when an LLM serves as decoder.
  • Joint Image-Semantic Modeling: HSQ achieves the highest codebook utilization and fastest convergence in image reconstruction and understanding, while VQ performs worst among the compared methods.HSQ obtains these results without randomized sampling, using angular-based routing and codebook updates.

8 Conclusion

dRAE uses Hyper-Spherical Codes to address codebook collapse in high-dimensional discrete visual tokenizers. By aligning codebook updates with the semantic latent distribution, it achieves high utilization, fast convergence, and robust reconstruction with minimal anti-collapsing tricks.

  • dRAE directly tackles codebook collapse in high-dimensional discrete visual tokenizers using Hyper-Spherical Codes.
  • Aligning codebook update dynamics with the underlying semantic latent distribution enables exceptionally high codebook utilization and fast convergence.
  • The approach delivers robust reconstruction quality while requiring minimal anti-collapsing tricks.

A Implementation Details · A.1 Datasets

The implementation uses BLIP3-o imagery for dRAE reconstruction and follows established alignment, pretraining, and instruction-tuning settings for image understanding. T2I generation uses 12M image-text pairs spanning generative pretraining and instruction fine-tuning.

  • A.1 Datasets: dRAE reconstructs 36M open-sourced BLIP3-o images, following UniLIP and VQRAE.These images provide the reconstruction-training data.
  • A.1 Datasets: Image understanding follows LLaVA-1.5 alignment and pretraining settings and includes instruction tuning.The instruction-tuning procedure follows [27].
  • A.1 Datasets: T2I generation uses 12M image-text pairs from Megalith-10M and T2I-2M for generative pretraining, plus BLIP3-o-60K for instruction fine-tuning.The datasets cover both generative pretraining and instruction fine-tuning.

A.2 Tokenizer Training … A.6 Evaluation

The appendix specifies HSQ tokenizer optimization, dRAE ablations, multimodal training setups, discrete diffusion training, and evaluation protocols. It also documents implementation choices for reconstruction and generation experiments, including token compression and sampling settings.

  • A.2 Tokenizer Training: HSQ training takes feature tensors, base codebook embeddings, a learnable projection, and commitment weight β as inputs, then optimizes angular codebook and commitment losses.The algorithm combines Lcodebook and βLcommit into LVQ and uses a straight-through estimator to produce Zout.
  • A.3 Additional Results: dRAE ablations use DINOv2-B at 256 × 256 resolution with a 12-layer ViT decoder, 20000 training steps, and global batch size 512.A RAEv2-style variant sums hidden states from DINOv2 layers {3,6,9,12} to improve reconstruction quality.
  • A.4 MLLM Training: The controlled multimodal comparison uses a 384-resolution encoder, Vicuna-7B backbone, and only LLaVA-1.5 data.The main setup combines publicly available LLaVA-1.5 datasets with 1M randomly sampled image-text pairs from LLaVA-OneVision.
  • A.5 Generative Model Training: DDM corrupts discrete tokens toward absorbing state m through a continuous-time Markov process governed by a decreasing survival function αt with α0 = 1 and α1 = 0.For conditional generation, the sampling derivation assumes perfect reconstruction of x0 and approximates denoising with πθ(xt).
  • A.5 Generative Model Training: At 512 resolution, each image requires 1024 tokens, so a SigLIP2 ViT-large-patch16-384 dRAE compressing images into 576 tokens is additionally trained for efficiency.Large vocabularies also increase output-layer memory, motivating vocabulary-size constraints during generative training.
  • A.5 Generative Model Training: C2I generation uses SigLIP2 ViT-base-patch16-256, 256-token image representations, and a 12-layer transformer trained for 300 epochs on ImageNet.Its ablation uses 65536-vocab tokenizers for 100 epochs with the default iREPA design for B-size DiT models.
  • A.6 Evaluation: Reconstruction follows RAE on ImageNet validation, while noting that UniLIP and VQRAE SSIM evaluation overestimates SSIM by ∼10% because the data range is set to 2.0 instead of 1.0.The stated correction corresponds to the actual image dynamic range.
  • A.6 Evaluation: Understanding uses LLaVA-1.5’s official evaluation script, and generation follows GenEval and DPG-Bench by sampling 4 images per prompt for model-based scoring.Generation samples use 50 steps, a cosine schedule, and constant classifier-free guidance of 3.5.
Loading 2607.22148v1…