Vector Quantization using Gaussian Variational Autoencoder
作者: Tongda Xu, Wendi Zheng, Jiajun He, Jose Miguel Hernandez-Lobato, Yan Wang, Ya-Qin Zhang, Jie Tang
分类: cs.LG, cs.CV
发布日期: 2025-12-07 (更新: 2026-02-05)
💡 一句话要点
提出高斯量化(GQ)方法,无需额外训练即可将高斯VAE转化为高效VQ-VAE
🎯 匹配领域: 支柱四:生成式动作 (Generative Motion)
关键词: 向量量化 变分自编码器 离散表示学习 图像生成 高斯量化
📋 核心要点
- VQ-VAE训练困难,主要挑战在于离散化过程带来的优化难题。
- GQ方法通过先训练高斯VAE,再利用高斯噪声码本进行量化转换,避免了直接训练离散码本。
- 实验表明,GQ在图像生成任务中优于现有VQ-VAE方法,且目标散度约束(TDC)能进一步提升性能。
📝 摘要(中文)
向量量化变分自编码器(VQ-VAE)是一种将图像压缩成离散token的离散自编码器。然而,由于离散化,它们难以训练。本文提出了一种简单而有效的技术,称为高斯量化(GQ),它首先在一定的约束条件下训练一个高斯VAE,然后将其转换为VQ-VAE,而无需额外的训练。对于转换,GQ生成随机高斯噪声作为码本,并找到最接近后验均值噪声向量。从理论上讲,我们证明了当码本大小的对数超过高斯VAE的bits-back编码率时,可以保证较小的量化误差。在实践中,我们提出了一种启发式方法来训练高斯VAE以进行有效的转换,称为目标散度约束(TDC)。在实验上,我们表明GQ在UNet和ViT架构上都优于之前的VQ-VAE,如VQGAN、FSQ、LFQ和BSQ。此外,TDC也改进了之前的高斯VAE离散化方法,如TokenBridge。源代码在补充材料中提供。
🔬 方法详解
问题定义:VQ-VAE旨在学习图像的离散表示,但由于离散码本的梯度问题,直接训练非常困难。现有方法通常需要复杂的训练技巧或损失函数来缓解这个问题,但仍然存在训练不稳定或性能瓶颈。
核心思路:GQ的核心在于避免直接训练离散码本。它首先训练一个连续的高斯VAE,然后利用高斯分布的特性,生成随机高斯噪声作为码本,并通过寻找与高斯VAE后验均值最接近的噪声向量来实现量化。这种方法将离散化过程与VAE的训练解耦,简化了训练流程。
技术框架:GQ包含两个主要阶段:1) 训练一个高斯VAE,该VAE学习将输入图像编码为高斯分布的参数(均值和方差);2) 使用生成的高斯噪声码本对VAE的输出进行量化。具体来说,对于每个图像的编码向量(后验均值),GQ在码本中找到与其欧氏距离最近的向量,并用该向量作为图像的离散表示。
关键创新:GQ的关键创新在于利用高斯分布的性质,将连续的VAE表示转换为离散的量化表示,而无需直接训练离散码本。此外,提出的目标散度约束(TDC)是一种新的训练策略,旨在优化高斯VAE的潜在空间,使其更适合后续的量化过程。
关键设计:目标散度约束(TDC)是GQ的关键设计之一。TDC通过约束高斯VAE的潜在空间分布,使其与目标分布(例如,标准高斯分布)的KL散度接近于某个预设值,从而提高量化性能。此外,码本大小的选择也很重要,论文中给出了理论分析,表明码本大小需要超过高斯VAE的bits-back编码率才能保证较小的量化误差。
🖼️ 关键图片
📊 实验亮点
实验结果表明,GQ在图像生成任务中优于现有的VQ-VAE方法,如VQGAN、FSQ、LFQ和BSQ。例如,在UNet和ViT架构上,GQ都取得了更好的生成质量和更低的重建误差。此外,TDC的引入进一步提升了GQ的性能,并改进了其他高斯VAE离散化方法,如TokenBridge。
🎯 应用场景
GQ方法可应用于图像压缩、图像生成、视频编码等领域。通过将图像或视频转换为离散token,可以实现高效的存储和传输。此外,GQ还可以作为一种有效的特征提取器,用于下游的分类、检测等任务。该方法有望推动多媒体内容处理和人工智能应用的发展。
📄 摘要(原文)
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 the supplementary materials.