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
视觉特征离散化的 codebook 扩到 1.6 万码字以上就会崩塌。这篇诊断根因是欧氏距离和特征各向异性几何错配,提出用余弦相似度做码字分配的超球面量化,词表扩到 13.1 万仍保持 90% 以上利用率。
要把高维视觉特征离散化、好跟语言模型对接,主流做法是向量量化(VQ),给它配一个 codebook(码本,一组可学习的离散向量),每个特征去找最近的码字代替。问题在于 codebook 会「崩塌」:扩到 1.6 万码字以上,大部分码字用不起来,活跃码字数量卡住上不去,语义也保不住。这篇要找崩塌的根因,并做出一个能继续扩、扩了还不崩的离散化方法。
作者把崩塌的根因诊断为「度量错配」:标准 VQ 用欧氏距离做 codebook 目标,但视觉特征实际集中在一个超球壳上,语义信息编码在方向(夹角)里,不在幅度(模长)里。用欧氏距离的结果是,模长大的向量不管语义对不对都会主导码字分配,codebook 嵌入的幅度尺度方差大、角度分布不均,一扩就崩。
解法是 Hyper-Spherical Quantization(HSQ,超球面量化)。核心是「角度路由」:码字分配不看欧氏距离,看特征和码字的余弦相似度(夹角),把语义内容(方向)和特征幅度解耦。codebook 损失用球面约束,但 commitment loss 还保留欧氏距离,因为重建需要幅度信息。最终的离散表征自编码器 dRAE 用冻结的 SigLIP2 ViT-So400M 当编码器,一个对称的 ViT 当解码器,端到端一阶段训练,不用 VQ 里那套随机采样的 trick。
HSQ 在扩词表时表现稳定:扩到 13.1 万码字,codebook 利用率仍高于 90%,而传统 VQ 在 1.6 万以上就涨不动了。
重建质量上,dRAE(HSQ,13.1 万码字)rFID 0.42、PSNR 24.52、SSIM 0.72,接近 SD-VAE 的水平(rFID 0.49)。同码字数(1.6 万)下,dRAE 的 rFID 0.69 好于 VQRAE 的 1.31。
一个关键诊断证明了解耦的必要性:把特征归一化(只留方向)后,理解类任务几乎不掉(MMBench 82.2→81.1),但重建质量大幅下滑(PSNR 22.5→20.6、rFID 4.62→9.57)。说明方向承载语义、幅度承载重建细节,两者都得留。
下游也验证了。多模态理解上,用 Qwen2.5-7B 配 dRAE 编码器,HSQ 扩词表时理解能力跟着涨(16K→65K,GQA 34.3→35.8、MMB 44.1→45.6),而 VQ 扩词表反而退化。文生图上,dRAE 只用约 1200 万训练数据、7 亿参数,GenEval 0.63、DPG 80.58,接近用了 10 亿+ 数据的 DALL-E 3(GenEval 0.67、DPG 83.50)。特征重建层面,HSQ 的 PSNR 8.24、相似度 0.92,明显优于 VQ 的 4.91、0.81。
对做统一多模态模型(理解加生成共用一套离散 token)的人,这套方法解决了「词表扩不动」这个长期卡点,而且不用 VQ 那套脆弱的采样 trick,训练流程更简单。它把语义理解和视觉重建这两件原本打架的事用一个干净的设计(角度管语义、欧氏管幅度)调和了,codebook 能扩到十万级还保持满利用率,这对把视觉 token 当语言 token 一样做自回归很关键。
作者自己点出,重建高维 patch 特征比重建像素难,因为编码器表征里有固有噪声、又没有感知监督。另一个实际问题是,13.1 万这么大的词表会显著撑大生成模型输出层的内存。另外要拎清楚:HSQ 的编码器是冻结的 SigLIP2,效果的底子很大程度来自这个预训练视觉基础模型;论文里和 VQ 系的对比是同码字数下的,但训练数据量、解码器规模这些变量是否完全对齐,正文没有逐项交代,跨方法的绝对优势有多大需要打折看。