Enhancing Vector Quantization with Distributional Matching: A Theoretical and Empirical Study

TL;DR

This paper introduces a distributional matching approach using Wasserstein distance to improve vector quantization in autoregressive models, addressing training instability and codebook collapse, with validated empirical and theoretical results.

Contribution

It proposes a novel distributional matching method with Wasserstein distance to enhance codebook utilization and reduce quantization errors in vector quantization.

Findings

Near 100% codebook utilization achieved.

Significant reduction in quantization error.

Validated effectiveness through empirical and theoretical analysis.

Abstract

The success of autoregressive models largely depends on the effectiveness of vector quantization, a technique that discretizes continuous features by mapping them to the nearest code vectors within a learnable codebook. Two critical issues in existing vector quantization methods are training instability and codebook collapse. Training instability arises from the gradient discrepancy introduced by the straight-through estimator, especially in the presence of significant quantization errors, while codebook collapse occurs when only a small subset of code vectors are utilized during training. A closer examination of these issues reveals that they are primarily driven by a mismatch between the distributions of the features and code vectors, leading to unrepresentative code vectors and significant data information loss during compression. To address this, we employ the Wasserstein distance to align these two distributions, achieving near 100\% codebook utilization and significantly reducing the quantization error. Both empirical and theoretical analyses validate the effectiveness of the proposed approach.