StableVQ: Practical Guidelines for Stable Vector-Quantized Tokenizer Training

Bao Tang; Jiahao Guo; Haoxiang Cao; Wenyu Liu; Changqian Yu; Kun Gai; Xinggang Wang

cs.CVSep 22, 2026

StableVQ: Practical Guidelines for Stable Vector-Quantized Tokenizer Training

Authors: Bao Tang, Jiahao Guo, Haoxiang Cao, Wenyu Liu, Changqian Yu, Kun Gai, Xinggang Wang

Organizations: Huazhong University of Science and Technology · KlingAI Research · South China Normal University

Abstract

Vector Quantization (VQ) is fundamental to discrete visual tokenizers that power modern autoregressive and masked image generation models. While recent shared-projection codebook methods have substantially advanced codebook utilization, training stability remains a critical and underexplored challenge. We argue that the root cause lies in the entanglement of the Encoder--Decoder and Codebook training: because neither module can reliably fulfill its own responsibility in isolation, the system can only function when the two subsystems happen to cooperate---a fragile condition that breaks down precisely when training is most stressed. We propose StableVQ, which revisits the proper learning objective of each module and resolves the problems that arise when each is trained to fulfill its own role independently. Concretely, (1) Dynamic STE corrects the instability in the Encoder's learning objective, enabling it to robustly optimize the reconstruction space under discrete regularization even when codebook utilization is low. (2) Region VQ Loss reconceives the Codebook's learning objective so that it can independently guarantee full tracking of the encoder output distribution, without relying on encoder oscillations to drive activation. (3) Decoupled Schedule recognizes that the distinct responsibilities of the Encoder--Decoder and the Codebook demand distinct optimization dynamics, and assigns each an independent learning rate schedule to ensure robust system-level behavior. Built on top of shared-projection codebooks, StableVQ is lightweight and introduces no learnable parameters. Experiments on ImageNet demonstrate consistent improvements in training stability, codebook utilization, and reconstruction quality across diverse codebook sizes and initialization settings.

Figures & tables

Figure 1: Three characteristic failure modes of VQ training, illustrated through the relationship between the token distribution (blue) and code distribution (red) at different training stages.

Figure 2: Left: Dynamic STE design. Relatively farther tokens have more unreliable gradients and receive proportionally suppressed contributions. Right: Pilot Study. Standard STE causes commitment loss spikes and training collapse, whereas Dynamic STE maintains stable losses.

Figure 3: Left: Region VQ Loss. Active codes propagate targets to nearby inactive codes proportionally to selection count. Right: Pilot Study. T-SNE visualizations at different training steps (blue: encoder outputs; red: codebook entries) with codebook utilization.

Figure 4: Pilot study on learning rate schedules for different modules. WU-AN denotes Warmup-plus-Annealing; C.B. denotes Codebook.

Method	Projector	Epochs	Codebook Size ( $n\times d$ )	rFID $\downarrow$	LPIPS $\downarrow$	Usage $\uparrow$	UR-AUC $\uparrow$
LlamaGen [ 27 ]	—	40	$16{,}384\times 8$	2.19	0.2281	97%	—
LlamaGen [ 27 ]	—	40	$16{,}384\times 256$	9.21	—	0.29%	—
IBQ † [ 25 ]	—	330	$16{,}384\times 256$	1.37	0.2235	96%	—
IBQ † [ 25 ]	—	330	$262{,}144\times 256$	1.00	0.2030	84%	—
Shared-Projection-Based Methods
VQGAN-LC [ 35 ]	Linear-1	20	$16{,}384\times 8$	3.01	0.2358	99%	—

Table 1: Reconstruction results on ImageNet

256\times 256

with

16\times 16

tokens. † denotes a larger encoder–decoder. UR-AUC denotes Usage Recovery AUC under the robustness test.

Figure 5: Codebook usage recovery under different codebook-token distribution relationships.

Type	Tokenizer	Generator	Param.	FID $\downarrow$	IS $\uparrow$	Pre. $\uparrow$	Rec. $\uparrow$
Vanilla AR	IBQ [ 25 ]	IBQ-B [ 25 ]	342M	2.88	254.7	0.84	0.51
Vanilla AR	IBQ [ 25 ]	IBQ-L [ 25 ]	649M	2.45	267.5	0.83	0.52
Vanilla AR	StableVQ	IBQ-B [ 25 ]	342M	2.35	256.0	0.82	0.58
Vanilla AR	StableVQ	IBQ-L [ 25 ]	649M	2.18	250.4	0.82	0.59

Table 2: Class-conditional image generation on ImageNet

256\times 256

Region VQ	Dyn. STE	De. Sch.	Peak Commit	Util. $\uparrow$	rFID $\downarrow$	LPIPS $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$
			53.11	49.13%	2.06	0.2297	21.65	0.5730
✓			$>$ 200 (NaN)	—	—	—	—	—
	✓		$<$ 0.1	1.27%	6.68	0.3100	19.89	0.5026
		✓	$<$ 0.1	83.77%	1.71	0.2209	21.79	0.5864
	✓	✓	$<$ 0.1	21.66%	2.07	0.2354	21.36	0.5665
✓	✓		$<$ 0.1	100%	1.72	0.2200	21.84	0.5854

Table 3: Ablation on proposed strategies under the codebook expansion setting.

Projector	Region VQ	Init	Util. $\uparrow$	rFID $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
ViTBlock-2 [ 2 ]		uniform	18.75%	2.4355	20.9789	0.5677	0.2414
ViTBlock-2 [ 2 ]		gaussian	62.5%	1.9530	21.3774	0.5836	0.2273
ViTBlock-2 [ 2 ]	✓	uniform	100%	1.8176	21.5497	0.5896	0.2210
ViTBlock-2 [ 2 ]	✓	gaussian	100%	1.8966	21.4603	0.5931	0.2255

Table 4: Ablation on methods under the codebook shrinkage setting.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Projector	Tokens	Codebook Size ( $n\times d$ )	rFID $\downarrow$	LPIPS $\downarrow$	Usage $\uparrow$
VQGAN [ 6 ]	—	$16\times 16$	$1{,}024\times 256$	7.94	—	44%
VQGAN [ 6 ]	—	$16\times 16$	$16{,}384\times 256$	4.98	0.2843	5.9%
SD-VQGAN [ 23 ]	—	$16\times 16$	$16{,}384\times 8$	5.15	—	—
MaskGIT [ 1 ]	—	$16\times 16$	$1{,}024\times 256$	2.28	—	—
LlamaGen [ 27 ]	—	$16\times 16$	$16{,}384\times 8$	2.19	0.2281	97%
LlamaGen [ 27 ]	—	$16\times 16$	$16{,}384\times 256$	9.21	—	0.29%

Appendix

Table 5: Complete reconstruction results on ImageNet

256\times 256

. Codebook Size is reported as

n\times d

(number of codes

\times

channel dimension). † denotes a larger encoder–decoder trained for up to 330 epochs. ‡ denotes training for 120 epochs.

Method	Codebook Size ( $n\times d$ )	Epochs	rFID $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	Usage $\uparrow$
SimVQ [ 36 ]	$16{,}384\times 256$	40	2.89	21.36	0.5551	0.2492	100%
SimVQ [ 36 ]	$262{,}144\times 256$	40	3.16	21.31	0.5527	0.2516	100%
FVQ [ 2 ]	$16{,}384\times 256$	40	1.70	21.85	0.5894	0.2176	100%
FVQ [ 2 ]	$262{,}144\times 256$	40	1.29	22.50	0.6146	0.2003	100%
FVQ [ 2 ]	$16{,}384\times 256$	120	1.46	21.91	0.5951	0.2156	100%
FVQ [ 2 ]	$262{,}144\times 256$	120	1.07	22.55	0.6244	0.1965	100%

Appendix

Table 6: Additional reconstruction metrics under the same evaluation script on ImageNet

256\times 256

Type	Tokenizer	Generator	Param.	FID $\downarrow$	IS $\uparrow$	Pre. $\uparrow$	Rec. $\uparrow$
Diff.	SD-VAE [ 23 ]	DiT-L/2 [ 20 ]	458M	5.02	167.2	0.75	0.57
Diff.	SD-VAE [ 23 ]	DiT-XL/2 [ 20 ]	675M	2.27	278.2	0.83	0.57
Mask.	VQGAN [ 6 ]	MaskGIT [ 1 ]	227M	6.18	182.1	0.80	0.51
VAR	VAR [ 28 ]	VAR-d16 [ 28 ]	310M	3.30	274.4	0.84	0.51
VAR	VAR [ 28 ]	VAR-d20 [ 28 ]	600M	2.57	302.6	0.83	0.56
AR	LlamaGen [ 27 ]	LlamaGen-L [ 27 ]	343M	3.80	248.3	0.83	0.51

Appendix

Table 7: Class-conditional image generation results on ImageNet

256\times 256

Figure 6: Additional analysis experiments. (a) Usage@5k under different Gaussian codebook initialization scales. (b) Usage@5k under different learning rates.

	Training Cost					Memory Usage
Codebook	Baseline	Dynamic STE	Region VQ @30%	Region VQ @50%	Region VQ @90%	Baseline	Region VQ
	(ms)	(ms)	(ms, $\Delta$ )	(ms, $\Delta$ )	(ms, $\Delta$ )	(GiB)	(GiB, $\Delta$ )
16K	438.1	437.2	441.2 ( $+0.71\%$ )	440.4 ( $+0.51\%$ )	440.2 ( $+0.46\%$ )	34.06	34.07 ( $+0.04\%$ )
262K	462.2	462.7	472.2 ( $+2.16\%$ )	469.2 ( $+1.50\%$ )	466.2 ( $+0.85\%$ )	35.12	35.55 ( $+1.24\%$ )

Appendix

Table 8: Training-time and memory overhead of StableVQ components. Parentheses denote changes from the corresponding baseline, and all memory values are per device.

Method	Training Status	Usage $\uparrow$	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\downarrow$
VQ-STE++ [ 11 ]	Collapsed	–	–	–	–
Rotation Trick [ 9 ]	Completed	100.00%	3.6437	20.4024	0.2568
Dynamic STE (ours)	Completed	100.00%	2.7942	21.1882	0.2359

Appendix

Table 9: Comparison of STE corrections in training stability and reconstruction quality.

Method	$\zeta=0$	$\zeta=1$	$\zeta=2$	$\zeta=3$	$\zeta=4$	Time (ms)	Relative Time
Wasserstein VQ [ 7 ]	99.9%	97.0%	62.7%	44.8%	34.8%	6.98	$1.00\times$
MMD VQ [ 8 ]	99.9%	99.8%	92.5%	85.7%	75.6%	295.57	$42.36\times$
Region VQ (ours)	99.8%	99.3%	93.2%	99.4%	99.4%	8.64	$1.24\times$

Appendix

Table 10: Codebook utilization and training efficiency on the non-Gaussian distribution-fitting benchmark. Utilization is measured at 10K steps. Relative time is normalized to Wasserstein VQ.

Config	SimVQ	FVQ	StableVQ
Base Batch Size	128	128	128
Training Epochs	40	40 / 120	40 / 120
Optimizer	Adam	Adam	Adam
Optimizer Parameters	$\beta_{1}=0.9,\beta_{2}=0.95$	$\beta_{1}=0.9,\beta_{2}=0.95$	$\beta_{1}=0.9,\beta_{2}=0.95$
Base Learning Rate	$1\mathrm{e}{-4}$	$1\mathrm{e}{-4}$	$1\mathrm{e}{-4}$
Codebook Init Type	Gaussian	Uniform	Uniform

Appendix

Table 11: Main experiment configurations.

Figure 7: Qualitative reconstruction comparison with a 16k codebook.

Figure 8: Qualitative reconstruction comparison with a 262k codebook.

Figure 9: Qualitative samples from the downstream class-conditional image generation task.

Straight-Through Estimator

See all 11 papers

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
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
VQ-Transplant: Efficient VQ-Module Integration for Pre-trained Visual Tokenizers
Jul 21, 2026Xianghong Fang, Yuan Yuan, Dehan Kong +1Visual Tokenization Image Reconstruction
University of Toronto · 2Boston College · 3Vijil

Jun 9, 2026J/K move · Enter open · S save

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

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

Wake Forest University School of Medicine / Advocate Health · Wake Forest University School of Medicine · Advocate Health

Vector Quantization

Open paper AskPDF arXiv

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

\times

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.

Abstract

Straight-Through Estimator

Explore similar work

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

Beyond Stationarity: Rethinking Codebook Collapse in Vector Quantization

VQ-Transplant: Efficient VQ-Module Integration for Pre-trained Visual Tokenizers