Gaussian Mixture Vector Quantization with Aggregated Categorical Posterior

TL;DR

This paper introduces a Gaussian mixture vector quantization framework that enhances data modeling in discrete latent spaces, improving codebook utilization and reducing information loss without heuristics.

Contribution

It generalizes VQ-VAE by integrating Gaussian mixtures into a variational Bayesian framework and proposes ALBO as a new optimization objective.

Findings

Improved codebook utilization in experiments.

Reduced information loss compared to baseline methods.

Avoided heuristics used in traditional VQ-VAE training.

Abstract

The vector quantization is a widely used method to map continuous representation to discrete space and has important application in tokenization for generative mode, bottlenecking information and many other tasks in machine learning. Vector Quantized Variational Autoencoder (VQ-VAE) is a type of variational autoencoder using discrete embedding as latent. We generalize the technique further, enriching the probabilistic framework with a Gaussian mixture as the underlying generative model. This framework leverages a codebook of latent means and adaptive variances to capture complex data distributions. This principled framework avoids various heuristics and strong assumptions that are needed with the VQ-VAE to address training instability and to improve codebook utilization. This approach integrates the benefits of both discrete and continuous representations within a variational Bayesian framework. Furthermore, by introducing the \textit{Aggregated Categorical Posterior Evidence Lower Bound} (ALBO), we offer a principled alternative optimization objective that aligns variational distributions with the generative model. Our experiments demonstrate that GM-VQ improves codebook utilization and reduces information loss without relying on handcrafted heuristics.