Restructuring Vector Quantization with the Rotation Trick (2410.06424v1)

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.

Restructuring Vector Quantization with the Rotation Trick

This paper addresses the challenges inherent within the use of Vector Quantized Variational AutoEncoders (VQ-VAEs) by proposing a novel method termed the Rotation Trick. VQ-VAEs, despite their prominent use in compressing continuous inputs into discrete latent space, suffer from training instabilities primarily due to the non-differentiable nature of vector quantization. Current methodologies commonly employ the Straight-Through Estimator (STE) to approximate gradients, but this can result in significant issues such as codebook collapse and suboptimal model performance.

Contributions

The paper makes significant technical contributions via the introduction of the Rotation Trick, which offers an alternative method for gradient propagation through the vector quantization layer. The Rotation Trick utilizes a smooth linear transformation involving rotation and rescaling. This transformation shifts the process from focusing merely on the closest codebook vector to considering the angular relationships between the encoder output and the codebook vectors. By encoding not only the magnitude but also the angle between the encoder output and its nearest codebook vector into the gradient, this method promises improvements in codebook utilization and quantization error.

Experimental Results

The experimental results presented indicate measurable improvements using the Rotation Trick. When evaluated on various datasets, including ImageNet, across multiple metrics—such as codebook usage, reconstruction FID, IS, and quantization error—the proposed method significantly enhances the reconstruction performance of VQ-VAEs. For instance, utilizing the Rotation Trick on a VQGAN model reduced FID from 5.0 to 1.6 and increased codebook usage from 2% to 9%, showcasing its efficacy in improving model performance.

Theoretical and Practical Implications

Theoretically, the Rotation Trick redefines the gradient flow through vector quantization layers by preserving relative angles and magnitudes, bypassing the shortcomings of the Straight-Through Estimator. Practically, this leads to enhanced model training stability, increased codebook utilization, and reduced quantization error, potentially allowing more information to be retained during the quantization process.

Future Directions

Given the demonstrated benefits of the Rotation Trick in various experimental setups, future work could explore its applicability beyond VQ-VAEs to other models utilizing vector quantization layers. Additionally, examining the efficacy of this approach in generative modeling applications where the Visual Quality of reconstructed outputs is paramount might provide further insights into its versatility and robustness.

Conclusion

In conclusion, the work effectively tackles the intrinsic limitations posed by vector quantization in VQ-VAEs through an innovative approach that restructures gradient propagation. By preserving angular information, the Rotation Trick not only mitigates the risk of codebook collapse but also substantially enhances model performance metrics, offering a promising avenue for further research and application in both established and emerging deep learning architectures.