Quantize-then-Rectify: Efficient VQ-VAE Training

📄 arXiv: 2507.10547v1 📥 PDF

作者: Borui Zhang, Qihang Rao, Wenzhao Zheng, Jie Zhou, Jiwen Lu

分类: cs.CV, cs.LG

发布日期: 2025-07-14

💡 一句话要点

提出ReVQ框架，通过量化修正加速VQ-VAE训练，降低计算成本。

🎯 匹配领域: 支柱四：生成式动作 (Generative Motion) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: VQ-VAE 视觉Tokenizer 量化 预训练VAE 图像压缩

📋 核心要点

  1. 高压缩率VQ-VAE训练计算成本高昂，需要大量GPU资源，限制了其应用。
  2. ReVQ框架利用预训练VAE，通过量化和修正策略，大幅降低VQ-VAE的训练成本。
  3. ReVQ在ImageNet上实现了高压缩率和高质量重建，训练成本显著降低。

📝 摘要（中文）

视觉tokenizer在多模态大模型中至关重要，它们充当连续输入和离散token之间的桥梁。然而，训练高压缩率的VQ-VAE仍然需要大量的计算资源，通常需要数千GPU小时。本文表明，通过控制VAE的容错阈值内的量化噪声，可以将预训练的VAE高效地转换为VQ-VAE。我们提出了Quantize-then-Rectify (ReVQ)框架，利用预训练的VAE以最小的计算开销实现快速VQ-VAE训练。通过集成通道多组量化来扩大码本容量，并使用后置修正器来减轻量化误差，ReVQ将ImageNet图像压缩为最多512个token，同时保持了具有竞争力的重建质量(rFID = 1.06)。值得注意的是，相对于最先进的方法，ReVQ将训练成本降低了两个数量级以上：ReVQ在单个NVIDIA 4090上大约22小时内完成完整训练，而类似方法在32个A100 GPU上需要4.5天。实验结果表明，ReVQ实现了卓越的效率-重建权衡。

🔬 方法详解

问题定义：论文旨在解决训练高压缩率VQ-VAE时计算资源需求过高的问题。现有方法通常需要大量的GPU时间和算力，这限制了VQ-VAE在资源受限环境下的应用，阻碍了其在大规模多模态模型中的部署。

核心思路：论文的核心思路是利用预训练的VAE模型，通过控制量化噪声，将其高效地转换为VQ-VAE。通过在VAE的潜在空间中引入量化操作，并使用修正机制来减轻量化误差，从而实现快速且高效的VQ-VAE训练。这种方法避免了从头开始训练VQ-VAE，显著降低了计算成本。

技术框架：ReVQ框架主要包含以下几个阶段：1) 预训练VAE：首先训练一个标准的VAE模型。2) 量化：在VAE的潜在空间中应用通道多组量化，将连续的潜在表示转换为离散的token。3) 修正：使用后置修正器来减轻量化误差，提高重建质量。整个框架利用预训练的VAE作为基础，通过量化和修正两个关键步骤，实现高效的VQ-VAE训练。

关键创新：ReVQ的关键创新在于其量化和修正策略。通道多组量化可以有效扩大码本容量，提高模型的表达能力。后置修正器能够减轻量化引入的误差，保证重建质量。与从头开始训练VQ-VAE的方法相比，ReVQ利用预训练的VAE，大大减少了训练时间和计算资源。

关键设计：ReVQ的关键设计包括：1) 通道多组量化：将潜在空间的通道划分为多个组，对每个组进行独立的量化，从而增加码本容量。2) 后置修正器：使用一个小的神经网络来学习量化误差的分布，并对量化后的潜在表示进行修正，提高重建质量。3) 损失函数：使用重建损失和量化损失的加权和作为训练目标，平衡重建质量和量化误差。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

ReVQ在ImageNet图像压缩任务中，使用至多512个token，实现了rFID=1.06的重建质量。与现有方法相比，ReVQ将训练成本降低了两个数量级以上，在单个NVIDIA 4090上仅需约22小时即可完成训练，而同类方法在32个A100 GPU上需要4.5天。实验结果表明，ReVQ在效率和重建质量之间取得了更好的平衡。

🎯 应用场景

ReVQ框架可应用于多模态大模型中的视觉tokenizer，实现高效的图像压缩和表示学习。其潜在应用领域包括图像生成、视频处理、以及需要高效数据压缩的边缘计算设备。该研究降低了VQ-VAE的训练成本，有望加速多模态模型的开发和部署，并促进其在资源受限环境下的应用。

📄 摘要（原文）

Visual tokenizers are pivotal in multimodal large models, acting as bridges between continuous inputs and discrete tokens. Nevertheless, training high-compression-rate VQ-VAEs remains computationally demanding, often necessitating thousands of GPU hours. This work demonstrates that a pre-trained VAE can be efficiently transformed into a VQ-VAE by controlling quantization noise within the VAE's tolerance threshold. We present \textbf{Quantize-then-Rectify (ReVQ)}, a framework leveraging pre-trained VAEs to enable rapid VQ-VAE training with minimal computational overhead. By integrating \textbf{channel multi-group quantization} to enlarge codebook capacity and a \textbf{post rectifier} to mitigate quantization errors, ReVQ compresses ImageNet images into at most 512 tokens while sustaining competitive reconstruction quality (rFID = 1.06). Significantly, ReVQ reduces training costs by over two orders of magnitude relative to state-of-the-art approaches: ReVQ finishes full training on a single NVIDIA 4090 in approximately 22 hours, whereas comparable methods require 4.5 days on 32 A100 GPUs. Experimental results show that ReVQ achieves superior efficiency-reconstruction trade-offs.