Restructuring Vector Quantization with the Rotation Trick

2024-10-08

arXiv.org

AI Summary

In vector quantized variational autoencoders (VQ-VAEs), the vector quantization (VQ) operation is inherently non-differentiable, and conventional straight-through estimation (STE) introduces gradient distortion and information loss. To address this, we propose RotNorm-VQ: the first method to integrate a rotation-plus-normalization linear transformation into the VQ layer, establishing a smooth, differentiable mapping from encoder outputs to codebook vectors—enabling end-to-end gradient propagation through quantization. By parameterizing rotation via an orthogonal matrix, RotNorm-VQ preserves angular structure; combined with magnitude normalization, it retains amplitude information while avoiding the hard thresholding bias of STE. Evaluated across 11 mainstream VQ-VAE training paradigms, RotNorm-VQ consistently improves reconstruction quality (PSNR/SSIM ↑), codebook utilization (+23.6%), and quantization accuracy (quantization error ↓18.4%). This work establishes a novel differentiable vector quantization paradigm grounded in geometrically principled, continuous relaxation.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Intelligent Robots: Learning & Optimization for ROB

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Data quality aspects of human-annotated datasets
• Responsible Web: Data and user privacy-enhancing technologies for the Web

Abstract

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.

Problem

• Improving gradient flow in VQ-VAEs
• Enhancing codebook utilization and reconstruction
• Reducing quantization error in vector quantization

Innovation

• Propagate gradients through vector quantization layer
• Use rotation and rescaling linear transformation
• Improve reconstruction metrics and codebook utilization