Balance of Number of Embedding and their Dimensions in Vector Quantization

TL;DR

The paper tackles the challenge of balancing the codebook size $N$ and embedding dimension $D$ in VQ-VAE under a fixed capacity $W=ND$ by showing that larger $N$ with smaller $D$ can enhance reconstruction. It introduces an adaptive dynamic quantization mechanism based on attention and $Gumbel$-Softmax to select an optimal codebook per data point, coupled with EMA-updated codebooks. Through six benchmark datasets, it demonstrates that the adaptive approach outperforms the best fixed-codebook configurations at the same capacity and reveals a two-stage learning dynamic where the model initially leverages larger codebooks and later specializes. This work highlights the value of flexible, per-instance discretization for improving discrete latent representations and offers a scalable path to better VQ-based reconstruction and generation tasks.

Abstract

The dimensionality of the embedding and the number of available embeddings ( also called codebook size) are critical factors influencing the performance of Vector Quantization(VQ), a discretization process used in many models such as the Vector Quantized Variational Autoencoder (VQ-VAE) architecture. This study examines the balance between the codebook sizes and dimensions of embeddings in VQ, while maintaining their product constant. Traditionally, these hyper parameters are static during training; however, our findings indicate that augmenting the codebook size while simultaneously reducing the embedding dimension can significantly boost the effectiveness of the VQ-VAE. As a result, the strategic selection of codebook size and embedding dimensions, while preserving the capacity of the discrete codebook space, is critically important. To address this, we propose a novel adaptive dynamic quantization approach, underpinned by the Gumbel-Softmax mechanism, which allows the model to autonomously determine the optimal codebook configuration for each data instance. This dynamic discretizer gives the VQ-VAE remarkable flexibility. Thorough empirical evaluations across multiple benchmark datasets validate the notable performance enhancements achieved by our approach, highlighting the significant potential of adaptive dynamic quantization to improve model performance.

Figures (9)

• Figure 1: Schematic diagram of VQ-VAE network with adaptive dynamic quantization structure: $(1)$Utilizes the multi-head "attention mechanism" and Gumbel-Softmax operation to achieve dynamic selection of the quantization codebook. $(2)$Both the encoder and decoder employ CNN+ResNet architectures. $(3)$The experimental model includes functions for dimension transformation before and after quantization within the quantization layer to facilitate data flow.
• Figure 2: Gradient gap and quantization loss under fixed codebook models.
• Figure 3: The frequency of adaptive selection of different-sized codebooks throughout the training process. The training can be divided into two stages: Initially, the model predominantly learns from the codebook with the most codewords; Subsequently, the model optimizes and increasingly selects the type of codebook with the minimum reconstruction loss in the fixed codebook model. The experimental results on the Diabetic Retinopathy dataset are slightly less conspicuous, but in the final stage of training, there is still a tendency to choose the most suitable codebook
• Figure 4: Comparison of quantization loss and gradient gap between adaptive dynamic quantization model and fixed codebook model. The blue color represents the model with the smallest gradient gap and the largest quantization loss in the fixed codebook model, the yellow color represents the fixed codebook model with the smallest reconstruction loss, and the red color represents the adaptive dynamic quantization model. The gradient gap and quantization loss of the adaptive quantizer are kept at low values.
• Figure 5: Sample images of Diabetic Retinopathy dataset. Left: Original image. Right: Reconstructed image.