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
Figures & tables
Figure 1: Illustration of codebook adaptation under non-stationary data. (a) In vanilla VQ, codewords fail to track the drifting data distribution, leading to representation collapse. (b) In the proposed NS-VQ, adaptive variance-controlled updates allow the codebook to follow distributional shifts over time, maintaining coverage and stability. (c) In the proposed TransVQ, a few codewords still lag behind the drifting data distribution, but the projector gradients drive all other codewords to move jointly toward the data, improving overall alignment. In the figures, purple dots denote the target distribution Yt (least visible), green dots represent the base data X , blue dots indicate the current batch Xb , and red crosses mark the codebook vectors C .
Figure 2: Scheme of the proposed (a) Non-Stationary Vector Quantization (NS-VQ) and (b) the Transformer-based Vector Quantization (TransVQ).
Encoder
Decoder
Conv2d (3 → 128, k =4, s =2, p =1), LeakyReLU
Conv2d (64 → 256, k =3, s =1, p =1), LeakyReLU
Conv2d (128 → 256, k =4, s =2, p =1), LeakyReLU
ResidualLayer × 6 (256 → 256)
Conv2d (256 → 256, k =3, s =1, p =1), LeakyReLU
LeakyReLU
ResidualLayer × 6 (256 → 256)
ConvTranspose2d (256 → 128, k =4, s =2, p =1), LeakyReLU
LeakyReLU
ConvTranspose2d (128 → 3, k =4, s =2, p =1), Tanh
Conv2d (256 → 64, k =1, s =1), LeakyReLU
–
Table 1: Encoder and decoder architecture of the proposed model.
Figure 3: Comparison of the proposed NS-VQVAE and TransVQVAE with VQGAN-FC [ 3 ] under varying codebook sizes. (a) rFID comparison showing both NS-VQVAE and TransVQVAE consistently reduce reconstruction error compared with VQGAN-FC [ 3 ] . (b) Codebook utilization indicating that both proposed methods maintain nearly full codebook usage, effectively preventing codebook collapse.
Method
64
1024
8912
rFID
Usage
rFID
Usage
rFID
Usage
VQGAN-FC [ 3 ]
24.10
1.00
18.48
1.00
17.57
0.34
VQVAE2 [ 18 ]
53.25
1.00
31.17
0.40
24.24
0.60
VQGAN-EMA [ 19 ]
60.86
1.00
17.86
1.00
17.93
1.00
SimVQ [ 19 ]
23.39
1.00
23.35
1.00
21.69
1.00
NS-VQ
24.55
1.00
17.86
1.00
16.89
1.00
Table 2: Comparison of different VQ methods with 3D code dimension under varying codebook sizes. Results are reported in terms of reconstruction FID (rFID, lower is better) and codebook usage (higher is better).
Method
64
1024
8912
rFID
Usage
rFID
Usage
rFID
Usage
VQGAN-FC [ 3 ]
24.41
0.66
25.77
0.03
26.00
0.00
VQVAE2 [ 18 ]
85.86
0.82
36.91
0.57
25.87
0.28
VQGAN-EMA [ 19 ]
159.21
1.00
20.81
1.00
19.39
1.00
SimVQ [ 13 ]
23.40
1.00
17.77
1.00
14.37
1.00
NS-VQ
23.39
1.00
17.44
1.00
14.01
1.00
Table 3: Comparison of different VQ methods with 64D code dimension under varying codebook sizes. Results are reported in terms of reconstruction FID (rFID, lower is better) and codebook usage (higher is better).
Method
Codebook Size
Code Dim
SSIM
LPIPS
rFID
VQGAN-FC [ 3 ]
8912
3
0.93
0.022
17.57
VQVAE2 [ 18 ] (3D)
8912
3
0.89
0.042
24.24
VQGAN-EMA [ 19 ]
1024
3
0.92
0.024
17.86
SimVQ [ 13 ]
8912
64
0.93
0.017
14.37
NS-VQ
8912
64
0.96
0.015
14.01
TransVQ
8912
64
0.94
0.015
13.70
Table 4: Detailed comparison of VQ methods at their best-performing configurations (as determined by rFID in Table 2 and Table 3 ). Results are reported in terms of SSIM (higher is better), LPIPS (lower is better), and rFID (lower is better).
Figure 4: rFID curves of standard VQ-VAE with codebook size and code dimension fixed at 64, evaluated under different batch sizes. Larger batch sizes lead to lower rFID values, consistent with our theoretical analysis that larger batches provide more stable codebook updates.
Table 5: VQ-VAE architecture (Taming-style, attention-free). The encoder and decoder are symmetric with residual connections.
Method
rFID
LPIPS
VQGAN [ 6 ]
1.060
0.113
VQGAN-FC [ 3 ]
0.600
0.087
VQGAN-EMA [ 18 ]
1.270
0.124
SimVQ [ 13 ]
0.551
0.088
NS-VQ (ours)
0.561
0.087
TransVQ (ours)
0.432
0.067
Table 6: Quantitative comparison on ImageNet-1K at 256×256 resolution. Lower rFID and LPIPS indicate better reconstruction quality and perceptual fidelity.
Figure 5: Qualitative reconstruction comparison. Each strip shows the same set of identities. NS-VQ and TransVQ maintain sharp edges and facial micro-structure while reducing artifacts observed in vanilla and EMA-VQ. This aligns with the quantitative gains in rFID/LPIPS/SSIM reported in the main paper.
Figure 6: Toy Demo Visualization for Vanilla VQ. Each row corresponds to one toy experiment: translation (1st row), shrink (2nd row), split (3rd row), and expansion (4th row). Each row contains five snapshots showing the convergence process across training steps.
Figure 7: Toy Demo Visualization for NS-VQ. Each row corresponds to one toy experiment: translation (1st row), shrink (2nd row), split (3rd row), and expansion (4th row). Each row contains five snapshots showing how NS-VQ maintains stability and codebook utilization during convergence.