PaperlayerView exact revision on arXiv

Source-linked AI summary

Restructuring Vector Quantization with the Rotation Trick

Christopher Fifty, Ronald G. Junkins, Dennis Duan, Aniketh Iyengar, Jerry W. Liu, Ehsan Amid, Sebastian Thrun, Christopher Ré

arXiv:2410.06424v2cs.LGcs.CV

TL;DR

Vector quantization makes VQ-VAE gradients difficult to propagate because the quantization step is non-differentiable, and the STE bypasses that operation. The rotation trick uses a rotation-and-rescaling transformation to carry gradients back through the quantization layer. Across 11 settings, the authors report improved reconstructions, codebook use, and quantization error.

  • Problem

    Vector quantization is non-differentiable, so the STE bypasses the quantization operation, an approximation associated with poor performance and codebook collapse or under-utilization.

    Source: §1 · p. 1

  • Method

    The rotation trick smoothly transforms each encoder output to its nearest codebook vector with a rotation and rescaling treated as constants during backpropagation.

    Source: §1 · p. 2

  • Results

    Across 11 settings, reconstructions improve; one VQGAN's r-FID changes from 5.0 to 1.1, r-IS from 141.5 to 200.2, codebook usage increases 13.5x, and quantization error falls two orders of magnitude.

    Source: §7 · p. 10

  • Takeaways & Limitations

    Preserving the angle rather than the direction between the codebook vector and gradient induces desirable updates within a quantized region and is associated with substantial model-performance improvement.

    Source: §4 · p. 5, §7 · p. 10

    When encoder outputs or codebook vectors are forced near zero norm, the angle between them may be obtuse, causing gradient over-rotation and likely worse performance than the STE.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

The introduction frames non-differentiable vector quantization and limitations of the straight-through estimator as challenges for VQ-VAE training. It presents the rotation trick as an angle-preserving way to propagate gradients, reporting improved reconstruction, codebook use, and quantization error.

1 INTRODUCTION

  • Vector quantization is non-differentiable, and the STE bypasses it by copying decoder-input gradients to encoder outputs, a shortcut associated with poor performance and codebook under-utilization.Under-utilization limits the information capacity of the VQ-VAE bottleneck.Source: §1 · p. 1
  • The rotation trick smoothly maps each encoder output to its nearest codebook vector through a rotation and rescaling, treated as constant during backpropagation.The transformed encoder output is sent to the decoder, allowing gradients to flow back to the encoder.Source: §1 · p. 2
  • Preserving the angle between the codebook vector and gradient makes updates depend on relative angles and magnitudes, allowing points in one codebook region to spread apart or move closer.The authors connect spreading points with increased codebook use and bringing them closer with lower quantization error.Source: §1 · p. 2
  • On ImageNet VQGAN, reconstruction FID improves from 5.0 to 1.1, reconstruction IS from 141.5 to 200.2, codebook usage from 2% to 27%, and quantization error by two orders of magnitude.The reported comparison uses the VQGAN from Rombach et al. (2022).Source: §1 · p. 2

2 RELATED WORK

Related work addresses vector-quantization training instability either by sidestepping the STE or by changing codebook-model interactions. These approaches target different parts of the training process, while the rotation trick replaces the STE used by many of them.

2 RELATED WORK

  • Prior approaches fall into two categories: methods that sidestep the STE and methods that improve codebook-model interactions.The review focuses on efforts addressing training instabilities caused by the vector-quantization layer.Source: §2 · pp. 2–3
  • STE alternatives include stochastic or soft quantization and alternating optimization, but can bring complications such as low inference-time codebook use or temperature-schedule tuning.The text notes that many applications and papers continue to train VQ-VAEs with the STE.Source: §2 · p. 3
  • Codebook-interaction methods change lookup, codebook learning, or utilization losses; the authors report that replacing the STE with the rotation trick further improves performance in most of their experiments.Examples include alternative distance measures, codebook-vector updates, and utilization penalties.Source: §2 · p. 3

3 STRAIGHT THROUGH ESTIMATOR (STE)

This section describes the VQ-VAE quantization and training setup, then examines the STE and alternatives that attempt to improve gradient propagation. It explains why the STE gives uniform updates within codebook regions and why Hessian-based or exact-gradient alternatives can create a mismatch between encoder and decoder objectives.

3 STRAIGHT THROUGH ESTIMATOR (STE)

  • The quantizer selects the codebook vector nearest to encoder output e, while the VQ-VAE objective combines reconstruction with codebook and commitment losses.The codebook loss learns codebook vectors; the commitment loss pulls encoder outputs toward them.Source: §3 · p. 3
  • The STE copies the gradient at codebook vector q to encoder output e, bypassing non-differentiable quantization and giving every point within a codebook region the same update.The resulting gradient field is partitioned into regions with uniform updates, regardless of where the encoder output lies within each region.Source: §3 · p. 4, Figure 2 · p. 4
  • A Hessian approximation estimates curvature around q to approximate the gradient at e, with codebook-only Hessian-vector products described as feasible without forming the full Hessian.The section introduces this as an apparent way to improve on the STE approximation.Source: §3 · pp. 4–5
  • The exact-gradient alternative uses two network passes, combining a quantized loss with a scaled loss from a network without quantization to obtain encoder gradients.The text specifies a small λ, such as 10^-6, and rescales the encoder’s gradient contribution by λ^-1.Source: §3 · p. 5
  • The authors report poor reconstruction with both alternatives, attributing it likely to AutoEncoder-like encoder training and a mismatch between encoder and decoder objectives.They state that the Hessian approximation and exact-gradient double pass train the encoder like an AutoEncoder and the decoder like a VQ-VAE.Source: §3 · p. 5

4 THE ROTATION TRICK

The rotation trick changes how gradients pass through vector quantization by preserving the angle between the codebook vector and its gradient. Its position-dependent updates can push encoder outputs into other codebook regions or pull them toward their selected vector.

4 THE ROTATION TRICK

  • The section frames the rotation trick as a way to use an encoder output’s location within its quantized region to transform gradients, rather than copying gradients as the STE does.Source: §4 · p. 5

4.1 THE ROTATION TRICK PRESERVES ANGLES

  • The rotation trick moves gradients from the codebook vector to the encoder output while preserving their angle to the codebook vector, unlike the STE, which preserves gradient direction and magnitude.Source: §4 · p. 5, §4.1 · p. 5

4.2 EFFICIENT ROTATION COMPUTATION

  • The rotation aligning encoder output e with codebook vector q can be computed using Householder matrix reflections; this parameterization avoids outer products and showed no detected wall-clock difference from the STE in the experiments.Source: §4.2 · p. 6

4.3 VORONOI PARTITION ANALYSIS

  • Across VQ-VAEs, the rotation trick often reduces quantization error by an order of magnitude and substantially increases codebook usage compared with the STE.Source: §4.3 · p. 6
  • Compared with the STE’s identical updates within a region, the rotation trick’s push-pull behavior can increase codebook utilization and reduce quantization error.Boundary points may move into other, possibly unused, regions; points loosely clustered around a codebook vector may be pulled closer.Source: §4.3 · pp. 7–8

4.4 FURTHER ANALYSIS

  • The appendix further analyzes a non-convex synthetic example, behavior far from the origin, reflection versus rotation, and scaling the gradient norm by the codebook-vector norm.Source: §4.4 · p. 8

5 EXPERIMENTS

The experiments test whether the rotation trick’s predicted effects extend across VQ-VAE paradigms. Results report improved reconstruction and codebook usage across VQ-VAE, VQGAN, and ViT-VQGAN evaluations, with quantization error varying by model.

5 EXPERIMENTS

  • The evaluation begins with ImageNet reconstruction VQ-VAEs and extends to more complex VQGAN and ViT-VQGAN paradigms.The section frames this broader evaluation as a test of how the rotation trick’s predicted effects translate to applications.Source: §5 · p. 8, §5.1 · p. 8, §5.2 · p. 9

5.1 VQ-VAE EVALUATION

  • The rotation trick substantially improves VQ-VAE reconstruction r-FID and r-IS, reduces quantization error, and improves low codebook utilization.The reported changes are consistent with the paper’s Voronoi-region analysis.Source: §5.1 · p. 9
  • The Hessian approximation and exact-gradient methods yield poor reconstruction performance; the authors suggest overfitting and mismatched training objectives may contribute.The proposed explanation is qualified as likely, not established causation.Source: §5.1 · p. 9

5.2 VQGAN EVALUATION

  • Across both VQGAN modeling paradigms, the rotation trick improves validation total loss, r-FID, and r-IS while decreasing quantization error and improving codebook usage.The evaluation covers autoregressive and latent-diffusion VQGAN settings.Source: Table 3 · p. 9, Table 2 · p. 9, §5.2 · p. 10

5.3 VIT-VQGAN EVALUATION

  • For ViT-VQGAN, codebook utilization and reconstruction metrics improve, while quantization error remains roughly the same.The results are summarized in Table 4 for ImageNet.Source: §5.3 · p. 10, Table 4 · p. 10

6 LIMITATIONS

The rotation trick has a limitation when encoder outputs or codebook vectors are near zero norm: the angle between them may be obtuse, causing gradient over-rotation and potentially worse performance than the STE.

6 LIMITATIONS

  • Near-zero encoder-output or codebook-vector norms can make the angle between them obtuse.The limitation arises when either vector is forced close to zero norm.Source: §6 · p. 10
  • With an obtuse angle, transporting the gradient can over-rotate it so the gradients at the encoder output and codebook vector point in different directions.Figure 6 illustrates this case.Source: §6 · p. 10, Figure 6 · p. 10
  • Under that condition, the rotation trick may perform worse than the STE.The authors describe this outcome as likely, not certain.Source: §6 · p. 10

7 CONCLUSION

The paper finds that preserving the angle between the codebook vector and gradient, rather than the direction, changes updates within codebook regions and improves model performance. Across 11 settings, rotation-trick VQ-VAEs improve reconstructions; a latent-diffusion VQGAN example also reports large gains in reconstruction metrics, quantization error, and codebook usage.

7 CONCLUSION

  • Preserving the angle rather than the direction between codebook vector and gradient changes how points within a codebook region are updated.The paper reports that these effects substantially improve model performance.Source: §7 · p. 10
  • Across 11 settings, rotation-trick VQ-VAEs improve reconstructions.The conclusion summarizes results across varied training configurations.Source: §7 · p. 10
  • For a latent-diffusion VQGAN, rotation-trick training improves r-FID from 5.0 to 1.1 and r-IS from 141.5 to 200.2.The same example reduces quantization error by two orders of magnitude and increases codebook usage by 13.5x.Source: §7 · p. 10

A APPENDIX

The appendix includes Table 5 results for TimeSformer-VQGAN trained on BAIR and UCF-101, with 1024 codebook vectors.

A APPENDIX

  • Table 5 reports TimeSformer-VQGAN results on BAIR and UCF-101.The models use 1024 codebook vectors.Source: Table 5 · p. 15
  • The table marks a model that suffers codebook collapse and diverges with †.The marker identifies the affected model in the table.Source: Table 5 · p. 15
  • The reported r-FVD is computed on the validation set.Source: Table 5 · p. 15

A.1 VIDEO EVALUATION

The video evaluation compares STE and rotation-trick training in small TimeSformer-VQGAN models on BAIR Robot and UCF101. STE training collapses codebook utilization, whereas the rotation trick yields stable training and high-quality reconstructions.

A.1 VIDEO EVALUATION

  • Compute limitations led to relatively small models with 8 layers, 256 hidden dimensions, 4 attention heads, and 768 MLP hidden dimensions.Both encoder and decoder use this TimeSformer configuration.Source: §A.1 · p. 15
  • With STE, codebook utilization falls to almost 0% within the first several epochs on both video datasets.The authors report this collapse across several hyperparameter settings.Source: §A.1 · p. 15
  • Rotation-trick models show no training instability and produce high-quality reconstructions as indicated by r-FVD.The study evaluates BAIR Robot and UCF101 using the ViT-VQGAN quantization paradigm with a TimeSformer in place of ViT.Source: §A.1 · p. 15

A.2 NON-CONVEX SYNTHETIC EXAMPLE

A synthetic vector-quantization simulation on Himmelblau’s function compares STE and rotation-trick gradient estimators over 100 updates. The rotation trick clusters points more tightly around codebook vectors and converges more rapidly toward the function’s four equal local minima.

A.2 NON-CONVEX SYNTHETIC EXAMPLE

  • The simulation runs 100 gradient updates with EMA codebook updates at decay 0.8 and a 1e-3 learning rate for pre-quantized points.STE and rotation-trick runs share the same random initialization; only the gradient estimator differs.Source: §A.2 · p. 15
  • The rotation trick clusters points more tightly around codebook vectors than STE, lowering distortion, and converges faster toward Himmelblau’s four equal minima.Across the simulated points, the rotation trick also yields a lower average objective value.Source: §A.2 · p. 15
  • Himmelblau’s function has four equal local minima, each with function value 0.0.The listed minima are approximately (3.0, 2.0), (-2.8, 3.1), (-3.7, -3.2), and (3.5, -1.8).Source: Figure 7 · p. 15

A.3 HESSIAN APPROXIMATION AND EXACT GRADIENT ANALYSIS

The analysis identifies a limitation of curvature-based exact-gradient and Hessian-approximation approaches: they can send gradients in very different directions even when encoder outputs are close to their codebook vectors. STE and the rotation trick preserve the close-pair gradient assumption in these cases.

A.3 HESSIAN APPROXIMATION AND EXACT GRADIENT ANALYSIS

  • The Hessian approximation treats the encoder update as if quantization had not occurred, approximating the gradient used in the original AutoEncoder model.The section introduces the Hessian approximation as a second-order approximation to exact gradients using a Taylor expansion around q.Source: §A.3 · pp. 15–17
  • The analysis frames the desired behavior as ∇eL ≈ ∇qL when e ≈ q, because otherwise a small encoder-to-codeword distortion can increase during an update.The authors describe this as an inductive bias for vector quantization to work well.Source: §A.3 · p. 17
  • Negative curvature or an indefinite Hessian can make exact-gradient and Hessian-approximation updates violate the close-pair assumption, even when e is near q.In the illustrated cases, gradients at e and q may point in very different directions, including near a saddle point.Source: §A.3 · p. 17
  • STE and the rotation trick preserve the close-pair gradient assumption because they do not use the loss-surface curvature to move ∇qL from q to e.The analysis contrasts these estimators with curvature-based exact-gradient and Hessian approaches.Source: §A.3 · p. 17

A.4 BEHAVIOR AWAY FROM THE ORIGIN

The rotation trick depends on the location of the origin, unlike STE. Shifting encoder outputs and codebook vectors far from the origin reduces the angle between them, and the rotation trick approaches the STE update in that limit.

A.4 BEHAVIOR AWAY FROM THE ORIGIN

  • As a shared displacement moves encoder outputs and codebook vectors far from the origin, the angle between them decreases.The section considers a displacement much larger than the original vector magnitudes.Source: §A.4 · p. 18, Figure 11 · p. 18
  • In the far-from-origin limit, the rotation matrix approaches identity and the rotation trick smoothly becomes the STE update.The figure description likewise notes that the trick’s effect decreases as the angle between e and q shrinks.Source: §A.4 · p. 18, Figure 10 · p. 18
  • After translation, the rotation trick’s push effect is reduced; other regions can shift from push to pull or show more rotation than contraction.The illustrated behavior varies with the gradient direction and region.Source: §A.4 · p. 18

A.5 HOUSEHOLDER REFLECTION TRANSFORMATION

This appendix examines the rotation trick’s gradient-transport construction and its relationship to the straight-through estimator, including alternative transformations, scaling choices, and evaluation behavior. It also discusses what reported training metrics do—and do not—establish about downstream generative modeling.

A.6 PROOF THE ROTATION TRICK PRESERVES ANGLES

  • A proof establishes that the angle between q and ∇qL is preserved as the gradient moves to e.The appendix describes the rotation as efficiently computed using two Householder reflections in the plane spanned by e and q.Source: §A.6 · pp. 19–20, §A.5 · p. 19

A.7 TREATING R AND ||q||

  • Treating R and the norm ratio as constants avoids differentiating through the non-differentiable codebook lookup, and the resulting approximation conveys more quantization information than the STE.The method drops the derivative term that would require differentiating through Q(e).Source: §A.7 · p. 20

A.8 THE REFLECTION TRICK

  • Validation loss with the reflection trick was approximately 3x higher than the rotation trick for both VQGAN codebook sizes after one epoch, and approximately 2x higher for the specified VQ-VAE after 15 epochs.The VQGANs used codebooks of 8192 and 16384; the VQ-VAE used a Euclidean codebook and latent shape 64 × 64 × 3.Source: §A.8 · p. 22

A.9 GRADIENT NORM SCALING IN THE ROTATION TRICK

  • VQ-VAE ablations showed no observed difference between the two rotation formulations, whereas applying gradient-norm scaling modestly improved VQGAN performance.The comparison uses the formulations ˜q = ||q||/||e||Re and ˜q = Re + (q − Re); Tables 6 and 7 summarize the VQ-VAE and VQGAN settings.Source: §A.9 · p. 22, Table 7 · p. 23, Table 6 · p. 23

A.9.2 GENERAL FAMILY OF ROTATION-BASED GRADIENT ESTIMATORS

  • The rotation trick’s additive and multiplicative formulations are instances of a broader family of rotation-based gradient estimators parameterized by a scaling factor γ(e).The appendix identifies fixed scaling based on an inductive bias and adaptive scaling as possible directions for future work.Source: §A.9.2 · p. 23

A.10.2 VQGAN EVALUATION

  • The appendix reports that VQGAN reconstructions trained with the rotation trick better reproduce original-image details, while also documenting the hyperparameter settings used for the VQ-VAE experiments.The reconstruction comparison is described alongside increased r-FID; the VQ-VAE settings use the CNN architecture of Esser et al. (2021) and EMA codebook updates with decay = 0.8.Source: §A.10.1 · p. 23, §A.10.2 · p. 24, Figure 15 · p. 25

A.10.3 VIT-VQGAN EVALUATION

  • The ViT-VQGAN rotation trick reduced r-FID from 29.2 to 11.2 and reproduced small image details missed by the STE model.The reported comparison is for the ViT-VQGAN evaluation; Figure 16 shows non-cherry-picked reconstructions.Source: §A.10.3 · p. 25, Figure 16 · p. 26

A.10.4 TIMESFORMER VIDEO EVALUATION

  • Rotation-trained video models used 43% of the codebook on BAIR and 30% on UCF101, while STE-trained models collapsed to 4 of 1024 and 2 of 2048 vectors, respectively.The rotation trick also produced reconstructions that accurately matched inputs despite the small video models.Source: §A.10.4 · p. 26, Figure 17 · p. 26, Figure 18 · p. 27

A.11 COMPARISON WITHIN GENERATIVE MODELING APPLICATIONS

  • The paper does not analyze downstream generative modeling or claim that better Stage 1 reconstruction, codebook usage, or quantization error improves Stage 2 performance.The authors leave that analysis to future work.Source: §A.11 · pp. 27–28
  • Prior work (Yu et al., 2023) suggests that better VQ-VAE reconstruction can make MaskGit likelihood-based modeling harder as codebook size grows, and preferences may differ across generative models.The authors suggest likelihood-based and diffusion models may have different desiderata, and Transformers (Vaswani, 2017) may differ from MaskGit (Chang et al., 2022).Source: §A.11 · p. 28

A.12.1 BACKGROUND ON HYPERSPHERICAL COORDINATES

  • Normalized hyperspherical basis vectors rotate with position and induce an identity metric, so constant components in that basis correspond to rotated vectors in Cartesian coordinates.The appendix presents normalization as equivalent to rotating the Cartesian tangent-space basis to align with the hyperspherical frame.Source: §A.12.1 · pp. 31–32

A.12.2 STE AS PARALLEL TRANSPORT

  • The STE preserves gradient direction and magnitude along any curve from q to e, equivalent to parallel transport with the identity metric in Cartesian coordinates.Its transport path is unspecified; the gradient is simply held constant along the curve.Source: §A.12.2 · pp. 32, 34

A.12.3 THE ROTATION TRICK AS PARALLEL TRANSPORT

  • The rotation trick parallel-transports gradients along any smooth path from q to e in normalized hyperspherical coordinates, corresponding to rotating them in Cartesian coordinates.The analysis treats only the rotational component, not the rotation trick’s rescaling by ∥q∥.Source: §A.12.3 · pp. 34, 36
  • The rotation trick and STE both path-independently parallel-transport ∇qL from q to e under the Euclidean metric, differing only in Cartesian versus hyperspherical normalized-basis coordinates.The STE uses Cartesian coordinates, while the rotation trick uses hyperspherical coordinates with normalized basis vectors.Source: §A.12.3 · p. 36
Loading 2410.06424v2…