Optimizing codebook training through control chart analysis

Frequent use of codebook resets to enhance the usage of the Vector Quantization Variational Autoencoder (VQ-VAE) may significantly alter the codebook distribution and consequently diminish the training efficiency. In this work, we introduce a novel codebook learning approach called Exponentially Weighted Moving Average Control VQ-VAE (ECVQ-VAE). This method considers the nearest neighbor distance of the codebook during training as a monitoring sample and constructs a control line. Our quantizer restricts the update process of codebook vectors based on whether the drift during monitoring exceeds the control line while simultaneously adjusting the overall usage distribution of the codebook by promoting competition. This process enables an optimization that sustains full codebook usage while reducing the training demands. We demonstrate that our approach achieves better results than the existing methods in lightweight scenarios and extensively validate the generalizability of our quantizer across various datasets, tasks, and architectures (VQ-VAE, VQ-GAN).