Vector quantization is a technique in machine learning that discretizes continuous representations into a set of discrete vectors. It is widely employed in tokenizing data representations for large language models, diffusion models, and other generative models. Despite its prevalence, the characteristics and behaviors of vector quantization in generative models remain largely underexplored. In this study, we systematically investigate the issue of collapses in vector quantization, where collapsed representations are observed across discrete codebook tokens and continuous latent embeddings. By leveraging both synthetic and real datasets, we identify the severity of each type of collapses and triggering conditions. Our analysis reveals that random initialization and limited encoder capacity result in tokens collapse and embeddings collapse. Building on these findings, we propose potential solutions aimed at mitigating each collapse. To the best of our knowledge, this is the first comprehensive study examining representation collapsing problems in vector quantization.
Figures & tables
Figure 1: Token representation shrinkage degrade the diversity of generative model. Vector quantization is a widely used technique to map continuous data into discrete tokens, which assists the transformer-based generative model’s generation. We observe that token representation shrinkage, manifested as narrow distribution in latent space, leads to a shrunk distribution of the generated data.
Figure 2: Token representation shrinkage shrinks latent support and induces reconstruction mode collapse while Deferred Quantization mitigates it. (a) w/o deferred quantization: tokens cluster into a narrow region of the embedding space, reducing coverage and causing reconstructions to collapse onto fewer modes across different input dimensions. (b) w/ deferred quantization: tokens spread across the embedding space, improving latent support coverage and yielding reconstructions with better modality coverage.
Figure 3: Early quantization with clustered initialization induces token representation shrinkage while Deferred Quantization mitigates it. Initializing the codebook from an untrained encoder yields a narrow, uninformative embedding distribution, causing tokens to cluster and shrink latent support at early stage. In contrast, Deferred Quantization first learns a dispersed continuous representation and then initializes the codebook with semantic embeddings from the pretrained encoder, yielding better token coverage and reducing shrinkage.
ID
Claim
Evidence
Metrics
C1
Root cause of shrinkage is clustered codebook initialization from untrained encoder biases.
Fig. 2 : Visualization of token shrinkage on synthetic dataset. Fig. 3 : Initialization induced shrinkage.
Latent and reconstruction visualization.
C2
Deferred Quantization reduces shrinkage and improves codebook coverage.
Sec. 5 : Introduction of Deferred Quantization. Fig. 2 & Fig. 4 & Tab. 3 & Tab. 5 : Deferred Quantization show better token spread and codebook utilization as well as improved reconstruction quality.
Generative diversity loss is caused by early quantization via latent support shrinkage, not discretization itself.
Fig. 1 : Illustration of token representation shrinkage limiting generative diversity. Tab. 4 and Tab. 6 : Quantified loss of diversity (ImageNet-100 & ODIR).
g-FID ↓ , Pixel Distance ↑ , LPIPS Diversity ↑
C5
Mitigating shrinkage improves generation quality and diversity across models and datasets.
Tab. 4 & 6 : Consistent fidelity and diversity gains on ImageNet-100 and ODIR benchmarks.
g-FID ↓ , LPIPS Diversity ↑ , Pixel Distance ↑
Table 1: Claim-Evidence map for token representation shrinkage and Deferred Quantization. This table summarizes our key theoretical claims and maps them to specific empirical evidence, visualizations, and quantitative metrics provided in the paper.
Figure 4: Deferred Quantization alleviates the token representation shrinkage on CIFAR-10. (a) Token representation shrinkage in standard VQ impairs the model’s ability to scale, leading to stagnating MSE. Deferred Quantization mitigates this shrinkage, allowing the model to achieve better reconstruction performance as number of token increase. (b)Without Deferred Quantization, perplexity remains low, indicating highly uneven token usage. Deferred Quantization resolves this by spreading representations across the embedding space, ensuring high perplexity and efficient utilization of the available latent capacity. (c) The sharp drop in codebook Euclidean Distance for standard VQ indicates that tokens are clustering into a narrow region (shrinkage). Deferred Quantization maintains a high distance between tokens, preserving codebook coverage.
VAR Tokenizer
r-FID ↓
MSE ↓
LPIPS ↓
Euc. ↑
Perp. ↑
w/o Deferred
5.39
2.60
1.85
1.19
2801.88
+ SimVQ
5.52
2.50
1.78
0.80
6920.23
w/ Deferred
5.04
2.22
1.63
7.56
7044.51
+ SimVQ
4.93
2.17
1.59
0.93
8222.83
Table 2: Performance comparison of VAR tokenizers on the ImageNet-100 dataset. SimVQ ( Zhu et al., 2025 ) is a well-known method to solve the code collapse problem. The results show that mitigating code collapse alone does not lead to a substantial improvement, indicating that token representation shrinkage is different from code collapse.
Tokenizer
w/ Deferred?
r-FID ↓
MSE ↓
LPIPS ↓
Euc. ↑
Perp. ↑
MaskGIT
✓
8.58
3.28
2.34
18.75
5311.88
✗
12.22
3.91
2.70
6.45
924.57
VAR
✓
5.04
2.22
1.19
7.56
7044.51
✗
5.39
2.60
1.85
1.19
2801.88
Table 3: Performance evaluation of various tokenizer on the ImageNet-100 dataset. ✓ indicates that Deferred Quantization is applied, whereas ✗ indicates that Deferred Quantization is not applied.
Model
w/ Deferred?
g-FID ↓
Pixel Dist. ↑
LPIPS D. ↑
MaskGIT
✓
14.60
80.77
0.677
✗
14.75
75.89
0.668
VAR
✓
10.44
75.75
0.667
✗
12.88
70.69
0.633
Table 4: Comparison of generative models in the presence and absence of token representation shrinkage on ImageNet-100. ✓ indicates that token representation shrinkage is mitigated, whereas ✗ indicates without mitigation. Token representation shrinkage impairs both generative quality (g-FID) and diversity (pixel distance and LPIPS diversity).
Model
w/ Deferred?
r-FID ↓
MSE ↓
LPIPS ↓
Euc. ↑
Perp. ↑
VAR
✓
11.04
2.05
6.79
8.50
5396.17
✗
10.91
2.57
8.79
1.27
940.55
Table 5: Performance evaluation of VAR tokenizer on the ODIR medical dataset. ✓ indicates that token representation shrinkage is mitigated, whereas ✗ indicates without mitigation. The results indicate that token representation shrinkage also exists in a real-world medical setting.
Model
w/ Deferred?
g-FID ↓
Pixel Dist. ↑
LPIPS D. ↑
VAR
✓
34.33
49.83
0.401
✗
37.65
49.01
0.390
Table 6: Generative results of VAR on the ODIR medical dataset. ✓ indicates that Deferred Quantization is applied, whereas ✗ indicates that Deferred Quantization is not applied. The results indicate that token representation shrinkage also exists in a real-world medical setting.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Model
w/ Deferred?
g-FID ↓
Pixel Dist. ↑
LPIPS D.↑
MaskGIT
✓
10.85
79.17
0.712
✗
12.25
78.83
0.714
VAR
✓
8.30
77.10
0.719
✗
7.83
73.37
0.714
Appendix
Table 7: Comparison of generative models based on VQGAN on ImageNet-100. ✓ indicates that token representation shrinkage is mitigated, whereas ✗ indicates without mitigation.
Figure 5: Images generated using VAR. (a) ImageNet (a.left) and real-world medical images of eyes (a.right) generated by VAR w/o deferred quantization . (b) ImageNet (b.left) and real-world medical images of eyes (b.right) generated by VAR w/ deferred quantization .
Figure 6: Distribution of untrained and trained encoder’s output on synthetic dataset. (a) Untrained encoder’s output has fewer peaks than 10 peaks of input and shrinks into a relatively small range. (b) Trained encoder’s output displays 10 peaks, which is the same as the input. This phenomena inspire us that pretraining might be an invention for token representation shrinkage.
Model
w/ Deferred?
g-FID ↓
Pixel Dist. ↑
LPIPS D.↑
VAR
✓
29.65
50.56
0.550
✗
27.89
45.14
0.570
Appendix
Table 8: ODIR generation based on VQGAN of VAR. ✓ indicates that token representation shrinkage is mitigated, whereas ✗ indicates without mitigation. Token representation shrinkage impairs generative quality and diversity.
Tokenizer
w/ Deferred?
r-FID ↓
MSE ↓
LPIPS ↓
Euc. ↑
Perp. ↑
MaskGIT
✓
5.28
3.96
2.40
18.75
5575.92
✗
6.40
4.58
2.70
6.52
920.95
VAR
✓
2.13
2.59
1.73
7.50
7143.41
✗
2.09
3.00
1.87
1.13
2814.91
Appendix
Table 9: Evaluation of tokenizers trained with GAN loss on the ImageNet-100 dataset. ✓ indicates that token representation shrinkage is mitigated, whereas ✗ indicates without mitigation. MSE and LPIPS values in the table are scaled. To recover the actual values, multiply them by 10−4 and 10−3 , respectively.
Tokenizer
w/ Deferred?
r-FID ↓
MSE ↓
LPIPS ↓
Euc. ↑
Perp. ↑
VAR
✓
6.90
2.19
7.15
8.50
5451.76
✗
8.12
2.66
9.39
1.27
801.07
Appendix
Table 10: Evaluation of VAR tokenizers trained with GAN loss on the ODIR dataset. ✓ indicates that token representation shrinkage is mitigated, whereas ✗ indicates without mitigation. To recover the actual MSE and LPIPS, multiply them by 10−5 and 10−4 , respectively.
LGAN
w/ Deferred?
r-FID ↓
MSE ↓
LPIPS ↓
Euc. ↑
Perp. ↑
g-FID ↓
Pixel Dist. ↑
LPIPS D. ↑
No
✓
10.33
3.16
0.90
12.70
3346.04
38.04
48.44
0.370
✗
11.96
4.19
1.16
4.92
438.36
27.95
49.73
0.397
Yes
✓
10.47
3.67
1.02
12.60
3211.72
38.30
48.15
0.371
✗
12.72
4.64
1.30
4.93
432.07
32.60
49.28
0.402
Appendix
Table 11: MaskGIT tokenizer and generative performance on medical dataset. LGAN denotes whether GAN loss was used during training. ✓ indicates that token representation shrinkage is mitigated, whereas ✗ indicates without mitigation. To recover the actual MSE and LPIPS values, multiply them by 10−5 and 10−4 .
Vector Quantization (VQ) underpins many modern generative frameworks such as VQ-VAE, VQ-GAN, and latent diffusion models. Yet, it suffers from the persistent problem of codebook collapse, where a large fraction of code vectors remains unused during training. This work provides a new theoretical explanation by identifying the nonstationary nature of encoder updates as the fundamental cause of this phenomenon. We show that as the encoder drifts, unselected code vectors fail to receive updates and gradually become inactive. To address this, we propose two new methods: Non-Stationary Vector Quantization (NSVQ), which propagates encoder drift to non-selected codes through a kernel-based rule, and Transformer-based Vector Quantization (TransVQ), which employs a lightweight mapping to adaptively transform the entire codebook while preserving convergence to the k-means solution. Experiments on the CelebA-HQ dataset demonstrate that both methods achieve near-complete codebook utilization and superior reconstruction quality compared to baseline VQ variants, providing a principled and scalable foundation for future VQ-based generative models. The code is available at: https://github.com/CAIR- LAB- WFUSM/NSVQ-TransVQ.git
Hao Lu, Onur C. Koyun, Yongxin Guo +3
Wake Forest University School of Medicine, Winston-Salem, NC, USA
Effective discrete representation learning with VQ-VAE serves as a fundamental component of modern autoregressive generative models. However, it is widely known to suffer from codebook collapse and inefficient code utilization. To address this issue, we propose FlowVQ, a simple yet effective alternative to standard vector quantization by formulating neural quantization as a joint modeling and clustering process. Specifically, we learn the latent space using a lightweight flow matching model alongside training, corresponding to a probability flow ODE that transports an often fixed well-behaved source distribution while preserving probability mass. Based on this, we define the codebook C_1 by transporting even-quantile partitioned prior codes C_0 from the source space through the learned flow, and obtaining the Voronoi centroids under the induced distribution, leading to a structured discretization. We evaluated on image tokenization and generation, as well as on two tasks requiring high-precision tokenization, demonstrating scalability and improved expressivity of the codebook with full utilization.
Vector quantization is central to modern generative modeling pipelines, but large-codebook VQ models often suffer from codebook collapse. We identify encoder drift as a key driver of this failure: as the encoder moves the latent distribution, sparsely updated code vectors can lag behind, lose assignments, and increase quantization error, creating a feedback loop through the straight-through estimator. We propose NSVQ, a non-stationary-aware VQ training strategy that combines a dense non-stationary embedding loss, codebook replacement, and stage-wise encoder freezing. NSVQ first helps the codebook track encoder drift during early training, then freezes the encoder to consolidate the codebook under a fixed latent geometry, and finally reintroduces adversarial refinement. Experiments on ImageNet-1k show that NSVQ improves reconstruction quality while maintaining full codebook utilization. On ImageNet-1k at 128×128 with 65,536 codes, NSVQ reduces rFID from 2.39 to 2.10 compared with SimVQ, while both methods maintain 100% utilization. Additional latent diffusion experiments show that NSVQ also improves downstream ImageNet generation FID.
Hao Lu, Yongxin Guo, Onur Koyun +3
Wake Forest University School of Medicine / Advocate Health · Wake Forest University School of Medicine · Advocate Health