Vector Quantization for PyTorch

A practical discrete-representation toolkit for PyTorch. This is an implementation library, not an accompanying paper: it brings together production-ready quantizers and training techniques inspired by published work so they can be used directly in tokenizers, autoencoders, audio codecs, and generative models.

The project is maintained by Aryan and builds on the original work by Phil Wang and the wider open-source contributor community. See License and attribution.

Why this library?

Quantizers turn continuous neural features into a finite vocabulary of discrete codes. That makes it possible to build compact tokenizers, train autoregressive or diffusion models over tokens, and inspect what a model has learned.

This package provides several approaches in one consistent PyTorch API:

Use case	Start with
Standard learned codebook	VectorQuantize
High-fidelity / multi-stage encoding	ResidualVQ or GroupedResidualVQ
No learned codebook	FSQ, FSP, or LFQ
Fixed random tokenizer	RandomProjectionQuantizer
Implicit or structured codebooks	SimVQ, LatentQuantize, or HierarchicalVQ

Install

Install this checkout during development:

pip install -e .

Requires Python 3.9+ and PyTorch 2.4+.

Quick start: Vector Quantization

import torch
from vector_quantize_pytorch import VectorQuantize

quantizer = VectorQuantize(
    dim = 256,
    codebook_size = 512,
    decay = 0.8,
    commitment_weight = 1.0,
)

features = torch.randn(2, 1024, 256)  # batch, tokens, features
quantized, indices, auxiliary_loss = quantizer(features)

assert quantized.shape == features.shape
assert indices.shape == (2, 1024)

Add auxiliary_loss to the loss of your encoder/decoder model during training. At evaluation time, reconstruct vectors from stored token IDs with quantizer.get_output_from_indices(indices).

New in Aryan's fork: codebook health diagnostics

Dead or overused codes are among the most common reasons a VQ model underperforms. VectorQuantize now exposes a lightweight, no-gradient CodebookStats report so your training loop can measure code usage without custom bookkeeping.

quantized, indices, auxiliary_loss, stats = quantizer(
    features,
    return_codebook_stats = True,
)

print(stats.usage_ratio)  # active codes / total codes
print(stats.perplexity)   # effective number of codes being used
print(stats.dead_codes)   # codebook entries unused in this batch
print(stats.counts)       # token count for each code

You can also inspect saved indices later:

stats = quantizer.codebook_stats(indices)

For a shared multi-head codebook, statistics are aggregated into one row. For separate codebooks, each head receives its own row. Masked padding (-1) is ignored. This diagnostic mode cannot be combined with topk output.

Choose a quantization strategy

Standard VQ

VectorQuantize learns a codebook and uses EMA updates by default. It works with sequences, image feature maps, and 3D feature maps.

quantizer = VectorQuantize(
    dim = 256,
    codebook_size = 1024,
    kmeans_init = True,
    kmeans_iters = 10,
    threshold_ema_dead_code = 2,
)

Useful controls include codebook_dim (a lower-dimensional codebook), use_cosine_sim, threshold_ema_dead_code, rotation_trick, directional_reparam, and orthogonal_reg_weight.

For image features, pass accept_image_fmap=True and use (batch, channels, height, width) inputs. For video or volumetric features use accept_3d_fmap=True with (batch, channels, depth, height, width).

Residual and grouped residual VQ

Residual VQ applies several quantizers in sequence, encoding the remaining residual at each stage. It is a strong fit for variable-rate tokenizers and neural audio codecs.

from vector_quantize_pytorch import ResidualVQ

quantizer = ResidualVQ(
    dim = 256,
    num_quantizers = 8,
    codebook_size = 1024,
    stochastic_sample_codes = True,
    sample_codebook_temp = 0.1,
)

quantized, indices, loss = quantizer(features)

Use GroupedResidualVQ to independently quantize groups of the feature dimension. ResidualFSQ, ResidualLFQ, and ResidualSimVQ offer corresponding residual variants for the other families.

Finite Scalar Quantization (FSQ)

FSQ quantizes each scalar to a fixed number of levels rather than learning a vector codebook. It is compact, simple, and avoids codebook-collapse management.

from vector_quantize_pytorch import FSQ

quantizer = FSQ(levels = [8, 5, 5, 5])
quantized, indices = quantizer(torch.randn(2, 1024, 4))

Finite Scalar Perturbation (FSP)

FSP uses structured perturbation during training and fixed scalar levels at inference. It exposes controls for the activation/CDF, perturbation rate, and optional statistical regularization.

from vector_quantize_pytorch import FSP

quantizer = FSP(
    levels = [8, 5, 5, 5],
    act_name = 'normal',
    quantize_rate = 0.5,
    vector_norm = 'var',
)
quantized, indices, regularization_loss, info = quantizer(torch.randn(2, 1024, 4))

Lookup-Free Quantization (LFQ)

LFQ represents tokens with independent binary latents, eliminating embedding lookup. It supports sequences, images, video, and multiple codebooks.

from vector_quantize_pytorch import LFQ

quantizer = LFQ(
    codebook_size = 65536,
    dim = 16,
    entropy_loss_weight = 0.1,
    diversity_gamma = 1.0,
)
quantized, indices, entropy_loss = quantizer(torch.randn(2, 16, 32, 32))

SimVQ and random-projection quantization

SimVQ uses an implicit codebook derived from a frozen base codebook and a learned projection; RandomProjectionQuantizer is a fixed, non-learning tokenizer useful for masked-speech-style objectives.

from vector_quantize_pytorch import SimVQ, RandomProjectionQuantizer

sim_vq = SimVQ(dim = 512, codebook_size = 1024, rotation_trick = True)
rpq = RandomProjectionQuantizer(
    dim = 512, num_codebooks = 16, codebook_dim = 256, codebook_size = 1024
)

Additional modules

Module	Purpose
LatentQuantize	Per-dimension learned scalar codebooks for organized latents
HierarchicalVQ	Multi-scale vector quantization for hierarchical representations
BinaryMapper	Maps tokens between binary and integer representations
Sequential	Small utility for optional module composition

Training notes

• Begin with VectorQuantize(dim, codebook_size) and a decoder reconstruction objective.
• Watch usage_ratio, perplexity, and dead_codes; if use is low, consider codebook_dim, cosine similarity, K-means initialization, or dead-code replacement.
• Keep padding out of the quantizer with mask or lens; output indices at padding locations are -1.
• Add the returned auxiliary loss to your primary model loss. The exact components depend on the quantizer configuration.
• Use eval() before exporting indices when you need deterministic tokens.

Testing

pytest

Research references

This implementation is informed by, among others:

• Neural Discrete Representation Learning (VQ-VAE)
• SoundStream
• Finite Scalar Quantization
• MAGVIT-v2 / Lookup-Free Quantization
• The Rotation Trick
• SimVQ

License and attribution

This repository is an Aryan-maintained fork of the original vector-quantize-pytorch implementation by Phil Wang. It remains available under the MIT License. The original copyright notice is retained in LICENSE, alongside copyright for Aryan's modifications.