Addressing Representation Collapse in Vector Quantized Models with One Linear Layer

## Abstract

Abstract Vector Quantization (VQ) is a widely used method for converting continuous representations into discrete codes, which has become fundamental in unsupervised representation learning and latent generative models. However, VQ models are often hindered by the problem of representation collapse in the latent space, which leads to low codebook utilization and limits the scalability of the codebook for large-scale training. Existing methods designed to mitigate representation collapse typically reduce the dimensionality of latent space at the expense of model capacity, which do not fully resolve the core issue. In this study, we conduct a theoretical analysis of representation collapse in VQ models and identify its primary cause as the disjoint optimization of the codebook, where only a small subset of code vectors are updated through gradient descent. To address this issue, we propose SimVQ , a novel method which reparameterizes the code vectors through a linear transformation layer based on a learnable latent basis. This transformation optimizes the entire linear space spanned by the codebook, rather than merely updating the code vector selected by the nearest-neighbor search in vanilla VQ models. Although it is commonly understood that the multiplication of two linear matrices is equivalent to applying a single linear layer, our approach works surprisingly well in resolving the collapse issue in VQ models with just one linear layer. We validate the efficacy of SimVQ through extensive experiments across various modalities, including image and audio data with different model architectures. The results show that SimVQ not only effectively addresses the problem of representation collapse but also proves highly adaptable and easy to implement, suggesting its broad applicability in diverse machine learning contexts. Our code is available at https://github.com/youngsheen/SimVQ .

## 1 Introduction

In recent years, vector quantization (VQ) (van den Oord et al., ; Razavi et al., ) has emerged as a foundational technique in unsupervised representation learning (Baevski et al., ; Bruce et al., ) and latent generative models (Rombach et al., ; Yu et al., ; ; Borsos et al., ; Wang et al., ; Zhu et al., ). By converting continuous representations into discrete codes, VQ models can effectively identify the inherent structure of data and enable various discrete modeling methods on continuous data, from high-quality image generation (Esser et al., ) to audio synthesis (Défossez et al., ). The recent success of Large Language Models (LLMs) (Achiam et al., ) has highlighted the effectiveness of next-token prediction as a powerful and versatile training objective. Consequently, VQ models are taken as the direct method to transform data from various modalities (Zhang et al., ; Sun et al., ; Team, ) or scientific domains (Gao et al., ) to discrete sequences for next token prediction training. However, attempts to integrate VQ models as multimodal tokenizers to leverage the scaling laws of LLMs face significant challenges because of the difficulty of expanding the codebook. For example, the Chameleon model (Team, ) constrains its codebook size to $8k$, which is significantly trailing behind the vocabulary size of LLMs (e.g., LLaMA3’s vocabulary size is $128k$ (Dubey et al., )).

Figure: Figure 1: Comparison of Vanilla VQ and SimVQ. (a): (left) Disjoint optimization in Vanilla VQ. Only the nearest codes are updated, resulting in a high percentage of “dead” codes that are not updated. (b): (right) Joint optimization in SimVQ. The entire codebook is updated with a latent basis, ensuring all codes remain active.
Refer to caption: x1.png

There is a broad agreement that increasing vocabulary size can consistently improve the performance of LLMs (Tao et al., ). However, recent studies (Zhu et al., ) indicate that traditional VQ models often fail to utilize the additional parameters introduced by codebook expansion, leaving most codes inactive during training. The contradiction between codebook expansion and low codebook utilization in VQ models is known as the representation collapse problem (Roy et al., ), where increasing the codebook size fails to improve the performance.
To address these discrepancies, we conduct a theoretical analysis of the optimization procedure of VQ models and identify that the disjoint optimization of the codebook is the root cause of representation collapse. As illustrated in Fig. 1(a), the core mechanism of VQ models involves a nearest-neighbor replacement strategy, where the encoder’s output features are replaced by the nearest vector in the codebook to serve as input to the decoder. The indices of the nearest vector are taken as the discrete representation of the data. This nearest-selection operator results in only a subset of codes being updated through gradient descent, while the remaining codes remain unchanged.

Recently, some approaches have been proposed to mitigate representation collapse. FSQ (Mentzer et al., ), LFQ (Yu et al., ) and ViTVQGAN (Yu et al., ) reduce the dimension of the latent space to a very small scale (e.g., 8 vs. 128) to alleviate the curse of dimensionality, thereby improving the overlap between the encoder’s features and the codebook. However, while these methods enhance codebook utilization, they do so at the cost of model capacity, leading to worse performance compared to vanilla VQ models when the codebook size is small and representation collapse is not severe. Another approach, VQGAN-LC (Zhu et al., ), initializes the codebook with features extracted from the pre-trained CLIP model (Radford et al., ) to create a well-structured latent space that better matches the distribution of the encoder output. Nevertheless, the latent space defined by an external pre-trained model limits the model’s ability to generalize to diverse datasets and reaches a performance plateau as the codebook size increases. These limitations highlight the need for a more effective method to improve codebook utilization without compromising model capacity or relying on external models.

We critically assess prevalent methodologies and reveal that optimizing the latent space rather than individual code vectors is key to preventing representation collapse. Building on this insight, we introduce a simple yet effective method, termed SimVQ, to directly update the latent space spanned by the codebook by linear transforming the code vectors via a learnable latent basis. Specifically, the vectors in the codebook are reparameterized as a linear combination of the basis in the learnable linear layer $\bm{W}$:

$$ $\bm{C}\in\mathbb{R}^{K\times d}\Rightarrow\bm{CW}~{}\text{with}~{}\bm{W}\in \mathbb{R}^{d\times d},$ (1) $$

where $K$ denotes the codebook size and $d$ represents the dimension of latent space. This reparameterization with linear transformation disentangles the optimization of the codebook into two components: the coefficient matrix $\bm{C}$ and the basis of linear space $\bm{W}$ respectively. As illustrated in Fig. 1(b), by optimizing the basis matrix $\bm{W}$, the latent space spanned by $\bm{CW}$ is rotated and stretched to match encoder’s output feature. The entire codebook is updated jointly to prevent the representation collapse problem. The simplicity of the proposed method makes it highly portable and easily adaptable for improving VQ-based models across a wide range of domains, requiring only one linear layer.

In summary, our contributions to vector quantized models are as follows:

- •
We theoretically analyze the representation collapse problem of VQ models and reveal that optimizing the latent space spanned by the codebook, rather than focusing on the individual code vectors, is crucial to addressing this issue.
- •
We propose a novel method, SimVQ, which reparameterizes the codebook vectors in VQ models via a linear transformation with a learnable latent basis. This simple yet effective approach is highly adaptable and easy to implement, making it broadly applicable across various machine learning contexts.
- •
We conduct an extensive evaluation of SimVQ across diverse modalities, including image and audio with different model architectures. The results show that SimVQ not only effectively addresses the representation collapse problem by achieving near-complete codebook utilization regardless of the codebook size, but also establishes new state-of-the-art performance. Furthermore, when scaling up the codebook size, SimVQ consistently delivers improved results.

## 2 Related Work

VQ-VAE (van den Oord et al., ) is the pioneering work to encode data into discrete representations, which is further improved by VQ-VAE2 (Razavi et al., ) by employing a hierarchical architecture to enable richer representations. Building on these developments, VQGAN (Esser et al., ) combines VQ-VAE with adversarial networks to improve the perceptual quality of generated samples and establish a fundamental quantization protocol in latent generative models (Rombach et al., ; Yu et al., ; Team, ). However, these methods suffer from a critical issue of representation collapse, as they struggle to scale the codebook size beyond 10k entries, limiting their scalability. In response to this challenge, several approaches have been proposed recently. DALLE (Ramesh et al., ) employs the gumbel-softmax trick (Jang et al., ) and stochastic sampling strategies to activate most codes during training. However, during inference, only a small subset of codes is utilized for quantization (Zhang et al., ). Huh et al. () proposes rescaling the vectors in the codebook during training to match the distributions in the latent space. VQGAN-FC (Yu et al., ) introduces a method to map latent vectors into a lower-dimensional space followed by $l_{2}$ normalization to alleviate representation collapse. FSQ (Mentzer et al., ) extends this idea by projecting representations into a reduced-dimensional space, where they are quantized into a small set of fixed values. LFQ (Yu et al., ), a variant of FSQ, uses binary values for quantized representations, thereby simplifying the encoding process. While these methods improve the codebook utilization, they do so at the cost of model capacity by significantly reducing the dimensionality of latent space (often to as low as 8), leading to worse performance compared to vanilla VQ models when the codebook size is small and representation collapse is not severe. Additionally, VQGAN-LC (Zhu et al., ) proposes to initialize the codebook using features extracted from the pre-trained CLIP model to avoid representation collapse. However, the reliance on the pre-trained model limits the VQ model’s ability to generalize to diverse datasets and results in a performance plateau as the codebook size increases. In contrast, our method, SimVQ, effectively addresses the representation collapse problem with a simple linear layer, without sacrificing model capacity or relying on external pre-trained models.

## 3 Representation Collapse in VQ Models

### 3.1 Preliminaries

A vector quantized model is typically a reconstructive encoder-decoder architecture that includes a vector quantization layer to convert continuous representations into discrete codes. For simplicity, we represent an image with a single random variable $x$. Formally, the encoder $f_{\theta}$ maps the input image into a latent space, producing a continuous representation $z_{e}=f_{\theta}(x)\in\mathbb{R}^{d}$. This representation is then quantized using a learnable codebook $\bm{C}=[q_{1},\ldots,q_{K}]\in\mathbb{R}^{K\times d}$, where $q_{i}$ is a codebook vector. We define $\delta_{k}\in\{0,1\}^{1\times K}$ as a characteristic (one-hot) vector where only the $k$-th element is $1$, such that $q_{k}=\delta_{k}\bm{C}\in\mathbb{R}^{1\times d}$. The quantization layer selects the nearest codebook vector $q_{k}$ by minimizing the Euclidean distance between $z_{e}$ and the codebook entries (van den Oord et al., ):

$$ $\displaystyle k=\arg\min_{j}\|z_{e}-q_{j}\|^{2}_{2}=\arg\min_{j}\|z_{e}-\delta _{j}\bm{C}\|^{2}_{2}.$ (2) $$

The selected vector $q_{k}$ is then passed to the decoder $g_{\phi}$ to reconstruct the input image.

To enable gradient propagation through the non-differentiable characteristic vector $\delta_{k}$, the straight-through estimator (STE) (Bengio et al., ) is applied. During the backward pass, the gradient of $z_{q}=\delta_{k}\bm{C}$ is copied to $z_{e}$ as follows,

$$ $\displaystyle z_{q}=\text{sg}(\delta_{k}\bm{C}-z_{e})+z_{e},\quad\Rightarrow \frac{\partial z_{q}}{\partial z_{e}}=1$ (3) $$

where sg is the stop gradient operator, ensuring the gradient for $\delta_{k}\bm{C}$ is discarded during the backward pass.

The learning objective is the combination of a reconstruction loss and commitment loss that ensures that the encoder commits to an embedding and the encoder’s output does not drift:

$$ $\mathcal{L}=\log p(x|z_{q})+\|\text{sg}(\delta_{k}\bm{C})-z_{e}\|^{2}_{2}+ \beta\|\delta_{k}\bm{C}-\text{sg}(z_{e})\|^{2}_{2},$ (4) $$

where $\log p(x|z_{q})$ is typically the mean squared error (MSE) loss $\|x-g_{\phi}(z_{q})\|^{2}_{2}$ for image and audio data.

### 3.2 Disjoint Optimization of Codebook

In VQ models, only the nearest code is selected and updated via gradient descent. Ideally, all codebook entries should be updated and utilized for decoding. However, experimental evidence shows that only a small fraction of the codebook gets updated and utilized, leading to what is known as the representation collapse problem (Roy et al., ). To investigate the root cause of this issue, we provide a theoretical analysis of the optimization dynamics in VQ models.

Due to the use of the straight-through estimator (STE) for gradient propagation, the codebook $\bm{C}$ can only be updated through the gradient of the commitment loss, which is defined as:

$$ $\mathcal{L}_{commit}(\bm{C})=\|z_{e}-\delta_{k}\bm{C}\|_{2}^{2}.$ (5) $$

The codebook $\bm{C}$ is updated according to the following equation, where $\eta$ is the learning rate:

$$ $\bm{C}^{(t+1)}=\bm{C}^{(t)}+\eta\mathbb{E}_{z_{e}}\left[\frac{\partial\mathcal {L}_{commit}(\bm{C}^{(t)})}{\partial\bm{C}^{(t)}}\right]=\bm{C}^{(t)}-\eta \mathbb{E}_{z_{e}}\left[\delta_{k}^{T}\delta_{k}\bm{C}^{(t)}\right]+\eta \mathbb{E}_{z_{e}}\left[\delta_{k}^{T}z_{e}\right]$ (6) $$

where $\delta_{k}^{T}\delta_{k}$ is the Kronecker delta matrix, defined as:

$$ $(\delta_{k}^{T}\delta_{k})_{ij}=\begin{cases}1&\text{if }i=j=k,\\ 0&\text{otherwise}.\end{cases}$ (7) $$

All vectors in $\bm{C}$ will be updated and utilized if and only if the expectation $\mathbb{E}_{z_{e}}\left[\delta_{k}^{T}\delta_{k}\right]$ converges to the identity matrix. Unlike variational autoencoders (VAEs) (Kingma & Welling, ), which enforce a Gaussian distribution on the latent space via a KL-divergence penalty, VQ models optimize $z_{e}$ towards the selected codebook vectors $\mathbb{E}_{z_{e}}\left[\delta_{k}^{T}\delta_{k}\bm{C}\right]$. At the same time, the selected codebook vectors are optimized towards the distribution of $z_{e}$ , resulting in the same selected subset of vectors moving closer to $z_{e}$, somewhat akin to a cocoon effect. However, this disjoint optimization of the codebook leads to part of the codebook, specifically $(\bm{I}-\mathbb{E}_{z_{e}}\left[\delta_{k}^{T}\delta_{k}\right])\bm{C}$, remaining unupdated and underutilized once the optimization process begins. This phenomenon occurs because the optimization focuses only on a subset of codebook vectors, leaving other vectors effectively stagnant.

This analysis reveals the fundamental cause of representation collapse in VQ models: the disjoint optimization process that updates only a subset of codebook vectors. This insight forms the basis for our proposed solution, SimVQ, which aims to address this issue by optimizing the entire latent space spanned by the codebook, rather than individual code vectors.

## 4 Addressing Collapse with Latent Linear Transformation

### 4.1 Reparameterize Codes with Latent Basis

Let $\bm{W}=\{\bm{w}_{1},\ldots,\bm{w}_{n}\}$ be a basis of a linear space. Any vector $\bm{v}$ in the space can be uniquely expressed as a linear combination of the basis vectors with coefficients $c_{1},\ldots,c_{n}\in\mathbb{R}$:

$$ $\bm{v}=c_{1}\bm{w}_{1}+\cdots+c_{n}\bm{w}_{n}=\bm{c}\bm{W}.$ (8) $$

Given the equivalence between $\bm{v}$ and $\bm{c}\bm{W}$ in the linear space, we can reparameterize each vector in the codebook of VQ models with a new basis matrix $\bm{W}\in\mathbb{R}^{d\times d}$. Specifically, the codebook $\bm{C}=\{\bm{c}_{1},\ldots,\bm{c}_{K}\}$ can be reparameterized as:

$$ $\{\bm{\hat{c}}_{1}\bm{W},\ldots,\bm{\hat{c}}_{N}\bm{W}\}=\bm{\hat{C}}\bm{W}\in \mathbb{R}^{K\times d}.$ (9) $$

This reparameterization introduces two components: the basis matrix $\bm{W}$ and the coefficient matrix $\bm{\hat{C}}$. In the following, we will discuss the optimization of both the basis matrix $\bm{W}$ and the coefficient matrix $\bm{\hat{C}}$. For simplicity, we slightly abuse $\bm{C}$ and $\bm{\hat{C}}$ below.

##### Asymmetric Optimization Dynamics

While it is commonly accepted that multiplying two linear matrices is equivalent
to a single linear layer, we argue that the disjoint optimization problem of the codebook in VQ models can be addressed by the basis transformation. In vanilla VQ models, only the codebook $\bm{C}$ is responsible for minimizing commitment loss, leading to the disjoint optimization problem where only the selected codes will be updated.

In contrast, when the codebook is reparameterized as $\bm{C}\bm{W}$, both the basis $\bm{W}$ and the coefficient matrix $\bm{C}$ contribute to minimizing the commitment loss. The gradients $\frac{\partial\mathcal{L}}{\partial\bm{W}}$ and $\frac{\partial\mathcal{L}}{\partial\bm{C}}$ can simultaneously reduce the loss. As a result, the optimization of the reparameterized codebook can be divided into three scenarios:

- •
Updating $\bm{C}$ with $\bm{W}$ frozen: Only the selected codes adapt to the latent distribution of $z_{e}$, as depicted on Fig. 1(a). The vanilla VQ is a special case of this scenario with $\bm{W}=\bm{I}$.
- •
Updating $\bm{W}$ with $\bm{C}$ frozen: The entire codebook $\bm{C}\bm{W}$ adjusts to the latent distribution of $z_{e}$. The basis matrix $\bm{W}$ rotates and stretches the codebook space as shown on Fig. 1(b).
- •
Updating both $\bm{C}$ and $\bm{W}$: The selected subset of codes moves towards $z_{e}$ while the space spanned by $W$ undergoes simultaneous rotation and stretching.

### 4.2 Toy Examples

Figure: Figure 2: (a): (left) The optimization trajectory of the objective $\|\bm{x}-\bm{q}\|^{2}_{2}$, which is the same as vanilla VQ. Only a small fraction of points are updated while others remain inactive. (b): (right) The optimization trajectory of the objective $\|\bm{x}-\bm{q}\bm{w}\|^{2}_{2}$ with $\bm{q}$ frozen, which is the same as SimVQ. All the points are updated towards targets $x$.
Refer to caption: x2.png

Figure: Figure 3: (a): (left) The optimization trajectory of the optimization objective: $\|\bm{x}-\bm{q}\bm{w}\|^{2}_{2}$ with both $\bm{q}$ and $\bm{w}$ unfrozen. (b): (right) The Frobenius norm of the projection matrix $\bm{w}$ and loss curves. The loss quickly converges to 0 with $\bm{w}$ almost unchanged.
Refer to caption: x4.png

To highlight the difference in optimization between $\bm{C}$ and $\bm{CW}$, we conduct a toy experiment in a two-dimensional setting and visualize the optimization process in Fig. 2 and Fig. 3. We randomly sample two target points $\bm{x}$ from Gaussian distribution as follows:

$$ $\bm{x}_{1}\sim\mathcal{N}(\left(\begin{matrix}2\\ 2\end{matrix}\right),\left(\begin{matrix}1&0\\ 0&1\end{matrix}\right)),\quad\bm{x}_{2}\sim\mathcal{N}(\left(\begin{matrix}-2 \\ -2\end{matrix}\right),\left(\begin{matrix}1&0\\ 0&1\end{matrix}\right)).$ (10) $$

Then we initialize $10$ learnable vectors $\bm{q}$ from a Gaussian distribution:

$$ $\{\bm{q}_{i}\}_{i=1}^{10}\sim\mathcal{N}(\left(\begin{matrix}0\\ 0\end{matrix}\right),\left(\begin{matrix}1&0\\ 0&1\end{matrix}\right)),$ (11) $$

During training with gradient descent, we introduce perturbation noise $\mathcal{N}(0,0.01)$ to the targets. In Fig. 2(a), the optimization objective is similar to vanilla VQ: $\|\bm{x}-\bm{q}\|^{2}_{2}$. Only the nearest points $\bm{q}_{4}$ and $\bm{q}_{10}$ are updated. In contrast, in Fig. 2(b), the optimization objective $\|\bm{x}-\bm{q}\bm{w}\|^{2}_{2}$ is similar to SimVQ with the points reparameterized by a learnable latent basis $\bm{w}$ and $\bm{q}$ frozen, resulting in the entire codebook $\{\bm{q}\}_{i=1}^{10}$ being jointly updated.

###### Remark 4.1 .

The simultaneous optimization of the latent basis $\bm{w}$ and the coefficient matrix $\bm{q}$ may lead to the collapse.

We provide an example in Fig. 3(a) where the optimization objective is $\|\bm{x}-\bm{q}\bm{w}\|^{2}_{2}$ with $\bm{q}$ unfrozen this time. In the training process, only the nearest point $\bm{q}_{1}$ and point $\bm{q}_{10}$ move towards the target point, while other points remain almost unchanged. We also visualize the loss curve in Fig. 3(b). The optimization objective with both $\bm{q}$ and $\bm{w}$ unfrozen converges quickly, where the norm of basis $\bm{w}$ is much smaller than the objective with $\bm{q}$ frozen. This indicates that the disjoint optimization of the codebook persists: $\bm{q}$ can directly commit to the loss and dominate the optimization process, with $\bm{w}$ being largely ignored, leading to the collapse quickly.

###### Remark 4.1 .

### 4.3 Joint Optimization of the Codebook

Figure: Algorithm 1 Training Procedure for SimVQ

We propose SimVQ by simply using a learnable basis $\bm{W}\in\mathbb{R}^{d\times d}$ to reparameterize the codebook such that the codebook is transformed into $\bm{C}\bm{W}$. The pseudo-code for this approach is provided in Algorithm 1. During training, we optimize only the latent basis matrix $\bm{W}$, while keeping the coefficient matrix $\bm{C}$ frozen. The commitment loss for SimVQ is defined as:

$$ $\mathcal{L}_{commit}(z_{e},q_{k})=\|z_{e}-\delta_{k}\bm{C}\bm{W}\|^{2}_{2}.$ (12) $$

The vanilla VQ model is a special case of SimVQ, where the linear basis matrix $\bm{W}$ is fixed as the identity matrix $\bm{I}$. The update for $\bm{W}$ with learning rate $\eta$ is:

$$ $\bm{W}^{(t+1)}=\bm{W}^{(t)}-\eta\frac{\partial\mathcal{L}_{commit}(z_{e},\bm{q }_{k})}{\partial\bm{W}^{(t)}}=(\bm{I}-\eta\mathbb{E}_{z_{e}}\left[\bm{C}^{T} \delta_{k}^{T}\delta_{k}\bm{C}\right])\bm{W}^{(t)}+\eta\mathbb{E}_{z_{e}}\left [\bm{C}^{T}\delta_{k}^{T}z_{e}\right].$ (13) $$

The term $\mathbb{E}\left[\bm{C}^{T}\delta_{k}^{T}\delta_{k}\bm{C}\right]$ represents the expectation of the quadratic form, and simplifies to $\mathbb{E}[\bm{q}_{k}^{T}\bm{q}_{k}]$. Since the codes are randomly sampled from a Gaussian distribution, we have:

$$ $\mathbb{E}\left[\bm{q}_{k}^{T}\bm{q}_{k}\right]=\bm{I},\text{where}~{}\bm{q} \sim\mathcal{N}(0,1),$ (14) $$

which ensures that all elements of $\bm{W}$ are updated. As training progresses, the latent basis $\bm{W}$ converges to:

$$ $\lim_{t\rightarrow\infty}\bm{W}^{(t)}=\mathbb{E}_{z_{e}}\left[\bm{q}_{k}^{T}z_ {e}\right]$ (15) $$

Thus, in the limit:

$$ $\lim_{t\rightarrow\infty}\bm{q}_{k}\bm{W}^{(t)}=\mathbb{E}\left[\bm{q}_{k}\bm{ q}_{k}^{T}\bm{e}\right]=\mathbb{E}\left[\bm{e}\right]$ (16) $$

At convergence, the product $\bm{q}_{k}\bm{W}$ equals the nearest feature.

### 4.4 Efficiency Analysis

SimVQ demonstrates greater efficiency than vanilla VQ due to its asymmetric training strategy, wherein the codebook $\bm{C}$ remains static and only the linear projection $\bm{W}$ is optimized. This approach results in a significant reduction in memory usage during the gradient backpropagation process. In vanilla VQ, the memory cost for the optimization of the codebook is $O(Kd)$, where $K$ is the number of vectors in the codebook, and $d$ is the dimension of each vector. In our experiments, $K=65,536$ is much larger than $d=128$. As the vocabulary size increases, the memory required for backpropagation grows proportionally, significantly impacting resource consumption. In contrast, SimVQ’s memory cost for backpropagation is only $O(d^{2})$ because the codebook $\bm{C}$ is fixed, and only the linear layer $\bm{W}$ is updated. This results in a constant memory requirement in backpropagation, independent of the vocabulary size. The $d\times d$ scaling becomes particularly advantageous as $K$ increases in practical applications. This structural design minimizes the computational overhead and improves training efficiency, especially when dealing with large vocabularies.

## 5 Experiments

To assess the efficacy and versatility of the proposed SimVQ, we conduct experiments across both image and audio modalities. Subsequently, we analyze the learned linear layer to investigate the latent basis. The experimental configurations are listed in Appendix A.1.

### 5.1 Vision Modality

#### 5.1.1 Implementation Details

To rigorously evaluate the proposed SimVQ, we reproduce all the VQ models listed in Tab. 1 using the same architecture of VQGAN (Esser et al., ) with the quantization layer different only.
Among the baselines, for VQGAN-FC (Yu et al., ), we follow the original setting to reduce the dimension of the latent space to $8$ followed by $l_{2}$ normalization to improve codebook utilization. For FSQ (Mentzer et al., ), we adopt a codebook size of $[8,8,8,5,5,5,]$ as recommended, to approximately match the default codebook size. For VQGAN-LC (Zhu et al., ), we follow them and leverage an external pre-trained CLIP model to extract features of the training dataset in advance for a well-defined latent space. All models are trained on the ImageNet (Deng et al., ) dataset for 50 epochs with a batch size of 256. Input images are processed at a resolution of $128\times 128$ pixels and downsampled by a factor of $8$, yielding a feature map of $16\times 16\times 128$, where $128$ is the dimension of the latent space. We set the default codebook size to a large number of $2^{16}=65536$ rather than the traditional number $8192$ to highlight the representation collapse problem. Performance is evaluated using rFID, LPIPS, PSNR, and SSIM metrics on the ImageNet validation set.

**Table 1: Reconstruction performance on ImageNet-1k with a resolution of $128\times 128$. All models are trained using images downsampled into $16\times 16$ tokens. $\dagger$ Results are reproduced using the codebook size of $[8,8,8,5,5,5]$ to approximately match $65,536$. $+$ Following VQGAN-LC, we extract CLIP features with the codebook frozen.**
| Method | Latent dim | Codebook size | Util$\uparrow$ | rFID$\downarrow$ | LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- |
| VQGAN (Esser et al., ) | 128 | 65,536 | 1.4% | 3.74 | 0.17 | 22.20 | 70.6 |
| VQGAN-EMA (Razavi et al., ) | 128 | 65,536 | 4.5% | 3.23 | 0.15 | 22.89 | 72.3 |
| VQGAN-FC (Yu et al., ) | 128 | 65,536 | 1.4% | 5.33 | 0.18 | 21.45 | 68.8 |
| VQGAN-FC (Yu et al., ) | 8 | 65,536 | 100.0% | 2.63 | 0.13 | 23.79 | 77.5 |
| FSQ^† (Mentzer et al., ) | 16 | 64,000 | 100.0% | 2.80 | 0.13 | 23.63 | 75.8 |
| LFQ (Yu et al., ) | 6 | 65,536 | 100.0% | 2.88 | 0.13 | 23.60 | 77.2 |
| VQGAN-LC-CLIP^+ (Zhu et al., ) | 768 | 65,536 | 100.0% | 2.40 | 0.13 | 23.98 | 77.3 |
| SimVQ (ours) | 128 | 65,536 | 100.0% | 2.24 | 0.12 | 24.15 | 78.4 |
| SimVQ (ours) | 128 | 262,144 | 100.0% | 1.99 | 0.11 | 24.68 | 80.3 |

#### 5.1.2 Main Results

Tab. 1 presents the reconstruction performance of various VQ models on image data. We make three key observations: 1) Traditional VQGAN models utilize only a very small subset of the codebook, with a utilization rate of just $1.4\%$. Although VQGAN-EMA is proposed to improve codebook utilization, especially when the codebook size scales up to $65k$, it still suffers from severe representation collapse. 2) Recently proposed methods, such as LFQ, FSQ, and VQGAN-FC, effectively improve codebook utilization to $100\%$. However, these methods require reducing the latent space to a very low dimension. For example, applying VQGAN-FC to the standard latent dimension of $128$ results in severe representation collapse and degraded reconstruction performance. Additionally, these models face limitations in model capacity due to the low-dimensional latent space. While they achieve full codebook utilization, their reconstruction quality on rFID score lags significantly behind SimVQ. 3) VQGAN-LC-CLIP leverages an external pre-trained CLIP model to provide a well-defined latent space. However, VQGAN-LC relies on CLIP features pre-trained on much larger datasets than ImageNet, which introduces generalization issues and a lower performance ceiling (degradation issue in Tab. 2). In contrast, SimVQ can be applied to a wide range of data types and achieves superior performance (rFID $2.40$ vs. $2.24$) without the limitations imposed by a pre-trained feature extraction model.

#### 5.1.3 Ablation Study

##### Codebook Size

In Tab. 2, we explore the impact of different codebook sizes, ranging from $1k$ to $262k$, which is the level of LLM’s vocabulary size. SimVQ consistently improves performance as the codebook size increases. For instance, the rFID score decreases to $1.99$, and SSIM surpasses $80.0$. In contrast, VQGAN-LC-CLIP encounters performance degradation, with the rFID score worsening from $2.62$ to $2.66$ when the codebook size is increased from $100,000$ to $200,000$.

##### Codebook Optimization Strategy

We investigate codebook initialization and the training of the linear layer in Tab. 3. Our findings are as follows: 1) The codebook is robust to different initialization strategies, yielding similar results with both Gaussian and uniform initialization. 2) When the codebook is updated during training, SimVQ continues to address the representation collapse issue, though with a slight degradation in performance.

**Table 2: Ablation study on the effect of various codebook sizes on ImageNet at a resolution of $128\times 128$. $\dagger$ We directly copy the reported results of VQGAN-LC from the original paper on ImageNet $256\times 256$ resolution.**
| Method | Codebook Size | Util$\uparrow$ | rFID$\downarrow$ | LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- |
| VQGAN-LC-CLIP^† | 50,000 | $99.9\%$ | 2.75 | 0.13 | 23.8 | 58.4 |
| VQGAN-LC-CLIP^† | 100,000 | $99.9\%$ | 2.62 | 0.12 | 23.8 | 58.9 |
| VQGAN-LC-CLIP^† | 200,000 | $99.8\%$ | 2.66 | 0.12 | 23.9 | 59.2 |
| SimVQ | 1,024 | 100.0% | 3.67 | 0.16 | 22.34 | 70.8 |
| SimVQ | 8,192 | 100.0% | 2.98 | 0.14 | 23.23 | 74.7 |
| SimVQ | 65,536 | 100.0% | 2.24 | 0.12 | 24.15 | 78.4 |
| SimVQ | 262,144 | 100.0% | 1.99 | 0.11 | 24.68 | 80.3 |

**Table 3: Ablation study of codebook optimization.**
| Initialization | Trainable | Util$\uparrow$ | rFID$\downarrow$ | LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- |
| Gaussian | Yes | 100.0% | 2.31 | 0.12 | 24.04 | 77.2 |
| Uniform | No | 100.0% | 2.31 | 0.12 | 24.15 | 78.4 |
| Gaussian | No | 100.0% | 2.24 | 0.12 | 24.15 | 78.4 |

### 5.2 Audio Modality

#### 5.2.1 Implementation Details

We use LibriTTS dataset (Zen et al., ) for audio-based VQ model training. The baselines such as Encodec (Défossez et al., ), Vocos (Siuzdak, ), and SpeechTokenizer (Zhang et al., ) are based on residual vector quantization. Our SimVQ model adopts the same architecture as WavTokenizer (Ji et al., ) with the only modification being the replacement of their EMA codebook with our one linear layer reparameterization method. We train SimVQ on LibriTTS-580h for 50 epochs with a batch size of 64. Note that WavTokenizer is trained with a 3-second window size for optimal performance, we train SimVQ using a 1-second window to accelerate training. For objective evaluation of the reconstructed audio, we follow Vocos (Siuzdak, ) and employ metrics such as UTMOS (Saeki et al., ), PESQ (Rix et al., ), STOI, and the F1 score for voiced/unvoiced classification (V/UV F1). UTMOS is particularly valuable as it produces scores highly correlated with human evaluations.

#### 5.2.2 Main Results

Tab. 4 presents the reconstruction performance of various VQ models on audio data. Baseline models using residual vector quantization perform significantly worse than SimVQ, even when utilizing much larger bandwidths. Despite using the same architecture as WavTokenizer, our model, which replaces the quantization layer with SimVQ, achieves superior performance with a 1-second window size and maintains nearly $100\%$ codebook utilization when scaling up to a size of 262,144. The consistent performance of the SimVQ model across both image and audio data demonstrates that SimVQ is a general method for addressing the representation collapse problem in VQ models and can be effectively applied across multiple modalities.

**Table 4: Reconstruction performance on LibriTTS test-clean/test-other dataset. $*$ WavTokenizer is trained with a window size of 3 seconds. The bandwidth of 0.9kbps, 0.975kbps, 1.2kbps, 1.35kbps means the codebook size of 4096, 8192, 65536, 262144 respectively.**
| Method | Bandwidth | Util$\uparrow$ | UTMOS$\uparrow$ | PESQ$\uparrow$ | STOI$\uparrow$ | V/UV F1$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- |
| GT | - | - | 4.06/3.48 | - | - | - |
| EnCodec (Défossez et al., ) | 3.0kbps | - | 2.31/2.09 | 2.05/2.05 | 0.90/0.88 | 0.92/0.89 |
| Vocos (Siuzdak, ) | 3.0kbps | - | 3.53/3.06 | 2.40/2.19 | 0.92/0.90 | 0.94/0.91 |
| SpeechTokenizer (Zhang et al., ) | 3.0kbps | - | 3.56/3.02 | 1.93/1.74 | 0.88/0.84 | 0.93/0.89 |
| WavTokenizer (Ji et al., ) | 0.9kbps | 100/100% | 3.74/3.43^∗ | 2.01/2.26^∗ | 0.89/0.89^∗ | 0.92/0.92^∗ |
| SimVQ (ours) | 0.9kbps | 100.0/100.0% | 4.00/3.51 | 2.33/2.08 | 0.91/0.88 | 0.94/0.91 |
| WavTokenizer (Ji et al., ) | 0.975kbps | 68/-% | 4.02^∗/- | 2.39^∗/- | 0.92^∗/- | 0.94^∗/- |
| WavTokenizer (Ji et al., ) | 1.05kbps | 27/-% | 4.00^∗/- | 2.36^∗/- | 0.81^∗/- | 0.94^∗/- |
| SimVQ (ours) | 0.975kbps | 99.4/99.4% | 4.03/3.52 | 2.42/2.15 | 0.92/0.88 | 0.94/0.92 |
| SimVQ (ours) | 1.2kbps | 99.4/99.0% | 4.03/3.52 | 2.54/2.26 | 0.93/0.90 | 0.94/0.92 |
| SimVQ (ours) | 1.35kbps | 95.6/94.7% | 4.03/3.53 | 2.61/2.31 | 0.93/0.90 | 0.95/0.93 |

### 5.3 Analysis

In Fig. 4(a), we plot the rank of the latent basis matrix over training epochs. Notably, SimVQ demonstrates the ability to adaptively adjust the dimensionality of the latent space. Specifically, when the codebook size is increased from $65,536$ to $262,144$, the rank of the latent basis matrix decreases more rapidly and converges to a lower value. This observation suggests that a larger codebook can effectively alleviate the pressure on the latent space dimensionality, allowing the model to represent data more efficiently. Additionally, despite the rank decreasing to a lower-dimensional space, SimVQ maintains $100\%$ codebook utilization, highlighting its superiority over VQGAN-FC, which struggles when increasing the latent dimension from 8 to 128.
We also calculate the Frobenius norm of the latent basis matrix, as shown in Fig. 4. The norm of a codebook size of $262,144$ is slightly large than for $65,536$, indicating that a larger codebook can span a broader area in the linear space.
For a comprehensive evaluation, we also provide the reconstruction loss curve on the ImageNet validation dataset in Appendix A.2. The results consistently show that SimVQ achieves improved performance, further validating the effectiveness of our approach.

Figure: Figure 4: (a):(left) The rank of latent basis matrix $\bm{W}$ over training epochs. (b):(right) The Frobenius norm of latent basis matrix $\bm{W}$ over training epochs.
Refer to caption: x6.png

### 5.4 Qualitative Evaluation

We qualitatively compare the reconstruction quality of both images and audio in Appendix A.3. SimVQ achieves better reconstruction quality with an enlarged codebook size. For images, SimVQ with a larger codebook effectively preserves fine details, such as ”eyes” and ”text,” which are challenging for vanilla VQ models. For audio, SimVQ retains more acoustic details in both spectrograms and waveforms, as demonstrated in Fig. 7 and Fig. 8.

## 6 Conclusion

In this paper, we explore the representation collapse problem in VQ models. We conduct a theoretical analysis of the optimization process in VQ models and propose a simple yet effective method, SimVQ, to address this issue. Our method addresses the representation collapse by jointly optimizing the latent space through linear transformation with one linear layer. Experimental results demonstrate that SimVQ outperforms previous approaches on both image and audio datasets, highlighting its broad applicability across diverse machine learning tasks.

## Appendix A Appendix

### A.1 Experimental Configurations

**Table 5: Experimental configurations on image and audio.**
| Config | Image | Audio |
| --- | --- | --- |
| inputs | pixels | window size |
| input size | 128 $\times$ 128 $\times$ 3 | 24,000 $\times$ 1 |
| batch size | 256 | 64 |
| training epochs | 50 | 50 |
| quantized sequence length | 16 $\times$ 16 | 75 |
| optimization |  |  |
| optimizer | AdamW | AdamW |
| learning rate | 1e-4 | 1e-4 |
| learning rate schedule | constant | constant |
| warmup epochs | 0 | 0 |
| commitment coefficient | 1.0 | 1000.0 |
| adversarial coefficient | 0.1 | 1.0 |
| data augmentations |  |  |
| random horizontal flip | true | false |

### A.2 Loss Curve

Figure: Figure 5: The loss curve over epochs of different models on the validation dataset.
Refer to caption: x8.png

### A.3 Qualitative Cases

Figure: Figure 6: Image reconstruction samples with different codebook sizes.
Refer to caption: extracted/5955606/material/origin/2.png

Figure: Figure 7: The spectrogram of audio reconstruction samples with different codebook sizes.
Refer to caption: extracted/5955606/material/spec.png

Figure: Figure 8: The waveform of audio reconstruction samples with different codebook sizes.
Refer to caption: extracted/5955606/material/waveform.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: na
soundness: 2
strengths: 1. The authors identify disjoint optimization of the codebook as the key reason for representation collapse in VQ
2. The proposed approach SimVQ only requires adding one linear layer for implementation
3. Maintains effectiveness with increasing codebook size
summary: The paper titled "Addressing Representation Collapse in Vector Quantized Models with One Linear Layer" describes an apporach which addresses the representation collapse problem in VQ where only a small subset of codebook vectors are effectively utilized, primarily due to disjoint optimization during training. The paper introduces SimVQ, which addresses this by reparameterizing code vectors through a learnable linear transformation layer that optimizes the entire linear space spanned by the codebook, rather than just updating individual selected code vectors.
weaknesses: 1. the main contribution of the paper can be summed up by the fact that if assigning one codeword to the encoder output causes representation collapse, one can use a weighted combination of the whole codebook. While this may work well in practise, I dont think a whole new paper is needed on this. 
2. toy example section section if included needs to be properly fleshed out where the reader can work out the example by hand
3. please add reference to Residual Quantization
with Implicit Neural Codebooks, ICML 2024 and other recent papers on VQ

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: Questions on Codebook Reparameterization: The new codebook is represented as CW, raising several questions:

- "Remark 4.1. The simultaneous optimization of the latent basis w and the coefficient matrix q may lead to the collapse.".  Why is this the case?

- It is unclear whether the analysis of codebook utilization in this paper is based on C or CW, as the codebook initialization ablation appears to be applied to C.
soundness: 2
strengths: The paper is  easy to follow. 

The proposed codebook reparameterization, which ensures updates across the entire codebook, is a promising approach. 

The solution is straightforward yet effectively addresses codebook collapse, and the experiments on reconstruction tasks are convincing.
summary: The paper argues that a fundamental cause of representation collapse in VQ models stems from the disjoint optimization process, where only a subset of codebook vectors is updated. To address this problem, the authors propose reparameterizing the code vectors through linear transformation layer based on a learnable latent basis. This transformation allows optimization across the entire linear space spanned by the codebook, rather than only updating the single selected code vectors. Experiments on image and audio reconstruction tasks with various model architectures demonstrate that this approach outperforms the baselines.
weaknesses: **W1:** Limited Literature Review: The paper lacks a comprehensive review of related work, omitting several relevant methods that tackle codebook collapse, such as SQ-VAE [1], VQ-WAE [2], HVQ-VAE [3], and CVQ-VAE [4].  I find that the insight from Section 3.2 isn’t particularly novel. Codebook collapse is often attributed to the initialization process, where latents tend to concentrate around a few codebook vectors. During training, latents can easily overfit to these vectors (referred to as the disjoint optimization process in this paper), leaving the remaining code vectors unused. Simple techniques like codebook reset [5], which randomly reinitializes unused or infrequently used codewords to one of the encoder outputs or to vectors near frequently used codewords, can help mitigate this issue.

&nbsp;

**W2:** Missing Generation Results: Since VQ-VAE aims to learn meaningful discrete representations for downstream tasks, it would be valuable to see generation results to verify the quality of the codebook. With the new codewords represented as linear combinations of W, it is unclear how these representations will impact codeword effectiveness. Generation results could clarify this.

&nbsp;

**W3:** Missing Baselines: Important baseline methods such as SQ-VAE [1], VQ-WAE [2], HVQ-VAE [3], and CVQ-VAE [4] are not included in the experimental comparisons.

&nbsp;


[1] Takida, Yuhta, et al. "Sq-vae: Variational bayes on discrete representation with self-annealed stochastic quantization." ICML (2022).

[2] Vuong, Tung-Long, et al. "Vector quantized wasserstein auto-encoder." ICML (2023).

[3] Williams, Will, et al. "Hierarchical quantized autoencoders." Neurips (2020)

[4] Zheng, Chuanxia, and Andrea Vedaldi. "Online clustered codebook." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

[5] Dhariwal, Prafulla, et al. "Jukebox: A generative model for music." arXiv preprint arXiv:2005.00341 (2020).

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. How does the method works for just simple vector quantization problem (1 sparse sparse coding or clustering)? Would the same logic apply?
soundness: 2
strengths: 1. The idea is simple! I really like it.
2. It solve a problem that's annoying, important but also relatively ignored by the community.
summary: The author proposes a simple way to solve the “representation collapse” problem in VQVAE. In VQVAE, the latent vector is being quantized to a one-hot vector with stop-gradient operator. Only one element of the dictionary is updated during training. The updated dictionary will likely to attract more latent vector which reinforce that dictionary element to be more activate. This results in only a subset off dictionary used in VQVAE is not utilized after training. The author proposes a simple way to solve this problem. They parametrized the dictionary C = C_hat W, where C_hat is a random matrix, and W is trainable parameter. In this case, all dictionary element must be updated during each backward pass.
weaknesses: 1. The method can be understood as a low rank re-parametrization. I felt like this method should have lots of connection to classic dictionary learning, sparse coding, and vector quantization. But this connection is not explored in the paper. I'm not that familiar with either the current VQVAE or dictionary learning literature. But I know this idea of "dead unit" is not new problem and people must tried something. Maybe other reviewers/AC can also help if they're familiar with these literatures.
2. The author claim other method that proposes to solve the similar problem require the dictionary element's dimension to be small. Is this true? I happen to read [1]. Correct me if I'm wrong but I don't think they has this assumption. And the author didn't compare their method to this method. Just to be clear, I don't mean to ask the author to benchmark against all the method. I think this is still very simple and cool method. But I simply don't really understand the logic behind WHY other method assume dictionary element's dimension to be small. 
3. I personally think claiming "theoretically analyze" the dead unit phenomenon is too much. The analysis mentioned in the paper is not rigorous. To do so, I think you want to prove or cite some result for single layer dictionary learning. 
4.By the way the index for equation 9 should be cK not cN?


[1] M, Huh, et al. Straightening Out the Straight-Through Estimator: Overcoming Optimization Challenges in Vector Quantized Networks

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 2
questions: The work missed the generation results, which is a very important component of the proposed model. I would recommend the authors include generation results.
soundness: 2
strengths: 1. A simple linear layer can be used to address the codebook collapse problem and obtain better reconstruction results.
2. Experiments were demonstrated on two modalities.
summary: This paper addresses the issue of representation collapse in Vector Quantization (VQ) models, which hampers codebook utilization and limits scalability during large-scale training. Representation collapse occurs because the optimization process updates only a small subset of code vectors, leading to inefficient use of the latent space. Existing solutions reduce the dimensionality of the latent space but at the cost of model capacity.
To solve this, the authors propose SimVQ, a novel method that reparameterizes code vectors through a linear transformation based on a learnable latent basis. This approach optimizes the entire linear space spanned by the codebook, rather than just updating selected code vectors as in traditional VQ models.

While showing good reconstruction results, there are some shortcomings for this work.
weaknesses: 1. The idea using a linear combination of "atoms" and "coefficients" in the latent space has been explored on existing works.
    1) Auto-regressive Image Generation Using Residual Quantization.
    2) SC-VAE: Sparse Coding-based Variational Autoencoder with Learned ISTA.
    No discussion about these relevant works on the related work section. No comparisons with these methods were shown on the experiments section.
2. Though showing better reconstruction results. A very important aspect of vq-vae based model is to generate high-quality data (image, video, audio).  No generation results were exhibited. If I understand correctly, the proposed model can not discrete the latent representations. The latent representations were transformed into coefficients of codebook. How will you train a transformer model based on these coefficients?