Online Clustered Codebook

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). CVQ-VAE can be easily integrated into the existing models with just a few lines of code.

1. Introduction

Vector Quantisation (VQ) [12] is a basic building block of many machine learning techniques. It is often used to help learning unsupervised representations for vision and language tasks, including data compression [1], [39], [36], recognition [26], [3], [44], [24], [23], and generation [37], [31], [11], [32], [47], [34], [33]. VQ quantises continuous feature vectors into a discrete space by embedding them to the closest vectors in a codebook of representatives or codevectors. Quantisation has been shown to simplify optimization problems by reducing a continuous search space to a discrete one.