DiVeQ: Differentiable Vector Quantization Using the Reparameterization Trick
About
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 input 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. In VQ-VAE image compression, VQGAN image generation, and DAC speech coding tasks across various data sets, our proposed methods improve reconstruction and sample quality over alternative quantization approaches.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Image Generation | CelebA-HQ | FID5.21 | 92 | |
| Image Generation | FFHQ (test) | FID5.11 | 77 | |
| Image Generation | LSUN Bedroom v1 (test) | FID4.87 | 56 | |
| Image Generation | AFHQ v1 (test) | FID5.4 | 56 | |
| Image Generation | LSUN Church v1 (test) | FID3.56 | 55 | |
| Speech Decompression | VCTK (test) | Log Spectral Distance1.01 | 28 |