[Paper Review] Fast Decoding in Sequence Models using Discrete Latent Variables

Łukasz Kaiser, Aurko Roy+5 more|arXiv (Cornell University)|Mar 9, 2018

Natural Language Processing TechniquesComputer Science34 references178 citations

TL;DR

The paper introduces a Latent Transformer that auto-encodes target sequences into a shorter discrete latent sequence, enabling parallelizable decoding and achieving faster decoding with competitive BLEU in neural machine translation.

ABSTRACT

Autoregressive sequence models based on deep neural networks, such as RNNs, Wavenet and the Transformer attain state-of-the-art results on many tasks. However, they are difficult to parallelize and are thus slow at processing long sequences. RNNs lack parallelism both during training and decoding, while architectures like WaveNet and Transformer are much more parallelizable during training, yet still operate sequentially during decoding. Inspired by [arxiv:1711.00937], we present a method to extend sequence models using discrete latent variables that makes decoding much more parallelizable. We first auto-encode the target sequence into a shorter sequence of discrete latent variables, which at inference time is generated autoregressively, and finally decode the output sequence from this shorter latent sequence in parallel. To this end, we introduce a novel method for constructing a sequence of discrete latent variables and compare it with previously introduced methods. Finally, we evaluate our model end-to-end on the task of neural machine translation, where it is an order of magnitude faster at decoding than comparable autoregressive models. While lower in BLEU than purely autoregressive models, our model achieves higher scores than previously proposed non-autoregressive translation models.

Motivation & Objective

• Motivate fast decoding for autoregressive sequence models.
• Propose discrete latent bottlenecks to compress target sequences.
• Develop end-to-end trainable Latent Transformer leveraging discrete latents.
• Evaluate on neural machine translation and compare to autoregressive and non-autoregressive baselines.

Proposed method

• Introduce discrete latent bottlenecks (Gumbel-Softmax, improved semantic hashing, VQ-VAE, and decomposed vector quantization) to encode y into l where m < n.
• Train a Latent Transformer consisting of an autoencoder ae(y, x) to produce l, a latent predictor lp(x) to autoregressively generate l from x, and a decoder ad(l, x) to reconstruct y in parallel.
• Use a Transformer-based lp for autoregressive latent prediction and a parallel ad decoder conditioned on x and l.
• Optimize the joint objective lr + llp, where lr is the autoencoder reconstruction loss and llp is the latent prediction loss.
• Experiment with different discretization bottlenecks and DVQ variants to manage large latent alphabets and avoid codeword collapse.

Experimental results

Research questions

• RQ1 Can discrete latent variables compress a target sequence enough to allow parallel decoding without severe loss in translation quality?
• RQ2 Which discretization bottlenecks (Gumbel-Softmax, improved semantic hashing, VQ-VAE, DVQ and DVQ variants) yield best speed-quality tradeoffs for NMT?
• RQ3 How does the Latent Transformer compare to autoregressive and non-autoregressive baselines in BLEU and decoding latency?

Key findings

• The Latent Transformer achieves decoding that is an order of magnitude faster than autoregressive models in translation tasks.
• Using DVQ or improved semantic hashing yields strong BLEU with faster decoding, outperforming the hand-tuned non-autoregressive approach in several settings.
• VQ-VAE alone underperforms for this task, while decomposed DVQ with two decompositions maximizes latent usage and performance.
• With latent sequences of length n/m = 8, LT achieves BLEU scores around 19.7–19.8 with competitive decoding latency.
• Rescoring top-k translations using the LT approach closes the gap toward the autoregressive baseline without beam search.
• The method enables a trade-off between speed and accuracy by adjusting the ratio n/m and the latent alphabet size K.