Accelerating Transformer Inference Through Speculative Decoding

Summary & Methodology Analysis

Inference from large autoregressive models, which are neural network architectures that process sequences using attention mechanisms to weigh context, like Transformers, is slow because decoding K tokens takes K serial runs of the model. To solve this bottleneck, the authors introduce a novel approach. First, completions are generated autoregressively using a more efficient approximation model called M_q. Then, all guesses and their respective probabilities from the approximation model are evaluated in parallel using the target model called M_p. A novel sampling method, termed speculative sampling, accepts all guesses that can lead to an identical distribution. Finally, an additional token is sampled from an adjusted distribution to fix the first rejected token or to add an additional one if all guesses are accepted. This architecture allows the system to process multiple candidate tokens simultaneously while preserving the output distribution of the target model.

The practical results show clear performance gains across translation tasks. The paper demonstrated a 2X-3X acceleration on T5-XXL compared to the standard T5X implementation with identical outputs. Furthermore, it achieved speedups of 2.6X at temperature 1 and 3.4X at temperature 0 on the translation task using T5-small as the approximation model. The evaluation leverages multiple models and datasets, including T5-XXL, T5X, GPT-3, LaMDA, Parti, PaLM, T5 version 1.1, WMT EnDe, CCN/DM, T5-large, T5-base, T5-small, Chinchilla 70B, lm1b, LaMDA 8B, LaMDA 2B, and LaMDA 100M. These figures highlight how parallelism can be exploited during generation phases when compute resources are appropriately allocated.

Despite the significant speedups, the method has clear resource constraints and limitations. Speculative decoding assumes that there are enough compute resources available to support increased concurrency from parallel evaluations. Consequently, the method is not helpful for configurations where additional computation resources are not available. The paper does not specify alternative strategies for memory-constrained environments where running both an approximation model and a target model concurrently exceeds hardware limits. Engineers considering this technique must ensure their serving infrastructure has sufficient headroom for parallel execution.