Early Quantization Shrinks Codebook: A Simple Fix for Diversity-Preserving Tokenization

Wenhao Zhao; Qiran Zou; Rushi Shah; Yudi Wu; Zhouhan Lin; Dianbo Liu

cs.LGMar 17, 2026

Early Quantization Shrinks Codebook: A Simple Fix for Diversity-Preserving Tokenization

Authors: Wenhao Zhao, Qiran Zou, Rushi Shah, Yudi Wu, Zhouhan Lin, Dianbo Liu

Organizations: National University of Singapore · Zhejiang University · Shanghai Jiao Tong University

Abstract

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.	r-FID $\downarrow$ , LPIPS $\downarrow$ , MSE $\downarrow$ , Perplexity $\uparrow$ , Codebook distance $\uparrow$
C3	Shrinkage is different from codebook collapse; fixing collapse alone is insufficient.	Tab. 2 : Deferred Quantization enables gains where collapse-fixing fails.	r-FID $\downarrow$ , LPIPS $\downarrow$ , MSE $\downarrow$ , Perplexity $\uparrow$ , Codebook distance $\uparrow$
C4	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 $\downarrow$ , Pixel Distance $\uparrow$ , LPIPS Diversity $\uparrow$
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 $\downarrow$ , LPIPS Diversity $\uparrow$ , Pixel Distance $\uparrow$

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
MaskGIT	✗	12.22	3.91	2.70	6.45	924.57
VAR	✓	5.04	2.22	1.19	7.56	7044.51
VAR	✗	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
MaskGIT	✗	14.75	75.89	0.668
VAR	✓	10.44	75.75	0.667
VAR	✗	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
VAR	✗	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
VAR	✗	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
MaskGIT	✗	12.25	78.83	0.714
VAR	✓	8.30	77.10	0.719
VAR	✗	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
VAR	✗	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
MaskGIT	✗	6.40	4.58	2.70	6.52	920.95
VAR	✓	2.13	2.59	1.73	7.50	7143.41
VAR	✗	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
VAR	✗	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.

$\mathcal{L}_{\text{GAN}}$	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
No	✗	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
Yes	✗	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.

\mathcal{L}_{\text{GAN}}

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

Vector quantization compresses vectors into discrete codebook entries for efficient storage, transmission, retrieval, and learned representations.

See all 121 papers

Beyond Stationarity: Rethinking Codebook Collapse in Vector Quantization
Feb 21, 2026Hao Lu, Onur C. Koyun, Yongxin Guo +3Vector Quantization
Wake Forest University School of Medicine, Winston-Salem, NC, USA
NSVQ: Mitigating Codebook Collapse by Stabilizing Encoder Drift in Vector Quantization
Jun 9, 2026Hao Lu, Yongxin Guo, Onur Koyun +3Vector Quantization
Wake Forest University School of Medicine / Advocate Health · Wake Forest University School of Medicine · Advocate Health
VP-VAE: Rethinking Vector Quantization via Adaptive Vector Perturbation
Feb 19, 2026Linwei Zhai, Han Ding, Mingzhi Lin +5Finite Scalar Quantization Vector-Quantized Variational Autoencoder
Xi’an Jiaotong University, China

Feb 21, 2026J/K move · Enter open · S save

Beyond Stationarity: Rethinking Codebook Collapse in Vector Quantization

Hao Lu, Onur C. Koyun, Yongxin Guo, Zhengjie Zhu+2

Wake Forest University School of Medicine, Winston-Salem, NC, USA

Vector Quantization

Open paper AskPDF arXiv

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

Abstract

Vector Quantization

Explore similar work

Beyond Stationarity: Rethinking Codebook Collapse in Vector Quantization

NSVQ: Mitigating Codebook Collapse by Stabilizing Encoder Drift in Vector Quantization

VP-VAE: Rethinking Vector Quantization via Adaptive Vector Perturbation