Addressing Representation Collapse in Vector Quantized Models with One Linear Layer

TL;DR

This paper introduces SimVQ, a simple linear transformation approach to prevent representation collapse in vector quantized models, improving codebook utilization and scalability across various tasks.

Contribution

We propose SimVQ, a novel linear reparameterization method that addresses codebook collapse by optimizing the entire linear space instead of individual code vectors.

Findings

Improves codebook utilization in VQ models

Effective across image and audio tasks

Easy to implement and generalizes well

Abstract

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.