Our new X account is live! Follow @wizwand_team for updates
WorkDL logo mark

Addressing Representation Collapse in Vector Quantized Models with One Linear Layer

About

Vector Quantization (VQ) is essential for discretizing continuous representations in unsupervised learning but suffers from representation collapse, causing low codebook utilization and limiting scalability. Existing solutions often rely on complex optimizations or reduce latent dimensionality, which compromises model capacity and fails to fully solve the problem. We identify the root cause as disjoint codebook optimization, where only a few code vectors are updated via gradient descent. To fix this, we propose \textbf{Sim}ple\textbf{VQ}, which reparameterizes code vectors through a learnable linear transformation layer over a latent basis, optimizing the \textit{entire linear space} rather than nearest \textit{individual code vectors}. Although the multiplication of two linear matrices is equivalent to applying a single linear layer, this simple approach effectively prevents collapse. Extensive experiments on image and audio tasks demonstrate that SimVQ improves codebook usage, is easy to implement, and generalizes well across modalities and architectures. The code is available at https://github.com/youngsheen/SimVQ.

Yongxin Zhu, Bocheng Li, Yifei Xin, Zhihua Xia, Linli Xu• 2024

Related benchmarks

TaskDatasetResultRank
Image ReconstructionImageNet
PSNR25.3304
43
Image ReconstructionCOCO (test)
CVU0.9429
24
Audio ReconstructionCommon Voice
CVU0.0041
21
Audio ReconstructionLibriSpeech (test-clean test-other)
CVU0.004
21
Showing 4 of 4 rows

Other info

Follow for update