Welcome to diveq

diveq (short for differentiable vector quantization) is a tool for implementing and training vector quantization (VQ) in deep neural networks (DNNs), such as a VQ-VAE. It allows end-to-end training of DNNs that contain the non-differentiable VQ module, without any auxiliary losses and hyperparameter tunings. diveq is implemented via PyTorch, and it requires python >= 3.11 and torch >= 2.0.0 .

The diveq method is published as a research paper entitled "DiVeQ: Differentiable Vector Quantization Using the Reparameterization Trick" in the International Conference on Learning Representations (ICLR) in 2026. You can find the original GitHub repository of the paper at https://github.com/AaltoML/DiVeQ.

The diveq package includes eight different vector quantization (VQ) methods:

1. from diveq import DIVEQ optimizes the VQ codebook via DiVeQ technique. DiVeQ is the first proposed method in the paper that works as an ordinary VQ by mapping the input to codebook vectors.
2. from diveq import SFDIVEQ optimizes the VQ codebook via Space-Filling DiVeQ technique. SF-DiVeQ is the second proposed method in the paper, different from ordinary VQ in a way that it maps the input to a space-filling curve constructed from codebook vectors.

VQ variants that use multiple codebooks for vector quantization, i.e., Residual VQ and Product VQ:

3. from diveq import ResidualDIVEQ optimizes the Residual VQ codebooks via DiVeQ technique.
4. from diveq import ResidualSFDIVEQ optimizes the Residual VQ codebooks via SF-DiVeQ technique.
5. from diveq import ProductDIVEQ optimizes the Product VQ codebooks via DiVeQ technique.
6. from diveq import ProductSFDIVEQ optimizes the Product VQ codebooks via SF-DiVeQ technique.

Variants of DiVeQ and SF-DiVeQ techniques that use deterministic quantization instead of stochastic quantization:

7. from diveq import DIVEQDetach optimizes the VQ codebook via DiVeQ_Detach technique.
8. from diveq import SFDIVEQDetach optimizes the VQ codebook via SF-DiVeQ_Detach technique.

For more details on these eight different VQ methods, please see the paper.

Installation

You can install diveq through pip by running:

pip install diveq

After installing diveq, you can verify the installation and package details by running:

python -m pip show diveq

Usage Example

Before using diveq, you have to install it using pip install diveq.

Below you see a minimal example of how to import and use the DIVEQ optimization method as a vector quantizer in a model.

from diveq import DIVEQ
vector_quantizer = DIVEQ(num_embeddings=512, # codebook size
                         embedding_dim=256   # embedding dimension
                         )

z = torch.randn(100, 256)

# when training
z_quantized, indices, perplexity = vector_quantizer(z) # shapes: (100, 256), (100, ), (1)

# at inference
z_q_hard, indices, perplexity = vector_quantizer.inference(z) # shapes: (100, 256), (100, ), (1)

• vector_quantizer is the vector quantization module that will be used for building the model.

• num_embeddings and embedding_dim are the codebook size and dimension of each codebook entry, respectively. In the following, you can find the list of all input parameters in different vector quantization modules incorporated in diveq package.

• z_quantized is the differentiable quantized version of input z used when training

• z_q_hard is the non-differentiable (hard) quantized version of input z used at inference

• indices is the selected codebook indices

• perplexity is the codebook perplexity/usage (entropy of selected indices)

In the example directory of the GitHub for diveq package, we provide a code example of how vector quantization modules in diveq can be used in a vector quantized variational autoencoder (VQ-VAE). You can create the required environment to run the code by running:

cd example  #change directory to the example folder
conda create --name diveq_example python=3.11
conda activate diveq_example
pip install -r requirements.txt

Then, you can train the VQ-VAE model by running:

python train.py

Tensor shape requirements

All quantization methods in diveq package take input with the shape of (N, D), where N is the number of latent vectors and D is the embedding dimension. Usually for image data, the input of vector quantizer has the shape of (B, C, H, W), where B: batch size, C: input channels, H and W: Height and Width. In this case, to use diveq quantization methods, you need permutation and reshaping of the latent z:

# extract latent z from input image x
z = self.encoder(x)

# first do the permutation and reshaping
z_permuted = z.permute(0, 2, 3, 1) # Convert BCHW -> BHWC
z_shape = z_permuted.shape
z_flattened = z_permuted.view(-1, embedding_dim)  # Flatten the latents

# apply vector quantization using any method in diveq package
z_quantized, indices, perplexity = self.vector_quantizer(z)

# apply permutation and reshaping back
z_quantized = (z_quantized.view(z_shape)).permute(0, 3, 1, 2) # Convert BHWC -> BCHW

# reconstruct the image from quantized latents
x_recon = self.decoder(z_quantized)

List of Parameters

Here, we provide the list of parameters that are used as inputs to eight different vector quantization methods included in diveq package.

• num_embeddings (integer): Codebook size or the number of codewords in the codebook.
• embedding_dim (integer): Dimensionality of each codebook entry or codeword.
• noise_var (float): Variance of the directional noise for stochastic DiVeQ- and SF-DiVeQ-based methods.
• replacement_iters (integer): Number of training iterations to apply codebook replacement.
• discard_threshold (float): Threshold to discard the codebook entries that are used less than this threshold after replacement_iters iterations.
• perturb_eps (float): Adjusts perturbation/shift magnitude from used codewords for codebook replacement.
• uniform_init (bool): Whether to use uniform initialization. If False, codebook is initialized from a normal distribution.
• verbose (bool): Whether to print codebook replacement status, i.e., to print how many unused codewords are replaced.
• codebook_reset (bool): Whether to apply codebook replacement to avoid codebook collapse.
• skip_iters (integer): Number of training iterations to skip quantization (for SF-DiVeQ and SF-DiVeQ_Detach) or to use DiVeQ quantization (for Residual_SF-DiVeQ and Product_SF-DiVeQ) in the custom initialization.
• avg_iters (integer): Number of recent training iterations to extract latents for custom codebook initialization in Space-Filling Versions.
• latents_on_cpu (bool): Whether to collect latents for custom initialization on CPU. If running out of CUDA memory, set it to True.
• allow_warning (bool): Whether to print the warnings. The warnings will warn if the user inserts unusual values for the parameters.
• num_codebooks (integer): Number of codebooks to be used for quantization in VQ variants of Residual VQ and Product VQ. All the codebooks will have the same size and dimensionality.

Important Notes about Parameters

1. Codebook Replacement: Note that to prevent codebook collapse, we include a codebook replacement function (in cases where it is required) inside different quantization modules. Codebook replacement function is called after each replacement_iters training iterations, and it replaces the codewords which are used less than discard_threshold with perturbation of actively used codewords which are shifted by perturb_eps magnitude. If verbose=True, the status of how many unused codewords are replaced will be printed by the module. Note that the number of unused codewords should decrease over training, and it might take a while.

2. Variants of Vector Quantization: Residual VQ and Product VQ are two variants of vector quantization, which are included in the diveq package. These variants utilize multiple codebooks for quantization, where num_codebooks determines the number of codebooks used in these VQ variants.

3. Space-Filling Methods: Quantization methods based on Space-Filling (i.e., SF-DiVeQ, SF-DiVeQ_Detach, Residual_SF-DiVeQ, Product_SF-DiVeQ) use a custom initialization. SF-DiVeQ and SF-DiVeQ_Detach skip quantizing the latents for skip_iters training iterations, and initialize the codebook with an average of latents captured from avg_iters recent training iterations. After this custom initialization, they start to quantize the latents. Residual_SF-DiVeQ and Product_SF-DiVeQ work in the same way, but they apply DiVeQ for the first skip_iters training iterations. Note that if avg_iters value is set to a large value, CUDA might run out of memory, as there should be a large pool of latents to be stored for custom initialization. Therefore, the user can set latents_on_cpu=True to store the latents on CPU, or set a smaller value for avg_iters.

4. Detach Methods: DiVeQ_Detach and SF-DiVeQ_Detach methods do not use directional noise. Therefore, they do not need to set the noise_var parameter.

For further details about different vector quantization methods in the diveq package and their corresponding parameters, please see the details provided in the Python codes in src directory of the diveq package.

Citation

If this package contributed to your work, please consider citing it:

@InProceedings{vali2026diveq,
 title={{DiVeQ}: {D}ifferentiable Vector Quantization Using the Reparameterization Trick},
 author={Vali, Mohammad Hassan and Bäckström, Tom and Solin, Arno},
 booktitle={International Conference on Learning Representations (ICLR)},
 year={2026}
}

License

diveq was developed by Mohammad Hassan Vali, part of the AaltoML research group from Aalto University and is licensed under MIT license.