The Rotation Trick for VQ: A General Extension of the Straight-Through Gradient Estimator

VQ的旋转技巧：梯度直通估计的一般推广 苏剑林 · 2024-10-24 · Machine Learning generative model encoding gradient discretization

With multimodal LLMs on the rise, VQ (Vector Quantization) has correspondingly gained in prominence — it can serve as a tokenizer for vision or indeed any modality, unifying multimodal data within the autoregressive generation framework. Unfortunately, since VQ was first proposed in VQ-VAE], its theory has not progressed much, and problems such as codebook collapse or low codebook utilization remain to be solved. Meanwhile, alternatives such as FSQ] have been proposed and have become formidable "competitors" to VQ.

However, FSQ cannot replace VQ in every scenario, so improving VQ itself is still worthwhile. Recently I read Restructuring Vector Quantization with the Rotation Trick], which proposes a rotation trick claiming to alleviate a range of VQ's problems. Let's take a look at it in this post.

Recap

Back five years ago, in the post A Concise Introduction to VQ-VAE: The Quantized Autoencoder], we already introduced VQ-VAE, and later, when introducing FSQ in Embarrassingly Simple FSQ: "Rounding" Surpasses VQ-VAE], we revisited VQ-VAE carefully once more. Readers unfamiliar with it may want to read those two posts first. more

Although VQ-VAE bears the name "VAE," it is in fact just an AE — it does not have the generative capability of a VAE. What distinguishes it from an ordinary AE is that its encoding result is a discrete sequence rather than a continuous vector, i.e., it can encode either continuous or discrete data into a discrete sequence, and it allows the decoder to reconstruct the original input from this discrete sequence — much like a text tokenizer, which converts input into another discrete sequence and then lets you recover the original text from that sequence. This is why VQ-VAE is regarded as a tokenizer for arbitrary modalities.

In formulas, an ordinary AE is:

\begin{equation}z = encoder(x),\quad \hat{x}=decoder(z),\quad \mathcal{L}=\Vert x - \hat{x}\Vert^2 \end{equation}

while VQ-VAE is

\begin{equation}\begin{aligned} z =&\, encoder(x)\\[5pt] z_q =&\, z + \text{sg}[q - z],\quad q = \mathop{\text{argmin}}_{e\in\{e_1,e_2,\cdots,e_K\}} \Vert z - e\Vert\\ \hat{x} =&\, decoder(z_q)\\[5pt] \mathcal{L} =&\, \Vert x - \hat{x}\Vert^2 + \beta\Vert q - \text{sg}[z]\Vert^2 + \gamma\Vert z - \text{sg}[q]\Vert^2 \end{aligned}\end{equation}

Here "VQ" mainly refers to the process of transforming $z$ into $q$, which maps $z$ to one of the $e_1,e_2,\cdots,e_K$; these $e_i$ are called the codebook, and they are also learnable vectors. The "masterstroke" for training VQ-VAE is the step $z_q = z + \text{sg}[q - z]$, known as the "Straight-Through Estimator" (STE) for the gradient.

STE

The need for a straight-through estimator arises because the transformation from $z$ to $q$ involves the non-differentiable $\text{argmin}$ operation, so gradients cannot be directly propagated back into the encoder — in other words, the encoder cannot be trained. To get around this, VQ-VAE devised a trick: using the stop_gradient operator together with the nearest-neighbor relationship between $q$ and $z$, it replaces $q$ with $z$ during backpropagation, i.e., $z_q = z + \text{sg}[q - z]$.

In this construction, the forward computation is equivalent to treating $\text{sg}$ as nonexistent, so $z_q = z + q - z = q$, meaning what is fed into the decoder is $q$; while when computing gradients, the gradient of $\text{sg}$ is zero, so $\nabla z_q = \nabla z$, and gradients can bypass the non-differentiable operator and flow straight through to the encoder — hence "straight-through estimator." However, this means that while the encoder can now be optimized, the codebook cannot, so VQ-VAE adds $\beta\Vert q - \text{sg}[z]\Vert^2$ to the loss function to optimize the codebook, with an intent similar to K-Means: it wants $q$ to equal the center of all the $z$ closest to it. The final term, $\gamma\Vert z - \text{sg}[q]\Vert^2$, is meant to encourage the encoder to actively cooperate in promoting this clustering property.

From the perspective of the chain rule for gradients, we have

\begin{equation}\frac{\partial \mathcal{L}}{\partial z} = \frac{\partial q}{\partial z}\frac{\partial \mathcal{L}}{\partial q}\end{equation}

Note that here $z,q$ are all vectors, so $\frac{\partial \mathcal{L}}{\partial z},\frac{\partial \mathcal{L}}{\partial q}$ are also vectors, while $\frac{\partial q}{\partial z}$ is a matrix. Because of the non-differentiability of the map from $z$ to $q$, the problem is that $\frac{\partial q}{\partial z}$ is not well-defined, and STE is essentially assuming that $\frac{\partial q}{\partial z}=I$ (the identity matrix), giving $\frac{\partial \mathcal{L}}{\partial z} = \frac{\partial \mathcal{L}}{\partial q}$. This choice is naturally reasonable to some extent, but is there room for improvement?

Intuitively, the consequence of STE is that for all $z$ belonging to the same $q$, their gradients are all the same $\frac{\partial \mathcal{L}}{\partial q}$, regardless of how far they are from $z,q$ — and this seems like something we could improve on: can we define a more general $\frac{\partial q}{\partial z}$ that depends on how different it is from $z,q$? To this end, let's first generalize STE to

\begin{equation}z_q = \text{sg}[G]z + \text{sg}[q - Gz]\end{equation}

where $G$ is a matrix. Again, using the principle that forward propagation treats $\text{sg}$ as nonexistent and backpropagation treats the gradient of $\text{sg}$ as zero, we obtain $z_q = q$, $\frac{\partial \mathcal{L}}{\partial z_q} = G\frac{\partial \mathcal{L}}{\partial z}$, which is equivalent to defining $\frac{\partial q}{\partial z}=G$.

Rotation

So how should we choose $G$? The paper mentioned at the beginning of this post proposes a reference scheme, constructing $G$ based on the rotation transformation from $z$ to $q$ — hence the "rotation trick" in the paper's title.

Specifically, the original paper considers the simple case where $Gz = q$, in which $\text{sg}[q - Gz]$ is automatically zero, simplifying things to $z_q = \text{sg}[G]z$. To find the matrix $G$, we first normalize $z,q$ into unit vectors $\tilde{z} = \frac{z}{\Vert z\Vert},\tilde{q} = \frac{q}{\Vert q\Vert}$, and then we can construct a rotation transformation from $\tilde{z}$ to $\tilde{q}$. We've already discussed the concrete construction in An Orthogonal Matrix That Transforms One Unit Vector into Another]; the answer is

\begin{equation}R = I + 2\tilde{q}\tilde{z}^{\top}- \frac{(\tilde{q} + \tilde{z})(\tilde{q} + \tilde{z})^{\top}}{1 + \cos\theta} = I + 2\tilde{q}\tilde{z}^{\top}- 2\left(\frac{\tilde{q} + \tilde{z}}{\Vert\tilde{q} + \tilde{z}\Vert}\right)\left(\frac{\tilde{q} + \tilde{z}}{\Vert\tilde{q} + \tilde{z}\Vert}\right)^{\top} \end{equation}

where $\theta$ is the angle between $q,z$. Using this result, we can write

\begin{equation}\tilde{q}=R\tilde{z}\quad\Rightarrow\quad q = \frac{\Vert q\Vert}{\Vert z\Vert} R z\quad\Rightarrow\quad G = \frac{\Vert q\Vert}{\Vert z\Vert} R\end{equation}

To improve the efficiency of computing $Gz$, we typically use the associativity of matrix multiplication to first compute $\tilde{z}^{\top}z$ and $\left(\frac{\tilde{q} + \tilde{z}}{\Vert\tilde{q} + \tilde{z}\Vert}\right)^{\top}z$, but note that what we actually need is $\text{sg}[G]z$, so we must remember to stop the gradient on $\tilde{q},\tilde{z},\frac{\Vert q\Vert}{\Vert z\Vert}$ before computing $Gz$.

Geometrically, $\frac{\partial q}{\partial z}=G=\frac{\Vert q\Vert}{\Vert z\Vert} R$ makes the geometric relationship between $\frac{\partial \mathcal{L}}{\partial q}$ and $\frac{\partial \mathcal{L}}{\partial z}$ exactly identical to that between $q$ and $z$ — for instance, the angle between $\frac{\partial \mathcal{L}}{\partial q}$ and $\frac{\partial \mathcal{L}}{\partial z}$ equals the angle between $q$ and $z$, and their norm ratios are also equal. These properties are naturally elegant from a theoretical standpoint, but does this actually improve VQ-VAE's performance? Let's move on to the experiments.

Experiments

Under the same configuration, the paper compares the original STE with the rotation trick, and finds the latter's performance to be quite "striking":

Performance of VQ-VAE + rotation trickPerformance of VQ-GAN + rotation trick

In short: metrics that should be high (codebook utilization, IS) go up, and metrics that should be low (reconstruction error, loss, FID) go down — exactly the profile of an ideal model. The paper's code has also been open-sourced, and interested readers can try running it themselves.

GitHub: https://github.com/cfifty/rotation_trick

Some Thoughts

Does this mean every VQ-VAE/VQ-GAN can blindly adopt the rotation trick? I added the rotation trick to a working VQ-VAE codebase of my own, and found that performance actually got worse — specifically, the reconstruction loss $\Vert x - \hat{x}\Vert^2$ went up, while the codebook loss $\Vert q - z\Vert^2$ went down.

After a bit of analysis, I found the issue lies in the choice of $\frac{\partial q}{\partial z}=G=\frac{\Vert q\Vert}{\Vert z\Vert} R$. The original STE is equivalent to $\frac{\partial q}{\partial z}=I$, where the scale of the identity matrix $I$ matches that of the identity itself, whereas with the rotation trick, the rotation matrix $R$ introduces an extra scaling factor $\frac{\Vert q\Vert}{\Vert z\Vert}$. If at initialization $\Vert q\Vert \ll \Vert z\Vert$ (which happens to be the case for my VQ-VAE), then under the rotation trick the gradient of the reconstruction loss becomes much smaller than under STE, so for the encoder, the gradient from the term $\gamma\Vert z - \text{sg}[q]\Vert^2$ ends up dominating.

In other words, in the early stage, this amounts to optimizing only $\beta\Vert q - \text{sg}[z]\Vert^2 + \gamma\Vert z - \text{sg}[q]\Vert^2$, which leads to $q,z\to 0$, i.e., codebook collapse — which explains why the codebook loss decreases while the reconstruction loss increases. So switching from STE to the rotation trick most likely requires at least retuning $\gamma$. Looking briefly at the paper's open-source code, it seems they initialize the codebook using K-Means on the initial encoder outputs, so that $\Vert q\Vert$ and $\Vert z\Vert$ don't differ too much in magnitude, allowing for a relatively smooth transition.

That said, even after tuning $\gamma$, I was unable to obtain better results on my own VQ-VAE code, so I remain skeptical about the effectiveness of the rotation trick. Setting aside practice, I also find it hard to grasp the theoretical justification for why the rotation trick should work. The paper's explanation is that when $q$ and $z$ are close, $G$ is close to $I$, in which case $\frac{\partial \mathcal{L}}{\partial z} \approx \frac{\partial \mathcal{L}}{\partial q}$ is reasonable; but when $q$ and $z$ are far apart — for instance, when $z$ sits near the boundary of category $q$ — the gap between $G$ and $I$ is large, i.e., $\frac{\partial \mathcal{L}}{\partial z}$ deviates noticeably from $\frac{\partial \mathcal{L}}{\partial q}$, so $z$ ends up "flying around erratically," which supposedly helps $z$ break out of its "cage" and move toward a new category, thereby improving codebook utilization. This explanation, though, admittedly feels rather shaky.

There is also another issue with the rotation trick: it establishes a privileged central point — the origin. It's easy to see that VQ itself is analogous to K-Means clustering, and K-Means is centerless, having translation invariance, whereas rotation inherently requires a center (the origin). So the rotation trick actually conflicts somewhat with the original spirit of VQ. Of course, one could modify VQ to use cosine similarity for finding nearest neighbors, which would fit better with the rotation trick — but that still doesn't explain why the rotation trick also helps Euclidean-distance-based VQ. Overall, the fundamental reason the rotation trick works remains a question worth pondering.

Finally, some readers may wonder: given all these problems with VQ, why bother studying it at all? Why not just use something simpler like FSQ? In my view, alternatives such as FSQ cannot replace VQ in every scenario. For example, in the Transformer-VQ introduced in Applying VQ to Keys Makes Transformer Complexity Linear], it would be very difficult to substitute FSQ for VQ, because VQ is applied at every layer there — meaning the effective model allotted per VQ operation is quite small, and in practice FSQ only outperforms VQ once the model is sufficiently large.

Summary

The rotation trick is a new technique recently proposed on arXiv for training VQ (Vector Quantization) models. It generalizes the original Straight-Through Estimator (STE) and claims to alleviate problems such as codebook collapse and low utilization. This post has given a brief introduction to it, along with some of my own thoughts and questions about it.