Vector Quantisation (VQ) is experiencing a comeback in machine learning, where it is increasingly used in representation learning. However, optimizing the codevectors in existing VQ-VAE is not entirely trivial. A problem is codebook collapse, where only a small subset of codevectors receive gradients useful for their optimisation, whereas a majority of them simply ``dies off'' and is never updated or used. This limits the effectiveness of VQ for learning larger codebooks in complex computer vision tasks that require high-capacity representations. In this paper, we present a simple alternative method for online codebook learning, Clustering VQ-VAE (CVQ-VAE). Our approach selects encoded features as anchors to update the ``dead'' codevectors, while optimising the codebooks which are alive via the original loss. This strategy brings unused codevectors closer in distribution to the encoded features, increasing the likelihood of being chosen and optimized. We extensively validate the generalization capability of our quantiser on various datasets, tasks (e.g. reconstruction and generation), and architectures (e.g. VQ-VAE, VQGAN, LDM). Our CVQ-VAE can be easily integrated into the existing models with just a few lines of code.

翻译：向量量化（Vector Quantisation, VQ）在机器学习中正经历复兴，越来越多地被用于表示学习。然而，现有VQ-VAE中码向量的优化并非完全简单。一个问题是码本崩溃，即只有少数码向量能收到对其优化有用的梯度，而大多数码向量会“消亡”并不再被更新或使用。这限制了VQ在需要高容量表示的复杂计算机视觉任务中学习更大码本的效果。本文提出了一种简单的在线码本学习方法——聚类VQ-VAE（Clustering VQ-VAE, CVQ-VAE）。该方法选择编码特征作为锚点来更新“死亡”码向量，同时通过原始损失优化活跃的码本。该策略使未使用的码向量在分布上更接近编码特征，从而增加其被选中和优化的可能性。我们广泛验证了所提量化器在不同数据集、任务（如重建和生成）及架构（如VQ-VAE、VQGAN、LDM）上的泛化能力。CVQ-VAE仅需几行代码即可轻松集成到现有模型中。