[论文解读] Robust Training of Vector Quantized Bottleneck Models

Adrian Łańcucki, Jan Chorowski+7 more|arXiv (Cornell University)|May 18, 2020

Speech Recognition and Synthesis参考文献 35被引用 8

一句话总结

本文通过将码本学习视为在线k-means聚类，提出改进的VQ-VAE训练技术。引入了更高的码本学习率、量化前的批量归一化，以及基于数据的周期性码本重初始化，实现了更稳定的训练、更广泛的码本利用，以及在语音、图像和单元发现任务中的性能提升。

ABSTRACT

In this paper we demonstrate methods for reliable and efficient training of discrete representation using Vector-Quantized Variational Auto-Encoder models (VQ-VAEs). Discrete latent variable models have been shown to learn nontrivial representations of speech, applicable to unsupervised voice conversion and reaching state-of-the-art performance on unit discovery tasks. For unsupervised representation learning, they became viable alternatives to continuous latent variable models such as the Variational Auto-Encoder (VAE). However, training deep discrete variable models is challenging, due to the inherent non-differentiability of the discretization operation. In this paper we focus on VQ-VAE, a state-of-the-art discrete bottleneck model shown to perform on par with its continuous counterparts. It quantizes encoder outputs with on-line $k$-means clustering. We show that the codebook learning can suffer from poor initialization and non-stationarity of clustered encoder outputs. We demonstrate that these can be successfully overcome by increasing the learning rate for the codebook and periodic date-dependent codeword re-initialization. As a result, we achieve more robust training across different tasks, and significantly increase the usage of latent codewords even for large codebooks. This has practical benefit, for instance, in unsupervised representation learning, where large codebooks may lead to disentanglement of latent representations.

研究动机与目标

• 解决因编码器输出非平稳性和码本初始化敏感性导致的VQ-VAE训练不稳定与收敛性差的问题。
• 提升大规模VQ-VAE模型中的码本利用率，以实现更好的解耦与表征质量。
• 通过将码本学习视为在线k-means聚类，稳定深度离散瓶颈模型的训练动态。
• 评估架构与优化改进对图像生成、语音识别和无监督单元发现等多样化下游任务性能的影响。
• 证明多种训练增强方法的组合可带来优于单一改进的更优且更一致的结果。

提出的方法

• 将VQ-VAE的码本学习视为在线k-means聚类，利用已知的k-means优化策略。
• 相对于编码器，提高码本学习率，以加速收敛并提升码本稳定性。
• 在量化层之前应用批量归一化，以平衡潜在特征的大小并改善梯度流动。
• 通过使用离线k-means对近期编码器输出进行聚类，实现周期性的、基于数据的码本重初始化。
• 使用直通估计器进行反向传播通过离散量化操作，保持端到端可微性。
• 将三种增强方法——提高码本学习率、批量归一化和重初始化——结合使用，以实现最优训练性能。

实验结果

研究问题

• RQ1 码本初始化规模如何影响VQ-VAE训练的稳定性和收敛性？
• RQ2 周期性的、基于数据的码本重初始化在多大程度上缓解了训练过程中编码器激活的非平稳性？
• RQ3 为码本和编码器分别设置学习率是否能改善训练动态与码本利用率？
• RQ4 在量化前应用批量归一化如何影响表征质量与训练稳定性？
• RQ5 多种训练增强方法的综合效应在图像、语音和单元发现任务的下游性能中表现如何？

主要发现

• 周期性的、基于数据的码本重初始化显著提升了训练收敛性，并降低了初始BPD值，尤其在与其他增强方法结合时效果更明显。
• 提高码本学习率可加快收敛速度并提升训练稳定性，尤其在码本较大时更为显著。
• 在量化前应用批量归一化在CIFAR-10的所有设置中均带来微小但一致的BPD（比特/维）改进。
• 将提高码本学习率、批量归一化和周期性重初始化相结合，可在所有评估任务中实现最佳性能与最鲁棒的训练。
• 码本重初始化有助于防止码本坍缩，并提升潜在码字的使用率，即使在大码本中亦能增强表征的解耦性。
• 所提出的训练策略在监督与无监督任务中均带来下游性能提升，包括语音识别与单元发现，且在BPD与聚类指标上均表现出一致增益。