dRAE：基于超球面码的表示自编码器

用角相似度替代欧氏距离做码本路由，解决高维视觉量化坍缩。

前置知识

向量量化（VQ）

向量量化是一种把连续向量映射到离散码本（codebook）中最近邻条目的离散化技术。给定一个含 $K$ 个可学习条目的码本 $C=\{c_i\}_{i=1}^K$，标准 VQ 用欧氏距离 $\arg\min_{c_i} \|z-c_i\|^2$ 找到最近码字，再用直通梯度估计器（STE）反向传播。它是离散图像 tokenizer（如 VQGAN）和语音、多模态离散建模的基础。

本文的研究对象就是 VQ 在高维表示空间中失效的问题，必须先理解 VQ 的赋值规则和损失（码本损失+承诺损失）才能看懂它为何会坍缩。

表示自编码器（RAE / VQRAE）

RAE 不像传统 VAE 那样把图像压缩到 8-64 维的低维瓶颈，而是直接使用冻结的视觉基础模型（如 SigLIP2 ViT-So400M）输出的高维（768-1536 维）、语义丰富特征作为潜在表示，用一个可训练的 ViT 解码器重建像素。其离散版本 VQRAE 再用 SimVQ 对这些高维 token 量化，以便用于自回归/离散扩散建模。

dRAE 是 VQRAE 的直接改进，理解 RAE 范式（重建作为代理任务、蒸馏保持语义、为何需要高维潜在）才能理解 HSQ 要量化的对象是什么。

码本坍缩（Codebook Collapse）

指码本中大部分条目从未被激活、只有少数码字垄断几乎所有赋值，导致有效词表远小于名义容量。表现是 active code 数量低、利用率差、增大词表也带不来增益。常见诱因包括码字初始化不当、指数移动平均更新失衡、目标度量与数据几何不匹配等。

本文的核心要解决的问题就是高维空间的码本坍缩，作者把根因归结为'度量失配'（欧氏距离 vs 角向语义），这是全文论证主线。

视觉基础模型与超球面几何

CLIP、SigLIP2、DINOv2、JEPA 等视觉基础模型经对比/自监督训练后，patch 嵌入近似集中在一个薄球壳上（norm 集中在 $\sqrt{d}$ 附近），语义信息主要编码在向量的方向（角向），半径（模长）则承载结构纹理信息。余弦相似度 $\cos\theta=\frac{z\cdot c}{\|z\|\|c\|}$ 度量的是方向而非尺度。

本文两个关键假设 H1（语义在球面上）和 H2（模长对重建必不可少）正是建立在这种各向异性、薄球壳几何之上，是 HSQ 设计的直接动机。

研究动机

要把视觉理解（高维语义编码）与生成（VAE 低维重建）统一，业界用表示自编码器 RAE 把 SigLIP2 等基础模型的高维（1536 维）特征直接解码为像素，再用向量量化得到离散 token。但现有 VQ 方法（如 VQRAE 用 SimVQ）在高维空间遭遇严重码本坍缩：作者实测在词表 $K=65536$、每 GPU batch 64 时，VQRAE 每次 forward 仅激活约 6K 个码，而 dRAE 能激活约 12K；公开基线显示 tokenizer 在词表超过 16K 代码后就触顶，无法继续从更大词表获益。坍缩使得有限的词表无法承载高维潜在空间的丰富语义，直接抵消了表示式设计的核心优势。作者进一步把根因诊断为'度量失配'：VQ 的赋值与更新都用欧氏 $\ell_2$ 距离，而视觉特征的语义编码在角向方向上，欧氏度量对模长敏感，会让大范数向量'劫持'赋值，产生高方差模长、聚集角向分布的码本，从而触发坍缩。

本文的目标是设计一种可扩展的离散表示 tokenizer，能在高维视觉特征空间上实现：(1) 抗坍缩的高码本利用率——支持词表扩大到 131072 仍保持 100% 利用率；(2) 同时保持语义方向（用于多模态理解）和模长（用于像素重建）两种信息，使理解与生成双赢；(3) 简化训练流程——无需特殊初始化、抗坍缩技巧、课程学习，端到端单阶段训练；(4) 在重建（rFID/PSNR/SSIM）、多模态理解（GQA/TextVQA/MMBench 等）和生成（C2I 的 gFID/IS，T2I 的 GenEval/DPG-Bench）上同时取得强结果。

与已有工作不同的是，作者的关键切入点是'解耦'：把用于码本路由（赋值）的度量与量化目标本身分离开来。以往要么整体用欧氏（VQ/SimVQ）导致坍缩，要么整体归一化到单位球面（SVQ）丢掉模长而损害重建；FSQ 则靠固定格点取整、维度爆炸而无法用于高维。本文只把'最近邻赋值'和'码本损失'换成角向/余弦，却刻意保留 $\ell_2$ 承诺损失来监督模长，并保留未归一化的码向量送入解码器。这样既顺应了语义的方向性几何，又留住了重建必需的尺度信息——这是区别于 FSQ、SVQ 的本质不同，也是突破坍缩的独到角度。

核心方法

直觉上：既然语义住在超球面的'方向'上、模长只服务于重建，量化时就该按方向找邻居、按方向更新码本，却把模长原样留给解码器。技术路线上 dRAE 沿用 VQRAE 框架——冻结 SigLIP2 ViT-So400M 作编码器 $E$，加可学习线性投影，用对称 ViT 作解码器 $D$，重建损失 $\mathcal{L}_{rec}=\mathcal{L}_1(\hat{x},x)+\omega_p\mathcal{L}_{perc}+\omega_d\mathcal{L}_{disc}$ 并配教师自蒸馏 $\lambda\|Z-T(x)\|^2$。区别只在量化算子 $Q(\cdot)$：赋值从欧氏 $\arg\min\|z-c_i\|^2$ 换成余弦 $\arg\max (Z_i\cdot C_j)/(\|Z_i\|\|C_j\|)$（角向路由），码本损失换成球面形式 $1-\mathrm{sg}[Z/\|Z\|]\cdot Z_q/\|Z_q\|$，但承诺损失仍用欧氏 $\|Z-\mathrm{sg}[Z_q]\|^2_2$。整套方法单阶段端到端训练，无需任何抗坍缩技巧。

核心创新是'度量解耦'：路由用角向、模长用欧氏。具体而言，赋值阶段只用余弦相似度比较方向（消除模长偏差带来的尺度劫持），码本更新约束在球面切空间（让码字在方向上均匀铺开、避免聚集，实现'幅度同质化'与'角向均匀性'，如图1c所示），而承诺损失保留欧氏形式以监督投影后特征的模长——因为表7表明若承诺损失也换成余弦（完全球面），rFID 反而恶化到 12.7、SSIM 跌到 0.33。换言之，作者发现'纯角向'会让模长完全无监督而发散收敛，必须留一个 $\ell_2$ 锚点；这种'半欧半角'的非对称设计是与 SimVQ（全欧氏）、SVQ（全归一化）、FSQ（无码本格点）的根本区别。

方法步骤详情

完整流程为：(1) 冻结 SigLIP2 ViT-So400M 对输入图像 $x$ 提取 patch 特征 $Z=E(x)\in\mathbb{R}^{N\times d}$；(2) 经可学习线性投影送入量化；(3) 角向路由：对每个 token 算 $S_{ij}=(Z_i\cdot C_j)/(\|Z_i\|\|C_j\|)$，取 $I_i=\arg\max_j S_{ij}$，再查表 $Z_q=C[I]$（取未归一化码向量）；(4) 直通估计把梯度从 $Z_q$ 复制到投影特征；(5) 计算重建损失（L1+感知+判别）、球面码本损失、欧氏承诺损失 $\mathcal{L}_{commit}=\|Z-\mathrm{sg}[Z_q]\|^2_2$ 与教师蒸馏损失；(6) 端到端单阶段反向传播更新解码器、投影层、码本与微调编码器。完整伪代码见 Algorithm 1，对照了经典 VQ 与 HSQ 的差异。理解与生成下游再分别接 Qwen2.5-7B（LLaVA 式指令微调）与离散扩散 Transformer（Flan-T5-XXL 条件）。

技术新颖性

新颖性体现在三方面：第一，诊断层面，首次系统论证坍缩源于'度量失配'而非单纯初始化或更新策略，并用图2的码本梯度范数曲线直接观测到 VQ 的坍缩模式；第二，方法层面，提出非对称度量解耦——这是介于'全欧氏 VQ'与'全归一化 SVQ'之间的新设计点，且证明完全球面会失效（表7第4行 rFID 12.7），凸显保留 $\ell_2$ 承诺项的必要性；第三，范式层面，把重建作为代理任务但验证 HSQ 可迁移到特征重建（表9）与联合图像-语义建模（图4，与 IBQ/IBQ-L2 比较），表明其思想不绑定具体代理任务。相对 FSQ（无码本、维度爆炸）与 SVQ（强加单位球先验、损害重建），HSQ 在原生高维空间保持高利用率又留住模长。

**Fig. 1**: Illustration of the code space defined by VQ (a) and HSQ (b), and distribution of their learned codebook embeddings (c).

**Algorithm 1**: Classic Vector Quantization vs Hyper-Spherical Quantization

实验结果

发现：(1) 诊断（表1）证明语义在球面——特征归一化后多模态理解仅降约 1%，但重建 PSNR 22.5→20.6、rFID 4.62→9.57，验证 H1 与 H2。(2) 重建（表2）：dRAE@16384 的 rFID 0.69/PSNR 24.03 优于 VQRAE 1.31/22.33；@131072 达 rFID 0.42，居离散 tokenizer 之首且超 SD-VAE(0.49) 与连续 RAE(0.62)。图3显示随词表 $2^{14}\to2^{17}$ dRAE 单调升优而 VQ 触顶。(3) 理解（表3）：dRAE@512 的 TextVQA 67.7、MMBench 81.5，超 VQRAE 58.8/67.6。(4) 生成：T2I 仅 12M 数据即获 GenEval 0.63、DPG 80.58，逼近 >100M 数据的 Janus；C2I@65536 gFID 4.45 优于 VQ。(5) 码本利用率：$K=65536$ 时 VQRAE 每 batch 约 6K 活跃码，dRAE 超 12K、全局 >90%。消融（表7）与联合建模（图4）印证角向路由优越。

**Table 1**: Diagnostic Analysis. Comparison of understanding (a) and reconstruction (b) performance between raw and normalized image features as inputs.

**Table 2**: Evaluation on reconstruction quality.

**Table 3**: Evaluation on multimodal understanding.

**Table 4**: Evaluation on text-to-image generation.

**Table 7**: Ablation of objective metric on reconstruction tasks.

**Table 8**: Ablation of tokenizer design on multimodal understanding tasks under identical visual instruction tuning setting.

**Table 9**: Comparison of VQ and HSQ on feature reconstruction.

**Fig. 3**: Scaling behavior with K=2^14~2^17 vocabulary size. From left to right: active code per batch, reconstruction PSNR, and rFID.

**Fig. 4**: Comparison of quantization methods under the joint image-semantic optimization.

查看结构化数据

任务	指标	本文	基线	提升
图像重建（ImageNet-1K）	rFID ↓	0.42（K=131072）/ 0.69（K=16384）	VQRAE 1.31（K=16384）	K=16384 时 rFID 几乎减半；扩到 131072 再降约 47%
图像重建	PSNR ↑ / SSIM ↑	24.52 / 0.72（K=131072）	VQRAE 22.33 / 0.67	PSNR +2.19 dB，SSIM +0.05
多模态理解	MMBench-en ↑ / TextVQA ↑	81.5 / 67.7（512分辨率）	VQRAE 67.6 / 58.8	MMBench +13.9、TextVQA +8.9
文生图 T2I	GenEval Overall / DPG-Bench	0.63 / 80.58（12M 数据）	Janus 0.61 / 79.68（>100M 数据）	用约 1/8 数据即略超 Janus
类别生成 C2I	gFID ↓ / IS ↑	4.45 / 287.3（K=65536）	VQ 5.51 / 264.3	gFID −1.06，IS +23
码本利用率	每 batch 活跃码数	约 12K（K=65536）	VQRAE 约 6K	活跃码翻倍，全局利用率 >90%

局限与改进

作者承认的局限：特征重建（表9）仍是开放难题——HSQ 虽把 PSNR 从 VQ 的 4.91 提到 8.24、相似度 0.81→0.92，但作者指出重建高维 patch 嵌入比重建像素难得多，受编码器表示内在噪声与缺乏感知监督所限，VAE 在这方面反而更受益。生成模型规模偏小（约 700M 可训练参数）、数据量克制（T2I 仅 12M 对），作者坦言把详尽的数据与模型规模调优留待未来。我额外观察到的局限：(1) 主要只在 SigLIP2 一种编码器上验证，虽然附录给了 DINOv2，但跨更多基础模型（如 EVA、InternVL-C）的普适性待证；(2) 实验框架高度依赖 RAE 范式与自蒸馏教师，迁移到非 RAE 体系的成本未量化；(3) 对'为何完全球面会发散'只给经验结论，缺理论分析；(4) 缺少在更大规模真实多模态评测集上的端到端结果。

独立分析的弱点

弱点一：方法效果部分依赖自蒸馏教师（原始 VFM 初始化的 $T$）与 SigLIP2 几何特性，若换到角向结构不规则的基础模型（如某些自监督模型薄球壳假设较弱），角向路由的增益可能下降——改进方向是给出更通用的几何度量自适应选择或可学习的度量参数。弱点二：联合图像-语义建模（图4）必须额外加熵正则损失（沿用 IBQ 的设计）才能保证利用率，说明 HSQ 在有 LLM 解码器、监督信号复杂的场景下并不天然免技巧——改进方向是设计针对 LLM 解码的稳定化项。弱点三：词表到 131072 时虽利用率 100%，但解码器/词表嵌入参数与显存开销线性增长，端侧部署不友好——改进方向是结合 FSQ 的格点思想做混合离散化以压缩码本存储。弱点四：完全球面变体崩溃（表7 rFID 12.7）暴露模长监督的脆弱性，若能引入显式模长正则或多尺度承诺项，可能进一步放宽对 $\ell_2$ 承诺的强依赖。弱点五：评测以代理任务为主，缺少与最新大一统 tokenizer（如 TokenFlow、Show-o 升级版）在同等数据/算力下的严格 apples-to-apples 对比。

未来方向

作者明确提出的方向：把 HSQ 推广到更广的代理任务，验证它作为'高维表示通用量化器'的普适性（文中已初步试了特征重建表9 与联合图像-语义建模图4）；并做更详尽的数据与模型规模 curation（T2I 部分原话'leave the exhaustive data and model curation for future study'）。基于成果可延伸的方向：(1) 在更大 LLM（如 70B+）与更大扩散 Transformer 上验证 dRAE token 的 scaling law；(2) 探索角向路由与残差量化、二进制量化（如 BSQ）结合，进一步推高词表或压缩比特率；(3) 给'度量失配→坍缩'建立理论模型，从优化几何推导临界维度与坍缩阈值；(4) 将 HSQ 用于视频/3D 表示 tokenizer，测试时序维度对模长/角向解耦的敏感度；(5) 研究熵正则与角向路由的协同，去掉对 IBQ 风格正则的依赖。

复现评估

复现评估：算法核心（HSQ 的赋值与三项损失）在 Algorithm 1 给出了与经典 VQ 并排的完整伪代码，公式清晰、含 STE 与码本投影说明，定性易复现。关键组件（SigLIP2 ViT-So400M 编码器、Qwen2.5-7B、Flan-T5-XXL、离散扩散）均为公开模型，提供了项目主页 https://drae-hsq.github.io（暗示可能有代码/模型发布）。但论文正文只给'简要实现概述'，完整超参与训练细节指向附录 A（本次未提供全文），训练规模较大（ViT-So400M+ViT 解码器+7B LLM+扩散），单卡难以复现，需显著算力。C2I/T2I 数据规模（12M~60M 对）也是门槛。词表/分辨率/码本利用率指标都有明确数值，便于横向核对。综合：核心方法可复现性中上，完整端到端复现（尤其生成部分）需要团队级算力与数据。

dRAE：基于超球面码的表示自编码器 dRAE: Representation Autoencoder with Hyper-Spherical Codes