Restructuring Vector Quantization with the Rotation Trick

Authors: Christopher Fifty, Ronald Guenther Junkins, Dennis Duan, Aniketh Iyengar, Jerry Weihong Liu, Ehsan Amid, Sebastian Thrun, Christopher Re
Primary area: generative models
Keywords: Vector Quantization, VQ-VAE

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.

Decision: accept (oral)
Reviewer scores: 8, 8, 8, 8

Reviews and discussion:
<thread>
creation date:
1730279291695

thread_type:
official_review

<thread_object>
creator: reviewer_dhux

<content>
summary:
This paper proposes a method, reffered to as the "rotation trick", to propagate approximated gradients through the vector quantization layer of VQ-VAEs.
The proposed method addresses the issue of a "straight-through estimator (STE)," which loses information about the positions of encoder outputs relative to codebook vectors.
The capability of the method to increase information capacity and reduce distortion in vector quantization is illustrated with examples and explanation of the "push-pull effect".
The authors conduct experiments across 11 VQ-VAE training settings to demonstrate that the proposed method improves reconstruction performance compared to baseline methods.

strengths:
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.

weaknesses:
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.

questions:
* 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730660354510

thread_type:
official_review

<thread_object>
creator: reviewer_fmln

<content>
summary:
The authors propose an alternative to using straight-through estimation for the gradients of a vector-quantized variational autoencoder (VQ-VAE) to stabilize training and improve codebook usage.

VQ-VAEs are a popular model for data compression and generation, consisting of an encoder and a decoder network. However, instead of passing the encoder output directly to the decoder, the output is first quantized to some codebook vector; these vectors are also learnt during training. However, vector quantization (VQ) is a non-differentiable operation; thus, practitioners usually ignore the quantization operation in the backward pass to enable gradient back-propagation; this is called straight-through estimation (STE). More formally, given some encoder output $e$ and its vector quantized version $q$, STE "sets" the Jacobian $\frac{\partial q}{\partial e} = I$, where $I$ is the identity matrix instead of the all-zero matrix (almost everywhere) as would be dictated by the standard rules of calculus.

In this paper, the authors identify some important shortcomings of using STE in practice, such as training instability and codebook collapse. Based on their observations, they instead suggest "setting" the Jacobian $\frac{\partial q}{\partial e} = \frac{\lVert q \rVert}{\lVert e \rVert} \cdot R$, where $R$ is an orthogonal matrix depending on $e$ such that $q = \frac{\partial q}{\partial e} \times e$.

The authors argue and demonstrate empirically across several experiments that this pick significantly improves codebook utilization, reconstruction loss, and reconstruction quality.

strengths:
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.

weaknesses:
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.

questions:
- 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730828841494

thread_type:
official_review

<thread_object>
creator: reviewer_akyx

<content>
summary:
The authors address the training instability of the straight-through estimator by proposing the rotation trick for vector quantization (VQ), i.e. rotating the gradient $\nabla_q\mathcal{L}$ at the codebook vector $q$ to become the gradient $\nabla_e\mathcal{L}$ at the encoder output $e$ such that the angle between $q$ and $\nabla_q\mathcal{L}$ after the VQ equals the angle between $e$ and $\nabla_e\mathcal{L}$. The rotation trick significantly improves the image reconstruction quality across 11 VQ(-VAE|GAN) training paradigms.

strengths:
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

weaknesses:
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.

questions:
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)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731045562271

thread_type:
official_review

<thread_object>
creator: reviewer_x3va

<content>
summary:
This paper proposes an improvement to the STE method for training VQ-VAE. While STE copies the gradient from the quantizer output to input, this paper proposes the rotation trick, which uses a rotated and scaled gradient according to the value of the quantizer input. This enables the gradient to depend on the value of the quantizer input, whereas in STE it is the same for all inputs that get mapped to the same codebook vector. Experimental results demonstrate superior reconstruction performance and codebook usage across a variety of tasks using VQ-VAE methods.

strengths:
- 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.

weaknesses:
- 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.

questions:
Please see the weaknesses above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732144224635

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Meta Response

comment:
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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732938091058

thread_type:
public_comment

<thread_object>
creator: ~xianghong_fang1

<content>
title:
Reported Quantization Errors in Tables 1 are Unfair and Meaningless.

comment:
Dear Program Chairs, Senior Area Chairs, Area Chairs, All Reviewers, Authors, and All Readers,
 
I am a PhD student in University of Toronto, and I have four-year's experience in Vector Quantization (VQ). I am highly skeptical about the reported quantization error improvements in Tables 1 and 2. My reservations are based on the following reasons:
- Suppose there are a feature vector distribution $\mathcal{P}_A$ and a code vector distribution $\mathcal{P}_B$. I have observed that the quantization error is minimized when feature vectors and code vectors are identically distributed, i.e., $\mathcal{P}_A=\mathcal{P}_B$.
- The scale of distribution has a significant impact on the quantization error. Suppose feature vectors and code vectors are sampled from $\mathcal{N}(0, \sigma^2 I)$ or $\text{Unif}(-\zeta, \zeta)$, The variation in quantization error lower bound can be significantly large when $\sigma$ varies. We also provide some experimental results as follow (We provide the code in the Github (https://github.com/sunset-clouds/Quantization-error), every one can justify this):
- For Gaussian distribution $\mathcal{N}(0, \sigma^2 I)$ (e-1=0.1), 
 | $\sigma$ | 0.0001 | 0.001 | 0.01 | 0.1 | 1.0|
| --- | --- | --- | --- | --- | --- | 
| Quantization Error |   1.25e-8  |   1.25e-6  |    1.25e-4  |    1.24e-2   |   1.25   | 

- For uniform distribution $\text{Unif}(-\zeta, \zeta)$ (e-1=0.1), 
 | $\zeta$ | 0.0001 | 0.001 | 0.01 | 0.1 | 1.0|
| --- | --- | --- | --- | --- | --- | 
| Quantization Error |   3.27e-9  |   3.27e-7  |    3.27e-5  |    3.27e-3   |   0.327  | 
- The core issue is **how the authors ensure that the feature vectors across different methods are identically distributed**. If the feature vector distributions are not the same, the reported improvements in quantization error are meaningless, as the reported improvements might be **attributable to differences in distribution scales rather than the effectiveness of the incorporation of the Rotation Trick**. For example, by using vanilla  STE in [1], **I can achieve a more significant reduction in quantization error through regularization of the distribution scales** from $\sigma=1$ to $\sigma=0.001$ (quantization error from $1.25$ to $1.25e-6$). Therefore, the authors should provide a clear explanation demonstrating that the quantization error improvement is indeed due to the effectiveness of the incorporation of the Rotation Trick, rather than merely a result of altered distribution scales.
- Based on my experience, the quantization error comparison in image reconstruction experiments is almost impossible. In different vector quantization (VQ) methods, encoders receive entirely different gradient updates, and thus feature vector distribution are totally different.

The authors need to explain **whether the incorporation of the Rotation Trick will not inflence the distribution scales**; otherwise, the quantization error comparisons provided in Tables 1 and 2 are entirely meaningless. If authors fail to address my concern, quantization error comparison should be removed in the Table 1,2,3,4,5.

We kindly ask the area chair, all reviewers, and readers to take note of my comments and request the authors to respond to this, as I have provided solid evidence to support my concerns. 

**I request that the authors provide a detailed explanation for the observed improvement in quantization error within the scope of their study, and I have not requested a comparison with other research papers.**

- [1] Neural Discrete Representation Learning, NeurIPS 2017

Best Regards
</content>
</thread_object>

</thread>


Metareview:
metareview:
All reviewers appreciate the paper’s contribution and strongly recommend accepting the paper. The AC concurs, after reading the paper quickly.