Residual-Quantized Variational Autoencoder

09-14-202609-14-2026 blog 9 minutes read (About 1396 words) visits

Introduction

In representation learning, capturing hierarchical semantic information is crucial for building effective models. An example of the semantic hierarchy to describe an item could be: “Electronics” $\rightarrow$ “Computers” $\rightarrow$ “Laptops” $\rightarrow$ “Gaming Laptops”. Each level of the hierarchy captures increasingly specific semantic information about the item.

In this blog post, I would like to quickly discuss how to use Residual-Quantized Variational Autoencoder (RQ-VAE) for capturing hierarchical semantic information into semantic indices with codebooks.

Vector-Quantized Variational Autoencoder

Vector Quantized Variational Autoencoder (VQ-VAE) is a generative model architecture introduced that pairs continuous feature extraction with discrete latent representations. Unlike a standard Variational Autoencoder (VAE), which maps inputs into a continuous, smooth Gaussian distribution, a VQ-VAE maps inputs to a discrete codebook of learned vector representations. Essentially, it performs tokenization of the features.

Suppose $\mathbf{z}_{e}(\mathbf{x}) \in \mathbb{R}^M$ is a latent variable of a sampled input $\mathbf{x}$, i.e., an output of the encoder for a sample input $\mathbf{x}$, and $\mathbf{e}_k \in \mathbb{R}^M$ represents the $\mathbf{k}$-th vector in the codebook of size $K$. The quantized latent representation $\mathbf{z}_q(\mathbf{x}) \in \mathbb{R}^M$ is obtained by finding the nearest codebook vector:

$$
\mathbf{z}_q(\mathbf{x}) = \mathbf{e}_{\mathbf{k}}
$$

where $\mathbf{k} = \arg\min_i \lVert \mathbf{z}_{e}(\mathbf{x}) - \mathbf{e}_i \rVert^2$.

This essentially means that a latent variable is just the index of the nearest codebook vector, which is just an integer $\mathbf{k}$. If there are multiple latent variables, multiple codebooks can exist and the sample input $\mathbf{x}$ can be encoded and represented using multiple integers, $\mathbf{k}_1, \mathbf{k}_2, \cdots, \mathbf{k}_n$, resulting in representation compression with the assistance from codebooks.

As we have learned from my previous article “Variational Autoencoder”, VAE is a generative model. It models a joint probability distribution $p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z}_{1}, \mathbf{z}_{2}, \cdots, \mathbf{z}_{n})$ of the observed data $\mathbf{x}$ and the latent variables $\mathbf{z}_{1}, \mathbf{z}_{2}, \cdots, \mathbf{z}_{n}$. Using chain rule, the joint probability distribution can be factorized as follows and the distribution $p_{\boldsymbol{\theta}}(\mathbf{z}_{1}, \mathbf{z}_{2}, \cdots, \mathbf{z}_{n})$ can be explicitly modeled in an autoregressive fashion.

$$
\begin{align}
p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z}_{1}, \mathbf{z}_{2}, \cdots, \mathbf{z}_{n})
&= p_{\boldsymbol{\theta}}(\mathbf{x} | \mathbf{z}_{1}, \mathbf{z}_{2}, \cdots, \mathbf{z}_{n}) p_{\boldsymbol{\theta}}(\mathbf{z}_{1}, \mathbf{z}_{2}, \cdots, \mathbf{z}_{n}) \\
&= p_{\boldsymbol{\theta}}(\mathbf{x} | \mathbf{z}_{1}, \mathbf{z}_{2}, \cdots, \mathbf{z}_{n}) p_{\boldsymbol{\theta}}(\mathbf{z}_{1}) p_{\boldsymbol{\theta}}(\mathbf{z}_{2} | \mathbf{z}_{1}) p_{\boldsymbol{\theta}}(\mathbf{z}_{3} | \mathbf{z}_{1}, \mathbf{z}_{2}) \cdots p_{\boldsymbol{\theta}}(\mathbf{z}_{n} | \mathbf{z}_{1}, \mathbf{z}_{2}, \cdots, \mathbf{z}_{n-1})
\end{align}
$$

The decoder models the conditional distribution $p_{\boldsymbol{\theta}}(\mathbf{x} | \mathbf{z}_{1}, \mathbf{z}_{2}, \cdots, \mathbf{z}_{n})$, which generates the observed data $\mathbf{x}$ given the latent variables.

In the case of VQ-VAE, the distribution $p_{\boldsymbol{\theta}}(\mathbf{z}_{1}, \mathbf{z}_{2}, \cdots, \mathbf{z}_{n})$ degenerates into a product of categorical distributions over the codebook indices:

$$
\begin{align}
p_{\boldsymbol{\theta}}(\mathbf{z}_{1}, \mathbf{z}_{2}, \cdots, \mathbf{z}_{n})
&= p_{\boldsymbol{\theta}}(\mathbf{k}_{1}, \mathbf{k}_{2}, \cdots, \mathbf{k}_{n}) \\
&= p_{\boldsymbol{\theta}}(\mathbf{k}_{1}) p_{\boldsymbol{\theta}}(\mathbf{k}_{2} | \mathbf{k}_{1}) p_{\boldsymbol{\theta}}(\mathbf{k}_{3} | \mathbf{k}_{1}, \mathbf{k}_{2}) \cdots p_{\boldsymbol{\theta}}(\mathbf{k}_{n} | \mathbf{k}_{1}, \mathbf{k}_{2}, \cdots, \mathbf{k}_{n-1})
\end{align}
$$

Taken together, we have learned how to encode a sample input $\mathbf{x}$ into a sequence of discrete latent variables $\mathbf{z}_{1}, \mathbf{z}_{2}, \cdots, \mathbf{z}_{n}$, how can the latent distribution be modeled and sampled autoregressively, and how the decoder can generate the observed data given these latent variables.

Residual-Quantized Variational Autoencoder

As we have discussed in my previous article “Vector Quantization and Product Quantization”, vector quantization has the problem of high quantization error when the codebook size is limited and the embedding semantics are rich, especially in high-dimensional spaces. As a consequence of this, many different input samples might be encoded into exactly the same codebook entry, resulting in significant information loss and reduced representation capacity. Using extremely large codebook makes the model size much larger, makes the model training unstable due to model collapse, and makes inference solutions sometimes infeasible due to the large model size.

To address this problem, based on VQ-VAE, Residual Quantized Variational Autoencoder (RQ-VAE) introduces a sequence of residual quantizers, where each quantizer encodes the residual error left by the previous quantizers. Given a latent variable $\mathbf{z}_{e}(\mathbf{x})$, the first quantizer encodes $\mathbf{z}_{e}(\mathbf{x})$ into a discrete code $\mathbf{k}_1$, and the residual $\mathbf{r}_1 = \mathbf{z}_{e}(\mathbf{x}) - \mathbf{e}_{\mathbf{k}_1}$ is passed to the next quantizer, which encodes it into $\mathbf{k}_2$, and so on. Formally, the quantization process can be described as:

$$
\begin{align}
\mathbf{r}_0 = \mathbf{z}_{e}(\mathbf{x}), \quad \mathbf{k}_i = \arg\min_{\mathbf{k}} \lVert \mathbf{r}_{i-1} - \mathbf{e}_{\mathbf{k}} \rVert^2, \quad \mathbf{r}_i = \mathbf{r}_{i-1} - \mathbf{e}_{\mathbf{k}_i}, \quad i = 1, 2, \ldots, D
\end{align}
$$

where $\mathbf{e}_{\mathbf{k}}$ denotes the embedding vector corresponding to code $\mathbf{k}$ in the codebook, and $D$ is the residual quantization depth.

The dequantization process in RQ-VAE involves reconstructing the latent variable $\mathbf{z}_{e}(\mathbf{x})$ from the sequence of discrete codes $\mathbf{k}_1, \mathbf{k}_2, \cdots, \mathbf{k}_D$. This is done by summing the corresponding embedding vectors from each quantizer:

$$
\begin{align}
\mathbf{z}_{q}(\mathbf{x}) = \sum_{i=1}^{D} \mathbf{e}_{\mathbf{k}_i}.
\end{align}
$$

RQ-VAE allows the model to achieve higher representation capacity without requiring an excessively large codebook, thereby reducing quantization error while maintaining manageable model size and stable training. For example, with RQ-VAE, if we have residual quantization depth of $D$, each level has a codebook of size $K$, the effective representation capacity becomes $K^D$, while the actual model size only grows linearly with the number of levels, i.e., $DK$. With VQ-VAE, to well represent $K^D$ different representations, the codebook size would need to go exponentially large. The additional cost to inference of RQ-VAE is that the latent representation becomes a sequence of $D$ discrete variables, instead of one single discrete variable as in standard VQ-VAE.

RQ-VAE introduces representation hierarchy through residual quantization, where each level captures the residual information not represented by the previous levels. This is somewhat natural in many semantic representation tasks. For example, to classify an animal, the first level might capture coarse features such as shape and size, while subsequent levels capture finer details such as fur texture, color patterns, and facial features. Consequently, RQ-VAE can effectively learn hierarchical representations that progressively refine the captured information, leading to more accurate and expressive latent codes. It is widely used in recommendation systems that encode items into semantic indices and embeddings.

Miscellaneous

Residual Quantization VS Product Quantization

Residual Quantization (RQ) and Product Quantization (PQ) are both techniques used to quantize high-dimensional vectors into compact representations, but they differ in their approach:

Residual Quantization (RQ): RQ decomposes the quantization process into multiple stages. Each stage quantizes the residual error left by the previous stages. This allows RQ to achieve high representation capacity with relatively small codebooks at each stage.
Product Quantization (PQ): PQ splits the original vector into multiple sub-vectors and quantizes each sub-vector independently using separate codebooks. The final representation is a concatenation of the quantized sub-vectors.

Why there is no Product-Quantized VAE (PQ-VAE) equivalent? Because PQ-VAE is mathematically equivalent to having a VQ-VAE with more latent variables, each having a separate codebook. From our previous analysis, we already know VQ-VAE is less expressive than RQ-VAE in terms of representation capacity.