Transporting Quantiles to Codebook: Scalable Vector Quantization without Codebook Collapse

Yuwei Zeng; Zekun Shi

cs.LGOct 26, 2026

Transporting Quantiles to Codebook: Scalable Vector Quantization without Codebook Collapse

Authors: Yuwei Zeng, Zekun Shi

Organizations: National University of Singapore · SEA AI Lab

Abstract

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.

Figures & tables

Figure 1: FlowVQ on a three-mode 2D latent distribution (K = 50): the learned flow $\psi^{\gamma}_{1}$ transports the approximately equal-mass quantile prior codebook $\mathcal{C}_{0}\in\mathcal{N}(\mathbf{0},I)$ , initialized using Sobol samples here and kept fixed thereafter, to the latent space. Self-correction refines each transported code using the pushed quantile samples $\mathcal{S}_{1}$ , yielding the density-adaptive posterior codebook $\mathcal{C}_{1}$ .

	d	$\|\mathcal{C}\|$	rFID $\downarrow$	IS $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	LPIPS $\downarrow$	MAE $\downarrow$	Entropy $\uparrow$	Util $\%\uparrow$
VQ-GAN+	10	1024	2.40	10.77	0.78	23.62	0.130	0.050	8.81	45.6
	10	10240	2.65	10.70	0.75	23.27	0.143	0.052	8.07	2.7
w/o $\mathcal{L}_{ent}$	14	16384	1.93	10.88	0.81	24.79	0.115	0.044	8.71	2.6
	14	16384	1.65	10.87	0.81	24.82	0.114	0.044	8.57	2.8
LFQ	10	1024	3.00	10.65	0.74	23.01	0.149	0.058	8.46	50.0
w/o $\mathcal{L}_{ent}$	14	16384	2.95	10.83	0.73	22.57	0.161	0.056	5.91	0.8

Table 1: Evaluation of image quality and codebook utilization on CIFAR-10. By default, all baselines optimize the auxiliary entropy loss while FlowVQ does not if not specified. As the entropy maximization is tightly coupled in the OptVQ optimization process, we did not explore removing it.

Figure 2: Visualization of (left) the latent vectors on the evaluation set

\bullet

; (middle) the codebook

\blacktriangle

; (right) the joint plot of latent vectors

\bullet

and codebook vectors

\blacktriangle

in each plot. For each model, the latent vectors and code vectors are randomly down-sampled to the same number, then jointly reduced to 3D using t-SNE for visualization.

Figure 3: class-conditioned generation on ImageNet-1K

256\times 256

trained model. Each image corresponds to a different class. More samples can be found in Appendix C.3 .

Model	Tokenizer	$\|\mathcal{C}\|$	#tokens	#params	rFID $\downarrow$	gFID $\downarrow$	IS $\uparrow$
ImageNet-1K val. data [ 21 ]						1.78	236.9
LDM-4-G [ 23 ]	Continuous ∗	8192	$64^{2}$	400M	-	3.60	247.7
LDM-8-G [ 23 ]	Continuous ∗	16384	$32^{2}$	506M	-	7.76	209.5
MAR-L-407M [ 38 ]	Continuous	-	$16^{2}$	407M	1.22	1.97	291.2
FSQ [ 11 ]	LFQ	1000	$16^{2}$	155M	-	4.53	-
MAGVIT-v2 [ 12 ]	LFQ	262144	$16^{2}$	307M	-	1.78	319.4

Table 2: Image reconstruction and generation on ImageNet-1K

256\times 256

with classifier free guidance. We compare FlowVQ with the existing works with similar model capacity and their guided or rejection sampling scores. * LDM reconstructs on continuous space, but generation discretizes input to the nearest token.

Figure 4: Reconstruction on unseen images from COCO-2017 validation data on ImageNet-1K

256\times 256

trained model.

	VQVAE	LFQ	OptVQ ( $\epsilon=10^{*}$ )	KM-VQVAE	KM-OptVQ ( $\epsilon=10$ )	KM-OptVQ ( $\epsilon=1$ )	FlowVQ	FlowVQ (w/ oc)
Match Rate $\uparrow$	47.16 %	24.63%	10.81%	78.38%	17.85%	53.27%	90.81 %	90.66 %

Table 3: Match rate between the reconstructed and ground-truth material on MP-20.

	FAST+	KM-VQVAE	FlowVQ
MSE(e-4) $\downarrow$	2.40	1.88	1.71
Utilization $\uparrow$	97.02%	10.89%	100%

Table 4: Trajectory reconstruction error of 50-step action chunks on DROID dataset.

	FID $\downarrow$	IS $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	LPIPS $\downarrow$	MAE $\downarrow$	Entropy $\uparrow$	Util $\uparrow$
$\mathcal{C}_{0}\rightarrow\mathcal{C}_{1}$	1.49	10.79	0.856	26.11	0.086	0.038	13.92	100%
$\mathcal{C}_{0}\rightarrow\mathcal{C}_{1}$ + sc	1.52	10.70	0.841	25.73	0.094	0.039	13.97	100%
$\mathcal{C}_{0}\rightarrow\mathcal{C}_{1}$ + sc + oc	1.40	10.72	0.861	26.50	0.086	0.036	13.97	100%

Table 5: Ablation study on FlowVQ components on CIFAR-10 image tokenization and MP-20 material tokenization tasks. With details in Sec 3.3 , sc refers to self-correction mechanism, and oc refers to the optional online-correction mechanism.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

\{\mathbf{C}_{0}^{i}\}_{i=1}^{K}=\mathcal{F}^{-1}(\textit{Sobol}(K,d)),\quad\{\mathbf{S}_{0}^{i}\}_{i=1}^{NK}=\mathcal{F}^{-1}(\textit{Sobol}(NK,d))

Appendix

Algorithm 1 Flow-based Vector Quantization (FlowVQ)

	VQ-GAN	LFQ	OptVQ	FlowVQ w/ sc
t (s) $\downarrow$	$535.2\pm 0.4$	$578.2\pm 0.4$	$596.4\pm 0.5$	$618.0\pm 13.5$

Appendix

Table 6: Runtime per 1,000 steps (one codebook update cycle) on a single NVIDIA A5000 GPU, with no other processes running and

|\mathcal{C}|=16384

Config	CIFAR-10 Reconstruction	ImageNet-1k Reconstruction
input size	$32\times 32\times 3$	$256\times 256\times 3$
latent size	$8\times 8\times d$	$16\times 16\times d$
batch size	128	256
Data Augmentation	False	True
random cropping		[0.8, 1.0]
random aspect ratio		[ $\dfrac{3}{4}$ , $\dfrac{4}{3}$ ]

Appendix

Table 7: Hyperparameters for image reconstruction.

Config
#Tokens	$16\times 16$
batch size	1024
Data Augmentation
random cropping	[0.8, 1.0]
resize to input size	True
random horizontal flip	True

Appendix

Table 8: Hyperparameters for masked language model on image generation.

Figure 9: king penguin, Aptenodytes patagonica

Figure 10: lion, king of beasts, Panthera leo

Vector Quantization

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

See all 122 papers

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
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

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

VP-VAE: Rethinking Vector Quantization via Adaptive Vector Perturbation

Linwei Zhai, Han Ding, Mingzhi Lin, Cui Zhao+4

Xi’an Jiaotong University, China

Finite Scalar Quantization Vector-Quantized Variational Autoencoder Vector Quantization

Open paper AskPDF arXiv

Vector Quantized Variational Autoencoders (VQ-VAEs) are fundamental to modern generative modeling, yet they often suffer from training instability and "codebook collapse" due to the inherent coupling of representation learning and discrete codebook optimization. In this paper, we propose VP-VAE (Vector Perturbation VAE), a novel paradigm that decouples representation learning from discretization by eliminating the need for an explicit codebook during training. Our key insight is that, from the neural network's viewpoint, performing quantization primarily manifests as injecting a structured perturbation in latent space. Accordingly, VP-VAE replaces the non-differentiable quantizer with distribution-consistent and scale-adaptive latent perturbations generated via Metropolis--Hastings sampling. This design enables stable training without a codebook while making the model robust to inference-time quantization error. Moreover, under the assumption of approximately uniform latent variables, we derive FSP (Finite Scalar Perturbation), a lightweight variant of VP-VAE that provides a unified theoretical explanation and a practical improvement for FSQ-style fixed quantizers. Extensive experiments on image and audio benchmarks demonstrate that VP-VAE and FSP improve reconstruction fidelity and achieve substantially more balanced token usage, while avoiding the instability inherent to coupled codebook training.

Abstract

Vector Quantization

Explore similar work

VP-VAE: Rethinking Vector Quantization via Adaptive Vector Perturbation

Beyond Stationarity: Rethinking Codebook Collapse in Vector Quantization

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