ガウス変分自己符号化器を用いたベクトル量子化

要旨

ベクトル量子化変分オートエンコーダ（VQ-VAE）は、画像を離散トークンに圧縮する離散型オートエンコーダである。離散化のため学習が困難という課題がある。本論文では、特定の制約を課したガウシアンVAEを学習なしでVQ-VAEに変換する、Gaussian Quant（GQ）と名付けた簡潔かつ効果的な手法を提案する。GQはランダムなガウシアンノイズをコードブックとして生成し、事後平均に最も近いノイズを選択する。理論的には、コードブックサイズの対数がガウシアンVAEのビットバック符号化レートを超える場合、小さな量子化誤差が保証されることを示す。実用的には、GQを効果的に機能させるためのヒューリスティックなガウシアンVAE学習法として、目標発散制約（TDC）を提案する。実験では、GQがUNetとViTの両アーキテクチャにおいて、VQGAN、FSQ、LFQ、BSQなどの従来VQ-VAEを凌駕することを示す。さらに、TDCはTokenBridgeなどの従来のガウシアンVAE離散化手法も改善する。ソースコードは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.

ガウス変分自己符号化器を用いたベクトル量子化

Vector Quantization using Gaussian Variational Autoencoder

要旨

Support