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
Venue: ICLR 2025 Oral (2025)
Keywords: Vector Quantization
VQ-VAE
Primary area: generative models

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.

Published as a conference paper at ICLR 2025
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é Stannorr University. Gooole DeepMind fi fiftyêcs. stanford. com
ABSTRACT
Vector Quantized Variational AutoEncoders (VQ-VABs) are designed to compress a continuous input to a discrete latent space and reconstruct with minimal distortion. They operate by maintaining a set of vecters-often referred to the codebook-and quantizing each encoder output to the nearest vector the codebook. However, vector quantization is non-differentiable, the gradient to the encoder flows around the vector quantization layer rather than through it straight- through approximation. This approximation may be undesirable information from the vector quantization operation is lost. In this work, we propose way to propagate gradients through the vector quantization layer of VQ-VAEs. We smoothly transform each encoder output into its corresponding codebook vector via rotation and rescaling linear transformation that 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-VAB training paradigms, we find this restructuring improves reconstruction metrics, codebook utilization, and quantization error. Our code is available at https///gitub..co///fffftttttttt trick.
all
INTRODUCTION
Vector quantization (Gray, 1984) is an approach to discretize a continuous vector space. It defines a finite of vectors- referred vectors-referree to as the codebook-and and maps any vector in the continuous vector space to the closest vector in the codebook. However, deep learning paradigms that use vector quantization are often difficult to train because replacing vector with its closest codebook counterpart is non- differentiable operation (Huh et al., 2023). This characteristic was not issue at its creation during the Renaissance of Information Theory for applications like noisy channel communication (Cover, 1 1999); however in the deep learning, it presents challenge gradients cannot directly flow through layers that use vector quantization during backpropagation.
In deep learning, vector quantization is largely used in the eponymous Vector Quantized-Varti Variational AutoEncoder (VQ-VAR) (Van Den Oord et al., 2017). VQ-VAE an AutoEncoder with a vector quantization layer between the encoder's output and decoder's input, thereby quantizing the learned representation at the bottleneck. While VQ- VQ-VAEs are ubiquitous in state-of-the-art generative modeling (Rombach et al., 2022; Dhariwal et al., 2020; Brooks al., 2024), their gradients cannot flow from the decoder to the encoder uninterrupted as they must pass through a non-difierentiable vector quantization layer.
A solution to the non-differentiability problem is to approximate gradients via a strriigt-through estimator" (STE) (Bengio et al., 2013). During backpropagation, the STE copies and pastes the gradients from the decoder's input to the encoder's output, thereby skipping the quantization operation altogether. However, this approximation can lead to poor-performing models and codebook collapse: phenomena where large percentage of the codebook converge to zero norm and are unused by the model (Mentzer et al., 2023). Even if codebook collapse does not occur, the codebook is often under-utilized thereby limiting the information capacity of the VQ- VO-V'ss's 's bottleneck (Dhariwal et al., 2020).

---

Published as a conference paper at ICLR 2025
Codebook
27
Rotate und Rescale to Codebook Vectons
Encoder
Decoder
Rry Cor
steps srodima
Backward Pass Codebook Regions before Update
Gradient Update
Codebook Regions after Updale
V,C
STE
V,C
Rotation Trick
C
a
Figure 1: Illustration of the rotation trick. In the forward pass, encoder output is rotated and rescaled to Qt- For simplicity, the rotations of other encoder outputs not shown. In the backward pass, the gradient Qı moves so that the angle between V,, and qI preserved. Now, points within the same codebook region receive different gradients depending on their relative angle magnitude to the codebook vector. example, points with high angular distance can be pushed into new codebook regions, thereby increasing codebook utilization.
In this work, propose alternate way propagate gradients through the vector quantization layer in VQ- VAEs. For given encoder output and nearest codebook vector we smoothly transform e to via a rotation and rescaling linear transformation and then send this output- rather than the direct result of the codebook lookup to the decoder. As the input to the decoder, 7. is now treated as smooth linear transformation of e, gradients flow back from the decoder to the encoder unimpeded. To avoid differentiating through the rotation and rescaling, treat both as constants with respect to E and q. We explain why this choice is necessary in Appendix A.7. Following the convention of Kingma & Welling (2013), call this restructuring "the rotation trick." " is illustrated in Figure and described in Algorithm 1.
Algorithm I The Rotation Trick
Require: input example x
<- Encoder(x:)
q f- nearest codebook vector to e
R rotation matrix that aligns to
i <- stop-gradient
|gl
R
|
I < Decoder(
loss + L(x, 2)
return loss
The rotation trick does not change the output of the VQ-VAE in the forward pass. However, during the backward pass, it transports the gradient V,C. at q to become the gradient VeL at so that the angle between q and V.L. after the vector quantization layer equals the angle between e and V.L. before the vector quantization layer. Preserving this angle encodes relative angular distances and magnitudes into the gradient and changes how points within the same codebook region are updated.
The STE applies the same update to all points within the same codebook region, maintaining their relative distances. However as we will show in Section 4.3, the rotation trick can push points within the same codebook region farther apart-or pull them closer together-depeee the gradient vector. The former capability can correspond to increased codebook usage while the latter to lower quantization error. In the context of lossy compression, both capabilities are desirable for reducing the distortion and increasing the information capacity of the vector quantization layer.
on the direction of
When applied to several open-source VQ-VAE repositories, we find the rotation trick substantively improves reconstruction performance, increases codebook usage, and decreases the distance between encoder outputs and their corresponding codebook vectors. For instance, training the VQGAN from Rombach et al. (2022) on ImageNet (Deng et al., 2009) with the rotation trick improves reconstruction FID from 5.0 to 1.1, reconstruction IS from 141.5 to 200.2, increases codebook usage from 2% to 27%, and decreases quantization error by two orders of magnitude.
RELATED WORK
Many researchers have built upon the seminal work of Van Den Oord et al. (2017) to improve VQ- VAE performance. While non-exhaustive, review focuses methods that address training
2

---

Published as a conference paper at ICLR 2025
instabilities caused by the vector quantization layer. We partition these efforts into two categories: (1) methods that sidestep the STE and (2) methods that improve codebook model interactions.
Sidestepping the STE. Several prior works have sought to fix the problems caused by the STE by avoiding deterministic vector quantization. Baevski et al. (2019) employ the Gumbel-Softmax trick (Jang et al., 2016) to fit categorical distribution over codebook vectors that converges to a one-hot distribution towards the end of training. Gautam et al. (2023) quantize using a convex combination of codebook vectors, and Takida et al. (2022) employ stochastic quantization. Unlike the above that cast vector quantization as a distribution over codebook vectors, Huh et al. (2023) propose an alternating optimization where the encoder is optimized to output representations close to the codebook vectors while the decoder minimizes reconstruction loss from fixed set of codebook vector inputs. While these approaches sidestep the training instabilities caused by the STE, they can introduce their own set of problems and complexities such as low codebook utilization at inference and the tuning of temperature schedule (Zhang et al., 2023). As result, many applications and research papers continue to employ VQ-VAEs that are trained using the STE (Rombach et al., 2022; Chang et al., 2022; Huang et al., 2023; Zhu et al., 2023; Dong et al., 2023).
Codebook-Modee Improvements. Another way to attack codebook collapse or under-utilization is to change the codebook lookup. Rather than use Euclidean distance, Yu et al. (2021) employ a cosine similarity measure, Goswami et al. (2024) hyperbolic metric, and Lee et al. (2022) stochastically sample codes as a function of the distance between the encoder output and codebook vectors. Another perspective examines the learning of the codebook. Kolesnikov et al. (2022) split high-usage codebook vectors, Dhariwal et al. (2020); Lańcucki et al. (2020); Zheng & Vedaldi (2023) resurrect low-usage codebook vectors throughout training, Chen et al. (2024) dynamically selects one of m codebooks for each datapoint, and Mentzer et al. (2023); Zhao et al. (2024); Yu et al. (2023); Chiu et al. (2022)
fix the codebook vectors to an a priori geometry and train the model without learning the codebook at all. Other works propose loss penalties to encourage codebook utilization. Zhang et al. (2023) add Kl.-divergence penalty between codebook utilization and uniform distribution while Yu et al. (2023) add an entropy loss term to penalize low codebook utilization. While effective at targeting specific training difficulties, as each of these methods continue to use the STE, the training instability caused by this estimator persist. Most of our experiments in Section 5 implement subset of these approaches, and we find that replacing the STE with the rotation trick further improves performance.
3 STRAIGHT THROUGH ESTIMATOR (STE)
In this section, we review the Straight-Through Estimator (STE) and visualize its effect on the gradi- ents. We then explore two STE alternatives that-at at first glance appear to correct the approximation made by the STE.
For notation, we define a sample space A' over the input data with probability distribution p. For input & E X, we define the encoder as a deterministic mapping that parameterizes a posterior distribution (e(z))) The vector quantization layer, 2(.). is function that selects the codebook vector q C nearest to the encoder output e. Under Euclidean distance, it has the form:
[ i arg min gig lle 4ااو2
Q(g i|e)
0 otherwise
The decoder is similarly defined as a deterministic mapping that parameterizes the conditional distribution over reconstructions pD (rlg) As in the VAE (Kingma & Welling, 2013), the loss function follows from the ELBO with the KL-divergence term zeroing out as pe (elx)) ) is deterministic and the utilization over codebook vectors is assumed to be uniform. Van Den Oord et al. (2017) additionally add "cooebook loss" term (|) 2 to learn the codebook vectors and "commitment loss" term Blle- sy(4) 2/ to pull the encoder's 's output towards the codebook vectors. ag stands for stop- gradient and /3 is hyperparameter, typically set to value in [0.25, 2]. For predicted reconstruction I, the optimization objective becomes:
( || a||2 + [|sg(e) u1l2 + Bl|e ~
In the subsequent analysis, we focus only on the || - - 2 term as the other two are not functions of the decoder. During backpropagation, the model must differentiate through the vector quantization
3

---

Published as a conference paper at ICLR 2025
(() x²
Gradient Field f(2,)) = =x² y/²
STE Gradient Field f(a.3) < + %
O X-axis
2
f(z, v)
log(1 2
(an)())))
((,,5) log( 2a tamb(y)
f(x, m0
logtl 2
tanb(g.))
-2
X-axis
Figure 2: Visualization of how the straight- estimator (STE) transforms the gradient field for 16 codebook vectors for (top) (z, x² + u and (bottom) fx, y) log (1 + tanh(y)l)- The STE takes the gradient the codebook vector (9=, 9.) and "copies-and-pastes" it to all other locations within the same codebook region, forming a "checker- board""--borr" pattern in the gradient field.
function ((-). We can break down the backward pass into three terms:
OC Oc 0g de
dr 0g de dr
where 0c represents backpropagation through the decoder, By
de represents backpropagation through the vector quantization layer, and äu represents backpropagation through the encoder. As vector quantization is not a smooth transformation, de cannot be computed and gradients cannot flow through this term update the encoder in backpropagation.
dg
To solve the issue of non-differentiabiiiyyy the STE copies the gradients from g to e, bypassing vector 2g quantization entirely. Simply, the STE sets de to the identity matrix / in the backward pass:
ac Oc de
I dz 0g
dr
The first terms
0C dg combine 8C which, somewhat misleadingly, does not actually depend 3q on As a consequence, the location of within the Voronoi partition generated by codebook vector qe it close to q or at the boundary of the region-has no impact on the gradient update to the encoder.
An example of this effect is visualized in Figure 2 for two example functions. In the STE approxi- mation, the "exact" gradient at the encoder output replaced by the gradient at the corresponding codebook vector for each Voronoi partition, irrespective of where in that region the encoder output lies. result, the exact gradient field becomes "partitioned" into 16 different regions- all with the same gradient update to the encoder for the 16 vectors in the codebook.
Returning to our question, is there a better way propagate gradients through the vector quantization layer? At first glance, one may be tempted to estimate the curvature g and use this information to transform 8g de moves to This is accomplished by taking second order expansion around to approximate the value of the loss at e:
L. RJ Ey + (V 4) + -) (V: -))
4

---

Published as a conference paper at ICLR 2025
Then compute the gradient the point instead of q up second order approximation with:
ac à
de ~ 4 + (V, a) (e s)
de
- v,C + (e:
While computing Hessians with respect model parameters are typically prohíbitive in modern deep learning architectures, computing them with respect to only the codebook is feasible. Moreover as we must only compute (V) - q), one may take advantage efficient Besssian-Vector products implementations in deep learning frameworks (Dagréou al., 2024) and avoid computing the full Hessian matrix.
Extending this idea step further, through the network. Let L be the loss with the vector quantization layer and L. be the loss without vector quantization, i.e. q rather than q Q(e). Then one may form the total loss Lug + ALe. where À is a small constant like 10
ac compute the exact gradient de at by making two passes
scale down the effect of Le on the decoder's parameters and use a gradient scaling multiplier of X to reweigh the effect of L. on the encoder's parameters to 1. As Be is dinedfffrretiible, gradients from Lg will not flow to the encoder.
While seeming to correct the encoder's gradients, replacing the STE with either approach will likely result in worse performance. This is because computing the exact gradient with respect to is actually the AutoEncoder (Hinton & Zemel, 1993) gradient, the model that VAEs (Kingma & Welling, 2013) and VQ-VAEs (Van Den Oord et al., 2017) were designed to replace given the AutoEncoder's propensity overfit and difficultly generalizing. Accordingly using either Hessian approximation or exact gradients via a double forward pass will cause the encoder to be trained like AutoEreoder and the decoder be trained like VQ-VAE. This mis-match in optimization objectives is likely another contributing factor to the poor performance we observe for both methods Table 1, and a deeper analysis into these characteristics presented in Appendix A.3.
THE ROTATION TRICK
As discussed in Section 3, updating the encoder's parameters by approximating, exactly, computing the gradient at the encoder's output is undesirable. Similarly, the STE appears to lose information: the location of e within the quantized region- be it close q or far away the boundary has no impact on the gradient update to the encoder. Capturing this information, i.e. using the location of in relation to transform the gradients through de updates and an improvement over the STE.
0y could be beneficial to the encoder's gradient
Viewed geometrically, we ask how to move the gradient V,C from to e. and what characteristics of V,C and q should be preserved during this movement. The STE offers one possible answer: move the gradient from & to that its direction and magnitude are preserved. However, this paper supplies a different answer: move the gradient so that the angle between V.C and q is preserved VgG moves to We term this approach "the rotation trick", and in Section 4.3 show that preserving the angle between g and V,L. conveys desirable properties to how points move within the same quantized region.
4.1 THE ROTATION TRICK PRESERVES ANGLES
In this section, we formally define the rotation trick. For encoder output e, let Q(e) represent the corresponding codebook vector. Q(-) non-differentiable so gradients cannot flow through this layer during the backward pass. The STE solves this goghlemmmaintaiiing the direction and magnitude of the gradient V.C-as as V,C moves from to e with some clever hacking of the backpropagation function in deep learning frameworks:
q = c
constant
which parameterization of vector quantization that sets the gradient at the encoder output to the gradient at the decoder's input. The rotation trick offers a different pa- rameterization: casting the forward pass as a rotation and rescaling that aligns with g:
5

---

Published as a conference paper at ICLR 2025
Gradient at
STE
Rolation Trick
V.C
V
R
V.C
c1
constant
R is the rotation' transformation that aligns e with rescales to have the same magnitude
q and Igll Tell
V.C
R.L
as Note that both R and ||]|| [[l] are functions of e.
VIP
To avoid differentiating through this dependency, treat them as fixed constants or detached from the computational graph in deep learning frameworks when differentiating. This choice is explained in Appendix A.7.
e1
Voll
RCC
While the rotation trick does not change the output of the forward pass, the backward pass changes. Rather than set dy de / as in the STE, the rotation trick sets 8g to a rotation and rescaling transformation: de
V
Figure 3: Illustration of how the gradient at
moves to via the STE (middle) and rolation trick (right). The STE "copies-and-pawee to preserve direction while the rotation trick moves the gradient so the angle between & and v,C is preserved (proved in Appendix A.6).
0ÿ | |ell
R
the gradient
de
As result, 8g changes based on the position of in codebook partition q, and notably, the angle between VgC and 9 is preserved as VgL moves to This effect is visualized in Figure 3. While the STE translates the gradient from q to e, the rotation trick rotates it so that the angle between VgC and 9 is preserved. In sense, the rotation trick and the STE sibilings. They choose different characteristics the gradient as desiderata and then preserve those characteristics as the gradient flows around the non-differentiable vector quantization operation the encoder.
de
4.2 EFFICIENT ROTATION COMPUTATION
The rotation transformation R that rotates to q can be efficiently computed with Householder matrix reflections. We define [e] 4
À
Igll and
2+4 Then the rotation and rescaling
Tall
[el
[0+11
that aligns to q simply:
ARe
X(I 2rr + 29e""
A(e- 2r & - 29e e]
Due to space constraints, we leave the derivation of this formula to Appendix A.5. Parameterizing the rotation this fashion avoids computing outer products and therefore consumes minimal GPU VRAM. Further, we did not detect difference in wall-clock time between VQ-VAEs trained with the STE and VQ- VAEV trained with the rotation trick for our experiments in Section 5.
4.3 VORONOI PARTITION ANALYSIS
In the context of lossy compression, vector quantization works well when the distortion, or equiv- alently quantization error |/ - all low and the information capacity praacit--equivlenntti codebook utilization- is high (Cover, 1999). Later in Section 5, we will see that VQ-VAEs trained with the rotation trick have this desiderata-ond often reducing quantization error by an order of magnitude and substantially increasing codebook usage-when compared to VQ-VAEs trained with the STE. However, the underlying reason why this occurs is less clear.
rotation is defined linear transformation that RRT I, R- R². and det( R) I.

---

Published as a conference paper at ICLR 2025
Voronoi Partition
STE Updates
Rotation Trick Updates
Figure 4: Depiction of how points within the same codebook region change after a gradient update (red arrow) the codebook vector (orange circle). The STE applies the same update each point in the same region. The rotation trick modifies the update based on the location of cach point with respect to the codebook vector.
In this section, we analyze the effect of the rotation trick by looking how encoder outputs that are mapped to the same Voronoi region updated. While the STE applies the same update to all points within the same partitionn, the rotation trick changes the update based the location of points within the Voronoi region. It can push points within the same region farther apart or pull them closer together depending former capability can correspond to increased codebook usage while the latter to lower quantization error.
Change in Distance Between e and q After an Update
the direction of the gradient vector. The
STE
Let be the angle between e and q and 6 be the angle between q and V.C. When V,C and g point in the same direction, i.c. - -/22 < o < n/22 encoder outputs with large angular distance to q are pushed farther away than they would otherwise be moved by the STE update. Figure illustrates this effect. The points with large angular distance (blue regions) move further away from q than the points with low angular distance (ivory regions).
0-
1.5
$
1.0
0.5
The top right partitions of Figure 4 present an example of this effect. The two clusters of points at the boundary- with relatively large angle to the codebook vector-arg pushed away while the cluster of points with small angle the codebook vector move with it. The ability to push points at the boundary out of a quantized region and into another is desirable for increasing codebook utiliza- tion. Specifically, codebook utilization improves when points are pushed into the Voronoi regions of previously unused codebook vectors. This capability is not shared by the STE, which moves all points in the same region by the same amount.
0.0
Rotation Trick
0.5
-1,0
2
-1.5
2.0
Figure With the STE, the dis- tances among points within the same region do not change. How- ever with the rotation trick, the distances among points do change. When 6 < n/2, points with large angular distance are pushed away (blue: increasing distance). When 7/2, points pulled to- wards the codebook vector (green: decreasing distance).
When V and q point in opposite directions, i.e. T/2 < Φ < 3n/ /2, the distance among points within the same Voronoi region decreases as they are pulled towards the location of the updated codebook vector. This effect visualized in Figure 5 (green regions) and the bottom partitions of Figure 4 show an exam- ple. Unlike the STE update-that maintains the distances among points-the the rotation trick pulls points with high angular distances closer towards the post-update codebook vector. This capability is desirable for reducing the quantization error and enabling the encoder to lock on (Van Den Oord et al., 2017) to a target codebook vector.
Taken together, both capabilities can form push-pull effect that achieves two desiderata of vector quantization: increasing information capacity and reducing distortion. Encoder outputs that have large
7

---

Published as a conference paper at ICLR 2025
Table 1: Comparison of VQ- VAEs trained on ImageNet following Van Den Oord et al. (2017). We use the
Vector Quantization layer from http:///lhu..oooluuiiiaannaeeee/ wetprr////eooooouiieeeeeeeee -uwanuie--piiiiiiii..... pytorch.
Training Metries
Validation Metrics
Approach
Codebook Usage (1)
Rec. Loss
Quantization Error (1)
Rec. Loss ())
r-FiD (1)
1-18 (1)
Codebook Lookup: Euclidean & Latent Shape: 32 32
32 Codebook Size: 1024
VQ-VAE
100%
0. 107
5.9c-3
0.115
106.1
11.7
VQ-VAE w/ Rotation Trick
97%
0.116
5.10-
0.122
85.7
17.0
Cudebook Lookup:
& Shape: 32
Codebook
1024
VQ-VAE
75%
0. 107
2.9c-3
0.114
84.3
17.7
VQ-VAE w/ Rotation Trick
0.105
2.70-3
0.111
82.9
18.1
Codebook Lookup: Euclidean Shape:
Codebook Size:
VQ-VAE
100%
0.028
1.0c-3
0.030
19.0
97.3
Gumbel VQ-VAE
0.054
0.058
VQ. VAE w/ Ilessian Approx.
39%
0.082
6.9e
0.112
35.6
65.1
VQ-VAE w/ Gradients
80.4
VQ-VAE VAE w/ Rotation Trick
99%
0.028
1.4e 4
0,030
16.5
106.3
Codchook Lookup: Cosine & Latent Shape: 64 64 x & Codehook Size: 8192
VQ-VAE
31%
0.034
1.20-4
0.038
77.8
VQ-VAE w/ Hessian Approx.
37%
0.035
3.8e-5
0.037
29.0
71.5
Gradients
75.0
VAEE w/ Rotation Trick
38%
0.033
9.6e-5
0.035
24.2
83.9
angular distance to the chosen codebook vector "pushed" to other, possibly unused, codebook regions by outwards- sowardd-poiintig gradients, thereby increasing codebook utilization. Concurrent with this effect, center-poik center-pointing gradients will "pull" points loosely clustered around the codebook vector closer together, locking on to the chosen codebook vector and reducing quantization error.
4.4 FURTHER ANALYSIS
The Appendix contains several supplementary analyses. Appendix A.2 compares the rotation trick with the STE for a non-convex synthetic example; Appendix A.4 looks at the behavior far away from the origin; and Appendix A.8 analyzes the effect of using a reflection rather than rotation. Finally, Appendix A.9 examines scaling the gradient's norm by [[]]
[i]l] and explores alternatives.
5 EXPERIMENTS
In Section 4.3, we showed the rotation trick enables behavior that would increase codebook utilization and reduce quantization error by changing how points within the same Voronoi region are updated. However, the extent to which these changes will affect applications is unclear. In this section, we evaluate the effect of the rotation trick across many different VQ-VAE paradigms.
We begin with image reconstruction: training VQ- VAE with the reconstruction objective of Van Den Oord al. (2017) and later extend our evaluation the more complex VQGANs (Esser et al., 2021), 1), the VQGANs designed for latent diffusion (Rombach et al., 2022), and then the ViT-VQGAN (Yu et al., 2021). Finally, we evaluate VQ-VAE reconstructions on videos using TimeSformer (Bertasius et al., 2021) encoder and decoder. Due to space constraints. the video results are presented in Appendix A.1. In total, our empirical analysis spans 11 different VQ-VAE configurations. For all experiments, aside from handling 8q differently, the models, hyperparameters, and training settings are identical and described in Appendix A.10.
ae
5.1 VQ- VAE EVALUATION
We begin with a straightforward evaluation: training a VQ-VAE to reconstruct examples from ImageNet (Deng al., 2009). Following Van Den Oord et al. (2017), training objective is linear combination of the reconstruction, codebook, and commitment loss:
L = || ã 1 Isg(e) 412 + Blle- s9(g)112 3)]e-sg(g]1
where 3 a hyperparameter scaling constant. Following convention, we drop the codebook loss term from the objective and instead
exponential moving average update the codebook vectors.
Evaluation Settings. For 256 x 256 x 3 input images, we evaluate two different settings: (1) compressing latent space of dimension 32 x 32 x 32 with codebook size of 1024 following Yu et al. (2021) and (2) compressing to 64 x 64 x with a codebook size of 8192 following Rombach al. (2022). In both settings, compare with Fuclidean and cosine similarity codebook lookup.
8

---

Published as a conference paper at ICLR 2025
Table 2: Results for VQGAN designed for autoregressive generation as implemented in https://github.co http:///gihub.oomm
Comp Vis/taming- transformers. Experiments on ImageNet and combined dataset FFH (Karras et al., 2019)
and CelebA-IIQ (Karras, 2017) use a latent bottleneck of dimension 16 x 16 x 256 with 1024 codebook vectors.
Approach
Dataset
Codebook Usage
Quantization Error ()) Valid Loss ())
r-FID ())
1-15 (?)
VQGAN (reported)
ImageNet
7.9
114,4
VQGAN run)
ImageNet
95%
0.134
0.594
118.2
VQGAN w/ Rotation Trick
ImageNet
98%
0.002
0.422
4.6
146.5
VQGAN
FFHQ CelebA-HQ
27% 99%
0.233
0.565
4.7
5.0
VQGAN w/ Rotation Trick
FFHQ ColebA-HQ
0.002
0.313
3,7
5.2
Table 3: Results for VQGAN designed for latent diffusion as implemented in http://giitubb.oo//omm
dilent-diffusion Both settings train on ImageNet.
Appronch
I Latent Shope
Cadchook Size
Codehook Usage
Quantization Eror (
Valid oss (J)
r-18 r-FI)>())
VQGAN
8192
15%
2.5e-3
0.183
0.53
220.6
Gumbel VQGAN
x
$192
4%
0.197
0.60
219.7
VOGAN w/ Rotation Trick
64 x 64 x 3
8192
86%
L.7e44
0.142
0.27
228.0
VQGAN
x
2%
1.2e-2
0.385
141.5
Gumbel VQGAN
32 x 82 x 1
16381
12%
0,3031
1.7
189.5
VQGAN w/ Rotation Trick
32 x
16384
27%
2.4c-4
0.269
1,1
200.2
Evaluation Metries. We log both training and validation set reconstruction metrics. Of note, we compute reconstruction FID (Heusel et al., 2017) and reconstruction IS (Salimans et al., 2016) on reconstructions from the full ImageNet validation set as a measure of reconstruction quality. We also compute codebook usage, or the percentage of codebook vectors that are used in each batch of data, as measure of the information capacity of the vector quantization layer and quantization error |le 41l2 as a measure of distortion.
Baselines. Our comparison spans the STE estimator (VQ-VAE), stochastic quantization with Gumbel- Softmax (Baeyski et al., 2019), (Gumbel VQ- VAE) the Hessian approximation described in Section (VQ-VAE w/ Hessian Approx), the exact gradient backward pass described in Section 3 (VQ- VAE w/ Exact Gradients), and the rotation trick (VQ-VAE w/ Rotation Trick). All methods share the same architecture, hyperparameters, and training settings, and these settings of the Appendix. There is no functional difference among methods in the forward pass; the only differences relates to how gradients are propagated through &y during backpropagation.
summarized in Table 8
de
Results. Table 1 displays our findings. We find that using the rotation trick reduces the quantization error-someting by an order of magnitude-and and improves low codebook utilization. Both results are expected given the Voronoi partition analysis in Section 4.3: points at the boundary of quantized regions are likely pushed to under-utilized codebook vectors while points loosely grouped around the codebook vector are condensed towards it. These two features appear to have meaningful effect on reconstruction metrics: training a VQ- VAE with the rotation trick substantially improves r-FID and r-IS.
We also see that the Hessian Approximation or using Exact Gradients results in poor reconstruc- tion performance. While the gradients to the encoder are, in a sense, "more accurate". training the encoder like an AutoEncoder (Hinton & Zemel, 1993) likely introduces overfitting and poor generalization. Moreover, the mismatch in training objectives between the encoder and decoder is likely an aggravating factor and partly responsible for both models' poor performance.
5.2 VQGAN EVALUATION
Moving to the next level of complexity, we evaluate the effect of the rotation trick on VQGANs (Esssr al., 2021). The VQGAN training objective is:
LVgGAN Lver + [|sg(e) -l12 + Blle -g(g))l A
where Lper the perceptual loss from Johnson et al. (2016) and replaces the L.2 loss used to train VQ-VAEs. Adv is patch-based adversarial loss similar to the adversarial loss in Conditional GAN (Isola et al., 2017). ,3 is a constant that weights the commitment loss while À is an adaptive weight based on the ratio of V Lyer to LAd, with respect to the last layer of the decoder.
Experimental Settings. We evaluate VQGANs under two settings: (1) the paradigm amenable to autoregressive modeling with Transformers described in Esser et al. (2021) and (2) the paradigm suitable to latent diffusion models as described in Rombach et al. (2022). The first setting follows the convolutional neural network and default hyperparameters described in Esser et al. (2021) while

---

Published as a conference paper at ICLR 2025
Table 4: Results for ViT-VQGAN (Yu al., 2021) trained ImageNet. The latent shape 8 8 32 with
8192 codebook vectors. r-FID and r-IS are reported on validation set.
Approach
Codebook Usage (1)
Train Loss ())
Quantization Error ())
Valid Loss ())
r-FID (L)
1-18 (1)
VEI-YQGAN [reported]
22.8
72.9
VIT-VOGAN [ours]
0.3%
0.124
6.7e-3
0.127
29.2
43.0
VIT-QQGAN Rotation Trick
2.2%
0.113
8.3e-3
0.113
11.2
93.1
the second follows those from Rombach al. (2022). A full description of both training settings is provided in Table 9 of the Appendix.
Results. Our results are listed in Table for the first setting and Table 3 for the second. Similar to our findings in Section 5.1, we find that training VQ-VAE with the rotation trick substantially decreases quantization error and improves codebook usage, Moreover, reconstruction performance as measured on the validation set by the total loss, r-FID, and r-IS are improved across both modeling paradigms.
5.3 VIT- VQGAN EVALUATION
Improving upon the VQGAN model, Yu et al. (2021) propose using a ViT (Dosovitskiy, 2020) rather than CNN parameterize the encoder and decoder. The ViT-VQGAN uses factorized codes and L2 normalization on the output and input to the vector quantization layer to improve performance and training stability. Additionally, the authors change the training objective, adding a logit-laplace loss and restoring the Lg reconstruction error to LVQGAN.
Experimental Settings. We follow the open source implementation of http::/ggtuub..oo//mieee 123/ enhancing- transformerss and use the default model and hyperparameter settings for the small ViT- VQGAN. A complete description of the training settings can be found in Table 10 of the Appendix.
Results. Table 4 summarizes our findings. Similar to our previous results for VQ- in Section 5.1 and VQGANs in Section 5.2, codebook utilization and reconstruction metries are significantly improved; however in this case, the quantization error is roughly the same.
6 LIMITATIONS
A limitation of the rotation trick can arise when the encoder outputs codebook vectors are forced to be close to norm (i.e., ||e]| z 0 or ||all 0). In this case, the angle between e and may be obtuse. When this happens, the rota- tion trick will "over-rotate" the gradient V.C as it is transported from g to so that V.C and and v.C now point in different directions (i.e. the cosine of the angle between V.L. and V V.L. will be negative). An example visualized in Fig- ure 6.
STE
Rotation Trick
0> 90
0 > 90
V e1
Figure 6: Illustration of the rotation trick "over-rotating" vectors when the angle between e and g obtuse.
This undesirable because-when the angle between and q is obtuse-the rotation trick will violate the assumption that when e = q.
v,L= = V.L, and it will likely result in worse performance than VQ- VA-V trained with the STE. While obtuse angles between and are very unlikely-by design, the codebook vectors should be "angularly close" to the vectors that are mapped to theeerhoweve, if there restriction that forces codewords to have near O norm, then the rotation trick will likely perform worse than the STE.
7 CONCLUSION
In this work, we explore different ways to propagate gradients through the vector quantization layer of VQ-VAEs and find that preserving the angle-rather than the direction-betweee the codebook vector and gradient induces desirable effects for how points within the same codebook region are updated. These effects cause substantial improvement in model performance. Across 11 different settings, we find that training VQ-VAEs with the rotation trick improves their reconstructions. For example, training one of the VQGANs used in latent diffusion with the rotation trick improves r-FID from 5.0 to 1.1 and r-IS from 141.5 to 200.2, reduces quantization error by two orders magnitude, and increases codebook usage by 13.5x.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank Henry Bosch, Benjamin Spector, Dan Biderman, Jordan Juravsky, Mayee Chen, Owen Dugan, Sabri Eyuboglu, and the Hazy Group as a whole for their invaluable feedback and help during revisions of this work. We gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF2247015 ( (Wardware-Aware),, CCF1763315 (Beyond Sparsity), CCF1 (Volume to Velocity), and 1937301 (RTML); US DEVCOM ARL under Nos. W911NF- (Interactive Human-Al Teaming); ONR under Nos. N00014231 16333 (Deep Signal Processing); Stanford HAI under No. 247183; NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASE, Accenture, Eriesson, Qualcomm, Analog Devices, Google Cloud, Salesforce, Total, the HAI-GCP Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI), and members of the Stanford DAWN project: Meta, Google, and VMWare. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH. ONR, or the U.S. Government.
23-2-0184 (Lonexcontex) and N911NF-21-2-022
REFERENCES
Alexei Baevski, Steffen Schneider, and Michael Auli. vq-wav2vee: Self-supervised learning of discrete speech representations. arXiv preprint arXiv:1910.05453, 2019.
Jonathan Baxter. A model of inductive bias learning. Journal of artificial intelligence research, 12: 149-198, 2000.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint 0rXiv:1308.3433 2013.
Gedas Bertasius, Heng Wang, and Lorenzo Torresani. space-time attention all you need for video understanding? In ICML, volume 2, pp. 4, 2021.
Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, Clarence Ng, Ricky Wang, and Aditya Ramesh. Video generation models as world simulators. 2024. URL htpp://ppenai.coo/rrseccchh
aallo-geeeatioon--odlll---------
Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 11315-11325, 2022.
Hang Chen, Sankepally Sainath Reddy, Ziwei Chen, and Dianbo Liu. Balance of number of embedding and their dimensions vector quantization. arXiv preprint arXiv:2407.04939 2024.
Zhao Chen, Vijay Badrinarayanan, Chen-Yu Lee, and Andrew Rabinovich. Gradnorm: Gradient normalization for adaptive loss balaneing in deep multitask networks. In International conference machine learning. pp. 794 803. PMLR, 2018
Chung-Cheng Chiu, James Qin, Yu Zhang, Jiahui Yu, and Yonghui Wu. Self-supervised learning with mandom-projecctonnn quantizer for speech recognition. In International Conference on Machine Learning, pp. 3915-3924. PMLR, 2022.
Thomas M Cover. Elements of information theory. John Wiley & Sons, 1999.
Mathieu Dagréou, Pierre Ablin, Samuel Vaiter, and Thomas Moreau. How to compute hessian- vector products? In ICLR Blogposts 2024. 2024. URL bench-hyp/.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 JEEE conference on computer vision and pattern recognition, pp. 248-255. leee, 2009.
11

---

Published as a conference paper at ICLR 2025
Prafulla Dhariwal, Heewoo Jun, Christine Payne, Jong Wook Kim, Alec Radford, and Ilya Sutskever. Jukebox: A generative model for music. arXiv preprint arXivv2005.00344, 2020.
Xiaoyi Dong, Jianmin Bao, Ting Zhang, Dongdong Chen, Weiming Zhang. Lu Yuan, Dong Chen, Fang Wen, Nenghai Yu, and Baining Guo. Peco: Perceptual codebook for bert pre-training of vision transformers. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37,
pp. 552-560, 2023.
Alexey Dosovitskiy. An image worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. 1929, 2020.
Frederik Ebert, Chelsea Finn, Alex X Lee, and Sergey Levine. Self-supervised visual planning with temporal skip connections. CoRL, 12(16):23, 2017.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEF/CVE conference on computer vision and pattern recognition, pp. 12873-12883, 2021.
Tanmay Gautam, Reid Pryzant, Ziyi Yang. Chenguang Zhu, and Somaych Sojoudi. Soft convex quan- tization: Revisiting vector quantization with convex optimization. arXiv preprint arXiv:2310.03004, 2023.
Nabarun Goswami, Yusuke Mukuta, and Tatsuya Harada. Hypervq: Mir-based vector quantization in hyperbolic space. arXiv preprint arXiv:2403. 13015, 2024.
Robert Gray. Vector quantization. IEEE Assp Magazine, 1(2)4 29, 1984.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Geoffrey E Hinton and Richard Zemel. Autoencoders, minimum description length and helmholtz free energy. Advances neural information processing systems, 6, 1993.
Mengqi Huang, Zhendong Mao, Zhuowei Chen, and Yongdong Zhang. Towards accurate image coding: Improved autoregressive image generation with dynamic vector quantization. In Proceed- ings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12596-22601 2023.
Minyoung Huh, Brian Cheung, Pulkit Agrawal, and Phillip Isola. Straightening out the straight- through estimator: Overcoming optimization challenges in vector quantized networks. In Interna- tional Conference on Machine Learning, pp. 14096-14113. PMLR, 2023.
Phillip Isola, Jun- Yan Zhu, Tinghui Zhou, and Alexei A Efros. toaggeto-image translation with conditional adversarial networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1124-1134, 2017.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint 11 144, 2016.
Justin Johnson, Alexandre Alahi, and Li Fei-Fei. Perceptual losses for real-time style transfer and super-resolution. In Computer Vision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part I1 14, pp. 694-711. Springer, 2016.
Tero Karras. Progressive growing of gans for improved quality, stability, and variation. arXiv preprint arXiv:1710.00196, 2017.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEF/CV conference on computer vision and pattern recognition, pp. 4401-4410, 2019.
Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7482-7491, 2018
12

---

Published as a conference paper at ICLR 2025
Diederik P Kingma and Max Welling. arXiv:13/2.611,, 2013.
Auto-encoding variational bayes. arXiv preprint
Alexander Kolesnikov, André Susano Pinto, Lucas Beyer, Xiaohua Zhai, Jeremiah Harmsen, and Neil Houlsby. Uvim: unified modeling approach for vision with learned guiding codes. Advances in Neural Information Processing Systems, 35:26295-26308, 2022.
Adrian Lańcucki, Jan Chorowski, Guillaume Sanchez, Ricard Marxer, Nanxin Chen, Hans JGA Dolfing, Sameer Khurana, Tanel Alumäe, and Antoine Laurent. Robust training of vector quantized bottleneck models. In 2020 International Joint Conference on Neural Networks (LJCNN), pp. 1-7. IEEE, 2020.
Doyup Lee, Chiheon Kim, Sachoon Kim, Minsu Cho, and Wook-Shin Han. Autoregressive image generation using residual quantization. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 1 11523-11532, 2022.
Fabian Mentzer, David Minnen, Eirikur Agustsson, and Michael Tschannen. Finite scalar quantization: Vq-vae made simple. arXiv preprint arXiv:2309.15505 2023.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Chen. Improved techniques for training gans, Advances in neural information processing systems, 29, 2016.
K Soomro. Ucf101: A dataset of 101 human actions classes from videos in the wild. arXiv preprint 121200222 2012.
Yuhta Takida, Takashi Shibuya, WeiHsiang Liao, Chieh-Hsin Lai, Junki Ohmura, Toshimitsu Uesaka, Naoki Murata, Shusuke Takahashi, Toshiyuki Kumakura, and Yuki Mitsufuji. Sq-vae: Varia- tional bayes on discrete representation with self-annealed stochastic quantization. arXiv preprint arXiv:2205.07547, 2022.
Thomas Unterthiner, Sjoerd Van Steenkiste, Karol Kurach, Raphael Marinier, Marcin Michalski, and Sylvain Gelly. Towards accurate generative models video: A new metric challenges. arXiv preprint arXiv:1812.011177 2018
Aaron Van Den Oord, Oriol Vinyals, al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Wilson Yan, Yunzhi Zhang, Pieter Abbeel, and Aravind Srinivas. Videogpt: Video generation using vq-vae and transformers. arXiv preprint 100i 10157, 2021.
Jiahui Yu, Xin Li, Jing Koh, Han Zhang, Ruoming Pang, James Qin, Alexander Ku, Yuanzhong Xu, Jason Baldridge, and Yonghui Wu. Vector-quantized image modeling with improved vggan. arXiv preprint orXiv:2110.04622 2021.
Lijun Yu, José Lezama, Nitesh Gundavarapu, Luca Versari, Kihyuk Sohn, David Minnen, Yong Cheng, Agrim Gupta, Xiuye Gu, Alexander G Hauptmann, al. Language model beats diffusion- tokenizer is key to visual generation. arXiv preprint 130.0577 2023.
Jiahui Zhang, Fangneng Zhan, Christian Theobalt, and Shijian Lu. Regularized vector quantization for tokenized image synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 18467-18476, 2023.
Yue Zhao, Yuanjun Xiong, and Philipp Krähenbühl. Image and video tokenization with binary spherical quantization. arXiv preprint arXiv:2406.07548, 2024.
Chuanxia Zheng and Andrea Vedaldi. Online clustered codebook. Proceedings of the IEEE/CVF International Conference Computer Vision, pp. 22798-22807, 2023.
13

---

Published as a conference paper at ICLR 2025
Zixin Zhu, Xuelu Feng, Dongdong Chen, Jianmin Bao, Le Wang. Yinpeng Chen, Lu Yuan, and Gang Hua. Designing a better asymmetric vqgan for stablediffusion. arXiv preprint arXiv:2306. 04632, 2023.
14

---

Published as a conference paper at ICLR 2025
A APPENDIX
Table 5: Results for TimeSforme- trained BAIR and UCF-101 with 1024 codebook vectors. t: model suffers from codcbook collapse and diverges. r-FVD is computed the validation set.
Approuch
Dataset
Codebook Usage
Train Less (
Quantization Error (J)
Valid Loss (L)
rFVD (U)
TimeSformer
BAIR BAIR
0.4%
0.221
0.03
0.28
1661.1
TimeStomerr Rotation Trick
43%
0.074
3.0e-3
0,074
21.4
TimeSformrr'
UCF-101 UCF-101
0.1%
0.190
0.006
0.169
2878.1
TimeStormer w/ Rolation Trick
30%
0.111
0.020
0.109
229.1
A.1 VIDEO EVALUATION
Expanding our analysis beyond the image modality, we evaluate the effect of the rotation trick on video reconstructions from the BAIR Robot dataset (Ebert et al., 2017) and from the UCF101 action recognition dataset (Soomro, 2012). We follow the quantization paradigm used by ViT-VQGAN, but replace the VIT with TimeSformer (Bertasius limitations, both encoder and decoder follow a relatively small TimeSformer model: 8 layers, 256 hidden dimensions, 4 attention heads, and 768 MIP hidden dimensions. A complete description of the architecture, training settings, and hyperparameters are provided in Appendix A.10.4.
al., 2021) video model. Due to compute
Results. Table 5 shows results. For both datasets, training TimeSformer-VQGAAl model with the STE results in codebook collapse. We explored several different hyperparameter settings; however in all cases, codebook utilization drops to almost 0% within the first several epochs. On the other hand. models trained with the rotation trick do not exhibit any training instability and produce high quality reconstructions as indicated by r-FVD (Unterthiner al., 2018). Several non-cherry picked video reconstructions are displayed in Appendix A.10.4.
A.2 NON-CONVEX SYNTHETIC EXAMPLE
To supplement our analysis in Section 4.3, we include a nu- merical simulation of vector quantization for minimizing Himmelblau's s function (Figure 7) across 100 gradient up- dates for the STE and rotation trick gradient estimators highlight the differences in their behaviors. Our simulation uses an EMA with a decay rate 0.8 as described in Van Den Oord et al. (2017) to update the codebook vectors and a learning rate 1e 3 to update the pre- rre-quantzzed points. Points for both the STE and the rotation trick simulation use the same random initialization for both codewords and pre-quantized vectors. The only difference is whether the STE or the rotation trick is used as the gradient estimator through the vector quantization operation.
f(. = (++ 1112
(
100
250
200
130
100
2
Figure 8 visualizes our results after 33, 66, and 100 gra- dient updates. The orange circles represent codebook vec- tors, the green dots the initial points, and the blue dots the updated points. Contour lines are drawn each diagram to indicate regions of equal loss, with blue representing regions of low loss and red indicating regions of high loss. Similar to findings in Section 5, we see that the rotation trick clusters points more tightly around each codebook vector when compared to the STE, resulting in lower dis- tortion. Moreover, the codebook vectors more rapidly converge to the four equal local minima in Himmelblau's function, resulting in lower objective function value when averaged across all points.
N-avis
4
4
Figure 7: Loss surface for Himmelblau's function. Himmelblau's function has four equal local minima: f(3.0. 2.0) 0.0. ʃ(- 2.8.., 3..... ...) = 0.0, f(-3.7.., -3.2..) 0.0, and f3.5.., -1.8..) = 0.0..
A.3 HESSIAN APPROXIMATION AND EXACT GRADIENT ANALYSIS
In this section, expand analysis in Section 3 and offer some intuition for why using exact gradients, or a Hessian approximation of the exact gradients, may convey undesirable characteristics. We begin by showing the Hessian approximates the exact gradient up second order term with
15

---

Published as a conference paper at ICLR 2025
Straight-Through Estimator
O Updates
33 Updates
66 Updates
100 Updates
The Rotation Trick
35
Synthetic experiment for minimizing Himmelblau's aus function with vector quantization using the STE gradient estimator (top row) and the rotation trick (bottom row). rotation trick more quickly converges to these minima and achieves substantively lower distortion between codewords and pre-quantized points.
Figure 8:
Taylor series expansion. We can write the loss L. exactly as an infinite series of around q:
C. = C, + (V + (e (=" (V) q) + (e'''''V' - q, -e--gl 4) +
so that the loss computed by the Hessian approximation differs from the loss computed with the exact gradients method by the remainder term from truncating the Taylor series expansion after the second term:
(Le 1esssian L, + (V ②) 2 (e-
(c-
((e--)=(V220 (V? 9)
When differentiating both of these losses to compute the gradients, the difference between the exact gradient update and the Hessian update is:
aL. ah.
O(lle - -1
Hessiim
de
de
de
where
1
O(lle -11²)=
''' L' g,e 9) +
16

---

Published as a conference paper at ICLR 2025
Loss Surface
STE Gradient
Rotation Trick Gradient
Hessian & Exact Gradient
Top-Left Partition Loss Surface
Top-Right Partition Loss Surface
Bottom Partition Loss Surface
30
12.2
22.8
Paa
9900
Figure 9: Examples of how the gradient can change due to the presence of negative curvature or an indefinite Hessian. As the loss in each partition is quadratic, the exact gradient will equal the Hessian approximation. Notice that when y As e, V,f V.L for both the STE the rotation trick. As the Hessian approximation and exact gradients use the curvature of the loss surface to move V, from g to e, the direction of the gradient can change substantively, even when == e.
The Hessian idea described in Section approximates the exact gradients to the encoder as quantization did not occur, i.e. it approximates the gradient used to update the encoder in the original AutoEncoder (Hinton & Zemel, 1993) model.
We now explore some instances where the exact gradients, or their Hessian approximation, may produce undesirable behavior vector quantization. An inductive bias (Baxter, 2000) for vector quantization to work well is that when is "close" to g, their gradients are also "close". i.c. if e then VeL R V,C. Intuitively, if the distortion between e and is small-ie. is a very good codeword for chenen these points should move together during a gradient update. If they do not, the distortion would increase.
This assumption holds for both the STE and Rotation Trick gradients; however, it can be violated by the Hessian approximation or the exact gradient approaches, especially when the curvature around negative or the Hessian is indefinite and forms saddle point.
Figure 9 illustrates three such cases. As both the STE and Rotation Trick do not use the loss surface to move VgC from q to e, when q = e, VgL= z V.L. However, approaches that use the curvature around such as the Hessian approximation or exact gradients, to either find or approximate the loss e can have V.C point in very different direction from V,C. even when q close to e. The lef-left and bottom partitions of Figure 9 scatter the gradients as they move from to the points in these partitions due to negative curvature. A similar effect occurs the top-right partition of Figure due to the presence of saddle point.
A.4 BEНAAIIOR AWAY FROM THE ORIGIN
Unlike the STE, the rotation trick is not invariant to the location of the origin. In this section, explore this characteristic and its effect on how points within the same Voronoi region are updated. For example, suppose each codebook vector and encoder output in Figure 4 were shifted by some
17

---

Published as a conference paper at ICLR 2025
Voronoi Partition
STE Updates
Rotation Trick Updates
Figure 10: Depiction how points within same codebook region change after gradient update (red arrow) at the codebook vector (orange cirele) when all points are far from the origin. The STE is invariant to the this translation; however as the angle between and decreases these vectors translated away from the origin, the effect of the rotation trick will decrease. In the limit, the rotation trick reduces to the STE.
constant vector so that each now has all positive components. How would this affect the rotation trick's 's gradient estimator?
Consider one codebook vector and one encoder output separated by angle 0. We define + and è e + where d is some large displacement vector. Let be the angle between I and é. We visualize this example in Figure 11. From the law cosines:
d
1q - e11 =lo/1 + ||e||2 [|llllle|||co///////
d
and
6
l1q 211- 11 el 2 = + ||@|| 2 ||| cos
Substituting, we find that
Illustration of code- book and encoder output shifted away from the origin by a con- stant vector d The angle after the shift is smaller than the angle be- fore the shift: < O.
Figure 11:
cos ||q1| + |e|| 2||a| ||l|| cos(() - |1q + d]| lee + d]2
21|g
d|l[le + d||
and consider the case when & and are far from the origin, i...||d| >> [|q]]. ||el]. Then we have:
2]|d]] 2
cos
2
So as d 0o, 0 This implies that points move away from the origin, the rotation trick smoothly transforms into the STE.
1ill 1 and Rt + I, which is exactly the STE update. As [|l
We visualize example of this effect in Figure 10, where each point from Figure 4 is translated by positive ten along each dimension. As illustrated above, the effect for the "push" gradient the top-right quadrant remains but it's effect is reduced, i.e., more similar to the STE update. The top-left partition becomes a "pull" because the gradient now points towards the origin, so points within this region move closer together. Finally, the gradient in the bottom region longer points towards the origin, but is now more orthogonal to the codebook vector. As a result, we see more of a rotation applied to the points in this region than the contraction that is depicted in Figure
18

---

Published as a conference paper at ICLR 2025
A.5 HOUSEHOLDER REFLECTION TRANSFORMATION
For any given and the rotation R that aligns with in the plane spanned by both vectors can be efficiently computed with Householder matrix reflections.
Definition (Householder Reflection Matrix), For a unit norm vector a E R.,, I 2aa E Rddd reflection matrix across the subspace (hyperplane) orthogonal
Remark 1. Let a, b E TRd that define hyperplanes at and b- respectively. Then reflection across a- followed by a reflection across b- is a rotation of in the plane spanned by o, b where 0 is the angle between a, b.
Remark 2.
Let a,b Rd with ||a| 1b11 1. Define c
| as the vector half-way
between a and that L(a 0 and L(a,c) (bb
From Definition 1, (I 2ec
encodes a reflection across et and (I 2bb²") encodes reflection across b-. From Remark (T-2MF (T 2oe"") then corresponds As the span(b, c) span(a, b), (I 2bb (I 2cc"") corresponds rotation of B in the plane spanned by and b. Therefore, (I - 26b 26b")(1 2ec") ) =
rotation of 2(5) 0 in the plane spanned by and
Returning to vector quantization with ||||| ItJe, we can write It as the product of two Householder reflection matrices that rotates in the plane spanned between them. Without loss of generality, assume e and y are unit norm, and let be the angle between and q. Setting = c+g and simplifying yields:
Tetgl
R = (1 12qq (1 2rTT)
299" -2r + 499" rr'
= - 290" -2rT" + Ag [q r] ,
T C c+0 |e + 91l
= I - 29q" -2rr" + 44q
r
9" + q"
1 - -2g9 -Tr + 449
r
lee + 9l
||allle|| cos 0
|19
|19
I 1-299 2rT + 44q
7
le 9
cos 0
1 - 2g9 -Tr + 4g
le + q
|e + all 2
1 - -2299 - + 4g
2 + 411
4|le + gll
= 2q9 - -Tr +
gr
2|le + gll
4/le + all
I - 299 -
g(e g)
2|le + gll
= I - 229- -2rr T + 29e T + 2gg"
I - 2r 2ge
A.6 PROOF THE ROTATION TRICK PRESERVES ANGLES
For encoder output e and corresponding codebook vector g. we provide formal proof that the rotation trick preserves the angle between V,C and q V,C moves to Unlike the notation in the main text, which assumes E Rdx1, we use batch notation in the following proof to illustrate how the rotation trick works when training neural networks. Specifically, E lixx and e dxxdxd where b is the number of examples in a batch and d is the dimension of the codebook vector.
Remark 3. The angle between q and V v,Lh is preserved V vgL. moves to e.
19

---

di Published as a conference paper at ICLR 2025
Proof. With loss of generality, suppose ||e|| ||q/l = 1. Then have
eR.
dq
R
de
The gradient at will then equal:
[Qq
V V.C
de
v [R]
Let be the angle between and v,L and & be the angle between and V,C. Via the Euclidean inner product, we have:
| IV, cos q VV,
eR' [V,C]'
e [V,CR)
=e eVVeL]"
1|V cos ạ
so 0 o and the angle between q and VgC is preserved as V Vg moves to
A.7 TREATING R AND |all
Te] AS CONSTANTS
In the rotation trick, we treat R and |gll| [e] as constants and detached from the computational graph
during the forward pass of the rotation trick. In this section, we explain why this is the case.
The rotation trick computes the input to the decoder after performing non-differentiable codebook
lookup on e to find It is defined as:
191 Tell
=
Re
As shown in Section 4, R is function of both and q. However, using the quantization function ||ll and single function of [e
(e)) g. we can rewrite both
[[(())l ||e]|
e Q(e) le 2(e)
e Q(e) T [(e)]
f(e)
- 2
+2@(E)e
|g||
R
Jell
The rotation trick then becomes
f(eee
and differentiating q with respect to e gives us:
0ğ de
f' f(e)
However, (c) cannot be computed would require differentiating through Q(e), which is non- differentiable codebook lookup. We therefore drop this term and use only f(e) as our approximation of the gradient through the vector quantization layer: 3e f(e). This approximation conveys more information about the vector quantization operation than the STE, which sets dg T.
20

---

Published as a conference paper at ICLR 2025
Gradient
STB
Rotation Trick
Reflection Trick
F C
C
Co
R.C.
AGE
aut
41
F,C
JGC
T R,C
W.LC
N.C
41
V.
7C
AV.CC
F.C
Illustration of how gradient at q moves to E via STE, rotation trick, and reflection trick. The reflection trick matches the behavior of the rotation trick when the gradient V.C is parallel to However, it will reverse the components of the gradients orthogonal to for points in q's partition. This effect is illustrated in the bottom two rows of the rightmost column.
Figure 12:
Voronoi Partition
STE Updates
Reflection Trick Updates
-3
Figure 13: Depiction of how points within the same eodebook region change after gradient update (red arrow) the codchook vector (orange cirele). The STE applies the update each point in the same region. The reflection trick (Appendix A.8) modifies the update based on the location each point with respect to codebook vector. Note the top-left region of the reflection trick update, where the points actually move in the opposite direction of the gradient update.
A.8 THE REFLECTION TRICK
One may also use single reflection align q, rather than rotation. For instance, using the c-g and reflecting across the plane orthogonal to this notation from Appendix A.5, setting r = -
vector via the Houscholder reflection (I - 2rr"" will reflect to We denote this reflection as i so that q le|
|g]| Re. We call this approach "the reflection trick."
21

---

Published as a conference paper at ICLR 2025
The reflection trick can result in undesirable behavior during the backward pass. While it replicates the rotation trick when VgLi is parallel to q, as illustrated in the top two rows of Figure 12 and the top-right and bottom regions of Figure 13, it reflects orthogonal components of the gradient across the hyperplane orthogonal to e q so that these components are reversed. Simply, if the quantized gradient points "left" then the reflected gradient will point "right", and vice-versa. This behavior is undesirable for points with low distortion, R q. because it will cause e to move away from along the components of the gradient orthogonal to q, thereby increasing distortion for two points that are a "good match". The top-left partition of Figure 13 illustrates one such example. In this case, the gradient pushes the codebook vector "left" while the points in this region are pushed in the opposite direction of the gradient.
We evaluate this effect experimentally following the VQ- VAE evaluation paradigm from Table and the VQGAN evaluation paradigm from Table 3. While we did not train these models to completion due GPU resource limitations, both paradigms exhibited poor convergence when trained with the reflection trick. Specifically, after one epoch, the validation Ioss was approximately 3x higher than the rotation trick for both 8192 and 16384 codebook VQGANs Table 3. For the Euclidean codebook model with latent Shape 64 64 x in Table 1, the validation loss was approximately 2x higher than the rotation trick after 15 epochs.
A.9 GRADIENT NORM SCALING IN THE ROTATION TRICK
In this section, we analyze the effect of the ||e|| term in the rotation trick. While this norm rescaling [|l]
necessary to transform into q during the forward pass, one could avoid the multiplicative factor by instead formulating the rotation trick as:
R ConsEant
(q Re)
possible benefit of this latter formulation is that dg R. an orthogonal transformation with ||g|| In this section, we analyze determinant one that does not shrink or expand space by factor of [[]
de
the differences between these two approaches and formulate both specific instantiations of more general family of rotation-based gradient approximations.
A.9.1 COMPARISON BETWEEN |ggl|
AND (q - Re)
||e]|
An inductive bias of vector quantization is that when e 2 q. then V R VgL. Simply, when the distortion between e and i is small, the gradient for both and q should be approximately the same. However when ||e|| & 0 and a Euclidean metric is used to determine the closest codebook vector, the angle between and q can be obtuse as illustrated in Figure 6. In this instance, the rotation trick will cause the gradient V.L to "over-rotate" and point away from VgC.
Using a grad scaling of |aهl| can fix this. When ||l|| z 0 and ||el| ||9//. the norm of the gradient [|[]|
will be scaled up to push away from the origin. Pushing away from the origin makes the angle between and q more of a factor when computing the Fuclidean distance:
||e wl1 = / ||e| ||al2 2||e|||lall cos
so e is more likely to map different y that forms acute angle with it as ||e|] increases.
Now consider if [|ql] ss 0 and ||e|| > ||/||- When this occurs, the update to e will vanish because Iall Tell R 0. This behavior may also be desirable because when is close to the origin, there's higher likelihood the angle between and g would be obtuse.
We also explore this factor in ablation experiments for VQ-VAE and VQGANs. Table 6 mirrors Table and summarizes our findings for VQ-VAES while Table 7 mirrors Table and summarizes our findings for the VQGANs used in latent diffusion. In Table 6, we do not observe a difference
between using q TT|
|gll| Re and q Re + (q Re). However, for the VQGAN results in Table 7, we find that using the grad scaling factor modestly improves performance.
22

---

Published as a conference paper at ICLR 2025
Table 6: Comparison of the rotation trick using q
[llll Re with using q = Re + (g Re) for VQ-VAI [|e]
models, The experimental setting follows Table
Rotation Trick Function
Training Metrics
Validation Metrics
Codebook Usage (1)
Rec. Loss (
Quantization Error (U)
Rec. Loss (U)
r-IDD (1)
r- (1)
Codebook Lookup: Euclidean & Latent Shape: 64 3 & Codebook Size: 8192
[cl
Nel Re
99%
0.028
1.4c4
0.030
16.5
106.3
Re (q Re)
100%
0.028
4.0c-4
0.030
16.5
106.1
Table 7: Comparison of the rotation trick using q Ilal.
Re with using q = Re (g Re) for VQGAN models.
[|eT
The models with codebook size were stopped after 2 epochs while the models with codebook size 16384 were stopped after 3 epochs.
Rotation Trick Function
Latent Shape
Codebook Size
Codehook Usage
Quantization Error (J)
Valid Loss (J)
r-FID (.)
r-18 (f)
Re te Re)
64 x 64 3
8192
45%
4.0c-4
0.161
0.46
225.0
1.5e-3
Lal
x
16394
3.3e4
196.1
9.40-
A.9.2 GENERAL FAMILY OF ROTATION- SHRATION-BASED GRADIENT ESTIMATORS
Generalizing the additive and multiplicative formulations of the rotation trick, we formulate both as specific instantiations a more general family:
(c) Re - Re)
where ?(e) determines the multiplicative scaling factor. For q
Iglll ||e||
Re, r(e)
lllll [|ell
and for
q = Re + (q (()) = L. However, one can explore other scaling factors such as
(e)
8|14 ell
We visualize the gradient fields for different formulations of 7(e) in Figure 14.
It is almost certain that other formulations of 7(() from the improve the training dynamics or performance of VQ-VAEs. In particular, a priori fixing (e)) to satisfy an inductive bias developing an adaptive scaling factor that dynamically sets <(e) similar the functions that adapt task weights in multi-task learning throughout training (Kendall et al., 2018; Chen et al., 2018) are exciting directions for future work.
explore in this work would
A. TRAINING SETTINGS
We detail the training settings used in our experimental analysis in Section 5. While text description can be helpful for understanding the experimental settings, our released code should be referenced to fully reproduce the results presented in this work.
A.10.1 VQ-VAE EVALUATION.
Table 8 summarizes the hyperparameters used for the experiments in Section 5.1. For the encoder and decoder architectures, we use the Convolutional Neural Network described by Esser et al. (2021), The hyperparameters for the cosine similarity codebook lookup follow from Yu et al. (2021) and the hyperparameters for the Fuclidean distance codebook lookup follow from the default values set in the Vector Quantization library from https://github.co https///gihbb..co///ddaaaa com//ucidrains/vctton qurioor--uuizzeee-- prtantie--pytrch... All models replace the codebook loss with the exponential moving average described in Van Den Oord al. (2017) with decay 0.8. The notation for both encoder and decoder architectures is adapted from Esser et al. (2021).
For the Gumbel VQ- VAE baseline, follow the implementation of https://github.c https://ithubb.co//aaathhh cow./arpath///////oaoooaooooaaat deep-vector wet--vecorrquuutiieee quantization and use the suggested schedule to attenuate the softmax temperature from 1.0 over the course of training. Aside from the difference in quantization, i.e. deterministic
23

---

Published as a conference paper at ICLR 2025
1191 ||c||
Ra.a) 2.
i
Re
5 Re
Re)
y-hii
A-axts
X-asis
kel tanbigl))
2
Figure 14: Visualization of how different choices (e)) rotation trick affect gradient field (top) f(z, v) x² + ' and (bottom) fx, y) = (og(1 (1 + tanh(y))). prevent cluttered visualizations, the maximum and minimum gradient norms are capped within the gradient field.
Table 8: Hyperparameters for the experiments Table 1. (1024, 32) indicates a model trained with codebook size of 1024 and codebook dimension of 32. Similarly, (8192, 3) indicates model trained with codebook size of 8192 and codebook dimension of
Cosine Similarity Lookup
Euclidean Lookup
(1024, 32)
(8192,
(1024, 32)
(8192,
Input size
256 x 256 x 3
256 256 3
256 256 3
256 256 x
Latent size
16 16 x
64
16 16 x
x 64
(commitment loss coefficient)
1.0
1.0
1.
1.0
encoder/decoder channels encoder/decoder channel mult.
128
128
128
128
[1, 1, 2, 4]
[1, 41 256
[1, 1, 2, 256
[1, 4] 256
[Effective] Batch size
256
Learning rate
1 X 10
1 10
5 10-6
x 10
Weight Decay
10 1024
1 10 10
0 1024
0
Codebook size Codebook dimension
8192
8192
32
Training epochs
25
20
25
20
versus stochastic, the architecture and optimization of the Gumbel VQ-VAE model are identical to the VQ-VAE baseline.
A.10.2 VQGAN EVALUATION
Table 9 summarizes the hyperparameters for the VQGAN experiments in Section 5.2. For the Gumbel VQGAN model, we follow the default hyperparameters and settings from Rombach al. (2022). Non-cherry picked reconstructions for the models trained in Table 2 and Table 3 are depicted in Figure As indicated by the increased r-FID score, the reconstructions out by the VQGAN trained with the rotation trick appear to better reproduce the original image, especially fine details.
A.10.3 VIT-VQGAN EVALUATION
Our experiments in Section 5.3 use the VIT-VQGAN implemented in the open source repository https//ggitub..co//maan hwwpaa//w//eeeeaaannnnnnannnnnnn The default hyperparameters follow those
24

---

Published as a conference paper at ICLR 2025
VQGAN from Taming Transformers
VQGAN from Latent Diffusion
ImageNet
FFI & CelebA-HQ
ImageNet [f-8]
[f-4]
Orig
STE
ROT
Orig
STE
ROT
Orig
STE
ROT
STE
ROT
Figure 15: Non-cherry picked reconstructions VQGAN results Table 2 and Table ROT an abbreviation for the rotation trick.
Table 9: Hyperparameters for the experiments in Table and Table 3. We implement the rotation trick in the open source https://githb. com/Comp Vis/tamingttranfformer for the experiments in Table and implement the rotation trick in https://github.c http://gghhuubooo//oooooe com/Comp Vis/latent-diffusion for Table 3. In both settings, we use the default hyperparameters. t: 18 epochs for ImageNet and 50 epochs for FFHQ Ce..bA-HQ.
transformerss
Table VQGAN
Table 3 VQGAN
Table VQGAN
Input size Latent size
256 x 256 x 3
256 x 256 x 3
256 x 256 x 3 32 x 32 x 1.0
16 x 16 x 256 1.0
64 x 64 x
Codehook weight
1.0
Discriminator weight
0.8
0.75
0.6
encoder/decoder channels
128
128
128
encoder/decoder channel mult.
2,,1,
[1,2,
[1,2,2,4
[Effective] Batch size
48
16
16
[Effective] Learning
4.5 10-6
4.5 x 10
4.5 10-6
Codebook size
1024
8192
16384
Codebook dimensions
256
3
4
Training Epochs
18/50'
specified by Yu et al. (2021), and our experiments use the default architecture settings specified by the VIT small model configuration file.
We depict several reconstructions in Figure 16 and see that the ViT-VQGAN trained with the rotation trick is able better replicate small details that the VIT- ViT-VQGAN trained with the STE misses. This expected as the rotation trick drops r-FID from 29.2 to 1 11.2 as shown in Table 4.
A.10.4 TIMESFORMER VIDEO EVALUATION
We use the Hugging Face implementation of the TimeSformer from https:/huuginggtee co/does/ transformers/en/modde httpr///ee//ooodeddod/tteesrnroo doc/timesformer and the ViT-VQGAN vector quantization layer from hitps: /github.cco com/too////////o comm/huanz// /23/ehhanci/ng transformers. We loosely follow the hyperparameters listed in Yu et al. (2021) and implement a small TimmSformer encoder and decoder due to GPU VRAM constraints. We reuse the dataloading functions of both BAIR Robot Pushing and UCF101 dataloaders from Yan et al. (2021) at https://github. wwwpght....o/wwwleeeeee com/wilson yan Van/VideoGPT. A complete description of the settings we use for the experiments in Appendix A.I are listed in Table 11.
We also visualize the reconstructions for the TimeSformer- VQGAN trained with the rotation trick and the STE. Figure 17 shows the reconstructions for BAIR Robot Pushing, and Figure 18 shows the
25

---

Published as a conference paper at ICLR 2025
ViT-VOGAN
Table 10: Hyperparameters for the experiments in Table 4.
Orig
STE
ROT
ViT-VQGAN Settings
Input
256 256
Patch size
Encoder Decoder Hidden Dim
512
Encoder / Decoder MI.P Dim
1024
Encoder Decoder Hidden Depth
8
Encoder Decoder Hidden Num Heads
Codebook Dimension
32
Codebook Size
8192
Codebook Loss Coefficient
1.0
Log Laplace loss Coefficient
0.0
Log Gaussian Coefficient Perceptual loss Coefficient Adversarial loss Coefficient [Effective] Batch size
1.0
0.1
0.1
32
Non-cherry picked reconstructions for VIT-VQGAN results in Ta- ble 4. ROT is an abbreviation for the rotation trick.
Figure 16:
Learning rate
1 10
Weight Decay
1 10-4
Training epochs
10
STE Reconstructions
Rotation Trick Reconstructions
Original Video
Figure 17: BAIR Robot Pushing reconstruction examples. While the model trains on 16 video frames at time, we only visualize 4 at time in this figure. The model trained with the STE undergoes codebook collapse, using 4 out of the 1024 codehook vectors for reconstruction and therefore crippling the information capacity of the vector quantization layer. On the other hand, the VQ- VAE trained with the rotation trick instead uses an average of 441 of the 1024 codebook vectors in each batch of example videos.
reconstructions for UCF101. For both datasets, the model trained with the STE undergoes codebook collapse early into training. Specifically, it learns to only use 1021 of the available codebook vectors for BAIR Robot Pushing and 2 for UCF101 in a batch of input examples. Small manual tweaks to the architecture and training hyperparameters did not fix this issue.
2048
In contrast, VQ-VAES VAEs trained with the rotation trick do not manifest this training instability, Instead, codebook usage is relatively high-ap at 43% for BAIR Robot Pushing and 30% for UCF101- and the reconstructions accurately match the input, even though both encoder and decoder are very small video models.
A.1 CREATION OF VORONOI REGION FIGURE
In this section, we describe the creation of Figure 4 as well as the other figures that use this format. For the tiphright and bottom partitions, fix the codebook pre-quantized points from four different Gaussian distributions. For the pre-quantized points in the lep-left partition, manually set them form crescent shape around the codeword.
set of preset values and sample
26

---

Published as a conference paper at ICLR 2025
STE Reconstructions
Rotation Trick
Original Vidco
Reconstructions
Figure 18: UCF-101 reconstruction examples. While the model trains on video frames time, we only visualize at time in this ligure. The model trained with the STE undergoes codebook collapse, using approximately 2 out of the 2048 codebook vectors for reconstruction and therefore crippling the information capacity of the vector quantization layer. The VQ- trained with the rotation trick instead uses average of 615 of the 2048 codebook vectors in each batch of example videos.
Table 11: Hyperparameters for the experiments in Table 5. A TimeSformer (Bertasius et al., 2021) is used for the Encoder and Decoder architecture implemented at https//hugginggfcce http:/ugggggggc.ccooocsssess/see transformerssen/moodd doc/timesformer. The vector quantization layer between Encoder and Decoder follow from Yu al. (2021) implemented in https://gtthu..comm/ wtt.///haaaan/////e/aaaaaaaaaaaa transforrrrs....
TimeSformer- VQGAN Settings
BAIR Robot Pushing
UCF101 Action Recognition
Input size
16 x 64 x 64 x 3
16 x 128 x 128 x 3
Patch size
2
4
Encoder Decoder Hidden Dim
256
256
Encoder Decoder MLP Dim
768
768
Encoder Decoder Hidden Depth
Encoder Decoder Hidden Num Heads
Codehook Dimension
32
32
Codehook Size
1024
2048
Codebook Loss Coefficient
1.0
Log Laplace loss Coefficient
0.0
0.0
Log Gaussian Coefficient Perceptual loss Coefficient Adversarial loss Coefficient [Effective] Batch size Learning rate
1.0
1.0
0.1
0.1
0.1
0.1
24
20
1 10
4.5 x 10-6
Weight Decay
x 10
x 10 3
Training epochs
30
We similarly fix constant gradient vectors for each partition, and apply them to the pre-quantized points after transformation by the STE, i.c. simply moving the gradient to each pre-quantized point in the quantized region, by the rotation trick, i.e. rotating the gradient based on the angle between the pre-quantized point and closest codebook vector and rescaling appropriately. We multiply the gradient by small constant-the learning rate-and then apply the gradient to each pre-quantized point. We repeat the above 25 times, at each point re-computing the angle and magnitude between the pre-quantized point and the codebook vector for the rotation trick update. For simplicity, do update the codebook vectors themselves or recompute codebook regions throughout the numerical simulation.
A.11 COMPARISON WITHIN GENERATIVE MODELING APPLICATIONS
Absent from our work is an analysis on the effect VQ- VAEs trained with the rotation trick on down-stream generative modeling applications. We this comparison as outside the scope of this
27

---

Published as a conference paper at ICLR 2025
work and do not claim that improving reconstruction metrics, codebook usage, quantization error in "Stage 1" VQ- VAE training will lead to improvements in "Stage 2" generative modeling applications.
While poor reconstruction performance will clearly lead to poor generative modeling, recent work (Yu et al., 2023) suggests that-at at least for autoregressive modeling of codebook sequences with MaskGit (Chang al., 2022)-th the connection between VQ- VQ-VAE reconstruction performance and downstream generative modeling performance is non-linear. Specifically, increasing the size of the codebook past a certain amount will improve VQ-VAE reconstruction performance but make geneative modeling of codebook vectors more difficult.
downstream likelihood-
We believe this nuance may extend beyond MaskGit, and that the desiderata for likelihood-baad generative models will likely be different than that for score-based generative models like diffusion. is even possible that different preferences appear within the same class. For example, left-to-right autoregressive modeling of codebook elements with Transformers (Vaswani, 2017) may exhibit different preferences for Stage VQ-VAL VAE models than those of MaskGit.
These topics deserve a deep, and rich, analysis that we would find difficult to include within this work as our focus is on propagating gradients through vector quantization layers. As result, we entrust the exploration of these questions to future work.
A.12 GRADIENT ESTIMATORS AS PARALLEL TRANSPORT
In this section, we analyze the STE and the rotation trick through the lens of differential geometry. specifically as the parallel transport of the gradient v,L vector from the codeword to the encoder output e. For this analysis in this section, we only consider the rotational component Ro of the rotation trick, not the rescaling by || [e]l
A.12.1 BACKGROUND ON HYPERSPIIERICAI COORDINATES
Hyperspherical coordinate systems are ubiquitous in applications of math and physics, where certain formulas become greatly simplified by parameterizing the location of points by the radius and angles to coordinate axes. An familiar instantiation of the hyperspherical coordinate system may be polar coordinates with radial component and polar angle 0:
x = r cos 0
У sin
or the instantiation of the hyperspherical coordinate system for three dimensions, otherwise known as spherical coordinates, with radial component r, polar angle 0 and azimuthal angle di
cos 0
= sin 0 cos d
sin 0 sin ø
More generally, hyperspherical coordinates composed by radial coordinate and d 1 angular coordinates ..... 0a-1 where .1,...0--22 are supported over [O. while 0 1 ranges from [o. 2x.. We outline one common conversion from Cartesian coordinates hyperspherical coordinates below, and other conversions are equivalent up to permutation of the coordinate axes:
cos
F2 sin(O1) cos(O2)
sin(()) sin((02) cos(O2)
Ød-1 = =r sin (1) sin(Au-2) cos cos(Oa-1)
Fa sin)())) sin(As-2) 2) sin(As-1)
28

---

Published as a conference paper at ICLR 2025
Cartesian Coordinate System
Spherical Coordinate System
bv
Figure 19: Visualization of basis vectors different points under Cartesian (left) and spherical (right) coordi- natate systems. Notice that the Cartesian basis vectors do not change from point-to-point; however, the spherical basis vectors change in both direction and magnitude. Even at the same radius, the & coordinate changes based on the azimuth angle & because the same infinitesimal change in & will result in longer (or smaller) change in arclength depending on the radius of the circle at latitude 0.
3e
and the reverse transform from Cartesian coordinates to hyperspherical coordinates:
r (x1)² + (x2)² + ... + (ra)²
0 arctan 2( / (xn) +... + (
02 = arctan 2( v (Ia)² + + (xr)² (73)2,33) ,22)
4 arctan 2 2(V (Ira2 + (Id-1)) 2d
0 = arctan 2( / (xa)² x.-11)
where arctan 2(x,, y) returns the angle measurement radians over the support (- between between a and y.
Unlike the Cartesian coordinate system, the hyperspherical basis vectors are not identical over the entire space; they change with position. For instance, moving outwards along r will increase the 5gt as an infinitesimal change in 0¹ will now cover larger arclength distance-i.e. the line segment traveled by changing the angle 8'=than than that same infinitesimal change with smaller r. This effect is visualized for three dimensions in Figure 19.
length of 0
At any given point in hyperspherical coordinates ÿ. the transformation from Cartesian basis vectors à
8 15 Tg "zz. to hyperspherical basis vectors Dr T0 can be computed with the multivariate chain rule:
i 00
Or θ
00
Fx
Ox' where 76. can be computed from the coordinate transform functions, ie. 21 = roos(f)). It is typical express these relationships in matrix that transforms arbitrary vector in Cartesian coordinates
29

---

Published as a conference paper at ICLR 2025
point p to its counterpart in hyperspherical coordinates
Ox
0x-
Dr.
86
0y4-T
dy
On.
8 Lar
8 Barit
dr
00
20
Ø Br
8
8
Ox 4
8x d
0r
88
003-1
The Jacohian
As illustrated in Figure 19, J does necessarily have determinant equal to and changes function of position, so the norms of the basis vectors spanning the hyperspherical tangent space change based on position. More generally, this notion of distance captured by the line element: the length of line segment resulting from an infinitesimal change along the coordinate axes. The Cartesian line element is given by:
ds
(dx*) + (dr²)² + + (2x²))
while the hyperspherical line element is:
ds² dr² + x² (dd + r sin² 01 (d0))) +
II sin? 0, (dod
which reflects that distance traveled by small changes in the hyperspherical coordinates "increases" with increasing radius and "decreases" with distance from the equator. To ensure that the norm of the basis vectors does not change during conversion, it is common renormalize hyperspherical basis vectors to have unit norm for all points. However, a notion of norm not defined a priori for hyperspherical vectors; the metric tensor imposed on this space defines the inner product which in turn defines a sense of arclength.
Using the induced metric from Cartesian coordinates, coordinates on the hyperspherical coordinate system by expressing hyperspherical basis vectors as linear combination of Cartesian basis vectors and then computing the norm of this resulting vector in the Cartesian tangent space:
inherit the inner product from Cartesian
8
0
00
a0
Dxk
Oxt ô ag Oxi
agi Ox
Oak Oak
0
30 Off
Dark
Oic
Oxr Og )2
The first fundamental form gives us the normalization constants:
2
0
2
0
2 sin
0
I
0
sin
i=
as the diagonal represents the inner product
( (0 ag and we would like renormalize each basis
vector to have unit norm:
||
V
à 2 ). Therefore, our normalized hyperspherical basis 09:
50
vectors 8 a
become:
87 99,
30

---

Published as a conference paper at ICLR 2025
0
Of 0
(I:)
00,
00;
Using our convention from earlier, we can now compute the transformation from Cartesian basis vectors normalized hyperspherical basis vectors:
Ox
(I:))
00
00'
Or
R=
to compose the normalized "Jacobian" J:
Ogt
Ow'
0z
BF
00
06
0x2
OF
00;
86
0 862-1
(1)
8F
801
On
801
86
JESO(d)
Rescaling the hyperspherical basis vectors to have unit norm at all points causes the matrix J to become the orthogonal matrix with determinant equal to one J. This set of d matrices belongs to the group SO(d), which represents the of d-dimensional rotations about the origin. Similarly, the backwards change-of-bass j-1 jJ converts vectors in hyperspherical coordinates to Cartesian coordinates.
As result, vectors from the tangent space at p in Cartesian coordinates simply rotate convert to the normalized tangent space at P in hyperspherical coordinates. Specifically, for point p (r, 01, ....... 01) and vector ēir + 020² +... 2z0d-1, converting e = c, ... + Cax from Cartesian to hyperspherical coordinates is the transformation:
= j 0
where J operates on vector v-i.e. Ju-by first rotating by angle 22 in the x² plane (i.e. the 01 axis of rotation), then by angle č3 the x2 1 plane (i.e. the 6² axis of rotation), so on and so forth until a final rotation by angle č in the -d-1 - x plane (i.e. the gd axis of rotation). Composing these rotations together leads to a rotation from Ao 0,,0,0....
...,.. 0) pt:
Ju = (Rgo = (Ro4)1
paa
R ) R
j-!u= f = -p)= = ( 1 1_
R
x 2
Tu = Rø
where define Rg- be the rotation from à as described above and R Re' -s' by angle 0, in the x² x plane. Important for our later discussion on the rotation trick, this rotational characteristic causes moving fixed vector along curve in hyperspherical coordinates rotate in Cartesian coordinates.
be the rotation
Remark 4. Using the renormalized transformation in Equation (1), a constant vector field ū in hyperspherical coordinates corresponds to a rotated vector field in Cartesian coordinates.
Proof. Cartesian point p and corresponding hyperspherical point ýt:
T
R5
T
Rg d
R
R2
T
Up = 0p
[Rg- 1. = Up
so a constant vector field in hyperspherical coordinates will correspond to a cartesian vector field where each vector point p rotated by the rotation that alights pn-
31

---

Published as a conference paper at ICLR 2025
Another important characteristic relates the metric tensor with normalized hyperspherical basis vectors. We can explicitly compute the induced metric in hyperspherical coordiantes in terms of our renormalized basis vectors:
0
8F
07
0
8
0
00,
302
862
80,3
80,
(In))
-1
Dr
(I22)
86
09
(Iaa
882
092
0
(Ida)
30
80
(In) -1 (In1)
0
(I22) (Inz)
0
(Isa)
(Iaa)_tu
(Iaz)
0
(Iaa) (((ae)
0
0
0
which yields the identity matrix. This is perhaps unsurprising: we normalize basis vectors so that a
Another way to view the renormalized tangent plane transformation is as a change-of-base in the Cartesian coordiante system, with the basis vectors spanning each tangent space in Cartesian coordinates rotated to align with the directions of hyperspherical basis vectors. Rotating the tangent space each point in Cartesian coordinates does not change the Euclidean metric tensor-it IR I-so it remains the identity. These two formulations are equivalent: renormalizing the basis vectors in the hyperspherical tangent space have unit norm corresponds rotating the basis vectors in the Cartesian coordinate system to align with the hyperspherical basis vectors at all points.
A.12.2 STE AS PARALLEL TRANSPORT
From the description of the STE Bengio et al. (2013), a gradient vector V, C is transported from during the backwards pass in such way that its direction and magnitude is preserved. Critically, the curve along which V.C. f is transported is not specified; the effect is to simply "copy-and sooy--an--aster pastt" the vector from g to e.
To use the machinery of calculus, we assume that v,C is transported from y to along any smooth curve (() running from y to Along this curve, define the transport of V,L position r() simply as V C to emulate how the STE would move v,L. from to 7(1). Therefore, the direction and magnitude of N does not change along the curve (t). An example of this transport visualized in Figure 20, and in Remark 5, we show this formulation is equivalent to the parallel transport of v,C along any curve (t) from q to with the Levi-Civita connection.
Remark 5. The Straight Through Estimator (STE) is equivalent to the parallel transport of V f. along any curve connecting to e with the identity metric tensor Cartesian coordinates using the Levi-Civite connection.
Proof. A vector field v is parallel transported along a curve y(t) the covariant derivative of v in the direction of (t) is zero. Informally, the change in the vector field v must exactly match how the
32

---

Published as a conference paper at ICLR 2025
Transport of VaC from q to e
in Cartesian Coordinates in Hyperspherical Coordinates
71 (t)
7i (t)
72
(f)
(t
(4)
Cartesian Vector Field for Y3
Hyperspherical Vector Field for 7t
8F
Figure 20: (top) Visualization of vector transport Cartesian coordinates and renormalized hyperspherical coordinates along curves (f), (f and (f). Notice hyperspherical basis changes from point to point. (bottom) Depiction of the transported vector in terms of the basis vectors @ and a for Cartesian coordinates and and 26 vector field in the Cartesian tangent space.
for hyperspherical coordinates. Notice how the components of and a6 change for constant
basis vectors of the tangent plane change along (t)) to remain "parallel" along the curve:
((t) 0
Parallel Transport Condition
Using the identity metric tensor:
8.g
with the Levi-Civita connection will result in all zero Christoffel symbols:
1
Ogik Ox
Dga Jxd
Ogi
mk
0
2
33

---

Published as a conference paper at ICLR 2025
where gT is the m, k entry of inverse metric tensor. Computing the covariant derivative for general curve (t):
V v(t)
V frer Fieet...
1 Ve
a 1/ Or
" Ox (V1c1 U262 + - 2yet)
must be equal for parallel transport
Considering the 7th term in this summation:
0 = (V o() = t Dr
(wie1 + 0202
vala)
% Oz
[omex]
Ook
ek + uk Jek Oxd
t
Ond
Ook
Oxe Ck + u² b*Imee em
Ookk
= 4
Ox Ckk
For this equation hold for arbitrary 7l), dz O for S k, S d. Therefore, ve must be constant, and vector fields along curves must be constant to satisfy the parallel transport criteria.
Pulling this back to the STE, holding v,C constant along the curve (t) from q to e results in constant vector field along 7t). The covariant derivative of this vector field is zero, and therefore the STE parallel transports vgC from q to e.
A.1 THE ROTATION TRICK AS PARALLEL TRANSPORT
In this section, analyze the rotation trick through the lens of geometry. As in Appendix A.12.2, extend the rotation trick to any smooth curve ytt) connecting to e and define the transport of V, C at 7() as the rotation trick applied to move V V.C from (t). This definition allows us to use the structure of calculus, without imposing any prohibitive restrictions on the path taken from to e.
To build visual intuition, Figure 21 illustrates how the rotation trick transforms initial vector along three different curves 71. 72. 7/3 in both Cartesian coordinates and hyperspherical coordinates with normalized basis vectors. In Cartesian coordinates, the rotation trick changes the components of the basis vectors during transport to follow rotation; however in normalized hyperspherical coordinates, the components of this vector during transport are constant because the basis vectors themselves rotate.
Remark The rotation trick is equivalent to the parallel transport of V along any curve connect- ing g to with the induced metric hyperspherical coordinates with the normalized transformation described in Equation (1) using the Levi-Civita connection.
Proof. From Equation (2), the metric tensor in hyperspherical coordinates with normalized basis vectors equivallllll the cartesian coordinate system with each tangent space rotated to align with the hyperspherical frame at every point-is the identity. Therefore, using the Levi-Civin connection
34

---

Published as a conference paper at ICLR 2025
Transport of g from q to in Cartesian Coordinates
in Hyperspherical Coordinates
(1)
OF
(1)
7:
(t
(1
(t
Hyperspherical Vector Field for (t)
Cartesian Vector Field for 7
BF
R,4
Figure 21: (top) Visualization of vector transport in hyperspherical coordinates with normalized basis vectors and Cartesian coordinates along curves ?/ (t), (t)() and (t()). The vectors along cach curve hyperspherical coordinates rotate to stay constant with respect to the natural rotation the basis vectors. This same rotation Cartesian coordinates yields a non-constant vector as Cartesian basis vectors do not change from point to 2 point. (bottom) Depiction of the transported vector in terms of the basis vectors by and for hyperspherical coordinates and à a år and By for Cartesian coordinates. In the former case, the transported vector remains constant with respect to the normalized basis vectors, while in Cartesian coordinates, the components change along ? (f))
leads to zero Christoffel symbols, and the parallel transport of vector along any curve keeps the vector constant.
We define T,C the tangent space of the Cartesian coordinate system at point p and T.HH as the tangent space of the hyperspherical coordinate system with normalized basis vectors at point p. remains to show that for vector V V,G E T2C and corresponding V,C E T7H, the transformation of VgL E T.H to TCC will yield Ry V where Ry is the rotation trick's 's transformation, i.e. the rotation that rotates q to e.
For vector VgC in hyperspherical coordinates point q (1, 02,,.., 0.-1) and using the normalized change-of-basie in Equation (I), the corresponding vector V, L in Cartesian coordinates is:
vyc² = =V [
vyC [J VgL.
= [R V V
[os Roa_2* Ro, VgC
35

---

Published as a conference paper at ICLR 2025
and the corresponding vector VgL at point è is:
V.C² VoL" Jo
[Je] v
VeC
[Rio va) VgL
[ Ro + VgC
Ri-
Reo 4 V 7
[Rg-tē] VgL
which is exactly how the rotation trick transforms the vector. Informally, "copy-and-pastingg the vector VgLi from I to è in hyperspherical coordinates with normalized basis vectors corresponds to rotating V, C by the rotation that aligns g in Cartesian coordinates.
In summary, we consider a geometry where the tangent space is spanned by unit norm basis vectors OF ag a0a-1 that match the direction of the typical hyperspherical basis vectors Br bg a6 The induced metric tensor is the identity, so the parallel transport of vec- along any curve holds its components constant. Converting a vector V,C to this tangent space via the normalized transformation in Equation (1), parallel transporting the resulting vector from è, and then converting it back to Cartesian coordinates corresponds exactly the rotation trick's transformation.
0
This is remarkably simple result; the rotation trick and the STE can be viewed the same operation. Both parallel transport the gradient V vgL. from g in path-independent manner with the Euclidean metric. The only difference is the coordinate system where parallel transport oceurs. The STE employs the Cartesian coordinate system while the rotation trick uses the hyperspherical coordinate system with normalized basis vectors.
36


Reviews:
[{"review_id": "DpDr22igpj", "reviewer": "Reviewer_x3Va", "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.", "questions": "Please see the weaknesses above.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The method is principled, and I think it is a good and simple idea to improve some of the issues of STE\n- 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. \n- 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.\n- 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. \n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oNJSLXC3Ps", "reviewer": "Reviewer_AkYx", "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.", "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.\n1. 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)?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The rotation trick is novel, intuitive, and mathematically principled\n1. Applying two Householder reflectors is computationally efficient, and incurs no additional training cost\n1. 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).\n1. 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$.\n1. 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.\n1. 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.\n1. There is an extra \"a\" on line 113.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iOL4tKtrl3", "reviewer": "Reviewer_fmLN", "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.\n\nVQ-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.\n\nIn 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$.\n\nThe authors argue and demonstrate empirically across several experiments that this pick significantly improves codebook utilization, reconstruction loss, and reconstruction quality.", "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!\n - 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?\n\n## References\n[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.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\nThe core idea in the paper is neat and simple, and the experimental verification appears reasonably extensive and rigorous.\n\nThe 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:\n1. 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?\n2. 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?\n\nMinor:\n - 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eTPZ1xAHO6", "reviewer": "Reviewer_DhUX", "summary": "This paper proposes a method, reffered to as the \"rotation trick\", to propagate approximated gradients through the vector quantization layer of VQ-VAEs.\nThe proposed method addresses the issue of a \"straight-through estimator (STE),\" which loses information about the positions of encoder outputs relative to codebook vectors.\nThe 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\".\nThe authors conduct experiments across 11 VQ-VAE training settings to demonstrate that the proposed method improves reconstruction performance compared to baseline methods.", "questions": "* Why is the rescaling factor $||q|| / ||e||$ used for the gradient approximation? How does it affect training of the encoder and the codebook?\n* Minor comment: In Definition 1 (Householder Reflection Matrix), isn't it required that $||a|| = 1$?\n* Minor comment: there is a typo: \"a a convex combination'\" in p3 L6.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "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.\nThe algorithm for the rotation trick is clearly formulated.\nThe explanation of the effects of the rotation trick, supported by figures, is illustrative.\nThe 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.\n\nThe 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$.)\nAdditionally, it seems that rescaling of gradients may hinder the propagation of gradients when $||q||$ is close to zero.\nProviding 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)