HQ-VAE: Hierarchical Discrete Representation Learning with Variational Bayes

📄 arXiv: 2401.00365v2 📥 PDF

作者: Yuhta Takida, Yukara Ikemiya, Takashi Shibuya, Kazuki Shimada, Woosung Choi, Chieh-Hsin Lai, Naoki Murata, Toshimitsu Uesaka, Kengo Uchida, Wei-Hsiang Liao, Yuki Mitsufuji

分类: cs.LG, cs.AI, cs.CV

发布日期: 2023-12-31 (更新: 2024-03-28)

备注: 34 pages with 17 figures, accepted for TMLR

💡 一句话要点

HQ-VAE：提出基于变分贝叶斯的层级离散表示学习框架，提升重建质量。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱四：生成式动作 (Generative Motion)

关键词: 分层表示学习 向量量化 变分自编码器 变分贝叶斯 离散表示

📋 核心要点

  1. 层级VQ-VAE存在码本崩溃问题，导致码本利用率低，重建精度下降。
  2. HQ-VAE基于变分贝叶斯框架，以随机方式学习层级离散表示，缓解码本崩溃。
  3. 实验表明，HQ-VAE提升了码本利用率，改善了图像和音频数据的重建性能。

📝 摘要（中文）

向量量化(VQ)是一种使用离散码本表示确定性地学习特征的技术。它通常与变分自编码模型VQ-VAE一起使用，VQ-VAE可以进一步扩展到分层结构，以实现高保真重建。然而，VQ-VAE的这种分层扩展通常受到码本/层崩溃问题的影响，即码本不能有效地用于表达数据，从而降低重建精度。为了缓解这个问题，我们提出了一个新的统一框架，在变分贝叶斯框架的基础上随机学习分层离散表示，称为分层量化变分自编码器(HQ-VAE)。HQ-VAE自然地推广了VQ-VAE的分层变体，如VQ-VAE-2和残差量化VAE (RQ-VAE)，并为它们提供了贝叶斯训练方案。我们在图像数据集上的综合实验表明，HQ-VAE增强了码本的使用，提高了重建性能。我们还在音频数据集上验证了HQ-VAE在不同模态中的适用性。

🔬 方法详解

问题定义：论文旨在解决分层VQ-VAE中存在的码本崩溃问题。现有方法在构建深层离散表示时，容易出现部分码本未被充分利用的情况，导致模型表达能力受限，最终影响重建质量。

核心思路：HQ-VAE的核心思路是将分层VQ-VAE置于变分贝叶斯框架下，通过引入随机性来鼓励码本的均匀使用。通过贝叶斯训练方案，模型能够学习到更鲁棒的离散表示，从而缓解码本崩溃问题。

技术框架：HQ-VAE的整体架构仍然是自编码器结构，包含编码器、量化器和解码器。与传统VQ-VAE不同的是，HQ-VAE在每一层量化器中引入了变分推断，将码本选择视为一个概率过程。编码器将输入映射到潜在空间，量化器将潜在表示离散化，解码器则根据离散码本重建输入。整个框架通过变分下界进行优化。

关键创新：HQ-VAE的关键创新在于将变分贝叶斯框架引入到分层离散表示学习中。通过对码本选择过程进行建模，HQ-VAE能够学习到更具表达能力的离散表示，并有效缓解码本崩溃问题。此外，HQ-VAE统一了VQ-VAE-2和RQ-VAE等分层VQ-VAE变体，并为它们提供了贝叶斯训练方案。

关键设计：HQ-VAE的关键设计包括：1) 使用变分推断对码本选择进行建模，引入KL散度项来约束码本的使用；2) 设计合适的先验分布，鼓励码本的均匀使用；3) 使用重参数化技巧进行梯度反向传播；4) 通过调整KL散度的权重来平衡重建误差和码本利用率。

📊 实验亮点

实验结果表明，HQ-VAE在图像数据集上显著提升了重建性能，并有效缓解了码本崩溃问题。与VQ-VAE-2和RQ-VAE等基线模型相比，HQ-VAE在重建质量和码本利用率方面均取得了更好的结果。此外，HQ-VAE在音频数据集上的实验也验证了其在不同模态上的适用性。

🎯 应用场景

HQ-VAE可应用于图像、音频等多种模态数据的压缩、生成和表示学习。在图像领域，可用于高保真图像重建、图像编辑和图像生成等任务。在音频领域，可用于音频压缩、语音合成和音乐生成等任务。该研究有助于提升多媒体数据的处理和理解能力。

📄 摘要（原文）

Vector quantization (VQ) is a technique to deterministically learn features with discrete codebook representations. It is commonly performed with a variational autoencoding model, VQ-VAE, which can be further extended to hierarchical structures for making high-fidelity reconstructions. However, such hierarchical extensions of VQ-VAE often suffer from the codebook/layer collapse issue, where the codebook is not efficiently used to express the data, and hence degrades reconstruction accuracy. To mitigate this problem, we propose a novel unified framework to stochastically learn hierarchical discrete representation on the basis of the variational Bayes framework, called hierarchically quantized variational autoencoder (HQ-VAE). HQ-VAE naturally generalizes the hierarchical variants of VQ-VAE, such as VQ-VAE-2 and residual-quantized VAE (RQ-VAE), and provides them with a Bayesian training scheme. Our comprehensive experiments on image datasets show that HQ-VAE enhances codebook usage and improves reconstruction performance. We also validated HQ-VAE in terms of its applicability to a different modality with an audio dataset.