解决VQ-VAE训练中表示维度坍缩问题，提升模型性能。
Continuous First, Discrete Later: VQ-VAEs Without Dimensional Collapse
- 先用自编码器预训练，再引入向量量化，避免低维坍缩。
- 在图像和音频任务中，重建损失降低17%-35%，感知质量显著提升。
- 理论分析指导暖启动时长，可自适应切换训练阶段，适合优化实践者。
尽管许多改进VQ-VAE的方法聚焦于码本大小与利用率，但其表征维度坍缩（仅占全秩的1-2%）问题仍未解决。本文从理论上和实证上证明，维度坍缩导致一个难以突破的损失下限，现有码本优化方法均无法跨越。通过结合率失真理论与序列学习机制，揭示了向量量化抑制低方差方向是坍缩根源。提出简单有效方案：先以无量化自编码器进行暖启动训练，再引入VQ。在合成数据及大规模图像（VQGAN）和音频（WavTokenizer）VQ-VAE实验中，该方法成功恢复表示维度，在相同训练预算下实现更低重构与感知损失。当码本大小$K \= \{2^{10}, 2^{14}, 2^{16}\}$时，有效维度由3-5提升至17-19，rFID下降17%-35%；在$K \u003d \{2^{13}, 2^{14}\}$的WavTokenizer上，维度从4升至17-19，PESQ提升11%-14%。实验还刻画了暖启动时长对最终损失的影响，理论预测与实验一致，支持基于暖启动长度的自适应切换策略。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。