Enhancing Vector Quantization with Distributional Matching: A Theoretical and Empirical Study
This addresses a critical bottleneck in vector quantization for machine learning, offering a solution to improve model performance, though it appears incremental as it builds on existing methods.
The paper tackled training instability and codebook collapse in vector quantization for autoregressive models by aligning feature and code vector distributions using the Wasserstein distance, achieving near 100% codebook utilization and reduced quantization error.
The success of autoregressive models largely depends on the effectiveness of vector quantization, a technique that discretizes continuous features by mapping them to the nearest code vectors within a learnable codebook. Two critical issues in existing vector quantization methods are training instability and codebook collapse. Training instability arises from the gradient discrepancy introduced by the straight-through estimator, especially in the presence of significant quantization errors, while codebook collapse occurs when only a small subset of code vectors are utilized during training. A closer examination of these issues reveals that they are primarily driven by a mismatch between the distributions of the features and code vectors, leading to unrepresentative code vectors and significant data information loss during compression. To address this, we employ the Wasserstein distance to align these two distributions, achieving near 100\% codebook utilization and significantly reducing the quantization error. Both empirical and theoretical analyses validate the effectiveness of the proposed approach.