Vector Quantized Variational Autoencoders (VQVAE) have produced remarkable results in multiple domains. VQVAE learns a prior distribution ze along with its mapping to a discrete number of K vectors (Vector Quantization). We propose applying VQ along the feature axis. We hypothesize that by doing so, we are learning a mapping between the codebook vectors and the marginal distribution of the prior feature space. Our approach leads to 33% improvement as compared to prevous discrete models and has similar performance to state of the art auto-regressive models (e.g. PixelSNAIL).
For exact benchmarks as reported in the paper please see branch tf1. Training for DVQ, final train loss 2.163272 as compared to VQVAE 3.2411757
Current training and evaluation for TF2 is in master branch Results It still requires additional fine-tuning and architectural changes to achieve same results as in TF1.
- Evalaute learned prior on conditioned PixelCNN and PixelCNN++