Source-linked AI summary
Blockwise Parallel Decoding for Deep Autoregressive Models
Mitchell Stern, Noam Shazeer, Jakob Uszkoreit
TL;DR
Autoregressive models remain sequential at generation time even when their architectures support parallel computation. The paper introduces blockwise parallel decoding, which proposes several future tokens in parallel and validates the longest greedy-consistent prefix. Across machine translation and image super-resolution, it approximately doubles speed without quality loss, with larger iteration and wall-clock speedups available at modest quality cost.
Problem
Autoregressive sequence-to-sequence models generate outputs one token at a time during inference despite architectures designed for parallel computation.
Method
Blockwise parallel decoding uses auxiliary models to propose multiple future tokens in parallel, then has the base model verify and accept the longest greedy-consistent prefix.
Results
Approximately 2x generation speed is achieved with no quality loss relative to greedy decoding, while combined strategies reach up to five-fold faster iterations for translation and seven-fold for super-resolution.
Takeaways & Limitations
The technique provides a simple, generic speedup for autoregressive models whose architectures can score output positions in parallel, with no or small quality losses demonstrated on two tasks.
Takeaways & Limitations
The method depends critically on architectures that can execute verification across output positions in parallel.
Abstract
from arXiv · showhide
Deep autoregressive sequence-to-sequence models have demonstrated impressive performance across a wide variety of tasks in recent years. While common architecture classes such as recurrent, convolutional, and self-attention networks make different trade-offs between the amount of computation needed per layer and the length of the critical path at training time, generation still remains an inherently sequential process. To overcome this limitation, we propose a novel blockwise parallel decoding scheme in which we make predictions for multiple time steps in parallel then back off to the longest prefix validated by a scoring model. This allows for substantial theoretical improvements in generation speed when applied to architectures that can process output sequences in parallel. We verify our approach empirically through a series of experiments using state-of-the-art self-attention models for machine translation and image super-resolution, achieving iteration reductions of up to 2x over a baseline greedy decoder with no loss in quality, or up to 7x in exchange for a slight decrease in performance. In terms of wall-clock time, our fastest models exhibit real-time speedups of up to 4x over standard greedy decoding.
1 Introduction
Autoregressive models generate outputs sequentially despite architectures that support parallel computation. Blockwise parallel decoding uses parallel proposals and base-model verification to accelerate generation, with speed gains demonstrated across tasks.
- Autoregressive sequence-to-sequence models generate outputs one token at a time during inference, creating a practical challenge despite faster parallel training architectures.
- Blockwise parallel decoding predicts multiple future positions independently, then scores them in parallel and retains the longest prefix matching greedy decoding.If the validated prefix exceeds one token, the greedy decoding loop skips iterations.
- Approximately 2x generation speed is achieved with no quality loss relative to greedy decoding.
- Combining the technique with knowledge distillation and approximate decoding reaches up to five-fold faster decoding iterations for machine translation and seven-fold for image super-resolution, with modest quality sacrifice.The corresponding wall-clock speedups are three-fold and four-fold, respectively.
- The approach can be implemented on existing models with minimal modifications.The authors state that code is publicly available in Tensor2Tensor.
2 Greedy Decoding
The paper formulates sequence-to-sequence inference as maximizing an autoregressive scoring model, then approximates this intractable search with left-to-right greedy decoding.
- The scoring model decomposes output likelihood left to right into conditional next-token probabilities.
- Sequence-to-sequence inference seeks an output sequence y maximizing the conditional score p(y | x).
- Greedy decoding repeatedly appends the highest-scoring next token conditioned on the input and the already generated prefix.
- Greedy decoding stops at an end-of-sequence token for language generation and after a fixed number of steps for image generation.
3 Blockwise Parallel Decoding
Blockwise parallel decoding proposes a multi-token extension, verifies its longest greedy-consistent prefix with the base model, and accepts only that prefix, preserving greedy output while reducing sequential work.
- Standard greedy decoding takes m sequential steps for an output of length m, motivating auxiliary models that propose multi-token extensions.
- The algorithm learns auxiliary models p2, . . . , pk to predict future positions beyond the base model’s next-token prediction.
- Predict independently generates k candidate tokens, verify finds the largest prefix matching the base model’s greedy predictions, and accept appends that prefix.
- Parallel prediction and verification are effective when the base architecture can process multiple output positions in parallel.
- Stopping acceptance at the first divergence guarantees recovery of the output produced by greedy decoding with the base model.
4 Combined Scoring and Proposal Model
A combined scoring-and-proposal model merges verification with the next prediction step, reducing model invocations while retaining parallel computation across positions.
- The basic Transformer implementation requires two model invocations per decoding step: one for prediction and one for verification.Even perfect auxiliary models reduce invocations from m to only 2m/k.
- Merging verification with the next prediction reduces invocations from 2m/k to m/k + 1.
- The combined model computes proposal probabilities for all future offsets in a constant number of operations, for example by expanding the final projection and using separate softmaxes.
- Figure 2 depicts one model call per iteration rather than two, halving the number of decoding model invocations.
- After verification computes the accepted prefix, the same outputs already supply the next iteration’s predictions, except for the first iteration.
5 Approximate Inference
Approximate block parallel decoding relaxes verification to gain additional speed, allowing outputs to diverge from standard greedy decoding. Minimum block-size constraints can guarantee that each step advances by multiple tokens.
- Approximate verification: Relaxing verification beyond exact greedy agreement enables additional speedups at the cost of potentially different outputs.The exact criterion reproduces standard greedy decoding; approximate criteria trade output fidelity for speed.
- Approximate verification: Top-k verification accepts a predicted item when it lies among the scoring model’s k highest-scoring items.
- Approximate verification: Natural distance metrics can replace exact agreement with approximate matching in suitable output spaces.For image generation, the distance is the absolute intensity difference within a color channel.
- Minimum block size: A non-greedy error can reduce a decoding step to adding only one token, motivating a minimum number of accepted tokens per step.
- Model modification: The modified Transformer produces k predictions through a multi-output feedforward layer with residual connections, then applies the original vocabulary projection to every output.
- Minimum block size: Setting ℓ = k enforces fixed-size parallel blocks, while requiring 1 < ℓ ≤ k guarantees at least ℓ tokens are added per step.
6 Implementation and Training
The implementation adds a multi-output feedforward layer to a pretrained Transformer and computes all proposal logits in parallel. Training uses a memory-conscious unbiased loss estimate, while experiments compare frozen and fine-tuned base parameters and investigate sequence-level distillation.
- Implementation: A single feedforward layer expands each decoder output into k predictions before the original projection produces logits for p1, ..., pk.The added layer has hidden size k × dhidden and output size k × dmodel, with residual connections to each output.
- Implementation: At inference time, all k logits are computed in parallel with marginal cost relative to the base model.
- Training: Randomly selecting one of the k cross-entropy sub-losses per minibatch provides an unbiased estimate of the full loss under training-time memory constraints.
- Training: Freezing pretrained parameters preserves the original model’s quality, whereas fine-tuning may improve internal consistency but can reduce final performance.
- Training: Sequence-level distillation is investigated because it can produce more predictable training sequences for blockwise parallel decoding.
7 Experiments
Experiments on machine translation and image super-resolution show that blockwise decoding reduces decoding iterations while offering quality–speed trade-offs. Wall-clock speedups reach 3.3x for translation and 4.0x for super-resolution, with translation quality remaining close to baseline in test-set comparisons.
- Experimental Setup: The experiments evaluate English-German translation on WMT 2014 and 32 × 32 image generation from 8 × 8 inputs on CelebA.Translation uses a Transformer baseline; super-resolution uses an Image Transformer with local attention.
- Experimental Setup: Combined scoring and proposal Transformer models are trained for various block sizes k, with additional runs using distilled data and fine-tuned baseline parameters.The translation models receive an additional 1,000,000 training steps, while super-resolution models use warm-started parameters with or without fine tuning.
- Machine Translation: 4.95 mean accepted block size is reached for translation, only 0.81 BLEU points below the initial distilled-data model.Distilled data also raises BLEU by 0.43 at mean block size 1.91, while fine tuning increases block size at some quality cost.
- Machine Translation: Distilled-data settings provide a smooth BLEU–block-size frontier, allowing selection of highest quality, fastest speed, or an intermediate operating point.Quality degradation at larger k is less pronounced with distilled data.
- Machine Translation: Top-2 approximate decoding reaches mean accepted block size 5.67 at k = 10, while larger block sizes trade additional iteration reductions for lower BLEU.Top-3 decoding follows a similar trend; minimum block-size constraints produce larger BLEU drops with only minor block-size gains.
- Image Super-Resolution: Human evaluation finds preference percentages close to 50%, indicating little perceived quality difference between baseline and fine-tuned blockwise outputs.Fine-tuned approximate decoding with k = 6 receives the highest overall preference, with slightly more noise and variation than smoothed baseline outputs.
- Wall-Clock Speedup: 4.0x wall-clock speedup is achieved for super-resolution and 3.3x for translation, although larger k eventually reduces wall-clock gains because of higher computational cost.The corresponding mean accepted block sizes are 5.3 for super-resolution and 4.7 for translation.
8 Conclusion
The paper presents blockwise parallel decoding as a simple, generic technique for accelerating deep autoregressive models, with no loss or only small losses in quality. Future work will examine combinations with potentially orthogonal approaches.
- 8 Conclusion: Blockwise parallel decoding improves decoding speed in machine translation and conditional image generation with no loss or only small losses in quality.The technique targets autoregressive models whose architectures support parallel scoring across output positions.
- 8 Conclusion: The technique is comparatively straightforward to add to existing models.
- 8 Conclusion: Combining blockwise parallel decoding with sequences of discrete latent variables remains planned future work.