SoftVQ-VAE: Efficient 1-Dimensional Continuous Tokenizer

📄 arXiv: 2412.10958v3 📥 PDF

作者: Hao Chen, Ze Wang, Xiang Li, Ximeng Sun, Fangyi Chen, Jiang Liu, Jindong Wang, Bhiksha Raj, Zicheng Liu, Emad Barsoum

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

发布日期: 2024-12-14 (更新: 2025-03-14)

备注: Code and model: https://github.com/Hhhhhhao/continuous_tokenizer

💡 一句话要点

SoftVQ-VAE：一种高效的1维连续图像令牌化方法，显著加速生成模型推理。

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

关键词: 图像生成 令牌化 VQ-VAE 生成模型 Transformer 软量化 高压缩率

📋 核心要点

  1. 现有图像令牌化方法在高压缩率下效率不足，限制了生成模型的训练和推理速度。
  2. SoftVQ-VAE利用软类别后验概率聚合多个码字，增加潜在空间表示能力，实现高效令牌化。
  3. 实验表明，SoftVQ-VAE显著提升图像生成速度，同时保持或提升生成质量，并已开源。

📝 摘要（中文）

本文提出SoftVQ-VAE，一种连续图像令牌化方法，它利用软类别后验概率将多个码字聚合到每个潜在令牌中，从而显著提高潜在空间的表示能力。当应用于基于Transformer的架构时，我们的方法可以使用少至32或64个1维令牌来压缩256x256和512x512的图像。SoftVQ-VAE不仅表现出一致且高质量的重建效果，更重要的是，它还在不同的基于去噪的生成模型中实现了最先进且显著更快的图像生成结果。值得注意的是，SoftVQ-VAE在生成256x256图像时，推理吞吐量提高了18倍，在生成512x512图像时提高了55倍，同时对于SiT-XL，实现了具有竞争力的FID分数1.78和2.21。它还通过减少2.3倍的训练迭代次数，同时保持相当的性能，提高了生成模型的训练效率。凭借其完全可微的设计和语义丰富的潜在空间，我们的实验表明SoftVQ-VAE实现了高效的令牌化，而不会影响生成质量，为更高效的生成模型铺平了道路。代码和模型已发布。

🔬 方法详解

问题定义：论文旨在解决图像生成模型中图像令牌化效率低下的问题。现有方法在实现高压缩率的同时，往往难以保证生成图像的质量和推理速度，限制了生成模型的应用。

核心思路：SoftVQ-VAE的核心思路是使用软类别后验概率来聚合多个码字，从而增加潜在空间的表示能力。与传统的VQ-VAE只选择一个最接近的码字不同，SoftVQ-VAE允许每个潜在令牌携带多个码字的信息，从而更有效地表示图像内容。

技术框架：SoftVQ-VAE的整体框架基于VQ-VAE，包括编码器、码本和解码器。编码器将输入图像映射到潜在空间，然后使用码本将潜在向量量化为离散的令牌。与传统VQ-VAE不同的是，SoftVQ-VAE使用softmax函数计算每个码字的概率，并使用这些概率来加权码本中的码字。解码器接收量化后的令牌，并将其解码为重建的图像。整个过程是端到端可微的。

关键创新：SoftVQ-VAE的关键创新在于使用软类别后验概率进行码字聚合。这种方法允许每个潜在令牌携带多个码字的信息，从而显著提高了潜在空间的表示能力。此外，SoftVQ-VAE使用1维令牌序列，进一步降低了计算复杂度。

关键设计：SoftVQ-VAE的关键设计包括：1) 使用Gumbel-Softmax技巧来近似argmax操作，使其可微；2) 使用KL散度来约束软类别后验概率的分布，防止其过于集中；3) 使用残差连接和注意力机制来提高模型的表达能力。具体的损失函数包括重建损失、量化损失和KL散度损失。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

SoftVQ-VAE在图像生成任务中取得了显著的性能提升。在256x256图像生成任务中，推理吞吐量提高了18倍；在512x512图像生成任务中，推理吞吐量提高了55倍。同时，对于SiT-XL模型，实现了具有竞争力的FID分数1.78和2.21。此外，SoftVQ-VAE还能够将生成模型的训练迭代次数减少2.3倍，同时保持相当的性能。

🎯 应用场景

SoftVQ-VAE可应用于各种图像生成任务，如图像合成、图像编辑、超分辨率等。其高效的令牌化能力可以显著加速生成模型的训练和推理，降低计算成本，并推动生成模型在移动设备和嵌入式系统上的应用。此外，该方法还可以扩展到其他领域，如音频生成和视频生成。

📄 摘要（原文）

Efficient image tokenization with high compression ratios remains a critical challenge for training generative models. We present SoftVQ-VAE, a continuous image tokenizer that leverages soft categorical posteriors to aggregate multiple codewords into each latent token, substantially increasing the representation capacity of the latent space. When applied to Transformer-based architectures, our approach compresses 256x256 and 512x512 images using as few as 32 or 64 1-dimensional tokens. Not only does SoftVQ-VAE show consistent and high-quality reconstruction, more importantly, it also achieves state-of-the-art and significantly faster image generation results across different denoising-based generative models. Remarkably, SoftVQ-VAE improves inference throughput by up to 18x for generating 256x256 images and 55x for 512x512 images while achieving competitive FID scores of 1.78 and 2.21 for SiT-XL. It also improves the training efficiency of the generative models by reducing the number of training iterations by 2.3x while maintaining comparable performance. With its fully-differentiable design and semantic-rich latent space, our experiment demonstrates that SoftVQ-VAE achieves efficient tokenization without compromising generation quality, paving the way for more efficient generative models. Code and model are released.