Generalized Radius and Integrated Codebook Transforms for Differentiable Vector Quantization

TL;DR

This work addresses instability and codebook under-utilization in vector quantization by proposing GRIT-VQ, a unified framework that combines a radius-based differentiable surrogate with a data-agnostic integrated transform of the codebook. The forward path preserves hard nearest-neighbor quantization, while the backward path enables stable gradient flow through a radius-controlled update and coordinated, shared-code updates. Theoretical analysis characterizes gradient structure, alignment toward the selected codeword, and stability, and experiments across image generation, reconstruction, and recommendation show improved reconstruction quality, generative performance, and codebook utilization, especially at large codebooks and low bitrates. Overall, GRIT-VQ demonstrates scalable, stable discrete representations with practical benefits for tokenization and downstream tasks across vision and recommendation domains.

Abstract

Vector quantization (VQ) underpins modern generative and representation models by turning continuous latents into discrete tokens. Yet hard nearest-neighbor assignments are non-differentiable and are typically optimized with heuristic straight-through estimators, which couple the update step size to the quantization gap and train each code in isolation, leading to unstable gradients and severe codebook under-utilization at scale. In this paper, we introduce GRIT-VQ (Generalized Radius and Integrated Transform-Vector Quantization), a unified surrogate framework that keeps hard assignments in the forward pass while making VQ fully differentiable. GRIT-VQ replaces the straight-through estimator with a radius-based update that moves latents along the quantization direction with a controllable, geometry-aware step, and applies a data-agnostic integrated transform to the codebook so that all codes are updated through shared parameters instead of independently. Our theoretical analysis clarifies the fundamental optimization dynamics introduced by GRIT-VQ, establishing conditions for stable gradient flow, coordinated codebook evolution, and reliable avoidance of collapse across a broad family of quantizers. Across image reconstruction, image generation, and recommendation tokenization benchmarks, GRIT-VQ consistently improves reconstruction error, generative quality, and recommendation accuracy while substantially increasing codebook utilization compared to existing VQ variants.

Figures (11)

• Figure 1: CelebA-HQ FID curves for different tokenizers.
• Figure 2: CelebA-HQ tokenization trade-offs and transformer burden.
• Figure 3: Qualitative comparison of CelebA-HQ samples at bitrate $B{=}9$ for different tokenizers. Each row corresponds to a VQ variant and each column to a shared random seed.
• Figure 4: Code utilization and dead-code rate over training for different VQ tokenizers on the Beauty dataset. GRIT-VQ maintains high utilization and low dead-code rate without collapse, while alternative tokenizers either under-utilize the codebook or exhibit late-stage collapse.
• Figure 5: Top-50 retrieval recall as a function of codebook size $K$ on a held-out retrieval dataset. All models share the same backbone and training data; only the vector quantizer changes. GRIT-VQ consistently achieves the best recall across codebook sizes and remains robust even with small $K$.