Vector Quantization and Product Quantization

Introduction

Quantization is a useful technique for high-dimensional data compression. The conventional straightforward scaling based quantization techniques typically used for neural network inference requires special low bitwidth data types and accelerator hardware support. In many extreme low bitwidth quantization scenarios, the accuracy of the quantized model using the conventional approach can be significantly degraded, usually caused by data outliers. There are other quantization techniques that are not based on scaling, such as vector quantization and product quantization.

In this blog post, I would like to quickly discuss vector quantization and product quantization, which uses codebook and indices to represent vectors in a finite set.

Vector Quantization

Given a high-dimensional vector space $\mathcal{X} \in \mathbb{R}^D$, vector quantization is a technique that maps these vectors to a finite set of representative vectors, often referred to as the codebook, $\mathcal{C} = [\mathbf{c}_1, \mathbf{c}_2, \ldots, \mathbf{c}_M]$, where the codebook can be trained from representative data from $\mathcal{X}$ using algorithms such as K-Means and codebook indices are usually represented by $k$ bits, i.e., $k = \lceil \log_2(M) \rceil$. Each vector $\mathbf{x} \in \mathcal{X}$ is usually assigned to the nearest codebook vector $\mathbf{c}_j$ based on a distance metric, such as the Euclidean distance. Consequently, with a codebook of size $M$, each vector can be represented by an index $j$ that requires only $k$ bits, significantly reducing the storage requirements.

In online computing, given a new vector $\mathbf{x} \in \mathcal{X}$, during the encoding process, its quantized value $\mathbf{x}_q$ is just to find the nearest codebook vector $\mathbf{c}_j$ and return its index $j$. The dequantization decoding process is to retrieve the codebook vector $\mathbf{c}_j$ corresponding to the index $j$. During the decodingp process, its dequantized value $\mathbf{x}^{\prime}$ is just to retrieve the codebook vector $\mathbf{c}_j$ corresponding to the index $j$. The quantization and dequantization processes can be mathematically expressed as follows:

$$
\begin{align}
\mathbf{x}_q &= j = \arg\min_{j} \Vert \mathbf{x} - \mathbf{c}_j \Vert^2 \\
\mathbf{x}^{\prime} &= \mathbf{c}_j \\
\end{align}
$$

Suppose the original vectors and the codebook vectors are in a $D$-dimensional space, each element of the original vector $\mathbf{x} \in \mathcal{X}$ is represented by $b$ bits. The total storage required for the original set of vectors is $N \cdot D \cdot b$ bits. After vector quantization, the storage required for the codebook is $M \cdot D \cdot b$ bits, and the storage required for the indices is $N \cdot k$ bits.

In online computing, because the transfer of codebook vectors is usually done ahead of time, what matters is the compression ratio of the original vectors to the quantized indices. The compression ratio $\text{CR}$ can be expressed as:

$$
\begin{align}
\text{CR} &= \frac{D \cdot b}{k} \\
\end{align}
$$

Usually $b > k$, because the value of $k$ and $M$ are usually constrained by computer resources, such as memory storage and encoding time. So it is very common to have a high compression ratio using vector quantization. For example, if $D = 32$, $b = 32$, and $M = 65536$, then $k = \lceil \log_2(65536) \rceil = 16$, and the compression ratio is $\text{CR} = \frac{32 \cdot 32}{16} = 64$. This level of compression is often theoretically not achievable with other compression techniques normally used in neural network quantizations. One of the trade-offs is that the storage of codebook vectors and the encoding process can be expensive. Consequently, if the vector data transfer in online computing is a bottleneck, vector quantization can be a very effective technique to reduce the data transfer size.

The accuracy of vector quantization is related to the ratio of $\frac{M}{D}$ and it can be controlled by the size of the codebook $M$. For a smaller dimensional vector space, a smaller codebook can be representative enough to provide good accuracy. For a larger dimensional vector space, it would require an extremely large codebook to provide good accuracy. Consequently, controlling the accuracy only using the value of $M$ is often not practical. For example, if $D$ is large and having the maximum $M$ that the platform can allow might still not be enough to provide good accuracy, there will be no further way to improve.

A natural and intuitive solution is to split the original vector into multiple sub-vectors, and then perform vector quantization on each sub-vector independently, resulting in multiple indices corresponding to multiple codebooks. This would sacrifice compression ratio because multiple indices are used for representing the original vector, but it provides an additional degree of freedom to control the accuracy by controlling the number of sub-vectors. This leads to the technique of product quantization.

Product Quantization

A high-dimensional vector space $\mathcal{X} \in \mathbb{R}^D$ is essentially a Cartesian product of $N = \frac{D}{d}$ sub-vector spaces $\mathcal{X}_i \in \mathbb{R}^d$, where $d$ is the dimension of each sub-vector space. Product quantization splits the original vector $\mathbf{x} \in \mathcal{X}$ into $N$ sub-vectors $\mathbf{x}_i \in \mathcal{X}_i$, and then performs vector quantization on each sub-vector independently, resulting in $N$ indices corresponding to $N$ codebooks $[\mathcal{C}_1, \mathcal{C}_2, \ldots, \mathcal{C}_N]$, where $\mathcal{C}_{i} = [\mathbf{c}_{i,1}, \mathbf{c}_{i,2}, \ldots, \mathbf{c}_{i,M}]$. The quantization and dequantization processes can be mathematically expressed as follows:

$$
\begin{align}
\mathbf{x}_{i,q} &= j_i = \arg\min_{j_i} \Vert \mathbf{x}_i - \mathbf{c}_{i,j_i} \Vert^2 \\
\mathbf{x}_i^{\prime} &= \mathbf{c}_{i,j_i} \\
\end{align}
$$

where $\mathbf{x} = [\mathbf{x}_1, \mathbf{x}_2, \ldots, \mathbf{x}_N]$ and $\mathbf{x}^{\prime} = [\mathbf{x}_1^{\prime}, \mathbf{x}_2^{\prime}, \ldots, \mathbf{x}_N^{\prime}]$. The total storage required for the codebooks is $N \cdot M \cdot d \cdot b = \frac{D}{d} \cdot M \cdot d \cdot b = M \cdot D \cdot b$ bits, which is the same as the storage required for a single codebook, and the total storage required for the indices is $N \cdot k$ bits, which is $N$ times larger than the storage required for a single index. Consequently, the compression ratio $\text{CR}$ can be expressed as:

$$
\begin{align}
\text{CR} &= \frac{D \cdot b}{N \cdot k} = \frac{D \cdot b}{\frac{D}{d} \cdot k} = \frac{d \cdot b}{k} \\
\end{align}
$$

The product quantization is a way of factorization that compresses information based on the structure of the data. Let’s assume we use $N$ codebooks and each subvector could be represented using $M$ codebook vectors without any loss of information. If we use a single codebook to represent the original vector, it would require $M^N$ codebook vectors to represent all the original vectors without any loss of information, which is usually not practical.

References

• Vector Quantization - Wikipedia
• Quantization for Neural Networks