MGVQ: Could VQ-VAE Beat VAE? A Generalizable Tokenizer with Multi-group Quantization
作者: Mingkai Jia, Wei Yin, Xiaotao Hu, Jiaxin Guo, Xiaoyang Guo, Qian Zhang, Xiao-Xiao Long, Ping Tan
分类: cs.CV
发布日期: 2025-07-10 (更新: 2025-07-14)
🔗 代码/项目: GITHUB
💡 一句话要点
MGVQ:一种基于多组量化的通用Tokenizer,显著提升VQ-VAE的图像重建质量。
🎯 匹配领域: 支柱四:生成式动作 (Generative Motion)
关键词: VQ-VAE 向量量化 图像重建 多组量化 自编码器
📋 核心要点
- 现有VQ-VAE方法在图像重建质量上与VAE存在差距,主要瓶颈在于离散码本的表示能力不足,导致信息损失。
- MGVQ通过保留潜在维度和引入多组子码本进行量化,增强了离散码本的表示能力,从而更容易优化码本并减少信息损失。
- 实验表明,MGVQ在ImageNet和多个零样本基准测试中均取得了SOTA性能,显著优于现有VQ-VAE方法,甚至超越了SD-VAE。
📝 摘要(中文)
向量量化变分自编码器(VQ-VAE)是将连续视觉数据压缩为离散token的基础模型。现有方法致力于改进量化策略以获得更好的重建质量,但VQ-VAE与VAE之间仍存在较大差距。为了缩小这一差距,我们提出了MGVQ,一种增强离散码本表示能力的新方法,它有助于码本的优化并最小化信息损失,从而提高重建质量。具体来说,我们建议保留潜在维度以保留编码特征,并结合一组子码本进行量化。此外,我们构建了包含512p和2k分辨率的全面零样本基准,以严格评估现有方法的重建性能。MGVQ在ImageNet和所有VQ-VAE的8个零样本基准上均实现了最先进的性能。值得注意的是,与SD-VAE相比,我们在ImageNet上显著优于它们,rFID分别为0.49和0.91,并在所有零样本基准上实现了卓越的PSNR。这些结果突出了MGVQ在重建方面的优越性,并为高清图像处理任务中保持保真度铺平了道路。
🔬 方法详解
问题定义:VQ-VAE旨在将连续的图像数据压缩成离散的token表示,以便进行高效的存储和处理。然而,现有的VQ-VAE方法在重建图像时,其质量与VAE相比仍有差距。主要痛点在于离散码本的表达能力有限,导致在量化过程中不可避免地丢失了部分信息,从而影响了重建质量。
核心思路:MGVQ的核心思路是通过增强离散码本的表达能力来缩小VQ-VAE与VAE之间的差距。具体来说,MGVQ保留了原始的潜在维度,以尽可能地保留编码后的特征信息。同时,引入了多组子码本进行量化,使得每个潜在特征可以被多个码本共同表示,从而提升了码本的整体表达能力。
技术框架:MGVQ的整体框架仍然基于VQ-VAE的自编码器结构。首先,输入图像通过编码器得到潜在表示。然后,潜在表示被划分为多个组,每个组对应一个子码本。对于每个组,选择与其最接近的码本条目进行量化。最后,量化后的潜在表示通过解码器重建图像。MGVQ的关键在于多组量化模块,它负责将潜在表示分配到不同的子码本,并进行量化操作。
关键创新:MGVQ最重要的创新点在于其多组量化策略。与传统的VQ-VAE只使用一个码本进行量化不同,MGVQ使用多个子码本,每个子码本负责量化潜在表示的不同部分。这种多组量化策略可以显著提高码本的表达能力,从而减少量化过程中的信息损失。此外,保留潜在维度也是一个关键创新,它避免了完全依赖码本进行信息编码,从而进一步提升了重建质量。
关键设计:MGVQ的关键设计包括子码本的数量、子码本的大小以及如何将潜在表示分配到不同的子码本。论文中可能采用了经验性的方法来选择这些参数。损失函数方面,除了VQ-VAE中常用的重建损失和码本损失外,可能还引入了额外的正则化项,以防止子码本之间的冗余。具体的网络结构细节(如编码器和解码器的结构)可能与标准的VQ-VAE类似,但针对多组量化模块进行了优化。
🖼️ 关键图片
📊 实验亮点
MGVQ在ImageNet和8个零样本基准测试中均取得了SOTA性能。在ImageNet上,MGVQ的rFID为0.49,显著优于SD-VAE的0.91。在零样本基准测试中,MGVQ在所有测试集上均取得了最高的PSNR值,表明其具有出色的泛化能力和重建质量。这些实验结果充分证明了MGVQ在图像重建方面的优越性。
🎯 应用场景
MGVQ具有广泛的应用前景,尤其是在需要高保真图像重建的领域。例如,它可以应用于高清视频压缩、图像修复、图像生成等任务。通过将图像压缩为离散token,MGVQ可以实现高效的存储和传输,同时保证较高的重建质量。此外,MGVQ还可以作为一种通用的tokenizer,用于训练各种下游任务,如图像分类、目标检测等。
📄 摘要(原文)
Vector Quantized Variational Autoencoders (VQ-VAEs) are fundamental models that compress continuous visual data into discrete tokens. Existing methods have tried to improve the quantization strategy for better reconstruction quality, however, there still exists a large gap between VQ-VAEs and VAEs. To narrow this gap, we propose MGVQ, a novel method to augment the representation capability of discrete codebooks, facilitating easier optimization for codebooks and minimizing information loss, thereby enhancing reconstruction quality. Specifically, we propose to retain the latent dimension to preserve encoded features and incorporate a set of sub-codebooks for quantization. Furthermore, we construct comprehensive zero-shot benchmarks featuring resolutions of 512p and 2k to evaluate the reconstruction performance of existing methods rigorously. MGVQ achieves the state-of-the-art performance on both ImageNet and 8 zero-shot benchmarks across all VQ-VAEs. Notably, compared with SD-VAE, we outperform them on ImageNet significantly, with rFID 0.49 v.s. 0.91, and achieve superior PSNR on all zero-shot benchmarks. These results highlight the superiority of MGVQ in reconstruction and pave the way for preserving fidelity in HD image processing tasks. Code will be publicly available at https://github.com/MKJia/MGVQ.