PaperlayerView exact revision on arXiv

Source-linked AI summary

Speculative Decoding with Big Little Decoder

Sehoon Kim, Karttikeya Mangalam, Suhong Moon, Jitendra Malik, Michael W. Mahoney, Amir Gholami, Kurt Keutzer

arXiv:2302.07863v4cs.CL

TL;DR

Large language models deliver strong NLP performance but incur high inference latency, particularly for sequential autoregressive generation. BiLD combines a small autoregressive model with occasional non-autoregressive correction by a large model, coordinated through fallback and rollback policies. Across translation and summarization scenarios, it reaches up to 2.12× speedup with minimal quality degradation and requires no training-pipeline or architecture changes.

  • Problem

    Large language models have high inference latency, worsened by autoregressive generation that produces tokens sequentially without token-level parallelization.

  • Method

    BiLD uses a small model for most autoregressive generation and occasionally invokes a large model non-autoregressively to refine inaccurate predictions through fallback and rollback policies.

  • Results

    2.12× speedup was achieved on an NVIDIA T4 GPU across varied text-generation scenarios while allowing approximately one point of generation-quality degradation.

  • Takeaways & Limitations

    BiLD provides a plug-and-play runtime approach for reducing text-generation latency without modifying existing training pipelines or model architectures.

  • Takeaways & Limitations

    BiLD assumes compatible vocabularies, and separately trained models may disagree at the vocabulary level, causing unnecessary rollbacks without improving generation quality.

Abstract

from arXiv · show

The recent emergence of Large Language Models based on the Transformer architecture has enabled dramatic advancements in the field of Natural Language Processing. However, these models have long inference latency, which limits their deployment and makes them prohibitively expensive for various real-time applications. The inference latency is further exacerbated by autoregressive generative tasks, as models need to run iteratively to generate tokens sequentially without leveraging token-level parallelization. To address this, we propose Big Little Decoder (BiLD), a framework that can improve inference efficiency and latency for a wide range of text generation applications. The BiLD framework contains two models with different sizes that collaboratively generate text. The small model runs autoregressively to generate text with a low inference cost, and the large model is only invoked occasionally to refine the small model's inaccurate predictions in a non-autoregressive manner. To coordinate the small and large models, BiLD introduces two simple yet effective policies: (1) the fallback policy that determines when to hand control over to the large model; and (2) the rollback policy that determines when the large model needs to correct the small model's inaccurate predictions. To evaluate our framework across different tasks and models, we apply BiLD to various text generation scenarios encompassing machine translation on IWSLT 2017 De-En and WMT 2014 De-En, and summarization on XSUM and CNN/DailyMail. On an NVIDIA T4 GPU, our framework achieves a speedup of up to 2.12x speedup with minimal generation quality degradation. Furthermore, our framework is fully plug-and-play and can be applied without any modifications in the training process or model architecture. Our code is open-sourced

1 Introduction

BiLD addresses high Transformer inference latency by combining a small autoregressive decoder with occasional parallel refinement by a large decoder. It coordinates this process through fallback and rollback policies and achieves substantial speedups across translation and summarization tasks.

  • Motivation: Large Transformer models face high inference cost and latency, especially for autoregressive generation that processes tokens sequentially and limits hardware utilization.Autoregressive decoding repeatedly loads model weights and cached keys and values, making inference memory-bandwidth constrained.
  • Motivation: Non-autoregressive decoding improves token-level parallelization but can reduce generation quality because it assumes conditional independence between output tokens.
  • BiLD Framework: BiLD coordinates small and large decoder models: the small model generates most text autoregressively, while the large model occasionally refines inaccurate predictions non-autoregressively.
  • BiLD Policies: BiLD uses fallback to transfer control when the small model lacks confidence and rollback to let the large model review and correct inaccurate predictions.
  • Evaluation: 1.85× speedup was achieved without generation quality degradation, while 2.12× speedup allowed approximately one point of degradation on an NVIDIA T4 GPU across four text-generation scenarios.

2 Related Work

Prior work accelerates Transformer inference through architectural changes, compression, shallow decoders, non-autoregressive decoding, or multiple-model coordination. BiLD differs by targeting decoding parallelism with a plug-and-play runtime framework that avoids training-pipeline and model-architecture modifications.

  • Efficient Transformer Decoding Inference: Efficient Transformer inference has been pursued through efficient architectures, quantization, pruning, neural architecture search, and specialized decoding mechanisms.
  • Efficient Transformer Decoding Inference: Non-autoregressive and multi-iteration decoding generate multiple tokens or refine predictions in parallel, but often require complex, task-dependent training strategies or auxiliary information.
  • Efficient Transformer Decoding Inference: Shallow-decoder methods reduce decoding computation, whereas BiLD improves decoding parallelism without modifying model structures or requiring training changes.
  • Use of Multiple Models: Knowledge distillation uses large models during training to improve smaller models, while BiLD coordinates models at runtime during decoding without adding training complexity.
  • Use of Multiple Models: Ensemble learning combines independently trained models at runtime, but generally requires running multiple models and can be computationally expensive and poorly optimized for latency.
  • Related Speculative Decoding: Compared with speculative-sampling approaches, BiLD reports better latency-performance trade-offs attributed to non-random rollback and dynamic fallback window size.

3 Methodology

BiLD accelerates autoregressive text generation by letting a small model generate most tokens and occasionally invoking a large model non-autoregressively to correct uncertain predictions. Its fallback and rollback policies coordinate model handoffs and corrections, while prediction alignment addresses vocabulary-level disagreements between independently trained models.

  • Motivating Examples: ∼20% substitution of incorrect small-model predictions can yield comparable or better generation quality than the large model.This trend is reported for WMT 2014 De-En translation and CNN/DailyMail summarization, including models roughly 10× smaller than the large models.
  • Problem Formulation: BiLD hands control to the large model when a lightweight policy estimates that the small model is likely to make an inaccurate prediction.The policy chooses between models at each decoding iteration rather than running the large model at every step.
  • Latency Reduction: Large-model non-autoregressive decoding improves hardware efficiency by processing multiple tokens in one memory operation, despite comparable FLOPs to sequential decoding.Token-level parallelism and increased arithmetic intensity can offset the small model’s additional cost when the large model is invoked selectively.
  • Fallback Policy: The fallback policy switches models when the small model’s maximum token probability falls below α_FB.The fallback decision uses max_y pS(y|y1:n−1) as a confidence measure and sets the next token to the large model’s prediction.
  • Rollback Policy: The rollback policy detects disagreement at the earliest step m exceeding α_RB, replaces that token with the large model’s prediction, and discards subsequent dependent tokens.The distance metric compares the small and large models’ probability distributions; rollback adds computation but can improve generation quality.
  • Model Prediction Alignment: BiLD permits independently selected small and large models that share a vocabulary, but vocabulary-level differences can cause unnecessary rollbacks without quality gains.Model prediction alignment uses a calibration dataset and large-model outputs to align predictions between separately trained models.

4 Evaluations

BiLD evaluations span translation, summarization, ablations, and early exiting, showing latency-quality trade-offs that depend on its fallback, rollback, and alignment mechanisms.

  • Evaluation Setup: BiLD was evaluated on IWSLT 2017 De-En, WMT 2014 De-En, XSUM, and CNN/DailyMail using large and small models.The translation models were mT5-large/small, while summarization used T5-large/small.
  • Main Results: 1.50× average speedup was achieved across benchmarks with baseline small models, reaching 1.71× on CNN/DailyMail without generation-quality degradation.Allowing approximately 1 point degradation increased average speedup to 1.70× and the maximum to 2.05×.
  • Main Results: 1.61× average speedup was achieved with aligned BiLD, reaching 1.85× with minimal degradation and 2.12× within approximately 1 point degradation.Aligned BiLD consistently improved over unaligned BiLD, and both variants exceeded baseline BLEU/ROUGE-L scores in high-latency regimes.
  • Main Results: BiLD performed better than speculative sampling on IWSLT 2017 De-En and XSUM under the reported comparison settings.On IWSLT, speculative sampling scored 39.93 BLEU at 1.28× speedup, while unaligned BiLD achieved a 0.61-point higher BLEU score.
  • Ablation Studies: Removing either rollback or fallback caused significant performance degradation, showing that both policies are critical components of BiLD.Rollback improved quality across latency regimes, while fixed-period fallback degraded performance compared with the adaptive policy.
  • Early Exiting Strategy: BiLD also supports early exiting by using a model’s first layer as the small model, achieving up to 1.60× and 1.74× speedup within less than one BLEU point drop.These results were obtained on IWSLT 2017 De-En and WMT 2014 De-En, respectively.

5 Conclusion

BiLD reduces end-to-end latency by combining autoregressive generation from a small decoder with occasional non-autoregressive refinement by a large decoder. Its fallback and rollback policies coordinate when the large model takes control and corrects inaccurate predictions.

  • Conclusion: BiLD reduces latency across machine translation, summarization, and language modeling without requiring training or modifications to existing models.The framework couples large and small decoder models for more efficient text generation.
  • Conclusion: The small model generates most tokens autoregressively, while the large model is invoked non-autoregressively to refine inaccurate predictions.The fallback policy hands over control when the small model is uncertain, and rollback lets the large model revert inaccurate predictions.
  • Conclusion: 1.52× average speedup was achieved without performance drop on an NVIDIA Titan Xp GPU, with improvements up to 2.18× on some tasks.The conclusion also reports results when a 1-point performance degradation was allowed.

7.1 Experimental Details

The experiments compare model configurations and training procedures across translation and summarization benchmarks, while also reporting an early-exit comparison with CALM.

  • Models and Datasets: Translation uses IWSLT 2017 De-En and WMT 2014 De-En with 8-layer mT5-small and 24-layer mT5-large models.Summarization uses XSUM and CNN/DailyMail with 6-layer T5-small and 24-layer T5-large models.
  • Training: Aligned small models are fine-tuned on calibration datasets generated by fully trained large models, using the same recipes and steps as baseline small models.The calibration outputs are generated through greedy sampling with beam size 1.
  • Inference: Inference evaluations use batch size 1 on a single NVIDIA T4 GPU in a GCP n1-standard-4 instance.The rollback distance metric uses cross-entropy between the small model’s hard label and the large model’s soft label.
  • Early-Exit Comparison: BiLD outperforms CALM across early-exit speedup regimes by up to approximately 2–2.5 BLEU points on IWSLT 2017 De-En and WMT 2014 De-En.The comparison reports latency-quality trade-off curves for both benchmarks.

7.2 Details of Early Exiting Strategy in the BiLD Framework

BiLD’s early-exiting evaluation uses a jointly fine-tuned mT5 model and compares its BLEU–latency trade-off with CALM. BiLD achieves higher BLEU at matched speedups, aided by rollback correction and more stable cache computation.

  • Training and evaluation details: BiLD fine-tunes mT5-small with the first and final layers’ averaged loss to align early-exit and final predictions.The prediction head is shared, and evaluation uses fallback and rollback thresholds.
  • Performance comparison with CALM: BiLD achieves up to around 2 points better BLEU than CALM in the approximately 1.5× speedup regime.The comparison covers both machine-translation tasks in the early-exiting setting.
  • Performance comparison with CALM: Rollback corrects incorrect early-exited predictions, reducing the propagation of errors into future predictions.Incorrect predictions from the smaller model can be replaced using the rollback policy.
  • Performance comparison with CALM: Filling skipped-layer key and value caches with actual values improves decoding stability by reducing error propagation.This avoids computing skipped-layer caches from the exiting layer’s hidden states.

7.3 Comparison with Other Speculative Decoding Frameworks

BiLD differs from rejection-sampling speculative decoding through non-random rollback, dynamic fallback windows, and model alignment. Across IWSLT and XSUM, it consistently improves quality at matched latency or latency at matched quality.

  • Methodological differences: BiLD’s rollback rejects small-model predictions using distance from the large model, rather than stochastic rejection sampling.The approach uses a non-random decision based on the disagreement between model predictions.
  • Methodological differences: BiLD dynamically adjusts its fallback window at runtime using the small model’s confidence, unlike methods with fixed window sizes.The ablation study evaluates the contribution of this dynamic window.
  • Methodological differences: BiLD adds a model-alignment method that aligns small-model predictions with the large model to reduce unnecessary rejections.The alignment method requires minimal training-pipeline adjustments.
  • Quantitative comparisons: Across IWSLT and XSUM, BiLD consistently achieves better BLEU/ROUGE-L at equivalent latency gains or better latency at matched scores.The comparison uses aligned latency- and quality-matched configurations against rejection-sampling methods.
  • Quantitative comparisons: BiLD’s comparison tracks fallback and rollback percentages, with fewer occurrences associated with its better latency speedup.Table 4 compares these rates while reporting better BLEU/ROUGE-L scores.
  • Quantitative comparisons: BiLD’s quantitative analysis reports simultaneous improvements in text-generation quality and latency over other speculative-decoding frameworks.The paper provides empirical comparisons and explanations for this outcome.

(1) Better text generation quality

BiLD improves generation quality by combining model outputs and using hard, distance-based rollback decisions. Its dynamic fallback window also reduces unnecessary interactions with the large model, lowering computation and latency.

  • Better text generation quality: Blending the large model’s prediction with the small model’s prediction can improve generation quality through complementary model outputs.The paper illustrates this ensembling effect with a 20% contribution from the large model.
  • Better text generation quality: BiLD’s hard rollback discards small-model predictions that significantly deviate from the large model according to a cross-entropy distance.This contrasts with stochastic rejection sampling and is empirically associated with better BLEU/ROUGE-L using fewer rollbacks.
  • Lower end-to-end latency: BiLD’s dynamic fallback window is longer when the small model is confident and preempts it when uncertainty is high.This policy differs from speculative decoding’s static window size.
  • Lower end-to-end latency: Minimizing fallbacks and rollbacks reduces unnecessary computation and directly improves end-to-end latency.The two policies jointly limit unnecessary engagement and correction work.

7.4 BiLD with Sampling

BiLD extends to sampling by replacing greedy token selection with random sampling while retaining the same fallback and rollback policies. On IWSLT and XSUM, it preserves the same quality–speedup pattern observed with greedy decoding.

  • BiLD with sampling: BiLD supports sampling by randomly sampling tokens from both models while keeping the fallback and rollback policies unchanged.The policies use maximum prediction probability as an uncertainty indicator.

7.5 Additional Analysis

BiLD reduces memory operations and increases arithmetic intensity by processing multiple tokens per weight load, producing an actual hardware speedup. Its fallback and rollback thresholds expose a latency–quality trade-off across summarization tasks.

  • Model Analysis: ∼5× lower memory operations and approximately 5× higher arithmetic intensity than vanilla inference were observed on CNN/DailyMail.BiLD processes multiple tokens with one weight load, improving token-level parallelism and data reuse.
  • Model Analysis: 1.85× latency speedup was achieved on actual hardware for CNN/DailyMail.
  • Qualitative Analysis: BiLD combines a small model's output with occasional large-model tokens to correct inaccurate vocabulary and semantics.
  • Threshold Analysis: Smaller rollback thresholds improve generation quality but increase latency because more small-model predictions are preempted.
  • Threshold Analysis: Smaller fallback thresholds reduce latency but worsen generation quality, reflecting a second controllable latency–quality trade-off.
Loading 2302.07863v4…