SQ-VAE: Variational Bayes on Discrete Representation with Self-annealed Stochastic Quantization

TL;DR

SQ-VAE introduces stochastic dequantization/quantization into the VAE/VQ-VAE framework, using trainable posteriors to realize self-annealing where quantization becomes increasingly deterministic during training. It provides Gaussian and von Mises–Fisher variants to handle continuous and categorical data, respectively, all within a standard variational Bayes objective, avoiding heuristic tricks like stop-gradient or EMA. Empirically, SQ-VAE improves codebook utilization and reconstruction quality across vision and speech tasks, and the vMF variant particularly excels on categorical data, underscoring the method’s versatility for discrete latent representations. The work suggests strong potential for improved data compression and scalable discrete latent modeling without manual hyperparameter tuning.

Abstract

One noted issue of vector-quantized variational autoencoder (VQ-VAE) is that the learned discrete representation uses only a fraction of the full capacity of the codebook, also known as codebook collapse. We hypothesize that the training scheme of VQ-VAE, which involves some carefully designed heuristics, underlies this issue. In this paper, we propose a new training scheme that extends the standard VAE via novel stochastic dequantization and quantization, called stochastically quantized variational autoencoder (SQ-VAE). In SQ-VAE, we observe a trend that the quantization is stochastic at the initial stage of the training but gradually converges toward a deterministic quantization, which we call self-annealing. Our experiments show that SQ-VAE improves codebook utilization without using common heuristics. Furthermore, we empirically show that SQ-VAE is superior to VAE and VQ-VAE in vision- and speech-related tasks.

Figures (10)

• Figure 1: Encoding and generative processes of SQ-VAE. The encoding path from $\mathbf{x}$ to $\mathbf{Z}_\mathrm{q}$ consists of (E1) deterministic encoding, (E2) stochastic dequantization, and (E3) quantization processes. For generation, in (G1) we first sample $\mathbf{Z}_\mathrm{q}\in\mathbf{B}^{d_z}$ from the prior $p(\mathbf{Z}_\mathrm{q})$. Then, in (G2) we feed $\mathbf{Z}_\mathrm{q}$ into the stochastic decoder to generate data samples.
• Figure 2: Empirical study on the dynamics related to $\sigma_{\bm\varphi}^2$ in Section \ref{['sec:sub_behavior_quantization']}. (a) The variance parameter $\sigma_{\bm\varphi}^2$ (blue) decreased with $\sigma^2$ (red), where $\sigma_0^2$ and $\sigma_{{\bm\varphi},0}^2$ are their initial values. (b) Average entropy of the quantization process w.r.t. the iteration, which is obtained by Monte Carlo estimation. (c) MSE for trainable $\sigma_{\bm\varphi}^2$ and various values of $\sigma_\mathrm{q}^2$ on the test set.
• Figure 3: vMF decoder.
• Figure 4: Comparison between vMF and NC decoders: (a) The concentration parameter of vMF decoder $\kappa_{\bm\varphi}$ increases with $\kappa$, whereas the growth of $\kappa_{\bm\varphi}$ of the NC decoder is relatively small. Here, $\kappa_0$ and $\kappa_{{\bm\varphi},0}$ indicate initial values. (b) Average entropy of probabilities of quantization processes.
• Figure 5: Empirical studies on the impact of codebook capacity examined on MNIST Fashion-MNIST and CIFAR10. (a)--(c) The size $K$ is swept with the dimension $d_b$ fixed to $64$. (d)--(f) Various $d_b$ values are tested with the size $K$ fixed as $128$, $256$, and $512$, respectively. The black lines with "$+$" marks indicate the upper bounds of the perplexities, i.e., $K$. All the y-axes are in log-scale.

Theorems & Definitions (4)

• Proposition 1
• Proposition 1
• proof
• Proposition 2