lucidrains/vector-quantize-pytorch

Provides multiple VQ variants including Residual VQ (stacking quantizers to compress residuals), Grouped Residual VQ (applying quantization to feature dimension groups), and k-means initialization for codebook setup. Uses exponential moving average dictionary updates with techniques to combat dead codebook entries: lower-dimensional codebooks, cosine similarity matching, stale code expiration, and orthogonal regularization. Supports both straight-through estimators and the rotation trick for gradient flow, integrating seamlessly into PyTorch models for image and audio generation tasks.

3,878 stars. Actively maintained with 3 commits in the last 30 days.