Векторное квантование с использованием гауссовского вариационного автоэнкодера
Vector Quantization using Gaussian Variational Autoencoder
December 7, 2025
Авторы: Tongda Xu, Wendi Zheng, Jiajun He, Jose Miguel Hernandez-Lobato, Yan Wang, Ya-Qin Zhang, Jie Tang
cs.AI
Аннотация
Векторно-квантованный вариационный автоэнкодер (VQ-VAE) — это дискретный автоэнкодер, сжимающий изображения в дискретные токены. Его обучение затруднено из-за дискретизации. В данной статье мы предлагаем простой, но эффективный метод, названный Gaussian Quant (GQ), который преобразует гауссовский VAE с определённым ограничением в VQ-VAE без обучения. GQ генерирует случайный гауссовский шум в качестве кодбука и находит ближайший шум к апостериорному среднему. Теоретически мы доказываем, что когда логарифм размера кодбука превышает скорость кодирования с возвращением битов (bits-back coding) гауссовского VAE, гарантируется малая ошибка квантования. Практически мы предлагаем эвристику для обучения гауссовского VAE с целью эффективного применения GQ, названную ограничением на целевую дивергенцию (TDC). Экспериментально мы показываем, что GQ превосходит предыдущие VQ-VAE, такие как VQGAN, FSQ, LFQ и BSQ, как на архитектуре UNet, так и на ViT. Кроме того, TDC также улучшает предыдущие методы дискретизации гауссовских VAE, такие как TokenBridge. Исходный код доступен по адресу https://github.com/tongdaxu/VQ-VAE-from-Gaussian-VAE.
English
Vector quantized variational autoencoder (VQ-VAE) is a discrete auto-encoder that compresses images into discrete tokens. It is difficult to train due to discretization. In this paper, we propose a simple yet effective technique, dubbed Gaussian Quant (GQ), that converts a Gaussian VAE with certain constraint into a VQ-VAE without training. GQ generates random Gaussian noise as a codebook and finds the closest noise 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 VAE for effective GQ, named 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 upon previous Gaussian VAE discretization methods, such as TokenBridge. The source code is provided in https://github.com/tongdaxu/VQ-VAE-from-Gaussian-VAE.