SQ-VAE: Variational Bayes on Discrete Representation with Self-annealed Stochastic Quantization
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.
Keywords
Cite
@article{arxiv.2205.07547,
title = {SQ-VAE: Variational Bayes on Discrete Representation with Self-annealed Stochastic Quantization},
author = {Yuhta Takida and Takashi Shibuya and WeiHsiang Liao and Chieh-Hsin Lai and Junki Ohmura and Toshimitsu Uesaka and Naoki Murata and Shusuke Takahashi and Toshiyuki Kumakura and Yuki Mitsufuji},
journal= {arXiv preprint arXiv:2205.07547},
year = {2022}
}
Comments
25 pages with 10 figures, accepted for publication in ICML 2022 (Our code is available at https://github.com/sony/sqvae)