# Restructuring Vector Quantization with the Rotation Trick

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

## Key facts

- Cited by 51 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 51 references that are in the map.
- First posted on arXiv 8 Oct 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [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)
- [Image-to-Image Translation with Conditional Adversarial Networks](https://mlsanity.com/paper/1611.07004/) (2016)
- [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)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Perceptual Losses for Real-Time Style Transfer and Super-Resolution](https://mlsanity.com/paper/1603.08155/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Is Space-Time Attention All You Need for Video Understanding?](https://mlsanity.com/paper/2102.05095/) (2021)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Towards Accurate Generative Models of Video: A New Metric & Challenges](https://mlsanity.com/paper/1812.01717/) (2018)
- [MaskGIT: Masked Generative Image Transformer](https://mlsanity.com/paper/2202.04200/) (2022)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Autoregressive Image Generation using Residual Quantization](https://mlsanity.com/paper/2203.01941/) (2022)
- [VideoGPT: Video Generation using VQ-VAE and Transformers](https://mlsanity.com/paper/2104.10157/) (2021)
- [Vector-quantized Image Modeling with Improved VQGAN](https://mlsanity.com/paper/2110.04627/) (2021)
- [Language Model Beats Diffusion -- Tokenizer is Key to Visual Generation](https://mlsanity.com/paper/2310.05737/) (2023)
- [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)
- [vq-wav2vec: Self-Supervised Learning of Discrete Speech Representations](https://mlsanity.com/paper/1910.05453/) (2019)
- [PeCo: Perceptual Codebook for BERT Pre-training of Vision Transformers](https://mlsanity.com/paper/2111.12710/) (2021)
- [Self-supervised Learning with Random-projection Quantizer for Speech Recognition](https://mlsanity.com/paper/2202.01855/) (2022)
- [Self-Supervised Visual Planning with Temporal Skip Connections](https://mlsanity.com/paper/1710.05268/) (2017)
- [Image and Video Tokenization with Binary Spherical Quantization](https://mlsanity.com/paper/2406.07548/) (2024)
- [Straightening Out the Straight-Through Estimator: Overcoming Optimization Challenges in Vector Quantized Networks](https://mlsanity.com/paper/2305.08842/) (2023)
- [UViM: A Unified Modeling Approach for Vision with Learned Guiding Codes](https://mlsanity.com/paper/2205.10337/) (2022)

## Built on by (most cited first, 30 of 51)

- [GigaTok: Scaling Visual Tokenizers to 3 Billion Parameters for Autoregressive Image Generation](https://mlsanity.com/paper/2504.08736/) (2025)
- [SoftVQ-VAE: Efficient 1-Dimensional Continuous Tokenizer](https://mlsanity.com/paper/2412.10958/) (2024)
- [Flow to the Mode: Mode-Seeking Diffusion Autoencoders for State-of-the-Art Image Tokenization](https://mlsanity.com/paper/2503.11056/) (2025)
- [Cube: A Roblox View of 3D Intelligence](https://mlsanity.com/paper/2503.15475/) (2025)
- [STAR: Learning Diverse Robot Skill Abstractions through Rotation-Augmented Vector Quantization](https://mlsanity.com/paper/2506.03863/) (2025)
- [Towards Unifying Understanding and Generation in the Era of Vision Foundation Models: A Survey from the Autoregression Perspective](https://mlsanity.com/paper/2410.22217/) (2024)
- [Diffusion Generative Recommendation with Continuous Tokens](https://mlsanity.com/paper/2504.12007/) (2025)
- [CODA: Repurposing Continuous VAEs for Discrete Tokenization](https://mlsanity.com/paper/2503.17760/) (2025)
- [Discrete Tokenization for Multimodal LLMs: A Comprehensive Survey](https://mlsanity.com/paper/2507.22920/) (2025)
- [Spherical Leech Quantization for Visual Tokenization and Generation](https://mlsanity.com/paper/2512.14697/) (2025)
- [Stop Treating Collisions Equally: Qualification-Aware Semantic ID Learning for Recommendation at Industrial Scale](https://mlsanity.com/paper/2603.00632/) (2026)
- [Instella-T2I: Pushing the Limits of 1D Discrete Latent Space Image Generation](https://mlsanity.com/paper/2506.21022/) (2025)
- [DiVeQ: Differentiable Vector Quantization Using the Reparameterization Trick](https://mlsanity.com/paper/2509.26469/) (2025)
- [Eliciting Chain-of-Thought Reasoning for Time Series Analysis using Reinforcement Learning](https://mlsanity.com/paper/2510.01116/) (2025)
- [Mofasa: A Step Change in Metal-Organic Framework Generation](https://mlsanity.com/paper/2512.01756/) (2025)
- [Scaling Image Tokenizers with Grouped Spherical Quantization](https://mlsanity.com/paper/2412.02632/) (2024)
- [Flow Autoencoders are Effective Protein Tokenizers](https://mlsanity.com/paper/2510.00351/) (2025)
- [Remote Inference over Dynamic Links via Adaptive Rate Deep Task-Oriented Vector Quantization](https://mlsanity.com/paper/2501.02521/) (2025)
- [PathVQ: Reforming Computational Pathology Foundation Model for Whole Slide Image Analysis via Vector Quantization](https://mlsanity.com/paper/2503.06482/) (2025)
- [WinQ: Accelerating Quantization-Aware Training of Language Models Around Saddle Points](https://mlsanity.com/paper/2605.17471/) (2026)
- [HH-Codec: High Compression High-fidelity Discrete Neural Codec for Spoken Language Modeling](https://mlsanity.com/paper/2507.18897/) (2025)
- [Graph is a Natural Regularization: Revisiting Vector Quantization for Graph Representation Learning](https://mlsanity.com/paper/2508.06588/) (2025)
- [JacQuant: STE-Free Quantization-Aware Training via Learned Jacobian Surrogates](https://mlsanity.com/paper/2605.25469/) (2026)
- [Spotlight-TTS: Spotlighting the Style via Voiced-Aware Style Extraction and Style Direction Adjustment for Expressive Text-to-Speech](https://mlsanity.com/paper/2505.20868/) (2025)
- [MuseTok: Symbolic Music Tokenization for Generation and Semantic Understanding](https://mlsanity.com/paper/2510.16273/) (2025)
- [Training-Free Vector Quantization via Gaussian VAEs](https://mlsanity.com/paper/2512.06609/) (2025)
- [VP-VAE: Rethinking Vector Quantization via Adaptive Vector Perturbation](https://mlsanity.com/paper/2602.17133/) (2026)
- [Beyond Stationarity: Rethinking Codebook Collapse in Vector Quantization](https://mlsanity.com/paper/2602.18896/) (2026)
- [Flow-Matching-Based Protein Structure Tokenizer Made Efficient and Easy](https://mlsanity.com/paper/2609.33129/) (2026)
- [Scaling Collider Event Generation with Residual-Quantized Tokens](https://mlsanity.com/paper/2610.00569/) (2026)

## Abstract (arXiv)

Vector Quantized Variational AutoEncoders (VQ-VAEs) are designed to compress a continuous input to a discrete latent space and reconstruct it with minimal distortion. They operate by maintaining a set of vectors -- often referred to as the codebook -- and quantizing each encoder output to the nearest vector in the codebook. However, as vector quantization is non-differentiable, the gradient to the encoder flows around the vector quantization layer rather than through it in a straight-through approximation. This approximation may be undesirable as all information from the vector quantization operation is lost. In this work, we propose a way to propagate gradients through the vector quantization layer of VQ-VAEs. We smoothly transform each encoder output into its corresponding codebook vector via a rotation and rescaling linear transformation that is treated as a constant during backpropagation. As a result, the relative magnitude and angle between encoder output and codebook vector becomes encoded into the gradient as it propagates through the vector quantization layer and back to the encoder. Across 11 different VQ-VAE training paradigms, we find this restructuring improves reconstruction metrics, codebook utilization, and quantization error. Our code is available at https://github.com/cfifty/rotation_trick.