Pushing Toward the Simplex Vertices: A Simple Remedy for Code Collapse in Smoothed Vector Quantization

## Abstract

Abstract Vector quantization, which discretizes a continuous vector space into a finite set of representative vectors (a codebook ), has been widely adopted in modern machine learning.
Despite its effectiveness, vector quantization poses a fundamental challenge:
the non-differentiable quantization step blocks gradient backpropagation. Smoothed vector quantization addresses this issue by relaxing the hard assignment of a codebook vector into a weighted combination of codebook entries, represented as the matrix product of a simplex vector and the codebook.
Effective smoothing requires two properties:
(1) smoothed quantizers should remain close to a onehot vector, ensuring tight approximation, and
(2) all codebook entries should be utilized, preventing code collapse .
Existing methods typically address these desiderata separately.
By contrast, the present study introduces a simple and intuitive regularization that promotes both simultaneously by minimizing the distance between each simplex vertex and its K K -nearest smoothed quantizers.
Experiments on representative benchmarks—including discrete image autoencoding and contrastive speech representation learning—demonstrate that the proposed method achieves more reliable codebook utilization and improves performance compared to prior approaches.

## 1 Introduction

Vector quantization is a method for discretizing a continuous vector space (Gray,, 1984; van den Oord et al.,, 2017).
It maps each vector in the continuous space to the nearest element of a finite set of representative vectors (a.k.a. a *codebook*).
The resulting discrete representations are easier to manipulate and interpret than the original continuous forms,
and have proven effective across diverse applications, including image generation (Esser et al.,, 2021; Ramesh et al.,, 2021; Rombach et al.,, 2022; Yu et al.,, 2022), speech recognition (Baevski et al., 2020a,; Baevski et al., 2020b,), and music generation (Hadjeres and Crestel,, 2020; Dhariwal et al.,, 2020).

When integrated into deep neural networks, however, vector quantization introduces a fundamental challenge:
quantization is a non-differentiable operation that blocks gradient backpropagation (van den Oord et al.,, 2017).
Accordingly, some approximation is required to enable learning through quantization.

One effective workaround is to *smooth* vector quantization (Jang et al.,, 2017).
The selection of a codebook vector can be expressed as multiplying the codebook matrix by its corresponding onehot vector $(0,\dots,1,\dots,0)^{\mathsf{T}}$, whose nonzero entry indexes the chosen vector.
*Smoothed* vector quantization relaxes this onehot vector to lie within the simplex $\Delta^{M-1}:=\{(p_{1},\dots,p_{M})\mid\sum_{m=1}^{M}p_{m}=1\}$, where $M$ denotes the number of codebook vectors.
Consequently, differentiable mappings (e.g., softmax) become available and can be incorporated into neural networks.

Figure: Figure 1: ([1](https://arxiv.org/html/2509.22161v1#S1.F1)) Four different distributions on the simplex $\Delta^{3-1}$. For effective smoothed vector quantization, samples should be concentrated near the vertices of the simplex (i.e., onehot-like vectors; orange), rather than centered (dark gray) or uniformly spread across the simplex (light gray). At the same time, each vertex must be neighbored by some samples to avoid code collapse (blue). ([1](https://arxiv.org/html/2509.22161v1#S1.F1)) Maximizing the perplexity of the sample mean (Baevski et al., 2020b,) penalizes code collapse but cannot discriminate among the other three distributions. ([1](https://arxiv.org/html/2509.22161v1#S1.F1)) The proposed $K$-nearest neighbor (KNN) distance minimization ($K=8$) favors the desired vertex-concentrated distribution while also preventing code collapse.
Refer to caption: x1.png

To successfully approximate onehot quantizers, smoothed quantizers must be distributed around the vertices of the simplex (orange “o” in Figure [1](https://arxiv.org/html/2509.22161v1#S1.F1)).
At the same time, they should not concentrate around a few vertices, leaving other codebook entries unused (blue “x”).
This latter issue—called *code collapse*—has been identified as a major challenge in vector quantization (Dieleman et al.,, 2018; Baevski et al., 2020b,; Dhariwal et al.,, 2020; Fifty et al.,, 2025).
A widely adopted workaround for code collapse (in smoothed quantization) is to introduce an auxiliary learning objective that maximizes the entropy or perplexity of the *mean* of the smoothed quantizers (Dieleman et al.,, 2018; Baevski et al., 2020b,, see §[2.3](https://arxiv.org/html/2509.22161v1#S2.SS3) for the formal definition).
However, maximizing the entropy/perplexity of the mean can be achieved by various distributions, not only the desired vertex-neighboring ones.
For example, both uniform and center-concentrated distributions have their mean at the simplex center (Figure [1](https://arxiv.org/html/2509.22161v1#S1.F1)), which also maximizes this objective (Figure [1](https://arxiv.org/html/2509.22161v1#S1.F1)).
Accordingly, an additional mechanism is needed to tighten the smoothing (e.g., by adjusting the temperature parameter of the (Gumbel-)softmax; see §[2.2](https://arxiv.org/html/2509.22161v1#S2.SS2)).

Beyond this standard approach, however, there exists a simple and unified strategy for simultaneously tightening smoothed quantization and maximizing codebook usage:
*Why don’t we directly encourage clustering around all simplex vertices?*
Specifically, minimizing the distance between each simplex vertex and its $K$-nearest neighbors (KNNs) satisfies both desiderata at once (Figure [1](https://arxiv.org/html/2509.22161v1#S1.F1)).
The present study investigates this intuitive yet underexplored approach, comparing it against existing alternatives on representative benchmarks.
The results indicate that the proposed method enables the exhaustive usage of the entire codebook, even when other approaches suffer from code collapse.

The contributions of this work are summarized as follows:

- •
Neural vector quantization is reformulated as a smoothing problem of onehot vectors.
This simple reformulation has been absent in the literature, which traditionally framed vector quantization as an extension of variational autoencoding (Kingma and Welling,, 2014; Jang et al.,, 2017; van den Oord et al.,, 2017).
- •
Under this reformulation, an effective regularization loss function is proposed to promote both tight smoothing and exhaustive code utilization. The method demonstrates robustness across different learning settings.

The remainder of this paper is organized as follows.
§[2](https://arxiv.org/html/2509.22161v1#S2) reviews related studies on vector quantization.
§[3](https://arxiv.org/html/2509.22161v1#S3) introduces the proposed method,
which is then evaluated on representative benchmarks in §[4](https://arxiv.org/html/2509.22161v1#S4).
Finally, §[5](https://arxiv.org/html/2509.22161v1#S5) discusses the results and the limitations of the proposed method.

## 2 Related Studies

The central challenge of vector quantization lies in its non-differentiability, which disrupts the backpropagation of gradients in neural networks.
To address this, either the gradient computation (backward path) or the quantization itself (forward path) must be approximated.

### 2.1 Approximation in the Backward Path (Gradient Estimation)

One line of work retains the original non-differentiable quantization in the forward path but replaces the gradient computation in the backward path.
Let $\mathbf{z}\in\mathbf{R}^{D}$ denote a pre-quantized feature vector, and $\{\mathbf{q}_{1},\dots,\mathbf{q}_{M}\}\subset\mathbf{R}^{D}$ the set of quantized vectors.
Quantization maps $\mathbf{z}$ to its “closest” codebook entry according under a distance metric $\mathcal{D}$:
i.e., $\mathbf{z}\mapsto\mathbf{q}_{\iota(\mathbf{z})}$ where $\iota(\mathbf{z}):=\operatorname{argmin}_{m}\mathcal{D}(\mathbf{z},\mathbf{q}_{m})$.
Consequently, the partial derivatives $\frac{\partial q_{\iota(\cdot),i}}{\partial z_{j}}$ are ill-defined and must be approximated.

A canonical approximation—known as the *straight-through estimation* (STE)—replaces the ill-defined Jacobian with the identity matrix (van den Oord et al.,, 2017):

$$ $\displaystyle\frac{\partial q_{\iota(\cdot),i}}{\partial z_{j}}\approx\begin{cases}1&i=j\\ 0&\text{otherwise}\end{cases}$ (1) $$

Algorithmically, the STE is implemented using the $\mathtt{detach}$ operation, which excludes its argument from gradient computation:

$$ $\displaystyle\operatorname{STE}(\mathbf{q}_{\iota(\mathbf{z})},\mathbf{z})$ $\displaystyle:=\mathbf{q}_{\iota(\mathbf{z})}+\mathbf{z}-\mathtt{detach}(\mathbf{z})$ (2) $$

Here, $z-\mathtt{detach}(z)$ evaluates to zero, while the gradient with respect to $z$ can still be propagated through the first term.

More recently, Fifty et al., (2025) proposed an alternative gradient approximation, and demonstrated its empirical superiority over the STE:

$$ $\displaystyle\textsc{RE}(\mathbf{q}_{\iota(\mathbf{z})},\mathbf{z})$ $\displaystyle:=\mathtt{detach}\left(\frac{\|\mathbf{q}_{\iota(\mathbf{z})}\|}{\|\mathbf{z}\|}R\right)\mathbf{z}$ (3) $$

where $R$ is the rotation matrix aligning $\mathbf{z}$ to $\mathbf{q}_{\iota(\mathbf{z})}$,(^1^11The rotation matrix is given by $R=I-2\hat{\mathbf{r}}\hat{\mathbf{r}}^{\mathsf{T}}+2\hat{\mathbf{q}}_{\iota(\mathbf{z})}\hat{\mathbf{z}}^{\mathsf{T}}$, where $\hat{\mathbf{v}}:=\mathbf{v}/\|\mathbf{v}\|$ is the L2-normalization of vector $\mathbf{v}$, and $\mathbf{r}:=\hat{\mathbf{q}}_{\iota(\mathbf{z})}+\hat{\mathbf{z}}$.)
and $\frac{\|\mathbf{q}_{\iota(\mathbf{z})}\|}{\|\mathbf{z}\|}$ rescales the rotated vector to match the amplitude of $\mathbf{q}_{\iota(\mathbf{z})}$.
In this formulation, the Jacobian of the quantization is approximated by the scaled rotation matrix:

$$ $\displaystyle\frac{\partial q_{\iota(\cdot),i}}{\partial z_{j}}\approx\frac{\|\mathbf{q}_{\iota(\mathbf{z})}\|}{\|\mathbf{z}\|}R$ (4) $$

### 2.2 Approximation in the Forward Path (Smoothing)

An alternative approach approximates the forward quantization itself.
Using the onehot representation $\mathbf{e}_{m}$ of the code index $m$, quantization can be expressed as:

$$ $\displaystyle\mathbf{q}_{\iota(\mathbf{z})}$ $\displaystyle=Q\mathbf{e}_{\iota(\mathbf{z})}$ (5) $$

where $Q:=(\mathbf{q}_{1},\dots,\mathbf{q}_{M})$.
*Smoothed* quantization extends the possible range of $\mathbf{e}_{m}$ to the simplex $\Delta^{M-1}$.
As noted in §[1](https://arxiv.org/html/2509.22161v1#S1), effective learning requires smoothed quantizers $\mathbf{p}\in\Delta^{M-1}$ to concentrate near the vertices of the simplex (i.e., $\mathbf{p}\approx\mathbf{e}_{m}$ for some $m$).

A widely studied instance of smoothed quantization is Gumbel-softmax sampling (Jang et al.,, 2017).
Given assignment probabilities $\pi_{m}$ of $\mathbf{z}$ to the $m$-th code—typically log-proportional to their dot-product $\mathbf{q}_{m}^{\mathsf{T}}\mathbf{z}$—categorical sampling can be implemented using Gumbel samples $g_{m}=-\log(-\log u_{m})$ with $u_{m}\sim\operatorname{Uniform}(0,1)$:

$$ $\displaystyle\iota(\mathbf{z})\sim\operatorname{Categorical}(\pi_{1},\dots,\pi_{M})$ $\displaystyle\Leftrightarrow$ $\displaystyle\iota(\mathbf{z})=\underset{m}{\operatorname{argmax}}(g_{m}+\log\pi_{m})$ (6) $$

Replacing $\operatorname{argmax}$ above with $\operatorname{softmax}$ yields a smoothed quantization:

$$ $\displaystyle p_{m}=$ $\displaystyle\frac{\exp\left(\left(g_{m}+\log\pi_{m}\right)/\tau\right)}{\sum_{m^{\prime}}\exp\left(\left(g_{m^{\prime}}+\log\pi_{m^{\prime}}\right)/\tau\right)}$ (7) $$

where lowering the temperature parameter $\tau$ produces a tighter approximation of categorical sampling.

The Gumbel-softmax sampling can also be combined with hard quantization using the STE:

$$ $\displaystyle\operatorname{STE}(\mathbf{e}_{\iota(\mathbf{z})},\mathbf{p})=\mathbf{e}_{\iota(\mathbf{z})}+\mathbf{p}-\mathtt{detach}(\mathbf{p})$ (8) $$

### 2.3 Regularization

Beyond approximation strategies, vector quantization also requires auxiliary regularization losses to ensure effective training.
For example, the STE alone does not guarantee alignment between pre-quantized features and codebook entries.
This alignment is instead fostered by the following regularization loss, $\mathcal{L}_{\text{reg}}$ (van den Oord et al.,, 2017; Fifty et al.,, 2025):

$$ $\displaystyle\mathcal{L}_{\text{total}}$ $\displaystyle=\mathcal{L}_{\text{main}}+\mathcal{L}_{\text{reg}}$ (9) $\displaystyle\mathcal{L}_{\text{reg}}$ $\displaystyle=\mathcal{L}_{\text{hard}}:=N^{-1}\sum_{i=1}^{N}\left(\beta\underbracket{\|\mathbf{z}_{i}-\mathtt{detach}(\mathbf{q}_{\iota(\mathbf{z}_{i})})\|^{2}}_{\text{Commitment Loss}}+\underbracket{\|\mathtt{detach}(\mathbf{z}_{i})-\mathbf{q}_{\iota(\mathbf{z}_{i})}\|^{2}}_{\text{Codebook Loss}}\right)$ (10) $$

where $\mathcal{L}_{\text{main}}$ is the primary task loss (e.g., L2 regression in autoencoding), and $\beta>0$ is a weighting hyperparameter.
As defined in in Equation [10](https://arxiv.org/html/2509.22161v1#S2.E10), $\mathcal{L}_{\text{hard}}$ consists of two components.
The first term (known as the *commitment loss*) aligns each pre-quantized feature $\mathbf{z}_{i}$ in a batch ($i=1,\dots,N$) with their nearest codebook entry $\mathbf{q}_{\iota(\mathbf{z}_{i})}$.
The second term (called the *codebook loss*) moves each codebook vector $\mathbf{q}_{m}$ toward the centroid of its assigned features whose nearest neighbor is $\mathbf{q}_{m}$ (i.e., $\{\mathbf{z}_{i}:\iota(\mathbf{z}_{i})=m\}$).

It should be noted, however, that the codebook loss in Equation [10](https://arxiv.org/html/2509.22161v1#S2.E10) does not inherently prevent code collapse.
Some codebook vectors may never serve as the nearest neighbor of any pre-quantized feature and therefore receive no updates.
Accordingly, previous studies have resorted to additional workarounds; for example, unused codebook vectors may be reset to the positions of pre-quantized features Dhariwal et al., (2020).

Similarly, smoothed quantization also requires auxiliary regularization to avoid code collapse.
A widely adopted option is the normalized perplexity of the mean assignment probability
(Dieleman et al.,, 2018; Baevski et al., 2020b,):

$$ $\displaystyle\mathcal{L}_{\text{reg}}$ $\displaystyle=\mathcal{L}_{\text{ppl}}:=\frac{\exp\left(-\sum_{m}\bar{\pi}_{m}\log\bar{\pi}_{m}\right)}{M}$ (11) $\displaystyle\bar{\pi}_{m}$ $\displaystyle:=N^{-1}\sum_{i=1}^{N}\pi_{i,m}$ (12) $$

As noted in §[1](https://arxiv.org/html/2509.22161v1#S1), this perplexity-based regularization does not promote the onehotness of $\boldsymbol{\pi}_{i}$, failing to distinguish onehot-like samples from uniform or centered ones (Figure [1](https://arxiv.org/html/2509.22161v1#S1.F1)).
Onehotness can instead be induced by annealing the temperature parameter of the Gumbel-softmax sampling ($\tau\to 0$).
However, manually scheduling this annealing is empirically challenging.
The next section therefore proposes an alternative regularization loss that automatically encourages onehotness within the framework of gradient-based learning, while simultaneously preventing code collapse.

## 3 Methods

As noted in previous sections, ideal smoothed quantizers $\mathbf{p}\in\Delta^{M-1}$ are distributed near the vertices of the simplex (Figure [1](https://arxiv.org/html/2509.22161v1#S1.F1)).
Moreover, each vertex should have at least some smoothed quantizers in its neighborhood; otherwise, the quantization suffers from code collapse.
A simple way to achieve these objectives simultaneously is to impose a loss penalizing the deviation of the KNNs from each vertex (Figure [1](https://arxiv.org/html/2509.22161v1#S1.F1)):

$$ $\displaystyle\mathcal{L}_{\text{reg}}$ $\displaystyle=\mathcal{L}_{\text{KNN}}:=(MK)^{-1}\sum_{m=1}^{M}\sum_{k=1}^{K}\mathcal{D}(\mathbf{e}_{m},\mathbf{p}^{(m,k)})$ (13) $$

where $\mathbf{p}^{(m,k)}$ denotes the $k$-th nearest neighbor of the simplex vertex (onehot vector) $\mathbf{e}_{m}$ according to a distance/deviation metric $\mathcal{D}$.
Two options for $\mathcal{D}$ are considered in this study: the squared L2 distance, $\|\mathbf{e}_{m}-\mathbf{p}^{(m,k)}\|^{2}$, and cross-entropy, $-\log p^{(m,k)}_{m}$.

At first glance, the proposed regularization may appear similar to the commitment and codebook losses used in the gradient-estimation approaches (Equation [10](https://arxiv.org/html/2509.22161v1#S2.E10)),
Both involve nearest neighbors, but
the two methods differ in their choice of anchors and neighbors.
While commitment and codebook losses take data points as anchors and identify their nearest codebook entries, the proposed method uses codebook entries as anchors and treats data as neighbors. Accordingly, the proposed regularization ensures that every codebook entry receives optimization feedback, whereas commitment/codebook losses may leave some entries untrained if they never become the nearest neighbor of any data point.

A further advantage of the proposed regularization is that Gumbel-softmax sampling is no longer required; smoothed quantizers can be obtained directly as $\mathbf{p}=\boldsymbol{\pi}=\operatorname{softmax}(Q^{\mathsf{T}}\mathbf{z})$.(^2^22The exact implementation of $\mathbf{p}$ and $\boldsymbol{\pi}$ involves normalization and rescaling; see Appendix [A.1](https://arxiv.org/html/2509.22161v1#A1.SS1) for details.)
At the same time, the proposed regularization is fully compatible with Gumbel-softmax sampling; one can simply replace $\mathbf{p}$ (Gumbel-softmax samples) in Equation [13](https://arxiv.org/html/2509.22161v1#S3.E13) with $\boldsymbol{\pi}$ (assignment probabilities).

During inference, hard quantization is applied by taking $\operatorname{argmax}_{m}p_{m}$ in the onehot representation.
In the following section, both the deterministic and stochastic approaches are evaluated on representative benchmarks.

## 4 Experiments

The proposed regularization for smoothed vector quantization was benchmarked on two tasks: discrete autoencoding (§[4.1](https://arxiv.org/html/2509.22161v1#S4.SS1)) and contrastive learning (§[4.2](https://arxiv.org/html/2509.22161v1#S4.SS2)).
The Python code used for these experiments will be published when this manuscript is accepted for a journal/conference.

### 4.1 Discrete Autoencoding

The first benchmark assessed the proposed regularization in the context of discrete autoencoding on the ImageNet dataset (Deng et al.,, 2009).
Input images were convolutionally encoded into latent feature maps, whose pixels were then quantized (Esser et al.,, 2021; Fifty et al.,, 2025).
The decoder convolutional network reconstructed the input images from these quantized feature maps, and the entire model was trained to minimize the L2 reconstruction loss ($L_{\text{main}}$ in Equation [9](https://arxiv.org/html/2509.22161v1#S2.E9)).
Further details about the network architecture and training setup are provided in Appendix [B.1](https://arxiv.org/html/2509.22161v1#A2.SS1).

Extending prior work (Esser et al.,, 2021; Fifty et al.,, 2025), three different combinations of feature map and codebook sizes were examined.
In addition to the previously used settings of $H\times W\times C=16\times 16\times 32$ ($M=1024$) and
$64\times 64\times 3$ ($M=8196$), an additional configuration increased the channel dimensionality of the latter to $32$
($H\times W\times C=64\times 64\times 32,M=8196$).

Each model was trained using four GPUs, and the proposed method computed the $K/4$-nearest neighbors of each simplex vertex per GPU (see §[5.3](https://arxiv.org/html/2509.22161v1#S5.SS3) for further discussion).
The number of neighbors was set to $K=2\times 4$ for the spatially smaller feature maps of size $16\times 16\times 32$, and to $K=8\times 4$ for the larger ones.(^3^33The value of $K$ was upper-bounded at $8\times 4$ by available computational resources; With a maximum batch size of $64$, the total number of latent pixels was $64\times 64\times 64=8\times 4\times 8196$, allowing only $8\times 4$ neighbors per vertex of $\Delta^{8196}$. This implementation constraint is further discussed in §[5.3](https://arxiv.org/html/2509.22161v1#S5.SS3).)

**Table 1: Performance of discrete autoencoding on the ImageNet validation set. Reported metrics are codebook usage and reconstruction quality scores: root mean squared error (rMSE), Fréchet Inception Distance (FID), and Inception Score (IS). The proposed method is denoted as “KNN-L2/CE”.**
| Method | Feature Map Size; Codebook Size |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 16$\times$16$\times$32; 1024 | 64$\times$64$\times$3; 8196 | 64$\times$64$\times$32; 8196 |  |  |  |  |  |  |  |  |  |  |  |
| Code Use ($\uparrow$) | rMSE ($\downarrow$) | FID ($\downarrow$) | IS ($\uparrow$) | Code Use | rMSE | FID | IS | Code Use | rMSE | FID | IS |  |  |
| STE | Euclid | 4.5% | 0.404 | 124.855 | 36.608 | 100.0% | 0.167 | 7.246 | 402.837 | 1.7% | 0.235 | 22.019 | 290.243 |
| Cosine | 3.0% | 0.381 | 117.950 | 41.475 | 70.9% | 0.197 | 13.972 | 348.307 | 2.8% | 0.186 | 12.108 | 363.862 |  |
| RE | Euclid | 3.1% | 0.460 | 170.300 | 19.253 | 78.85% | 0.171 | 10.2086 | 377.5555 | 0.5% | 0.271 | 40.575 | 203.515 |
| Cosine | 2.8% | 0.423 | 157.767 | 24.327 | 99.5% | 0.19 | 14.672 | 344.166 | 4.4% | 0.180 | 10.833 | 372.889 |  |
| H-Gmb | PPL | 100.0% | 0.349 | 100.286 | 54.042 | 100.0% | 0.189 | 19.069 | 321.164 | 100.0% | 0.163 | 10.331 | 384.041 |
| KNN-L2 | 20.5% | 0.354 | 100.467 | 56.053 | 99.9% | 0.222 | 25.002 | 280.972 | 23.6% | 0.185 | 10.718 | 376.354 |  |
| KNN-CE | 100.0% | 0.390 | 110.974 | 44.817 | 100.0% | 0.226 | 12.787 | 349.677 | 100.0% | 0.211 | 9.343 | 371.964 |  |
| S-Gmb | PPL | 55.0% | 0.826 | 173.593 | 9.605 | 100.0% | 0.183 | 10.278 | 377.787 | 48.4% | 0.386 | 45.606 | 167.142 |
| KNN-L2 | 100.0% | 0.371 | 81.479 | 71.946 | 100.0% | 0.201 | 10.168 | 373.830 | 100.0% | 0.193 | 5.727 | 404.184 |  |
| KNN-CE | 100.0% | 0.569 | 194.736 | 16.595 | 100.0% | 0.227 | 13.647 | 345.019 | 100.0% | 0.218 | 9.205 | 366.454 |  |
| Softmax | PPL | 100.0% | 1.112 | 309.603 | 4.552 | 81.6% | 0.701 | 37.987 | 196.723 | 99.8% | 0.725 | 83.070 | 77.981 |
| KNN-L2 | 100.0% | 0.366 | 81.691 | 71.750 | 100.0% | 0.199 | 12.644 | 361.228 | 100.0% | 0.204 | 7.908 | 381.654 |  |
| KNN-CE | 100.0% | 0.380 | 106.048 | 50.750 | 100.0% | 0.219 | 14.241 | 343.228 | 100.0% | 0.199 | 5.267 | 404.098 |  |

Table [3](https://arxiv.org/html/2509.22161v1#A2.T3) reports the codebook usage and reconstruction quality scores—including root mean squared error (rMSE), Inception Score (IS; Salimans et al.,, 2016), and Fréchet Inception Distance (FID; Heusel et al.,, 2017)–for the proposed method (“KNN-L2/CE”) and the baselines.
The results indicate that reconstruction through softmax-smoothed quantization (i.e., without Gumbel randomness) deteriorated substantially when combined with the perplexity-based regularization (“PPL”; Equation [11](https://arxiv.org/html/2509.22161v1#S2.E11)).
This degradation stems from the mismatch between soft quantization during the training and hard quantization at inference,
supporting the argument made in the Introduction that perplexity-based regularization alone does not promote tight smoothing (i.e., vertex-concentrated distributions).

The perplexity-based regularization was only effective when combined with the Gumbel-softmax sampling,
and onehot quantization in the forward computation by STE was necessary to ensure full codebook usage (“Hard-Gumbel”).
Otherwise, code collapse was not prevented when the channel dimensionality was large ($C=64$; “Soft-Gumbel”).

By contrast, the proposed KNN-based regularization successfully prevented code collapse and achieved near-complete codebook utilization without resorting to Gumbel-softmax sampling.
The reconstruction quality was also superior or competitive with the Gumbel-softmax + perplexity approach across all feature map sizes.

The choice of deviation metric (L2 distance vs. cross-entropy, CE) did not result in consistent global superiority, although noticeable code collapse occurred when KNN-L2 regularization was combined with hard Gumbel-softmax sampling, which nonetheless had limited effect on the reconstruction quality.

It is also worth noting that both STE and rotational gradient estimation (RE) suffered from code collapse when the channel dimensionality was large ($C=64$),
illustrating the difficulty of achieving robust prevention of code collapse across different settings.
The next section further highlights this challenge under a different training objective.

### 4.2 Contrastive Learning

The second experiment evaluated vector quantization methods within Wav2Vec 2.0 pretraining for speech feature extraction (Baevski et al., 2020b,).
Unlike autoencoding, this pretraining integrates vector quantization directly into the main loss function.

Two codebook configurations were investigated.
The first used a single codebook of size 1024.
By contrast, the second followed the original work of Baevski et al., 2020b, combining two smaller codebooks—each of size 320—to implement rich code diversity efficiently via *product quantization* (Jégou et al.,, 2011, see §[5.3](https://arxiv.org/html/2509.22161v1#S5.SS3) for more information).
The dimensionality of the codebook vectors was set to 256 for the single-codebook configuration and to 128 for the dual-codebook configuration.
The number of nearest neighbors was set to $K=2\times 4$, meaning that two neighbors were identified per codebook entry on each of the four GPUs.
In the dual-codebook condition, the two codebooks were treated independently.

Models were trained on the LibriSpeech dataset, combining all training splits (train-clean-100 + train-clean-360 + train-other-500; Panayotov et al.,, 2015).
Further details on the network architecture and learning objective are provided in Appendix [B.2](https://arxiv.org/html/2509.22161v1#A2.SS2).

**Table 2: Codebook usage in Wav2Vec 2.0 pretraining, evaluated on the LibriSpeech dev-clean split (similar results were observed for the other dev/test splits). The proposed method is denoted as “KNN-L2/CE”.**
|  |  | #Codebooks $\times$ Codebook Size |  |  |
| --- | --- | --- | --- | --- |
| Method | 1$\times$1024 | 2$\times$320 |  |  |
|  |  |  | Codebook#1 | Codebook#2 |
| STE | Euclid | 0.2% | 0.9% | 0.9% |
| Cosine | 0.2% | 0.6% | 0.6% |  |
| RE | Euclid | 0.2% | 0.6% | 0.6% |
| Cosine | 0.2% | 0.6% | 0.6% |  |
| H-Gmb | PPL | 0.7% | 0.6% | 0.6% |
| KNN-L2 | 0.2% | 0.6% | 0.6% |  |
| KNN-CE | 99.7% | 100.0% | 100.0% |  |
| S-Gmb | PPL | 0.3% | 0.6% | 0.6% |
| KNN-L2 | 90.1% | 0.6% | 0.6% |  |
| KNN-CE | 100.0% | 100.0% | 100.0% |  |
| Softmax | PPL | 0.2% | 1.2% | 1.2% |
| KNN-L2 | 60.4% | 100.0% | 100.0% |  |
| KNN-CE | 100.0% | 100.0% | 100.0% |  |

Table [4](https://arxiv.org/html/2509.22161v1#A2.T4) reports codebook usage. The perplexity-based regularization failed to prevent code collapse in both single- and dual-codebook settings.
Likewise, the STE and rotational gradient estimation approaches exhibited the same failure.

By contrast, the proposed KNN-based regularization ensured complete code utilization in both conditions when cross-entropy was used as the divergence metric on the simplex.
The L2 metric, on the other hand, proved insufficient to prevent code collapse, especially when combined with Gumbel sampling.

## 5 Discussions

### 5.1 Summary of Findings & Contributions

This study introduced a simple and unified regularization method that simultaneously tightens smoothed vector quantization and promote effective code utilization.
The proposed method successfully prevented code collapse in two representative applications of vector quantization:
a middle layer in discrete autoencoding (§[4.1](https://arxiv.org/html/2509.22161v1#S4.SS1)) and target construction in contrastive learning (§[4.2](https://arxiv.org/html/2509.22161v1#S4.SS2)).
This robustness is noteworthy, as prior approaches were effective only in specific settings and remained vulnerable to code collapse in others.

The proposed method is geometrically intuitive and straightforward, yet appears unaddressed in the existing literature.
Research on neural vector quantization has traditionally been rooted in variational autoencoding (Kingma and Welling,, 2014), primarily aiming to extend this stochastic framework to discrete variables (Jang et al.,, 2017; van den Oord et al.,, 2017).
The issue of code collapse was recognized (or documented) later, and workarounds were developed independently of the quantization methods themselves (Dieleman et al.,, 2018; Baevski et al., 2020b,; Dhariwal et al.,, 2020).

By contrast, the present work reformulates neural vector quantization as a simple smoothing problem: onehot vectors are approximated by elements of the simplex.
Within this perspective, concentrating the approximators near the simplex vertices naturally arises as a desirable property.
This reformulation, together with the proposed regularization strategy, represents a key conceptual contribution of the study.

### 5.2 Alternative Implementations of the Intended Regularization

Alternative regularization strategies could also achieve the intended distribution of smoothed quantizers ($\mathbf{p}$ or $\boldsymbol{\pi}$) around the simplex vertices.
For example, one could align smoothed quantizers with a Dirichlet distribution with concentration parameters $\alpha_{1}=\cdots=\alpha_{K}<1.0$ (Figure [2](https://arxiv.org/html/2509.22161v1#S5.F2)).
The probability density of such a distribution is highest at the vertices of the simplex, thus matching the ideal distribution of smoothed quantizers.
A possible formalization of this alignment is based on the Kullback-Leibler (KL) divergence between the Dirichlet prior ($\mathbb{P}$) and the distribution inferred from smoothed quantizers ($\mathbb{Q}$).

$$ $\displaystyle\mathcal{D}_{\text{KL}}\left(\mathbb{P}\mid\mathbb{Q}\right):=$ $\displaystyle\int_{\Delta^{M-1}}\mathbb{P}(\mathbf{p})\log\frac{\mathbb{P}(\mathbf{p})}{\mathbb{Q}(\mathbf{p})}d\mathbf{p}$ (14) $$

However, Equation [14](https://arxiv.org/html/2509.22161v1#S5.E14) is difficult to use directly as a regularization loss, due to the complexity of estimating $\mathbb{Q}$ from sample quantizers, $\mathbf{p}$.
Although one could constrain $\mathbb{Q}$ as another Dirichlet to make the KL divergence tractable,
maximum likelihood estimation of its parameters requires iterative algorithms (e.g., the Newton-Raphson method; Ronning,, 1989; Sklar,, 2014; Wicker et al.,, 2008), complicating and slowing gradient-based optimization in deep learning frameworks.
Moreover, this estimation involves digamma and trigamma functions (Sklar,, 2014), whose derivatives can explode when the concentration parameters approach small values—as desired for tight smoothing—during the course of learning.

Figure: Figure 2: Dirichlet distributions on the simplex $\Delta^{3-1}$ with concentration parameters $\alpha_{1}=\alpha_{2}=\alpha_{3}=\alpha$, where $\alpha\in\{0.5,1.0,2.0\}$.
Refer to caption: x2.png

A more practical approach is to approximate Equation [14](https://arxiv.org/html/2509.22161v1#S5.E14) itself in a tractable manner.
For instance, Perez-Cruz, (2008) proposed a KNN-based estimation of the KL divergence that relies solely on *samples* from the two distributions, using the $k$-th nearest neighbor from $\mathbb{Q}$ for each sample from $\mathbb{P}$.
The regularization method proposed in this study can thus be interpreted as minimizing this estimated KL divergence, with the samples from $\mathbb{P}$ constrained to onehot vectors.

### 5.3 Limitations

A primary limitation of the proposed method is its memory requirement.
When training across multiple GPUs, the KNN-based regularization identifies $K$ nearest smoothed quantizers ($\mathbf{p}$ or $\boldsymbol{\pi}$) per simplex vertex *on each GPU*, rather than finding global neighbors across all GPUs.
Consequently, each GPU must have sufficient VRAM to store at least $KM$ latent pixels/frames, where $M$ denotes the codebook size.
This requirement can become prohibitive when $M$ is large (i.e., for fine-grained quantization).

One possible workaround is to randomly select a subset of simplex vertices when computing the regularization loss, rather than using all vertices in a single iteration.
In expectation, this achieves the same effect as the original implementation, although its empirical effectiveness remains to be assessed in future studies.

Additionally, fine-grained quantization can be achieved more efficiently using smaller $G$ codebooks in combination (product quantization; Jégou et al.,, 2011), as explored in the Wav2Vec 2.0 pretraining (§[4.2](https://arxiv.org/html/2509.22161v1#S4.SS2)).
This approach represents $\prod_{g=1}^{G}M_{g}$ distinct quantized vectors while requiring only $K\sum_{g=1}^{G}M_{g}$ smoothed quantizers per GPU.
Leveraging these strategies, the proposed method can overcome its limitation and become applicable to real-world scenarios.

#### Acknowledgments

This study was supported by
JSPS Grant-in-Aid for Scientific Research
A (JP24H00774),
B (JP22H03914),
and C (JP24K15087);
JST AIP Accelerated Program (JPMJCR25U6) and Core Research for Evolutional Science and Technology (JPMJCR22P5);
and Kayamori Foundation of Informational Science Advancement (K35XXVIII620).
The author also gratefully acknowledges the support of the
Academic Center for Computing and Media Studies, Kyoto University,
regarding the use of their supercomputer system.

## References

- (1)
Baevski, A., Schneider, S., and Auli, M. (2020a).
vq-wav2vec: Self-supervised learning of discrete speech representations.
In Proceedings of the 8th International Conference on Learning Representations (ICLR). OpenReview.net.
- (2)
Baevski, A., Zhou, Y., Mohamed, A., and Auli, M. (2020b).
wav2vec 2.0: A framework for self-supervised learning of speech representations.
In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M. F., and Lin, H., editors, Advances in Neural Information Processing Systems, volume 33, pages 12449–12460. Curran Associates, Inc.
- Chollet, (2017)
Chollet, F. (2017).
Xception: Deep learning with depthwise separable convolutions.
In 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 1800–1807.
- Deng et al., (2009)
Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. (2009).
ImageNet: A large-scale hierarchical image database.
In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248–255.
- Dhariwal et al., (2020)
Dhariwal, P., Jun, H., Payne, C., Kim, J. W., Radford, A., and Sutskever, I. (2020).
Jukebox: A generative model for music.
- Dieleman et al., (2018)
Dieleman, S., van den Oord, A., and Simonyan, K. (2018).
The challenge of realistic music generation: modelling raw audio at scale.
In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc.
- Esser et al., (2021)
Esser, P., Rombach, R., and Ommer, B. (2021).
Taming Transformers for high-resolution image synthesis.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 12868–12878, Los Alamitos, CA, USA. IEEE Computer Society.
- Fifty et al., (2025)
Fifty, C., Junkins, R. G., Duan, D., Iyengar, A., Liu, J. W., Amid, E., Thrun, S., and Ré, C. (2025).
Restructuring vector quantization with the rotation trick.
In The Thirteenth International Conference on Learning Representations, ICLR 2025, Singapore, April 24-28, 2025. OpenReview.net.
- Gray, (1984)
Gray, R. (1984).
Vector quantization.
IEEE ASSP Magazine, 1(2):4–29.
- Hadjeres and Crestel, (2020)
Hadjeres, G. and Crestel, L. (2020).
Vector quantized contrastive predictive coding for template-based music generation.
- He et al., (2016)
He, K., Zhang, X., Ren, S., and Sun, J. (2016).
Deep residual learning for image recognition.
In The IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 770–778.
- Heusel et al., (2017)
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. (2017).
Gans trained by a two time-scale update rule converge to a local nash equilibrium.
In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.
- Jang et al., (2017)
Jang, E., Gu, S., and Poole, B. (2017).
Categorical reparameterization with Gumbel-softmax.
In Proceedings of the 5th International Conference on Learning Representations (ICLR).
- Jégou et al., (2011)
Jégou, H., Douze, M., and Schmid, C. (2011).
Product quantization for nearest neighbor search.
IEEE Transactions on Pattern Analysis and Machine Intelligence, 33(1):117–128.
- Kingma and Welling, (2014)
Kingma, D. P. and Welling, M. (2014).
Auto-encoding variational bayes.
The International Conference on Learning Representations (ICLR) 2014.
- Panayotov et al., (2015)
Panayotov, V., Chen, G., Povey, D., and Khudanpur, S. (2015).
Librispeech: An ASR corpus based on public domain audio books.
In 2015 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 5206–5210, South Brisbane, Queensland, Australia.
- Perez-Cruz, (2008)
Perez-Cruz, F. (2008).
Kullback-leibler divergence estimation of continuous distributions.
In 2008 IEEE International Symposium on Information Theory, pages 1666–1670.
- Ramesh et al., (2021)
Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C., Radford, A., Chen, M., and Sutskever, I. (2021).
Zero-shot text-to-image generation.
In Meila, M. and Zhang, T., editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8821–8831. PMLR.
- Rombach et al., (2022)
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. (2022).
High-resolution image synthesis with latent diffusion models.
In 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 10674–10685.
- Ronning, (1989)
Ronning, G. (1989).
Maximum likelihood estimation of dirichlet distributions.
Journal of Statistical Computation and Simulation, 32(4):215–221.
- Salimans et al., (2016)
Salimans, T., Goodfellow, I., Zaremba, W., Cheung, V., Radford, A., Chen, X., and Chen, X. (2016).
Improved techniques for training GANs.
In Lee, D., Sugiyama, M., Luxburg, U., Guyon, I., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc.
- Sklar, (2014)
Sklar, M. (2014).
Fast MLE computation for the dirichlet multinomial.
- van den Oord et al., (2017)
van den Oord, A., Vinyals, O., and Kavukcuoglu, K. (2017).
Neural discrete representation learning.
In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R., editors, Advances in Neural Information Processing Systems 30, pages 6306–6315. Curran Associates, Inc.
- Wicker et al., (2008)
Wicker, N., Muller, J., Kalathur, R. K. R., and Poch, O. (2008).
A maximum likelihood approximation method for dirichlet’s parameter estimation.
Computational Statistics & Data Analysis, 52(3):1315–1322.
- Yu et al., (2022)
Yu, J., Xu, Y., Koh, J. Y., Luong, T., Baid, G., Wang, Z., Vasudevan, V., Ku, A., Yang, Y., Ayan, B. K., Hutchinson, B., Han, W., Parekh, Z., Li, X., Zhang, H., Baldridge, J., and Wu, Y. (2022).
Scaling autoregressive models for content-rich text-to-image generation.
Transactions on Machine Learning Research.

## Appendix A Implementation of the Quantization Methods

This section provides details on the implementation of the quantization methods.

### A.1 Smoothed Quantization

Smoothed quantizers $\mathbf{p}$ were computed as $\mathbf{p}=\operatorname{softmax}(\hat{Q}^{\mathsf{T}}\hat{\mathbf{z}}/\tau)$.
In other words, the codebook vectors $(\mathbf{q}_{1},\dots,\mathbf{q}_{M})=Q$ and the feature vectors $\mathbf{z}$ were first L2-normalized, and their product (i.e., cosine similarity) was rescaled by a learnable temperature $\tau$.
This temperature was shared across the codebook so that all the logits had the same amplitude.
Assignment probabilities $\mathcal{\pi}$ for Gumbel-softmax sampling were computed in the same way.(^4^44Previous studies manually annealed the temperature $\tau$ instead of treating it as a learnable parameter (Baevski et al., 2020b,). This approach was also tested in the experiments here but did not yield improvements over the learnable temperature.)

### A.2 Hard Quantization

The weight $\beta$ on the commitment loss in Equation [10](https://arxiv.org/html/2509.22161v1#S2.E10) was set to 1.0, based on the previous observations that its value does not significantly affect learning outcomes within the range 0.1–2.0. (van den Oord et al.,, 2017).

## Appendix B Details of the Experiments

### B.1 Discrete Autoencoding

This section provides implementation details for the autoencoding experiment described in §[4.1](https://arxiv.org/html/2509.22161v1#S4.SS1).

The network architecture followed prior work on discrete autoencoding of ImageNet (Esser et al.,, 2021; Fifty et al.,, 2025).
Input images were center-cropped to $H\times W\times C=256\times 256\times 3$.
The encoder first expanded the channel dimensionality of the input images from 3 to 256 by convolution, and then progressively downsampled them through a series of strided convolutions (see Table [3](https://arxiv.org/html/2509.22161v1#A2.T3) for the spatial and channel sizes at each layer).
Each downsampling layer was followed by a residual block (He et al.,, 2016).
The decoder reconstructed the input images by upsampling the latent feature maps with a sequence of interpolations and residual blocks.

To improve memory efficiency—particularly important for the proposed KNN-based regularization—all but the input and output layers were implemented as depthwise separable convolutions (Chollet,, 2017).
All convolutional kernels had size $3\times 3$.

Training employed the AdamW optimizer with $(\beta_{1},\beta_{2})=(0.9,0.99)$ and
a weight decay coefficient of $10^{-4}$,
except for Euclidean-based STE/rotational hard quantization, where weight decay was set to zero.
The learning rate was linearly warmed up from $0.0$ to $\rho_{\mathrm{max}}$, and subsequently annealed to $0.5\rho_{\mathrm{max}}$ by cosine scheduling.
The maximum learning rate $\rho_{\mathrm{max}}$ was set to $5\times 10^{-5}$ for Euclidean-based STE/rotational hard quantization, and to $10^{-4}$ for all other configurations (Fifty et al.,, 2025).

**Table 3: Hyperparameters for discrete autoencoding.**
| Feature Map Size | $16\times 16\times 32$ | $64\times 64\times 3$ | $64\times 64\times 32$ |
| --- | --- | --- | --- |
| Codebook Size | 1024 | 8196 | 8196 |
| Latent Channels | $128\to 128\to 64\to 64\to 32\to 32$ | $128\to 64\to 32\to 3$ | $128\to 64\to 32\to 32$ |
| Height & Width | $256\to 128\to 64\to 32\to 16\to 16$ | $256\to 128\to 64\to 64$ | $256\to 128\to 64\to 64$ |
| Batch Size | 64 | 64 | 64 |
| Training Epochs | 25 | 20 | 20 |
| Warmup Iterations | 16,000 | 16,000 | 16,000 |

Inception Score (IS; Salimans et al.,, 2016) and Fréchet Inception Distance (FID; Heusel et al.,, 2017) were estimated using the ImageNet-pretrained Inception V3 provided in torchvision.

### B.2 Wav2Vec 2.0

This section provides details for the Wav2Vec 2.0 pretraining discussed in §[4.2](https://arxiv.org/html/2509.22161v1#S4.SS2).

The model consisted of a convolutional feature encoder followed by a Transformer module (Baevski et al., 2020b,).
The convolutional encoder extracted latent feature sequences from input waveforms (16kHz$\to$50Hz).
Then, a subset of these latent vectors was masked and fed into the Transformer, whose outputs $\mathbf{y}_{t}$ were trained to predict the quantized version $\mathbf{q}_{t}$ of the masked vectors.
The masking scheme followed Baevski et al., 2020b; random 6.5% of the latent vectors were masked, together with the following 10 time steps.
The learning objective was:

$$ $\displaystyle\mathcal{L}_{\text{main}}=$ $\displaystyle-\log\frac{\exp(\hat{\mathbf{y}}_{t}\hat{\mathbf{q}}_{t}^{\mathsf{T}}/\mathcal{T})}{\sum_{\hat{\mathbf{q}}\sim\mathcal{Q}}\exp(\hat{\mathbf{y}}_{t}\hat{\mathbf{q}}^{\mathsf{T}}/\mathcal{T})}$ (15) $$

where $\hat{\cdot}$ denotes the L2-normalization of vectors (i.e., measuring the cosine similarity), and $\mathcal{T}:=0.1$ is the temperature parameter.
For each masked vector, a set of distractors $\mathbf{\tilde{q}}$ was sampled from the other quantized vectors according to a distribution $\mathcal{Q}$.

Both the convolutional encoder and Transformer were implemented using the publicly available code in torchaudio, and only the quantization module was implemented from scratch.

Input waveforms were randomly cropped to the length of 250k samples. Both stages employed the AdamW optimizer with $(\beta_{1},\beta_{2})=(0.9,0.99)$ and zero weight decay.
The learning rate was warmed up from $0.0$ to $5.0\times 10^{-4}$, and then annealed to $5.0\times 10^{-6}$ by cosine scheduling.

**Table 4: Hyperparameters for Wav2Vec 2.0.**
| Input Frequency | 16kHz |  |
| --- | --- | --- |
| Latent Frequency | 50Hz |  |
| Codebook Dimensionality | 256 |  |
| Codebook Size | 1024 |  |
| CNN |  |  |
| Latent Channels | 512 |  |
| Kernel Sizes | 10$\to$3$\to$3$\to$3$\to$3$\to$2$\to$2 |  |
| Strides | 5$\to$2$\to$2$\to$2$\to$2$\to$2$\to$2 |  |
| Transformer |  |  |
| # Layers | 12 |  |
| Model Dimensionality | 768 |  |
| # Heads | 8 |  |
| Feed-Forward Dimensionality | 4096 |  |
| Dropout Rate | 0.1 |  |
| Layer Drop | 0.05 |  |
|  | #Codebooks$\times$Codebook Size |  |
|  | 1$\times$1024 | 2$\times$320 |
| Batch Size | 128 | 64 |
| Training Epochs | 128 | 20 |
| Warmup Iterations | 32,000 | 10,000 |

When the single-codebook condition was first examined, training was run for 128 epochs, following the original schedule (Baevski et al., 2020b,).
However, since convergence occurred rapidly, the number of epochs was reduced in the dual-codebook condition to improve time efficiency.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 4
presentation: 3
questions: - The experimental results show that the proposed method combined with L2 distance performs noticeably worse than KNN + CE, particularly on the contrastive learning task. My main concern is about the method’s robustness to the choice of distance metric:

Is the proposed regularizer sensitive to the form of the distance function?

Could you provide theoretical intuition or gradient-level analysis explaining why CE behaves better than L2 in contrastive settings?

Additional experiments or ablation studies comparing both distances under controlled settings would strengthen the claim.

- Please provide further analysis on the influence of number of neighbors (K):

When K is small, can the method still mitigate the soft/hard assignment mismatch observed in Perplexity regularization?

Please include an ablation study on varying K and compare it with other baselines to measure effectiveness of the proposed method.

- Because the value of K directly affects GPU memory usage, please report experiments for both baselines and the proposed method under similar GPU resource.
soundness: 2
strengths: - Empirical results support that KNN is superior to Perplexity regularizer, as it effectively reduce the mismatch between soft / hard assignment when training and testing.

- Experiment results in representation learning highlights the robustness of the proposed method when adopted with cross-entropy (KNN-CE), which achieves 100% usage in both single and dual codebooks, where other methods fail.

- The paper is generally clear. The Related Work section provides not only a concise overview of recent advances in vector quantization but also a coherent narrative that situates the proposed method within the broader literature. It effectively summarizes the fundamental components of VQ, such as the quantization step, gradient estimation, and smoothing strategies, making it accessible and informative even to readers new to the topic.
summary: - This paper proposes a simple yet effective KNN-based regularization method to stabilize smoothed vector quantization. A well-behaved quantizer should produce one-hot–like assignments while maintaining balanced codebook usage. The proposed approach achieves both by penalizing the distance between each simplex vertex and its K nearest soft assignments. This encourages all codebook entries to be actively used while keeping each assignment sharply concentrated near a single vertex.


- As a result, the method mitigates the train-test mismatch that occurs in perplexity-based regularization, where training uses soft assignments but inference requires hard quantization.


- Experiments on discrete image autoencoding and contrastive speech representation learning demonstrate that the proposed regularizer effectively prevents code collapse, achieves near-complete code utilization, and delivers improved reconstruction and representation quality compared with other baselines.
weaknesses: - The paper lacks an ablation study on the number of neighbors (K), which is a key hyperparameter of the proposed regularizer.


- There is limited qualitative results in the experiment section. Beyond the simplex visualizations, it would be valuable to include additional qualitative examples such as reconstructed images or assignment-map visualizations to illustrate the impact of KNN-based over Perplexity regularization.


- Lack of experimental settings: details such as the exact loss components (reconstruction, codebook, commitment) and weighting used for baselines like STE, RE, and Perplexity regularization should be declared.

---

confidence: 4
contribution: 3
initial_rating: 4
presentation: 2
questions: 1. The paper might improve by adding more baseline methods to show how their method compares with current state-of-the-art methods
2. How does the method perform on Image generation (diffusion, flow-based, or autoregressive) or language models?
3. How sensitive is the performance of the method in terms of K
4. A dedicated ablation study and an explanation for the choice of hyperparameters can help in reproducibility.
soundness: 3
strengths: 1. Their method naturally incorporates all the entries of the codebook, which ensures maximum codebook utilization by design.
2. They achieve 100% codebook utilization among all the tasks they train, while performing on par or better than baseline methods
3. The regularization loss they propose is quite flexible and can be combined with other backbones, such as Gumbel softmax, L2 or cross entropy losses, etc.
summary: This paper proposes a new regularization method for smoothed vector quantization that simultaneously promotes near one-hot assignments and prevents codebook collapse. Their method primarily varies from prior methods in its choice of anchors and neighbours. More specifically, they use codebook entries as anchors and data as neighbours. Using this, the KNN-based regularizer encourages samples to cluster near all vertices, ensuring full codebook utilization and tight smoothing. They test their method on two VQ tasks (auto-encoding and contrastive learning) and show that their method performs on par with or better than the baseline methods in different codebook settings.
weaknesses: 1. Missing comparison with other baseline methods such as Group VQ, SimVQ etc.,
2. Since they use the codebook entries as anchors, some of the datapoints might not contribute to the regularization loss
3. Their method does not perform the best in some settings like in table 1, 64x64x3 setting, where the STE method performs the best

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. How does the proposed regularization scale to very large codebooks (e.g., M = 16k ~ 64k)? Have you explored memory-efficient approximations such as partial vertex sampling or approximate KNN search? 
2. How sensitive are the results to the choice of K (number of nearest neighbors)? Is there a principled way to select K relative to the codebook size or batch size? 
3. While the paper demonstrates improved code utilization and reconstruction performance, does this lead to improvements in downstream tasks (e.g., image generation or text-to-speech synthesis)?
soundness: 3
strengths: 1. The illustration and explanation of smoothed vector quantization from a simplex perspective provide a clear geometric intuition that unifies and clarifies connections with prior approaches. 
2. The proposed KNN regularization is an intuitive idea that offers strong empirical benefits compared to prior entropy or perplexity-based regularization methods.
summary: The paper addresses the issue of code collapse in smoothed vector quantization by proposing a simple K-nearest neighbor (KNN) regularization that simultaneously enforces tight smoothing and uniform code utilization. Empirical evaluations on discrete image autoencoding and contrastive speech representation learning demonstrate that the proposed approach achieves nearly complete codebook utilization and competitive or superior reconstruction quality compared to baselines including straight-through estimation, Gumbel-softmax, and perplexity-based regularization.
weaknesses: 1. While the method is designed to improve both smoothing tightness and codebook utilization, the paper does not provide quantitative analysis or visualization to support the claim about improved smoothing tightness. 
2. The KNN loss is acknowledged to be memory-inefficient, but the paper does not include any experiments or analysis to quantify the computational overhead or explore mitigation strategies.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - Please refer to the Weaknesses section. Overall I appreciate the technical contribution, but because the experiments appear insufficient and the discussion of recent related work is limited, I will keep the current score.
soundness: 3
strengths: - The paper is well written. The technical parts are clear and detailed.
- Smoothing the VQ process is an interesting idea. It provides new insight and motivates a simple KNN regularization to reduce code collapse. The loss is simple and effective.
- Experiments confirm the effectiveness of the proposed method across tasks.
summary: This paper addresses the non-differentiability of hard quantization by smoothing assignments while also aiming to avoid code collapse. It reformulates VQ as smoothing of one-hot vectors and adds a simple KNN regularization that pulls smoothed assignments toward all simplex vertices (using L2 or cross-entropy), compatible with softmax or Gumbel-softmax. On ImageNet autoencoding and wav2vec2 pretraining, the method shows reasonable performance against the reported baselines.
weaknesses: - First, the experimental scale is limited, and the paper mostly compares its own variants. There are few comparisons to recent methods, and the related work leans on pre‑2023 citations, with the rotation trick as the only recent item. As a result, the basis feels thin for an active area like VQ codebook learning.
- Second, more experiments are needed to study how batch size influences the KNN regularization. For example, compute the loss using groups within the same batch and analyze the trends. This would show whether the method scales to larger settings.
- Third, include an ablation of the Gumbel‑Softmax temperature to make the study more complete and to clarify its interaction with the proposed loss.


Decision: Reject