DiVeQ: Differentiable Vector Quantization Using the Reparameterization Trick
This addresses the gradient flow issue in vector quantization for deep learning models, enabling better end-to-end training in applications like compression and generation, though it appears incremental as it builds on existing quantization frameworks.
The authors tackled the problem of vector quantization blocking gradients in deep models by proposing DiVeQ, which treats quantization as adding an error vector to allow end-to-end training, and a space-filling variant that reduces quantization error and improves codebook usage. They demonstrated improved reconstruction and sample quality over alternative methods on VQ-VAE compression and VQGAN generation across various datasets.
Vector quantization is common in deep models, yet its hard assignments block gradients and hinder end-to-end training. We propose DiVeQ, which treats quantization as adding an error vector that mimics the quantization distortion, keeping the forward pass hard while letting gradients flow. We also present a space-filling variant (SF-DiVeQ) that assigns to a curve constructed by the lines connecting codewords, resulting in less quantization error and full codebook usage. Both methods train end-to-end without requiring auxiliary losses or temperature schedules. On VQ-VAE compression and VQGAN generation across various data sets, they improve reconstruction and sample quality over alternative quantization approaches.