[Paper Review] Theory and Experiments on Vector Quantized Autoencoders

Aurko Roy, Ashish Vaswani+2 more|arXiv (Cornell University)|May 28, 2018

Generative Adversarial Networks and Image SynthesisComputer Science27 references57 citations

TL;DR

The paper improves VQ-VAE training using an EM-inspired approach, yielding better image generation on CIFAR-10 and a fast, non-autoregressive translation model with distillation that closely approaches autoregressive Transformer performance.

ABSTRACT

Deep neural networks with discrete latent variables offer the promise of better symbolic reasoning, and learning abstractions that are more useful to new tasks. There has been a surge in interest in discrete latent variable models, however, despite several recent improvements, the training of discrete latent variable models has remained challenging and their performance has mostly failed to match their continuous counterparts. Recent work on vector quantized autoencoders (VQ-VAE) has made substantial progress in this direction, with its perplexity almost matching that of a VAE on datasets such as CIFAR-10. In this work, we investigate an alternate training technique for VQ-VAE, inspired by its connection to the Expectation Maximization (EM) algorithm. Training the discrete bottleneck with EM helps us achieve better image generation results on CIFAR-10, and together with knowledge distillation, allows us to develop a non-autoregressive machine translation model whose accuracy almost matches a strong greedy autoregressive baseline Transformer, while being 3.3 times faster at inference.

Motivation & Objective

• Motivate discrete latent representations for symbolic reasoning and data compression.
• Explore training strategies for discrete bottlenecks in VQ-VAE beyond existing heuristics.
• Leverage EM-inspired updates to improve discrete latent code learning.
• Demonstrate gains in image generation and machine translation through EM and distillation.

Proposed method

• Describe VQ-VAE with a discrete bottleneck and nearest-neighbor codebook lookup.
• Establish hard EM and K-means connections to VQ-VAE updates and EMA-based codebook learning.
• Introduce soft EM with Monte-Carlo EM updates for the discrete latent variables.
• Train Latent Predictor autoregressively on learned discrete latents and decode via the decoder.
• Apply sequence-level knowledge distillation to improve non-autoregressive translation.
• Evaluate on CIFAR-10 for unconditional image generation and WMT English-German for supervised translation.

Experimental results

Research questions

• RQ1 Can EM-inspired training improve the learning of discrete latents in VQ-VAE compared to prior heuristics?
• RQ2 Does soft EM with Monte-Carlo updates provide more stable and higher-quality learning than hard EM in VQ-VAE?
• RQ3 How do EM-trained VQ-VAE models perform on image generation tasks (CIFAR-10) and translation tasks (WMT English-German) relative to autoregressive and non-autoregressive baselines?
• RQ4 What is the impact of code-book size and distillation on translation quality and decoding speed?

Key findings

• EM training improves CIFAR-10 image generation over baseline VQ-VAE without product quantization.
• On WMT14 English-German, EM with distillation achieves BLEU of 26.7, close to a greedy Transformer (27.0) and faster by 3.3x.
• Soft EM with Monte-Carlo updates yields more stable and higher BLEU than hard EM in translation experiments.
• Code-book size of 2^12 yields best translation BLEU among tested sizes; larger sizes did not improve results.
• Non-autoregressive translation with EM and distillation achieves competitive BLEU while significantly reducing decoding latency (e.g., 81 ms vs 331 ms for autoregressive).
• EM-enabled VQ-VAE achieves 22.4 BLEU on WMT14 with EM, and 26.7 BLEU with distillation, illustrating strong gains from the proposed training.