Preventing Local Pitfalls in Vector Quantization via Optimal Transport
This addresses training instability for researchers and practitioners using vector-quantized networks, representing a novel method for a known bottleneck rather than an incremental improvement.
The paper tackled training instability in vector-quantized networks by identifying local minima as the cause and integrating optimal transport with the Sinkhorn algorithm to improve assignments, achieving 100% codebook utilization and superior reconstruction quality compared to state-of-the-art methods.
Vector-quantized networks (VQNs) have exhibited remarkable performance across various tasks, yet they are prone to training instability, which complicates the training process due to the necessity for techniques such as subtle initialization and model distillation. In this study, we identify the local minima issue as the primary cause of this instability. To address this, we integrate an optimal transport method in place of the nearest neighbor search to achieve a more globally informed assignment. We introduce OptVQ, a novel vector quantization method that employs the Sinkhorn algorithm to optimize the optimal transport problem, thereby enhancing the stability and efficiency of the training process. To mitigate the influence of diverse data distributions on the Sinkhorn algorithm, we implement a straightforward yet effective normalization strategy. Our comprehensive experiments on image reconstruction tasks demonstrate that OptVQ achieves 100% codebook utilization and surpasses current state-of-the-art VQNs in reconstruction quality.