dRAE: Representation Autoencoder with Hyper-Spherical Codes
Tianren Ma, Lin Long, Chuyan Chen, Mu Zhang, Junbo Zhao, Tong Zhang, Qixiang Ye
cs.CV, cs.AI
2026-07-24
Visual codebooks collapse past 16K tokens. dRAE traces this to a Euclidean-vs-anisotropic geometry mismatch and fixes it with hyper-spherical quantization, scaling to 131K codes at 90%+ utilization.
Discretizing high-dimensional visual features so they interface with language models is usually done with vector quantization (VQ), which pairs features with a codebook (a set of learnable discrete vectors) and replaces each feature with its nearest code. The trouble is that codebooks collapse: past about 16K codes, most codes go unused, the count of active codes stalls, and semantic coherence breaks. This paper hunts the root cause and builds a discretizer that keeps scaling without collapsing.
The authors diagnose collapse as "metric mismatch." Standard VQ uses Euclidean distance for the codebook objective, but visual features actually concentrate on a hyper-spherical shell, with semantic information encoded in orientation (angle), not magnitude (norm). Under Euclidean distance, vectors with large norms dominate code assignment regardless of semantic fit; the codebook embeddings end up with high-variance magnitude scales and uneven angular distributions, and scaling tips it into collapse.
The fix is Hyper-Spherical Quantization (HSQ). Its core is angular routing: code assignment ignores Euclidean distance and uses cosine similarity (angle) between feature and code, decoupling semantic content (direction) from feature magnitude. The codebook loss applies a spherical constraint, while the commitment loss keeps Euclidean distance, because reconstruction needs magnitude. The resulting discrete Representation Autoencoder, dRAE, uses a frozen SigLIP2 ViT-So400M as encoder, a symmetric ViT as decoder, and trains end-to-end in one stage without VQ's stochastic-sampling tricks.
HSQ holds up as the vocabulary grows: at 131,072 codes it keeps over 90% codebook utilization, while traditional VQ stalls past 16K.
On reconstruction, dRAE (HSQ, 131K codes) reaches rFID 0.42, PSNR 24.52, SSIM 0.72, close to SD-VAE (rFID 0.49). At the same codebook size (16K), dRAE's rFID of 0.69 beats VQRAE's 1.31.
A diagnostic proves the decoupling is necessary: normalizing features (keeping only direction) barely hurts understanding (MMBench 82.2 to 81.1) but tanks reconstruction (PSNR 22.5 to 20.6, rFID 4.62 to 9.57). Direction carries semantics; magnitude carries reconstruction detail; both must stay.
Downstream tasks confirm it. For multimodal understanding with Qwen2.5-7B plus the dRAE encoder, understanding climbs as HSQ's vocabulary grows (16K to 65K: GQA 34.3 to 35.8, MMB 44.1 to 45.6), whereas VQ degrades when scaled. For text-to-image, dRAE using only about 12M training samples and 700M parameters reaches GenEval 0.63 and DPG 80.58, close to DALL-E 3 (GenEval 0.67, DPG 83.50), which trained on 1B+ samples. At the feature level, HSQ hits PSNR 8.24 and similarity 0.92, clearly above VQ's 4.91 and 0.81.
For anyone building unified multimodal models that share one set of discrete tokens across understanding and generation, this method removes the long-standing "vocabulary won't scale" blocker, and it drops VQ's fragile sampling tricks for a simpler training pipeline. It reconciles two goals that usually fight, semantic understanding and visual reconstruction, with one clean design: angle governs semantics, Euclidean governs magnitude. A codebook that scales to six figures while staying fully utilized matters a lot for treating visual tokens like language tokens in autoregressive models.
The authors note that reconstructing high-dimensional patch features is harder than reconstructing pixels, owing to intrinsic noise in encoder representations and the absence of perceptual supervision. A practical concern: a 131K vocabulary inflates the output-layer memory of generative models considerably. One more thing to keep straight: the encoder is frozen SigLIP2, so a large share of the performance rests on that pretrained vision foundation. The comparisons against VQ methods hold codebook size fixed, but whether training-data volume and decoder scale are fully aligned is not itemized in the text, so the size of the cross-method absolute advantage should be read with some discount.