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

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$-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.

💡 Research Summary

Vector quantization (VQ) is a cornerstone technique for discretizing continuous latent spaces into a finite set of codebook vectors, enabling applications ranging from image generation to speech recognition. The primary obstacle in integrating VQ into deep learning pipelines is its non‑differentiable assignment step, which blocks gradient flow. Smoothed VQ alleviates this by relaxing the hard one‑hot assignment to a probability vector that lives inside the simplex Δ^{M‑1}. Effective smoothing must satisfy two complementary desiderata: (1) the softened assignment should stay close to a one‑hot vector (tight approximation), and (2) every codebook entry should be used during training (preventing code collapse).

Existing remedies treat these goals separately. Entropy‑ or perplexity‑maximizing regularizers push the average assignment distribution toward uniformity, thereby encouraging code usage, but they do not guarantee that individual assignments are one‑hot‑like. Conversely, annealing the temperature τ of a Gumbel‑Softmax sampler forces assignments to become sharper, yet it does not automatically ensure that all codes receive updates. Consequently, a unified approach that simultaneously enforces one‑hotness and full codebook utilization has been missing.

The present paper introduces a conceptually simple regularization term that directly addresses both objectives. For each simplex vertex e_m (the one‑hot vector for code m), the method finds its K nearest smoothed assignments p(m,k) (according to a distance metric D) and penalizes the average distance:

L_KNN = (1/(M·K)) Σ_{m=1}^M Σ_{k=1}^K D(e_m, p(m,k)).

Two distance choices are explored: squared L2 distance ‖e_m – p(m,k)‖² and cross‑entropy –log p(m,k)_m. By anchoring the loss on the codebook entries themselves, the regularizer guarantees that every code receives gradient feedback, unlike the conventional commitment/codebook losses that only affect codes that happen to be the nearest neighbor of a data point. Moreover, the approach works with deterministic softmax assignments p = softmax(Qᵀz) and does not require stochastic Gumbel‑softmax sampling, though it remains compatible with it.

Related work is organized into two families. The first approximates gradients (e.g., Straight‑Through Estimator, rotation‑based RE) while keeping the hard quantizer in the forward pass. The second smooths the forward quantizer (e.g., Gumbel‑Softmax, softmax) and adds auxiliary regularizers (entropy, commitment loss). The paper argues that prior regularizers either fail to enforce one‑hotness (entropy) or leave some codes unused (commitment), motivating the KNN‑based loss.

Empirical evaluation covers two domains:

1. Discrete Autoencoding on ImageNet – Various encoder‑decoder configurations are tested, differing in spatial resolution (16×16, 64×64) and channel depth (C = 3, 32, 2048), all with a codebook size M = 8196. K is varied among {1,2,4,8}. Results show:

   • Codebook utilization consistently exceeds 99 % for both L2‑KNN and CE‑KNN, whereas perplexity‑only regularization (with softmax) yields severe collapse.
   • Reconstruction quality (rMSE, FID, Inception Score) matches or surpasses baselines that combine Gumbel‑Softmax with perplexity maximization.
   • The method remains effective without any stochastic sampling; hard‑quantization via argmax at inference works seamlessly.
   • At extreme channel dimensionality (C = 2048), CE‑KNN retains full utilization while L2‑KNN occasionally drops a few codes, indicating cross‑entropy may be more robust when vectors become high‑dimensional.

2. Contrastive Speech Representation Learning – A VQ‑based front‑end is trained with a contrastive loss (e.g., CPC). The KNN regularizer again prevents code collapse and improves downstream phoneme classification accuracy compared to entropy‑only baselines. Importantly, the method works equally well with deterministic softmax assignments, simplifying the training pipeline.

The paper acknowledges computational constraints: finding K‑nearest neighbors per codebook entry scales with batch size and codebook size; in the experiments K was capped at 8 due to GPU memory limits. Nevertheless, the authors argue that even modest K values suffice to enforce the desired geometry of the simplex.

Limitations and future directions are discussed. Larger K would increase memory and runtime, suggesting the need for efficient approximate nearest‑neighbor algorithms or hierarchical sampling. Additionally, overly strong KNN regularization could diminish the natural annealing effect of τ, potentially requiring a separate temperature schedule for very low‑τ regimes.

In summary, the authors propose a unified, easy‑to‑implement regularizer that simultaneously tightens the smoothing of VQ assignments and guarantees exhaustive codebook usage. Extensive experiments across image and speech domains demonstrate that the KNN‑based loss outperforms or matches state‑of‑the‑art methods while eliminating the need for stochastic Gumbel‑Softmax sampling. This contribution offers a practical remedy to the long‑standing code collapse problem and paves the way for more stable and efficient vector‑quantized models.