Group-Wise Optimization for Self-Extensible Codebooks in Vector Quantized Models

📄 arXiv: 2510.13331v2 📥 PDF

作者: Hong-Kai Zheng, Piji Li

分类: cs.CV

发布日期: 2025-10-15 (更新: 2025-10-16)

💡 一句话要点

提出Group-VQ，通过分组优化自扩展码书解决VQ-VAE中的码书坍塌问题

🎯 匹配领域: 支柱四：生成式动作 (Generative Motion)

关键词: 向量量化 变分自编码器 码书学习 分组优化 自监督学习

📋 核心要点

  1. VQ-VAE中码书坍塌限制了模型性能，现有方法或采用静态码书，或全局联合优化，均约束了码书的学习能力。
  2. Group-VQ通过分组优化码书，组内联合优化，组间独立优化，提升码书利用率和重建效果的平衡。
  3. 引入免训练的码书重采样方法，支持训练后灵活调整码书大小，实验表明Group-VQ在图像重建任务上性能提升。

📝 摘要（中文）

向量量化变分自编码器(VQ-VAE)利用自监督学习，通过重建任务，使用码书中最近的向量来表示连续向量。然而，码书坍塌等问题仍然存在于VQ模型中。为了解决这些问题，现有的方法采用隐式的静态码书或联合优化整个码书，但这些方法限制了码书的学习能力，导致重建质量下降。本文提出了Group-VQ，它对码书进行分组优化。每个组独立优化，组内进行联合优化。这种方法改善了码书利用率和重建性能之间的权衡。此外，我们还引入了一种无需训练的码书重采样方法，允许在训练后调整码书大小。在各种设置下的图像重建实验中，Group-VQ证明了重建指标的性能有所提高，并且训练后的码书采样方法实现了调整码书大小的期望灵活性。

🔬 方法详解

问题定义：VQ-VAE中的码书坍塌问题，即部分码字被过度使用，而另一些码字则很少或根本不被使用，导致码书利用率低，影响重建质量。现有方法要么使用静态码书，限制了码书的学习能力；要么全局联合优化整个码书，计算复杂度高，且容易陷入局部最优。

核心思路：将码书分成多个组，每个组独立优化，组内进行联合优化。这样既能保证组内码字的协同学习，又能避免全局优化带来的计算负担和局部最优问题。通过分组优化，可以更好地平衡码书利用率和重建性能。

技术框架：Group-VQ沿用VQ-VAE的整体框架，包括编码器、量化层和解码器。关键在于量化层，它将编码器的输出映射到码书中的最近向量。不同之处在于，Group-VQ将码书划分为多个组，每个组独立进行量化和优化。训练过程中，每个组的码字只与该组内的样本进行关联和更新。

关键创新：核心创新在于分组优化策略。与全局优化相比，分组优化降低了计算复杂度，并允许每个组专注于学习特定的特征表示。与静态码书相比，分组优化赋予了码书更强的学习能力和适应性。此外，提出的免训练码书重采样方法，允许在训练后灵活调整码书大小，进一步提升了模型的灵活性。

关键设计：分组数量是关键参数，需要根据数据集和任务进行调整。损失函数包括重建损失和码书更新损失。重建损失衡量解码器的输出与原始输入的差异。码书更新损失鼓励码字向其最近的编码器输出移动，并防止码书坍塌。免训练码书重采样方法基于码字的使用频率进行采样，保留使用频率高的码字，并随机初始化新的码字。

📊 实验亮点

实验结果表明，Group-VQ在图像重建任务上优于现有的VQ-VAE变体。在不同数据集和设置下，Group-VQ均取得了更高的重建质量，并有效缓解了码书坍塌问题。此外，免训练码书重采样方法能够灵活调整码书大小，并在保持重建质量的同时，降低计算复杂度。

🎯 应用场景

Group-VQ可应用于图像、音频、视频等多种模态数据的压缩、重建和生成任务。例如，在图像压缩领域，Group-VQ可以学习到更有效的图像表示，从而实现更高的压缩率和更好的重建质量。在生成模型领域，Group-VQ可以作为一种有效的先验知识，引导生成过程，生成更逼真的图像或音频。

📄 摘要（原文）

Vector Quantized Variational Autoencoders (VQ-VAEs) leverage self-supervised learning through reconstruction tasks to represent continuous vectors using the closest vectors in a codebook. However, issues such as codebook collapse persist in the VQ model. To address these issues, existing approaches employ implicit static codebooks or jointly optimize the entire codebook, but these methods constrain the codebook's learning capability, leading to reduced reconstruction quality. In this paper, we propose Group-VQ, which performs group-wise optimization on the codebook. Each group is optimized independently, with joint optimization performed within groups. This approach improves the trade-off between codebook utilization and reconstruction performance. Additionally, we introduce a training-free codebook resampling method, allowing post-training adjustment of the codebook size. In image reconstruction experiments under various settings, Group-VQ demonstrates improved performance on reconstruction metrics. And the post-training codebook sampling method achieves the desired flexibility in adjusting the codebook size.