Title: value: Restructuring Vector Quantization with the Rotation Trick
Authors: value: Christopher Fifty
Ronald Guenther Junkins
Dennis Duan
Aniketh Iyengar
Jerry Weihong Liu
Ehsan Amid
Sebastian Thrun
Christopher Re
Venue: ICLR 2025
Keywords: value: Vector Quantization
VQ-VAE

Abstract: value: 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.


Review discussion:
This paper is well-written and well-motivated. Its review of VQ-VAE, STE, and related work is concise and easy to follow, providing a clear motivation for the study.
The algorithm for the rotation trick is clearly formulated.
The explanation of the effects of the rotation trick, supported by figures, is illustrative.
The effectiveness of the proposed method is investigated through experiments in 11 different settings, which can be considered sufficient.
Overall, this paper is well-motivated and well-structured, but I would like to raise a point that requires clarification.
The algorithm is designed to preserves the angle between a vector and a gradient, which leads to the push-pull effect described in Section 4.3.
However, it appears that using only a rotation, $\frac{\partial \tilde{q}}{\partial e} = R$, satisfies this requirement, and rescaling with $\frac{||q||}{||e||}$ is not necessary for the gradient approximation. (I understand that rescaling is necessary for transforming $e$ into $q$.)
Additionally, it seems that rescaling of gradients may hinder the propagation of gradients when $||q||$ is close to zero.
Providing an explanation for why the rescaling factor $\frac{||q||}{||e||}$ is used for the gradient approximation would clarify the motivation behind this algorithm's design.
* Why is the rescaling factor $||q|| / ||e||$ used for the gradient approximation? How does it affect training of the encoder and the codebook?
* Minor comment: In Definition 1 (Householder Reflection Matrix), isn't it required that $||a|| = 1$?
* Minor comment: there is a typo: "a a convex combination'" in p3 L6.
No ethics review needed.
Thank you for championing our work!
Like many lines of research, our path to developing the rotation trick was quite convoluted, and we hoped to convey some of our learnings from the things that didn’t work (Section 3) as well as an understanding of why the rotation trick does work (Section 4). From your review and positive feedback, we’re encouraged that these points came through and were clearly communicated.
Specifically, you note that “the paper is well-written and well-motivated”, “the explanation of the effects of the rotation trick is illustrative”, and that our experimental evaluation “can be considered sufficient.”
### Addressing your questions:
>**Q1: Why is the rescaling factor $||q|| / ||e||$ used for the gradient approximation?**
We were mistaken to not include an explanation of this design choice; both you and Reviewer fmLN raise this point. Afterall, one could use $\tilde{q} = Re + (q-Re)$.detach() instead to transform $e$ into $\tilde{q}$ without scaling the norm of the gradients in the backward pass.
A _desideratum_ of vector quantization is that when $e \approx q$, then $\nabla_e \mathcal{L} \approx \nabla_q \mathcal{L}$ (i.e.
when distortion is low, the gradients for both $e$ and $q$ are approximately the same).
However when $||e|| \approx 0$ and a Euclidean metric is used to determine the closest codebook vector, the angle between $e$ and $q$ can be obtuse (see Figure 6 of A.1). In this instance, the rotation trick will cause the gradient $\nabla_e \mathcal{L}$ to “over-rotate” and point away from $\nabla_q \mathcal{L}$.
Using a grad scaling of $\frac{||q||}{||e||}$ can fix this.
When $||e|| \approx 0$ and $||e|| < ||q||$, the norm of the gradient will be scaled up to push $e$ away from the origin. Pushing $e$ away from the origin makes the angle between $e$ and $q$ more of a factor when computing the Euclidean distance:
\begin{align*}
||e-q|| &= \sqrt{||e||^2 + ||q||^2 - 2 ||e||||q||\cos \theta}
\end{align*}
so $e$ is more likely to map to a different $q$ that forms an acute angle with it as $||e||$ increases.
Now consider if $||q|| \approx 0$ and $||e|| > ||q||$.
When this occurs, the update to $e$ will vanish because $\frac{||q||}{||e||} \approx 0$. This behavior may also be desirable because when $q$ is close to the origin, there's a higher likelihood the angle between $e$ and $q$ would be obtuse. $||e||$ remains high and can be mapped to a different codebook vector after the update without moving in a (potentially) bad direction due to an obtuse angle with a very low norm $q$.
We’ve also added an ablation experiment to evaluate this difference in Table 8 and Table 9 of Appendix A.10.
Specifically, in the typical VQ-VAE paradigm presented in Table 8, we do not observe a difference between using $\tilde{q} = \frac{||q||}{||e||}Re$ and $\tilde{q} = Re + (q-Re)$.
However, for the VQGAN results in Table 9, we find that using the multiplicative gradient scaling factor modestly improves performance. We preview the results from Table 9 below:
| Rotation Trick Function            | Latent Shape      | Codebook Size | Codebook Usage (↑) | Quantization Error (↓) | Valid Loss (↓) | r-FID (↓) | r-IS (↑)  |
|------------------------------------|-------------------|---------------|---------------------|-------------------------|----------------|-----------|-----------|
| `\|q\|/\|e\|Re`           | `64×64×3`         | `8192`        | 45%                 | 4.0e-4                  | 0.161          | 0.46      | 225.0     |
| `Re - (q - Re)`                   | `64×64×3`         | `8192`        | 28%                 | 1.5e-3                  | 0.183          | 0.6       | 220.0     |
| `\|q\|/\|e\|Re`           | `32×32×4`         | `16384`       | 18%                 | 3.3e-4                  | 0.292          | 1.5       | 196.1     |
| `Re - (q - Re)`                   | `32×32×4`         | `16384`       | 13%                 | 9.4e-4                  | 0.292          | 1.5       | 191.5     |
We also examine other gradient scaling formulations, such as $\tilde{q} = \frac{1}{8||q-e||^2}Re + (q - \frac{1}{8||q-e||^2}Re)$, and discuss possible directions for future research in A.10.2 of the Appendix.
To summarize, we believe there is almost certainly a stronger approach to gradient scaling than a fixed $||q||/||e||$ factor throughout training—perhaps by exploring dynamic gradient scaling as is often done in the multi-task learning literature [1,2]---and this is an exciting direction for future work.
>**Q2: Definition 1 correction.**
Thank you for pointing this out; we’ve corrected the manuscript.
>**Q3: p3 L6 typo.**
Another good catch; thanks!
### References:
[1] Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks, Chen et al., 2018. [ICML]
[2] Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics, Kendall et al., 2018 [CVPR]
We would like to thank you again for championing our work – really, it encouraged our entire team. Repeating our remarks to Reviewer fmLN, we’d also like to thank you for prompting us to dig deeper into the effects of gradient scaling. These additional analyses led to a more thorough analysis of the method, adding 2 additional Tables, 2 additional Figures, and several sections to our analysis. Most importantly, it fills in a blank about the approach that was previously missing (i.e., the effect of the norm rescaling).
Dear Reviewer DhUX,
Thank you again for reviewing our work! As the discussion period ends shortly, we wanted to check if you have any further questions or found our responses helpful?
Please let us know and thank you for your time!
Thank you for answering my question about the rescaling factor.
The explanation in Section A.1 and the results in Table 9 are convincing to me.
I do not have any more questions.
I apologize for not having enough confidence to raise my rating to "should be highlighted at the conference", but I have raised the Soundness 3 --> 4.
Thank you again for the extensive investigation.
Great! We're glad the added analysis answers your question, and we appreciate your engagement during the discussion period!
The paper's core idea to set the Jacobian to the quantity suggested by the authors is reasonably well-motivated. It can be considered a natural dual to STE, as it uses the multiplicative/polar structure over Euclidean space as opposed to the additive structure used by STE.
The core idea in the paper is neat and simple, and the experimental verification appears reasonably extensive and rigorous.
The paper is well-written, with very nice illustrative figures.
While I do not think the paper suffers from any major weakness, there are two aspects in which it could be strengthened:
1. The authors do not discuss the limitations of their proposed method.
While they present good arguments for when their proposed rotation trick is preferable to STE, have they considered the scenarios in which it might be undesirable?
2. The theoretical motivation for the rotation trick is unclear, other than the fact that it is the natural "multiplicative counterpart" to STE.
While I really like the "post hoc" argument that the rotation trick results in a non-constant Jacobian, which might help improve codebook usage, I do not see a good mathematical reason for ensuring that the angles between the vectors and their gradients are preserved, which is better than magnitude and direction.
For example, one could argue that an advantage of the identity Jacobian is that it does not shrink or expand space since its determinant is one. This could be solved by combining the STE and the rotation trick by setting (in the authors' notation) $\tilde{q} = R e + (q - R e)$ for an appropriate rotation matrix $R$ and using an appropriate stop-gradient; this would now preserve both magnitude and angle as I understand.
In fact, we can even create an entire two-parameter family of gradient approximations by setting (in the authors' notation) $\tilde{q} = \lambda^\alpha R_{\alpha, \beta} e + \beta (q - R_{\alpha, \beta} e)$. Have the authors considered such a solution? If not, could the authors run some small-scale experiments at least to check its performance?
Minor:
- Fig 2: Increase the tick and axis label font size to match the font size of the caption, as the text is illegible at 100% zoom right now.
- Suggestion: Figure 2 beautifully illustrates the difference between the original and the STE gradient. To make it even nicer, I think the authors should consider adding the "rotation trick" gradient field as well!
- Recently, the STE was connected to the deeper theory of numerical integration in [1]. Could the authors comment on whether their method might also have such an interpretation?
## References
[1] Liu, L., Dong, C., Liu, X., Yu, B., & Gao, J. (2024). Bridging discrete and back-propagation: Straight-through and beyond. Advances in Neural Information Processing Systems, 36.
No ethics review needed.
We appreciate your positive review of our work! We really like the precision you employ to describe the rotation trick and thoughtfulness in proposing to generalize the algorithm to a larger family of rotation-based gradient estimators. We also *_really_* empathize with your point about searching for an underlying mathematical reason for gradient estimators to non-differentiable functions. This is something we’ve wrestled with quite a bit throughout the research process and hope our explanations below shed some light on how we’ve thought about this.
Moreover, we're encouraged by your thoughtful and positive feedback, particularly your recognition that the method is “neat and simple,” “reasonably well-motivated,” the experimental results are “reasonably extensive and rigorous,” and the manuscript is “well-written, with very nice illustrative figures.”
### Addressing your questions and concerns below:
>**W1: Limitations of the Rotation Trick.**
One potential limitation is when high numerical precision is required on the quantized vectors during training. As the rotation matrix does not have infinite precision, it will rotate $e$ to $q$ up to a small error term so that the input to the decoder will not be exactly $q$ during training, but rather offset by a very small factor depending on the precision used to compute $R$.
Nevertheless, this limitation does not arise during inference as a non-differentiable codebook lookup to transform $e$ into $q$ is used since gradients are not needed during test-time.
A second possible limitation is when either the encoder outputs or codebook vectors are forced to be close to 0 norm (i.e., $||e|| \approx 0$ and/or $||q|| \approx 0$).  In this case, the angle between $e$ and $q$ may be obtuse.
When this happens, $\nabla_q \mathcal{L}$ will “over-rotate” the gradient as it is transported from $e$ to $q$ so that $\nabla_q \mathcal{L}$ and $\nabla_e \mathcal{L}$ now point in different directions (i.e. the cosine of the angle between $\nabla_e \mathcal{L}$ and $\nabla_q \mathcal{L}$ will be negative). An example is visualized in Figure 6.
This is undesirable because---when the angle between $e$ and $q$ is obtuse---the rotation trick will violate the assumption that when $e \approx q$, $\nabla_q \mathcal{L} \approx \nabla_e \mathcal{L}$, and it will likely result in worse performance than VQ-VAEs trained with the STE.
While obtuse angles between $e$ and $q$ are very unlikely---by design, the codebook vectors should be ``angularly close'' to the vectors that are mapped to them---however, if there is a restriction that forces codewords to have near $0$ norm, then the rotation trick will likely perform worse than the STE.
We appreciate you prompting us to think about the limitations of this approach and have added this discussion to the manuscript in Appendix A.1.
>**W2: Theoretical Motivation and Preserving Gradient Magnitude.**
*Theoretical Motivation is unclear*
Adding a bit of context to how this work developed, we wanted to develop a method that incorporated information from the loss landscape around $q$ into the gradient at $e$.
Despite our best efforts for the better part of a month, we couldn’t get this to work, but found the learnings incredibly valuable. These learnings served as the basis for our presentation in Section 3, as well as the foundation of our analyses in Section 4.3, and we hope (and anticipate) they’ll be valuable to other researchers looking at vector quantization or similar problems.
We then started looking at how other fields moved vectors from point a to point b, and took inspiration from parallel transport in differential geometry to develop the rotation trick. As an exercise, imagine a circle with two points $a$ and $b$. Let $\vec{v}$ be a vector tangent to the circle at $a$. How should $\vec{v}$ be moved to point $b$?
There are many ways to answer this. We can translate $\vec{v}$ from $a$ to $b$, thereby preserving its direction and magnitude.
Alternatively, we can parallel transport $\vec{v}$ so that its orientation (i.e. tangent to the circle) is preserved. Neither solution is correct per se, but both can be useful in different circumstances (and even become equal to each other in a Euclidean space).
This analogy—along with the post hoc analysis in Section 4.3 and strong empirical results in Section 5—gave us the confidence that this result should be written up, even though we were unable to fix a precise notion of correctness for the vector quantization gradient estimator.
>**W2: Theoretical Motivation and Preserving Gradient Magnitude (cont.)**
*Preserving Gradient Magnitude in the Rotation Trick*
This is an excellent point (one that Reviewer DhUX also mentions); we decided to devote a large portion of our compute during the rebuttal period into running large-scale benchmarks to evaluate it. Empirically, we observe that setting $\tilde{q} = Re + (q - Re)$ results in slightly worse performance than $\tilde{q} = \frac{||q||}{||e||}Re$. We summarize these findings in Table 8 and Table 9, and we have also added A.10, A.10.1, and A.10.2 to the Appendix to dig into the differences between these two formulations. We preview the results from Table 9 below:
| Rotation Trick Function            | Latent Shape      | Codebook Size | Codebook Usage (↑) | Quantization Error (↓) | Valid Loss (↓) | r-FID (↓) | r-IS (↑)  |
|------------------------------------|-------------------|---------------|---------------------|-------------------------|----------------|-----------|-----------|
| `\|q\|/\|e\|Re`           | `64×64×3`         | `8192`        | 45%                 | 4.0e-4                  | 0.161          | 0.46      | 225.0     |
| `Re - (q - Re)`                   | `64×64×3`         | `8192`        | 28%                 | 1.5e-3                  | 0.183          | 0.6       | 220.0     |
| `\|q\|/\|e\|Re`           | `32×32×4`         | `16384`       | 18%                 | 3.3e-4                  | 0.292          | 1.5       | 196.1     |
| `Re - (q - Re)`                   | `32×32×4`         | `16384`       | 13%                 | 9.4e-4                  | 0.292          | 1.5       | 191.5     |
One difference is that the gradient scaling term will encourage acute angles between $e$ and $q$ by pushing $e$ away from the origin when $||e|| \approx 0$. We describe this effect in detail in A.10.1 of the Appendix. As discussed in Appendix A.1, obtuse angles cause the rotation trick to “over-rotate” the gradient, so acute angles between $e$ and $q$ are good for the rotation trick.
As you point out, we can consider an entire family of rotation-based gradient approximations by setting $\tilde{q} = \gamma(e) Re +(q - \gamma(e)Re)$. This is a really neat generalization: we create Section A.10.2 of the Appendix to review it in more depth.
Specifically, we discuss this generalization, visualize two cases we find especially interesting in Figure 14, and discuss how future work might build on this knowledge to develop better gradient estimators.
>**W3 [minor] Increase tick and axes label font:**
Done!
>**Q1: Reproduce Figure 2 for Rotation Trick:**
This is a good suggestion. We’ve added this visualization in Figure 14.
>**Q2: Connection of the Rotation Trick to Deeper Theory of Numerical Integration.**
This is a very cool question. We took a second, in-depth look at the Reinmax paper. Due to the multiplicative nature of our rotation trick, it is at this point unclear whether, like STE, there exists a connection between the rotation trick and numerical integration.
However, it is possible that such a connection may be established in the log-domain where the multiplication can be converted into addition. This parallels the connection between gradient descent and exponentiated gradient (mirror descent) updates. Establishing such connections more rigorously would be a future research direction for us.
We’d like to thank you again for taking the time to provide a thorough and through-provoking review of our work. You prompting us to dig deeper into the effects of gradient scaling led to a more thorough analysis of the rotation trick, adding 2 additional Tables, 2 additional Figures, and several sections to our analysis. Most importantly, it fills in a blank about the method that was previously missing.
If there are any other weaknesses that prevent this submission from being a strong contribution to the conference, please let us know.
I thank the authors for their detailed rebuttal. I checked the updated manuscript and am mostly happy with it; I have raised my score accordingly.
My one issue with the updated manuscript is that much of the appendix is not referenced in the main text: in the main text, I only found references to sections A.6, A.8 and A.11, and a few tables. Especially given that the appendix sections are also very nicely written, I would like to ask the authors to ensure that every section of the appendix is referenced in the main text so that the interested reader is aware of further discussion. It is particularly important that the limitations section is referenced. This section should ideally be in the main text, but I understand that the page limit does not allow for this.
Thank you for your quick response and engagement during the discussion period!
We also appreciate this latest suggestion to improve the presentation of our work and ensure that relevant analyses and learnings are not lost in the Appendix. **We commit to ensuring that every section of the Appendix is referenced in the main text, and in particular, referencing or pulling up the Limitations section.** For additional context, we’ve been discussing pushing the ViT-VQGAN and Video Experimental results to the Appendix to pull up excerpts from the Limitations Section and A.6 (Householder Reflection Transformation). We’d also like to add a subsection to Section 4 previewing the analyses in the Appendix and combine Tables 3, 7, and 9. We have held off on making these changes thus far in case other reviewers request additional analyses during the remainder of the review period (that would affect how these sections are organized).
Thank you again for the quick response, and we’re glad you find the added Appendix sections very nicely written.
1. The rotation trick is novel, intuitive, and mathematically principled
1. Applying two Householder reflectors is computationally efficient, and incurs no additional training cost
1. Strong empirical success across various VQ(-VAE|GAN) training paradigms
1. The method, i.e. how to achieve the rotation, is not explained clearly in the title or main text. It's probably a good idea to move Appendix A.4 into the main text.
In fact, consider renaming this paper to "Restructuring Vector Quantization with the *Reflection* Trick", because the rotation $R$ is really two Householder reflectors chained together (maybe we can get away with one reflection? see below).
1. Appendix A.4 appears to be incorrect: in Remark 3, I think the final reflection/rotation matrix should be $(I-2cc^T)(I-2aa^T)$, or equivalently $(I-2bb^T)(I-2cc^T)$, but not $(I-cc^T)(I-bb^T)$. To make this part easier to understand, perhaps you should re-use the $e, q, r$ notation instead of introducing new symbols $a, b, c$, and include a detailed derivation of why $R = (I-2cc^T)(I-2aa^T) = (I-2bb^T)(I-2cc^T) = I - 2rr^T + 2qe^T$.
1. The effect of the rotation trick is only verified on a high level by codebook usage and quantization error.
In particular, Figure 4 is only a diagram ("how we think it works"), not generated from an actual numerical simulation ("how it really works").
Moreover, Figure 4 only shows one gradient update step, and it's unclear how the rotation trick affect the codebook in the long run.
1. The notations are a little confusing.
For example, Section 4.3 allows $\tilde{\theta}$ to be negative or exceed $\pi$, but the angle between two vectors is typically defined within $[0, \pi]$. The top panel in Figure 5 also defines $\tilde{\theta}$ in a nonstandard way (when $q$ and $\nabla_q\mathcal{L}$ points at opposite directions, $\tilde{\theta} = 0$ according to the diagram, but normally we say $\tilde{\theta} = \pi$ for this case).
In addition, renaming $\tilde{\theta}$ to $\phi$ might avoid some misread.
1. There is an extra "a" on line 113.
1. What if we just do one Householder reflection, i.e. $R = I - 2rr^T$? It still preserves the angle ($\tilde{\theta} \in [0, \pi]$), but the "push" effect in Section 4.3 becomes "pull", and vice versa.
However, as you pointed out, both push and pull are beneficial, so the result should still be good. The mirror symmetry makes me think this will be as effective as $R = I - 2rr^T + 2qe^T$, and it would be great if you could verify this hypothesis experimentally.
1. The rotation/reflection transformation depends on the location of a center, which you implicitly choose to be the origin. Do you need to shift the $q$ and $e$ such that their center/mean is the origin, or does the rotation trick work even if all $q$ and $e$ fall on one side of the origin (e.g. imagine Figure 4, but on a shifted coordinates system $x' = x + 10, y' = y + 10$, such that the coordinate of all points are positive)?
No ethics review needed.
This is an excellent review; it’s clear you invested quite a bit of time into thoroughly reading and understanding our work. Your questions about a single reflection and invariance with respect to the origin are exceptional, and we’ve added several pages to the Appendix to address them.
Similarly, we appreciate your positive regard for our work, noting that “the rotation trick is novel, intuitive, and mathematically principled” as well as “the strong empirical success across various VQ(-VAE|GAN) training paradigms”. We address your questions below:
### Improvements to the Manuscript / Writing / Presentation:
>**W1: Describe how to Compute R in the Main Text.**
As you likely inferred, we moved the discussion of Householder reflections to the Appendix from the main text before the submission due to space limitations. Our thinking was that Householder transformations are not our contribution, and that interested readers could reference Appendix A.6 as suggested in Section 4.2. For the camera ready version, we’ll look into refactoring the paper—perhaps by moving the video results into the appendix—to pull up some of the background on Householder reflections.
>**W2: Remark 3 Incorrect & Detailed Derivation.**
You’re entirely correct; there are multiple mistakes in Remark 3. We sincerely thank you for catching these and have corrected them (see the revised manuscript). We also include a detailed derivation of why $R = I - 2rr^T + 2qe^T$.
>**W3: Figure 4 only a Diagram Not a Numerical Simulation.**
We were mistaken not to include a description of how Figure 4 was created. It is in fact a numerical simulation by applying a small gradient update over 25 steps. It’s not a perfect example: the gradient does not change across batches and we do not update the codebook regions;
however, it should faithfully highlight the differences between the STE and the rotation trick. We’ve added Appendix A.11.5 to describe this simulation.
As recommended in your review, we have also added an additional synthetic on Himmelblau’s function across 100 training steps staying as faithful as possible to the actual training paradigm used in VQ-VAEs. This simulation reflects what we see experimentally in our (ViT-)VQ(-VAE|GAN) results: the pre-quantized vectors cluster much more tightly around the codebook vectors (lower distortion) and the objective function’s loss is substantially lower. We have added this analysis to the manuscript in Figure 7, Figure 8, and Appendix A.2.
>**W4: Figure 5 Angle Axes.**
Another excellent catch, as you point out, Figure 5 is incorrect in the angle showing $\hat{\theta}$. We’ve updated Figure 5, and based on your later suggestion, have also changed $\tilde{\theta}$ to $\phi$.
For your point about the angle between vectors usually being in $[0, \pi]$, this makes sense for an undirected angle.
But if you want to distinguish between the angle needed to rotate $q$ to $\nabla_q \mathcal{L}$ and the angle needed to rotate $\nabla_q \mathcal{L}$ to $q$ (i.e. to differentiate between the clockwise and counterclockwise rotation), then you need the full range from $[0, 2\pi]$. Saying $\frac{-\pi}{2} < \phi < \frac{\pi}{2}$ is a precise way to indicate an acute angle, whereas saying $\phi < \frac{\pi}{2}$ may only indicate the angle lies within the first quadrant of the circle along which radians are typically defined.
>**W5: Typo.**
We appreciate you and reviewer DhUX both pointing out this typo. We have fixed it.
### Key Questions / Additional Analyses
>**Q1: The “Reflection Trick”**
This is an excellent insight. We’ve repeated the analysis in Section 4.3 for using a single reflection rather than a rotation to align $e$ to $q$ in the forward pass to create Figure 12, Figure 13, and A.9 of the Appendix.
Specifically, we define the Householder reflection $R = (I - 2 \frac{e-q}{||e-q||}(\frac{e-q}{||e-q||})^T)$ to reflect $e$ to $q$ in the hyperplane orthogonal to $\frac{e-q}{||e-q||}$ (where $e$ and $q$ are rescaled to have unit norm when computing R).
We refer you to the updated manuscript for the full analysis, but list the key findings below:
1.
When $\nabla_q \mathcal{L}$ is tangent to $q$, the behavior is identical to the rotation trick.
2. The component of $\nabla_q \mathcal{L}$ that is orthogonal to $q$ becomes reflected (i.e. reversed) when it is transported to $e$.
In particular, if the gradient points “left” at $q$, then it will flip to pointing “right” when it moves to $e$.
As a result, the update to $e$ can move in the opposite direction of the direction of the gradient, even when the angle between $q$ and $e$ is very small. This is undesirable as we would hope $\nabla_q \mathcal{L} \approx \nabla_e \mathcal{L}$ when $e \approx q$.
3. We evaluated this approach on both settings from Table 3 and the Euclidean codebook with latent Shape 64x64x3 and Codebook size of 8192 from Table 1. Both settings exhibited poor convergence using just a reflection rather than a rotation.
Specifically, after one epoch, the validation loss was approximately 3x higher than the rotation trick for both 8192 and 16384 codebook VQGANs in Table 3. For the Euclidean codebook model with latent Shape 64x64x3 in Table 1, the validation loss was approximately 2x higher than the rotation trick after 15 epochs.
This turned out to be quite the nuanced finding; thank you for suggesting this analysis.
>**Q2: Scaling and Invariance to the Origin**
Another excellent point that you note is the rotation trick—unlike the STE—is not invariant to the origin. For all of our empirical results, we use the origin of the vector space at the encoder output: we do not apply any translations or renormalizations.
Based on your suggestion, we directly analyze the case when all points are translated to be farther away from the origin (i.e. Figure 4 but with x’ = x+10 and y’ = y+10). We visualize this case in Figure 10 and discuss the general effect of translations in Appendix A.5. Summarizing this section, as points are translated far away from the origin, the effect of the rotation will become diminished, so that the rotation trick converges to the STE as $d \rightarrow \infty$ in $x’ = x+d$ and $y’ = y+d$.
Again, we truly appreciate your thorough review, helping us identify issues with the presentation, factual mistakes, and posing relevant (and interesting) questions. We hope our changes improve the paper so that it is now a strong contribution to the conference.
Dear Reviewer AkYx,
Thank you again for reviewing our work! As the discussion period ends shortly, we wanted to check if you have any further questions or found our responses helpful?
Please let us know and thank you for your time!
Dear authors, thanks for your detailed response and update to the manuscript. I'm also glad that you squashed that bug in your software implementation! The presentation is much better now, and most of my concerns have been addressed, so I'm raising my score.
Finally, I would appreciate it if you could clean up your code base and release it under a permissive license, so that others can build upon your work.
Thank you! We appreciate your engagement during the discussion period as well as your thorough review — really; you pointed out two subtle mistakes and suggested several notable follow-up analyses. Incorporating your feedback has markedly improved this work.
With respect to your request, **we commit to cleaning up our codebase and releasing it under an open-source (MIT) license.** We will include a link to the github repository at the end of the Abstract for the camera-ready version.
- The method is principled, and I think it is a good and simple idea to improve some of the issues of STE
- The paper motivates the issues with STE well, as well as the rotation trick itself. The explanations with figures are helpful to understand the intuition of the method.
- Experimental results are fairly extensive, encompassing many different scenarios where VQ-based models are used, each showing the superiority of the rotation trick over STE.
- I have some trouble understanding why the Hessian idea is not good. The explanations provided in the paper are that it is similar to a vanilla auto encoder training, with no quantization.
However, this seems only applicable to the "exact gradient" idea and not the Hessian idea. It seems the Hessian idea sets the gradient at e to the gradient at q (what STE does), plus the additional curvature term.
Whereas the rotation trick multiples the gradient at q by the scaled rotation matrix.
So it seems that both capture the location of e.
- Moreover, I wonder if the analysis in section 4.3 could be extended to the Hessian idea. It seems the Hessian idea also gets the benefits that Section 4.3 argues in favor of the rotation trick. There seems to be some disconnect between this and the experimental results where it was shown the Hessian idea does not perform well. If a stronger explanation is provided, that would be helpful to understand.
- While the experimental settings are fairly extensive, they primarily demonstrate how the rotation trick is superior to STE. I think some other baseline methods that improve upon STE are also necessary for comparison.
Namely, those mentioned in Section 2, such as Huh et al 2023, Gautam et al, 2023, and Takida et al, 2022. This would provide comparison between the rotation trick and other methods that go beyond STE.
Please see the weaknesses above.
No ethics review needed.
Thank you very much for taking the time to review our work and for sharing your thoughtful questions and concerns. You make an excellent point about adding additional baselines for comparison with stochastic quantization methods, we have added a section on this to the paper (see more below).
Your observations about the Hessian approximation also resonate deeply. As you likely inferred from the structure of the discussion in Section 3, we initially shared your intuition: we believed that estimating the curvature around $q$ and incorporating this information into the gradient at $e$ (referred to in the paper as "the Hessian approximation”) might lead to significant performance improvements for VQ-VAEs. We devoted considerable effort—weeks of synthetic experiments and benchmarking—to test this hypothesis rigorously.
However, despite our best efforts, the approach did not yield the results we hoped for.
While this was initially disappointing, it became a pivotal moment in developing our understanding to allow us to produce the algorithm we present today, the rotation trick. We expand on this transition from the Hessian idea to the rotation trick in Appendix A.3 and add Figure 9 to illustrate several instances where the Hessian idea can lead to undesirable behavior. We also touch on these concepts in our response below.
We are also encouraged by your positive reception of our work and noting its strengths.
Specifically, we appreciate your remarks that our “method is principled” and a “good and simple idea to improve upon the STE”, that our paper “motivates the issues with the STE well, as well as the rotation trick itself”, and finally that our “experimental results are fairly extensive, encompassing many different scenarios” and show “the superiority of the rotation trick over the STE”.
### Addressing your specific questions below:
>**“I have some trouble understanding why the Hessian idea is not good. The explanations provided in the paper are that it is similar to a vanilla auto encoder training, with no quantization.
However, this seems only applicable to the ‘exact gradient’ idea and not the Hessian idea."**
Interestingly, the Hessian idea actually approximates the exact gradient up to second order terms. By Taylor’s theorem, we can compute the loss at $e$:
\begin{align*}
\mathcal{L}_e &= \mathcal{L}_q + (\nabla_q \mathcal{L})^T(e-q) + \frac{1}{2}(e-q)^T(\nabla^2_q \mathcal{L})(e-q) + \frac{1}{6}(e-q)^T \nabla^3_q \mathcal{L} (e-q,e-q) + ….
\end{align*}
So the loss computed by the Hessian idea is different from the loss computed by the exact gradients method by $\mathcal{O}(||e-q||^3)$ which is the remainder term from truncating the Taylor series expansion at the second term.
Moreover, since the Hessian idea approximates the auto encoder gradients, we believe that both methods function similarly and are subject to the same limitations. **We added section A.3 to expand on this point: please let us know if you have any additional suggestions that would help further clarify this section.**
>**“It seems the Hessian idea sets the gradient at e to the gradient at q (what the STE does), plus the additional curvature term.
Whereas the rotation trick multiplies the gradient at q by the scaled rotation matrix.
So it seems that both capture the location of e.”**
This is exactly correct;
however, the way the Hessian idea transports the gradient is very different from the Rotation Trick.
The Hessian idea approximates the exact gradient at $e$ with a second-order Taylor expansion around $q$ (i.e. by adding the curvature term to the gradient at $q$). It approximates the loss landscape at $e$ and transforms the gradient to be more aligned with the direction of steepest descent. This is similar to the exact gradients idea, that exactly computes the loss landscape at $e$ to find the direction of steepest descent.
To summarize, the Hessian and exact gradient methods move the gradient based on the loss landscape, whereas the rotation trick (and the STE) transport the gradient in a way that is independent of the loss landscape. We have added Appendix A.3 to discuss several instances where using the loss surface to transport $\nabla_q \mathcal{L}$ from $q$ to $e$ can lead to undesirable behavior and visualize these effects in Figure 9.
>**"It seems the Hessian Idea also gets the benefits that Section 4.3 argues in favor of the rotation trick. There seems to be some disconnect between this and the experimental results where it was shown that the Hessian idea does not perform well. If a stronger explanation is provided, that would be helpful to understand."**
A desirable property of vector quantization is that when $e \approx q$ (i.e. low distortion), $\nabla_e \mathcal{L} \approx \nabla_q \mathcal{L}$. Informally, if a codeword matches a vector very well, then they should move together because this codeword approximates this vector with very little distortion. In the STE and the rotation trick, this will always be the case;
however in exact gradients and the Hessian approximation, it is not guaranteed since it depends on the loss surface in the Voronoi partition. This is undesirable because $q$ can move in one direction (the direction that minimizes the loss for the decoder), but $e$ can move in the opposite direction, even when it is very close to $q$ to begin with (i.e. a “good” fit with low distortion). We added Figure 9 to visualize this case.
Another undesirable effect is that—when using the loss surface to move gradients—even points that are very close to each other within the same Voronoi region can have drastically different gradients. This effect is visualized in Figure 9 of the Appendix, specifically the bottom partition.
We have added an in-depth discussion of this point to A.3 of the Appendix, and as you suggested, reproduced the analysis of Section 4.3 for the Hessian idea in Figure 9.
>**[Aside]**
As an aside, [1] finds that estimating the gradient to second order accuracy (i.e. the Hessian approximation) is beneficial for VAEs that use a discrete distribution (rather than the typical isotropic Gaussian). This is a great result, and our findings do not contradict theirs since the desiderata for vector quantization are very different from discrete VAEs.
>**Comparison with VQ-VAEs that use Non-Deterministic Quantization.**
We’ve added Table 6, Table 7, and A.4 of the Appendix to include comparisons to VQ-VAEs that use stochastic quantization, specifically focusing on industry-relevant benchmarks (i.e. the VQGANs used in latent diffusion). We preview the results of Tables 6 and 7 below:
| Approach                     | Latent Shape      | Codebook Size | Codebook Usage | Quantization Error (↓) | Valid Loss (↓) | r-FID (↓) | r-IS (↑)  |
|------------------------------|-------------------|---------------|----------------|-------------------------|----------------|-----------|-----------|
| VQGAN                        | `64×64×3`        | `8192`        | 15%            | 2.5e-3                  | 0.183          | 0.53      | 220.6     |
| Gumbel VQGAN                 | `64×64×3`        | `8192`        | 4%             | ---                     | 0.197          | 0.60      | 219.7     |
| VQGAN w/ Rotation Trick      | `64×64×3`        | `8192`        | 86%            | 1.7e-4                  | 0.142          | 0.27      | 228.0     |
| VQGAN                        | `32×32×4`        | `16384`       | 2%             | 1.2e-2                  | 0.385          | 5.0       | 141.5     |
| Gumbel VQGAN                 | `32×32×4`        | `16384`       | 12%            | ---                     | 0.3031         | 1.7       | 189.5     |
| VQGAN w/ Rotation Trick      | `32×32×4`        | `16384`       | 27%            | 2.4e-4                  | 0.269          | 1.1       | 200.2     |
| Approach                  | Codebook Usage (↑) | Rec. Loss (↓) | Quantization Error (↓) | Rec. Loss (↓) | r-FID (↓) | r-IS (↑)  |
|---------------------------|--------------------|---------------|-------------------------|---------------|-----------|-----------|
| **Codebook Lookup: Euclidean**, **Latent Shape:** `64×64×3`, **Codebook Size:** `8192`                                                                           |
| VQ-VAE                   | 100%              | 0.028         | 1.0e-3                 | **0.030**     | 19.0      | 97.3      |
| Gumbel VQ-VAE            | 39%               | 0.054         | ---                     | 0.058         | 28.6      | 74.9      |
| VQ-VAE w/ Rotation Trick | 99%               | 0.028         | 1.4e-4                 | **0.030**     | **16.5**  | **106.3** |
### References
[1] Liu, L., Dong, C., Liu, X., Yu, B., & Gao, J. (2024). Bridging discrete and back-propagation: Straight-through and beyond. Advances in Neural Information Processing Systems, 36.
We sincerely appreciate you pointing out areas where our presentation could be improved and asking us to highlight the differences between the Hessian approximation and the rotation trick gradient estimators. We hope our changes improve the paper so that it is now a strong contribution to the conference.
Dear Reviewer x3VA,
Thank you again for reviewing our work! As the discussion period ends shortly, we wanted to check if you have any further questions or found our responses helpful?
Please let us know and thank you for your time!
Dear authors,
Your response has addressed all my concerns. I find the explanations regarding the Hessian/exact gradients idea, along with the figures comparing how their gradients behave versus the rotation trick, extremely helpful, and provide a convincing explanation as to why they may not work in practice. The experiments demonstrating superiority of rotation trick over stochastic quantization are also convincing that the proposed method makes significant progress over other methods attempting to go beyond STE. I personally find the inclusion of these results to have strengthened the paper, and hope the authors agree. I have raised my score.
This is not necessary for the review period, but I would recommend the authors to integrate A.4 in main text and reference A.3 within sections 3 and 4 of the main text to help the reader.
Thank you! We appreciate your engagement during the discussion period and wholeheartedly agree that the changes you suggested—namely, going deeper with the Hessian/exact gradient analysis and adding a comparison to stochastic quantization methods—have decidedly strengthened this work.
Addressing your request, **we commit to integrating A.4 into the main text and referencing or pulling up the discussions from A.3 to Section 3 and 4** for the camera-ready version. We will begin working on this reorganization today.
Thank you again for reviewing this work, leaving constructive feedback that improves the paper, and participating in the discussion period!
We would like to thank each reviewer for considering our work, and their helpful suggestions for improving the clarity and presentation of our research. We are also encouraged by its positive reception, with the reviewers agreeing that our work is:
1. “Mathematically principled and well-motivated” (Reviewers x3Va, AkYx, fmLN, and DhUX)
2. “Extensive, rigorous, and very strong experimental results” (Reviewers x3Va, AkYx, fmLN, and DhUX)
3. “Well-written, easy to understand, illustrative figures” (Reviewers x3VA, fmLN, and DhUX)
Based on their questions and feedback, we have updated the manuscript and added 8 additional pages to the Appendix—changes highlighted in red—with new analyses, experimental results, and clarifications.
### Changes to the manuscript
* Reviewer x3VA’s suggestions:
* Clarify the Hessian approximation baseline (Figure 9, Appendix A.3)
* Benchmark stochastic quantization methods (Table 6, Table 7, Appendix A.4)
* Reviewer AkYx’s suggestions:
* Fix the Householder section (Appendix A.6)
* Run a numerical simulation comparing STE to rotation trick (Figure 7, Figure 8, Appendix A.2)
* Improve Section 4.3 notation (Figure 5, Section 4.3)
* Describe creation of Figure 4 (Appendix A.11.5)
* Use a single reflection rather than a rotation (Figure 12, Figure 13, Appendix A.9)
* Evaluate shifting the origin (Figure 11, Figure 10, Appendix A.5)
* Reviewer fmLN’s suggestions:
* Discuss possible limitations (Figure 6, Appendix A.1)
* Analysis into gradient scaling by $\frac{||q||}{||e||}$ (Table 8, Table 9, Figure 14, Appendix A.10, Appendix A.10.1, Appendix A.10.2)
* Add gradient fields for the rotation trick (Figure 14)
* Reviewer DhUX’s suggestions:
* Analysis into gradient scaling by $\frac{||q||}{||e||}$ (Table 8, Table 9, Figure 14, Appendix A.10, Appendix A.10.1, Appendix A.10.2)
We hope these changes improve the presentation of our work and turn it into a strong contribution to the conference.
### A Note Regarding Table 3 Results
After the initial submission deadline, we found a bug in our implementation of the rotation trick for the experiments in Table 3.
Specifically, we divided by $||e||$ twice in the forward pass (so the input to the decoder would be a codebook vector with its norm arbitrarily scaled up or down).
Fixing this bug further improves the performance of the rotation trick. For the 16384 codebook model, codebook utilization improves from 2% to 27%, r-FID from 1.6 to 1.1, and r-IS from 190.3 to 200.2.
Similarly, for the 8192 codebook model, codebook utilization improves from 10% to 86%, r-FID from 0.38 to 0.27, and r-IS from 222.4 to 228.0.

Meta-review:
role: AC
data: metareview: value: All reviewers appreciate the paper’s contribution and strongly recommend accepting the paper. The AC concurs, after reading the paper quickly.
additional_comments_on_reviewer_discussion: value: NA
role: PC

Decision: Accept (Oral)