Vector Quantization - Pytorch
A PyTorch library of vector and scalar quantization techniques, transcribed from DeepMind's TensorFlow implementation, used for models such as VQ-VAE-2 and Jukebox.
Vector Quantization - Pytorch is a library of vector and scalar quantization methods for PyTorch, originally transcribed from DeepMind's TensorFlow implementation and packaged for convenient reuse. It updates its codebook dictionary using exponential moving averages, and has been used by DeepMind and OpenAI for image generation (VQ-VAE-2) and music generation (Jukebox).
Beyond the base VectorQuantize class, it includes a ResidualVQ class that recursively quantizes the residuals of a signal, with an optional DiVeQ update technique that updates codebooks by gradients instead of auxiliary losses, and a GroupedResidualVQ class that runs residual quantization on separate groups of the feature dimension, matching results reported for Encodec with fewer codebooks. It also supports codebook initialization from k-means centroids of the first batch, and stochastic code sampling with a shared codebook across quantizers.
It installs as a Python package.