Generalized Radius and Integrated Codebook Transforms for Differentiable Vector Quantization

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.

💡 Research Summary

The paper addresses a fundamental limitation of vector quantization (VQ) in modern deep generative and representation models: the non‑differentiable hard nearest‑neighbor assignment. Existing approaches rely on Straight‑Through Estimators (STE) that treat the assignment as an identity during back‑propagation. This couples the update magnitude to the quantization gap, forces each code to be updated independently, and often leads to unstable gradients and severe codebook under‑utilization, especially for large vocabularies.

GRIT‑VQ (Generalized Radius and Integrated Transform‑Vector Quantization) proposes a unified surrogate that retains hard assignments in the forward pass while providing a fully differentiable backward path. The surrogate replaces the STE with a radius‑based update: \