PaperlayerView exact revision on arXiv

Source-linked AI summary

Fast Decoding in Sequence Models using Discrete Latent Variables

Łukasz Kaiser, Aurko Roy, Ashish Vaswani, Niki Parmar, Samy Bengio, Jakob Uszkoreit, Noam Shazeer

arXiv:1803.03382v6cs.LG

TL;DR

Autoregressive sequence models remain sequential during decoding, limiting parallel computation. This paper uses shorter discrete latent sequences to enable faster decoding, achieving an order-of-magnitude speedup while retaining lower BLEU than autoregressive models and outperforming a hand-tuned non-autoregressive model.

  • Problem

    Autoregressive sequence models require at least n sequential decoding steps, limiting their ability to exploit parallel computation.

  • Method

    The paper extends autoregressive models with discrete latent variables, including the improved DVQ technique, to decode outputs from shorter latent sequences more efficiently.

  • Results

    An order of magnitude faster at decoding than autoregressive models, the model achieves better BLEU than a heavily hand-tuned task-specific non-autoregressive model while not recovering full autoregressive performance.

  • Takeaways & Limitations

    The latent transformer lays groundwork for fast decoding in sequence models beyond machine translation while remaining more generic than hand-tuned translation methods.

  • Takeaways & Limitations

    Latents are still generated autoregressively, taking most of the time for longer sentences, and the latent transformer does not yet recover full autoregressive performance.

Abstract

from arXiv · show

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.

1. Introduction

Autoregressive sequence models remain sequential during decoding, limiting parallel computation even when training is parallelizable. The paper introduces shorter discrete latent sequences and the Latent Transformer to decode outputs in parallel, achieving faster translation decoding.

  • Motivation: RNNs, WaveNet, ByteNet, and Transformers remain limited by sequential dependencies during autoregressive decoding.During decoding, predicting each output requires previously generated symbols, requiring at least n sequential steps for a sequence of length n.
  • Method: The proposed method generates a shorter discrete latent sequence autoregressively, then reconstructs the full output sequence in parallel.The latent sequence has length m < n, reducing the number of autoregressive steps before parallel reconstruction.
  • Discrete representations: Discrete latents are preferred for the autoencoder, but training such variables end-to-end is challenging.The paper compares Gumbel-Softmax, VQ-VAE, and improved semantic hashing, and introduces decomposed vector quantization.
  • Discrete representations: DVQ is introduced as an improved discretization technique that performs better than VQ-VAE for large latent alphabet sizes.
  • Results: The resulting Latent Transformer achieves good translation quality while decoding an order of magnitude faster than autoregressive models.The model is evaluated end-to-end on neural machine translation using DVQ or improved semantic hashing.

2. Discretization Techniques

The paper reviews discretization bottlenecks for discrete autoencoders, including Gumbel-Softmax, semantic hashing, and VQ-VAE, then develops decomposed vector-quantization variants. These variants address index collapse by decomposing the latent representation into smaller subspaces or projections.

  • General setup: Discrete autoencoders encode continuous representations into discrete latent codes before decoding.The encoder produces enc(y), a discretization bottleneck produces z_d(y) in [K], and the decoder receives z_q(y).
  • Gumbel-Softmax: Gumbel-Softmax projects encoder outputs into logits and uses differentiable samples during training before selecting near-one-hot codes at low temperature.The decoder input is formed from the embedding matrix, while the differentiable relaxation permits backpropagation.
  • Vector Quantization: VQ-VAE assigns each encoder output to the nearest embedding vector and trains the embeddings with reconstruction loss and exponential moving averages.The EMA tracks embeddings and the number of encoder states assigned to each embedding, using decay λ = 0.999 in the experiments.
  • Decomposed Vector Quantization: Large VQ-VAE codebooks can suffer index collapse, motivating decomposed vector quantization to use embedding vectors more efficiently.The paper describes a rich-getting-richer effect in which only a few embeddings receive substantial training signals and remain used.
  • Sliced Vector Quantization: Sliced vector quantization divides encoder outputs into smaller slices, assigns each slice its own embedding space, and preserves the overall code size through K′ = 2^(log2 K)/nd.When nd = 1 it reduces to VQ-VAE; when nd = log2 K it is loosely analogous to improved semantic hashing.
  • Projected Vector Quantization: Projected vector quantization uses fixed random projections to map encoder outputs into lower-dimensional subspaces before applying the same decomposed quantization procedure.With nd = 1 it reduces to VQ-VAE with an additional projection layer, while nd = log2 K gives the analogous semantic-hashing limit.

3. Latent Transformer

The Latent Transformer compresses target sequences into shorter discrete latent sequences, predicts those latents autoregressively, and decodes outputs from them in parallel. Its autoencoder uses input-aware convolutional compression and decompression, while training combines reconstruction and latent-prediction losses.

  • The Latent Transformer autoencodes y into a shorter discrete latent sequence l using a discretization bottleneck.
  • A Transformer latent-prediction model autoregressively predicts l from the input x.
  • The parallel decoder reconstructs y from l and x without direct dependencies among output tokens.
  • Training minimizes the sum of reconstruction loss lr and latent-prediction loss llp.The reconstruction loss compares decoded targets with y, while latent-prediction loss compares ae(y, x) with generated lp(x).
  • The encoder compresses targets with residual convolutions, input attention, and repeated stride-2 convolutions that reduce sequence length.The compressed representation is passed through the discretization bottleneck to produce ae(y, x) = zq(y).
  • The decoder reverses compression through repeated length-doubling up-convolutions, input attention, and a final self-attention decoder.Each up-convolution step doubles the sequence length before the result enters the Transformer-style decoder.

4. Related Work

Prior work established successful autoregressive neural sequence models but retained sequential decoding, motivating non-autoregressive alternatives and the Latent Transformer’s conditional-independence formulation. Existing non-autoregressive translation methods improved speed but were specialized to translation and required reinforcement-learning-based fine-tuning.

  • RNN sequence models are inherently sequential, whereas WaveNet, ByteNet, and Transformer improve training parallelism but remain sequential during decoding.
  • Discrete-latent training has used Gumbel-Softmax, VQ-VAE, and improved semantic hashing, while this work introduces decomposed vector quantization.The paper reports that DVQ performs better than VQ-VAE for large latent alphabet sizes.
  • Gu et al. combine a self-attention Transformer with REINFORCE to model word fertilities for non-autoregressive translation.
  • That approach requires extensive policy-gradient fine-tuning and is limited to machine translation rather than generic sequence learning.
  • The Latent Transformer assumes output tokens are conditionally independent given the latent sequence l and input x.This graphical-model structure enables parallel prediction of y1 ... yn from l and x.

5. Experiments

Experiments evaluate the Latent Transformer on WMT English-German translation against autoregressive and non-autoregressive baselines. Rescoring nearly matches the autoregressive baseline without beam search, while unrescored results show batching and discretization choices affect practical performance.

  • The experiments use the WMT English-German newstest2014 task and compare Latent Transformer results with reported Transformer and non-autoregressive baselines.The implementation uses around 33K subword units and is released with its hyperparameters and reproduction materials.
  • The reported baseline and comparison results distinguish LT, NAT, fertility training, and noisy parallel decoding on the WMT test set.
  • Noisy parallel decoding rescoring lets the method almost match the baseline autoregressive model without beam search.
  • Higher batch sizes substantially reduce Latent Transformer latency, while its non-batch decoding is slower than the simple NAT baseline.The authors attribute the non-batch difference possibly to system and implementation differences.
  • DVQ and improved semantic hashing yield good BLEU scores, whereas VQ-VAE fails in this setting.The discretization bottleneck has only a small apparent impact on decoding speed.

6. Discussion

The discussion examines discretization bottlenecks, latent-space usage, and speed–quality trade-offs in the Latent Transformer. DVQ with two decompositions improves latent utilization, while greater compression increases speed but makes reconstruction harder.

  • Translation evaluation: Table 2 compares BLEU scores and decoding times across Latent Transformer variants on WMT English–German newstest2014.The comparison includes different discretization bottlenecks and reports decoding speed alongside translation quality.
  • Latent utilization: DVQ with nd = 2 uses the available discrete latent space more evenly than vanilla VQ-VAE, which suffers from index collapse.The histogram comparison shows only a few latents used for vanilla VQ-VAE, versus diverse latent usage for sliced DVQ with nd = 2.
  • Latent utilization: The optimal number of DVQ decompositions is nd = 2 for latent vocabulary sizes log2 K = 14 and 16.nd = 1 performs noticeably worse, while higher nd values reduce performance, possibly because each decomposition has less expressive power.
  • Speed–quality trade-off: Increasing the n/m compression ratio improves parallelism and decoding speed but requires latents to encode more information for parallel reconstruction.The Latent Transformer therefore exposes a tunable trade-off between shorter latent sequences and reconstruction demands.
  • Speed–quality trade-off: Reconstruction improves when latent states contain more bits or compress fewer subword units.This relationship is measured using the perplexity of reconstructed outputs relative to the original sequence.

7. Conclusions

The paper targets slow decoding in sequence models despite improved training parallelism. It uses latent representations to support faster decoding, achieving an order-of-magnitude speedup while not matching autoregressive performance and identifying further speed and accuracy improvements as future work.

  • Problem: Autoregressive models remain slow during decoding because generating each output requires the preceding symbols, even when training is parallelized.At least n sequential steps are needed to generate a sequence of length n.
  • Scope: The paper presents fast decoding for sequence models as a general goal, while evaluating the approach primarily on machine translation.The conclusion frames the method as groundwork beyond the immediate translation application.
  • Discretization: Only 5% of the available 2^16 discrete latents are used by vanilla VQ-VAE, with latent usage maximized at nd = 2.This conclusion is based on measurements after 500,000 training steps with D = 512.
  • Main outcome: The Latent Transformer is an order of magnitude faster than the autoregressive model but does not yet recover its full performance.It nevertheless performs better than a heavily hand-tuned, task-specific non-autoregressive model.
  • Future work: Future work proposes hierarchical latent generation for greater speed and sampling or partial autoregression for improved BLEU scores.The latent sequence is still generated autoregressively, which takes most of the time for longer sentences.
Loading 1803.03382v6…