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×d )
rFID ↓
LPIPS ↓
Usage ↑
UR-AUC ↑
LlamaGen [ 27 ]
—
40
16,384×8
2.19
0.2281
97%
—
LlamaGen [ 27 ]
—
40
16,384×256
9.21
—
0.29%
IBQ † [ 25 ]
—
330
16,384×256
1.37
0.2235
96%
—
IBQ † [ 25 ]
—
330
262,144×256
1.00
0.2030
84%
Shared-Projection-Based Methods
VQGAN-LC [ 35 ]
Linear-1
20
16,384×8
3.01
0.2358
99%
—
Table 1: Reconstruction results on ImageNet 256×256 with 16×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 ↓
IS ↑
Pre. ↑
Rec. ↑
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×256 .
Region VQ
Dyn. STE
De. Sch.
Peak Commit
Util. ↑
rFID ↓
LPIPS ↓
PSNR ↑
SSIM ↑
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. ↑
rFID ↓
PSNR ↑
SSIM ↑
LPIPS ↓
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×d )
rFID ↓
LPIPS ↓
Usage ↑
VQGAN [ 6 ]
—
16×16
1,024×256
7.94
—
44%
VQGAN [ 6 ]
—
16×16
16,384×256
4.98
0.2843
5.9%
SD-VQGAN [ 23 ]
—
16×16
16,384×8
5.15
—
—
MaskGIT [ 1 ]
—
16×16
1,024×256
2.28
—
—
LlamaGen [ 27 ]
—
16×16
16,384×8
2.19
0.2281
97%
LlamaGen [ 27 ]
—
16×16
16,384×256
9.21
—
0.29%
Appendix
Table 5: Complete reconstruction results on ImageNet 256×256 . Codebook Size is reported as n×d (number of codes × channel dimension). † denotes a larger encoder–decoder trained for up to 330 epochs. ‡ denotes training for 120 epochs.
Method
Codebook Size ( n×d )
Epochs
rFID ↓
PSNR ↑
SSIM ↑
LPIPS ↓
Usage ↑
SimVQ [ 36 ]
16,384×256
40
2.89
21.36
0.5551
0.2492
100%
SimVQ [ 36 ]
262,144×256
40
3.16
21.31
0.5527
0.2516
100%
FVQ [ 2 ]
16,384×256
40
1.70
21.85
0.5894
0.2176
100%
FVQ [ 2 ]
262,144×256
40
1.29
22.50
0.6146
0.2003
100%
FVQ [ 2 ]
16,384×256
120
1.46
21.91
0.5951
0.2156
100%
FVQ [ 2 ]
262,144×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×256 .
Type
Tokenizer
Generator
Param.
FID ↓
IS ↑
Pre. ↑
Rec. ↑
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×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, Δ )
(ms, Δ )
(ms, Δ )
(GiB)
(GiB, Δ )
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 ↑
rFID ↓
PSNR ↑
LPIPS ↓
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
ζ=0
ζ=1
ζ=2
ζ=3
ζ=4
Time (ms)
Relative Time
Wasserstein VQ [ 7 ]
99.9%
97.0%
62.7%
44.8%
34.8%
6.98
1.00×
MMD VQ [ 8 ]
99.9%
99.8%
92.5%
85.7%
75.6%
295.57
42.36×
Region VQ (ours)
99.8%
99.3%
93.2%
99.4%
99.4%
8.64
1.24×
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
β1=0.9,β2=0.95
β1=0.9,β2=0.95
β1=0.9,β2=0.95
Base Learning Rate
1e−4
1e−4
1e−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.
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
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
Vector Quantization (VQ) underpins modern discrete visual tokenization. However, training quantization modules for state-of-the-art VQ-based models requires significant computational resources which, in practice, all but prevents the development of novel, cutting-edge VQ techniques under resource constraints. To address this limitation, we propose {\bf VQ-Transplant}, a simple framework that enables plug-and-play integration of new VQ modules into frozen, pre-trained tokenizers by replacing their native VQ modules. Crucially, the proposed transplantation process preserves all encoder-decoder parameters, obviating the need for costly end-to-end retraining when modifying the quantization method. To mitigate decoder-quantization mismatch, we introduce a lightweight decoder adaptation strategy (trained for only 5 epochs on ImageNet-1k) to align feature priors with the new quantization space. In our empirical evaluation, we find that VQ-Transplant allows obtaining near state-of-the-art reconstruction fidelity for industry-level models like VAR while reducing the training cost by 95%. VQ-Transplant democratizes quantization research by enabling resource-efficient integration of novel VQ techniques while matching industry-level reconstruction performance.