LooC: Effective Low-Dimensional Codebook for Compositional Vector Quantization

TL;DR

LooC introduces a low-dimensional, shared codebook for compositional vector quantization, enabling a large effective representation space with a compact parameter budget. By treating codevectors as compositional units within feature vectors and applying a parameter-free extrapolation-by-interpolation step, LooC achieves 100% codevector usage and state-of-the-art performance on reconstruction and generation tasks across multiple datasets, while using far smaller codebooks than prior methods. The approach is plug-and-play for existing VQ-based systems and demonstrates strong generalization and efficiency benefits, with extensive ablations confirming the advantages of increased compositional granularity (larger m) and the extrapolation mechanism. Overall, LooC offers a practical, scalable pathway to high-capacity VQ in vision tasks, reducing storage and computation without sacrificing quality.

Abstract

Vector quantization (VQ) is a prevalent and fundamental technique that discretizes continuous feature vectors by approximating them using a codebook. As the diversity and complexity of data and models continue to increase, there is an urgent need for high-capacity, yet more compact VQ methods. This paper aims to reconcile this conflict by presenting a new approach called LooC, which utilizes an effective Low-dimensional codebook for Compositional vector quantization. Firstly, LooC introduces a parameter-efficient codebook by reframing the relationship between codevectors and feature vectors, significantly expanding its solution space. Instead of individually matching codevectors with feature vectors, LooC treats them as lower-dimensional compositional units within feature vectors and combines them, resulting in a more compact codebook with improved performance. Secondly, LooC incorporates a parameter-free extrapolation-by-interpolation mechanism to enhance and smooth features during the VQ process, which allows for better preservation of details and fidelity in feature approximation. The design of LooC leads to full codebook usage, effectively utilizing the compact codebook while avoiding the problem of collapse. Thirdly, LooC can serve as a plug-and-play module for existing methods for different downstream tasks based on VQ. Finally, extensive evaluations on different tasks, datasets, and architectures demonstrate that LooC outperforms existing VQ methods, achieving state-of-the-art performance with a significantly smaller codebook.

Figures (9)

• Figure 1: Codebook size and reconstruction performance.Left: Typical configurations (green dots) of codevector number $K$ and dimension $d^*$ in a codebook. LooC (red dot) stands out with a significantly smaller codebook size of $32\times 4$. Right: Reconstruction results on CIFAR10 krizhevsky2009cifar10. LooC performs significantly better with a much smaller codebook than other SOTA methods.
• Figure 2: Framework of Low-dimensional codebook for Compositional vector quantization (LooC). The encoder transforms the input image into a continuous latent feature map $z$. $z$ is then upsampled using bilinear interpolation with scale factor $\beta$. Simultaneously, each feature vector in $z$ is divided into $m$ units and quantized using a shared codebook $\mathcal{C}$ containing $K$ codevectors of dimension $d^* = d/m$. The quantized units are then reassembled and smoothed using average pooling to restore the shape as $z$. Finally, the decoder converts the feature map back to the image.
• Figure 3: Qualitative results. Reconstructed images using VQGAN esser2021taming_vq_gan, CVQ zheng2023online_CVQ, and LooC. LooC significantly enhances reconstruction quality by preserving image details and restoring texture structures, as highlighted in the red boxes (best viewed in PDF with zoom).
• Figure 4: Unconditional image generation on LSUN yu15lsun and class-conditional image generation on Imagenet deng2009imagenet.
• Figure 5: Codebook visualization with t-SNE for models trained on CIFAR10 and evaluated on the validation set. VQ-VAE has unused codevectors (green points) with only $24.12\%$ useage. LooC achieves $100\%$ usage at both $K=256$ and $K=1024$.