Distributional Matching for Vector Quantization: A Unified Theoretical and Empirical Framework

A unified theoretical and empirical framework that traces training instability and codebook collapse to one cause: a mismatch between feature and code vector distributions.

Xianghong Fang¹Litao Guo²Hengchao Chen¹ Yuxuan Zhang¹Xiaofan Xia¹Dingjie Song⁴ Yexin Liu²Hao Wang⁵Harry Yang² Qiang Sun¹Yuan Yuan³

¹University of Toronto²The Hong Kong University of Science and Technology³Boston College⁴Lehigh University⁵Southern University of Science and Technology

Read paper Code BibTeX

01 Motivation

Two failures. One distributional cause.

Vector quantization relies on the straight-through estimator and a learned codebook. When features and codes occupy different regions of the latent space, assignments concentrate on only a few codes: quantization error grows, gradient estimates become unreliable, and most of the codebook goes unused.

Comparison of distribution mismatch, where features cluster around a few code vectors, and distribution match, where code vectors cover the feature distribution. — **Matching changes the geometry of quantization.** Aligned code vectors cover the feature space, reducing distortion while activating the full codebook.

Quantization error

Average distortion between each feature and its nearest code. Lower error tightens the bound on the STE gradient gap.

Utilization rate

The fraction of code vectors selected by at least one feature. Full utilization is the first defense against collapse.

Codebook perplexity

The effective number of uniformly selected codes. It reveals imbalance that utilization alone cannot capture.

02 Theoretical insight

In high-dimensional latent spaces, the asymptotically optimal codebook distribution approaches the feature distribution.

Support. Matching supports is necessary and sufficient for full utilization and vanishing worst-case quantization error as the codebook grows.

Density. Optimal code density is proportional to f^(d+2)/d. As dimension d grows, the exponent tends to 1, motivating direct distribution alignment.

03 Methodology

Match globally. Quantize locally.

Wasserstein VQ adds a global distribution alignment objective to the local nearest-neighbor VQ loss. Alignment gradients update only the codebook; the encoder remains free to learn expressive feature distributions.

Wasserstein VQ architecture showing an encoder, vector quantizer, decoder, and Wasserstein distribution matching loss between encoder features and codebook vectors. — Wasserstein VQ integrates into standard encoder-quantizer-decoder pipelines as an auxiliary codebook objective.

Distribution matching objective

L_match = D(P_A, P_B)

Here, P_A is the empirical encoder-feature distribution and P_B is the distribution of code vectors.

A practical closed form

Under a mild Gaussian approximation, the quadratic Wasserstein distance aligns the mean and covariance of features and codes. It avoids pairwise comparisons and adds only modest overhead at large codebook sizes.

Wasserstein VQ

Efficient first- and second-order matching for approximately Gaussian latent spaces.

MMD VQ

Nonparametric matching for strongly non-Gaussian features, with higher quadratic cost.

04 Results

Complete codebooks. Sharper reconstructions.

Across FFHQ and ImageNet, distributional matching prevents collapse at unprecedented codebook scales. The same principle transfers to multi-scale VAR, improving quantization and reconstruction under a controlled architecture.

Tables 1 & 2

VQ-VAE at 100,000 codes

Wasserstein VQ is the only method to use the entire codebook on both datasets, while also delivering the best reconstruction metrics.

Method	Utilization ↑	Perplexity ↑	PSNR ↑	SSIM ↑	Rec. loss ↓
Vanilla VQ	0.6%	481.0	27.86	74.2	0.0118
STE++	0.5%	450.7	27.52	72.4	0.0130
EMA VQ	2.7%	2,087.5	28.43	74.8	0.0105
Online VQ	3.6%	1,556.8	27.12	71.1	0.0142
Wasserstein VQ	100%	93,152.7	29.53	78.0	0.0083

FFHQ · 256 tokens · codebook size K = 100,000

Method	Utilization ↑	Perplexity ↑	PSNR ↑	SSIM ↑	Rec. loss ↓
Vanilla VQ	0.4%	337.0	24.43	57.4	0.0295
STE++	0.9%	730.6	24.86	59.1	0.0269
EMA VQ	3.0%	2,170.0	25.13	60.1	0.0257
Online VQ	4.1%	1,709.9	24.95	59.1	0.0267
Wasserstein VQ	100%	93,264.7	25.88	63.0	0.0223

ImageNet-1K · 256 tokens · codebook size K = 100,000

Table 6

Distribution matching strengthens VAR

With the VAR encoder-decoder and latent features held fixed, Wasserstein matching lowers quantization error, increases perplexity, and improves reconstruction fidelity.

Method	Codes	Error ↓	Utilization ↑	Perplexity ↑	r-FID ↓
VAR	4,096	0.283	100%	2,981.4	0.92
Wasserstein VAR-c	4,096	0.255	100%	3,286.2	0.81
Wasserstein VAR-c	8,192	0.240	100%	6,518.2	0.73

0.73r-FID at K = 8,192

Wasserstein VAR-c improves from 0.83 after 5 adaptation epochs to 0.73 after 15.

Original FFHQ images in the top row and reconstructions from Wasserstein VQ in the bottom row. — FFHQ / Wasserstein VQ Original 256 x 256 inputs are shown above; reconstructions are shown below.

Original ImageNet images in the top row and reconstructions from Wasserstein VAR in the bottom row. — ImageNet / Wasserstein VAR Original 256 x 256 inputs are shown above; reconstructions are shown below.

05 Takeaways

Distributional matching turns a fragile lookup table into a faithful map of the latent space.

It unifies optimization stability and codebook collapse under one geometric explanation.

It scales to codebooks with 100,000 entries while maintaining complete utilization.

It is architecture-agnostic, supporting VQ-VAE, VQGAN, and downstream token-based models.

06 Citation

Build on this work.

Please cite the paper if distributional matching helps your research.

@article{fang2026distributional,
  title   = {Distributional Matching for Vector Quantization:
             A Unified Theoretical and Empirical Framework},
  author  = {Fang, Xianghong and Guo, Litao and Chen, Hengchao and
             Zhang, Yuxuan and Xia, Xiaofan and Song, Dingjie and
             Liu, Yexin and Wang, Hao and Yang, Harry and Sun, Qiang and
             Yuan, Yuan},
  journal = {arXiv},
  year    = {2026}
}