Another VQ Trick: Adding a Linear Transformation to the Codebook

VQ的又一技巧：给编码表加一个线性变换 苏剑林 · 2024-11-06 · Machine Learning generative model encoding gradient discretization

In The Rotation Trick for VQ: A General Extension of the Straight-Through Gradient Estimator], we introduced the Rotation Trick for VQ (Vector Quantization), whose core idea is to design better gradients for VQ by generalizing its STE (Straight-Through Estimator), thereby alleviating problems such as codebook collapse and low codebook utilization.

Coincidentally, a paper posted on arXiv yesterday, Addressing Representation Collapse in Vector Quantized Models with One Linear Layer], proposes yet another trick for improving VQ: adding a linear transformation to the codebook. This trick merely changes how the codebook is parameterized, without altering the underlying theoretical framework of VQ—yet the empirical results are excellent, making it a textbook case of "simple but effective." More below.

Background

Since we've already covered VQ and VQ-VAE multiple times in posts such as A Concise Introduction to VQ-VAE: The Quantized Autoencoder] and The Embarrassingly Simple FSQ: "Rounding" Surpasses VQ-VAE], I won't rehash the basics here and will jump straight to the mathematical formulation of a plain AE versus VQ-VAE:

\begin{align} \text{AE:}&\qquad z = encoder(x),\quad \hat{x}=decoder(z),\quad \mathcal{L}=\Vert x - \hat{x}\Vert^2 \\[12pt] \text{VQ-VAE:}&\qquad\left\{\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}\right.\label{eq:vqvae} \end{align}

Let me reiterate the well-worn point: VQ-VAE is not a VAE—it's simply an AE with VQ added on top, and it lacks the generative capability of a VAE. VQ itself is the operation of mapping an arbitrary vector to its nearest neighbor in the codebook. This operation is inherently non-differentiable, so gradients for the encoder are constructed via STE, and two additional loss terms $\beta,\gamma$ are introduced to supply gradients to the codebook, while also serving to regularize the encoder's representations.

The Modification

The paper calls its proposed method SimVQ, without explaining what "Sim" stands for—my guess is that it's short for "Simple," because the modification really is that simple:

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

Indeed, all that's done is multiplying the codebook by an extra matrix $W$, leaving everything else untouched.

If VQ was originally trained using formula $\eqref{eq:vqvae}$, then SimVQ can be adopted directly and straightforwardly. If instead the codebook was originally updated via EMA (i.e., $\beta=0$, with a separate moving-average process used to update the codebook—the approach taken by VQ-VAE-2 and some subsequent models, which is mathematically equivalent to optimizing the codebook loss with SGD while other losses can use non-SGD optimizers like Adam), then this EMA mechanism needs to be dropped, and the $\beta$ term reintroduced for end-to-end optimization.

Some readers might immediately object: isn't this just reparameterizing the codebook from $E$ to $EW$? Since $EW$ can be merged into a single matrix, equivalent to a new $E$, shouldn't this leave the model's theoretical capacity unchanged? Indeed, SimVQ leaves the model's capacity unchanged—but it does change things for SGD and Adam. It alters the optimizer's learning dynamics, and thereby affects the quality of the learned solution.

Experiments

Before diving further into analysis, let's first look at SimVQ's experimental results. SimVQ conducted experiments on both vision and audio, and a particularly representative one is Table 1:

SimVQ's experimental results

According to the paper, SimVQ's code is modified directly from the first-row VQGAN codebase, and the only change is inserting a linear transformation into the VQ layer—yet the improvement is very substantial. Not only does it achieve the best reconstruction quality at the same codebook size, it can further improve reconstruction quality by enlarging the codebook, which perfectly illustrates SimVQ's appeal: simple yet effective.

I also tried this on my own previously written VQ-VAE code, and my tests confirm that adding this linear transformation noticeably speeds up VQ-VAE's convergence, and also reduces the final reconstruction loss somewhat. I also experimented with a variant where $W$ is restricted to a diagonal matrix, which amounts to element-wise multiplying each codebook vector by a parameter vector (initialized to all ones). The results show that this variant achieves a similar effect, falling somewhere between plain VQ and SimVQ.

Analysis

Intuitively, VQ's updates to the codebook are rather "isolated": if a sample $z$ is quantized to $q$, then the gradient from that sample only affects $q$, leaving the rest of the codebook untouched. SimVQ is different: it not only updates $q$, but also updates $W$. Geometrically, $W$ acts as a basis for the codebook, so updating $W$ updates the entire codebook at once. In other words, SimVQ makes the whole codebook "move together" much more tightly, giving it a better chance of finding a superior solution rather than getting stuck in a local optimum where each code fends for itself.

So why does SimVQ improve codebook utilization? This is also fairly easy to understand. Building again on the interpretation that $W$ is the codebook's basis: if codebook utilization is too low, then $W$ becomes "anisotropic"—that is, the basis tilts toward the codes that are actually being used. But once the basis shifts in this way, its linear combinations will likewise tend to favor the codes that are being used, so utilization won't stay too low. Put simply, a learnable basis will automatically push its own utilization upward, which in turn raises the utilization of the entire codebook.

We can also describe this process from the perspective of the mathematical formulas. Suppose the optimizer is SGD; then, under plain VQ, the update to code $e_i$ is

\begin{equation}e_i^{(t+1)} = e_i^{(t)} - \eta\frac{\partial \mathcal{L}}{\partial e_i^{(t)}}\end{equation}

In this case, if $e_i$ isn't selected in the current batch, then $\frac{\partial \mathcal{L}}{\partial e_i^{(t)}}$ is zero, and that entry of the codebook simply doesn't get updated. But if $e_i$ is parameterized as $q_i W$, then

\begin{equation}\begin{aligned} q_i^{(t+1)} =&\, q_i^{(t)} - \eta\frac{\partial \mathcal{L}}{\partial q_i^{(t)}} = q_i^{(t)} - \eta \frac{\partial \mathcal{L}}{\partial e_i^{(t)}} W^{(t)}{}^{\top}\\ W^{(t+1)} =&\, W^{(t)} - \eta\frac{\partial \mathcal{L}}{\partial W^{(t)}} = W^{(t)} - \eta \sum_i q_i^{(t)}{}^{\top}\frac{\partial \mathcal{L}}{\partial e_i^{(t)}} \\ e_i^{(t+1)}=&\,q_i^{(t+1)}W^{(t+1)}\approx e_i^{(t)} - \eta\left(\frac{\partial \mathcal{L}}{\partial e_i^{(t)}} W^{(t)}{}^{\top}W^{(t)} + q_i^{(t)}\sum_i q_i^{(t)}{}^{\top}\frac{\partial \mathcal{L}}{\partial e_i^{(t)}}\right) \end{aligned}\end{equation}

We can see that:

1. $W$ is updated based on the sum of gradients over all selected codes, so it naturally tends toward directions of high utilization;
2. Because of the presence of $q_i^{(t)}\sum_i q_i^{(t)}{}^{\top}\frac{\partial \mathcal{L}}{\partial e_i^{(t)}}$, the update for code $i$ is almost never exactly zero, regardless of whether that code was selected;
3. $q_i^{(t)}\sum_i q_i^{(t)}{}^{\top}\frac{\partial \mathcal{L}}{\partial e_i^{(t)}}$ acts as a projection onto the direction of high utilization, pushing every code toward that direction.

However, too much of a good thing can backfire: if every single code is pushed hard toward the high-utilization direction, this could actually cause codebook collapse. For this reason, SimVQ by default adopts a conservative strategy: only $W$ is updated, while all $q$ are frozen after random initialization, which essentially eliminates the risk of codebook collapse. The good news is that, given an appropriate coding dimension, experiments show that updating both $q,W$ and updating only $W$ perform about equally well, so readers are free to choose whichever form they prefer.

Further Thoughts

Setting aside the VQ context, this style of trick—introducing extra parameters that are mathematically equivalent (i.e., that don't change the model's theoretical fitting capacity) but that change the dynamics of the optimization process—is what we call "overparameterization."

Overparameterization is far from rare in neural networks. For instance, in today's mainstream Pre-Norm architecture, i.e., $x + f(\text{RMSNorm}(x))$, the vector $\gamma$ that RMSNorm multiplies by at the end is typically overparameterized, because the first layer of $f$ is usually itself a linear transformation—for example, attention linearly projects into Q, K, V, and the FFN uses a linear transformation to expand dimensionality, and so on. During inference, these models could in principle fold the $\gamma$ vector entirely into the linear transformation of $f$, yet it's rare to see anyone simply dropping $\gamma$ during training.

This is because many people believe that overparameterization plays a non-negligible role in making deep learning models "easy to train," so removing the overparameterization of a well-validated model on a whim carries real risk. Here, "easy to train" mainly refers to the somewhat astonishing fact that gradient descent—a method that in theory could easily get stuck in local optima—often manages to find solutions that perform remarkably well in practice. There is also work such as On the Optimization of Deep Networks: Implicit Acceleration by Overparameterization], which shows that overparameterization implicitly accelerates training, playing a role similar to momentum in SGD.

Finally, VQ can fundamentally be understood as a form of sparse training scheme, so the insight and modification introduced by SimVQ might also be applicable to other sparse training models, such as MoE (Mixture of Experts). In current MoE training schemes, updates across experts are also fairly independent—only the experts selected by the router get their parameters updated. So could it be possible, following SimVQ's lead, to attach a shared linear transformation after every expert in order to improve expert utilization efficiency? Of course, MoE differs from VQ in many respects, so this remains purely a conjecture for now.

Summary

This post introduced another training trick for VQ (Vector Quantization)—SimVQ—which simply adds a single linear transformation to the VQ codebook, with no other changes required, and yet achieves faster convergence, improved codebook utilization, and reduced reconstruction loss. It's about as simple and effective as it gets.