Restructuring Vector Quantization with the Rotation Trick

TL;DR

This work tackles the non-differentiability of vector quantization in VQ-VAEs by introducing the rotation trick, which propagates gradients through the quantization layer via a rotation and scaling that align encoder outputs with their nearest codebook vectors. By treating the rotation and scaling as constants during backpropagation, the method preserves angular information in gradients, enabling more informative updates within Voronoi regions. Across 11 diverse VQ-based configurations, the rotation trick consistently improves reconstruction metrics and codebook utilization while reducing quantization error, outperforming STE and some gradient-augmentation baselines. The approach is practical for improving lossy compression performance in discrete latent representations and is accompanied by efficient implementation via Householder reflections, with code available for reproducibility.

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.

Figures (23)

• Figure 1: Illustration of the rotation trick. In the forward pass, encoder output $e$ is rotated and rescaled to $q_1$. For simplicity, the rotations of other encoder outputs are not shown. In the backward pass, the gradient at $q_1$ moves to $e$ so that the angle between $\nabla_{q_1} \mathcal{L}$ and $q_1$ is preserved. Now, points within the same codebook region receive different gradients depending on their relative angle and magnitude to the codebook vector. For example, points with high angular distance can be pushed into new codebook regions, thereby increasing codebook utilization.
• Figure 2: Visualization of how the straight-through estimator (STE) transforms the gradient field for $16$ codebook vectors for (top) $f(x,y) = x^2 + y^2$ and (bottom) $f(x,y) = \log(|\frac{1}{2}x + \tanh(y)|)$. The STE takes the gradient at the codebook vector $(q_x,q_y)$ and "copies-and-pastes" it to all other locations within the same codebook region, forming a "checker-board" pattern in the gradient field.
• Figure 3: Illustration of how the gradient at $q$ moves to $e$ via the STE (middle) and rotation trick (right). The STE "copies-and-pastes" the gradient to preserve its direction while the rotation trick moves the gradient so the angle between $q$ and $\nabla_q \mathcal{L}$ is preserved (proved in \ref{['appx:angle_preservation']}).
• Figure 4: Depiction of how points within the same codebook region change after a gradient update (red arrow) at the codebook vector (orange circle). The STE applies the same update to each point in the same region. The rotation trick modifies the update based on the location of each point with respect to the codebook vector.
• Figure 5: With the STE, the distances among points within the same region do not change. However with the rotation trick, the distances among points do change. When $\phi < \pi / 2$, points with large angular distance are pushed away (blue: increasing distance). When $\phi > \pi / 2$, points are pulled towards the codebook vector (green: decreasing distance).

Theorems & Definitions (11)

• Definition 1: Householder Reflection Matrix
• Remark 1
• Remark 2
• Remark 3
• proof
• Remark 4
• proof
• Remark 5
• proof
• Remark 6