Kernel Quantization for Efficient Network Compression

The pipeline of kernel-level quantization is shown in Figure 2. For a convolution layer with weight tensor $\mathcal{W}\in\mathbb{R}^{\omega\times\omega\times p\times q}$, where $\omega$ denotes the kernel size, $p$ and $q$ are input and output channel, respectively. We reshape it as a matrix $\bm{W}=[\bm{m}_{0},\bm{m}_{1},\cdots,\bm{m}_{n-1}]\in\mathbb{R}^{\omega^{2}\times n}$, where $n=pq$ denotes the total number of kernels in $\mathcal{W}$, $\bm{m}_{i}\in\mathbb{R}^{\omega^{2}}$ denotes the $i$-th kernel in $\mathcal{W}$. We generate a $k$ entry kernel codebook $\bm{C}=\{\bm{c}_{0},\bm{c}_{1},\cdots,\bm{c}_{k-1}\}$. The kernel set corresponding to entry $\bm{c}_{i}$ is $\bm{S}_{i}=\{\bm{m}_{0}^{i},\bm{m}_{1}^{i},\cdots,\bm{m}_{z-1}^{i}\}$ where $z$ is the number of kernels assigned to $\bm{c}_{i}$. We adopt k-means to minimize the following distance:

All kernels in $\mathcal{W}$ are mapped to corresponding entries in the kernel codebook. We represent each kernel with an index to the corresponding entry. Kernel-level quantization is performed in a layer by layer manner. During retraining, back propagation on codebook is done in a scheme like conventional quantization. We first calculate the gradient of each parameter in $\mathcal{W}$. Then, we calculate the elementwise average of gradients that are mapped to the same entry. The average gradients are used to update the kernel codebook values as follows,

where $L$ is the network loss, $\bm{c}^{t}_{i}$ and $\bm{c}^{t+1}_{i}$ are values of entry $\bm{c}_{i}$ in the codebook after updating for $t$ and $t+1$ times, $\gamma$ is the learning rate.

In conventional quantization methods, only the bit length $b$ needs to be carefully selected. Given $b$, the codebook size is usually set as $2^{b}$, because the size of the codebook is relatively small. But for kernel-level quantization, the codebook size is significantly larger. Carefully selecting the codebook size to balance accuracy and storage saves plenty of space. Thus, it is important to find the appropriate codebook size with adequate time complexity.

A naive approach is setting the codebook size to different orders of 2 in descending order until it reaches a certain threshold accuracy. This method makes full use of index bits while failed to select the appropriate size of codebook precisely. Fitting the codebook size to orders of 2 may either waste space to store redundant entries or lower the accuracy because of insufficient codebook size. Therefore, we propose a binary search approach searching for appropriate codebook size. We notice that network accuracy increases along with codebook size increasing. Thus, we set up a tolerable maximal accuracy drop for kernel-level quantization on each layer. Then we use a binary search to find the best codebook size that close to the target accuracy. In this way, simply given the parameters, the network adaptively finds the suitable codebook size and compression ratio. Thus, we precisely control the size of the codebook and the trade-off between the codebook size and network accuracy.

We define the initial entry ratio as $\alpha$ and the threshold ratio for target accuracy $r$. We test the reference accuracy $A_{ori}$ without quantization on the current layer. The initial codebook size is $N_{init}=\alpha\times n$, and the corresponding baseline quantized accuracy is $A_{base}$. We define the target accuracy $A_{target}$ as

Then we search for the codebook size with binary search to reach the target accuracy $A_{target}$. The upper bound of codebook size $B_{upper}$ is initialized as $N_{init}$ and lower bound $B_{lower}$ is initialized to zero. The testing codebook size is $N_{curr}=0.5\times(B_{upper}+B_{lower})$. Then kernel-level quantization generates a codebook with $N_{curr}$ entries and quantizes the layer to test the validation accuracy. If the validation accuracy is higher than $A_{target}$, we set $B_{upper}=N_{curr}$, else $B_{lower}=N_{curr}$. These steps are repeated until reaching the target accuracy $A_{target}$ or the maximum number of iteration. After finished binary search and quantized current convolution layer, the whole CNN is retrained for one epoch to finetune the model before quantizing next layer. The overall pseudo-code for the kernel quantization step is shown in Algorithm 1.

We further compress the kernel codebook after quantized kernels on all convolution layers. The model storage after kernel-level quantization includes two parts, codebook and indexes. Total number of bits $B$ needed to store the quantized convolution layer is

where $b_{c}$ is the bit length for each parameter in the kernel codebook. When $k$ is large, storing entries in the codebook consumes massive space. So quantizing the kernel codebook provides us with additional compression ratio. We adopt a simple yet effective 6-bit parameter quantization method for codebook quantization.

We use the layer by layer strategy to preserve performance. First, we do k-means clustering on all parameters in the codebook. A small trick is that since different entries appear for different times in $\mathcal{W}$, the more times an entry appears, the more important it is. So we use the entry appearance time as the weight of the kernel and perform weighted k-means. This helps the algorithm to pay more attention to preserving important parameters. Retraining of codebook quantization follows the same scheme as kernel quantization. We update the parameters with the average gradients from different kernels mapped to the same entry. Retraining is conducted after quantizing codebook for every two layers and iterate for one epoch. After quantizing all codebooks of the network, we further run a 6-bit quantization on the fully-connected layer with the same layer by layer k-means clustering method.

It is worth to notice that KQ does not add much extra burden on testing. In the testing phase, the procedure is the same as conventional quantization methods with extra mapping. KQ first recovers the kernel codebook, then recovers the weight tensor $\mathcal{W}$ from kernel codebook.

Conventional quantization methods use an index to map each parameter to an entry in the codebook. When the length of the index is shortened by reducing the size of the codebook, the total storage is reduced. The compression ratio of conventional quantization (denoted as $C_{con}$) is

where $b_{FP}$ is the bit length of a full precision parameter, which is 32-bit in most cases, $u$ is the size of the codebook, $b_{1}$ is the length of each entry in the codebook. We obtain the theoretical limit when there are only two entries in the codebook and each parameter is represented with 1-bit. The theoretical limit of conventional quantization is 32.

KQ bypasses this limit by mapping each kernel, instead of each parameter, to an entry in the codebook. Thus, KQ uses one index to represent nine parameters (for a $3\times 3$ kernel). The compression ratio of KQ (denoted as $C_{KQ}$) is

When the limit is approached, $k$ is 2. The equation turns to be $C_{KQ}=\frac{\omega^{2}\times b_{FP}}{n\times\text{ceil}(\log_{2}k)}$ ($k$ is too small compared to $n$), the theoretical compression ratio is $\frac{n\times 3\times 3\times 32}{n\times log_{2}2}=288$.

We evaluate KQ on the image classification task. All of the images are resized to $256\times 256$. The images are then randomly cropped to $224\times 224$ patches with mean subtraction and random flipping without any data augmentation. We report the top-1 accuracy on the validation set under single-center-crop testing.

We implement KQ on PyTorch [17] platform and the referenced full precision CNN models are from the torchvision package. During retraining, we use SGD optimizer with learning rate 0.001, momentum 0.9. In the experiments, we only conduct kernel quantization on kernels of size $3\times 3$. We adopt the Yinyang k-means [5] to deal with massive samples and cluster centers. Yinyang k-means has exactly the same results compared with conventional k-means but provides a significant boost in clustering speed.

For the sake of narrative, we always name the first $3\times 3$ convolution layer as conv1, the second $3\times 3$ convolution layer as conv2 and so on. When computing the compression ratio, we count all convolution layers regardless of the kernel size. In the rest of this section, we use ”K” for Kernel-Level Quantization, “C” for Codebook Quantization, and ”K+C” for apply Codebook Quantization after Kernel-Level Quantization.

To better compare KQ with other quantization methods, we use the average number of bits per parameter (denoted as $\beta$) as the measurement. It is defined as the total number of bits needed to represent the convolution weight divided by the total number of parameters in the convolution weight:

It is the reciprocal of compression ratio in Equation 6 multiplied by 32 (full precision bits).

VGG is widely used in a variety of computer vision tasks. It has 13 convolution layers and 3 fully connected layers connected in a sequential manner. All internal convolution layers (the first $7\times 7$ layer excluded) in VGG use $3\times 3$ convolution kernel. So we apply KQ on all $3\times 3$ convolution layers.

As shown in Table 3, we compare our method with some of the state-of-the-art methods. Kernel-level quantization itself is able to achieve on average 1.19 bits per parameter, which is about 26.9 times compression. Comparing with the methods which use approximately 2 bits per parameter, we outperform all the methods with fewer bits. As for LQ-Net [28] and SYQ [6] with 1 bit, KQ uses similar bits while has much better performance than both methods. After further applying Codebook Quantization, our method is able to get 30.5 times compression with only another 0.2% accuracy drop. Overall speaking, KQ achieves a great balance between compression and accuracy, and outperforms all other methods.

To better understand how KQ works on different layers, we list the layerwise compression ratio in Table 4. An obvious observation is deeper layers has smaller kernel codebook than shallower layers. This observation, on the one hand, proves there is more redundancy in the deeper layers. On the other hand, it shows the shallower layers of VGG are more important, and network compression tasks exploit major compression ratio gain in deeper layers. Since the max iteration is set to 8, $\alpha=0.5$, and the last five layers all have $512\times 512$ kernels, the kernel codebook size of 512 presents that it is halved from $512\times 256$ consecutively for eight times and the validation accuracy is still above the target accuracy. This implies with more iterations in binary search,KQ has potential to boost the performance further.

We further analysis the distribution of index after the kernel quantization. As shown in Figure 4, we count the appearances of each entry in the codebook on the first 4 layers of VGG. It is clear that the distributions of the statistics are non-uniform, which makes it possible to further losslessly compress the network with coding methods like Huffman coding in [9].

We also evaluate KQ on ResNet18 architecture. Unlike the VGG, ResNet18 has batch normalization layers and skip connections directly connecting a deeper layer with a shallower layer to prevent the gradient vanishing. For ResNet18 network, we set $\alpha=0.3$, $r=0.5$, and the maximum iteration is 8.

We compare KQ with some of the state-of-the-art methods and show the results in Table 6. Compared to INQ [29], although with only kernel-level quantization, INQ has less accuracy loss with the same average number of bits, but after applying the Codebook Quantization and further lower the average bit length, KQ significantly outperforms the INQ with 0.38 fewer average number of bits and $1.3\%$ less accuracy loss. And KQ has better performance with fewer bits comparing to other 2-bit quantization methods. Under 1-bit setting, KQ achieves significant improvement with a little more bits.

In Table 5, we report the layerwise compression ratio of KQ on ResNet18. Comparing with VGG, we need larger codebook size on most of the layers. The most obvious difference is deeper layers of ResNet18 need larger codebook size than shallower layers. This is an aspect demonstrating the compactness of ResNet18. Meanwhile, KQ achieves 1.62 bits per parameter without noticeable accuracy loss on such a compact network, proving the effectiveness of KQ.

As the increasing of codebook size, we need more bits to represent each index. However, in KQ, each index represents nine parameters. For 1-bit increase of index bit length, the average number of bits per parameter only increase by $\frac{1}{9}$ but the codebook size is doubled. With this merit, KQ achieves an extremely low average number of bits per parameter on most large networks.

To further demonstrate the superiority of KQ, we evaluate KQ with the state-of-the-art structural quantization method deep k-means [13] and structural compression method GreBdec [27] using GoogLeNet [23] on ImageNet dataset. We set $\alpha=0.25$ and $r=1.0$ for KQ on GoogLeNet. Table 7 shows the compression ratio of the above methods on convolution layers. Our method outperforms both of them with higher accuracy and compression ratio. This further proves the high efficiency of the proposed KQ algorithm. Compared with deep k-means, we credit the superiority in performance to KQ’s ability in preserving the tendency of the kernel, while [13] proposed to use each row of the kernel as quantization unit failed to do so. As one of key functions of the convolution kernel is that it can extract the texture feature as a filter from the signal. The tendency of parameters in the convolution kernel is critical in preserving this ability. Given a row $[2,3,2]$ in kernel, it has the same distance to $[3,2,3]$ and $[1,2,1]$, but quantizing $[2,3,2]$ into $[1,2,1]$ is obviously better in preserving performance,  [13] failed in such case while our KQ considers the kernel as a whole and better deals with these kinds of cases. Moreover, in [13], each index represents three parameters while in KQ, each parameter represents nine parameters (for a kernel with size $3\times 3$). This provides a higher theoretical compression ratio for KQ.