Continuous First, Discrete Later: VQ-VAEs Without Dimensional Collapse
作者: Xinyu Zhao, Nikita Karagodin, Hamed Hassani, Sinan Hersek, Paul Pu Liang, Yury Polyanskiy
分类: cs.LG
发布日期: 2026-05-07
💡 一句话要点
提出AE预热策略以解决VQ-VAE中的维度坍缩问题，显著提升表征质量与重建性能
🎯 匹配领域: 支柱四：生成式动作 (Generative Motion)
关键词: VQ-VAE 维度坍缩 表征学习 向量量化 生成模型 自编码器
📋 核心要点
- 核心问题：VQ-VAE训练中普遍存在维度坍缩现象，导致表征仅占据极低维子空间，限制了模型性能的进一步提升。
- 方法要点：提出“先连续后离散”的训练策略，在正式引入向量量化前，先以无量化自编码器模式进行预热训练。
- 实验效果：在VQGAN和WavTokenizer上，该方法显著提升了表征的有效维度，并大幅改善了图像重建质量与音频感知指标。
📝 摘要（中文）
尽管许多改进VQ-VAE性能的方法集中于码本大小和利用率，但维度坍缩问题——即训练后的VQ-VAE表征仅存在于极低维子空间（仅占满秩的1-2%）——尚未得到解决。本文从理论和实验层面证明，维度坍缩导致了硬性的损失下界，使得现有的码本改进技术难以突破。我们扩展了Saxe等人（2014）的序贯学习理论，结合率失真理论，解释了VQ操作如何通过抑制低方差方向导致潜在空间坍缩。基于此，我们提出了一种简单有效的解决方案：在引入量化之前，先将模型作为无量化自编码器（AE）进行“预热”。在合成实验、大规模图像（VQGAN）和音频（WavTokenizer）任务中，AE预热成功恢复了表征维度，在相同训练预算下显著降低了重建和感知损失。实验表明，该方法将VQGAN和WavTokenizer的有效维度从个位数提升至17-19，并大幅优化了rFID和PESQ指标。
🔬 方法详解
问题定义：论文旨在解决VQ-VAE训练过程中出现的“维度坍缩”问题。现有方法多关注码本容量，却忽视了潜在表征在训练过程中逐渐退化至极低维子空间，导致模型无法充分利用高维信息，从而形成难以逾越的性能瓶颈。
核心思路：通过理论分析发现，VQ操作在训练初期会过早抑制低方差方向，导致表征空间坍缩。因此，核心思路是“先连续后离散”，即在训练初期通过无量化的自编码器（AE）学习阶段，让编码器先在连续空间内构建高质量、高维度的表征，待表征稳定后再引入量化层。
技术框架：整体流程分为两个阶段：第一阶段为AE预热期，模型作为标准自编码器训练，不进行向量量化；第二阶段为VQ-VAE微调期，引入码本并开启量化操作。模型架构保持不变，仅通过控制量化层的激活时机实现训练策略的平滑过渡。
关键创新：首次从理论上揭示了VQ操作与维度坍缩的因果关系，并提出了一种无需修改网络结构、仅通过训练策略调整即可恢复表征维度的通用方案，打破了传统VQ-VAE训练的性能下界。
关键设计：引入了基于预热时长（Warm-up duration）的自适应切换准则。理论分析预测了下游性能与预热长度的函数关系，使得研究者可以根据训练预算和任务需求，精确控制从AE到VQ-VAE的转换时机，从而实现最优的重建效果。
🖼️ 关键图片
📊 实验亮点
实验证明，AE预热策略效果显著：在VQGAN任务中，将码本有效维度从3-5提升至17-19，rFID指标降低17-35%；在WavTokenizer任务中，有效维度从4提升至17-19，PESQ指标提升11-14%。该方法在不同码本规模（$K=2^{10}$至$2^{16}$）下均表现出鲁棒的性能增益，证明了其在处理大规模生成任务时的优越性。
🎯 应用场景
该方法可广泛应用于生成式AI领域，特别是基于离散表征的生成模型，如图像生成（VQGAN）、音频压缩与合成（WavTokenizer）以及多模态大模型中的视觉/音频Token化器。通过提升表征的有效维度，该技术能显著改善生成内容的保真度与感知质量，在高质量媒体压缩与生成任务中具有极高的实际应用价值。
📄 摘要（原文）
While many approaches to improve VQ-VAE performance focus on codebook size and utilization, the effect of dimensional collapse, where trained VQ-VAE representations live in an extremely low-dimensional subspace (1-2% of full rank), remains unaddressed. We show theoretically and empirically that dimension collapse causes a hard loss lower bound that various codebook improvement techniques fail to surpass. Our analytic framework extends the sequential learning effect of Saxe et al. [2014] by introducing ideas from rate-distortion theory and explains how the latent collapse is caused by the VQ suppressing lower-variance directions. Our theory justifies a simple solution: a "warm-up phase" that trains the model as an (unquantized) autoencoder before introducing VQ. On both synthetic experiments and large-scale image (VQGAN) and audio (WavTokenizer) VQ-VAEs, we show that AE Warm-Up successfully restores representation dimension, leading to lower reconstruction and perceptual loss at the same training budget. Across codebook sizes $K \in$ {$2^{10}, 2^{14}, 2^{16}$}, AE warm-up raises VQGAN codebook effective dimension from 3-5 to 17-19 and reduces rFID by 17-35%; on WavTokenizer at $K \in$ {$2^{13}, 2^{14}$}, it raises codebook dimension from 4 to 17-19 and improves PESQ by 11-14%. We empirically characterize how warm-up duration governs the achievable final loss. In agreement with experiment, our theoretical analysis predicts downstream performance as a function of warm-up length, enabling an adaptive criterion for switching from AE Warm-up to VQ-VAE training.