[论文解读] SQ-VAE: Variational Bayes on Discrete Representation with Self-annealed Stochastic Quantization

Yuhta Takida, Takashi Shibuya+8 more|arXiv (Cornell University)|May 16, 2022

Domain Adaptation and Few-Shot Learning被引用 13

一句话总结

本文提出 SQ-VAE，一种具有随机量化和自适应退火训练的变分自编码器，通过避免依赖如停止梯度或码书重置等启发式方法，提升了码书利用率。通过引入可学习的随机去量化和量化机制，SQ-VAE 在训练过程中自然地向确定性量化收敛，相较于 VQ-VAE 和标准 VAE，在视觉和语音任务上实现了更优的重建与生成性能。

ABSTRACT

One noted issue of vector-quantized variational autoencoder (VQ-VAE) is that the learned discrete representation uses only a fraction of the full capacity of the codebook, also known as codebook collapse. We hypothesize that the training scheme of VQ-VAE, which involves some carefully designed heuristics, underlies this issue. In this paper, we propose a new training scheme that extends the standard VAE via novel stochastic dequantization and quantization, called stochastically quantized variational autoencoder (SQ-VAE). In SQ-VAE, we observe a trend that the quantization is stochastic at the initial stage of the training but gradually converges toward a deterministic quantization, which we call self-annealing. Our experiments show that SQ-VAE improves codebook utilization without using common heuristics. Furthermore, we empirically show that SQ-VAE is superior to VAE and VQ-VAE in vision- and speech-related tasks.

研究动机与目标

• 解决 VQ-VAE 中的码书坍缩问题，即由于确定性量化和启发式训练方案，仅使用码书向量的一小部分。
• 开发一种符合标准变分贝叶斯范式的训练框架，避免依赖停止梯度、EMA 或码书重置。
• 通过自适应退火过程实现更高的码书利用率，其中量化过程的随机性随时间逐渐降低。
• 设计一种可泛化于多种数据类型的方法，并为类别数据设计专用变体，采用冯·米塞斯-费舍尔分布。
• 在无需超参数调优或辅助技术的前提下，于视觉和语音基准上实现更优的重建与生成性能。

提出的方法

• 引入一种随机去量化过程，利用可微且可学习的概率分布，将连续潜在变量映射到离散码书向量。
• 应用具有可学习后验类别分布的随机量化，实现通过离散选择的梯度传播。
• 采用自适应退火机制，使量化过程的随机性在训练过程中逐渐降低，自然收敛至确定性量化。
• 将目标形式化为证据下界（ELBO），可通过标准反向传播进行优化，无需使用直通估计。
• 提出两种变体：适用于一般数据的高斯 SQ-VAE 和适用于类别数据的 vMF SQ-VAE，后者使用冯·米塞斯-费舍尔分布以提升在离散分布上的性能。
• 使用单一关键超参数（控制随机性的温度）进行训练，遵循 Gumbel-Softmax 或 Concrete 松弛方法。

实验结果

研究问题

• RQ1 在 VAE 框架中使用随机量化是否能相比确定性 VQ-VAE 提升码书利用率？
• RQ2 SQ-VAE 中的自适应退火过程是否能在无需启发式干预的情况下实现更好的表征学习与重建？
• RQ3 SQ-VAE 在视觉和语音任务上的表现相较于标准 VAE 和 VQ-VAE 如何？
• RQ4 vMF SQ-VAE 是否能在图像像素等类别数据分布上显著优于高斯 SQ-VAE？
• RQ5 在 SQ-VAE 中增大码书规模是否能提升性能，而不同于 VQ-VAE 中因坍缩导致的性能停滞或下降？

主要发现

• SQ-VAE 在不使用停止梯度、EMA 或码书重置的情况下，实现了高于 VQ-VAE 的码书利用率，表明其对码书的有效探索。
• SQ-VAE 中的自适应退火过程逐步降低量化中的随机性，实现稳定收敛并提升表征学习能力。
• 在 MNIST 和 Gray-CelebA 数据集上，SQ-VAE 的重建性能优于 VAE 和 VQ-VAE，重建损失更低。
• 在 CelebA-HQ 256×256 数据集上，SQ-VAE 生成了高保真度的重建与生成样本，证明其在大规模数据集上的可扩展性。
• vMF SQ-VAE 在类别数据上显著优于高斯 SQ-VAE，验证了冯·米塞斯-费舍尔分布在离散分布上的有效性。
• 在 SQ-VAE 中增大码书规模可提升性能，而 VQ-VAE 中更大的码书通常导致坍缩且无性能增益。