DiVeQ: A Remarkably Concise Training Scheme for VQ

DiVeQ：一种非常简洁的VQ训练方案 苏剑林 · 2025-10-08 · Machine Learning generative model encoding gradient discretization

For researchers who remain committed to the discretization route, VQ (Vector Quantization) is a crucial part of visual understanding and generation, playing the role of a "tokenizer" for vision. It was first proposed in the 2017 paper Neural Discrete Representation Learning, and I also introduced it back in 2019 in my blog post A Brief Introduction to VQ-VAE: Quantized Autoencoders.

However, all these years later, we find that the training techniques for VQ have barely changed: it's always STE (Straight-Through Estimator) plus an extra auxiliary loss. STE itself is not really a problem — it's arguably the standard way of designing gradients for discretization operations — but the existence of the auxiliary loss always leaves one with a feeling that things aren't quite end-to-end, and it also introduces an extra hyperparameter to tune.

Fortunately, this situation may be coming to an end. Last week's paper DiVeQ: Differentiable Vector Quantization Using the Reparameterization Trick proposes a new STE trick, and its biggest highlight is that it doesn't need an auxiliary loss at all, which makes it look particularly clean and elegant! more

Discrete Encoding

As usual, let's first review the existing training scheme for VQ. First, it should be pointed out that VQ (Vector Quantization) itself is actually a very old concept, dating back to the 1980s, originally meaning to cluster vectors and replace them with their corresponding cluster centers, thereby achieving data compression.

But the VQ we're talking about here mainly refers to the VQ used in VQ-VAE, as proposed in the paper Neural Discrete Representation Learning. Of course, the definition of VQ itself hasn't changed — it's always a mapping from vectors to cluster centers. The core contribution of the VQ-VAE paper is that it provides an end-to-end training scheme in which a latent variable is VQ'd and then decoded for reconstruction. The difficulty is that the VQ step is a discretization operation with no readily available gradient, so a gradient needs to be designed for it.

In formulas, an ordinary AE (AutoEncoder) is:

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

where $x$ is the original input, $z$ is the encoded vector, and $\hat{x}$ is the reconstruction result. What VQ-VAE wants to do, based on the idea of VQ, is turn $z$ into one of the entries of a codebook $E=\{e_1,e_2,\cdots,e_K\}$:

\begin{equation}q = \newcommand{argmin}{\mathop{\text{argmin}}}\argmin_{e\in E} \Vert z - e\Vert\end{equation}

Then $decoder$ takes $q$ as input to perform the reconstruction. Since $q$ is a one-to-one correspondence with the codebook indices, $q$ is effectively an integer encoding of $x$. Of course, to preserve reconstruction quality, in practice one certainly doesn't encode things into a single vector, but rather into multiple vectors, so after VQ we end up with multiple integers. Thus, what VQ-VAE is trying to do is encode the input into a sequence of integers, which is essentially the same idea as a text tokenizer.

Gradient Design

Now the modules we need to train are $encoder$, $decoder$, and the codebook $E$. Since the VQ operation involves an $\argmin$ operation, the gradient is cut off at $q$ and cannot be propagated back into $encoder$.

VQ-VAE uses a trick called STE, which says that what is fed into $decoder$ is still the post-VQ $q$, but when computing the gradient during backpropagation, we use the pre-VQ $z$, so that the gradient can be passed back to $encoder$. This can be implemented with the stop_gradient operator ($\newcommand{sg}{\mathop{\text{sg}}}\sg$):

\begin{equation}z = encoder(x),\quad q = \argmin_{e\in E} \Vert z - e\Vert,\quad z_q = z + \sg[q - z],\quad \hat{x} = decoder(z_q)\end{equation}

In simple terms, the effect STE achieves is $z_q=q$ but $\nabla z_q = \nabla z$, so that $encoder$ now has a gradient, but $q$ has none, meaning the codebook can't be optimized this way. To solve this problem, VQ-VAE adds two auxiliary loss terms:

\begin{equation}\mathcal{L} = \Vert x - \hat{x}\Vert^2 + \beta\Vert q - \sg[z]\Vert^2 + \gamma\Vert z - \sg[q]\Vert^2 \end{equation}

These two loss terms represent, respectively, pulling $q$ toward $z$ and pulling $z$ toward $q$, which is consistent with the original idea of VQ. STE combined with these two auxiliary loss terms constitutes the standard VQ-VAE. There's also a simple variant, which sets $\beta=0$ directly, but updates the codebook using an exponential moving average of $z$ — this is equivalent to specifying that the auxiliary loss for $q$ is updated with SGD.

Incidentally, it's worth noting that although VQ-VAE is labeled "VAE" in the original paper, it is actually just an AE, so calling it "VQ-AE" would in principle be more accurate — but that name has stuck, so we're stuck with it too. The later VQGAN builds on VQ-VAE by adding techniques like a GAN loss to improve the sharpness of reconstructions.

Alternative Works

For me, these two extra auxiliary loss terms have always been the uncomfortable part. I imagine quite a few people in the field feel the same way, so there have been various improvement efforts along these lines from time to time.

Among them, the most "root-and-branch" approach is to abandon VQ altogether in favor of some other discretization scheme, such as the FSQ introduced in Embarrassingly Simple FSQ: "Rounding" Surpasses VQ-VAE, which requires no auxiliary loss at all. If VQ is about clustering high-dimensional vectors, then FSQ is about "rounding" low-dimensional vectors to achieve discretization. However, as I discussed in this earlier article, FSQ cannot replace VQ in every scenario, so improving VQ itself is still valuable.

Before proposing DiVeQ, the original authors actually had already proposed a scheme called "NSVQ", which took a small step toward "abolishing" the auxiliary loss. It changes $z_q$ to:

\begin{equation}z_q = z + \Vert q - z\Vert \times \frac{\varepsilon}{\Vert \varepsilon\Vert},\qquad \varepsilon\sim\mathcal{N}(0, I)\label{eq:nsvq}\end{equation}

Here $\varepsilon$ is a vector of the same size as $z,q$, with components drawn from a standard normal distribution. After substituting in this new $z_q$, thanks to the differentiability of $\Vert q - z\Vert$, $q$ now also has a gradient, so in principle the codebook can be trained without any auxiliary loss. The geometric meaning of NSVQ is quite intuitive: it amounts to sampling uniformly on "a sphere centered at $z$ with radius $\Vert q-z\Vert$". The downside is that what gets fed to $decoder$ at this point is not $q$, whereas at inference time what we actually care about is the reconstruction quality of $q$ — so there's an inconsistency between NSVQ's training and inference behavior.

Enter the Protagonist

Starting from NSVQ, if we want to keep the forward pass as $q$ while still retaining the gradient contribution from $\Vert q - z\Vert$, it's natural to propose an improved version:

\begin{equation}z_q = z + \Vert q - z\Vert \times \sg\left[\frac{q - z}{\Vert q - z\Vert}\right]\label{eq:diveq0}\end{equation}

In the forward pass this strictly satisfies $z_q = q$, but in the backward pass it retains gradients from both $z$ and $\Vert q - z\Vert$. This is the "DiVeQ-detach" described in the appendix of the DiVeQ paper. The main-text version of DiVeQ, on the other hand, is essentially some kind of interpolation between equations $\eqref{eq:diveq0}$ and $\eqref{eq:nsvq}$:

\begin{equation}z_q = z + \Vert q - z\Vert \times \sg\left[\frac{q - z + \varepsilon}{\Vert q - z + \varepsilon\Vert}\right],\qquad \varepsilon\sim\mathcal{N}(0, \sigma^2 I)\label{eq:diveq}\end{equation}

Clearly, when $\sigma=0$, the result is "DiVeQ-detach", and when $\sigma\to\infty$, the result is "NSVQ". The paper's appendix runs a search over $\sigma$, and roughly concludes that $\sigma^2 = 10^{-3}$ is a generally good choice.

The paper's experimental results show that although equation $\eqref{eq:diveq}$ introduces randomness, and also comes with a certain degree of training-inference inconsistency, it performs better than equation $\eqref{eq:diveq0}$. However, as far as my own aesthetic preference goes, performance shouldn't come at the cost of elegance, so equation $\eqref{eq:diveq0}$'s "DiVeQ-detach" is, in my mind, the ideal scheme. In the analysis below, "DiVeQ" refers specifically to "DiVeQ-detach".

Theoretical Analysis

Unfortunately, the original paper doesn't offer much theoretical analysis, so in this section I'll try to give a basic account of why DiVeQ is effective and how it relates to VQ's original training scheme. First, consider the general form of equation $\eqref{eq:diveq0}$:

\begin{equation}z_q = z + r(q, z) \times \sg\left[\frac{q - z}{r(q, z)}\right]\end{equation}

where $r(q,z)$ is an arbitrary differentiable scalar function of $q,z$, which can be regarded as some arbitrary distance function on $q,z$. Denoting the loss function as $\mathcal{L}(z_q)$, its differential is:

\begin{equation}d\mathcal{L} = \langle\nabla_{z_q} \mathcal{L},d z_q\rangle = \left\langle\nabla_{z_q} \mathcal{L},dz + dr \times\frac{q-z}{r}\right\rangle = \langle\nabla_{z_q} \mathcal{L},d z\rangle + \langle\nabla_{z_q} \mathcal{L}, q-z\rangle d(\ln r)\end{equation}

Here $\langle\nabla_{z_q} \mathcal{L},d z\rangle$ is what ordinary VQ already has, while DiVeQ contributes the extra term $\langle\nabla_{z_q} \mathcal{L}, q-z\rangle d(\ln r)$ — or put another way, this is equivalent to introducing an auxiliary loss $\sg[\langle\nabla_{z_q} \mathcal{L}, q-z\rangle] \ln r$. If $r$ represents some kind of distance function of $q,z$, then this term is pulling $q,z$ closer together, in the same spirit as the auxiliary loss introduced by VQ. This gives a successful theoretical explanation of DiVeQ.

But let's not celebrate too soon — this explanation only holds under the premise that the coefficient $\langle\nabla_{z_q} \mathcal{L}, q-z\rangle > 0$, otherwise it would actually be pushing the distance apart. To demonstrate this, let's consider the first-order approximation of the loss function $\mathcal{L}(z)$ at $z_q$:

\begin{equation}\mathcal{L}(z) \approx \mathcal{L}(z_q) + \langle\nabla_{z_q} \mathcal{L}, z - z_q\rangle = \mathcal{L}(z_q) + \langle\nabla_{z_q} \mathcal{L}, z - q\rangle\end{equation}

That is, $\langle\nabla_{z_q} \mathcal{L}, q-z\rangle\approx \mathcal{L}(z_q) - \mathcal{L}(z)$. Note that $z,z_q$ are the features before and after VQ respectively. VQ is a process that loses information, so using $z$ for the downstream task (e.g., reconstruction) should be easier than using $z_q$. Therefore, as long as training has begun to converge, we can expect the loss for $z$ to be lower, i.e., $\mathcal{L}(z_q) - \mathcal{L}(z) > 0$, which shows that $\langle\nabla_{z_q} \mathcal{L}, q-z\rangle > 0$ holds with high probability.

Directions for Improvement

Strictly speaking, $\langle\nabla_{z_q} \mathcal{L}, q-z\rangle > 0$ can only be regarded as a necessary condition for DiVeQ's effectiveness; to fully justify its effectiveness we'd also need to show that this coefficient is "just right." Because $r(q,z)$ is arbitrary, we can only analyze specific choices of function on a case-by-case basis. Let's consider $r(q,z)=\Vert q-z\Vert^{\alpha}$, in which case it is equivalent to introducing the following auxiliary loss:

\begin{equation}\sg[\langle\nabla_{z_q} \mathcal{L}, q-z\rangle] \ln \Vert q-z\Vert^{\alpha}\approx \sg[\mathcal{L}(z_q) - \mathcal{L}(z)]\times \alpha\ln \Vert q-z\Vert\end{equation}

The coefficient $\mathcal{L}(z_q) - \mathcal{L}(z)$ is homogeneous with the main loss $\mathcal{L}(z_q)$, which means it can adapt well to the scale of the main loss, and can also adjust the weight of the auxiliary loss based on the performance gap before and after VQ. As for what value of $\alpha$ works best, I think that depends on experiments — from my own tuning attempts, $\alpha=1$ does indeed seem to perform well quite generally. Interested readers are welcome to try adjusting $\alpha$ themselves, or even try swapping in a different $r(q, z)$.

It should be pointed out that DiVeQ only offers a new auxiliary-loss-free training scheme for VQ; in principle it doesn't solve VQ's other problems, such as low codebook utilization or codebook collapse. Any enhancement techniques that were effective under the "STE + auxiliary loss" setting can in principle be layered on top of DiVeQ as well. The original paper itself combines DiVeQ with SFVQ, proposing SF-DiVeQ to help alleviate issues like codebook collapse.

That said, I personally find SFVQ a bit convoluted, so I won't go into it in detail here — and I suspect the authors' choice to combine it with SFVQ has more to do with the fact that SFVQ was their own earlier work, so it's a natural lineage for them. What I personally prefer is the linear-transformation trick introduced in Another Trick for VQ: Adding a Linear Transformation to the Codebook, i.e., adding an extra linear transformation after the codebook. In my own tests, this also noticeably improves the performance of DiVeQ.

Summary

This post introduced a new training scheme for VQ (Vector Quantization) that can be implemented purely through STE, without requiring any additional auxiliary loss, which makes it particularly clean and elegant.