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

2026-07-17

AI Summary

This work addresses the instability and codebook collapse in vector quantization caused by distributional mismatch between features and the codebook by introducing distribution matching as a central principle and proposing a unified theoretical framework. The approach explicitly aligns the two distributions using either the Wasserstein distance—admitting a closed-form solution under Gaussian approximation—or a non-parametric maximum mean discrepancy (MMD). This alignment significantly enhances codebook utilization and stabilizes training. Experimental results demonstrate that the proposed method substantially outperforms existing approaches on visual tokenization benchmarks, exhibiting strong effectiveness, robustness, and efficient codebook usage.

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.

Problem

• vector quantization
• distributional mismatch
• codebook collapse
• training instability
• feature discretization

Innovation

• distributional matching
• vector quantization
• Wasserstein distance
• maximum mean discrepancy
• codebook collapse