Research map / Probabilistic and scientific modeling / Variational and generative sampling
Restructuring Vector Quantization with the Rotation Trick
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 the rotation trick, a gradient estimator for VQ-VAE quantization that transports gradients from the codebook vector to the encoder output via a detached rotation-and-rescaling, improving reconstruction, codebook usage, and quantization error across 11 VQ-VAE/VQGAN/ViT-VQGAN/TimeSformer settings.
Problem
Straight-through estimators for vector quantization in VQ-VAEs discard information about where the encoder output lies within its Voronoi cell, contributing to codebook collapse, low codebook utilization, and high quantization error.
Approach
Replace the straight-through pass-through with a linear map q̃ = (||q||/||e||) R e, where R is a rotation aligning e to q treated as constant in backprop, so the gradient at q is rotated to e while preserving the angle between gradient and codebook vector; this lets the location of e within its quantization region affect the encoder update (push points apart or pull them together).
Key results
- On ImageNet VQGAN for latent diffusion (32x32x4, 16384 codes), the rotation trick improves r-FID from 5.0 to 1.1 and r-IS from 141.5 to 200.2, with codebook usage rising from 2% to 27%.
- For a VQ-VAE with 64x64x3 latent and 8192 Euclidean codes, quantization error drops from 1.0e-3 to 1.4e-4 while r-FID improves from 19.0 to 16.5.
- The STE causes TimeSformer-VQGAN codebook collapse on BAIR and UCF-101 (utilization near 0%), while the rotation trick reaches 43% and 30% utilization with r-FVD 21.4 and 229.1.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace STE with the rotation trick | ImageNet VQ-VAE, Euclidean, 32×32×32 · ImageNet | r-FID: 106.1 -> 85.7 | Helped |
| replace STE with the rotation trick | ImageNet VQ-VAE, Euclidean, 64×64×3 · ImageNet | r-FID: 19.0 -> 16.5 | Helped |
| replace STE with the rotation trick | VQGAN, autoregressive · ImageNet | r-FID: 7.3 -> 4.6 | Helped |
| replace STE with the rotation trick | VQGAN, autoregressive · FFHQ & CelebA-HQ | codebook usage: 27% -> 99% | Helped |
| replace STE with the rotation trick | VQGAN, latent diffusion, 32×32×4 · ImageNet | r-FID: 5.0 -> 1.1 | Helped |
| replace STE with the rotation trick | VQGAN, latent diffusion, 32×32×4 · ImageNet | r-IS: 141.5 -> 200.2 | Helped |
| replace STE with the rotation trick | ViT-VQGAN · ImageNet | r-FID: 29.2 -> 11.2 | Helped |
| replace STE with the rotation trick | TimeSformer VQ-VAE · BAIR video | r-FVD: 1661.1 -> 21.4 | Helped |
| replace STE with the rotation trick | TimeSformer VQ-VAE · UCF-101 video | r-FVD: 2878.1 -> 229.1 | Helped |
| replace norm scaling with the additive rotation formulation (ablation) | VQ-VAE, Euclidean, 64×64×3 · ImageNet | r-IS: 106.3 -> 106.1 | No clear effect |
| use norm scaling ||q||/||e|| rather than the additive rotation formulation (ablation) | VQGAN, latent diffusion, 64×64×3, 8192 codes · ImageNet | r-FID: 0.6 -> 0.46 | Helped |
| replace rotation with the reflection trick (ablation) | VQGAN · ImageNet, latent diffusion | validation loss: approximately 3x higher | Hurt |
Surprising findings
- The straight-through estimator leads to codebook collapse for TimeSformer-VQGAN on video in all tried hyperparameter settings, while the rotation trick trains stably.
- Replacing the rotation trick with a reflection, or using exact-gradient or Hessian-based approaches, generally degrades performance relative to the STE.
Limitations
- Undesirable over-rotation when encoder outputs or codebook vectors have near-zero norm, making the angle between e and q obtuse.
- Gains are reported on reconstruction and codebook metrics rather than downstream generative performance.
- Some experiments (e.g., reflection-trick ablations) were not trained to completion due to GPU limits.
Experimental setup
Datasets: ImageNet FFHQ CelebA-HQ BAIR Robot UCF-101
Benchmarks: ImageNet validation reconstruction FFHQ & CelebA-HQ reconstruction BAIR video reconstruction UCF-101 video reconstruction
Metrics: r-FID r-IS codebook usage quantization error validation reconstruction loss r-FVD
Baselines: VQ-VAE (straight-through) Gumbel VQ-VAE Hessian approximation exact gradients (double forward pass) VQGAN ViT-VQGAN
Builds on: VQ-VAE VQGAN ViT-VQGAN straight-through estimator Gumbel-Softmax
Closest papers by meaning
- Training-Free Vector Quantization via Gaussian VAEs (ICML 2026 · similarity 0.91)
- Diffusion bridges vector quantized variational autoencoders (ICML 2022 · similarity 0.90)
- Learning to Quantize for Training Vector-Quantized Networks (ICML 2025 · similarity 0.90)
- Finite Scalar Quantization: VQ-VAE Made Simple (ICLR 2024 · similarity 0.89)
- Information-theoretic Generalization Analysis for VQ-VAEs: A Role of Latent Variables (NeurIPS 2025 · similarity 0.88)
- LG-VQ: Language-Guided Codebook Learning (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.