# Generalized Radius and Integrated Codebook Transforms for Differentiable Vector Quantization

arXiv: https://arxiv.org/abs/2602.01140 · HTML: https://mlsanity.com/paper/2602.01140/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 54 of its 72 references that are in the map.
- First posted on arXiv 1 Feb 2026.

## Builds on (most cited first, 30 of 54)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [A Style-Based Generator Architecture for Generative Adversarial Networks](https://mlsanity.com/paper/1812.04948/) (2018)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [Progressive Growing of GANs for Improved Quality, Stability, and Variation](https://mlsanity.com/paper/1710.10196/) (2017)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [LSUN: Construction of a Large-scale Image Dataset using Deep Learning with Humans in the Loop](https://mlsanity.com/paper/1506.03365/) (2015)
- [Demystifying MMD GANs](https://mlsanity.com/paper/1801.01401/) (2018)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [StarGAN v2: Diverse Image Synthesis for Multiple Domains](https://mlsanity.com/paper/1912.01865/) (2019)
- [Autoregressive Image Generation using Residual Quantization](https://mlsanity.com/paper/2203.01941/) (2022)
- [End-to-end Optimized Image Compression](https://mlsanity.com/paper/1611.01704/) (2016)
- [Language Model Beats Diffusion -- Tokenizer is Key to Visual Generation](https://mlsanity.com/paper/2310.05737/) (2023)
- [BinaryConnect: Training Deep Neural Networks with binary weights during propagations](https://mlsanity.com/paper/1511.00363/) (2015)
- [Jukebox: A Generative Model for Music](https://mlsanity.com/paper/2005.00341/) (2020)
- [Finite Scalar Quantization: VQ-VAE Made Simple](https://mlsanity.com/paper/2309.15505/) (2023)
- [An Image is Worth 32 Tokens for Reconstruction and Generation](https://mlsanity.com/paper/2406.07550/) (2024)
- [Recommender Systems with Generative Retrieval](https://mlsanity.com/paper/2305.05065/) (2023)
- [Soft-to-Hard Vector Quantization for End-to-End Learning Compressible Representations](https://mlsanity.com/paper/1704.00648/) (2017)
- [MoVQ: Modulating Quantized Vectors for High-Fidelity Image Generation](https://mlsanity.com/paper/2209.09002/) (2022)
- [Neural Variational Inference and Learning in Belief Networks](https://mlsanity.com/paper/1402.0030/) (2014)
- [Learning Vector-Quantized Item Representation for Transferable Sequential Recommenders](https://mlsanity.com/paper/2210.12316/) (2022)
- [Straightening Out the Straight-Through Estimator: Overcoming Optimization Challenges in Vector Quantized Networks](https://mlsanity.com/paper/2305.08842/) (2023)
- [Addressing Representation Collapse in Vector Quantized Models with One Linear Layer](https://mlsanity.com/paper/2411.02038/) (2024)
- [Shopping Queries Dataset: A Large-Scale ESCI Benchmark for Improving Product Search](https://mlsanity.com/paper/2206.06588/) (2022)

## Abstract (arXiv)

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.