Distributional Matching for Vector Quantization: A Unified Theoretical and Empirical Framework
Computer Vision and Pattern Recognition
Summary
The authors looked at a key step in some visual AI models called vector quantization, which turns continuous data into discrete code vectors. They found that problems like unstable training and codebook collapse happen because the feature data and code vectors don't match well in distribution. To fix this, they proposed a new method that aligns these distributions using mathematical tools like the Wasserstein distance, improving stability and performance. Their tests showed this approach works well on visual tokenization tasks.
Authors
Xianghong Fang, Litao Guo, Hengchao Chen, Yuxuan Zhang, XiaofanXia, Dingjie Song, Yexin Liu, Hao Wang, Harry Yang, Qiang Sun, Yuan Yuan
Abstract
The effectiveness of modern visual representation learning and autoregressive models critically depends on vector quantization (VQ), which discretizes continuous feature representations using a learnable codebook. Despite its widespread use, existing VQ methods often suffer from training instability and codebook collapse, arising from gradient mismatch induced by the straight-through estimator and the under-utilization of code vectors. In this work, we show that both issues can be traced to a fundamental mismatch between the distributions of feature vectors and code vectors, leading to inefficient representation and information loss. Building on this observation, we propose a distributional matching framework for vector quantization. We introduce principled criteria for desirable VQ behavior and demonstrate through theoretical analysis and empirical evaluation that aligning feature and code vector distributions provides a unifying mechanism for mitigating training instability and codebook collapse. We instantiate this framework using a Wasserstein-based objective with an efficient closed-form under a mild Gaussian approximation, and further show that a nonparametric alternative based on maximum mean discrepancy yields comparable performance. Extensive experiments on visual tokenization benchmarks support the effectiveness and robustness of the proposed approach.