vector-quantize-pytorch

To prevent 'dead' codebook entries in VectorQuantize, you can use several techniques:

1. Lower codebook dimension: Project encoder values to a lower dimension before quantization using codebook_dim.
2. Cosine similarity: Use L2 normalization for codes and encoded vectors by setting use_cosine_sim = True.
3. Expiring stale codes: Replace codes with low EMA cluster sizes using threshold_ema_dead_code (e.g., set to 2 to replace codes with fewer than 2 hits).

import torch
from vector_quantize_pytorch import VectorQuantize

# Example: Lower codebook dimension
vq = VectorQuantize(
    dim = 256,
    codebook_size = 256,
    codebook_dim = 16
)

# Example: Cosine similarity
vq = VectorQuantize(
    dim = 256,
    codebook_size = 256,
    use_cosine_sim = True
)

# Example: Expiring stale codes
vq = VectorQuantize(
    dim = 256,
    codebook_size = 512,
    threshold_ema_dead_code = 2
)

The ResidualVQ class implements residual vector quantization, which uses multiple quantizers to recursively quantize residuals. This is useful for high-quality audio or image generation.

Parameters:

• dim: The dimensionality of the input vectors.
• num_quantizers: Number of quantizers to use.
• codebook_size: The size of each codebook.
• diveq: Set to True to use the DiVeQ technique (gradient-based codebook updates) instead of EMA.
• stochastic_sample_codes: Set to True to stochastically sample codes (RQ-VAE style).
• sample_codebook_temp: Temperature for stochastic sampling.
• shared_codebook: Set to True to share the same codebook across all quantizers.
• kmeans_init: Set to True to initialize the codebook using k-means centroids of the first batch.
• kmeans_iters: Number of k-means iterations for initialization.

Returns (on call):

• quantized: The quantized vectors.
• indices: The indices for each quantization layer.
• commit_loss: The commitment loss.
• all_codes: (Optional) If return_all_codes=True is passed to the call, returns all codes across quantization layers.

import torch
from vector_quantize_pytorch import ResidualVQ

residual_vq = ResidualVQ(
    dim = 256,
    num_quantizers = 8,
    codebook_size = 1024,
)

x = torch.randn(1, 1024, 256)

# Standard call
quantized, indices, commit_loss = residual_vq(x)

# Get all codes across layers
quantized, indices, commit_loss, all_codes = residual_vq(x, return_all_codes = True)

The VectorQuantize class implements standard vector quantization using exponential moving averages (EMA) to update the dictionary. It is suitable for tasks like VQ-VAE.

Parameters:

• dim: The dimensionality of the input vectors.
• codebook_size: The number of vectors in the codebook.
• decay: The EMA decay (lower values make the dictionary change faster).
• commitment_weight: The weight applied to the commitment loss.
• rotation_trick: (Optional) Set to True to use the rotation trick for gradient computation; otherwise, it uses the straight-through estimator (STE).
• directional_reparam: (Optional) Set to True to use the DiVeQ method for gradient computation.
• directional_reparam_variance: (Optional) Variance for the DiVeQ method.

Returns (on call):

• quantized: The quantized vectors.
• indices: The indices of the chosen codebook vectors.
• commit_loss: The commitment loss.

import torch
from vector_quantize_pytorch import VectorQuantize

vq = VectorQuantize(
    dim = 256,
    codebook_size = 512,
    decay = 0.8,
    commitment_weight = 1.
)

x = torch.randn(1, 1024, 256)
quantized, indices, commit_loss = vq(x) # (1, 1024, 256), (1, 1024), (1)

LFQ eliminates the codebook and embedding lookup entirely using independent binary latents. It supports image, video, and sequence inputs.

• codebook_size: Must be a power of 2.
• dim: Input feature dimension (defaults to log2(codebook_size) if not defined).
• entropy_loss_weight: Weight for entropy loss.
• diversity_gamma: Weight for diversity of codes within entropy loss.
• num_codebooks: Supports multiple codebooks.
• ResidualLFQ: An improvised residual version for audio compression.

import torch
from vector_quantize_pytorch import LFQ, ResidualLFQ

# Standard LFQ
quantizer = LFQ(
    codebook_size = 65536,
    dim = 16,
    entropy_loss_weight = 0.1,
    diversity_gamma = 1.
)

image_feats = torch.randn(1, 16, 32, 32)
quantized, indices, entropy_aux_loss = quantizer(image_feats, inv_temperature=100.)

# Residual LFQ
residual_lfq = ResidualLFQ(
    dim = 256,
    codebook_size = 256,
    num_quantizers = 8
)