# Straightening Out the Straight-Through Estimator: Overcoming Optimization Challenges in Vector Quantized Networks

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

## Key facts

- Cited by 82 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 41 references that are in the map.
- First posted on arXiv 15 May 2023.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [Unsupervised Learning of Visual Features by Contrasting Cluster Assignments](https://mlsanity.com/paper/2006.09882/) (2020)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [Contrastive Multiview Coding](https://mlsanity.com/paper/1906.05849/) (2019)
- [MaskGIT: Masked Generative Image Transformer](https://mlsanity.com/paper/2202.04200/) (2022)
- [Deep Clustering for Unsupervised Learning of Visual Features](https://mlsanity.com/paper/1807.05520/) (2018)
- [SoundStream: An End-to-End Neural Audio Codec](https://mlsanity.com/paper/2107.03312/) (2021)
- [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)
- [Muse: Text-To-Image Generation via Masked Generative Transformers](https://mlsanity.com/paper/2301.00704/) (2023)
- [Jukebox: A Generative Model for Music](https://mlsanity.com/paper/2005.00341/) (2020)
- [Self-labelling via simultaneous clustering and representation learning](https://mlsanity.com/paper/1911.05371/) (2019)
- [When Vision Transformers Outperform ResNets without Pre-training or Strong Data Augmentations](https://mlsanity.com/paper/2106.01548/) (2021)
- [Fast Decoding in Sequence Models using Discrete Latent Variables](https://mlsanity.com/paper/1803.03382/) (2018)
- [Vector-Quantized Autoregressive Predictive Coding](https://mlsanity.com/paper/2005.08392/) (2020)
- [SQ-VAE: Variational Bayes on Discrete Representation with Self-annealed Stochastic Quantization](https://mlsanity.com/paper/2205.07547/) (2022)
- [Hierarchical Quantized Autoencoders](https://mlsanity.com/paper/2002.08111/) (2020)

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

- [Emerging Properties in Unified Multimodal Pretraining](https://mlsanity.com/paper/2505.14683/) (2025)
- [Finite Scalar Quantization: VQ-VAE Made Simple](https://mlsanity.com/paper/2309.15505/) (2023)
- [Learning to Act without Actions](https://mlsanity.com/paper/2312.10812/) (2023)
- [Copilot4D: Learning Unsupervised World Models for Autonomous Driving via Discrete Diffusion](https://mlsanity.com/paper/2311.01017/) (2023)
- [GFT: Graph Foundation Model with Transferable Tree Vocabulary](https://mlsanity.com/paper/2411.06070/) (2024)
- [GIVT: Generative Infinite-Vocabulary Transformers](https://mlsanity.com/paper/2312.02116/) (2023)
- [Addressing Representation Collapse in Vector Quantized Models with One Linear Layer](https://mlsanity.com/paper/2411.02038/) (2024)
- [Scaling the Codebook Size of VQGAN to 100,000 with a Utilization Rate of 99%](https://mlsanity.com/paper/2406.11837/) (2024)
- [SoftVQ-VAE: Efficient 1-Dimensional Continuous Tokenizer](https://mlsanity.com/paper/2412.10958/) (2024)
- [Restructuring Vector Quantization with the Rotation Trick](https://mlsanity.com/paper/2410.06424/) (2024)
- [Masked Particle Modeling on Sets: Towards Self-Supervised High Energy Physics Foundation Models](https://mlsanity.com/paper/2401.13537/) (2024)
- [OmniJet-α: The first cross-task foundation model for particle physics](https://mlsanity.com/paper/2403.05618/) (2024)
- [Adaptive Length Image Tokenization via Recurrent Allocation](https://mlsanity.com/paper/2411.02393/) (2024)
- [Autoregressive Diffusion Transformer for Text-to-Speech Synthesis](https://mlsanity.com/paper/2406.05551/) (2024)
- [Residual Quantization with Implicit Neural Codebooks](https://mlsanity.com/paper/2401.14732/) (2024)
- [EdVAE: Mitigating Codebook Collapse with Evidential Discrete Variational Autoencoders](https://mlsanity.com/paper/2310.05718/) (2023)
- [VQDNA: Unleashing the Power of Vector Quantization for Multi-Species Genomic Sequence Modeling](https://mlsanity.com/paper/2405.10812/) (2024)
- [When Worse is Better: Navigating the compression-generation tradeoff in visual tokenization](https://mlsanity.com/paper/2412.16326/) (2024)
- [Cube: A Roblox View of 3D Intelligence](https://mlsanity.com/paper/2503.15475/) (2025)
- [Aspen Open Jets: Unlocking LHC Data for Foundation Models in Particle Physics](https://mlsanity.com/paper/2412.10504/) (2024)
- [OmniJet-αC: Learning point cloud calorimeter simulations using generative transformers](https://mlsanity.com/paper/2501.05534/) (2025)
- [Zero-Shot Styled Text Image Generation, but Make It Autoregressive](https://mlsanity.com/paper/2503.17074/) (2025)
- [Efficient Autoregressive Shape Generation via Octree-Based Adaptive Tokenization](https://mlsanity.com/paper/2504.02817/) (2025)
- [Soft Convex Quantization: Revisiting Vector Quantization with Convex Optimization](https://mlsanity.com/paper/2310.03004/) (2023)
- [Learning the Language of Protein Structure](https://mlsanity.com/paper/2405.15840/) (2024)
- [Scalable Training for Vector-Quantized Networks with 100% Codebook Utilization](https://mlsanity.com/paper/2509.10140/) (2025)
- [Foundation models for high-energy physics](https://mlsanity.com/paper/2509.21434/) (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)
- [SpidR: Learning Fast and Stable Linguistic Units for Spoken Language Models Without Supervision](https://mlsanity.com/paper/2512.20308/) (2025)

## Abstract (arXiv)

This work examines the challenges of training neural networks using vector quantization using straight-through estimation. We find that a primary cause of training instability is the discrepancy between the model embedding and the code-vector distribution. We identify the factors that contribute to this issue, including the codebook gradient sparsity and the asymmetric nature of the commitment loss, which leads to misaligned code-vector assignments. We propose to address this issue via affine re-parameterization of the code vectors. Additionally, we introduce an alternating optimization to reduce the gradient error introduced by the straight-through estimation. Moreover, we propose an improvement to the commitment loss to ensure better alignment between the codebook representation and the model embedding. These optimization methods improve the mathematical approximation of the straight-through estimation and, ultimately, the model performance. We demonstrate the effectiveness of our methods on several common model architectures, such as AlexNet, ResNet, and ViT, across various tasks, including image classification and generative modeling.