Research map / Image generation / Autoregressive and latent generation
Learning to Quantize for Training Vector-Quantized Networks
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Meta-Quantization (MQ), a bi-level meta-learning framework that trains VQ codebooks via a hyper-network (MLP) generating codebook entries from a learnable embedding, with a meta-objective giving task-aware gradients to the codebook, evaluated with VQVAE and VQGAN on reconstruction and generation.
Problem
Vector-quantized networks suffer from index collapse (low codebook utilization) and codebook training that is agnostic to the task loss because the straight-through estimator bypasses the codebook.
Approach
Reparameterize the codebook with a hyper-network that generates all codes from a learnable embedding, and train the codebook via a one-step-unrolled bi-level meta-objective so the task loss directly shapes the codebook through hypergradients, replacing the standard VQ codebook loss.
Key results
- On CelebA, MQVAE reaches MSE 3.10e-3 and LPIPS 0.14 vs 10.02e-3 and 0.27 for VQVAE, and MaskGIT generation FID 70.5 vs 90.4 baseline.
- On ImageNet-1K 128x128, MQGAN achieves rFID 2.13 with 100% codebook utilization, beating SimVQ (2.24) and VQGAN-LC (2.40).
- On FFHQ, MQGAN gets rFID 4.25 and generation FID 2.87 with 99.9% utilization, outperforming VQGAN-LC (3.05 FID).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use Meta-Quantization instead of standard VQVAE | VQVAE · CIFAR10 | MSE (10−3): 3.36 -> 3.05 | Helped |
| use Meta-Quantization instead of standard VQVAE | VQVAE · CelebA | MSE (10−3): 4.77 -> 3.10 | Helped |
| use Meta-Quantization instead of Affine + OPT | MaskGIT on CelebA · CelebA | FID: 74.8 -> 70.5 | Helped |
| use Meta-Quantization instead of VQGAN-LC | VQGAN · FFHQ | rFID: 4.65 -> 4.25 | Helped |
| use Meta-Quantization instead of SimVQ | VQGAN · ImageNet-1k | rFID: 2.24 -> 2.13 | Helped |
| use Meta-Quantization instead of VQGAN-LC | VQGAN generation · FFHQ | FID: 3.05 -> 2.87 | Helped |
| add bi-level optimization to hyper-net reparameterization (ablation) | MQVAE · CelebA | MSE: 3.65 -> 3.10 | Helped |
| use bi-level optimization with ξ rather than setting ξ = 0 (ablation) | MQVAE · CelebA | MSE: 3.61 -> 3.10 | Helped |
| replace identity hyper-net with MLP hyper-net (ablation) | MQVAE · CelebA | MSE: 5.54 -> 3.10 | Helped |
| replace Linear-1 hyper-net with MLP hyper-net (ablation) | MQVAE · CelebA | MSE: 3.50 -> 3.10 | Helped |
| replace Linear-2 hyper-net with MLP hyper-net (ablation) | MQVAE · CelebA | MSE: 4.25 -> 3.10 | Helped |
| replace VQVAE with MQVAE training | VQVAE · CelebA | Wall time to reach VQVAE LPIPS: 3.6 -> 2.4 | Helped |
Surprising findings
- Hypergradient ablation shows direct and indirect gradient norms are similar in magnitude with near-zero cosine similarity, suggesting both contribute distinct signals.
Limitations
- Total wall time of MQVAE (12.2 h) is much longer than VQVAE (3.6 h) due to bi-level optimization overhead.
- Experiments limited to image tasks with VQVAE and VQGAN; no evaluation on language or speech.
- Gradient approximation relies on one-step unrolling and finite-difference approximation, a practical surrogate rather than exact hypergradients.
Experimental setup
Datasets: CIFAR10 CelebA FFHQ ImageNet-1K
Benchmarks: CIFAR10 CelebA FFHQ ImageNet-1K
Metrics: MSE LPIPS perplexity rFID FID codebook utilization PSNR SSIM
Baselines: VQVAE SQVAE Gumbel-VQVAE l2 normalization LRU Affine OPT VQGAN-FC VQGAN-EMA VQGAN-LC SimVQ FSQ
Builds on: VQVAE VQGAN MAML Betty straight-through estimator
Closest papers by meaning
- Scalable Training for Vector-Quantized Networks with 100% Codebook Utilization (ICLR 2026 · similarity 0.90)
- Restructuring Vector Quantization with the Rotation Trick (ICLR 2025 · similarity 0.90)
- LG-VQ: Language-Guided Codebook Learning (NeurIPS 2024 · similarity 0.88)
- Finite Scalar Quantization: VQ-VAE Made Simple (ICLR 2024 · similarity 0.88)
- Codebook Transfer with Part-of-Speech for Vector-Quantized Image Modeling (CVPR 2024 · similarity 0.87)
- Scaling the Codebook Size of VQ-GAN to 100,000 with a Utilization Rate of 99% (NeurIPS 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.