无需训练即可将高斯VAE转为VQ-VAE,性能超越多个现有方法。
Training-Free Vector Quantization via Gaussian VAEs
- 用约束训练的高斯VAE生成噪声作为码本,直接转换为VQ-VAE
- 在UNet和ViT上均优于VQGAN、FSQ等主流VQ-VAE方法
- 适用于希望快速部署离散编码器的研究者与工程师
向量量化变分自编码器(VQ-VAEs)能将图像压缩为离散标记,但因量化过程难以训练。本文提出简单有效的Gaussian Quant(GQ)技术:先在特定约束下训练一个高斯VAE,再无需额外训练将其转化为VQ-VAE。转换时,以随机高斯噪声构成码本,并选取与后验均值最接近的向量作为码本项。理论上证明,当码本大小的对数超过高斯VAE的bits-back编码率时,可保证小量化误差。实践中,提出目标分歧约束(TDC)以优化高斯VAE训练,提升转换效果。实验表明,GQ在UNet与ViT架构上均优于VQGAN、FSQ、LFQ及BSQ等方法。此外,TDC亦可提升如TokenBridge等已有高斯VAE离散化方法的表现。源代码已开源于https://github.com/tongdaxu/VQ-VAE-from-Gaussian-VAE。
原文摘要 · Abstract (English)
Vector-quantized variational autoencoders (VQ-VAEs) are discrete autoencoders that compress images into discrete tokens. However, they are difficult to train due to discretization. In this paper, we propose a simple yet effective technique dubbed Gaussian Quant (GQ), which first trains a Gaussian VAE under certain constraints and then converts it into a VQ-VAE without additional training. For conversion, GQ generates random Gaussian noise as a codebook and finds the closest noise vector to the posterior mean. Theoretically, we prove that when the logarithm of the codebook size exceeds the bits-back coding rate of the Gaussian VAE, a small quantization error is guaranteed. Practically, we propose a heuristic to train Gaussian VAEs for effective conversion, named the target divergence constraint (TDC). Empirically, we show that GQ outperforms previous VQ-VAEs, such as VQGAN, FSQ, LFQ, and BSQ, on both UNet and ViT architectures. Furthermore, TDC also improves previous Gaussian VAE discretization methods, such as TokenBridge. The source code is provided in https://github.com/tongdaxu/VQ-VAE-from-Gaussian-VAE.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。